fix(tests): stop tests leaving Postgres session settings on pooled connections (#16129)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

backend/alembic/versions/9d8f310ac672_add_chat_execution_metadata.py

@@ -42,35 +42,44 @@ def upgrade() -> None:
     bind.commit()
     with bind.engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
         conn.exec_driver_sql(f"SET lock_timeout = '{LOCK_TIMEOUT}'")
-        quoted_schema = conn.dialect.identifier_preparer.quote_identifier(schema)
-        for name, statement in (
-            (
-                "ix_chat_session_spawned_by_message_id",
-                'CREATE INDEX CONCURRENTLY "ix_chat_session_spawned_by_message_id" '
-                f"ON {quoted_schema}.chat_session (spawned_by_message_id)",
-            ),
-            (
-                "ix_chat_message_invoking_tool_call_id",
-                'CREATE INDEX CONCURRENTLY "ix_chat_message_invoking_tool_call_id" '
-                f"ON {quoted_schema}.chat_message (invoking_tool_call_id)",
-            ),
-        ):
-            qualified_name = f'{quoted_schema}."{name}"'
-            state = _index_state(conn, qualified_name)
-            if state is True:
-                continue
-            if state is False:
-                conn.exec_driver_sql(f"DROP INDEX CONCURRENTLY {qualified_name}")
-            conn.exec_driver_sql(statement)
+        # A plain SET outlives the transaction on this connection. Reset it so a
+        # pooled connection does not carry the timeout to its next user.
+        try:
+            _build_indexes_and_validate(conn, schema)
+        finally:
+            conn.exec_driver_sql("RESET lock_timeout")
 
-        # Validate only after releasing the initial ALTER TABLE locks.
-        for table, constraint in (
-            ("chat_session", "fk_chat_session_spawned_by_message"),
-            ("chat_message", "fk_chat_message_invocation"),
-        ):
-            conn.exec_driver_sql(
-                f"ALTER TABLE {quoted_schema}.{table} VALIDATE CONSTRAINT {constraint}"
-            )
+
+def _build_indexes_and_validate(conn: sa.engine.Connection, schema: str) -> None:
+    quoted_schema = conn.dialect.identifier_preparer.quote_identifier(schema)
+    for name, statement in (
+        (
+            "ix_chat_session_spawned_by_message_id",
+            'CREATE INDEX CONCURRENTLY "ix_chat_session_spawned_by_message_id" '
+            f"ON {quoted_schema}.chat_session (spawned_by_message_id)",
+        ),
+        (
+            "ix_chat_message_invoking_tool_call_id",
+            'CREATE INDEX CONCURRENTLY "ix_chat_message_invoking_tool_call_id" '
+            f"ON {quoted_schema}.chat_message (invoking_tool_call_id)",
+        ),
+    ):
+        qualified_name = f'{quoted_schema}."{name}"'
+        state = _index_state(conn, qualified_name)
+        if state is True:
+            continue
+        if state is False:
+            conn.exec_driver_sql(f"DROP INDEX CONCURRENTLY {qualified_name}")
+        conn.exec_driver_sql(statement)
+
+    # Validate only after releasing the initial ALTER TABLE locks.
+    for table, constraint in (
+        ("chat_session", "fk_chat_session_spawned_by_message"),
+        ("chat_message", "fk_chat_message_invocation"),
+    ):
+        conn.exec_driver_sql(
+            f"ALTER TABLE {quoted_schema}.{table} VALIDATE CONSTRAINT {constraint}"
+        )
 
 
 def _add_execution_metadata() -> None:

backend/alembic/versions/e0ea2ae62e51_add_index_on_chat_session_user_id_and_.py

@@ -126,7 +126,12 @@ def upgrade() -> None:
     with bind.engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
         if _index_state(conn, schema) is True:
             return
-        _build_index(conn, schema)
+        # _build_index's plain SET outlives the transaction on this connection.
+        # Reset it so a pooled connection does not carry the timeout to its next user.
+        try:
+            _build_index(conn, schema)
+        finally:
+            conn.exec_driver_sql("RESET lock_timeout")
 
 
 def downgrade() -> None:

backend/tests/README.md

@@ -81,6 +81,20 @@ pytest --count=10 backend/tests/unit/path/to/test_file.py
 
 ## Best Practices
 
+### Leave nothing behind
+
+A test may reset or wipe shared state before it runs. It must not leave its own
+state behind for a later test to see: rows, changed global defaults, Redis keys,
+OpenSearch indexes, running threads, or settings on a pooled database connection.
+Undo each change in a `finally` or a yield-fixture teardown.
+
+External dependency unit tests check one case of this automatically. After each
+test, an autouse fixture (`no_leaked_db_session_settings`) fails the test if it
+returned a pooled connection with a session-level Postgres setting on it, such as
+a plain `SET lock_timeout`. To change session state, use `SET LOCAL`, reset the
+setting, or run on the `isolated_engine` fixture. Migration tests that run
+migrations in-process must use `isolated_engine`.
+
 ### Use `enable_ee` fixture instead of inlining
 
 Enables EE mode for a test, with proper teardown and cache clearing.

backend/tests/external_dependency_unit/conftest.py

@@ -4,8 +4,10 @@ from uuid import uuid4
 
 import pytest
 from fastapi_users.password import PasswordHelper
-from sqlalchemy import delete
+from sqlalchemy import create_engine, delete
+from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session
+from sqlalchemy.pool import NullPool
 
 from onyx.db.engine.sql_engine import SqlEngine, get_session_with_current_tenant
 from onyx.db.enums import AccountType
@@ -15,6 +17,9 @@ from onyx.file_store.file_store import get_default_file_store
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 from tests.external_dependency_unit.full_setup import ensure_full_deployment_setup
+from tests.external_dependency_unit.session_state_guard import (
+    reset_leaked_session_settings,
+)
 
 # Opt into the shared @pytest.mark.secrets / test_secrets infrastructure.
 from tests.utils.pytest_secrets import (
@@ -38,6 +43,20 @@ def audit_stream(caplog: pytest.LogCaptureFixture) -> Generator[None, None, None
         audit_logger.removeHandler(caplog.handler)
 
 
+@pytest.fixture(autouse=True)
+def no_leaked_db_session_settings() -> Generator[None, None, None]:
+    """Fail the test that returns a pooled connection with session settings on
+    it, instead of a later test that happens to check that connection out."""
+    yield
+    leaks = reset_leaked_session_settings()
+    if leaks:
+        pytest.fail(
+            "This test left Postgres session settings on pooled connections: "
+            f"{leaks}. Use SET LOCAL, reset the setting, or run on the "
+            "`isolated_engine` fixture."
+        )
+
+
 @pytest.fixture(scope="function")
 def db_session() -> Generator[Session, None, None]:
     """Create a database session for testing using the actual PostgreSQL database"""
@@ -50,6 +69,20 @@ def db_session() -> Generator[Session, None, None]:
         yield session
 
 
+@pytest.fixture
+def isolated_engine(
+    db_session: Session,  # noqa: ARG001 - initializes SqlEngine
+) -> Generator[Engine, None, None]:
+    """An engine with no pool, for tests that change session state (search_path,
+    lock_timeout) or hand a connection to migration code. Each connection closes
+    when the test returns it, so none of that state reaches the shared app pool."""
+    engine = create_engine(SqlEngine.get_engine().url, poolclass=NullPool)
+    try:
+        yield engine
+    finally:
+        engine.dispose()
+
+
 @pytest.fixture(scope="session")
 def full_deployment_setup() -> Generator[None, None, None]:
     """Optional fixture to perform full deployment-like setup on demand.

backend/tests/external_dependency_unit/craft/conftest.py

@@ -65,7 +65,7 @@ def _best_effort_delete(model: type[Any], ids: Iterable[Any]) -> None:
         token = CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE)
         try:
             with get_session_with_current_tenant() as session:
-                session.execute(text("SET lock_timeout = '10s'"))
+                session.execute(text("SET LOCAL lock_timeout = '10s'"))
                 for row_id in ids:
                     row = session.get(model, row_id)
                     if row is not None:

backend/tests/external_dependency_unit/db/test_chat_execution_metadata_migration.py

@@ -23,7 +23,7 @@ from sqlalchemy import (
     text,
 )
 from sqlalchemy.dialects.postgresql import UUID
-from sqlalchemy.engine import Connection
+from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.orm import Session
 from sqlalchemy.schema import CreateSchema, DropSchema
@@ -41,11 +41,11 @@ from onyx.db.tools import create_tool_call_no_commit
 
 
 def test_execution_metadata_upgrade_retry_downgrade_upgrade(
-    db_session: Session,
+    isolated_engine: Engine,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     schema = f"response_migration_{uuid4().hex}"
-    with db_session.connection().engine.connect() as connection:
+    with isolated_engine.connect() as connection:
         try:
             connection.execute(CreateSchema(schema))
             connection.execute(

backend/tests/external_dependency_unit/db/test_response_checkpoint_migration.py

@@ -7,13 +7,13 @@ from alembic.migration import MigrationContext
 from alembic.operations import Operations
 from alembic.script import ScriptDirectory
 from sqlalchemy import inspect, text
-from sqlalchemy.orm import Session
+from sqlalchemy.engine import Engine
 from sqlalchemy.schema import CreateSchema, DropSchema
 
 
-def test_checkpoint_upgrade_downgrade_upgrade(db_session: Session) -> None:
+def test_checkpoint_upgrade_downgrade_upgrade(isolated_engine: Engine) -> None:
     schema = f"response_checkpoint_{uuid4().hex}"
-    with db_session.connection().engine.connect() as connection:
+    with isolated_engine.connect() as connection:
         try:
             connection.execute(CreateSchema(schema))
             connection.execute(

backend/tests/external_dependency_unit/db/test_response_storage_migration.py

@@ -18,6 +18,7 @@ from sqlalchemy import (
     select,
     text,
 )
+from sqlalchemy.engine import Engine
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 from sqlalchemy.schema import CreateSchema, DropSchema
@@ -34,9 +35,9 @@ from onyx.db.models import ChatMessage, ChatSession, ToolCall
 from onyx.db.tools import create_tool_call_no_commit
 
 
-def test_response_items_upgrade_downgrade_upgrade(db_session: Session) -> None:
+def test_response_items_upgrade_downgrade_upgrade(isolated_engine: Engine) -> None:
     schema = f"response_items_{uuid4().hex}"
-    with db_session.connection().engine.connect() as connection:
+    with isolated_engine.connect() as connection:
         try:
             connection.execute(CreateSchema(schema))
             connection.execute(

backend/tests/external_dependency_unit/session_state_guard.py

@@ -0,0 +1,62 @@
+"""Find Postgres session settings that a test left on the app's pooled connections.
+
+A plain `SET` (or `set_config(..., false)`) that commits, or runs on an AUTOCOMMIT
+connection, stays on the connection after the pool takes it back. The next test
+that checks the connection out then runs with that setting: for example a
+`lock_timeout` of 100ms makes a blocking insert fail at once.
+"""
+
+from sqlalchemy import text
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.pool import QueuePool
+
+from onyx.db.engine.sql_engine import SqlEngine
+
+# SQLAlchemy sets this itself when it restores the isolation level of a
+# connection that a caller switched to AUTOCOMMIT.
+_SET_BY_SQLALCHEMY = frozenset({"transaction_isolation"})
+
+_SESSION_SETTINGS = text(
+    "SELECT name, setting FROM pg_settings WHERE source = 'session' ORDER BY name"
+)
+
+
+def _app_engines() -> list[Engine]:
+    engines: list[Engine] = []
+    for get_engine in (SqlEngine.get_engine, SqlEngine.get_readonly_engine):
+        try:
+            engines.append(get_engine())
+        except RuntimeError:
+            # Not initialized in this process.
+            continue
+    return engines
+
+
+def reset_leaked_session_settings() -> list[str]:
+    """Check out every idle pooled connection of the app engines, reset any
+    session settings left on it, and return a description of each one found."""
+    leaks: list[str] = []
+    for engine in _app_engines():
+        pool = engine.pool
+        if not isinstance(pool, QueuePool):
+            continue
+        # Holding them all at once makes the pool hand out each idle one.
+        connections: list[Connection] = []
+        try:
+            connections.extend(engine.connect() for _ in range(pool.checkedin()))
+            for connection in connections:
+                found = [
+                    f"{name}={setting}"
+                    for name, setting in connection.execute(_SESSION_SETTINGS).all()
+                    if name not in _SET_BY_SQLALCHEMY
+                ]
+                if found:
+                    leaks.append(", ".join(found))
+                    connection.exec_driver_sql("RESET ALL")
+                    connection.commit()
+                else:
+                    connection.rollback()
+        finally:
+            for connection in connections:
+                connection.close()
+    return leaks