Shattering

.github/workflows/test-backend.yml

@@ -51,7 +51,7 @@ jobs:
         working-directory: backend
 
       - name: Run tests
-        run: pytest -v --cov=app
+        run: pytest -v --cov=.
         working-directory: backend
         env:
           DOMAIN: postgres

backend/README.md

@@ -87,7 +87,7 @@ If needed, create a user for yourself.
 
 ### Testing
 
-`pytest --cov=app --cov-report=html`
+`pytest --cov=.`
 
 Or with full coverage report:
 

backend/app/alembic/versions/ccb2a6b81a8b_shattering.py

@@ -0,0 +1,118 @@
+"""shattering
+
+Revision ID: ccb2a6b81a8b
+Revises: 8437ce954087
+Create Date: 2024-09-13 09:44:34.534198
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+import app.models
+import sqlmodel.sql.sqltypes
+from pathlib import Path
+
+SQL_PATH = Path(__file__).parent.parent.parent / "sql"
+
+
+# revision identifiers, used by Alembic.
+revision: str = "ccb2a6b81a8b"
+down_revision: Union[str, None] = "8437ce954087"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "districtrmap",
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("CURRENT_TIMESTAMP"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("CURRENT_TIMESTAMP"),
+            nullable=False,
+        ),
+        sa.Column("uuid", app.models.UUIDType(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column(
+            "gerrydb_table_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column("num_districts", sa.Integer(), nullable=True),
+        sa.Column("tiles_s3_path", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("parent_layer", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("child_layer", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["child_layer"],
+            ["gerrydbtable.name"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["parent_layer"],
+            ["gerrydbtable.name"],
+        ),
+        sa.PrimaryKeyConstraint("uuid"),
+        sa.UniqueConstraint("uuid"),
+        sa.UniqueConstraint("gerrydb_table_name"),
+    )
+    op.create_table(
+        "parentchildedges",
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("CURRENT_TIMESTAMP"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("CURRENT_TIMESTAMP"),
+            nullable=False,
+        ),
+        sa.Column("districtr_map", app.models.UUIDType(), nullable=False),
+        sa.Column("parent_path", sa.String(), nullable=False),
+        sa.Column("child_path", sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["districtr_map"],
+            ["districtrmap.uuid"],
+        ),
+        sa.PrimaryKeyConstraint("districtr_map", "parent_path", "child_path"),
+    )
+    op.drop_column("gerrydbtable", "tiles_s3_path")
+    # ### end Alembic commands ###
+
+    for file_name in [
+        "parent_child_relationships.sql",
+        "create_shatterable_gerrydb_view.sql",
+        "create_districtr_map_udf.sql",
+        "shatter_parent.sql",
+    ]:
+        with open(SQL_PATH / file_name, "r") as f:
+            sql = f.read()
+            op.execute(sql)
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "gerrydbtable",
+        sa.Column("tiles_s3_path", sa.VARCHAR(), autoincrement=False, nullable=True),
+    )
+    op.drop_table("parentchildedges")
+    op.drop_table("districtrmap")
+    # ### end Alembic commands ###
+    op.execute("DROP PROCEDURE IF EXISTS add_parent_child_relationships(TEXT)")
+    op.execute(
+        "DROP PROCEDURE IF EXISTS create_shatterable_gerrydb_view(TEXT, TEXT, TEXT)"
+    )
+    for func_name in [
+        "create_districtr_map",
+        "shatter_parent",
+    ]:
+        op.execute(f"DROP FUNCTION IF EXISTS {func_name}")

backend/app/main.py

@@ -2,23 +2,27 @@
 from pydantic import UUID4
 from sqlalchemy import text
 from sqlalchemy.exc import ProgrammingError
-from sqlmodel import Session, select
+from sqlmodel import Session, String, select
 from starlette.middleware.cors import CORSMiddleware
 from sqlalchemy.dialects.postgresql import insert
 import logging
+from sqlalchemy import bindparam
+from sqlmodel import ARRAY
 
 import sentry_sdk
 from app.core.db import engine
 from app.core.config import settings
 from app.models import (
     Assignments,
     AssignmentsCreate,
+    DistrictrMap,
     Document,
     DocumentCreate,
     DocumentPublic,
+    GEOIDS,
+    UUIDType,
     ZonePopulation,
-    GerryDBTable,
-    GerryDBViewPublic,
+    DistrictrMapPublic,
 )
 
 if settings.ENVIRONMENT == "production":
@@ -82,18 +86,9 @@ async def create_document(
         {"gerrydb_table_name": data.gerrydb_table},
     )
     document_id = results.one()[0]  # should be only one row, one column of results
-    stmt = (
-        select(
-            Document.document_id,
-            Document.created_at,
-            Document.updated_at,
-            Document.gerrydb_table,
-            GerryDBTable.tiles_s3_path,
-        )
-        .where(Document.document_id == document_id)
-        .join(GerryDBTable, Document.gerrydb_table == GerryDBTable.name, isouter=True)
-        .limit(1)
-    )
+    stmt = select(Document).where(Document.document_id == document_id).limit(1)
+    # Document id has a unique constraint so I'm not sure we need to hit the DB again here
+    # more valuable would be to check that the assignments table
     doc = session.exec(
         stmt
     ).one()  # again if we've got more than one, we have problems.
@@ -139,11 +134,39 @@ async def update_assignments(
     stmt = stmt.on_conflict_do_update(
         constraint=Assignments.__table__.primary_key, set_={"zone": stmt.excluded.zone}
     )
-    session.exec(stmt)
+    session.execute(stmt)
     session.commit()
     return {"assignments_upserted": len(data.assignments)}
 
 
+@app.patch(
+    "/api/update_assignments/{document_id}/shatter_parents",
+    response_model=list[Assignments],
+)
+async def shatter_parent(
+    document_id: str, data: GEOIDS, session: Session = Depends(get_session)
+):
+    stmt = text("""SELECT *
+        FROM shatter_parent(:input_document_id, :parent_geoids)""").bindparams(
+        bindparam(key="input_document_id", type_=UUIDType),
+        bindparam(key="parent_geoids", type_=ARRAY(String)),
+    )
+    results = session.execute(
+        statement=stmt,
+        params={
+            "input_document_id": document_id,
+            "parent_geoids": data.geoids,
+        },
+    )
+    # :( was getting validation errors so am just going to loop
+    assignments = [
+        Assignments(document_id=str(document_id), geo_id=geo_id, zone=zone)
+        for document_id, geo_id, zone in results
+    ]
+    session.commit()
+    return assignments
+
+
 # called by getAssignments in apiHandlers.ts
 @app.get("/api/get_assignments/{document_id}", response_model=list[Assignments])
 async def get_assignments(document_id: str, session: Session = Depends(get_session)):
@@ -160,10 +183,10 @@ async def get_document(document_id: str, session: Session = Depends(get_session)
             Document.created_at,
             Document.gerrydb_table,
             Document.updated_at,
-            GerryDBTable.tiles_s3_path.label("tiles_s3_path"),
-        )
+            DistrictrMap.tiles_s3_path.label("tiles_s3_path"),  # pyright: ignore
+        )  # pyright: ignore
         .where(Document.document_id == document_id)
-        .join(GerryDBTable, Document.gerrydb_table == GerryDBTable.name, isouter=True)
+        .join(DistrictrMap, Document.gerrydb_table == DistrictrMap.name, isouter=True)
         .limit(1)
     )
     result = session.exec(stmt)
@@ -195,16 +218,16 @@ async def get_total_population(
             )
 
 
-@app.get("/api/gerrydb/views", response_model=list[GerryDBViewPublic])
+@app.get("/api/gerrydb/views", response_model=list[DistrictrMapPublic])
 async def get_projects(
     *,
     session: Session = Depends(get_session),
     offset: int = 0,
     limit: int = Query(default=100, le=100),
 ):
     gerrydb_views = session.exec(
-        select(GerryDBTable)
-        .order_by(GerryDBTable.created_at.asc())
+        select(DistrictrMap)
+        .order_by(DistrictrMap.created_at.asc())  # pyright: ignore
         .offset(offset)
         .limit(limit)
     ).all()

backend/app/models.py

@@ -3,13 +3,15 @@
 from pydantic import UUID4, BaseModel
 from sqlmodel import (
     Field,
+    ForeignKey,
     SQLModel,
     UUID,
     TIMESTAMP,
     UniqueConstraint,
     text,
     Column,
     MetaData,
+    String,
 )
 from app.constants import DOCUMENT_SCHEMA
 
@@ -40,15 +42,58 @@ class TimeStampMixin(SQLModel):
     )
 
 
-class GerryDBTable(TimeStampMixin, SQLModel, table=True):
+class DistrictrMap(TimeStampMixin, SQLModel, table=True):
     uuid: str = Field(sa_column=Column(UUIDType, unique=True, primary_key=True))
-    name: str = Field(nullable=False, unique=True)
+    name: str = Field(nullable=False)
+    # This is intentionally not a foreign key on `GerryDBTable` because in some cases
+    # this may be the GerryDBTable but in others the pop table may be a materialized
+    # view of two GerryDBTables in the case of shatterable maps.
+    # We'll want to enforce the constraint tha the gerrydb_table_name is either in
+    # GerrydbTable.name or a materialized view of two GerryDBTables some other way.
+    gerrydb_table_name: str | None = Field(nullable=True, unique=True)
+    # Null means default number of districts? Should we have a sensible default?
+    num_districts: int | None = Field(nullable=True, default=None)
     tiles_s3_path: str | None = Field(nullable=True)
+    parent_layer: str = Field(
+        sa_column=Column(UUIDType, ForeignKey("gerrydbtable.uuid"), nullable=False)
+    )
+    child_layer: str | None = Field(
+        sa_column=Column(
+            UUIDType, ForeignKey("gerrydbtable.uuid"), default=None, nullable=True
+        )
+    )
+    # schema? will need to contrain the schema
+    # where does this go?
+    # when you create the view, pull the columns that you need
+    # we'll want discrete management steps
 
 
-class GerryDBViewPublic(BaseModel):
+class DistrictrMapPublic(BaseModel):
     name: str
-    tiles_s3_path: str | None
+    gerrydb_table_name: str
+    parent_layer: str
+    child_layer: str | None = None
+    tiles_s3_path: str | None = None
+    num_districts: int | None = None
+
+
+class GerryDBTable(TimeStampMixin, SQLModel, table=True):
+    uuid: str = Field(sa_column=Column(UUIDType, unique=True, primary_key=True))
+    # Must correspond to the layer name in the tileset
+    name: str = Field(nullable=False, unique=True)
+
+
+class ParentChildEdges(TimeStampMixin, SQLModel, table=True):
+    districtr_map: str = Field(
+        sa_column=Column(
+            UUIDType,
+            ForeignKey("districtrmap.uuid"),
+            nullable=False,
+            primary_key=True,
+        )
+    )
+    parent_path: str = Field(sa_column=Column(String, nullable=False, primary_key=True))
+    child_path: str = Field(sa_column=Column(String, nullable=False, primary_key=True))
 
 
 class Document(TimeStampMixin, SQLModel, table=True):
@@ -91,6 +136,10 @@ class AssignmentsCreate(BaseModel):
     assignments: list[Assignments]
 
 
+class GEOIDS(BaseModel):
+    geoids: list[str]
+
+
 class ZonePopulation(BaseModel):
     zone: int
     total_pop: int