59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
"""Phase 3: audit trail columns and data classification fields.
|
|
|
|
Revision ID: b029phase3
|
|
Revises: b028phase0
|
|
"""
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
|
|
revision: str = "b029phase3"
|
|
down_revision: Union[str, None] = "b028phase0"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def _column_names(table: str) -> set[str]:
|
|
bind = op.get_bind()
|
|
insp = sa.inspect(bind)
|
|
return {c["name"] for c in insp.get_columns(table)}
|
|
|
|
|
|
def upgrade() -> None:
|
|
audit_cols = _column_names("audit_logs")
|
|
if "ip_address" not in audit_cols:
|
|
op.add_column("audit_logs", sa.Column("ip_address", sa.String(45), nullable=True))
|
|
if "user_agent" not in audit_cols:
|
|
op.add_column("audit_logs", sa.Column("user_agent", sa.String(500), nullable=True))
|
|
if "integrity_hash" not in audit_cols:
|
|
op.add_column("audit_logs", sa.Column("integrity_hash", sa.String(64), nullable=True))
|
|
if "session_id" not in audit_cols:
|
|
op.add_column("audit_logs", sa.Column("session_id", sa.String(100), nullable=True))
|
|
|
|
for table in ("tests", "evidences", "campaigns"):
|
|
cols = _column_names(table)
|
|
if "data_classification" not in cols:
|
|
op.add_column(
|
|
table,
|
|
sa.Column(
|
|
"data_classification",
|
|
sa.String(20),
|
|
nullable=False,
|
|
server_default="internal",
|
|
),
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
for table in ("campaigns", "evidences", "tests"):
|
|
cols = _column_names(table)
|
|
if "data_classification" in cols:
|
|
op.drop_column(table, "data_classification")
|
|
|
|
audit_cols = _column_names("audit_logs")
|
|
for col in ("session_id", "integrity_hash", "user_agent", "ip_address"):
|
|
if col in audit_cols:
|
|
op.drop_column("audit_logs", col)
|