Sync assignments endpoint

backend/app/main.py

@@ -132,7 +132,7 @@ async def update_assignments(
     return {"assignments_upserted": len(data.assignments)}
 
 
-@app.get("/get_assignments/{document_id}", response_model=list[Assignments])
+@app.get("/api/get_assignments/{document_id}", response_model=list[Assignments])
 async def get_assignments(document_id: str, session: Session = Depends(get_session)):
     stmt = select(Assignments).where(Assignments.document_id == document_id)
     results = session.exec(stmt)
@@ -141,6 +141,33 @@ async def get_assignments(document_id: str, session: Session = Depends(get_sessi
     return results
 
 
+@app.post("/api/sync_assignments/{document_id}")
+async def sync_assignemnts(
+    document_id: str, data: AssignmentsCreate, session: Session = Depends(get_session)
+):
+    doc_partition = f"document.assignments_{document_id}"
+
+    stmt = text(f"""
+        TRUNCATE TABLE "{doc_partition}";
+    """)
+    session.execute(stmt)
+    assignments_string = ",".join(
+        [
+            f"({assignment.document_id}, {assignment.geoid}, {assignment.zone} )"
+            for assignment in data.model_dump()["assignments"]
+        ]
+    )
+    # doing this as text so we can insert into a specific partition
+    stmt = text(f"""
+        INSERT INTO "{doc_partition}" (document_id, geoid, zone)
+        VALUES 
+        {assignments_string}
+    """)
+    session.execute(stmt)
+    session.commit()
+    return {"assignments_inserted": len(data.assignments)}
+
+
 @app.get("/api/document/{document_id}/total_pop", response_model=list[ZonePopulation])
 async def get_total_population(
     document_id: str, session: Session = Depends(get_session)

backend/tests/test_main.py

@@ -239,6 +239,21 @@ def test_patch_assignments(client, document_id):
     assert response.json() == {"assignments_upserted": 3}
 
 
+def test_sync_assignments(client, document_id):
+    response = client.post(
+        f"/api/sync_assignments/{document_id}",
+        json={
+            "assignments": [
+                {"document_id": document_id, "geo_id": "202090416004010", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090416003004", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090434001003", "zone": 2},
+            ]
+        },
+    )
+    assert response.status_code == 200
+    assert response.json() == {"assignments_inserted": 3}
+
+
 def test_patch_assignments_twice(client, document_id):
     response = client.patch(
         "/api/update_assignments",
@@ -265,7 +280,7 @@ def test_patch_assignments_twice(client, document_id):
     assert response.json() == {"assignments_upserted": 2}
     # Check that the assignments were updated and not inserted
     doc_uuid = str(uuid.UUID(document_id))
-    response = client.get(f"/get_assignments/{doc_uuid}")
+    response = client.get(f"/api/get_assignments/{doc_uuid}")
     assert response.status_code == 200
     data = response.json()
     assert data is not None