diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py index e5160bb0..3b5688b9 100644 --- a/tests/integration/test_sqlalchemy_integration.py +++ b/tests/integration/test_sqlalchemy_integration.py @@ -751,10 +751,16 @@ def test_version_is_lazy(trino_connection): assert isinstance(version_info, tuple) num_queries = _num_queries_containing_string(conn, "SELECT version()") assert num_queries == 1 + # Reading server_version_info again should be served from the cache and must + # not issue another SELECT version() query. + version_info = conn.dialect.server_version_info + assert isinstance(version_info, tuple) + num_queries = _num_queries_containing_string(conn, "SELECT version()") + assert num_queries == 1 def _num_queries_containing_string(connection, query_string): - statement = sqla.text("select query from system.runtime.queries order by query_id desc offset 1 limit 1") + statement = sqla.text("select query from system.runtime.queries") result = connection.execute(statement) rows = result.fetchall() return len(list(filter(lambda rec: query_string in rec[0], rows))) diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py index d247a536..ff90ee0b 100644 --- a/tests/unit/sqlalchemy/test_dialect.py +++ b/tests/unit/sqlalchemy/test_dialect.py @@ -314,3 +314,46 @@ def test_trino_connection_oauth2_auth(): _, cparams = dialect.create_connect_args(url) assert isinstance(cparams['auth'], OAuth2Authentication) + + +def test_server_version_info_does_not_recurse_and_is_cached(): + # Regression test for https://github.com/trinodb/trino-python-client/issues/559 + # + # aws-xray-sdk reads `dialect.server_version_info` before every execute. Since + # computing that property itself issues a `SELECT version()` query through + # `connection.execute`, a naive implementation causes aws-xray to read the + # property again while the query is in flight, recursing forever. + dialect = TrinoDialect() + + class FakeResult: + def __init__(self, value): + self._value = value + + def scalar(self): + return self._value + + class FakeConnection: + def __init__(self, dialect): + self.dialect = dialect + self.execute_count = 0 + + def execute(self, *args, **kwargs): + # Mimic aws-xray-sdk's re-entrant read of server_version_info before + # every execute call. + self.dialect.server_version_info + self.execute_count += 1 + return FakeResult("455") + + fake_conn = FakeConnection(dialect) + + dialect._get_server_version_info(fake_conn) + + version_info = dialect.server_version_info + assert version_info == ("455",) + assert fake_conn.execute_count == 1 + + # Subsequent reads should be served from the cache and must not trigger + # another query. + assert dialect.server_version_info == ("455",) + assert dialect.server_version_info == ("455",) + assert fake_conn.execute_count == 1 diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py index ad05aeec..98e05ed2 100644 --- a/trino/sqlalchemy/dialect.py +++ b/trino/sqlalchemy/dialect.py @@ -421,15 +421,23 @@ def has_sequence(self, connection: Connection, sequence_name: str, schema: str = @classmethod def _get_server_version_info(cls, connection: Connection) -> Any: - def get_server_version_info(_): + def get_server_version_info(self): + if "_trino_server_version_info" in self.__dict__: + return self.__dict__["_trino_server_version_info"] + # Seed the cache before running the query so that a re-entrant read + # triggered while the version query is in flight returns immediately + # instead of recursing infinitely. + self.__dict__["_trino_server_version_info"] = None query = "SELECT version()" try: res = connection.execute(sql.text(query)) version = res.scalar() - return tuple([version]) + value = tuple([version]) except exc.ProgrammingError as e: logger.debug(f"Failed to get server version: {e.orig.message}") - return None + value = None + self.__dict__["_trino_server_version_info"] = value + return value # Make server_version_info lazy in order to only make HTTP calls if user explicitly requests it. cls.server_version_info = property(get_server_version_info, lambda instance, value: None)