Skip to content

Commit 7184f3a

Browse files
authored
fix: broker shutdown before connection pool creation (#6)
2 parents be4f0c4 + 8f460ff commit 7184f3a

File tree

4 files changed

+8
-11
lines changed

4 files changed

+8
-11
lines changed

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ dependencies = [
3939
"Bug Tracker" = "https://github.com/danfimov/taskiq-postgres/issues"
4040
"Repository" = "https://github.com/danfimov/taskiq-postgres/"
4141

42-
[dependency-groups]
42+
[project.optional-dependencies]
4343
aiopg = [
4444
"aiopg>=1.4.0",
4545
]
@@ -49,6 +49,7 @@ asyncpg = [
4949
psqlpy = [
5050
"psqlpy>=0.10.1",
5151
]
52+
[dependency-groups]
5253
dev = [
5354
"ruff>=0.11.9",
5455
"pytest>=8.3.5",

src/taskiq_pg/aiopg/result_backend.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,8 @@ async def startup(self) -> None:
7676

7777
async def shutdown(self) -> None:
7878
"""Close the connection pool."""
79-
async with self._database_pool.acquire() as connection:
80-
await connection.close()
79+
if getattr(self, "_database_pool", None) is not None:
80+
self._database_pool.close()
8181

8282
async def set_result(
8383
self,

src/taskiq_pg/asyncpg/result_backend.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
from taskiq import AsyncResultBackend, TaskiqResult
77
from taskiq.compat import model_dump, model_validate
88
from taskiq.serializers import PickleSerializer
9-
from typing_extensions import override
109

1110
from taskiq_pg.asyncpg.queries import (
1211
CREATE_INDEX_QUERY,
@@ -67,7 +66,6 @@ def dsn(self) -> str | None:
6766
return self._dsn()
6867
return self._dsn
6968

70-
@override
7169
async def startup(self) -> None:
7270
"""
7371
Initialize the result backend.
@@ -96,12 +94,11 @@ async def startup(self) -> None:
9694
),
9795
)
9896

99-
@override
10097
async def shutdown(self) -> None:
10198
"""Close the connection pool."""
102-
await self._database_pool.close()
99+
if getattr(self, "_database_pool", None) is not None:
100+
await self._database_pool.close()
103101

104-
@override
105102
async def set_result(
106103
self,
107104
task_id: str,
@@ -121,7 +118,6 @@ async def set_result(
121118
self.serializer.dumpb(model_dump(result)),
122119
)
123120

124-
@override
125121
async def is_result_ready(self, task_id: str) -> bool:
126122
"""
127123
Returns whether the result is ready.
@@ -139,7 +135,6 @@ async def is_result_ready(self, task_id: str) -> bool:
139135
),
140136
)
141137

142-
@override
143138
async def get_result(
144139
self,
145140
task_id: str,

src/taskiq_pg/psqlpy/result_backend.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,8 @@ async def startup(self) -> None:
8484

8585
async def shutdown(self) -> None:
8686
"""Close the connection pool."""
87-
self._database_pool.close()
87+
if getattr(self, "_database_pool", None) is not None:
88+
self._database_pool.close()
8889

8990
async def set_result(
9091
self,

0 commit comments

Comments
 (0)