Skip to content

Commit

Permalink
Fix foreign key in lastalerttoincident for complex primary key in las…
Browse files Browse the repository at this point in the history 
…talert
  • Loading branch information
@VladimirFilonov
VladimirFilonov committed Nov 24, 2024
1 parent 3a3ad06 commit 6eae448
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 35 deletions.
2 changes: 1 addition & 1 deletion keep/api/models/db/alert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ class LastAlert(SQLModel, table=True):


class LastAlertToIncident(SQLModel, table=True):
tenant_id: str = Field(foreign_key="tenant.id", nullable=False)
tenant_id: str = Field(foreign_key="tenant.id", nullable=False, primary_key=True)
timestamp: datetime = Field(default_factory=datetime.utcnow)

fingerprint: str = Field(foreign_key="lastalert.fingerprint", primary_key=True)
Expand Down
82 changes: 48 additions & 34 deletions keep/api/models/db/migrations/versions/2024-11-13-22-48_bdae8684d0b4.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

migration_metadata = sa.MetaData()


def populate_db():
session = Session(op.get_bind())

Expand All @@ -32,8 +33,8 @@ def populate_db():
select
alert.fingerprint, max(alert.timestamp) as last_received
from alert
group by fingerprint
) as a ON alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received
group by fingerprint, tenant_id
) as a ON alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received and alert.tenant_id = a.tenant_id
on conflict
do nothing
"""
Expand All @@ -44,48 +45,58 @@ def populate_db():
from alerttoincident as ati
join
(
select alert.id, alert.fingerprint
select alert.tenant_id, alert.id, alert.fingerprint
from alert
join (
select
alert.fingerprint, max(alert.timestamp) as last_received
alert.tenant_id, alert.fingerprint, max(alert.timestamp) as last_received
from alert
group by fingerprint
) as a on alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received
group by fingerprint, tenant_id
) as a on alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received and alert.tenant_id = a.tenant_id
) as lf on ati.alert_id = lf.id
on conflict
do nothing
"""

else:
migrate_lastalert_query = """
replace into lastalert (tenant_id, fingerprint, alert_id, timestamp)
select alert.tenant_id, alert.fingerprint, alert.id as alert_id, alert.timestamp
INSERT INTO lastalert (tenant_id, fingerprint, alert_id, timestamp)
SELECT
grouped_alerts.tenant_id,
grouped_alerts.fingerprint,
MAX(grouped_alerts.alert_id) as alert_id, -- Using MAX to consistently pick one alert_id
grouped_alerts.timestamp
FROM (
select alert.tenant_id, alert.fingerprint, alert.id as alert_id, alert.timestamp
from alert
join (
select
alert.tenant_id, alert.fingerprint, max(alert.timestamp) as last_received
from alert
join (
select
alert.fingerprint, max(alert.timestamp) as last_received
from alert
group by fingerprint
) as a ON alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received;
"""
group by fingerprint, tenant_id
) as a ON alert.fingerprint = a.fingerprint
and alert.timestamp = a.last_received
and alert.tenant_id = a.tenant_id
) as grouped_alerts
GROUP BY grouped_alerts.tenant_id, grouped_alerts.fingerprint, grouped_alerts.timestamp;
"""

migrate_lastalerttoincident_query = """
replace into lastalerttoincident (incident_id, tenant_id, timestamp, fingerprint, is_created_by_ai, deleted_at)
select ati.incident_id, ati.tenant_id, ati.timestamp, lf.fingerprint, ati.is_created_by_ai, ati.deleted_at
from alerttoincident as ati
join
(
select alert.id, alert.fingerprint
REPLACE INTO lastalerttoincident (incident_id, tenant_id, timestamp, fingerprint, is_created_by_ai, deleted_at)
select ati.incident_id, ati.tenant_id, ati.timestamp, lf.fingerprint, ati.is_created_by_ai, ati.deleted_at
from alerttoincident as ati
join
(
select alert.id, alert.fingerprint, alert.tenant_id
from alert
join (
select
alert.tenant_id,alert.fingerprint, max(alert.timestamp) as last_received
from alert
join (
select
alert.fingerprint, max(alert.timestamp) as last_received
from alert
group by fingerprint
) as a on alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received
) as lf on ati.alert_id = lf.id
"""
group by fingerprint, tenant_id
) as a on alert.fingerprint = a.fingerprint and alert.timestamp = a.last_received and alert.tenant_id = a.tenant_id
) as lf on ati.alert_id = lf.id;
"""

session.execute(migrate_lastalert_query)
session.execute(migrate_lastalerttoincident_query)
Expand All @@ -102,12 +113,14 @@ def upgrade() -> None:
["alert_id"],
["alert.id"],
),
sa.PrimaryKeyConstraint("fingerprint"),
sa.PrimaryKeyConstraint("tenant_id", "fingerprint"),
)
with op.batch_alter_table("lastalert", schema=None) as batch_op:
batch_op.create_index(
batch_op.f("ix_lastalert_timestamp"), ["timestamp"], unique=False
)
# Add index for the fingerprint column that will be referenced by foreign key
batch_op.create_index("ix_lastalert_fingerprint", ["fingerprint"], unique=False)

op.create_table(
"lastalerttoincident",
Expand All @@ -122,22 +135,23 @@ def upgrade() -> None:
sa.Column("is_created_by_ai", sa.Boolean(), nullable=False),
sa.Column("deleted_at", sa.DateTime(), nullable=True),
sa.ForeignKeyConstraint(
["fingerprint"],
["lastalert.fingerprint"],
["tenant_id", "fingerprint"],
["lastalert.tenant_id", "lastalert.fingerprint"],
),
sa.ForeignKeyConstraint(["incident_id"], ["incident.id"], ondelete="CASCADE"),
sa.ForeignKeyConstraint(
["tenant_id"],
["tenant.id"],
),
sa.PrimaryKeyConstraint("incident_id", "fingerprint", "deleted_at"),
sa.PrimaryKeyConstraint("tenant_id", "incident_id", "fingerprint", "deleted_at"),
)

populate_db()


def downgrade() -> None:
op.drop_table("lastalerttoincident")
with op.batch_alter_table("lastalert", schema=None) as batch_op:
batch_op.drop_index(batch_op.f("ix_lastalert_timestamp"))

op.drop_table("lastalert")
op.drop_table("lastalert")

0 comments on commit 6eae448

Please sign in to comment.