From d1be5082b5f47549307e53fc89a08ab44f2c546a Mon Sep 17 00:00:00 2001 From: Dima Anfimov Date: Mon, 14 Sep 2026 01:59:20 +0200 Subject: [PATCH] feat: strip sqlalchemy-like dsn prefix --- pyproject.toml | 4 ++ src/taskiq_pg/_internal/broker.py | 17 +---- src/taskiq_pg/_internal/result_backend.py | 15 +---- src/taskiq_pg/_internal/schedule_source.py | 16 ++--- src/taskiq_pg/_internal/utils.py | 17 +++++ src/taskiq_pg/asyncpg/result_backend.py | 4 ++ tests/integration/test_dsn_helper.py | 75 ++++++++++++++++++++++ 7 files changed, 110 insertions(+), 38 deletions(-) create mode 100644 src/taskiq_pg/_internal/utils.py create mode 100644 tests/integration/test_dsn_helper.py diff --git a/pyproject.toml b/pyproject.toml index b64f585..c71a952 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,6 +101,10 @@ markers = [ "unit: marks unit tests", "integration: marks tests with real infrastructure env", ] +filterwarnings = [ + # psqlpy upstream shutdown problem + "ignore:coroutine 'PSQLPyBroker._notification_handler' was never awaited:RuntimeWarning", +] [tool.coverage.report] exclude_lines = [ diff --git a/src/taskiq_pg/_internal/broker.py b/src/taskiq_pg/_internal/broker.py index 4eb6e1b..8cab583 100644 --- a/src/taskiq_pg/_internal/broker.py +++ b/src/taskiq_pg/_internal/broker.py @@ -3,6 +3,8 @@ from taskiq import AsyncBroker, AsyncResultBackend +from taskiq_pg._internal.utils import DsnHelper + if tp.TYPE_CHECKING: import asyncio @@ -11,7 +13,7 @@ _T = tp.TypeVar("_T") -class BasePostgresBroker(AsyncBroker, abc.ABC): +class BasePostgresBroker(AsyncBroker, DsnHelper, abc.ABC): """Base class for Postgres brokers.""" def __init__( @@ -49,16 +51,3 @@ def __init__( self.write_kwargs: dict[str, tp.Any] = write_kwargs or {} self.max_retry_attempts: int = max_retry_attempts self._queue: asyncio.Queue[str] | None = None - - @property - def dsn(self) -> str: - """ - Get the DSN string. - - Returns: - A string with dsn or None if dsn isn't set yet. - - """ - if callable(self._dsn): - return self._dsn() - return self._dsn diff --git a/src/taskiq_pg/_internal/result_backend.py b/src/taskiq_pg/_internal/result_backend.py index ea3a5e0..bc72bce 100644 --- a/src/taskiq_pg/_internal/result_backend.py +++ b/src/taskiq_pg/_internal/result_backend.py @@ -5,11 +5,13 @@ from taskiq.abc.serializer import TaskiqSerializer from taskiq.serializers import PickleSerializer +from taskiq_pg._internal.utils import DsnHelper + ReturnType = tp.TypeVar("ReturnType") -class BasePostgresResultBackend(AsyncResultBackend[ReturnType], abc.ABC): +class BasePostgresResultBackend(AsyncResultBackend[ReturnType], DsnHelper, abc.ABC): """Base class for PostgreSQL result backends.""" def __init__( @@ -38,14 +40,3 @@ def __init__( self.field_for_task_id: tp.Final = field_for_task_id self.connect_kwargs: tp.Final = connect_kwargs self.serializer = serializer or PickleSerializer() - - @property - def dsn(self) -> str | None: - """ - Get the DSN string. - - Returns the DSN string or None if not set. - """ - if callable(self._dsn): - return self._dsn() - return self._dsn diff --git a/src/taskiq_pg/_internal/schedule_source.py b/src/taskiq_pg/_internal/schedule_source.py index 4c438b6..499fc74 100644 --- a/src/taskiq_pg/_internal/schedule_source.py +++ b/src/taskiq_pg/_internal/schedule_source.py @@ -1,3 +1,4 @@ +import abc import typing as tp import uuid from logging import getLogger @@ -7,11 +8,13 @@ from taskiq.abc.broker import AsyncBroker from taskiq.scheduler.scheduled_task import ScheduledTask +from taskiq_pg._internal.utils import DsnHelper + logger = getLogger("taskiq_pg") -class BasePostgresScheduleSource(ScheduleSource): +class BasePostgresScheduleSource(ScheduleSource, DsnHelper, abc.ABC): def __init__( self, broker: AsyncBroker, @@ -38,17 +41,6 @@ def __init__( self._table_name: tp.Final = table_name self._connect_kwargs: tp.Final = connect_kwargs - @property - def dsn(self) -> str | None: - """ - Get the DSN string. - - Returns the DSN string or None if not set. - """ - if callable(self._dsn): - return self._dsn() - return self._dsn - def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]: """ Extract schedules from tasks that were registered in broker. diff --git a/src/taskiq_pg/_internal/utils.py b/src/taskiq_pg/_internal/utils.py new file mode 100644 index 0000000..9d3f64b --- /dev/null +++ b/src/taskiq_pg/_internal/utils.py @@ -0,0 +1,17 @@ +import re +import typing as tp + + +class DsnHelper: + _sqlalchemy_dialect_suffix_re: tp.ClassVar[re.Pattern[str]] = re.compile(r"^(postgres(?:ql)?)\+[^:]+(://)") + + def _preformat_dsn(self, raw_dsn: str) -> str: + """Method to prepare the DSN string (strips the SQLAlchemy-style driver suffix if present).""" + return self._sqlalchemy_dialect_suffix_re.sub(r"\1\2", raw_dsn, count=1) + + @property + def dsn(self) -> str: + """Get the DSN string.""" + if callable(self._dsn): # type: ignore[attr-defined] + return self._preformat_dsn(self._dsn()) # type: ignore[attr-defined] + return self._preformat_dsn(self._dsn) # type: ignore[attr-defined] diff --git a/src/taskiq_pg/asyncpg/result_backend.py b/src/taskiq_pg/asyncpg/result_backend.py index 73573c4..73b7452 100644 --- a/src/taskiq_pg/asyncpg/result_backend.py +++ b/src/taskiq_pg/asyncpg/result_backend.py @@ -1,3 +1,4 @@ +import logging import typing as tp import asyncpg @@ -10,6 +11,9 @@ from taskiq_pg.asyncpg import queries +logger = logging.getLogger("taskiq.asyncpg_result_backend") + + class AsyncpgResultBackend(BasePostgresResultBackend): """Result backend for TaskIQ based on asyncpg.""" diff --git a/tests/integration/test_dsn_helper.py b/tests/integration/test_dsn_helper.py new file mode 100644 index 0000000..a437887 --- /dev/null +++ b/tests/integration/test_dsn_helper.py @@ -0,0 +1,75 @@ +import uuid + +import pytest + +from taskiq_pg.aiopg import AiopgResultBackend +from taskiq_pg.asyncpg import AsyncpgBroker, AsyncpgResultBackend +from taskiq_pg.psqlpy import PSQLPyBroker, PSQLPyResultBackend +from taskiq_pg.psycopg import PsycopgBroker, PsycopgResultBackend + + +@pytest.mark.integration +@pytest.mark.parametrize( + "dialect_suffix", + [ + "postgres+asyncpg", + "postgresql+asyncpg", + "postgres+psycopg", + "postgresql+psqlpy", + ], +) +@pytest.mark.parametrize( + "broker_class", + [ + AsyncpgBroker, + PSQLPyBroker, + PsycopgBroker, + ], +) +async def test_when_dsn_has_sqlalchemy_style_driver_suffix__then_broker_still_connects( + pg_dsn: str, + broker_class: type[AsyncpgBroker | PSQLPyBroker | PsycopgBroker], + dialect_suffix: str, +) -> None: + dsn_with_driver_suffix = pg_dsn.replace("postgres://", f"{dialect_suffix}://", 1) + broker = broker_class(dsn=dsn_with_driver_suffix) + + try: + await broker.startup() + finally: + await broker.shutdown() + + +@pytest.mark.integration +@pytest.mark.parametrize( + "dialect_suffix", + [ + "postgres+asyncpg", + "postgresql+asyncpg", + "postgres+psycopg", + "postgresql+psqlpy", + "postgres+aiopg", + ], +) +@pytest.mark.parametrize( + "result_backend_class", + [ + AsyncpgResultBackend, + AiopgResultBackend, + PSQLPyResultBackend, + PsycopgResultBackend, + ], +) +async def test_when_dsn_has_sqlalchemy_style_driver_suffix__then_result_backend_still_connects( + pg_dsn: str, + result_backend_class: type[AsyncpgResultBackend | AiopgResultBackend | PSQLPyResultBackend | PsycopgResultBackend], + dialect_suffix: str, +) -> None: + dsn_with_driver_suffix = pg_dsn.replace("postgres://", f"{dialect_suffix}://", 1) + table_name = f"taskiq_results_{uuid.uuid4().hex}" + backend = result_backend_class(dsn=dsn_with_driver_suffix, table_name=table_name) + + try: + await backend.startup() + finally: + await backend.shutdown()