Cache server_version_info and guard against re-entrant reads

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)