feat: add helper to format dsn before passing it to driver

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 = [

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

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

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.

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]

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."""
 

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()