[MAINTENANCE] Dispose of mssql connections in integration tests (#11663)

Author: tyler-hoffman

tests/integration/sql_session_manager.py

@@ -76,6 +76,20 @@ def get_engine(
             logger.info(f"Cache hit for engine: {cache_key}")
         return self._engine_cache[cache_key]
 
+    def dispose_engine(self, connection_details: ConnectionDetails) -> None:
+        """Dispose a specific cached engine and remove it from the cache.
+
+        Useful before teardown operations (e.g. DROP SCHEMA) that require all
+        connections to be released, such as MSSQL which holds schema locks.
+        """
+        if connection_details in self._engine_cache:
+            engine = self._engine_cache.pop(connection_details)
+            logger.info(f"Disposing engine for teardown: {connection_details}")
+            try:
+                engine.dispose()
+            except Exception:
+                logger.exception(f"Error disposing engine '{connection_details}'")
+
     def dispose_all_engines(self):
         logger.info("Disposing all cached SQLAlchemy engines.")
         for key, engine in self._engine_cache.items():

tests/integration/test_utils/data_source_config/mssql.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Mapping, Optional
 
 import pandas as pd
@@ -9,14 +10,16 @@
 from great_expectations.datasource.fluent.sql_server_datasource import (
     SQLServerAuthConnectionDetails,
 )
-from tests.integration.sql_session_manager import SessionSQLEngineManager
+from tests.integration.sql_session_manager import ConnectionDetails, SessionSQLEngineManager
 from tests.integration.test_utils.data_source_config.base import (
     BatchTestSetup,
     DataSourceTestConfig,
 )
 from tests.integration.test_utils.data_source_config.sql import SQLBatchTestSetup
 from tests.test_utils import get_default_mssql_url
 
+logger = logging.getLogger(__name__)
+
 
 class MSSQLDatasourceTestConfig(DataSourceTestConfig):
     @property
@@ -79,3 +82,23 @@ def make_asset(self) -> TableAsset:
             table_name=self.table_name,
             schema_name=self.schema,
         )
+
+    @override
+    def teardown(self) -> None:
+        """Override teardown to dispose cached engines before DROP SCHEMA.
+
+        MSSQL holds schema locks on connections. Disposing the session manager's
+        cached engine releases all pool connections before we run DROP, avoiding
+        hangs. We use a fresh engine for the drop since the cached one was disposed.
+        """
+        for datasource in self.context.data_sources.all().values():
+            execution_engine = datasource.execution_engine
+            if execution_engine:
+                execution_engine.close()
+
+        if self.engine_manager:
+            self.engine_manager.dispose_engine(
+                ConnectionDetails(connection_string=self.connection_string)
+            )
+
+        super().teardown()

tests/integration/test_utils/data_source_config/sql.py

@@ -232,24 +232,8 @@ def setup(self) -> None:
             self._safe_commit(conn)
         cleanup()
 
-    def _close_execution_engines(self) -> None:
-        """Close all execution engines in the data context.
-
-        Some dialects (MSSQL, SQLite, BigQuery, Databricks) use persisted
-        connections that are kept alive for temp-table support. These
-        connections hold locks that prevent DDL operations like DROP SCHEMA.
-        Closing them before teardown avoids hangs.
-        """
-        for datasource in self.context.data_sources.all().values():
-            execution_engine = datasource.execution_engine
-            if execution_engine:
-                execution_engine.close()
-
     @override
     def teardown(self) -> None:
-        # Close any execution engine connections that may hold schema locks
-        self._close_execution_engines()
-
         engine, cleanup = self._get_engine()
         with engine.connect() as conn:
             for table in self.tables: