main
py 119 lines 5.85 KB
Raw
1 """Add ai report tables
2
3 Revision ID: d0da7a384546
4 Revises: d000000000000
5 Create Date: 2026-04-07 11:01:35.704501
6
7 """
8 from typing import Sequence
9 from typing import Union
10
11 import sqlalchemy as sa
12 from sqlalchemy.dialects import mysql
13
14 from alembic import op
15
16 # revision identifiers, used by Alembic.
17 revision: str = "d0da7a384546"
18 down_revision: Union[str, None] = "d000000000000"
19 branch_labels: Union[str, Sequence[str], None] = None
20 depends_on: Union[str, Sequence[str], None] = None
21
22
23 def upgrade() -> None:
24 # ### commands auto generated by Alembic - please adjust! ###
25 op.create_table(
26 "ai_analyst_job",
27 sa.Column("error_message", sa.Text(), nullable=True),
28 sa.Column("id", sa.String(length=64), nullable=False),
29 sa.Column("alert_id", sa.Integer(), nullable=False),
30 sa.Column("customer_code", sa.String(length=64), nullable=False),
31 sa.Column("status", sa.String(length=50), nullable=False),
32 sa.Column("alert_type", sa.String(length=64), nullable=True),
33 sa.Column("triggered_by", sa.String(length=50), nullable=False),
34 sa.Column("template_used", sa.String(length=128), nullable=True),
35 sa.Column("created_at", sa.DateTime(), nullable=False),
36 sa.Column("started_at", sa.DateTime(), nullable=True),
37 sa.Column("completed_at", sa.DateTime(), nullable=True),
38 sa.ForeignKeyConstraint(
39 ["customer_code"],
40 ["customers.customer_code"],
41 ),
42 sa.PrimaryKeyConstraint("id"),
43 )
44 op.create_index(op.f("ix_ai_analyst_job_alert_id"), "ai_analyst_job", ["alert_id"], unique=False)
45 op.create_index(op.f("ix_ai_analyst_job_created_at"), "ai_analyst_job", ["created_at"], unique=False)
46 op.create_index(op.f("ix_ai_analyst_job_customer_code"), "ai_analyst_job", ["customer_code"], unique=False)
47 op.create_index(op.f("ix_ai_analyst_job_status"), "ai_analyst_job", ["status"], unique=False)
48 op.create_table(
49 "ai_analyst_report",
50 sa.Column("report_markdown", mysql.LONGTEXT(), nullable=True),
51 sa.Column("summary", sa.Text(), nullable=True),
52 sa.Column("recommended_actions", sa.Text(), nullable=True),
53 sa.Column("id", sa.Integer(), nullable=False),
54 sa.Column("job_id", sa.String(length=64), nullable=False),
55 sa.Column("alert_id", sa.Integer(), nullable=False),
56 sa.Column("customer_code", sa.String(length=64), nullable=False),
57 sa.Column("severity_assessment", sa.String(length=50), nullable=True),
58 sa.Column("created_at", sa.DateTime(), nullable=False),
59 sa.ForeignKeyConstraint(
60 ["customer_code"],
61 ["customers.customer_code"],
62 ),
63 sa.ForeignKeyConstraint(
64 ["job_id"],
65 ["ai_analyst_job.id"],
66 ),
67 sa.PrimaryKeyConstraint("id"),
68 )
69 op.create_index(op.f("ix_ai_analyst_report_alert_id"), "ai_analyst_report", ["alert_id"], unique=False)
70 op.create_index(op.f("ix_ai_analyst_report_created_at"), "ai_analyst_report", ["created_at"], unique=False)
71 op.create_index(op.f("ix_ai_analyst_report_customer_code"), "ai_analyst_report", ["customer_code"], unique=False)
72 op.create_index(op.f("ix_ai_analyst_report_job_id"), "ai_analyst_report", ["job_id"], unique=False)
73 op.create_table(
74 "ai_analyst_ioc",
75 sa.Column("details", sa.Text(), nullable=True),
76 sa.Column("id", sa.Integer(), nullable=False),
77 sa.Column("report_id", sa.Integer(), nullable=False),
78 sa.Column("alert_id", sa.Integer(), nullable=False),
79 sa.Column("customer_code", sa.String(length=64), nullable=False),
80 sa.Column("ioc_value", sa.String(length=512), nullable=False),
81 sa.Column("ioc_type", sa.String(length=50), nullable=False),
82 sa.Column("vt_verdict", sa.String(length=50), nullable=False),
83 sa.Column("vt_score", sa.String(length=32), nullable=True),
84 sa.Column("created_at", sa.DateTime(), nullable=False),
85 sa.ForeignKeyConstraint(
86 ["customer_code"],
87 ["customers.customer_code"],
88 ),
89 sa.ForeignKeyConstraint(
90 ["report_id"],
91 ["ai_analyst_report.id"],
92 ),
93 sa.PrimaryKeyConstraint("id"),
94 )
95 op.create_index(op.f("ix_ai_analyst_ioc_alert_id"), "ai_analyst_ioc", ["alert_id"], unique=False)
96 op.create_index(op.f("ix_ai_analyst_ioc_created_at"), "ai_analyst_ioc", ["created_at"], unique=False)
97 op.create_index(op.f("ix_ai_analyst_ioc_customer_code"), "ai_analyst_ioc", ["customer_code"], unique=False)
98 op.create_index(op.f("ix_ai_analyst_ioc_report_id"), "ai_analyst_ioc", ["report_id"], unique=False)
99 # ### end Alembic commands ###
100
101
102 def downgrade() -> None:
103 # ### commands auto generated by Alembic - please adjust! ###
104 op.drop_index(op.f("ix_ai_analyst_ioc_report_id"), table_name="ai_analyst_ioc")
105 op.drop_index(op.f("ix_ai_analyst_ioc_customer_code"), table_name="ai_analyst_ioc")
106 op.drop_index(op.f("ix_ai_analyst_ioc_created_at"), table_name="ai_analyst_ioc")
107 op.drop_index(op.f("ix_ai_analyst_ioc_alert_id"), table_name="ai_analyst_ioc")
108 op.drop_table("ai_analyst_ioc")
109 op.drop_index(op.f("ix_ai_analyst_report_job_id"), table_name="ai_analyst_report")
110 op.drop_index(op.f("ix_ai_analyst_report_customer_code"), table_name="ai_analyst_report")
111 op.drop_index(op.f("ix_ai_analyst_report_created_at"), table_name="ai_analyst_report")
112 op.drop_index(op.f("ix_ai_analyst_report_alert_id"), table_name="ai_analyst_report")
113 op.drop_table("ai_analyst_report")
114 op.drop_index(op.f("ix_ai_analyst_job_status"), table_name="ai_analyst_job")
115 op.drop_index(op.f("ix_ai_analyst_job_customer_code"), table_name="ai_analyst_job")
116 op.drop_index(op.f("ix_ai_analyst_job_created_at"), table_name="ai_analyst_job")
117 op.drop_index(op.f("ix_ai_analyst_job_alert_id"), table_name="ai_analyst_job")
118 op.drop_table("ai_analyst_job")
119 # ### end Alembic commands ###