Cache server_version_info and guard against re-entrant reads

The SQLAlchemy dialect exposed server_version_info as an un-cached
property whose getter runs "SELECT version()" through the connection.
Because the query executes via the same SQLAlchemy machinery any code
that reads server_version_info while a query is in flight re-enters the
getter and issues another query, recursing forever.

aws-xray-sdk is an example of this pattern.

Cache the resolved value on the dialect instance and seed the cache with
None before issuing the query. A re-entrant read that arrives while the
version query is in flight now returns immediately from the seeded cache
instead of recursing and subsequent reads are served from the cache
without additional HTTP calls.

Author: hashhar

tests/integration/test_sqlalchemy_integration.py

@@ -818,18 +818,27 @@ def test_get_view_names_raises(trino_connection):
 @pytest.mark.skipif(trino_version() == 351, reason="version() not supported in older Trino versions")
 def test_version_is_lazy(trino_connection):
     _, conn = trino_connection
+    # The suite can run against a reused, long-lived cluster. Count queries relative to
+    # what the server already has in its history.
+    baseline = _num_queries_containing_string(conn, "SELECT version()")
     result = conn.execute(sqla.text("SELECT 1"))
     result.fetchall()
     num_queries = _num_queries_containing_string(conn, "SELECT version()")
-    assert num_queries == 0
+    assert num_queries == baseline
     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
+    assert num_queries == baseline + 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 == baseline + 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

@@ -1,3 +1,5 @@
+import threading
+import time
 from typing import Any
 from typing import Dict
 from typing import List
@@ -349,3 +351,109 @@ def test_url_rejects_non_hostname_host(host):
 def test_url_accepts_bare_hostname(host):
     # A valid host must round-trip through make_url without error.
     assert make_url(trino_url(host=host, port=443, user="user")).host is not None
+
+
+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
+
+
+def test_server_version_info_is_resolved_for_concurrent_readers():
+    # The re-entrancy guard must be scoped to the thread doing the lookup. A reader on another
+    # thread has to wait for the real version instead of observing the in-flight placeholder.
+    dialect = TrinoDialect()
+    query_in_flight = threading.Event()
+    release_query = threading.Event()
+
+    class FakeResult:
+        def __init__(self, value):
+            self._value = value
+
+        def scalar(self):
+            return self._value
+
+    class FakeConnection:
+        def __init__(self):
+            self.execute_count = 0
+
+        def execute(self, *args, **kwargs):
+            self.execute_count += 1
+            query_in_flight.set()
+            release_query.wait(timeout=10)
+            return FakeResult("455")
+
+    fake_conn = FakeConnection()
+    dialect._get_server_version_info(fake_conn)
+
+    results = {}
+
+    def read(name):
+        results[name] = dialect.server_version_info
+
+    first = threading.Thread(target=read, args=("first",))
+    second = threading.Thread(target=read, args=("second",))
+
+    first.start()
+    assert query_in_flight.wait(timeout=10)
+    second.start()
+    # Give the second thread a chance to reach the getter while the query is still in flight.
+    time.sleep(0.1)
+    release_query.set()
+
+    first.join(timeout=10)
+    second.join(timeout=10)
+
+    assert results == {"first": ("455",), "second": ("455",)}
+    assert fake_conn.execute_count == 1
+
+
+def test_server_version_info_handles_orig_without_message_attribute():
+    # Only TrinoQueryError exposes `.message`. The DBAPI error wrapped by SQLAlchemy could be
+    # any exception. Failing to read the version must not raise AttributeError from the log call.
+    dialect = TrinoDialect()
+
+    class FailingConnection:
+        def execute(self, *args, **kwargs):
+            raise exc.ProgrammingError("SELECT version()", None, Exception("boom"))
+
+    dialect._get_server_version_info(FailingConnection())
+
+    assert dialect.server_version_info is None

trino/sqlalchemy/dialect.py

@@ -10,6 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import threading
 from collections.abc import Mapping
 from collections.abc import Sequence
 from textwrap import dedent
@@ -423,15 +424,35 @@ 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(_):
-            query = "SELECT version()"
-            try:
-                res = connection.execute(sql.text(query))
-                version = res.scalar()
-                return tuple([version])
-            except exc.ProgrammingError as e:
-                logger.debug(f"Failed to get server version: {e.orig.message}")
+        lock = threading.Lock()
+        resolving = threading.local()
+
+        def get_server_version_info(self):
+            if "_trino_server_version_info" in self.__dict__:
+                return self.__dict__["_trino_server_version_info"]
+            # The version query goes through the same SQLAlchemy machinery as any other query. A caller
+            # that reads this property before every execute re-enters the getter while that query is
+            # still in flight. Recursing there would never terminate, so the re-entrant read returns
+            # None instead. Reads from other threads are not re-entrant. Those wait on the lock and
+            # get the real value.
+            if getattr(resolving, "in_progress", False):
                 return None
+            with lock:
+                if "_trino_server_version_info" in self.__dict__:
+                    return self.__dict__["_trino_server_version_info"]
+                resolving.in_progress = True
+                query = "SELECT version()"
+                try:
+                    res = connection.execute(sql.text(query))
+                    version = res.scalar()
+                    value = tuple([version])
+                except exc.ProgrammingError as e:
+                    logger.debug("Failed to get server version: %s", e)
+                    value = None
+                finally:
+                    resolving.in_progress = False
+                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)