Fix Postgres live snapshot queries

Author: worldofhacks

.github/workflows/ci.yml

@@ -77,6 +77,8 @@ jobs:
       - run: uv run mypy python services
       - run: env -u DATABASE_URL uv run pytest tests/backend
       - run: uv run python -m seekandscore.db.migrate upgrade
+      - name: Postgres repository integration
+        run: SEEKANDSCORE_TEST_DATABASE_URL="$DATABASE_URL" env -u DATABASE_URL uv run pytest tests/backend/test_postgres_repository_integration.py
 
   web:
     name: Web quality gates

python/seekandscore/acquisition/repository.py

@@ -249,11 +249,11 @@ def latest_complete_run(
             sa.select(source_run_table.c.payload)
             .where(
                 source_run_table.c.source_id == source_id,
-                source_run_table.c.payload["run_profile"].astext == run_profile.value,
-                source_run_table.c.payload["status"].astext.in_(
-                    ("succeeded", "succeeded_unchanged")
-                ),
-                source_run_table.c.payload["partial"].astext == "false",
+                source_run_table.c.payload["run_profile"].as_string() == run_profile.value,
+                source_run_table.c.payload["status"]
+                .as_string()
+                .in_(("succeeded", "succeeded_unchanged")),
+                source_run_table.c.payload["partial"].as_boolean().is_(False),
             )
             .order_by(source_run_table.c.started_at.desc(), source_run_table.c.id.desc())
         )
@@ -373,15 +373,15 @@ def list_latest_observations(
         if artifact_ids is not None:
             filters.append(observation_table.c.artifact_id.in_(artifact_ids))
         if cities is not None:
-            filters.append(observation_table.c.payload["situs_city"].astext.in_(cities))
+            filters.append(observation_table.c.payload["situs_city"].as_string().in_(cities))
         ranked = (
             sa.select(
                 observation_table.c.payload,
                 sa.func.row_number()
                 .over(
                     partition_by=observation_table.c.source_record_id,
                     order_by=sa.cast(
-                        observation_table.c.payload["observed_at"].astext,
+                        observation_table.c.payload["observed_at"].as_string(),
                         sa.DateTime(timezone=True),
                     ).desc(),
                 )
@@ -394,13 +394,8 @@ def list_latest_observations(
             sa.select(ranked.c.payload)
             .where(ranked.c.version_rank == 1)
             .order_by(
-                sa.cast(
-                    ranked.c.payload["tcad_acres"].astext,
-                    sa.Float(),
-                )
-                .desc()
-                .nullslast(),
-                ranked.c.payload["local_parcel_id"].astext,
+                ranked.c.payload["tcad_acres"].as_float().desc().nullslast(),
+                ranked.c.payload["local_parcel_id"].as_string(),
             )
             .limit(limit)
             .offset(offset)
@@ -424,7 +419,7 @@ def count_latest_observations(
             statement = statement.where(observation_table.c.artifact_id.in_(artifact_ids))
         if cities is not None:
             statement = statement.where(
-                observation_table.c.payload["situs_city"].astext.in_(cities)
+                observation_table.c.payload["situs_city"].as_string().in_(cities)
             )
         with self.engine.connect() as connection:
             return int(connection.scalar(statement) or 0)

tests/backend/test_live_candidates.py

@@ -109,6 +109,35 @@ def test_postgres_insert_count_uses_returned_ids_not_indeterminate_rowcount() ->
     assert inserted == 2
 
 
+def test_postgres_snapshot_queries_use_supported_sqlalchemy_json_accessors() -> None:
+    engine = MagicMock()
+    connection = engine.connect.return_value.__enter__.return_value
+    connection.scalars.return_value.all.return_value = []
+    connection.scalar.return_value = 0
+    repository = PostgresAcquisitionRepository(engine)
+    artifact_id = UUID("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
+
+    assert (
+        repository.latest_complete_run("travis_tcad_parcels", run_profile=SourceRunProfile.COHORT)
+        is None
+    )
+    assert (
+        repository.list_latest_observations(
+            limit=25,
+            artifact_ids=(artifact_id,),
+            cities=("DEL VALLE", "MANOR"),
+        )
+        == ()
+    )
+    assert (
+        repository.count_latest_observations(
+            artifact_ids=(artifact_id,),
+            cities=("DEL VALLE", "MANOR"),
+        )
+        == 0
+    )
+
+
 def test_live_projection_labels_assessor_values_and_oz_as_unverified() -> None:
     repository = LiveCandidateRepository(
         seeded_repository(),

tests/backend/test_postgres_repository_integration.py

@@ -0,0 +1,130 @@
+"""Optional real-Postgres regression for JSON-backed live snapshot queries."""
+
+import os
+from datetime import UTC, datetime
+from uuid import uuid4
+
+import pytest
+import sqlalchemy as sa
+
+from seekandscore.acquisition.models import (
+    NormalizedParcelObservation,
+    RawArtifact,
+    SourceRun,
+    SourceRunProfile,
+    SourceRunStatus,
+)
+from seekandscore.acquisition.repository import (
+    PostgresAcquisitionRepository,
+    observation_table,
+    raw_artifact_table,
+    source_run_table,
+)
+from seekandscore.readmodels import LiveCandidateRepository
+from seekandscore.registry import InMemorySourceRegistry
+
+TEST_DATABASE_URL = os.getenv("SEEKANDSCORE_TEST_DATABASE_URL")
+
+
+@pytest.mark.skipif(not TEST_DATABASE_URL, reason="SEEKANDSCORE_TEST_DATABASE_URL is not set")
+def test_artifact_scoped_snapshot_queries_execute_on_postgres() -> None:
+    assert TEST_DATABASE_URL is not None
+    database_name = sa.engine.make_url(TEST_DATABASE_URL).database or ""
+    if not database_name.endswith("_test"):
+        pytest.fail("SEEKANDSCORE_TEST_DATABASE_URL must name a database ending in _test")
+
+    engine = sa.create_engine(TEST_DATABASE_URL)
+    repository = PostgresAcquisitionRepository(engine)
+    now = datetime.now(UTC)
+    artifact_id = uuid4()
+    observation_id = uuid4()
+    run_id = uuid4()
+    sha256 = artifact_id.hex * 2
+    artifact = RawArtifact(
+        id=artifact_id,
+        source_id="travis_tcad_parcels",
+        sha256=sha256,
+        byte_count=2,
+        media_type="application/json",
+        storage_uri=f"s3://private-test/{sha256}.json",
+        original_uri="https://example.invalid/query",
+        request_params={},
+        retrieved_at=now,
+    )
+    observation = NormalizedParcelObservation(
+        id=observation_id,
+        source_id="travis_tcad_parcels",
+        source_record_id=str(observation_id.int),
+        artifact_id=artifact_id,
+        artifact_sha256=sha256,
+        parser_version="travis-tcad-parcel-v1",
+        jurisdiction_id="us-tx-travis",
+        local_parcel_id=str(observation_id.int),
+        geographic_id=None,
+        situs_address="100 TEST RD",
+        situs_city="DEL VALLE",
+        situs_zip="78617",
+        tcad_acres=2.0,
+        observed_at=now,
+    )
+    run = SourceRun(
+        id=run_id,
+        source_id="travis_tcad_parcels",
+        status=SourceRunStatus.SUCCEEDED,
+        requested_at=now,
+        started_at=now,
+        completed_at=now,
+        retrieved_at=now,
+        records_fetched=1,
+        observations_created=1,
+        artifact_ids=(artifact_id,),
+        adapter_version="travis-tcad-arcgis-v1",
+        parser_version="travis-tcad-parcel-v1",
+        run_profile=SourceRunProfile.COHORT,
+        configuration_hash="a" * 64,
+        activation_id="test-only",
+    )
+
+    try:
+        repository.save_artifact(artifact)
+        assert repository.save_observations((observation,)) == 1
+        repository.save_run(run)
+
+        assert (
+            repository.latest_complete_run(
+                "travis_tcad_parcels", run_profile=SourceRunProfile.COHORT
+            )
+            == run
+        )
+        assert repository.list_latest_observations(
+            limit=25,
+            artifact_ids=(artifact_id,),
+            cities=("DEL VALLE", "MANOR"),
+        ) == (observation,)
+        assert (
+            repository.count_latest_observations(
+                artifact_ids=(artifact_id,),
+                cities=("DEL VALLE", "MANOR"),
+            )
+            == 1
+        )
+        candidates = LiveCandidateRepository(
+            repository,
+            InMemorySourceRegistry(),
+            display_enabled=True,
+        )
+        page = candidates.list(limit=25, cursor=None, dataset_mode="live")
+        assert page.dataset_status == "current"
+        assert page.total == 1
+        assert len(page.items) == 1
+        assert candidates.get(page.items[0].id) == page.items[0]
+    finally:
+        with engine.begin() as connection:
+            connection.execute(sa.delete(source_run_table).where(source_run_table.c.id == run_id))
+            connection.execute(
+                sa.delete(observation_table).where(observation_table.c.id == observation_id)
+            )
+            connection.execute(
+                sa.delete(raw_artifact_table).where(raw_artifact_table.c.id == artifact_id)
+            )
+        engine.dispose()