From 88600cc9895d59b37e66d5c3903fefecb3bbc947 Mon Sep 17 00:00:00 2001
From: ovsds <ovsds@yandex-team.ru>
Date: Thu, 30 Nov 2023 13:24:09 +0100
Subject: [PATCH] fix: sqlalchemy server version info without request

---
 .../dl_sqlalchemy_postgres/base.py            | 11 ++++++
 .../db/__init__.py                            |  0
 .../db/conftest.py                            | 14 ++++++++
 .../db/test_server_version.py                 | 34 +++++++++++++++++++
 lib/dl_sqlalchemy_postgres/docker-compose.yml | 11 ++++++
 lib/dl_sqlalchemy_postgres/pyproject.toml     |  7 ++--
 6 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/__init__.py
 create mode 100644 lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/conftest.py
 create mode 100644 lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/test_server_version.py
 create mode 100644 lib/dl_sqlalchemy_postgres/docker-compose.yml

diff --git a/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py
index 229e9ceea..f3a4c30fc 100644
--- a/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py
+++ b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py
@@ -122,6 +122,17 @@ def __init__(self, enforce_collate=None, **kwargs):
     type_compiler = BICustomPGTypeCompiler
     statement_compiler = BIPGCompilerBasic
     ischema_names = bi_pg_ischema_names
+    forced_server_version_string: str | None = None
+
+    def connect(self, *cargs, **cparams):
+        self.forced_server_version_string = cparams.pop("server_version", self.forced_server_version_string)
+        return super().connect(*cargs, **cparams)
+
+    def _get_server_version_info(self, connection) -> tuple[int, ...]:
+        if self.forced_server_version_string is not None:
+            return tuple(int(part) for part in self.forced_server_version_string.split("."))
+
+        return super()._get_server_version_info(connection)
 
 
 class BIPGDialect(BIPGDialectBasic):
diff --git a/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/__init__.py b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/conftest.py b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/conftest.py
new file mode 100644
index 000000000..f9847f814
--- /dev/null
+++ b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/conftest.py
@@ -0,0 +1,14 @@
+import pytest
+import sqlalchemy
+
+from dl_testing.containers import get_test_container_hostport
+
+
+@pytest.fixture(scope="session")
+def engine_url():
+    return f'bi_postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres", fallback_port=52301).as_pair()}/test_data'
+
+
+@pytest.fixture(scope="session")
+def sa_engine(engine_url):
+    return sqlalchemy.create_engine(engine_url)
diff --git a/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/test_server_version.py b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/test_server_version.py
new file mode 100644
index 000000000..fb7c13d9d
--- /dev/null
+++ b/lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/db/test_server_version.py
@@ -0,0 +1,34 @@
+import mock
+import sqlalchemy
+import sqlalchemy.dialects.postgresql.psycopg2 as sqlalchemy_dialect_psycopg2
+import sqlalchemy.orm as sqlalchemy_orm
+
+
+SERVER_VERSION_INFO = (123, 45, 67, 89)
+SERVER_VERSION = ".".join(map(str, SERVER_VERSION_INFO))
+
+
+@mock.patch.object(sqlalchemy_dialect_psycopg2.PGDialect_psycopg2, "_get_server_version_info")
+def test_server_version_default(patched_server_version: mock.Mock, engine_url: str):
+    patched_server_version.return_value = SERVER_VERSION_INFO
+
+    engine = sqlalchemy.create_engine(engine_url)
+    session_maker = sqlalchemy_orm.sessionmaker(bind=engine)
+    session = session_maker()
+
+    session.scalar("select 1")
+
+    patched_server_version.assert_called_once()
+    assert session.get_bind().dialect.server_version_info == SERVER_VERSION_INFO
+
+
+@mock.patch.object(sqlalchemy_dialect_psycopg2.PGDialect_psycopg2, "_get_server_version_info")
+def test_server_version_overwritten(patched_server_version: mock.Mock, engine_url: str):
+    engine = sqlalchemy.create_engine(engine_url, connect_args=dict(server_version=SERVER_VERSION))
+    session_maker = sqlalchemy_orm.sessionmaker(bind=engine)
+    session = session_maker()
+
+    session.scalar("select 1")
+
+    patched_server_version.assert_not_called()
+    assert session.get_bind().dialect.server_version_info == SERVER_VERSION_INFO
diff --git a/lib/dl_sqlalchemy_postgres/docker-compose.yml b/lib/dl_sqlalchemy_postgres/docker-compose.yml
new file mode 100644
index 000000000..6cf281e41
--- /dev/null
+++ b/lib/dl_sqlalchemy_postgres/docker-compose.yml
@@ -0,0 +1,11 @@
+version: '3.7'
+
+services:
+  db-postgres:
+    image: "postgres:13-alpine@sha256:b9f66c57932510574fb17bccd175776535cec9abcfe7ba306315af2f0b7bfbb4"
+    environment:
+      - POSTGRES_DB=test_data
+      - POSTGRES_USER=datalens
+      - POSTGRES_PASSWORD=qwerty
+    ports:
+     - "50319:5432"
diff --git a/lib/dl_sqlalchemy_postgres/pyproject.toml b/lib/dl_sqlalchemy_postgres/pyproject.toml
index 06342858c..ca001f962 100644
--- a/lib/dl_sqlalchemy_postgres/pyproject.toml
+++ b/lib/dl_sqlalchemy_postgres/pyproject.toml
@@ -28,9 +28,10 @@ build-backend = "poetry.core.masonry.api"
 [tool.pytest.ini_options]
 minversion = "6.0"
 addopts = "-ra"
-testpaths = ["dl_sqlalchemy_postgres_tests/unit"]
-
-
+testpaths = [
+    "dl_sqlalchemy_postgres_tests/unit",
+    "dl_sqlalchemy_postgres_tests/db",
+]
 
 [tool.mypy]
 warn_unused_configs = true