Fix zombie $partitions queries left running by SQLAlchemy reflection

tests/integration/test_sqlalchemy_integration.py

@@ -10,6 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 import math
+import time
 import uuid
 from decimal import Decimal
 
@@ -788,6 +789,34 @@ def _num_queries_containing_string(connection, query_string):
     return len(list(filter(lambda rec: query_string in rec[0], rows)))
 
 
+def _assert_no_lingering_query(cur, query_substring, timeout_seconds=10):
+    # Cancellation moves a query to a terminal state asynchronously, so poll.
+    non_terminal_states = {"QUEUED", "PLANNING", "STARTING", "RUNNING", "FINISHING", "WAITING_FOR_RESOURCES"}
+    deadline = time.time() + timeout_seconds
+    last_state = None
+    while time.time() < deadline:
+        # The "?" is rendered as a literal into this statement's own text. The second
+        # predicate excludes it, or the query would always match its own running row.
+        cur.execute(
+            "SELECT state FROM system.runtime.queries "
+            "WHERE query LIKE ? AND query NOT LIKE '%system.runtime.queries%' "
+            "ORDER BY created DESC LIMIT 1",
+            (f"%{query_substring}%",),
+        )
+        rows = cur.fetchall()
+        if rows:
+            last_state = rows[0][0]
+            if last_state not in non_terminal_states:
+                return
+        time.sleep(0.5)
+    if last_state is None:
+        pytest.fail(f"no query matching {query_substring!r} appeared in system.runtime.queries")
+    pytest.fail(
+        f"query matching {query_substring!r} did not reach a terminal state within "
+        f"{timeout_seconds}s (last observed state: {last_state})"
+    )
+
+
 @pytest.mark.skipif(trino_version() == 351, reason="Dynamic catalogs not supported")
 def test_get_indexes_returns_empty_for_iceberg_table(run_trino):
     host, port = run_trino
@@ -824,6 +853,7 @@ def test_get_indexes_returns_empty_for_iceberg_table(run_trino):
         )
         indexes = sqla.inspect(engine).get_indexes(table_name, schema=schema_name)
         assert indexes == []
+        _assert_no_lingering_query(cur, f'{table_name}$partitions')
     finally:
         cur = conn.cursor()
         cur.execute(f"DROP TABLE IF EXISTS {catalog_name}.{schema_name}.{table_name}")
@@ -873,6 +903,7 @@ def test_get_indexes_returns_partitions_for_hive_table(run_trino):
         assert len(indexes) == 1
         assert indexes[0]["name"] == "partition"
         assert indexes[0]["column_names"] == ["name", "region"]
+        _assert_no_lingering_query(cur, f'{table_name}$partitions')
     finally:
         cur = conn.cursor()
         cur.execute(f"DROP TABLE IF EXISTS {catalog_name}.{schema_name}.{table_name}")

tests/unit/sqlalchemy/test_dialect.py

@@ -258,6 +258,74 @@ def test_isolation_level(self):
         isolation_level = self.dialect.get_isolation_level(dbapi_conn)
         assert isolation_level == "SERIALIZABLE"
 
+    class _FakeCursor:
+        def __init__(self, description):
+            self.description = description
+
+    class _FakeResult:
+        """Hand-written fake standing in for a SQLAlchemy CursorResult."""
+
+        def __init__(self, partition_names, data_types):
+            self.cursor = TestTrinoDialect._FakeCursor(list(zip(partition_names, data_types)))
+            self.close_calls = 0
+
+        def close(self):
+            self.close_calls += 1
+
+    @staticmethod
+    def _partitions_connection(partition_names, data_types):
+        result = TestTrinoDialect._FakeResult(partition_names, data_types)
+        connection = mock.Mock()
+        connection.execute.return_value = result
+        return connection, result
+
+    def test_get_partitions_hive_table_returns_partition_names(self):
+        partition_names = ["year", "month"]
+        data_types = ["integer", "integer"]
+        connection, result = self._partitions_connection(partition_names, data_types)
+
+        returned = self.dialect._get_partitions(connection, "some_table", "some_schema")
+
+        assert returned == partition_names
+        assert result.close_calls == 1
+
+    def test_get_partitions_iceberg_table_returns_none(self):
+        # This is the exact shape of an Iceberg $partitions table.
+        partition_names = ["partition", "record_count", "file_count", "total_size", "data"]
+        data_types = ["row(...)", "bigint", "bigint", "bigint", "row(...)"]
+        connection, result = self._partitions_connection(partition_names, data_types)
+
+        returned = self.dialect._get_partitions(connection, "some_table", "some_schema")
+
+        assert returned is None
+        assert result.close_calls == 1
+
+    def test_get_partitions_closes_result_even_on_error(self):
+        connection = mock.Mock()
+        result = mock.Mock()
+        result.cursor.description = None  # triggers a TypeError while iterating
+        connection.execute.return_value = result
+
+        with pytest.raises(TypeError):
+            self.dialect._get_partitions(connection, "some_table", "some_schema")
+
+        result.close.assert_called_once()
+
+    def test_get_partitions_defaults_schema_when_not_given(self):
+        partition_names = ["year"]
+        data_types = ["integer"]
+        connection, _ = self._partitions_connection(partition_names, data_types)
+
+        with mock.patch.object(
+            self.dialect, "_get_default_schema_name", return_value="default_schema"
+        ) as get_default_schema:
+            self.dialect._get_partitions(connection, "some_table")
+
+        get_default_schema.assert_called_once_with(connection)
+        # The defaulted schema must appear in the query.
+        query = str(connection.execute.call_args[0][0])
+        assert "default_schema" in query
+
 
 def test_trino_connection_basic_auth():
     dialect = TrinoDialect()

trino/sqlalchemy/dialect.py

@@ -223,9 +223,14 @@ def _get_partitions(
             SELECT * FROM {schema}."{table_name}$partitions"
         """
         ).strip()
+
         res = connection.execute(sql.text(query))
-        partition_names = [desc[0] for desc in res.cursor.description]
-        data_types = [desc[1] for desc in res.cursor.description]
+        try:
+            partition_names = [desc[0] for desc in res.cursor.description]
+            data_types = [desc[1] for desc in res.cursor.description]
+        finally:
+            res.close()
+
         # Compare the column names and types to the shape of an Iceberg $partitions table
         if (partition_names == ['partition', 'record_count', 'file_count', 'total_size', 'data']
                 and data_types[0].startswith('row(')