From 8fa5d948b6601b99e530ee51cfe6a7045b1ddcb0 Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Fri, 14 Aug 2026 18:11:10 +0530 Subject: [PATCH] Fix zombie $partitions queries left running by SQLAlchemy reflection _get_partitions never closes the query result after reading cursor.description, so the underlying query is never cancelled and lingers on the coordinator until it eventually times out on its own. Close the result in a finally block so it's always cancelled, whether or not reading cursor.description succeeds. Co-authored-by: viniolivieri --- .../test_sqlalchemy_integration.py | 31 +++++++++ tests/unit/sqlalchemy/test_dialect.py | 68 +++++++++++++++++++ trino/sqlalchemy/dialect.py | 9 ++- 3 files changed, 106 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py index 8e6d9cdc..f153834e 100644 --- a/tests/integration/test_sqlalchemy_integration.py +++ b/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}") diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py index d247a536..5aafca57 100644 --- a/tests/unit/sqlalchemy/test_dialect.py +++ b/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() diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py index a3191e03..5becaa84 100644 --- a/trino/sqlalchemy/dialect.py +++ b/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(')