Cache server_version_info and guard against re-entrant reads

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)))

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

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)