Show server_version

Author: vchan

sqlmesh/core/engine_adapter/postgres.py

@@ -3,7 +3,7 @@
 import logging
 import re
 import typing as t
-from functools import partial
+from functools import cached_property, partial
 from sqlglot import exp
 
 from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
@@ -113,7 +113,7 @@ def merge(
         **kwargs: t.Any,
     ) -> None:
         # Merge isn't supported until Postgres 15
-        major, minor = self.get_server_version()
+        major, minor = self.server_version
         merge_impl = super().merge if major >= 15 else partial(logical_merge, self)
         merge_impl(  # type: ignore
             target_table,
@@ -124,22 +124,11 @@ def merge(
             merge_filter=merge_filter,
         )
 
-    def get_server_version(self) -> t.Tuple[int, int]:
-        """Return major and minor server versions of the connection"""
-        connection = self._connection_pool.get()
-        connection_module = connection.__class__.__module__
-        if connection_module.startswith("pg8000"):
-            server_version = connection.parameter_statuses.get("server_version")
-            # pg8000 server version contains version as well as packaging and distribution information
-            # e.g. 15.13 (Debian 15.13-1.pgdg120+1)
-            match = re.search(r"(\d+)\.(\d+)", server_version)
-            if match:
-                return int(match.group(1)), int(match.group(2))
-        elif connection_module.startswith("psycopg"):
-            # This handles both psycopg and psycopg2 connection objects
-            server_version = connection.info.server_version
-            # Since major version 10, PostgreSQL represents the server version with an integer by
-            # multiplying the server's major version number by 10000 and adding the minor version number
-            # See https://www.postgresql.org/docs/current/libpq-status.html#LIBPQ-PQSERVERVERSION
-            return server_version // 10000, server_version % 100
+    @cached_property
+    def server_version(self) -> t.Tuple[int, int]:
+        """Lazily fetch and cache major and minor server version"""
+        server_version, *_ = self.fetchone("SHOW server_version")
+        match = re.search(r"(\d+)\.(\d+)", server_version)
+        if match:
+            return int(match.group(1)), int(match.group(2))
         return 0, 0

tests/core/engine_adapter/test_postgres.py

@@ -94,8 +94,7 @@ def test_create_table_like(make_mocked_engine_adapter: t.Callable):
 
 def test_merge_version_gte_15(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
-    adapter._connection_pool.get().__class__.__module__ = "psycopg2.extensions"
-    adapter._connection_pool.get().info.server_version = 150000
+    adapter.server_version = (15, 0)
 
     adapter.merge(
         target_table="target",
@@ -118,8 +117,7 @@ def test_merge_version_lt_15(
     make_mocked_engine_adapter: t.Callable, make_temp_table_name: t.Callable, mocker: MockerFixture
 ):
     adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
-    adapter._connection_pool.get().__class__.__module__ = "psycopg2.extensions"
-    adapter._connection_pool.get().info.server_version = 140000
+    adapter.server_version = (14, 0)
 
     temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
     table_name = "test"
@@ -165,15 +163,17 @@ def table_columns(table_name: str) -> t.Dict[str, exp.DataType]:
     ]
 
 
-def test_get_server_version(make_mocked_engine_adapter: t.Callable):
+def test_server_version(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
 
-    adapter._connection_pool.get().__class__.__module__ = "psycopg2.extensions"
-    adapter._connection_pool.get().info.server_version = 150013
-    assert adapter.get_server_version() == (15, 13)
+    fetchone_mock = mocker.patch.object(adapter, "fetchone")
+    fetchone_mock.return_value = ("14.0",)
+    assert adapter.server_version == (14, 0)
 
-    adapter._connection_pool.get().__class__.__module__ = "pg8000.native"
-    adapter._connection_pool.get().parameter_statuses = {
-        "server_version": "15.13 (Debian 15.13-1.pgdg120+1)"
-    }
-    assert adapter.get_server_version() == (15, 13)
+    del adapter.server_version
+    fetchone_mock.return_value = ("15.8",)
+    assert adapter.server_version == (15, 8)
+
+    del adapter.server_version
+    fetchone_mock.return_value = ("15.13 (Debian 15.13-1.pgdg120+1)",)
+    assert adapter.server_version == (15, 13)