From 038f82a1c83d617e857cceadc24162b82aa57fc0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:03:50 +0400 Subject: [PATCH 1/2] Add failing tests for the MySQL FOUND_ROWS client flag The v1 engine opens its MySQL connections through the Cloud SQL connector, passed to create_engine as creator=. SQLAlchemy then never applies the connection arguments its MySQL dialect would otherwise give pymysql, so the CLIENT.FOUND_ROWS flag is missing. Without it MySQL reports the rows an UPDATE changed rather than the rows it matched, and SQLAlchemy raises StaleDataError when the ORM writes values the row already holds. - Unit tests build the remote engine with a fake connector and require the connect call to pass client_flag with CLIENT.FOUND_ROWS. - A real MySQL test runs the engine the Public API builds against the disposable MySQL database and has the ORM write a status that a SQL UPDATE already set in the same transaction. A control test shows the same steps raise StaleDataError on a connection without the flag. - The cross-database integration workflow runs the new MySQL test file. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/v2-integration-check.yml | 2 + .../test_mysql_update_row_counts.py | 175 ++++++++++++++++++ tests/unit/data/test_sqlalchemy_v2.py | 86 ++++++--- tests/unit/test_alembic_workflows.py | 3 +- 4 files changed, 244 insertions(+), 22 deletions(-) create mode 100644 tests/integration/test_mysql_update_row_counts.py diff --git a/.github/workflows/v2-integration-check.yml b/.github/workflows/v2-integration-check.yml index f6ebc8262..9ea69eff4 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.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: diff --git a/tests/integration/test_mysql_update_row_counts.py b/tests/integration/test_mysql_update_row_counts.py new file mode 100644 index 000000000..932a4bd57 --- /dev/null +++ b/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) diff --git a/tests/unit/data/test_sqlalchemy_v2.py b/tests/unit/data/test_sqlalchemy_v2.py index 94979f197..fede43295 100644 --- a/tests/unit/data/test_sqlalchemy_v2.py +++ b/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 ( diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 6bb8c4032..04e81c8fb 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/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 From 568dd94b0e7d4d19534b9defd83ceefa1652bca4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:05:08 +0400 Subject: [PATCH 2/2] Set the MySQL FOUND_ROWS client flag on Cloud SQL connections The v1 engine passes the Cloud SQL connector to create_engine as creator=, so SQLAlchemy never adds the CLIENT.FOUND_ROWS flag its MySQL dialect sets on connections it opens from a URL. Without the flag MySQL reports the rows an UPDATE changed rather than the rows it matched, so an ORM UPDATE that writes values the row already holds reports 0 rows and SQLAlchemy raises StaleDataError. Report saves (PATCH //report) failed this way with a 500. The connector call now passes client_flag with the FOUND_ROWS bit, which the connector forwards to pymysql. The flag's value is written out in orm.py because application code may not import database drivers; the unit tests check it against pymysql.constants.CLIENT.FOUND_ROWS. Co-Authored-By: Claude Opus 5.5 --- changelog.d/mysql-found-rows.fixed.md | 1 + docs/engineering/skills/alembic-migrations.md | 6 +++++- policyengine_api/data/orm.py | 8 ++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) create mode 100644 changelog.d/mysql-found-rows.fixed.md diff --git a/changelog.d/mysql-found-rows.fixed.md b/changelog.d/mysql-found-rows.fixed.md new file mode 100644 index 000000000..ae77c4402 --- /dev/null +++ b/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 //report`) with a 500. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index f5c72523c..9bcf0139a 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/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 diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 8342eb3b7..c64099222 100644 --- a/policyengine_api/data/orm.py +++ b/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(