Перейти к содержанию

taskiq_pg

aiopg

AiopgResultBackend 

AiopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AiopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: Pool,
)

AiopgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on Aiopg.

Construct a new AiopgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (Pool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/aiopg/result_backend.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: Pool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AiopgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )

keep_results instance-attribute 

keep_results = keep_results

table_name instance-attribute 

table_name = table_name

field_for_task_id instance-attribute 

field_for_task_id = field_for_task_id

connect_kwargs instance-attribute 

connect_kwargs = connect_kwargs

serializer instance-attribute 

serializer = serializer or PickleSerializer()

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 

startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/aiopg/result_backend.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table for results if not exists.
    """
    try:
        if self._owns_pool:
            self._database_pool = await create_pool(
                self.dsn,
                **self.connect_kwargs,
            )

        async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
            await cursor.execute(
                queries.CREATE_TABLE_QUERY.format(
                    self.table_name,
                    self.field_for_task_id,
                ),
            )
            await cursor.execute(
                queries.ADD_PROGRESS_COLUMN_QUERY.format(
                    self.table_name,
                    self.field_for_task_id,
                ),
            )
            await cursor.execute(
                queries.CREATE_INDEX_QUERY.format(
                    self.table_name,
                    self.table_name,
                ),
            )
    except Exception as error:
        raise exceptions.DatabaseConnectionError(str(error)) from error

shutdown async 

shutdown()

Close the connection pool if created by this backend.

Source code in src/taskiq_pg/aiopg/result_backend.py 
116
117
118
119
async def shutdown(self) -> None:
    """Close the connection pool if created by this backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()

set_result async 

set_result(task_id, result)

Set result to the PostgreSQL table.

Parameters:

  • task_id (Any) –

    ID of the task.

  • result (TaskiqResult[_ReturnType]) –

    result of the task.

Source code in src/taskiq_pg/aiopg/result_backend.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
async def set_result(
    self,
    task_id: tp.Any,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    Args:
        task_id (Any): ID of the task.
        result (TaskiqResult[_ReturnType]):  result of the task.

    """
    dumped_result = self.serializer.dumpb(result)
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.INSERT_RESULT_QUERY.format(
                self.table_name,
            ),
            (
                task_id,
                dumped_result,
                dumped_result,
            ),
        )

is_result_ready async 

is_result_ready(task_id)

Return whether the result is ready.

Parameters:

  • task_id (Any) –

    ID of the task.

Returns:

  • bool ( bool ) –

    True if the result is ready else False.

Source code in src/taskiq_pg/aiopg/result_backend.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
async def is_result_ready(
    self,
    task_id: tp.Any,
) -> bool:
    """
    Return whether the result is ready.

    Args:
        task_id (Any): ID of the task.

    Returns:
        bool: True if the result is ready else False.

    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        result = await cursor.fetchone()
        return bool(result[0]) if result else False

get_result async 

get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/aiopg/result_backend.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
async def get_result(
    self,
    task_id: tp.Any,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        result = await cursor.fetchone()

        if not result:
            msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
            raise exceptions.ResultIsMissingError(
                msg,
            )

        result_in_bytes: bytes = result[0]

        if not self.keep_results:
            await cursor.execute(
                queries.DELETE_RESULT_QUERY.format(
                    self.table_name,
                ),
                (task_id,),
            )

        taskiq_result: TaskiqResult[ReturnType] = self.serializer.loadb(
            result_in_bytes,
        )

        if not with_logs:
            taskiq_result.log = None

        return taskiq_result

set_progress async 

set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/aiopg/result_backend.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    dumped_progress = self.serializer.dumpb(progress)
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.INSERT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            (
                task_id,
                dumped_progress,
                dumped_progress,
            ),
        )

get_progress async 

get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/aiopg/result_backend.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.SELECT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        progress = await cursor.fetchone()
        if not progress or progress[0] is None:
            return None
        progress_in_bytes: bytes = progress[0]
        taskiq_progress: TaskProgress[tp.Any] = self.serializer.loadb(progress_in_bytes)
        return taskiq_progress

AiopgScheduleSource 

AiopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AiopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: Pool,
)

AiopgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses aiopg to store schedules in PostgreSQL.

Construct a new AiopgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (Pool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: Pool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AiopgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

extract_scheduled_tasks_from_broker 

extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

startup async 

startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for schedules if not exists
    and fill table with schedules from task labels.
    """
    try:
        if self._owns_pool:
            self._database_pool = await create_pool(
                dsn=self.dsn,
                **self._connect_kwargs,
            )
        async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
            await cursor.execute(CREATE_SCHEDULES_TABLE_QUERY.format(self._table_name))
        scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
        await self._update_schedules_on_startup(scheduled_tasks_for_creation)
    except Exception as error:
        raise exceptions.DatabaseConnectionError(str(error)) from error

shutdown async 

shutdown()

Close the connection pool if it was created by this schedule source.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
115
116
117
118
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()

get_schedules async 

get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
async def get_schedules(self) -> list[ScheduledTask]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
        schedules, rows = [], await cursor.fetchall()
    for schedule_id, task_name, schedule in rows:
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(schedule_id),
                    "task_name": task_name,
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules

add_schedule async 

add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
async def add_schedule(self, schedule: ScheduledTask) -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            INSERT_SCHEDULE_QUERY.format(self._table_name),
            [
                schedule.schedule_id,
                schedule.task_name,
                schedule.model_dump_json(
                    exclude={"schedule_id", "task_name"},
                ),
            ],
        )

delete_schedule async 

delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            DELETE_SCHEDULE_QUERY.format(self._table_name),
            [schedule_id],
        )

post_send async 

post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
179
180
181
182
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)

broker

queries

CREATE_TABLE_QUERY module-attribute 

CREATE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    task_id {} UNIQUE,\n    result BYTEA,\n    progress BYTEA\n)\n"

ADD_PROGRESS_COLUMN_QUERY module-attribute 

ADD_PROGRESS_COLUMN_QUERY = "\nALTER TABLE {} ADD COLUMN IF NOT EXISTS progress BYTEA;\n"

CREATE_INDEX_QUERY module-attribute 

CREATE_INDEX_QUERY = "\nCREATE INDEX IF NOT EXISTS {}_task_id_idx ON {} USING HASH (task_id)\n"

INSERT_RESULT_QUERY module-attribute 

INSERT_RESULT_QUERY = "\nINSERT INTO {} VALUES (%s, %s, NULL)\nON CONFLICT (task_id)\nDO UPDATE\nSET result = %s\n"

INSERT_PROGRESS_QUERY module-attribute 

INSERT_PROGRESS_QUERY = "\nINSERT INTO {} VALUES (%s, NULL, %s)\nON CONFLICT (task_id)\nDO UPDATE\nSET progress = %s\n"

SELECT_PROGRESS_QUERY module-attribute 

SELECT_PROGRESS_QUERY = (
    "\nSELECT progress FROM {} WHERE task_id = %s\n"
)

IS_RESULT_EXISTS_QUERY module-attribute 

IS_RESULT_EXISTS_QUERY = "\nSELECT EXISTS(\n    SELECT 1 FROM {} WHERE task_id = %s and result IS NOT NULL\n)\n"

SELECT_RESULT_QUERY module-attribute 

SELECT_RESULT_QUERY = (
    "\nSELECT result FROM {} WHERE task_id = %s\n"
)

DELETE_RESULT_QUERY module-attribute 

DELETE_RESULT_QUERY = (
    "\nDELETE FROM {} WHERE task_id = %s\n"
)

CREATE_SCHEDULES_TABLE_QUERY module-attribute 

CREATE_SCHEDULES_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id UUID PRIMARY KEY,\n    task_name VARCHAR(100) NOT NULL,\n    schedule JSONB NOT NULL,\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),\n    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_SCHEDULE_QUERY module-attribute 

INSERT_SCHEDULE_QUERY = "\nINSERT INTO {} (id, task_name, schedule)\nVALUES (%s, %s, %s)\nON CONFLICT (id) DO UPDATE\nSET task_name = EXCLUDED.task_name,\n    schedule = EXCLUDED.schedule,\n    updated_at = NOW();\n"

SELECT_SCHEDULES_QUERY module-attribute 

SELECT_SCHEDULES_QUERY = (
    "\nSELECT id, task_name, schedule\nFROM {};\n"
)

DELETE_ALL_SCHEDULES_QUERY module-attribute 

DELETE_ALL_SCHEDULES_QUERY = '\nDELETE FROM {};\n'

DELETE_SCHEDULE_QUERY module-attribute 

DELETE_SCHEDULE_QUERY = '\nDELETE FROM {} WHERE id = %s;\n'

result_backend

AiopgResultBackend 

AiopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AiopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: Pool,
)

AiopgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on Aiopg.

Construct a new AiopgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (Pool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/aiopg/result_backend.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: Pool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AiopgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )
keep_results instance-attribute 
keep_results = keep_results
table_name instance-attribute 
table_name = table_name
field_for_task_id instance-attribute 
field_for_task_id = field_for_task_id
connect_kwargs instance-attribute 
connect_kwargs = connect_kwargs
serializer instance-attribute 
serializer = serializer or PickleSerializer()
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/aiopg/result_backend.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table for results if not exists.
    """
    try:
        if self._owns_pool:
            self._database_pool = await create_pool(
                self.dsn,
                **self.connect_kwargs,
            )

        async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
            await cursor.execute(
                queries.CREATE_TABLE_QUERY.format(
                    self.table_name,
                    self.field_for_task_id,
                ),
            )
            await cursor.execute(
                queries.ADD_PROGRESS_COLUMN_QUERY.format(
                    self.table_name,
                    self.field_for_task_id,
                ),
            )
            await cursor.execute(
                queries.CREATE_INDEX_QUERY.format(
                    self.table_name,
                    self.table_name,
                ),
            )
    except Exception as error:
        raise exceptions.DatabaseConnectionError(str(error)) from error
shutdown async 
shutdown()

Close the connection pool if created by this backend.

Source code in src/taskiq_pg/aiopg/result_backend.py 
116
117
118
119
async def shutdown(self) -> None:
    """Close the connection pool if created by this backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()
set_result async 
set_result(task_id, result)

Set result to the PostgreSQL table.

Parameters:

  • task_id (Any) –

    ID of the task.

  • result (TaskiqResult[_ReturnType]) –

    result of the task.

Source code in src/taskiq_pg/aiopg/result_backend.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
async def set_result(
    self,
    task_id: tp.Any,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    Args:
        task_id (Any): ID of the task.
        result (TaskiqResult[_ReturnType]):  result of the task.

    """
    dumped_result = self.serializer.dumpb(result)
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.INSERT_RESULT_QUERY.format(
                self.table_name,
            ),
            (
                task_id,
                dumped_result,
                dumped_result,
            ),
        )
is_result_ready async 
is_result_ready(task_id)

Return whether the result is ready.

Parameters:

  • task_id (Any) –

    ID of the task.

Returns:

  • bool ( bool ) –

    True if the result is ready else False.

Source code in src/taskiq_pg/aiopg/result_backend.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
async def is_result_ready(
    self,
    task_id: tp.Any,
) -> bool:
    """
    Return whether the result is ready.

    Args:
        task_id (Any): ID of the task.

    Returns:
        bool: True if the result is ready else False.

    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        result = await cursor.fetchone()
        return bool(result[0]) if result else False
get_result async 
get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/aiopg/result_backend.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
async def get_result(
    self,
    task_id: tp.Any,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        result = await cursor.fetchone()

        if not result:
            msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
            raise exceptions.ResultIsMissingError(
                msg,
            )

        result_in_bytes: bytes = result[0]

        if not self.keep_results:
            await cursor.execute(
                queries.DELETE_RESULT_QUERY.format(
                    self.table_name,
                ),
                (task_id,),
            )

        taskiq_result: TaskiqResult[ReturnType] = self.serializer.loadb(
            result_in_bytes,
        )

        if not with_logs:
            taskiq_result.log = None

        return taskiq_result
set_progress async 
set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/aiopg/result_backend.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    dumped_progress = self.serializer.dumpb(progress)
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.INSERT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            (
                task_id,
                dumped_progress,
                dumped_progress,
            ),
        )
get_progress async 
get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/aiopg/result_backend.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            queries.SELECT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            (task_id,),
        )
        progress = await cursor.fetchone()
        if not progress or progress[0] is None:
            return None
        progress_in_bytes: bytes = progress[0]
        taskiq_progress: TaskProgress[tp.Any] = self.serializer.loadb(progress_in_bytes)
        return taskiq_progress

schedule_source

logger module-attribute 

logger = getLogger('taskiq_pg.aiopg_schedule_source')

AiopgScheduleSource 

AiopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AiopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: Pool,
)

AiopgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses aiopg to store schedules in PostgreSQL.

Construct a new AiopgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (Pool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: Pool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AiopgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for schedules if not exists
    and fill table with schedules from task labels.
    """
    try:
        if self._owns_pool:
            self._database_pool = await create_pool(
                dsn=self.dsn,
                **self._connect_kwargs,
            )
        async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
            await cursor.execute(CREATE_SCHEDULES_TABLE_QUERY.format(self._table_name))
        scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
        await self._update_schedules_on_startup(scheduled_tasks_for_creation)
    except Exception as error:
        raise exceptions.DatabaseConnectionError(str(error)) from error
shutdown async 
shutdown()

Close the connection pool if it was created by this schedule source.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
115
116
117
118
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()
get_schedules async 
get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
async def get_schedules(self) -> list[ScheduledTask]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
        schedules, rows = [], await cursor.fetchall()
    for schedule_id, task_name, schedule in rows:
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(schedule_id),
                    "task_name": task_name,
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules
add_schedule async 
add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
async def add_schedule(self, schedule: ScheduledTask) -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            INSERT_SCHEDULE_QUERY.format(self._table_name),
            [
                schedule.schedule_id,
                schedule.task_name,
                schedule.model_dump_json(
                    exclude={"schedule_id", "task_name"},
                ),
            ],
        )
delete_schedule async 
delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.acquire() as connection, connection.cursor() as cursor:
        await cursor.execute(
            DELETE_SCHEDULE_QUERY.format(self._table_name),
            [schedule_id],
        )
post_send async 
post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/aiopg/schedule_source.py 
179
180
181
182
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)
extract_scheduled_tasks_from_broker 
extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

asyncpg

AsyncpgBroker 

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: Pool,
    read_connection: None = ...,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: Connection,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: Pool,
    read_connection: Connection,
)

AsyncpgBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses asyncpg as driver and PostgreSQL with LISTEN/NOTIFY mechanism.

Construct a new AsyncpgBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or a callable returning one.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to asyncpg.connect().

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to asyncpg.create_pool().

  • write_pool (Pool | None, default: None ) –

    An existing connection pool to reuse for writes.

  • read_connection (Connection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/asyncpg/broker.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: asyncpg.pool.Pool | None = None,
    read_connection: asyncpg.Connection | None = None,
) -> None:
    """
    Construct a new AsyncpgBroker.

    Args:
        dsn: PostgreSQL connection string or a callable returning one.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `asyncpg.connect()`.
        write_kwargs: Extra kwargs forwarded to `asyncpg.create_pool()`.
        write_pool: An existing connection pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False

channel_name instance-attribute 

channel_name = channel_name

table_name instance-attribute 

table_name = table_name

read_kwargs instance-attribute 

read_kwargs = read_kwargs or {}

write_kwargs instance-attribute 

write_kwargs = write_kwargs or {}

max_retry_attempts instance-attribute 

max_retry_attempts = max_retry_attempts

dsn property 

dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 

startup()

Initialize the broker.

Source code in src/taskiq_pg/asyncpg/broker.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if not self._read_conn:
        self._read_conn = await asyncpg.connect(self.dsn, **self.read_kwargs)

    if not self._write_pool:
        self._write_pool = await asyncpg.create_pool(self.dsn, **self.write_kwargs)

    if self._read_conn is None:
        raise RuntimeError("Read connection is not initialized")
    if self._write_pool is None:
        raise RuntimeError("Write pool is not initialized")

    async with self._write_pool.acquire() as conn:
        await conn.execute(CREATE_MESSAGE_TABLE_QUERY.format(self.table_name))

    await self._read_conn.add_listener(self.channel_name, self._notification_handler)
    self._queue = asyncio.Queue()

shutdown async 

shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/asyncpg/broker.py 
167
168
169
170
171
172
173
174
175
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._read_conn is not None:
        await self._read_conn.remove_listener(self.channel_name, self._notification_handler)
        if self._owns_read_conn:
            await self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        await self._write_pool.close()

kick async 

kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/asyncpg/broker.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    if self._write_pool is None:
        msg = "Please run startup before kicking."
        raise ValueError(msg)

    async with self._write_pool.acquire() as conn:
        # Insert the message into the database
        message_inserted_id = tp.cast(
            "int",
            await conn.fetchval(
                INSERT_MESSAGE_QUERY.format(self.table_name),
                message.task_id,
                message.task_name,
                message.message.decode(),
                json.dumps(message.labels),
            ),
        )

        delay_value = message.labels.get("delay")
        if delay_value is not None:
            delay_seconds = int(delay_value)
            _ = asyncio.create_task(  # noqa: RUF006
                self._schedule_notification(message_inserted_id, delay_seconds),
            )
        else:
            # Send a NOTIFY with the message ID as payload
            _ = await conn.execute(
                f"NOTIFY {self.channel_name}, '{message_inserted_id}'",
            )

listen async 

listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/asyncpg/broker.py 
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    if self._write_pool is None:
        msg = "Call startup before starting listening."
        raise ValueError(msg)
    if self._queue is None:
        msg = "Startup did not initialize the queue."
        raise ValueError(msg)

    while True:
        try:
            payload = await self._queue.get()
            message_id = int(payload)
            async with self._write_pool.acquire() as conn:
                claimed = await conn.fetchrow(
                    CLAIM_MESSAGE_QUERY.format(self.table_name),
                    message_id,
                )
            if claimed is None:
                continue
            message_str = claimed["message"]
            if not isinstance(message_str, str):
                msg = "message is not a string"
                raise TypeError(msg)
            message_data = message_str.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                if self._write_pool is None:
                    msg = "Call startup before starting listening."
                    raise ValueError(msg)

                async with self._write_pool.acquire() as conn:
                    _ = await conn.execute(
                        DELETE_MESSAGE_QUERY.format(self.table_name),
                        _message_id,
                    )

            yield AckableMessage(data=message_data, ack=ack)
        except Exception:
            logger.exception("Error processing message")
            continue

AsyncpgResultBackend 

AsyncpgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AsyncpgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: Pool[Any],
)

AsyncpgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on asyncpg.

Construct a new AsyncpgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (Pool[Any] | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: "asyncpg.Pool[tp.Any] | None" = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AsyncpgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )

keep_results instance-attribute 

keep_results = keep_results

table_name instance-attribute 

table_name = table_name

field_for_task_id instance-attribute 

field_for_task_id = field_for_task_id

connect_kwargs instance-attribute 

connect_kwargs = connect_kwargs

serializer instance-attribute 

serializer = serializer or PickleSerializer()

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 

startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table
    for results if not exists.
    """
    if self._owns_pool:
        _database_pool = await asyncpg.create_pool(
            dsn=self.dsn,
            **self.connect_kwargs,
        )
        self._database_pool = _database_pool

    await self._database_pool.execute(
        queries.CREATE_TABLE_QUERY.format(
            self.table_name,
            self.field_for_task_id,
        ),
    )
    await self._database_pool.execute(
        queries.ADD_PROGRESS_COLUMN_QUERY.format(
            self.table_name,
        ),
    )
    await self._database_pool.execute(
        queries.CREATE_INDEX_QUERY.format(
            self.table_name,
            self.table_name,
        ),
    )

shutdown async 

shutdown()

Close the connection pool if it was created by this backend.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()

set_result async 

set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    _ = await self._database_pool.execute(
        queries.INSERT_RESULT_QUERY.format(
            self.table_name,
        ),
        task_id,
        self.serializer.dumpb(model_dump(result)),
    )

is_result_ready async 

is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task. :returns: True if the result is ready else False.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.
    :returns: True if the result is ready else False.
    """
    return tp.cast(
        "bool",
        await self._database_pool.fetchval(
            queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            task_id,
        ),
    )

get_result async 

get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. (deprecated in taskiq) :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs. (deprecated in taskiq)
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    result_in_bytes = tp.cast(
        "bytes",
        await self._database_pool.fetchval(
            queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            task_id,
        ),
    )
    if not self.keep_results:
        await self._database_pool.execute(
            queries.DELETE_RESULT_QUERY.format(
                self.table_name,
            ),
            task_id,
        )
    taskiq_result: tp.Final = model_validate(
        TaskiqResult[ReturnType],
        self.serializer.loadb(result_in_bytes),
    )
    if not with_logs:
        taskiq_result.log = None
    return taskiq_result

set_progress async 

set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    await self._database_pool.execute(
        queries.INSERT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        task_id,
        self.serializer.dumpb(model_dump(progress)),
    )

get_progress async 

get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    progress_in_bytes = await self._database_pool.fetchval(
        queries.SELECT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        task_id,
    )
    if progress_in_bytes is None:
        return None
    return model_validate(
        TaskProgress[tp.Any],
        self.serializer.loadb(progress_in_bytes),
    )

AsyncpgScheduleSource 

AsyncpgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AsyncpgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: Pool[Record],
)

AsyncpgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses asyncpg to store schedules in PostgreSQL.

Construct a new AsyncpgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (Pool[Record] | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: "asyncpg.Pool[asyncpg.Record] | None" = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AsyncpgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

extract_scheduled_tasks_from_broker 

extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

startup async 

startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for schedules if not exists
    and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = await asyncpg.create_pool(
            dsn=self.dsn,
            **self._connect_kwargs,
        )
    await self._database_pool.execute(
        CREATE_SCHEDULES_TABLE_QUERY.format(
            self._table_name,
        ),
    )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)

shutdown async 

shutdown()

Close the connection pool if it created by this schedule source.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()

get_schedules async 

get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as conn:
        rows_with_schedules = await conn.fetch(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
    schedules = []
    for row in rows_with_schedules:
        schedule = json.loads(row["schedule"])
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(row["id"]),
                    "task_name": row["task_name"],
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules

add_schedule async 

add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    await self._database_pool.execute(
        INSERT_SCHEDULE_QUERY.format(self._table_name),
        str(schedule.schedule_id),
        schedule.task_name,
        schedule.model_dump_json(
            exclude={"schedule_id", "task_name"},
        ),
    )

delete_schedule async 

delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    await self._database_pool.execute(
        DELETE_SCHEDULE_QUERY.format(self._table_name),
        schedule_id,
    )

post_send async 

post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
174
175
176
177
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)

broker

logger module-attribute 

logger = getLogger('taskiq.asyncpg_broker')

AsyncpgBroker 

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: Pool,
    read_connection: None = ...,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: Connection,
)

AsyncpgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: Pool,
    read_connection: Connection,
)

AsyncpgBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses asyncpg as driver and PostgreSQL with LISTEN/NOTIFY mechanism.

Construct a new AsyncpgBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or a callable returning one.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to asyncpg.connect().

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to asyncpg.create_pool().

  • write_pool (Pool | None, default: None ) –

    An existing connection pool to reuse for writes.

  • read_connection (Connection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/asyncpg/broker.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: asyncpg.pool.Pool | None = None,
    read_connection: asyncpg.Connection | None = None,
) -> None:
    """
    Construct a new AsyncpgBroker.

    Args:
        dsn: PostgreSQL connection string or a callable returning one.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `asyncpg.connect()`.
        write_kwargs: Extra kwargs forwarded to `asyncpg.create_pool()`.
        write_pool: An existing connection pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False
channel_name instance-attribute 
channel_name = channel_name
table_name instance-attribute 
table_name = table_name
read_kwargs instance-attribute 
read_kwargs = read_kwargs or {}
write_kwargs instance-attribute 
write_kwargs = write_kwargs or {}
max_retry_attempts instance-attribute 
max_retry_attempts = max_retry_attempts
dsn property 
dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 
startup()

Initialize the broker.

Source code in src/taskiq_pg/asyncpg/broker.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if not self._read_conn:
        self._read_conn = await asyncpg.connect(self.dsn, **self.read_kwargs)

    if not self._write_pool:
        self._write_pool = await asyncpg.create_pool(self.dsn, **self.write_kwargs)

    if self._read_conn is None:
        raise RuntimeError("Read connection is not initialized")
    if self._write_pool is None:
        raise RuntimeError("Write pool is not initialized")

    async with self._write_pool.acquire() as conn:
        await conn.execute(CREATE_MESSAGE_TABLE_QUERY.format(self.table_name))

    await self._read_conn.add_listener(self.channel_name, self._notification_handler)
    self._queue = asyncio.Queue()
shutdown async 
shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/asyncpg/broker.py 
167
168
169
170
171
172
173
174
175
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._read_conn is not None:
        await self._read_conn.remove_listener(self.channel_name, self._notification_handler)
        if self._owns_read_conn:
            await self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        await self._write_pool.close()
kick async 
kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/asyncpg/broker.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    if self._write_pool is None:
        msg = "Please run startup before kicking."
        raise ValueError(msg)

    async with self._write_pool.acquire() as conn:
        # Insert the message into the database
        message_inserted_id = tp.cast(
            "int",
            await conn.fetchval(
                INSERT_MESSAGE_QUERY.format(self.table_name),
                message.task_id,
                message.task_name,
                message.message.decode(),
                json.dumps(message.labels),
            ),
        )

        delay_value = message.labels.get("delay")
        if delay_value is not None:
            delay_seconds = int(delay_value)
            _ = asyncio.create_task(  # noqa: RUF006
                self._schedule_notification(message_inserted_id, delay_seconds),
            )
        else:
            # Send a NOTIFY with the message ID as payload
            _ = await conn.execute(
                f"NOTIFY {self.channel_name}, '{message_inserted_id}'",
            )
listen async 
listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/asyncpg/broker.py 
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    if self._write_pool is None:
        msg = "Call startup before starting listening."
        raise ValueError(msg)
    if self._queue is None:
        msg = "Startup did not initialize the queue."
        raise ValueError(msg)

    while True:
        try:
            payload = await self._queue.get()
            message_id = int(payload)
            async with self._write_pool.acquire() as conn:
                claimed = await conn.fetchrow(
                    CLAIM_MESSAGE_QUERY.format(self.table_name),
                    message_id,
                )
            if claimed is None:
                continue
            message_str = claimed["message"]
            if not isinstance(message_str, str):
                msg = "message is not a string"
                raise TypeError(msg)
            message_data = message_str.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                if self._write_pool is None:
                    msg = "Call startup before starting listening."
                    raise ValueError(msg)

                async with self._write_pool.acquire() as conn:
                    _ = await conn.execute(
                        DELETE_MESSAGE_QUERY.format(self.table_name),
                        _message_id,
                    )

            yield AckableMessage(data=message_data, ack=ack)
        except Exception:
            logger.exception("Error processing message")
            continue

queries

CREATE_TABLE_QUERY module-attribute 

CREATE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    task_id {} UNIQUE,\n    result BYTEA,\n    progress BYTEA\n)\n"

ADD_PROGRESS_COLUMN_QUERY module-attribute 

ADD_PROGRESS_COLUMN_QUERY = "\nALTER TABLE {} ADD COLUMN IF NOT EXISTS progress BYTEA;\n"

CREATE_INDEX_QUERY module-attribute 

CREATE_INDEX_QUERY = "\nCREATE INDEX IF NOT EXISTS {}_task_id_idx ON {} USING HASH (task_id)\n"

INSERT_RESULT_QUERY module-attribute 

INSERT_RESULT_QUERY = "\nINSERT INTO {} VALUES ($1, $2, NULL)\nON CONFLICT (task_id)\nDO UPDATE\nSET result = $2\n"

INSERT_PROGRESS_QUERY module-attribute 

INSERT_PROGRESS_QUERY = "\nINSERT INTO {} VALUES ($1, NULL, $2)\nON CONFLICT (task_id)\nDO UPDATE\nSET progress = $2\n"

SELECT_PROGRESS_QUERY module-attribute 

SELECT_PROGRESS_QUERY = (
    "\nSELECT progress FROM {} WHERE task_id = $1\n"
)

IS_RESULT_EXISTS_QUERY module-attribute 

IS_RESULT_EXISTS_QUERY = "\nSELECT EXISTS(\n    SELECT 1 FROM {} WHERE task_id = $1 and result IS NOT NULL\n)\n"

SELECT_RESULT_QUERY module-attribute 

SELECT_RESULT_QUERY = (
    "\nSELECT result FROM {} WHERE task_id = $1\n"
)

DELETE_RESULT_QUERY module-attribute 

DELETE_RESULT_QUERY = (
    "\nDELETE FROM {} WHERE task_id = $1\n"
)

CREATE_MESSAGE_TABLE_QUERY module-attribute 

CREATE_MESSAGE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id SERIAL PRIMARY KEY,\n    task_id VARCHAR NOT NULL,\n    task_name VARCHAR NOT NULL,\n    message TEXT NOT NULL,\n    labels JSONB NOT NULL,\n    status TEXT NOT NULL DEFAULT 'pending',\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_MESSAGE_QUERY module-attribute 

INSERT_MESSAGE_QUERY = "\nINSERT INTO {} (task_id, task_name, message, labels)\nVALUES ($1, $2, $3, $4)\nRETURNING id\n"

CLAIM_MESSAGE_QUERY module-attribute 

CLAIM_MESSAGE_QUERY = "UPDATE {} SET status = 'processing' WHERE id = $1 AND status = 'pending' RETURNING id, message"

DELETE_MESSAGE_QUERY module-attribute 

DELETE_MESSAGE_QUERY = 'DELETE FROM {} WHERE id = $1'

CREATE_SCHEDULES_TABLE_QUERY module-attribute 

CREATE_SCHEDULES_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id UUID PRIMARY KEY,\n    task_name VARCHAR(100) NOT NULL,\n    schedule JSONB NOT NULL,\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),\n    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_SCHEDULE_QUERY module-attribute 

INSERT_SCHEDULE_QUERY = "\nINSERT INTO {} (id, task_name, schedule)\nVALUES ($1, $2, $3)\nON CONFLICT (id) DO UPDATE\nSET task_name = EXCLUDED.task_name,\n    schedule = EXCLUDED.schedule,\n    updated_at = NOW();\n"

SELECT_SCHEDULES_QUERY module-attribute 

SELECT_SCHEDULES_QUERY = (
    "\nSELECT id, task_name, schedule\nFROM {};\n"
)

DELETE_ALL_SCHEDULES_QUERY module-attribute 

DELETE_ALL_SCHEDULES_QUERY = '\nDELETE FROM {};\n'

DELETE_SCHEDULE_QUERY module-attribute 

DELETE_SCHEDULE_QUERY = '\nDELETE FROM {} WHERE id = $1;\n'

result_backend

AsyncpgResultBackend 

AsyncpgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AsyncpgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: Pool[Any],
)

AsyncpgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on asyncpg.

Construct a new AsyncpgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (Pool[Any] | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: "asyncpg.Pool[tp.Any] | None" = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AsyncpgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )
keep_results instance-attribute 
keep_results = keep_results
table_name instance-attribute 
table_name = table_name
field_for_task_id instance-attribute 
field_for_task_id = field_for_task_id
connect_kwargs instance-attribute 
connect_kwargs = connect_kwargs
serializer instance-attribute 
serializer = serializer or PickleSerializer()
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table
    for results if not exists.
    """
    if self._owns_pool:
        _database_pool = await asyncpg.create_pool(
            dsn=self.dsn,
            **self.connect_kwargs,
        )
        self._database_pool = _database_pool

    await self._database_pool.execute(
        queries.CREATE_TABLE_QUERY.format(
            self.table_name,
            self.field_for_task_id,
        ),
    )
    await self._database_pool.execute(
        queries.ADD_PROGRESS_COLUMN_QUERY.format(
            self.table_name,
        ),
    )
    await self._database_pool.execute(
        queries.CREATE_INDEX_QUERY.format(
            self.table_name,
            self.table_name,
        ),
    )
shutdown async 
shutdown()

Close the connection pool if it was created by this backend.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()
set_result async 
set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    _ = await self._database_pool.execute(
        queries.INSERT_RESULT_QUERY.format(
            self.table_name,
        ),
        task_id,
        self.serializer.dumpb(model_dump(result)),
    )
is_result_ready async 
is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task. :returns: True if the result is ready else False.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.
    :returns: True if the result is ready else False.
    """
    return tp.cast(
        "bool",
        await self._database_pool.fetchval(
            queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            task_id,
        ),
    )
get_result async 
get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. (deprecated in taskiq) :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs. (deprecated in taskiq)
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    result_in_bytes = tp.cast(
        "bytes",
        await self._database_pool.fetchval(
            queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            task_id,
        ),
    )
    if not self.keep_results:
        await self._database_pool.execute(
            queries.DELETE_RESULT_QUERY.format(
                self.table_name,
            ),
            task_id,
        )
    taskiq_result: tp.Final = model_validate(
        TaskiqResult[ReturnType],
        self.serializer.loadb(result_in_bytes),
    )
    if not with_logs:
        taskiq_result.log = None
    return taskiq_result
set_progress async 
set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    await self._database_pool.execute(
        queries.INSERT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        task_id,
        self.serializer.dumpb(model_dump(progress)),
    )
get_progress async 
get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/asyncpg/result_backend.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    progress_in_bytes = await self._database_pool.fetchval(
        queries.SELECT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        task_id,
    )
    if progress_in_bytes is None:
        return None
    return model_validate(
        TaskProgress[tp.Any],
        self.serializer.loadb(progress_in_bytes),
    )

schedule_source

logger module-attribute 

logger = getLogger('taskiq_pg.asyncpg_schedule_source')

AsyncpgScheduleSource 

AsyncpgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

AsyncpgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: Pool[Record],
)

AsyncpgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses asyncpg to store schedules in PostgreSQL.

Construct a new AsyncpgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (Pool[Record] | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: "asyncpg.Pool[asyncpg.Record] | None" = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new AsyncpgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for schedules if not exists
    and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = await asyncpg.create_pool(
            dsn=self.dsn,
            **self._connect_kwargs,
        )
    await self._database_pool.execute(
        CREATE_SCHEDULES_TABLE_QUERY.format(
            self._table_name,
        ),
    )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)
shutdown async 
shutdown()

Close the connection pool if it created by this schedule source.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()
get_schedules async 
get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as conn:
        rows_with_schedules = await conn.fetch(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
    schedules = []
    for row in rows_with_schedules:
        schedule = json.loads(row["schedule"])
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(row["id"]),
                    "task_name": row["task_name"],
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules
add_schedule async 
add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    await self._database_pool.execute(
        INSERT_SCHEDULE_QUERY.format(self._table_name),
        str(schedule.schedule_id),
        schedule.task_name,
        schedule.model_dump_json(
            exclude={"schedule_id", "task_name"},
        ),
    )
delete_schedule async 
delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    await self._database_pool.execute(
        DELETE_SCHEDULE_QUERY.format(self._table_name),
        schedule_id,
    )
post_send async 
post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/asyncpg/schedule_source.py 
174
175
176
177
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)
extract_scheduled_tasks_from_broker 
extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

exceptions

BaseTaskiqPgError

Bases: Exception

Base error for all possible exception in the lib.

DatabaseConnectionError

Bases: BaseTaskiqPgError

Error if cannot connect to PostgreSQL.

ResultIsMissingError

Bases: BaseTaskiqPgError

Error if cannot retrieve result from PostgreSQL.

psqlpy

PSQLPyBroker 

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: ConnectionPool,
    read_connection: None = ...,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: Connection,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: ConnectionPool,
    read_connection: Connection,
)

PSQLPyBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses PostgreSQL and PSQLPy with LISTEN/NOTIFY.

Construct a new PSQLPyBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to psqlpy.connect()

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to psqlpy.ConnectionPool()

  • write_pool (ConnectionPool | None, default: None ) –

    An existing pool to reuse for writes.

  • read_connection (Connection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/psqlpy/broker.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: psqlpy.ConnectionPool | None = None,
    read_connection: psqlpy.Connection | None = None,
) -> None:
    """
    Construct a new PSQLPyBroker.

    Args:
        dsn: PostgreSQL connection string or callable.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `psqlpy.connect()`
        write_kwargs: Extra kwargs forwarded to `psqlpy.ConnectionPool()`
        write_pool: An existing pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False

channel_name instance-attribute 

channel_name = channel_name

table_name instance-attribute 

table_name = table_name

read_kwargs instance-attribute 

read_kwargs = read_kwargs or {}

write_kwargs instance-attribute 

write_kwargs = write_kwargs or {}

max_retry_attempts instance-attribute 

max_retry_attempts = max_retry_attempts

dsn property 

dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 

startup()

Initialize the broker.

Source code in src/taskiq_pg/psqlpy/broker.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if self._owns_read_conn:
        self._read_conn = await psqlpy.connect(
            dsn=self.dsn,
            **self.read_kwargs,
        )

    if self._owns_write_pool:
        self._write_pool = psqlpy.ConnectionPool(
            dsn=self.dsn,
            **self.write_kwargs,
        )

    # create messages table if it doesn't exist
    async with self._write_pool.acquire() as conn:
        await conn.execute(CREATE_MESSAGE_TABLE_QUERY.format(self.table_name))

    # listen to notification channel
    self._listener = self._write_pool.listener()
    await self._listener.add_callback(self.channel_name, self._notification_handler)
    await self._listener.startup()
    self._listener.listen()

    self._queue = asyncio.Queue()

shutdown async 

shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/psqlpy/broker.py 
192
193
194
195
196
197
198
199
200
201
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._read_conn is not None and self._owns_read_conn:
        self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        self._write_pool.close()
    if self._listener is not None:
        self._listener.abort_listen()
        await self._listener.shutdown()

kick async 

kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/psqlpy/broker.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    async with self._write_pool.acquire() as conn:
        # insert message into db table
        message_inserted_id = tp.cast(
            "int",
            await conn.fetch_val(
                INSERT_MESSAGE_QUERY.format(self.table_name),
                [
                    message.task_id,
                    message.task_name,
                    message.message.decode(),
                    JSONB(message.labels),
                ],
            ),
        )

        delay_value = tp.cast("str | None", message.labels.get("delay"))
        if delay_value is not None:
            delay_seconds = int(delay_value)
            asyncio.create_task(  # noqa: RUF006
                self._schedule_notification(message_inserted_id, delay_seconds),
            )
        else:
            # Send NOTIFY with message ID as payload
            _ = await conn.execute(
                f"NOTIFY {self.channel_name}, '{message_inserted_id}'",
            )

listen async 

listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/psqlpy/broker.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    while True:
        try:
            payload = await self._queue.get()
            message_id = int(payload)  # payload is the message id
            try:
                async with self._write_pool.acquire() as conn:
                    claimed_message = await conn.fetch_row(
                        CLAIM_MESSAGE_QUERY.format(self.table_name),
                        [message_id],
                    )
            except ConnectionExecuteError:  # message was claimed by another worker
                continue
            message_row_result = tp.cast(
                "MessageRow",
                tp.cast("object", claimed_message.as_class(MessageRow)),
            )
            message_data = message_row_result.message.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                async with self._write_pool.acquire() as conn:
                    _ = await conn.execute(
                        DELETE_MESSAGE_QUERY.format(self.table_name),
                        [_message_id],
                    )

            yield AckableMessage(data=message_data, ack=ack)
        except Exception:
            logger.exception("Error processing message")
            continue

PSQLPyResultBackend 

PSQLPyResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PSQLPyResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: ConnectionPool,
)

PSQLPyResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on PSQLPy.

Construct a new PSQLPyResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (ConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: ConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PSQLPyResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )

keep_results instance-attribute 

keep_results = keep_results

table_name instance-attribute 

table_name = table_name

field_for_task_id instance-attribute 

field_for_task_id = field_for_task_id

connect_kwargs instance-attribute 

connect_kwargs = connect_kwargs

serializer instance-attribute 

serializer = serializer or PickleSerializer()

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 

startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table for results if not exists.
    """
    if self._owns_pool:
        self._database_pool = ConnectionPool(
            dsn=self.dsn,
            **self.connect_kwargs,
        )
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.CREATE_TABLE_QUERY.format(
            self.table_name,
            self.field_for_task_id,
        ),
    )
    await connection.execute(
        querystring=queries.ADD_PROGRESS_COLUMN_QUERY.format(
            self.table_name,
        ),
    )
    await connection.execute(
        querystring=queries.CREATE_INDEX_QUERY.format(
            self.table_name,
            self.table_name,
        ),
    )

shutdown async 

shutdown()

Close the connection pool if it was created by this result backend.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this result backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()

set_result async 

set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.INSERT_RESULT_QUERY.format(
            self.table_name,
        ),
        parameters=[
            task_id,
            self.serializer.dumpb(model_dump(result)),
        ],
    )

is_result_ready async 

is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task.

:returns: True if the result is ready else False.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.

    :returns: True if the result is ready else False.
    """
    connection: tp.Final = await self._database_pool.connection()
    return tp.cast(
        "bool",
        await connection.fetch_val(
            querystring=queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        ),
    )

get_result async 

get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    connection: tp.Final = await self._database_pool.connection()
    try:
        result_in_bytes: tp.Final[bytes] = await connection.fetch_val(
            querystring=queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )
    except BaseConnectionError as exc:
        msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
        raise ResultIsMissingError(msg) from exc

    if not self.keep_results:
        await connection.execute(
            querystring=queries.DELETE_RESULT_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )

    taskiq_result: tp.Final = model_validate(
        TaskiqResult[ReturnType],
        self.serializer.loadb(result_in_bytes),
    )

    if not with_logs:
        taskiq_result.log = None

    return taskiq_result

set_progress async 

set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.INSERT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        parameters=[
            task_id,
            self.serializer.dumpb(model_dump(progress)),
        ],
    )

get_progress async 

get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    connection: tp.Final = await self._database_pool.connection()
    try:
        progress_in_bytes = await connection.fetch_val(
            querystring=queries.SELECT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )
    except BaseConnectionError as exc:
        msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
        raise ResultIsMissingError(msg) from exc
    if progress_in_bytes is None:
        return None
    return model_validate(
        TaskProgress[tp.Any],
        self.serializer.loadb(progress_in_bytes),
    )

PSQLPyScheduleSource 

PSQLPyScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PSQLPyScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: ConnectionPool,
)

PSQLPyScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses psqlpy to store schedules in PostgreSQL.

Construct a new PSQLPyScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (ConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    broker: AsyncBroker,
    dsn: "str | tp.Callable[[], str]" = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: ConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PSQLPyScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool
    else:
        self._owns_pool = True

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

extract_scheduled_tasks_from_broker 

extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

startup async 

startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for
    schedules if not exists and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = ConnectionPool(
            dsn=self.dsn,
            **self._connect_kwargs,
        )
    async with self._database_pool.acquire() as connection:
        await connection.execute(
            CREATE_SCHEDULES_TABLE_QUERY.format(
                self._table_name,
            ),
        )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)

shutdown async 

shutdown()

Close the connection pool if was created by this schedule source.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
125
126
127
128
async def shutdown(self) -> None:
    """Close the connection pool if was created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()

get_schedules async 

get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as connection:
        rows_with_schedules = await connection.fetch(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
    schedules = []
    for row in rows_with_schedules.result():
        schedule = row["schedule"]
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(row["id"]),
                    "task_name": row["task_name"],
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules

add_schedule async 

add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.acquire() as connection:
        schedule_dict = schedule.model_dump(
            mode="json",
            exclude={"schedule_id", "task_name"},
        )
        await connection.execute(
            INSERT_SCHEDULE_QUERY.format(self._table_name),
            [
                uuid.UUID(schedule.schedule_id),
                schedule.task_name,
                JSONB(schedule_dict),
            ]
        )

delete_schedule async 

delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.acquire() as connection:
        await connection.execute(
            DELETE_SCHEDULE_QUERY.format(self._table_name),
            [uuid.UUID(schedule_id)],
        )

post_send async 

post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
192
193
194
195
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)

broker

logger module-attribute 

logger = getLogger('taskiq.psqlpy_broker')

MessageRow dataclass 

MessageRow(
    id,
    task_id,
    task_name,
    message,
    labels,
    status,
    created_at,
)

Message in db table.

id instance-attribute 
id
task_id instance-attribute 
task_id
task_name instance-attribute 
task_name
message instance-attribute 
message
labels instance-attribute 
labels
status instance-attribute 
status
created_at instance-attribute 
created_at

PSQLPyBroker 

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: ConnectionPool,
    read_connection: None = ...,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: Connection,
)

PSQLPyBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: ConnectionPool,
    read_connection: Connection,
)

PSQLPyBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses PostgreSQL and PSQLPy with LISTEN/NOTIFY.

Construct a new PSQLPyBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to psqlpy.connect()

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to psqlpy.ConnectionPool()

  • write_pool (ConnectionPool | None, default: None ) –

    An existing pool to reuse for writes.

  • read_connection (Connection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/psqlpy/broker.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: psqlpy.ConnectionPool | None = None,
    read_connection: psqlpy.Connection | None = None,
) -> None:
    """
    Construct a new PSQLPyBroker.

    Args:
        dsn: PostgreSQL connection string or callable.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `psqlpy.connect()`
        write_kwargs: Extra kwargs forwarded to `psqlpy.ConnectionPool()`
        write_pool: An existing pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False
channel_name instance-attribute 
channel_name = channel_name
table_name instance-attribute 
table_name = table_name
read_kwargs instance-attribute 
read_kwargs = read_kwargs or {}
write_kwargs instance-attribute 
write_kwargs = write_kwargs or {}
max_retry_attempts instance-attribute 
max_retry_attempts = max_retry_attempts
dsn property 
dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 
startup()

Initialize the broker.

Source code in src/taskiq_pg/psqlpy/broker.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if self._owns_read_conn:
        self._read_conn = await psqlpy.connect(
            dsn=self.dsn,
            **self.read_kwargs,
        )

    if self._owns_write_pool:
        self._write_pool = psqlpy.ConnectionPool(
            dsn=self.dsn,
            **self.write_kwargs,
        )

    # create messages table if it doesn't exist
    async with self._write_pool.acquire() as conn:
        await conn.execute(CREATE_MESSAGE_TABLE_QUERY.format(self.table_name))

    # listen to notification channel
    self._listener = self._write_pool.listener()
    await self._listener.add_callback(self.channel_name, self._notification_handler)
    await self._listener.startup()
    self._listener.listen()

    self._queue = asyncio.Queue()
shutdown async 
shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/psqlpy/broker.py 
192
193
194
195
196
197
198
199
200
201
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._read_conn is not None and self._owns_read_conn:
        self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        self._write_pool.close()
    if self._listener is not None:
        self._listener.abort_listen()
        await self._listener.shutdown()
kick async 
kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/psqlpy/broker.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    async with self._write_pool.acquire() as conn:
        # insert message into db table
        message_inserted_id = tp.cast(
            "int",
            await conn.fetch_val(
                INSERT_MESSAGE_QUERY.format(self.table_name),
                [
                    message.task_id,
                    message.task_name,
                    message.message.decode(),
                    JSONB(message.labels),
                ],
            ),
        )

        delay_value = tp.cast("str | None", message.labels.get("delay"))
        if delay_value is not None:
            delay_seconds = int(delay_value)
            asyncio.create_task(  # noqa: RUF006
                self._schedule_notification(message_inserted_id, delay_seconds),
            )
        else:
            # Send NOTIFY with message ID as payload
            _ = await conn.execute(
                f"NOTIFY {self.channel_name}, '{message_inserted_id}'",
            )
listen async 
listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/psqlpy/broker.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    while True:
        try:
            payload = await self._queue.get()
            message_id = int(payload)  # payload is the message id
            try:
                async with self._write_pool.acquire() as conn:
                    claimed_message = await conn.fetch_row(
                        CLAIM_MESSAGE_QUERY.format(self.table_name),
                        [message_id],
                    )
            except ConnectionExecuteError:  # message was claimed by another worker
                continue
            message_row_result = tp.cast(
                "MessageRow",
                tp.cast("object", claimed_message.as_class(MessageRow)),
            )
            message_data = message_row_result.message.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                async with self._write_pool.acquire() as conn:
                    _ = await conn.execute(
                        DELETE_MESSAGE_QUERY.format(self.table_name),
                        [_message_id],
                    )

            yield AckableMessage(data=message_data, ack=ack)
        except Exception:
            logger.exception("Error processing message")
            continue

queries

CREATE_TABLE_QUERY module-attribute 

CREATE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    task_id {} UNIQUE,\n    result BYTEA,\n    progress BYTEA\n)\n"

ADD_PROGRESS_COLUMN_QUERY module-attribute 

ADD_PROGRESS_COLUMN_QUERY = "\nALTER TABLE {} ADD COLUMN IF NOT EXISTS progress BYTEA;\n"

CREATE_INDEX_QUERY module-attribute 

CREATE_INDEX_QUERY = "\nCREATE INDEX IF NOT EXISTS {}_task_id_idx ON {} USING HASH (task_id)\n"

INSERT_RESULT_QUERY module-attribute 

INSERT_RESULT_QUERY = "\nINSERT INTO {} VALUES ($1, $2, NULL)\nON CONFLICT (task_id)\nDO UPDATE\nSET result = $2\n"

INSERT_PROGRESS_QUERY module-attribute 

INSERT_PROGRESS_QUERY = "\nINSERT INTO {} VALUES ($1, NULL, $2)\nON CONFLICT (task_id)\nDO UPDATE\nSET progress = $2\n"

SELECT_PROGRESS_QUERY module-attribute 

SELECT_PROGRESS_QUERY = (
    "\nSELECT progress FROM {} WHERE task_id = $1\n"
)

IS_RESULT_EXISTS_QUERY module-attribute 

IS_RESULT_EXISTS_QUERY = "\nSELECT EXISTS(\n    SELECT 1 FROM {} WHERE task_id = $1 and result IS NOT NULL\n)\n"

SELECT_RESULT_QUERY module-attribute 

SELECT_RESULT_QUERY = (
    "\nSELECT result FROM {} WHERE task_id = $1\n"
)

DELETE_RESULT_QUERY module-attribute 

DELETE_RESULT_QUERY = (
    "\nDELETE FROM {} WHERE task_id = $1\n"
)

CREATE_MESSAGE_TABLE_QUERY module-attribute 

CREATE_MESSAGE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id SERIAL PRIMARY KEY,\n    task_id VARCHAR NOT NULL,\n    task_name VARCHAR NOT NULL,\n    message TEXT NOT NULL,\n    labels JSONB NOT NULL,\n    status TEXT NOT NULL DEFAULT 'pending',\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_MESSAGE_QUERY module-attribute 

INSERT_MESSAGE_QUERY = "\nINSERT INTO {} (task_id, task_name, message, labels)\nVALUES ($1, $2, $3, $4)\nRETURNING id\n"

CLAIM_MESSAGE_QUERY module-attribute 

CLAIM_MESSAGE_QUERY = "UPDATE {} SET status = 'processing' WHERE id = $1 AND status = 'pending' RETURNING *"

DELETE_MESSAGE_QUERY module-attribute 

DELETE_MESSAGE_QUERY = 'DELETE FROM {} WHERE id = $1'

CREATE_SCHEDULES_TABLE_QUERY module-attribute 

CREATE_SCHEDULES_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id UUID PRIMARY KEY,\n    task_name VARCHAR(100) NOT NULL,\n    schedule JSONB NOT NULL,\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),\n    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_SCHEDULE_QUERY module-attribute 

INSERT_SCHEDULE_QUERY = "\nINSERT INTO {} (id, task_name, schedule)\nVALUES ($1, $2, $3)\nON CONFLICT (id) DO UPDATE\nSET task_name = EXCLUDED.task_name,\n    schedule = EXCLUDED.schedule,\n    updated_at = NOW();\n"

SELECT_SCHEDULES_QUERY module-attribute 

SELECT_SCHEDULES_QUERY = (
    "\nSELECT id, task_name, schedule\nFROM {};\n"
)

DELETE_ALL_SCHEDULES_QUERY module-attribute 

DELETE_ALL_SCHEDULES_QUERY = '\nDELETE FROM {};\n'

DELETE_SCHEDULE_QUERY module-attribute 

DELETE_SCHEDULE_QUERY = '\nDELETE FROM {} WHERE id = $1;\n'

result_backend

PSQLPyResultBackend 

PSQLPyResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PSQLPyResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: ConnectionPool,
)

PSQLPyResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on PSQLPy.

Construct a new PSQLPyResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (ConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: ConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PSQLPyResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )
keep_results instance-attribute 
keep_results = keep_results
table_name instance-attribute 
table_name = table_name
field_for_task_id instance-attribute 
field_for_task_id = field_for_task_id
connect_kwargs instance-attribute 
connect_kwargs = connect_kwargs
serializer instance-attribute 
serializer = serializer or PickleSerializer()
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table for results if not exists.
    """
    if self._owns_pool:
        self._database_pool = ConnectionPool(
            dsn=self.dsn,
            **self.connect_kwargs,
        )
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.CREATE_TABLE_QUERY.format(
            self.table_name,
            self.field_for_task_id,
        ),
    )
    await connection.execute(
        querystring=queries.ADD_PROGRESS_COLUMN_QUERY.format(
            self.table_name,
        ),
    )
    await connection.execute(
        querystring=queries.CREATE_INDEX_QUERY.format(
            self.table_name,
            self.table_name,
        ),
    )
shutdown async 
shutdown()

Close the connection pool if it was created by this result backend.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
113
114
115
116
async def shutdown(self) -> None:
    """Close the connection pool if it was created by this result backend."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()
set_result async 
set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.INSERT_RESULT_QUERY.format(
            self.table_name,
        ),
        parameters=[
            task_id,
            self.serializer.dumpb(model_dump(result)),
        ],
    )
is_result_ready async 
is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task.

:returns: True if the result is ready else False.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.

    :returns: True if the result is ready else False.
    """
    connection: tp.Final = await self._database_pool.connection()
    return tp.cast(
        "bool",
        await connection.fetch_val(
            querystring=queries.IS_RESULT_EXISTS_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        ),
    )
get_result async 
get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    connection: tp.Final = await self._database_pool.connection()
    try:
        result_in_bytes: tp.Final[bytes] = await connection.fetch_val(
            querystring=queries.SELECT_RESULT_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )
    except BaseConnectionError as exc:
        msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
        raise ResultIsMissingError(msg) from exc

    if not self.keep_results:
        await connection.execute(
            querystring=queries.DELETE_RESULT_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )

    taskiq_result: tp.Final = model_validate(
        TaskiqResult[ReturnType],
        self.serializer.loadb(result_in_bytes),
    )

    if not with_logs:
        taskiq_result.log = None

    return taskiq_result
set_progress async 
set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    connection = await self._database_pool.connection()
    await connection.execute(
        querystring=queries.INSERT_PROGRESS_QUERY.format(
            self.table_name,
        ),
        parameters=[
            task_id,
            self.serializer.dumpb(model_dump(progress)),
        ],
    )
get_progress async 
get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/psqlpy/result_backend.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    connection: tp.Final = await self._database_pool.connection()
    try:
        progress_in_bytes = await connection.fetch_val(
            querystring=queries.SELECT_PROGRESS_QUERY.format(
                self.table_name,
            ),
            parameters=[task_id],
        )
    except BaseConnectionError as exc:
        msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
        raise ResultIsMissingError(msg) from exc
    if progress_in_bytes is None:
        return None
    return model_validate(
        TaskProgress[tp.Any],
        self.serializer.loadb(progress_in_bytes),
    )

schedule_source

logger module-attribute 

logger = getLogger('taskiq_pg.psqlpy_schedule_source')

PSQLPyScheduleSource 

PSQLPyScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PSQLPyScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: ConnectionPool,
)

PSQLPyScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses psqlpy to store schedules in PostgreSQL.

Construct a new PSQLPyScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (ConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    broker: AsyncBroker,
    dsn: "str | tp.Callable[[], str]" = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: ConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PSQLPyScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool
    else:
        self._owns_pool = True

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for
    schedules if not exists and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = ConnectionPool(
            dsn=self.dsn,
            **self._connect_kwargs,
        )
    async with self._database_pool.acquire() as connection:
        await connection.execute(
            CREATE_SCHEDULES_TABLE_QUERY.format(
                self._table_name,
            ),
        )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)
shutdown async 
shutdown()

Close the connection pool if was created by this schedule source.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
125
126
127
128
async def shutdown(self) -> None:
    """Close the connection pool if was created by this schedule source."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        self._database_pool.close()
get_schedules async 
get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    async with self._database_pool.acquire() as connection:
        rows_with_schedules = await connection.fetch(
            SELECT_SCHEDULES_QUERY.format(self._table_name),
        )
    schedules = []
    for row in rows_with_schedules.result():
        schedule = row["schedule"]
        schedules.append(
            ScheduledTask.model_validate(
                {
                    "schedule_id": str(row["id"]),
                    "task_name": row["task_name"],
                    "labels": schedule["labels"],
                    "args": schedule["args"],
                    "kwargs": schedule["kwargs"],
                    "cron": schedule["cron"],
                    "cron_offset": schedule["cron_offset"],
                    "time": schedule["time"],
                    "interval": schedule["interval"],
                },
            ),
        )
    return schedules
add_schedule async 
add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.acquire() as connection:
        schedule_dict = schedule.model_dump(
            mode="json",
            exclude={"schedule_id", "task_name"},
        )
        await connection.execute(
            INSERT_SCHEDULE_QUERY.format(self._table_name),
            [
                uuid.UUID(schedule.schedule_id),
                schedule.task_name,
                JSONB(schedule_dict),
            ]
        )
delete_schedule async 
delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.acquire() as connection:
        await connection.execute(
            DELETE_SCHEDULE_QUERY.format(self._table_name),
            [uuid.UUID(schedule_id)],
        )
post_send async 
post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/psqlpy/schedule_source.py 
192
193
194
195
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)
extract_scheduled_tasks_from_broker 
extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

psycopg

PsycopgBroker 

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: AsyncConnectionPool,
    read_connection: None = ...,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: AsyncConnection,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: AsyncConnectionPool,
    read_connection: AsyncConnection,
)

PsycopgBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses PostgreSQL and psycopg with LISTEN/NOTIFY.

Construct a new PsycopgBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to AsyncConnection.connect().

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to AsyncConnectionPool().

  • write_pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse for writes.

  • read_connection (AsyncConnection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/psycopg/broker.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: AsyncConnectionPool | None = None,
    read_connection: AsyncConnection | None = None,
) -> None:
    """
    Construct a new PsycopgBroker.

    Args:
        dsn: PostgreSQL connection string or callable.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `AsyncConnection.connect()`.
        write_kwargs: Extra kwargs forwarded to `AsyncConnectionPool()`.
        write_pool: An existing connection pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False

channel_name instance-attribute 

channel_name = channel_name

table_name instance-attribute 

table_name = table_name

read_kwargs instance-attribute 

read_kwargs = read_kwargs or {}

write_kwargs instance-attribute 

write_kwargs = write_kwargs or {}

max_retry_attempts instance-attribute 

max_retry_attempts = max_retry_attempts

dsn property 

dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 

startup()

Initialize the broker.

Source code in src/taskiq_pg/psycopg/broker.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if self._owns_read_conn:
        self._read_conn = await AsyncConnection.connect(
            conninfo=self.dsn,
            **self.read_kwargs,
            autocommit=True,
            cursor_factory=AsyncRawCursor,
        )

    if self._owns_write_pool:
        self._write_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self.write_kwargs,
        )

    if self._write_pool.closed:
        await self._write_pool.open()

    async with self._write_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(sql.SQL(CREATE_MESSAGE_TABLE_QUERY).format(sql.Identifier(self.table_name)))

    await self._read_conn.execute(sql.SQL("LISTEN {}").format(sql.Identifier(self.channel_name)))
    self._notifies_iter = self._read_conn.notifies()

shutdown async 

shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/psycopg/broker.py 
178
179
180
181
182
183
184
185
186
187
188
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._notifies_iter is not None:
        with suppress(RuntimeError):  # RuntimeError: aclose(): asynchronous generator is already running
            await self._notifies_iter.aclose()  # type: ignore[attr-defined]
    if self._read_conn is not None and self._owns_read_conn:
            await self._read_conn.notifies().aclose()
            await self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        await self._write_pool.close()

kick async 

kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/psycopg/broker.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    async with self._write_pool.connection() as connection, connection.cursor() as cursor:
        # insert message into db table
        await cursor.execute(
            sql.SQL(INSERT_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
            [
                message.task_id,
                message.task_name,
                message.message.decode(),
                json.dumps(message.labels),
            ],
        )
        row = await cursor.fetchone()
        if row is None:
            msg = "failed to insert message"
            raise RuntimeError(msg)
        message_inserted_id = int(row[0])

        delay_value = tp.cast("str | None", message.labels.get("delay"))
        if delay_value is not None:
            delay_seconds = int(delay_value)
            await self._schedule_notification(message_inserted_id, delay_seconds)
        else:
            # Send NOTIFY with message ID as payload
            await cursor.execute(
                sql.SQL("NOTIFY {}, {}").format(
                    sql.Identifier(self.channel_name),
                    sql.Literal(str(message_inserted_id)),
                ),
            )

listen async 

listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/psycopg/broker.py 
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    while True:
        async for message_id_str in self._listen_context():
            message_id = int(message_id_str)  # payload is the message id
            try:
                async with self._write_pool.connection() as connection, connection.cursor() as cursor:
                    await cursor.execute(
                        sql.SQL(CLAIM_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
                        [message_id],
                    )
                    claimed_message = await cursor.fetchone()
                    if claimed_message is None:
                        continue
            except psycopg.OperationalError:  # message was claimed by another worker
                continue
            message_str = claimed_message[3]
            if not isinstance(message_str, str):
                msg = "Message is not a string"
                raise TypeError(msg)
            message_data = message_str.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                async with self._write_pool.connection() as connection, connection.cursor() as cursor:
                    await cursor.execute(
                        sql.SQL(DELETE_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
                        [_message_id],
                    )

            yield AckableMessage(data=message_data, ack=ack)

PsycopgResultBackend 

PsycopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PsycopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: AsyncConnectionPool,
)

PsycopgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on psycopg.

Construct a new PsycopgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psycopg/result_backend.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: AsyncConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PsycopgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )

keep_results instance-attribute 

keep_results = keep_results

table_name instance-attribute 

table_name = table_name

field_for_task_id instance-attribute 

field_for_task_id = field_for_task_id

connect_kwargs instance-attribute 

connect_kwargs = connect_kwargs

serializer instance-attribute 

serializer = serializer or PickleSerializer()

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 

startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/psycopg/result_backend.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table
    for results if not exists.
    """
    if self._owns_pool:
        self._database_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self.connect_kwargs,
        )

    if self._database_pool.closed:
        await self._database_pool.open()

    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.CREATE_TABLE_QUERY).format(
                sql.Identifier(self.table_name),
                sql.SQL(self.field_for_task_id),
            ),
        )
        await cursor.execute(
            query=sql.SQL(queries.ADD_PROGRESS_COLUMN_QUERY).format(
                sql.Identifier(self.table_name),
            ),
        )
        await cursor.execute(
            query=sql.SQL(queries.CREATE_INDEX_QUERY).format(
                sql.Identifier(self.table_name + "_task_id_idx"),
                sql.Identifier(self.table_name),
            ),
        )

shutdown async 

shutdown()

Close the connection pool (only if owned by this backend).

Source code in src/taskiq_pg/psycopg/result_backend.py 
119
120
121
122
async def shutdown(self) -> None:
    """Close the connection pool (only if owned by this backend)."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()

set_result async 

set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/psycopg/result_backend.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.INSERT_RESULT_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[
                task_id,
                self.serializer.dumpb(model_dump(result)),
            ],
        )

is_result_ready async 

is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task.

:returns: True if the result is ready else False.

Source code in src/taskiq_pg/psycopg/result_backend.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.

    :returns: True if the result is ready else False.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.IS_RESULT_EXISTS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        row = await execute_result.fetchone()
        return bool(row and row[0])

get_result async 

get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/psycopg/result_backend.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.SELECT_RESULT_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        result = await execute_result.fetchone()
        if result is None:
            msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
            raise ResultIsMissingError(msg)
        result_in_bytes: tp.Final = result[0]

        if not self.keep_results:
            await cursor.execute(
                query=sql.SQL(queries.DELETE_RESULT_QUERY).format(
                    sql.Identifier(self.table_name),
                ),
                params=[task_id],
            )

        taskiq_result: tp.Final = model_validate(
            TaskiqResult[ReturnType],
            self.serializer.loadb(result_in_bytes),
        )

        if not with_logs:
            taskiq_result.log = None

        return taskiq_result

set_progress async 

set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/psycopg/result_backend.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.INSERT_PROGRESS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[
                task_id,
                self.serializer.dumpb(model_dump(progress)),
                self.serializer.dumpb(model_dump(progress)),
            ],
        )

get_progress async 

get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/psycopg/result_backend.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.SELECT_PROGRESS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        progress_in_bytes = await execute_result.fetchone()
        if progress_in_bytes is None or progress_in_bytes[0] is None:
            return None
        return model_validate(
            TaskProgress[tp.Any],
            self.serializer.loadb(progress_in_bytes[0]),
        )

PsycopgScheduleSource 

PsycopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PsycopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: AsyncConnectionPool,
)

PsycopgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses psycopg to store schedules in PostgreSQL.

Construct a new PsycopgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: AsyncConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PsycopgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.

    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )

dsn property 

dsn

Get the DSN string.

Returns the DSN string or None if not set.

extract_scheduled_tasks_from_broker 

extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation

startup async 

startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for
    schedules if not exists and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self._connect_kwargs,
        )

    if self._database_pool.closed:
        await self._database_pool.open()

    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(CREATE_SCHEDULES_TABLE_QUERY).format(sql.Identifier(self._table_name)),
        )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)

shutdown async 

shutdown()

Close the connection pool (only if owned by this source).

Source code in src/taskiq_pg/psycopg/schedule_source.py 
124
125
126
127
async def shutdown(self) -> None:
    """Close the connection pool (only if owned by this source)."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()

get_schedules async 

get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    schedules = []
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        rows_with_schedules = await cursor.execute(
            sql.SQL(SELECT_SCHEDULES_QUERY).format(sql.Identifier(self._table_name)),
        )
        rows = await rows_with_schedules.fetchall()
        for schedule_id, task_name, schedule in rows:
            schedules.append(
                ScheduledTask.model_validate(
                    {
                        "schedule_id": str(schedule_id),
                        "task_name": task_name,
                        "labels": schedule["labels"],
                        "args": schedule["args"],
                        "kwargs": schedule["kwargs"],
                        "cron": schedule["cron"],
                        "cron_offset": schedule["cron_offset"],
                        "time": schedule["time"],
                        "interval": schedule["interval"],
                    },
                ),
            )
    return schedules

add_schedule async 

add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(INSERT_SCHEDULE_QUERY).format(sql.Identifier(self._table_name)),
            [
                uuid.UUID(schedule.schedule_id),
                schedule.task_name,
                schedule.model_dump_json(
                    exclude={"schedule_id", "task_name"},
                ),
            ],
        )

delete_schedule async 

delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(DELETE_SCHEDULE_QUERY).format(sql.Identifier(self._table_name)),
            [schedule_id],
        )

post_send async 

post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
189
190
191
192
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)

broker

logger module-attribute 

logger = getLogger('taskiq.psycopg_broker')

PsycopgBroker 

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: None = ...,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: AsyncConnectionPool,
    read_connection: None = ...,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: None = ...,
    read_connection: AsyncConnection,
)

PsycopgBroker(
    dsn: str | Callable[[], str] = ...,
    result_backend: AsyncResultBackend[_T] | None = ...,
    task_id_generator: Callable[[], str] | None = ...,
    channel_name: str = ...,
    table_name: str = ...,
    max_retry_attempts: int = ...,
    read_kwargs: dict[str, Any] | None = ...,
    write_kwargs: dict[str, Any] | None = ...,
    *,
    write_pool: AsyncConnectionPool,
    read_connection: AsyncConnection,
)

PsycopgBroker(
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend=None,
    task_id_generator=None,
    channel_name="taskiq",
    table_name="taskiq_messages",
    max_retry_attempts=5,
    read_kwargs=None,
    write_kwargs=None,
    *,
    write_pool=None,
    read_connection=None,
)

Bases: BasePostgresBroker

Broker that uses PostgreSQL and psycopg with LISTEN/NOTIFY.

Construct a new PsycopgBroker.

Parameters:

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable.

  • result_backend (AsyncResultBackend[_T] | None, default: None ) –

    Custom result backend.

  • task_id_generator (Callable[[], str] | None, default: None ) –

    Custom task_id generator.

  • channel_name (str, default: 'taskiq' ) –

    Name of the LISTEN/NOTIFY channel.

  • table_name (str, default: 'taskiq_messages' ) –

    Name of the table used to store messages.

  • max_retry_attempts (int, default: 5 ) –

    Maximum number of message processing attempts.

  • read_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to AsyncConnection.connect().

  • write_kwargs (dict[str, Any] | None, default: None ) –

    Extra kwargs forwarded to AsyncConnectionPool().

  • write_pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse for writes.

  • read_connection (AsyncConnection | None, default: None ) –

    An existing connection to reuse for LISTEN.

Source code in src/taskiq_pg/psycopg/broker.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def __init__(  # noqa: PLR0913
    self,
    dsn: str | Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    result_backend: AsyncResultBackend[_T] | None = None,
    task_id_generator: Callable[[], str] | None = None,
    channel_name: str = "taskiq",
    table_name: str = "taskiq_messages",
    max_retry_attempts: int = 5,
    read_kwargs: dict[str, tp.Any] | None = None,
    write_kwargs: dict[str, tp.Any] | None = None,
    *,
    write_pool: AsyncConnectionPool | None = None,
    read_connection: AsyncConnection | None = None,
) -> None:
    """
    Construct a new PsycopgBroker.

    Args:
        dsn: PostgreSQL connection string or callable.
        result_backend: Custom result backend.
        task_id_generator: Custom task_id generator.
        channel_name: Name of the LISTEN/NOTIFY channel.
        table_name: Name of the table used to store messages.
        max_retry_attempts: Maximum number of message processing attempts.
        read_kwargs: Extra kwargs forwarded to `AsyncConnection.connect()`.
        write_kwargs: Extra kwargs forwarded to `AsyncConnectionPool()`.
        write_pool: An existing connection pool to reuse for writes.
        read_connection: An existing connection to reuse for LISTEN.
    """
    super().__init__(
        dsn=dsn,
        result_backend=result_backend,
        task_id_generator=task_id_generator,
        channel_name=channel_name,
        table_name=table_name,
        max_retry_attempts=max_retry_attempts,
        read_kwargs=read_kwargs,
        write_kwargs=write_kwargs,
    )

    self._owns_write_pool = True
    if write_pool is not None:
        self._write_pool = write_pool
        self._owns_write_pool = False

    self._owns_read_conn = True
    if read_connection is not None:
        self._read_conn = read_connection
        self._owns_read_conn = False
channel_name instance-attribute 
channel_name = channel_name
table_name instance-attribute 
table_name = table_name
read_kwargs instance-attribute 
read_kwargs = read_kwargs or {}
write_kwargs instance-attribute 
write_kwargs = write_kwargs or {}
max_retry_attempts instance-attribute 
max_retry_attempts = max_retry_attempts
dsn property 
dsn

Get the DSN string.

Returns:

  • str –

    A string with dsn or None if dsn isn't set yet.

startup async 
startup()

Initialize the broker.

Source code in src/taskiq_pg/psycopg/broker.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
async def startup(self) -> None:
    """Initialize the broker."""
    await super().startup()

    if self._owns_read_conn:
        self._read_conn = await AsyncConnection.connect(
            conninfo=self.dsn,
            **self.read_kwargs,
            autocommit=True,
            cursor_factory=AsyncRawCursor,
        )

    if self._owns_write_pool:
        self._write_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self.write_kwargs,
        )

    if self._write_pool.closed:
        await self._write_pool.open()

    async with self._write_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(sql.SQL(CREATE_MESSAGE_TABLE_QUERY).format(sql.Identifier(self.table_name)))

    await self._read_conn.execute(sql.SQL("LISTEN {}").format(sql.Identifier(self.channel_name)))
    self._notifies_iter = self._read_conn.notifies()
shutdown async 
shutdown()

Close all connections on shutdown.

Source code in src/taskiq_pg/psycopg/broker.py 
178
179
180
181
182
183
184
185
186
187
188
async def shutdown(self) -> None:
    """Close all connections on shutdown."""
    await super().shutdown()
    if self._notifies_iter is not None:
        with suppress(RuntimeError):  # RuntimeError: aclose(): asynchronous generator is already running
            await self._notifies_iter.aclose()  # type: ignore[attr-defined]
    if self._read_conn is not None and self._owns_read_conn:
            await self._read_conn.notifies().aclose()
            await self._read_conn.close()
    if self._write_pool is not None and self._owns_write_pool:
        await self._write_pool.close()
kick async 
kick(message)

Send message to the channel.

Inserts the message into the database and sends a NOTIFY.

:param message: Message to send.

Source code in src/taskiq_pg/psycopg/broker.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
async def kick(self, message: BrokerMessage) -> None:
    """
    Send message to the channel.

    Inserts the message into the database and sends a NOTIFY.

    :param message: Message to send.
    """
    async with self._write_pool.connection() as connection, connection.cursor() as cursor:
        # insert message into db table
        await cursor.execute(
            sql.SQL(INSERT_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
            [
                message.task_id,
                message.task_name,
                message.message.decode(),
                json.dumps(message.labels),
            ],
        )
        row = await cursor.fetchone()
        if row is None:
            msg = "failed to insert message"
            raise RuntimeError(msg)
        message_inserted_id = int(row[0])

        delay_value = tp.cast("str | None", message.labels.get("delay"))
        if delay_value is not None:
            delay_seconds = int(delay_value)
            await self._schedule_notification(message_inserted_id, delay_seconds)
        else:
            # Send NOTIFY with message ID as payload
            await cursor.execute(
                sql.SQL("NOTIFY {}, {}").format(
                    sql.Identifier(self.channel_name),
                    sql.Literal(str(message_inserted_id)),
                ),
            )
listen async 
listen()

Listen to the channel.

Yields messages as they are received.

:yields: AckableMessage instances.

Source code in src/taskiq_pg/psycopg/broker.py 
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
async def listen(self) -> AsyncGenerator[AckableMessage, None]:
    """
    Listen to the channel.

    Yields messages as they are received.

    :yields: AckableMessage instances.
    """
    while True:
        async for message_id_str in self._listen_context():
            message_id = int(message_id_str)  # payload is the message id
            try:
                async with self._write_pool.connection() as connection, connection.cursor() as cursor:
                    await cursor.execute(
                        sql.SQL(CLAIM_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
                        [message_id],
                    )
                    claimed_message = await cursor.fetchone()
                    if claimed_message is None:
                        continue
            except psycopg.OperationalError:  # message was claimed by another worker
                continue
            message_str = claimed_message[3]
            if not isinstance(message_str, str):
                msg = "Message is not a string"
                raise TypeError(msg)
            message_data = message_str.encode()

            async def ack(*, _message_id: int = message_id) -> None:
                async with self._write_pool.connection() as connection, connection.cursor() as cursor:
                    await cursor.execute(
                        sql.SQL(DELETE_MESSAGE_QUERY).format(sql.Identifier(self.table_name)),
                        [_message_id],
                    )

            yield AckableMessage(data=message_data, ack=ack)

queries

CREATE_TABLE_QUERY module-attribute 

CREATE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    task_id {} UNIQUE,\n    result BYTEA,\n    progress BYTEA\n)\n"

ADD_PROGRESS_COLUMN_QUERY module-attribute 

ADD_PROGRESS_COLUMN_QUERY = "\nALTER TABLE {} ADD COLUMN IF NOT EXISTS progress BYTEA;\n"

CREATE_INDEX_QUERY module-attribute 

CREATE_INDEX_QUERY = "\nCREATE INDEX IF NOT EXISTS {} ON {} USING HASH (task_id)\n"

INSERT_RESULT_QUERY module-attribute 

INSERT_RESULT_QUERY = "\nINSERT INTO {} VALUES (%s, %s)\nON CONFLICT (task_id)\nDO UPDATE\nSET result = EXCLUDED.result;\n"

INSERT_PROGRESS_QUERY module-attribute 

INSERT_PROGRESS_QUERY = "\nINSERT INTO {} VALUES (%s, NULL, %s)\nON CONFLICT (task_id)\nDO UPDATE\nSET progress = %s\n"

SELECT_PROGRESS_QUERY module-attribute 

SELECT_PROGRESS_QUERY = (
    "\nSELECT progress FROM {} WHERE task_id = %s\n"
)

IS_RESULT_EXISTS_QUERY module-attribute 

IS_RESULT_EXISTS_QUERY = "\nSELECT EXISTS(\n    SELECT 1 FROM {} WHERE task_id = %s and result IS NOT NULL\n);\n"

SELECT_RESULT_QUERY module-attribute 

SELECT_RESULT_QUERY = (
    "\nSELECT result FROM {} WHERE task_id = %s;\n"
)

DELETE_RESULT_QUERY module-attribute 

DELETE_RESULT_QUERY = (
    "\nDELETE FROM {} WHERE task_id = %s;\n"
)

CREATE_MESSAGE_TABLE_QUERY module-attribute 

CREATE_MESSAGE_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id SERIAL PRIMARY KEY,\n    task_id VARCHAR NOT NULL,\n    task_name VARCHAR NOT NULL,\n    message TEXT NOT NULL,\n    labels JSONB NOT NULL,\n    status TEXT NOT NULL DEFAULT 'pending',\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_MESSAGE_QUERY module-attribute 

INSERT_MESSAGE_QUERY = "\nINSERT INTO {} (task_id, task_name, message, labels)\nVALUES (%s, %s, %s, %s)\nRETURNING id\n"

CLAIM_MESSAGE_QUERY module-attribute 

CLAIM_MESSAGE_QUERY = "UPDATE {} SET status = 'processing' WHERE id = %s AND status = 'pending' RETURNING *"

DELETE_MESSAGE_QUERY module-attribute 

DELETE_MESSAGE_QUERY = 'DELETE FROM {} WHERE id = %s'

CREATE_SCHEDULES_TABLE_QUERY module-attribute 

CREATE_SCHEDULES_TABLE_QUERY = "\nCREATE TABLE IF NOT EXISTS {} (\n    id UUID PRIMARY KEY,\n    task_name VARCHAR(100) NOT NULL,\n    schedule JSONB NOT NULL,\n    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),\n    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()\n);\n"

INSERT_SCHEDULE_QUERY module-attribute 

INSERT_SCHEDULE_QUERY = "\nINSERT INTO {} (id, task_name, schedule)\nVALUES (%s, %s, %s)\nON CONFLICT (id) DO UPDATE\nSET task_name = EXCLUDED.task_name,\n    schedule = EXCLUDED.schedule,\n    updated_at = NOW();\n"

SELECT_SCHEDULES_QUERY module-attribute 

SELECT_SCHEDULES_QUERY = (
    "\nSELECT id, task_name, schedule\nFROM {};\n"
)

DELETE_ALL_SCHEDULES_QUERY module-attribute 

DELETE_ALL_SCHEDULES_QUERY = '\nDELETE FROM {};\n'

DELETE_SCHEDULE_QUERY module-attribute 

DELETE_SCHEDULE_QUERY = '\nDELETE FROM {} WHERE id = %s;\n'

result_backend

PsycopgResultBackend 

PsycopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PsycopgResultBackend(
    dsn: Callable[[], str] | str | None = ...,
    keep_results: bool = ...,
    table_name: str = ...,
    field_for_task_id: Literal[
        "VarChar", "Text", "Uuid"
    ] = ...,
    serializer: TaskiqSerializer | None = ...,
    *,
    pool: AsyncConnectionPool,
)

PsycopgResultBackend(
    dsn="postgres://postgres:postgres@localhost:5432/postgres",
    keep_results=True,
    table_name="taskiq_results",
    field_for_task_id="VarChar",
    serializer=None,
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresResultBackend

Result backend for TaskIQ based on psycopg.

Construct a new PsycopgResultBackend.

Parameters:

  • dsn (Callable[[], str] | str | None, default: 'postgres://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Can be None if pool is provided.

  • keep_results (bool, default: True ) –

    Whether to keep results after reading.

  • table_name (str, default: 'taskiq_results' ) –

    Table to store results in.

  • field_for_task_id (Literal['VarChar', 'Text', 'Uuid'], default: 'VarChar' ) –

    Column type for task_id.

  • serializer (TaskiqSerializer | None, default: None ) –

    Serializer for task results.

  • pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psycopg/result_backend.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    dsn: tp.Callable[[], str] | str | None = "postgres://postgres:postgres@localhost:5432/postgres",
    keep_results: bool = True,
    table_name: str = "taskiq_results",
    field_for_task_id: tp.Literal["VarChar", "Text", "Uuid"] = "VarChar",
    serializer: TaskiqSerializer | None = None,
    *,
    pool: AsyncConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PsycopgResultBackend.

    Args:
        dsn: PostgreSQL connection string or callable. Can be None if pool is provided.
        keep_results: Whether to keep results after reading.
        table_name: Table to store results in.
        field_for_task_id: Column type for task_id.
        serializer: Serializer for task results.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.
    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        dsn=dsn,
        keep_results=keep_results,
        table_name=table_name,
        field_for_task_id=field_for_task_id,
        serializer=serializer,
        **connect_kwargs,
    )
keep_results instance-attribute 
keep_results = keep_results
table_name instance-attribute 
table_name = table_name
field_for_task_id instance-attribute 
field_for_task_id = field_for_task_id
connect_kwargs instance-attribute 
connect_kwargs = connect_kwargs
serializer instance-attribute 
serializer = serializer or PickleSerializer()
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the result backend.

Construct new connection pool (if not provided externally) and create new table for results if not exists.

Source code in src/taskiq_pg/psycopg/result_backend.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
async def startup(self) -> None:
    """
    Initialize the result backend.

    Construct new connection pool (if not provided externally) and create new table
    for results if not exists.
    """
    if self._owns_pool:
        self._database_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self.connect_kwargs,
        )

    if self._database_pool.closed:
        await self._database_pool.open()

    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.CREATE_TABLE_QUERY).format(
                sql.Identifier(self.table_name),
                sql.SQL(self.field_for_task_id),
            ),
        )
        await cursor.execute(
            query=sql.SQL(queries.ADD_PROGRESS_COLUMN_QUERY).format(
                sql.Identifier(self.table_name),
            ),
        )
        await cursor.execute(
            query=sql.SQL(queries.CREATE_INDEX_QUERY).format(
                sql.Identifier(self.table_name + "_task_id_idx"),
                sql.Identifier(self.table_name),
            ),
        )
shutdown async 
shutdown()

Close the connection pool (only if owned by this backend).

Source code in src/taskiq_pg/psycopg/result_backend.py 
119
120
121
122
async def shutdown(self) -> None:
    """Close the connection pool (only if owned by this backend)."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()
set_result async 
set_result(task_id, result)

Set result to the PostgreSQL table.

:param task_id: ID of the task. :param result: result of the task.

Source code in src/taskiq_pg/psycopg/result_backend.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
async def set_result(
    self,
    task_id: str,
    result: TaskiqResult[ReturnType],
) -> None:
    """
    Set result to the PostgreSQL table.

    :param task_id: ID of the task.
    :param result: result of the task.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.INSERT_RESULT_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[
                task_id,
                self.serializer.dumpb(model_dump(result)),
            ],
        )
is_result_ready async 
is_result_ready(task_id)

Returns whether the result is ready.

:param task_id: ID of the task.

:returns: True if the result is ready else False.

Source code in src/taskiq_pg/psycopg/result_backend.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
async def is_result_ready(self, task_id: str) -> bool:
    """
    Returns whether the result is ready.

    :param task_id: ID of the task.

    :returns: True if the result is ready else False.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.IS_RESULT_EXISTS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        row = await execute_result.fetchone()
        return bool(row and row[0])
get_result async 
get_result(task_id, with_logs=False)

Retrieve result from the task.

:param task_id: task's id. :param with_logs: if True it will download task's logs. :raises ResultIsMissingError: if there is no result when trying to get it. :return: TaskiqResult.

Source code in src/taskiq_pg/psycopg/result_backend.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
async def get_result(
    self,
    task_id: str,
    with_logs: bool = False,
) -> TaskiqResult[ReturnType]:
    """
    Retrieve result from the task.

    :param task_id: task's id.
    :param with_logs: if True it will download task's logs.
    :raises ResultIsMissingError: if there is no result when trying to get it.
    :return: TaskiqResult.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.SELECT_RESULT_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        result = await execute_result.fetchone()
        if result is None:
            msg = f"Cannot find record with task_id = {task_id} in PostgreSQL"
            raise ResultIsMissingError(msg)
        result_in_bytes: tp.Final = result[0]

        if not self.keep_results:
            await cursor.execute(
                query=sql.SQL(queries.DELETE_RESULT_QUERY).format(
                    sql.Identifier(self.table_name),
                ),
                params=[task_id],
            )

        taskiq_result: tp.Final = model_validate(
            TaskiqResult[ReturnType],
            self.serializer.loadb(result_in_bytes),
        )

        if not with_logs:
            taskiq_result.log = None

        return taskiq_result
set_progress async 
set_progress(task_id, progress)

Saves progress.

:param task_id: task's id. :param progress: progress of execution.

Source code in src/taskiq_pg/psycopg/result_backend.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
async def set_progress(
    self,
    task_id: str,
    progress: TaskProgress[tp.Any],
) -> None:
    """
    Saves progress.

    :param task_id: task's id.
    :param progress: progress of execution.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            query=sql.SQL(queries.INSERT_PROGRESS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[
                task_id,
                self.serializer.dumpb(model_dump(progress)),
                self.serializer.dumpb(model_dump(progress)),
            ],
        )
get_progress async 
get_progress(task_id)

Gets progress.

:param task_id: task's id.

Source code in src/taskiq_pg/psycopg/result_backend.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
async def get_progress(
    self,
    task_id: str,
) -> TaskProgress[tp.Any] | None:
    """
    Gets progress.

    :param task_id: task's id.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        execute_result = await cursor.execute(
            query=sql.SQL(queries.SELECT_PROGRESS_QUERY).format(
                sql.Identifier(self.table_name),
            ),
            params=[task_id],
        )
        progress_in_bytes = await execute_result.fetchone()
        if progress_in_bytes is None or progress_in_bytes[0] is None:
            return None
        return model_validate(
            TaskProgress[tp.Any],
            self.serializer.loadb(progress_in_bytes[0]),
        )

schedule_source

logger module-attribute 

logger = getLogger('taskiq_pg.psycopg_schedule_source')

PsycopgScheduleSource 

PsycopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: None = ...,
    **connect_kwargs: Any,
)

PsycopgScheduleSource(
    broker: AsyncBroker,
    dsn: str | Callable[[], str] = ...,
    table_name: str = ...,
    *,
    pool: AsyncConnectionPool,
)

PsycopgScheduleSource(
    broker,
    dsn="postgresql://postgres:postgres@localhost:5432/postgres",
    table_name="taskiq_schedules",
    *,
    pool=None,
    **connect_kwargs,
)

Bases: BasePostgresScheduleSource

Schedule source that uses psycopg to store schedules in PostgreSQL.

Construct a new PsycopgScheduleSource.

Parameters:

  • broker (AsyncBroker) –

    The TaskIQ broker instance.

  • dsn (str | Callable[[], str], default: 'postgresql://postgres:postgres@localhost:5432/postgres' ) –

    PostgreSQL connection string or callable. Ignored in pool mode.

  • table_name (str, default: 'taskiq_schedules' ) –

    Table to store schedules in.

  • pool (AsyncConnectionPool | None, default: None ) –

    An existing connection pool to reuse.

  • **connect_kwargs (Any, default: {} ) –

    Extra kwargs for connection pool creation.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    broker: AsyncBroker,
    dsn: str | tp.Callable[[], str] = "postgresql://postgres:postgres@localhost:5432/postgres",
    table_name: str = "taskiq_schedules",
    *,
    pool: AsyncConnectionPool | None = None,
    **connect_kwargs: tp.Any,
) -> None:
    """
    Construct a new PsycopgScheduleSource.

    Args:
        broker: The TaskIQ broker instance.
        dsn: PostgreSQL connection string or callable. Ignored in pool mode.
        table_name: Table to store schedules in.
        pool: An existing connection pool to reuse.
        **connect_kwargs: Extra kwargs for connection pool creation.

    """
    self._owns_pool = True
    if pool is not None:
        self._owns_pool = False
        self._database_pool = pool

    super().__init__(
        broker=broker,
        dsn=dsn,
        table_name=table_name,
        **connect_kwargs,
    )
dsn property 
dsn

Get the DSN string.

Returns the DSN string or None if not set.

startup async 
startup()

Initialize the schedule source.

Construct new connection pool (if not provided externally), create new table for schedules if not exists and fill table with schedules from task labels.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
async def startup(self) -> None:
    """
    Initialize the schedule source.

    Construct new connection pool (if not provided externally), create new table for
    schedules if not exists and fill table with schedules from task labels.
    """
    if self._owns_pool:
        self._database_pool = AsyncConnectionPool(
            conninfo=self.dsn if self.dsn is not None else "",
            open=False,
            **self._connect_kwargs,
        )

    if self._database_pool.closed:
        await self._database_pool.open()

    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(CREATE_SCHEDULES_TABLE_QUERY).format(sql.Identifier(self._table_name)),
        )
    scheduled_tasks_for_creation = self.extract_scheduled_tasks_from_broker()
    await self._update_schedules_on_startup(scheduled_tasks_for_creation)
shutdown async 
shutdown()

Close the connection pool (only if owned by this source).

Source code in src/taskiq_pg/psycopg/schedule_source.py 
124
125
126
127
async def shutdown(self) -> None:
    """Close the connection pool (only if owned by this source)."""
    if self._owns_pool and getattr(self, "_database_pool", None) is not None:
        await self._database_pool.close()
get_schedules async 
get_schedules()

Fetch schedules from the database.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
async def get_schedules(self) -> list["ScheduledTask"]:
    """Fetch schedules from the database."""
    schedules = []
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        rows_with_schedules = await cursor.execute(
            sql.SQL(SELECT_SCHEDULES_QUERY).format(sql.Identifier(self._table_name)),
        )
        rows = await rows_with_schedules.fetchall()
        for schedule_id, task_name, schedule in rows:
            schedules.append(
                ScheduledTask.model_validate(
                    {
                        "schedule_id": str(schedule_id),
                        "task_name": task_name,
                        "labels": schedule["labels"],
                        "args": schedule["args"],
                        "kwargs": schedule["kwargs"],
                        "cron": schedule["cron"],
                        "cron_offset": schedule["cron_offset"],
                        "time": schedule["time"],
                        "interval": schedule["interval"],
                    },
                ),
            )
    return schedules
add_schedule async 
add_schedule(schedule)

Add a new schedule.

Parameters:

  • schedule (ScheduledTask) –

    schedule to add.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
async def add_schedule(self, schedule: "ScheduledTask") -> None:
    """
    Add a new schedule.

    Args:
        schedule: schedule to add.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(INSERT_SCHEDULE_QUERY).format(sql.Identifier(self._table_name)),
            [
                uuid.UUID(schedule.schedule_id),
                schedule.task_name,
                schedule.model_dump_json(
                    exclude={"schedule_id", "task_name"},
                ),
            ],
        )
delete_schedule async 
delete_schedule(schedule_id)

Method to delete schedule by id.

This is useful for schedule cancelation.

Parameters:

  • schedule_id (str) –

    id of schedule to delete.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
async def delete_schedule(self, schedule_id: str) -> None:
    """
    Method to delete schedule by id.

    This is useful for schedule cancelation.

    Args:
        schedule_id: id of schedule to delete.
    """
    async with self._database_pool.connection() as connection, connection.cursor() as cursor:
        await cursor.execute(
            sql.SQL(DELETE_SCHEDULE_QUERY).format(sql.Identifier(self._table_name)),
            [schedule_id],
        )
post_send async 
post_send(task)

Delete a task after it's completed.

Source code in src/taskiq_pg/psycopg/schedule_source.py 
189
190
191
192
async def post_send(self, task: ScheduledTask) -> None:
    """Delete a task after it's completed."""
    if task.time is not None:
        await self.delete_schedule(task.schedule_id)
extract_scheduled_tasks_from_broker 
extract_scheduled_tasks_from_broker()

Extract schedules from tasks that were registered in broker.

Returns:

  • list[ScheduledTask] –

    A list of ScheduledTask instances extracted from the task's labels.

Source code in src/taskiq_pg/_internal/schedule_source.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def extract_scheduled_tasks_from_broker(self) -> list[ScheduledTask]:
    """
    Extract schedules from tasks that were registered in broker.

    Returns:
        A list of ScheduledTask instances extracted from the task's labels.
    """
    scheduled_tasks_for_creation: list[ScheduledTask] = []
    for task_name, task in self._broker.get_all_tasks().items():
        if "schedule" not in task.labels:
            logger.debug("Task %s has no schedule, skipping", task_name)
            continue
        if not isinstance(task.labels["schedule"], list):
            logger.warning(
                "Schedule for task %s is not a list, skipping",
                task_name,
            )
            continue
        for schedule in task.labels["schedule"]:
            try:
                new_schedule = ScheduledTask.model_validate(
                    {
                        "task_name": task_name,
                        "labels": schedule.get("labels", {}),
                        "args": schedule.get("args", []),
                        "kwargs": schedule.get("kwargs", {}),
                        "schedule_id": str(uuid.uuid4()),
                        "cron": schedule.get("cron", None),
                        "cron_offset": schedule.get("cron_offset", None),
                        "interval": schedule.get("interval", None),
                        "time": schedule.get("time", None),
                    },
                )
                scheduled_tasks_for_creation.append(new_schedule)
            except ValidationError:  # noqa: PERF203
                logger.exception(
                    "Schedule for task %s is not valid, skipping",
                    task_name,
                )
                continue
    return scheduled_tasks_for_creation