Database Migrations

AuroraSOC uses Alembic for PostgreSQL schema migrations, integrated with SQLAlchemy 2.0 async models.

Setup

Alembic Configuration (alembic.ini)

[alembic]
script_location = alembic
sqlalchemy.url = postgresql+asyncpg://aurora:aurora@localhost:5432/aurorasoc

Migration Environment (alembic/env.py)

from aurorasoc.core.models import Base

target_metadata = Base.metadata

def run_migrations_online():
    connectable = create_async_engine(config.get_main_option("sqlalchemy.url"))
    
    async def do_run():
        async with connectable.connect() as connection:
            await connection.run_sync(do_run_migrations)
    
    asyncio.run(do_run())

Initial Schema (001_initial_schema.py)

The initial migration creates all 11 tables:

PostgreSQL Extensions

CREATE EXTENSION IF NOT EXISTS "uuid-ossp";   -- UUID generation
CREATE EXTENSION IF NOT EXISTS "pg_trgm";     -- Trigram similarity for IOC search

note

Migration 004 adds vector (pgvector) and pgcrypto extensions. See below.

Key Indexes

Table	Index	Type	Purpose
alerts	idx_alerts_severity	B-tree	Filter by severity
alerts	idx_alerts_status	B-tree	Filter by status
alerts	idx_alerts_created_at	B-tree	Time-range queries
alerts	idx_alerts_dedup_hash	B-tree	Deduplication lookup
cases	idx_cases_status	B-tree	Filter by status
cases	idx_cases_severity	B-tree	Filter by severity
cps_devices	idx_cps_device_id	B-tree	Device lookup
cps_devices	idx_cps_attestation_status	B-tree	Filter by attestation
iocs	idx_iocs_type_value	B-tree (unique)	Dedup + fast lookup
human_approvals	idx_approvals_status	B-tree	Pending approval filter

Auto-Update Trigger

CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = NOW();
    RETURN NEW;
END;
$$ language 'plpgsql';

-- Applied to all tables with updated_at columns
CREATE TRIGGER update_alerts_updated_at
    BEFORE UPDATE ON alerts
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

Schema Hardening (004_harden_schema_enums_fks_pgvector.py)

Migration 004 is a comprehensive hardening pass that introduces PostgreSQL enums, foreign keys, CHECK constraints, pgvector, and auto-update triggers across all tables.

New PostgreSQL Extensions

CREATE EXTENSION IF NOT EXISTS "vector";      -- pgvector for embedding storage & ANN search
CREATE EXTENSION IF NOT EXISTS "pgcrypto";    -- gen_random_uuid(), cryptographic functions

PostgreSQL Enums

Seven CREATE TYPE ... AS ENUM statements replace unconstrained VARCHAR columns with type-safe enums:

Enum Type	Values	Applied To
severity_enum	LOW, MEDIUM, HIGH, CRITICAL	alerts.severity, cases.severity
alert_status_enum	new, investigating, resolved, escalated, false_positive	alerts.status
case_status_enum	open, investigating, resolved, escalated, closed	cases.status
approval_status_enum	pending, approved, denied, expired	human_approvals.status
attestation_status_enum	unknown, trusted, untrusted, revoked	cps_devices.attestation_status, attestation_results.status
user_role_enum	admin, analyst, viewer, service	users.role
execution_status_enum	pending, running, completed, failed, cancelled	playbook_executions.status

The migration uses ALTER TABLE ... ALTER COLUMN ... TYPE enum_name USING column::enum_name to convert existing VARCHAR data in-place.

Foreign Key Constraints

ALTER TABLE alerts
    ADD CONSTRAINT fk_alerts_case_id FOREIGN KEY (case_id) REFERENCES cases(id);
ALTER TABLE case_timeline
    ADD CONSTRAINT fk_timeline_case_id FOREIGN KEY (case_id) REFERENCES cases(id);
ALTER TABLE human_approvals
    ADD CONSTRAINT fk_approvals_case_id FOREIGN KEY (case_id) REFERENCES cases(id);
ALTER TABLE playbook_executions
    ADD CONSTRAINT fk_executions_case_id FOREIGN KEY (case_id) REFERENCES cases(id);
ALTER TABLE reports
    ADD CONSTRAINT fk_reports_case_id FOREIGN KEY (case_id) REFERENCES cases(id);
ALTER TABLE attestation_results
    ADD CONSTRAINT fk_attestation_device_id FOREIGN KEY (device_id) REFERENCES cps_devices(id);

CHECK Constraints

ALTER TABLE iocs ADD CONSTRAINT chk_iocs_confidence
    CHECK (confidence >= 0.0 AND confidence <= 1.0);
ALTER TABLE cps_devices ADD CONSTRAINT chk_cps_risk_score
    CHECK (risk_score >= 0.0 AND risk_score <= 100.0);
ALTER TABLE iocs ADD CONSTRAINT chk_iocs_risk_score
    CHECK (risk_score >= 0.0 AND risk_score <= 100.0);

Vector Embeddings Table (pgvector)

CREATE TABLE IF NOT EXISTS vector_embeddings (
    id              UUID DEFAULT gen_random_uuid() PRIMARY KEY,
    collection      VARCHAR(100)  NOT NULL,
    source_id       VARCHAR(200)  NOT NULL,
    embedding       vector(384)   NOT NULL,
    payload         JSONB         NOT NULL DEFAULT '{}',
    created_at      TIMESTAMPTZ   NOT NULL DEFAULT now(),
    CONSTRAINT uq_vec_collection_source UNIQUE (collection, source_id)
);

CREATE INDEX IF NOT EXISTS idx_vec_emb_hnsw
    ON vector_embeddings
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 200);

• vector(384): Matches the 384-dimensional output of all-MiniLM-L6-v2
• HNSW index: m=16 (graph connectivity), ef_construction=200 (build-time search depth)
• vector_cosine_ops: Enables the <=> cosine distance operator
• UNIQUE (collection, source_id): Enables upsert via ON CONFLICT

See Database Models — vector_embeddings for usage details.

Auto-Update Triggers

Migration 004 applies updated_at triggers to all tables that have an updated_at column, using the same update_updated_at_column() function from migration 001.

Migration Commands

# Generate auto-migration from model changes
alembic revision --autogenerate -m "Description of change"

# Apply all pending migrations
alembic upgrade head

# Rollback one migration
alembic downgrade -1

# Rollback to specific revision
alembic downgrade 001

# Show current migration state
alembic current

# Show migration history
alembic history --verbose

Creating New Migrations

When you modify SQLAlchemy models in aurorasoc/core/models.py:

1. Make the model change:

class AlertModel(Base):
    # ... existing columns ...
    priority: Mapped[int] = mapped_column(Integer, default=3)  # NEW

2. Generate migration:

alembic revision --autogenerate -m "Add priority to alerts"

3. Review the generated migration (Alembic generates in alembic/versions/):

def upgrade():
    op.add_column('alerts', sa.Column('priority', sa.Integer(), nullable=True, default=3))

def downgrade():
    op.drop_column('alerts', 'priority')

4. Apply:

alembic upgrade head

Docker Integration

In Docker Compose, migrations run automatically on API startup:

# aurorasoc/core/database.py
async def init_db():
    engine = create_async_engine(settings.postgres.url)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

For production deployments, use explicit Alembic commands:

# docker-compose.yml
api:
  command: >
    sh -c "alembic upgrade head && uvicorn aurorasoc.api.main:app --host 0.0.0.0"

JSONB Migration Patterns

JSONB columns don't require migrations for schema changes within the JSON structure. Adding a new key to an IOC's context field just works:

# Before: {"source": "virustotal"}
# After:  {"source": "virustotal", "confidence": 0.9}
# No migration needed — JSONB is schema-less

This is why AuroraSOC uses JSONB for flexible fields like iocs, mitre_techniques, metadata, and context.