7-Add ORM tables models

alembic/env.py

@@ -5,7 +5,7 @@
 
 from alembic import context
 
-from app.db.models import Base
+from app.db.models import Base, Vehicle, VehiclePosition
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -67,10 +67,14 @@ def run_migrations_online() -> None:
 
     with connectable.connect() as connection:
         context.configure(
-            connection=connection, target_metadata=target_metadata
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table_schema=target_metadata.schema,
+            include_schemas=True
         )
 
         with context.begin_transaction():
+            context.execute('SET search_path TO public')
             context.run_migrations()
 
 

alembic/versions/2025_04_29_1155-7076062a2ab9.py

@@ -0,0 +1,61 @@
+"""Create vehicles & vehicle_positions tables
+
+Revision ID: 7076062a2ab9
+Revises: 
+Create Date: 2025-04-29 11:55:10.699738
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = '7076062a2ab9'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Create enum type explicitly
+    op.execute("""
+            CREATE TYPE line_product_enum AS ENUM (
+                'BUS', 'SUBWAY', 'TRAMWAY', 'SUBURBAN', 'FERRY', 'EXPRESS', 'REGIONAL'
+            );
+        """)
+
+    # Create partitioned 'vehicles' table
+    op.execute("""
+               CREATE TABLE vehicles
+               (
+                   id           UUID              NOT NULL DEFAULT gen_random_uuid(),
+                   trip_id      TEXT              NOT NULL,
+                   line_product line_product_enum NOT NULL,
+                   line_name    TEXT              NOT NULL,
+                   partition_dt DATE              NOT NULL,
+                   PRIMARY KEY (id, partition_dt)
+               ) PARTITION BY RANGE (partition_dt);
+               """)
+    # Create partitioned 'vehicle_positions' table
+    op.execute("""
+               CREATE TABLE vehicle_positions
+               (
+                   vehicle_id   UUID            NOT NULL,
+                   timestamp    TIMESTAMP       NOT NULL,
+                   latitude     NUMERIC(38, 18) NOT NULL,
+                   longitude    NUMERIC(38, 18) NOT NULL,
+                   partition_dt DATE            NOT NULL,
+                   PRIMARY KEY (vehicle_id, timestamp, partition_dt),
+                   FOREIGN KEY (vehicle_id, partition_dt) REFERENCES vehicles (id, partition_dt) ON DELETE CASCADE
+               ) PARTITION BY RANGE (partition_dt);
+               """)
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.execute("DROP TABLE IF EXISTS vehicle_positions CASCADE;")
+    op.execute("DROP TABLE IF EXISTS vehicles CASCADE;")
+    op.execute("DROP TYPE IF EXISTS line_product_enum;")

app/db/models.py

@@ -1,3 +1,71 @@
-from sqlalchemy.ext.declarative import declarative_base
+import enum
+import uuid
+from datetime import date, datetime
+from typing import List
 
-Base = declarative_base()
+from sqlalchemy import DECIMAL, Date, Enum, ForeignKey, String, func
+from sqlalchemy.dialects.postgresql import TIMESTAMP, UUID
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class LineProductEnum(enum.Enum):
+    BUS = "bus"
+    SUBWAY = "subway"
+    TRAMWAY = "tram"
+    SUBURBAN = "suburban"  # S-Bahn
+    FERRY = "ferry"
+    EXPRESS = "express"  # IC/ICE trains
+    REGIONAL = "regional"  # Regio trains
+
+
+class Vehicle(Base):
+    """
+    Table representing a vehicle, partitioned by day.
+    A vehicle is represented by its trip id as the API doesn't provide the actual id of vehicles themselves, thus
+    we can have multiple trip_ids which represent the same real-world vehicle.
+    """
+    __tablename__ = "vehicles"
+    __table_args__ = (
+        {'postgresql_partition_by': 'RANGE (partition_dt)'}
+    )
+
+    id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True),
+        primary_key=True,
+        server_default=func.gen_random_uuid()
+    )
+    trip_id: Mapped[str] = mapped_column(String, nullable=False)
+    line_product: Mapped[LineProductEnum] = mapped_column(Enum(LineProductEnum, name="line_product_enum"), nullable=False)
+    line_name: Mapped[str] = mapped_column(String, nullable=False)
+    partition_dt: Mapped[date] = mapped_column(Date, primary_key=True)
+
+    positions: Mapped[List["VehiclePosition"]] = relationship(
+        back_populates="vehicle",
+        cascade="all, delete-orphan"
+    )
+
+
+class VehiclePosition(Base):
+    """
+    Table representing the position of a vehicle at a given time, partitioned by day.
+    """
+    __tablename__ = "vehicle_positions"
+    __table_args__ = (
+        {'postgresql_partition_by': 'RANGE (partition_dt)'},
+    )
+
+    vehicle_id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True),
+        ForeignKey('vehicles.id', ondelete='CASCADE'),
+        primary_key=True
+    )
+    timestamp: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=False), primary_key=True)
+    latitude: Mapped[float] = mapped_column(DECIMAL(38, 18), nullable=False)
+    longitude: Mapped[float] = mapped_column(DECIMAL(38, 18), nullable=False)
+    partition_dt: Mapped[date] = mapped_column(Date, primary_key=True)
+
+    vehicle: Mapped["Vehicle"] = relationship(back_populates="positions")

docker-compose.yaml

@@ -1,6 +1,8 @@
 services:
   postgres:
     image: postgres:15.6
+    ports:
+      - "5432:5432"
     environment:
       POSTGRES_USER: ${POSTGRES_USER:-postgres}
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-S3cr3t_P4ssw0rd}