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/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( 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