Track1_BugBusters/database.py at main · VectorFlow-vvce/Track1_BugBusters · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""
SQLAlchemy async database setup for the Predictive Routing Engine.
"""

from __future__ import annotations

from datetime import datetime, timezone

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

import config

engine = create_async_engine(config.DATABASE_URL, echo=False)
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

metadata = sa.MetaData()

incidents_table = sa.Table(
    "incidents",
    metadata,
    sa.Column("id", sa.String, primary_key=True),
    sa.Column("created_at", sa.DateTime, nullable=False),
    sa.Column("status", sa.String, nullable=False, default="active"),
    sa.Column("degraded_mode", sa.Boolean, nullable=False, default=False),
)

routes_table = sa.Table(
    "routes",
    metadata,
    sa.Column("id", sa.String, primary_key=True),
    sa.Column("incident_id", sa.String, sa.ForeignKey("incidents.id"), nullable=False),
    sa.Column("destination_hospital_id", sa.String, nullable=False),
    sa.Column("eta_seconds", sa.Float, nullable=False),
    sa.Column("diversion_override", sa.Boolean, nullable=False, default=False),
    sa.Column("computed_at", sa.DateTime, nullable=False),
    sa.Column("traffic_snapshot_age_s", sa.Float, nullable=False),
    sa.Column("segments_json", sa.Text, nullable=False),
    sa.Column("traffic_snapshot_json", sa.Text, nullable=True),
    sa.Column("er_snapshot_json", sa.Text, nullable=True),
)

route_events_table = sa.Table(
    "route_events",
    metadata,
    sa.Column("id", sa.String, primary_key=True),
    sa.Column("incident_id", sa.String, sa.ForeignKey("incidents.id"), nullable=False),
    sa.Column("previous_route_id", sa.String, nullable=False),
    sa.Column("new_route_id", sa.String, nullable=False),
    sa.Column("reason", sa.String, nullable=False),
    sa.Column("replaced_at", sa.DateTime, nullable=False),
)

driver_acks_table = sa.Table(
    "driver_acks",
    metadata,
    sa.Column("incident_id", sa.String, sa.ForeignKey("incidents.id"), primary_key=True),
    sa.Column("route_id", sa.String, primary_key=True),
    sa.Column("acked_at", sa.DateTime, nullable=False),
)

audit_logs_table = sa.Table(
    "audit_logs",
    metadata,
    sa.Column("id", sa.String, primary_key=True),
    sa.Column("incident_id", sa.String, nullable=False),
    sa.Column("event_type", sa.String, nullable=False),
    sa.Column("payload_json", sa.Text, nullable=False),
    sa.Column("created_at", sa.DateTime, nullable=False),
)


async def init_db() -> None:
    """Create all tables if they don't exist."""
    async with engine.begin() as conn:
        await conn.run_sync(metadata.create_all)


async def get_db() -> AsyncSession:  # type: ignore[misc]
    """Async generator yielding a database session for dependency injection."""
    async with async_session() as session:
        yield session