Feat: Re-introduce merge for updating auto restatements

Author: vchan

sqlmesh/core/engine_adapter/postgres.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import re
 import typing as t
 from functools import partial
 from sqlglot import exp
@@ -112,11 +113,8 @@ def merge(
         **kwargs: t.Any,
     ) -> None:
         # Merge isn't supported until Postgres 15
-        merge_impl = (
-            super().merge
-            if self._connection_pool.get().server_version >= 150000
-            else partial(logical_merge, self)
-        )
+        major, minor = self.get_server_version()
+        merge_impl = super().merge if major >= 15 else partial(logical_merge, self)
         merge_impl(  # type: ignore
             target_table,
             source_table,
@@ -125,3 +123,23 @@ def merge(
             when_matched=when_matched,
             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
+        return 0, 0

tests/core/engine_adapter/test_postgres.py

@@ -94,7 +94,8 @@ 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().server_version = 150000
+    adapter._connection_pool.get().__class__.__module__ = "psycopg2.extensions"
+    adapter._connection_pool.get().info.server_version = 150000
 
     adapter.merge(
         target_table="target",
@@ -117,7 +118,8 @@ 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().server_version = 140000
+    adapter._connection_pool.get().__class__.__module__ = "psycopg2.extensions"
+    adapter._connection_pool.get().info.server_version = 140000
 
     temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
     table_name = "test"
@@ -161,3 +163,17 @@ def table_columns(table_name: str) -> t.Dict[str, exp.DataType]:
     assert to_sql_calls(adapter) == [
         'ALTER TABLE "test_table" DROP COLUMN "test_column" CASCADE',
     ]
+
+
+def test_get_server_version(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 = 150013
+    assert adapter.get_server_version() == (15, 13)
+
+    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)