Skip to content

Commit

Permalink
fix: sqlalchemy server version info without request
Browse files Browse the repository at this point in the history
  • Loading branch information
ovsds committed Dec 5, 2023
1 parent 5836008 commit 88600cc
Show file tree
Hide file tree
Showing 6 changed files with 74 additions and 3 deletions.
11 changes: 11 additions & 0 deletions lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Empty file.
Original file line number Diff line number Diff line change
@@ -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)
Original file line number Diff line number Diff line change
@@ -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
11 changes: 11 additions & 0 deletions lib/dl_sqlalchemy_postgres/docker-compose.yml
Original file line number Diff line number Diff line change
@@ -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"
7 changes: 4 additions & 3 deletions lib/dl_sqlalchemy_postgres/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 88600cc

Please sign in to comment.