Set the MySQL FOUND_ROWS client flag on Cloud SQL connections

.github/workflows/v2-integration-check.yml

@@ -90,6 +90,8 @@ jobs:
         run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_stage12_comparison_persistence.py tests/integration/test_v2_privilege_reconciliation.py
       - name: Test v2 household persistence and immediate v1 copying
         run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_household_persistence.py tests/integration/test_v1_household_dual_write.py tests/integration/test_mysql_household_spm_dual_write.py
+      - name: Test the row counts v1 ORM updates report on MySQL
+        run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_mysql_update_row_counts.py
       - name: Qualify production-scale v2 metadata publication
         run: uv run --frozen --extra dev coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication_qualification.py
         env:

changelog.d/mysql-found-rows.fixed.md

@@ -0,0 +1 @@
+Cloud SQL connections now set MySQL's `CLIENT_FOUND_ROWS` flag, so an UPDATE reports the rows it matched rather than the rows whose values changed. Without it, SQLAlchemy raised `StaleDataError` when the ORM wrote values a row already held, which failed report saves (`PATCH /<country_id>/report`) with a 500.

docs/engineering/skills/alembic-migrations.md

@@ -59,7 +59,11 @@ Alembic revisions.
 
 - API v1 builds its Cloud SQL engine in `policyengine_api/data/orm.py`. The
   Cloud SQL connector is passed to `create_engine` as `creator=`, so SQLAlchemy
-  opens and pools the connections.
+  opens and pools the connections. SQLAlchemy does not add its MySQL connection
+  arguments to connections a `creator=` function opens, so the connector call
+  passes MySQL's `CLIENT_FOUND_ROWS` flag in `client_flag` itself; with it an
+  UPDATE reports the rows it matched rather than the rows whose values changed,
+  which SQLAlchemy needs to check UPDATE row counts.
 - API v2 builds its SQLModel engine from a `postgresql+psycopg` URL in
   `policyengine_api/data/v2/database.py`.
 - Take connections and sessions from those engines, never from

policyengine_api/data/orm.py

@@ -24,6 +24,13 @@
 DATABASE_POOL_SIZE = 5
 DATABASE_POOL_MAX_OVERFLOW = 2
 DATABASE_POOL_TIMEOUT_SECONDS = 30
+# MySQL's CLIENT_FOUND_ROWS capability flag (pymysql.constants.CLIENT.FOUND_ROWS).
+# With it, an UPDATE reports the rows it matched rather than the rows whose
+# values changed, which SQLAlchemy's MySQL dialects rely on to check UPDATE row
+# counts. SQLAlchemy sets it only on connections it opens from a URL, not on
+# those a creator= function opens. The value is written out because application
+# code may not import database drivers (rule api.data.sqlalchemy-and-alembic-only).
+MYSQL_CLIENT_FOUND_ROWS_FLAG = 2
 
 
 @dataclass
@@ -76,6 +83,7 @@ def get_connection() -> DBAPIConnection:
             db=config["db_name"],
             user=config["db_user"],
             password=password,
+            client_flag=MYSQL_CLIENT_FOUND_ROWS_FLAG,
         )
 
     engine = create_engine(

tests/integration/test_mysql_update_row_counts.py

@@ -0,0 +1,175 @@
+"""Real MySQL checks of the row counts the v1 engine reports for ORM UPDATEs.
+
+SQLAlchemy raises StaleDataError when an ORM UPDATE reports fewer rows than it
+expected. MySQL counts only the rows whose values changed unless the connection
+sets the CLIENT.FOUND_ROWS flag, so an UPDATE that writes values the row already
+holds reports 0 rows without it.
+
+Use the local, Alembic-migrated target required by test_mysql_policy_dual_write.
+The Cloud SQL connector is replaced by one that opens pymysql connections to
+that server, passing on the connection arguments the engine gives it as the
+real connector does. Only the row each test creates is removed during cleanup.
+"""
+
+from __future__ import annotations
+
+import pymysql
+import pytest
+from sqlalchemy import Engine, create_engine, delete, select, update
+from sqlalchemy.engine import URL, make_url
+from sqlalchemy.orm.exc import StaleDataError
+from sqlalchemy.pool import NullPool
+
+from policyengine_api.data import orm
+from policyengine_api.data.v1_models import ReportOutput
+from tests.integration.test_mysql_policy_dual_write import _mysql_url
+
+MYSQL_DEFAULT_PORT = 3306
+
+
+class LocalMySqlConnector:
+    """Stands in for the Cloud SQL connector, connecting to the local test server.
+
+    The real connector's pymysql path passes every argument other than the
+    instance name and driver on to pymysql, so this one does the same.
+    """
+
+    def __init__(self, database_url: URL, **connector_options):
+        self.database_url = database_url
+        self.connector_options = connector_options
+
+    def connect(self, *, instance_connection_string, driver, **driver_arguments):
+        assert driver == "pymysql"
+        return pymysql.connect(
+            host=self.database_url.host,
+            port=self.database_url.port or MYSQL_DEFAULT_PORT,
+            **driver_arguments,
+        )
+
+    def close(self):
+        pass
+
+
+@pytest.fixture
+def mysql_database_url() -> URL:
+    return make_url(_mysql_url())
+
+
+@pytest.fixture
+def v1_remote_engine(monkeypatch, mysql_database_url):
+    """Yield the engine orm._build_remote_engine builds, connected to local MySQL."""
+    monkeypatch.setattr(orm, "_cloud_sql_state", orm._CloudSqlState())
+    monkeypatch.setattr(
+        orm,
+        "Connector",
+        lambda **connector_options: LocalMySqlConnector(
+            mysql_database_url, **connector_options
+        ),
+    )
+    monkeypatch.setenv("POLICYENGINE_DB_INSTANCE_CONNECTION_NAME", "local:test:mysql")
+    monkeypatch.setenv("POLICYENGINE_DB_USER", mysql_database_url.username)
+    monkeypatch.setenv("POLICYENGINE_DB_NAME", mysql_database_url.database)
+    monkeypatch.setenv("POLICYENGINE_DB_PASSWORD", mysql_database_url.password)
+    engine = orm._build_remote_engine()
+    try:
+        yield engine
+    finally:
+        engine.dispose()
+
+
+@pytest.fixture
+def engine_without_found_rows(mysql_database_url):
+    """Yield an engine whose creator opens pymysql connections with no client flags."""
+
+    def connect_without_found_rows():
+        return pymysql.connect(
+            host=mysql_database_url.host,
+            port=mysql_database_url.port or MYSQL_DEFAULT_PORT,
+            user=mysql_database_url.username,
+            password=mysql_database_url.password,
+            database=mysql_database_url.database,
+        )
+
+    engine = create_engine(
+        "mysql+pymysql://",
+        creator=connect_without_found_rows,
+        poolclass=NullPool,
+    )
+    try:
+        yield engine
+    finally:
+        engine.dispose()
+
+
+def _create_pending_report_output(engine: Engine) -> int:
+    with orm.build_session_factory(engine).begin() as session:
+        report_output = ReportOutput(
+            country_id="us",
+            simulation_1_id=0,
+            api_version="found-rows",
+            status="pending",
+        )
+        session.add(report_output)
+        session.flush()
+        return report_output.id
+
+
+def _delete_report_output(engine: Engine, report_output_id: int) -> None:
+    with engine.begin() as connection:
+        connection.execute(
+            delete(ReportOutput).where(ReportOutput.id == report_output_id)
+        )
+
+
+def _complete_report_output_already_completed_by_sql(
+    engine: Engine, report_output_id: int
+) -> None:
+    """Have the ORM write a status that SQL set earlier in the same transaction.
+
+    The ORM loads the row while it is pending, a Core UPDATE then sets it
+    complete without the ORM seeing it, and the ORM sets complete again and
+    flushes. MySQL finds the row but changes nothing.
+    """
+    with orm.build_session_factory(engine).begin() as session:
+        report_output = session.scalar(
+            select(ReportOutput)
+            .where(ReportOutput.id == report_output_id)
+            .with_for_update()
+        )
+        assert report_output is not None
+        assert report_output.status == "pending"
+        session.connection().execute(
+            update(ReportOutput.__table__)
+            .where(ReportOutput.__table__.c.id == report_output_id)
+            .values(status="complete")
+        )
+        report_output.status = "complete"
+        session.flush()
+
+
+def test_orm_update_to_unchanged_values_raises_without_found_rows(
+    engine_without_found_rows,
+):
+    report_output_id = _create_pending_report_output(engine_without_found_rows)
+    try:
+        with pytest.raises(StaleDataError, match="0 were matched"):
+            _complete_report_output_already_completed_by_sql(
+                engine_without_found_rows, report_output_id
+            )
+    finally:
+        _delete_report_output(engine_without_found_rows, report_output_id)
+
+
+def test_v1_engine_orm_update_to_unchanged_values_succeeds(v1_remote_engine):
+    report_output_id = _create_pending_report_output(v1_remote_engine)
+    try:
+        _complete_report_output_already_completed_by_sql(
+            v1_remote_engine, report_output_id
+        )
+
+        with orm.build_session_factory(v1_remote_engine)() as session:
+            report_output = session.get(ReportOutput, report_output_id)
+            assert report_output is not None
+            assert report_output.status == "complete"
+    finally:
+        _delete_report_output(v1_remote_engine, report_output_id)

tests/unit/data/test_sqlalchemy_v2.py

@@ -1,36 +1,49 @@
 from pathlib import Path
+from typing import NamedTuple
 from unittest.mock import Mock
 
+import pytest
 import sqlalchemy
+from pymysql.constants import CLIENT
 
 from policyengine_api.data import orm
 
 
-def test_sqlalchemy_v2_or_newer_is_installed():
-    assert int(sqlalchemy.__version__.split(".")[0]) >= 2
+class FakeCloudSqlConnector:
+    """Records the options it is built with and every connect call it receives."""
 
+    def __init__(self, **connector_options):
+        self.connector_options = connector_options
+        self.connect_calls = []
+
+    def connect(self, **connect_options):
+        self.connect_calls.append(connect_options)
+        return object()
+
+    def close(self):
+        pass
 
-def test_remote_engine_delegates_connections_and_pooling_to_sqlalchemy(monkeypatch):
-    connector_calls = []
-    engine_calls = []
 
-    class FakeConnector:
-        def __init__(self, **options):
-            self.options = options
+class RemoteEngineBuild(NamedTuple):
+    """What _build_remote_engine returned and what it passed to its dependencies."""
 
-        def connect(self, **options):
-            connector_calls.append(options)
-            return object()
+    engine: object
+    connector: FakeCloudSqlConnector
+    create_engine_calls: list
 
-        def close(self):
-            pass
 
-    connector = FakeConnector()
+@pytest.fixture
+def remote_engine_build(monkeypatch):
+    """Build the remote engine with a fake connector and a recording create_engine."""
+    connectors = []
+    engine_calls = []
 
-    def connector_factory(**options):
-        connector.options = options
+    def connector_factory(**connector_options):
+        connector = FakeCloudSqlConnector(**connector_options)
+        connectors.append(connector)
         return connector
 
+    monkeypatch.setattr(orm, "_cloud_sql_state", orm._CloudSqlState())
     monkeypatch.setattr(orm, "Connector", connector_factory)
     monkeypatch.setattr(
         orm,
@@ -46,31 +59,62 @@ def connector_factory(**options):
 
     engine = orm._build_remote_engine()
 
-    assert engine == "engine"
-    assert connector.options == {
+    [connector] = connectors
+    return RemoteEngineBuild(engine, connector, engine_calls)
+
+
+def test_sqlalchemy_v2_or_newer_is_installed():
+    assert int(sqlalchemy.__version__.split(".")[0]) >= 2
+
+
+def test_remote_engine_delegates_connections_and_pooling_to_sqlalchemy(
+    remote_engine_build,
+):
+    connector = remote_engine_build.connector
+
+    assert remote_engine_build.engine == "engine"
+    assert connector.connector_options == {
         "ip_type": orm.IPTypes.PUBLIC,
         "refresh_strategy": "LAZY",
     }
-    url, options = engine_calls[0]
+    url, options = remote_engine_build.create_engine_calls[0]
     assert url == "mysql+pymysql://"
     assert options["pool_pre_ping"] is True
     assert options["pool_recycle"] == 1800
     assert options["pool_size"] == 5
     assert options["max_overflow"] == 2
     assert options["pool_timeout"] == 30
-    assert connector_calls == []
+    assert connector.connect_calls == []
     options["creator"]()
-    assert connector_calls == [
+    assert connector.connect_calls == [
         {
             "instance_connection_string": "project:region:instance",
             "driver": "pymysql",
             "db": "database",
             "user": "user",
             "password": "password",
+            "client_flag": CLIENT.FOUND_ROWS,
         }
     ]
 
 
+def test_remote_engine_connections_report_rows_matched_by_updates(
+    remote_engine_build,
+):
+    """SQLAlchemy checks UPDATE row counts and needs MySQL to count matched rows.
+
+    SQLAlchemy sets CLIENT.FOUND_ROWS itself only when it opens a connection
+    from a URL. Through creator= it does not, so the connector call must.
+    """
+    _url, options = remote_engine_build.create_engine_calls[0]
+
+    options["creator"]()
+
+    [connect_options] = remote_engine_build.connector.connect_calls
+    client_flag = connect_options.get("client_flag", 0)
+    assert client_flag & CLIENT.FOUND_ROWS == CLIENT.FOUND_ROWS
+
+
 def test_production_orm_has_no_sqlite_debug_or_local_schema_path():
     source = Path(orm.__file__).read_text(encoding="utf-8")
     for prohibited in (

tests/unit/test_alembic_workflows.py

@@ -233,12 +233,13 @@ def test_reusable_v2_integration_check_uses_disposable_databases_and_coverage():
     assert "test_v2_household_persistence.py" in workflow
     assert "test_v1_household_dual_write.py" in workflow
     assert "test_mysql_household_spm_dual_write.py" in workflow
+    assert "test_mysql_update_row_counts.py" in workflow
     assert "test_v2_catalog_publication_qualification.py" in workflow
     assert "RUN_V2_CATALOG_PUBLICATION_QUALIFICATION" in workflow
     assert "test_runtime_cache_redis.py" in workflow
     assert "uv sync --frozen" in workflow
     assert workflow.count("coverage run --branch") == 1
-    assert workflow.count("coverage run -a --branch") == 7
+    assert workflow.count("coverage run -a --branch") == 8
     assert "coverage xml -i -o coverage-v2.xml" in workflow
     assert "codecov/codecov-action@v5" in workflow
     assert "files: coverage-v2.xml" in workflow