goatbeat-postgres 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Louis Midson LAJEANTY <midsonlajeanty@proton.me>
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,29 @@
1
+ Metadata-Version: 2.4
2
+ Name: goatbeat-postgres
3
+ Version: 0.1.0
4
+ Summary: PostgreSQL storage, fenced leases and LISTEN/NOTIFY for Goatbeat.
5
+ Keywords: scheduler,cron,periodic-tasks,distributed,django-tasks,leader-election
6
+ Author: Louis Midson LAJEANTY
7
+ Author-email: Louis Midson LAJEANTY <midsonlajeanty@proton.me>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
15
+ Classifier: Topic :: System :: Distributed Computing
16
+ Classifier: Typing :: Typed
17
+ Requires-Dist: goatbeat-core==0.1.0
18
+ Requires-Dist: psycopg>=3.2
19
+ Requires-Dist: psycopg-pool>=3.2
20
+ Requires-Python: >=3.13
21
+ Project-URL: Homepage, https://github.com/goatlabshq/goatbeat
22
+ Project-URL: Documentation, https://goatlabshq.github.io/goatbeat/
23
+ Project-URL: Issues, https://github.com/goatlabshq/goatbeat/issues
24
+ Description-Content-Type: text/markdown
25
+
26
+ # goatbeat-postgres
27
+
28
+ PostgreSQL storage and distributed coordination for [Goatbeat](https://github.com/goatlabshq/goatbeat):
29
+ schema migrations, the schedule store, the fenced leader lease and LISTEN/NOTIFY wake-ups.
@@ -0,0 +1,4 @@
1
+ # goatbeat-postgres
2
+
3
+ PostgreSQL storage and distributed coordination for [Goatbeat](https://github.com/goatlabshq/goatbeat):
4
+ schema migrations, the schedule store, the fenced leader lease and LISTEN/NOTIFY wake-ups.
@@ -0,0 +1,49 @@
1
+ [project]
2
+ name = "goatbeat-postgres"
3
+ version = "0.1.0"
4
+ description = "PostgreSQL storage, fenced leases and LISTEN/NOTIFY for Goatbeat."
5
+ readme = "README.md"
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ requires-python = ">=3.13"
9
+ keywords = [
10
+ "scheduler",
11
+ "cron",
12
+ "periodic-tasks",
13
+ "distributed",
14
+ "django-tasks",
15
+ "leader-election",
16
+ ]
17
+ classifiers = [
18
+ "Development Status :: 3 - Alpha",
19
+ "Intended Audience :: Developers",
20
+ "Programming Language :: Python :: 3",
21
+ "Programming Language :: Python :: 3.13",
22
+ "Programming Language :: Python :: 3.14",
23
+ "Topic :: System :: Distributed Computing",
24
+ "Typing :: Typed",
25
+ ]
26
+ dependencies = [
27
+ "goatbeat-core==0.1.0",
28
+ "psycopg>=3.2",
29
+ "psycopg-pool>=3.2",
30
+ ]
31
+
32
+ [[project.authors]]
33
+ name = "Louis Midson LAJEANTY"
34
+ email = "midsonlajeanty@proton.me"
35
+
36
+ [project.urls]
37
+ Homepage = "https://github.com/goatlabshq/goatbeat"
38
+ Documentation = "https://goatlabshq.github.io/goatbeat/"
39
+ Issues = "https://github.com/goatlabshq/goatbeat/issues"
40
+
41
+ [build-system]
42
+ requires = ["uv_build>=0.12,<0.13"]
43
+ build-backend = "uv_build"
44
+
45
+ [tool.uv.build-backend]
46
+ module-name = "goatbeat.postgres"
47
+
48
+ [tool.uv.sources.goatbeat-core]
49
+ workspace = true
@@ -0,0 +1,39 @@
1
+ [project]
2
+ name = "goatbeat-postgres"
3
+ version = "0.1.0" # x-release-please-version
4
+ description = "PostgreSQL storage, fenced leases and LISTEN/NOTIFY for Goatbeat."
5
+ readme = "README.md"
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ authors = [{ name = "Louis Midson LAJEANTY", email = "midsonlajeanty@proton.me" }]
9
+ requires-python = ">=3.13"
10
+ keywords = ["scheduler", "cron", "periodic-tasks", "distributed", "django-tasks", "leader-election"]
11
+ classifiers = [
12
+ "Development Status :: 3 - Alpha",
13
+ "Intended Audience :: Developers",
14
+ "Programming Language :: Python :: 3",
15
+ "Programming Language :: Python :: 3.13",
16
+ "Programming Language :: Python :: 3.14",
17
+ "Topic :: System :: Distributed Computing",
18
+ "Typing :: Typed",
19
+ ]
20
+ dependencies = [
21
+ "goatbeat-core==0.1.0", # x-release-please-version
22
+ "psycopg>=3.2",
23
+ "psycopg-pool>=3.2",
24
+ ]
25
+
26
+ [project.urls]
27
+ Homepage = "https://github.com/goatlabshq/goatbeat"
28
+ Documentation = "https://goatlabshq.github.io/goatbeat/"
29
+ Issues = "https://github.com/goatlabshq/goatbeat/issues"
30
+
31
+ [build-system]
32
+ requires = ["uv_build>=0.12,<0.13"]
33
+ build-backend = "uv_build"
34
+
35
+ [tool.uv.build-backend]
36
+ module-name = "goatbeat.postgres"
37
+
38
+ [tool.uv.sources]
39
+ goatbeat-core = { workspace = true }
@@ -0,0 +1,74 @@
1
+ import logging
2
+ import math
3
+ import time
4
+ from datetime import timedelta
5
+ from types import TracebackType
6
+ from typing import Final, Protocol, Self
7
+
8
+ import psycopg
9
+ from psycopg import sql
10
+ from psycopg.conninfo import make_conninfo
11
+
12
+ from goatbeat.core.observability import LOGGER_NAME, Event, log_event
13
+ from goatbeat.postgres.provider import KEEPALIVES
14
+
15
+ CHANNEL: Final = "goatbeat_schedule_changed"
16
+ # libpq rounds connect_timeout below 2 seconds up to 2.
17
+ MIN_CONNECT_TIMEOUT_SECONDS: Final = 2
18
+
19
+
20
+ class Notifier(Protocol):
21
+ def wait(self, timeout: timedelta) -> bool:
22
+ """Block until a change notification or the timeout; True when notified."""
23
+ ...
24
+
25
+
26
+ class PsycopgNotifier:
27
+ """Owns one dedicated connection: LISTEN cannot share a pooled, transactional one.
28
+
29
+ Losing it only degrades wake-ups to polling, so errors are logged and retried lazily.
30
+ """
31
+
32
+ def __init__(self, conninfo: str, logger: logging.Logger | None = None) -> None:
33
+ self._conninfo = conninfo
34
+ self._logger = logger or logging.getLogger(LOGGER_NAME)
35
+ self._connection: psycopg.Connection[tuple[object, ...]] | None = None
36
+
37
+ def wait(self, timeout: timedelta) -> bool:
38
+ started = time.monotonic()
39
+ try:
40
+ connection = self._listening_connection(timeout)
41
+ remaining = timeout.total_seconds() - (time.monotonic() - started)
42
+ for _ in connection.notifies(timeout=max(0.0, remaining), stop_after=1):
43
+ return True
44
+ except psycopg.Error as error:
45
+ log_event(self._logger, Event.STORE_UNAVAILABLE, logging.WARNING, error=str(error))
46
+ self.close()
47
+ time.sleep(max(0.0, timeout.total_seconds() - (time.monotonic() - started)))
48
+ return False
49
+
50
+ def close(self) -> None:
51
+ if self._connection is not None:
52
+ self._connection.close()
53
+ self._connection = None
54
+
55
+ def _listening_connection(self, timeout: timedelta) -> psycopg.Connection[tuple[object, ...]]:
56
+ if self._connection is None or self._connection.closed:
57
+ # Without a bound, an unroutable host blocks for the OS TCP timeout (minutes).
58
+ connect_timeout = max(MIN_CONNECT_TIMEOUT_SECONDS, math.ceil(timeout.total_seconds()))
59
+ conninfo = make_conninfo(self._conninfo, connect_timeout=connect_timeout, **KEEPALIVES)
60
+ connection = psycopg.connect(conninfo, autocommit=True)
61
+ connection.execute(sql.SQL("LISTEN {}").format(sql.Identifier(CHANNEL)))
62
+ self._connection = connection
63
+ return self._connection
64
+
65
+ def __enter__(self) -> Self:
66
+ return self
67
+
68
+ def __exit__(
69
+ self,
70
+ exc_type: type[BaseException] | None,
71
+ exc: BaseException | None,
72
+ traceback: TracebackType | None,
73
+ ) -> None:
74
+ self.close()
@@ -0,0 +1,112 @@
1
+ from collections.abc import Generator, Sequence
2
+ from contextlib import AbstractContextManager, contextmanager
3
+ from types import TracebackType
4
+ from typing import Final, LiteralString, Protocol, Self
5
+
6
+ import psycopg
7
+ from psycopg.conninfo import conninfo_to_dict
8
+ from psycopg.rows import TupleRow
9
+ from psycopg_pool import ConnectionPool, PoolTimeout
10
+
11
+ from goatbeat.core.errors import StoreUnavailableError
12
+
13
+ # Errors meaning "this database cannot serve us right now": a promoted replica still read-only,
14
+ # a session the server ended, a lost connection.
15
+ UNAVAILABLE_ERRORS: Final = (
16
+ psycopg.OperationalError,
17
+ psycopg.InterfaceError,
18
+ psycopg.errors.ReadOnlySqlTransaction,
19
+ psycopg.errors.IdleInTransactionSessionTimeout,
20
+ PoolTimeout,
21
+ )
22
+ CONNECT_TIMEOUT_SECONDS: Final = 5
23
+ POOL_WAIT_SECONDS: Final = 10.0
24
+ # Detect half-open connections (network partition) in about a minute instead of hours.
25
+ KEEPALIVES: Final = {
26
+ "keepalives": 1,
27
+ "keepalives_idle": 30,
28
+ "keepalives_interval": 10,
29
+ "keepalives_count": 3,
30
+ }
31
+
32
+
33
+ class Cursor(Protocol):
34
+ """The DB-API subset Goatbeat uses; psycopg and Django cursors both satisfy it."""
35
+
36
+ @property
37
+ def rowcount(self) -> int: ...
38
+
39
+ def execute(
40
+ self, query: LiteralString, params: Sequence[object] | None = None, /
41
+ ) -> object: ...
42
+
43
+ def fetchone(self) -> Sequence[object] | None: ...
44
+
45
+ def fetchall(self) -> Sequence[Sequence[object]]: ...
46
+
47
+
48
+ class ConnectionProvider(Protocol):
49
+ def transaction(self) -> AbstractContextManager[Cursor]:
50
+ """Yield a cursor inside a transaction committed on success, rolled back on error."""
51
+ ...
52
+
53
+
54
+ class PsycopgProvider:
55
+ def __init__(self, pool: ConnectionPool[psycopg.Connection[TupleRow]]) -> None:
56
+ self._pool = pool
57
+
58
+ @classmethod
59
+ def connect(
60
+ cls, conninfo: str, *, min_size: int = 1, max_size: int = 4, wait: bool = False
61
+ ) -> Self:
62
+ """``wait`` fails fast with StoreUnavailableError when the database is unreachable."""
63
+ defaults: dict[str, object] = {"connect_timeout": CONNECT_TIMEOUT_SECONDS, **KEEPALIVES}
64
+ given = conninfo_to_dict(conninfo)
65
+ pool = ConnectionPool(
66
+ conninfo,
67
+ connection_class=psycopg.Connection[TupleRow],
68
+ kwargs={key: value for key, value in defaults.items() if key not in given},
69
+ min_size=min_size,
70
+ max_size=max_size,
71
+ timeout=POOL_WAIT_SECONDS,
72
+ check=ConnectionPool.check_connection,
73
+ configure=_without_prepared_statements,
74
+ open=False,
75
+ )
76
+ try:
77
+ pool.open(wait=wait, timeout=POOL_WAIT_SECONDS)
78
+ except PoolTimeout as error:
79
+ pool.close()
80
+ raise StoreUnavailableError(f"cannot connect to PostgreSQL: {error}") from error
81
+ return cls(pool)
82
+
83
+ @contextmanager
84
+ def transaction(self) -> Generator[Cursor]:
85
+ try:
86
+ with (
87
+ self._pool.connection() as connection,
88
+ connection.transaction(),
89
+ connection.cursor() as cursor,
90
+ ):
91
+ yield cursor
92
+ except UNAVAILABLE_ERRORS as error:
93
+ raise StoreUnavailableError(str(error)) from error
94
+
95
+ def close(self) -> None:
96
+ self._pool.close()
97
+
98
+ def __enter__(self) -> Self:
99
+ return self
100
+
101
+ def __exit__(
102
+ self,
103
+ exc_type: type[BaseException] | None,
104
+ exc: BaseException | None,
105
+ traceback: TracebackType | None,
106
+ ) -> None:
107
+ self.close()
108
+
109
+
110
+ # Server-side prepared statements are session state: PgBouncer transaction mode breaks them.
111
+ def _without_prepared_statements(connection: psycopg.Connection[TupleRow]) -> None:
112
+ connection.prepare_threshold = None
File without changes
@@ -0,0 +1,149 @@
1
+ from datetime import datetime
2
+ from typing import Final, LiteralString
3
+
4
+ from goatbeat.core.errors import ConfigurationError, StoreUnavailableError
5
+ from goatbeat.core.time import Instant
6
+ from goatbeat.postgres.provider import ConnectionProvider
7
+
8
+ # Transaction-scoped, so it also works behind PgBouncer in transaction pooling mode.
9
+ _MIGRATION_LOCK: Final = "SELECT pg_advisory_xact_lock(hashtext('goatbeat.migrate'))"
10
+
11
+ _VERSION_TABLE: Final = """
12
+ CREATE TABLE IF NOT EXISTS goatbeat_schema_version (
13
+ version integer PRIMARY KEY,
14
+ applied_at timestamptz NOT NULL DEFAULT now()
15
+ )
16
+ """
17
+
18
+ _V1: Final[tuple[LiteralString, ...]] = (
19
+ """
20
+ CREATE TABLE goatbeat_schedule (
21
+ id text PRIMARY KEY,
22
+ name text NOT NULL UNIQUE,
23
+ source text NOT NULL,
24
+ task_path text NOT NULL,
25
+ task_args jsonb NOT NULL,
26
+ task_kwargs jsonb NOT NULL,
27
+ trigger_kind text NOT NULL,
28
+ interval_us bigint,
29
+ start_at timestamptz,
30
+ cron text,
31
+ timezone text,
32
+ misfire_policy text NOT NULL,
33
+ misfire_grace_us bigint NOT NULL,
34
+ max_catch_up integer NOT NULL,
35
+ queue text,
36
+ priority integer,
37
+ backend text,
38
+ enabled boolean NOT NULL,
39
+ next_run_at timestamptz NOT NULL,
40
+ created_at timestamptz NOT NULL DEFAULT statement_timestamp(),
41
+ updated_at timestamptz NOT NULL DEFAULT statement_timestamp()
42
+ )
43
+ """,
44
+ "CREATE INDEX goatbeat_schedule_due ON goatbeat_schedule (next_run_at) WHERE enabled",
45
+ "CREATE INDEX goatbeat_schedule_task ON goatbeat_schedule (task_path)",
46
+ """
47
+ CREATE TABLE goatbeat_execution (
48
+ id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
49
+ schedule_id text NOT NULL REFERENCES goatbeat_schedule (id) ON DELETE CASCADE
50
+ DEFERRABLE INITIALLY DEFERRED,
51
+ scheduled_at timestamptz NOT NULL,
52
+ idempotency_key text NOT NULL UNIQUE,
53
+ status text NOT NULL,
54
+ scheduler_id text NOT NULL,
55
+ fencing_token bigint NOT NULL,
56
+ claimed_at timestamptz NOT NULL,
57
+ missed_count integer NOT NULL,
58
+ retry_count integer NOT NULL,
59
+ dispatched_at timestamptz,
60
+ completed_at timestamptz,
61
+ next_retry_at timestamptz,
62
+ backend_result_id text,
63
+ error text
64
+ )
65
+ """,
66
+ (
67
+ "CREATE INDEX goatbeat_execution_history "
68
+ "ON goatbeat_execution (schedule_id, scheduled_at DESC)"
69
+ ),
70
+ (
71
+ "CREATE INDEX goatbeat_execution_retry "
72
+ "ON goatbeat_execution (next_retry_at) WHERE status = 'submit_failed'"
73
+ ),
74
+ (
75
+ "CREATE INDEX goatbeat_execution_requested "
76
+ "ON goatbeat_execution (claimed_at) WHERE status = 'requested'"
77
+ ),
78
+ (
79
+ "CREATE INDEX goatbeat_execution_claimed "
80
+ "ON goatbeat_execution (claimed_at) WHERE status = 'claimed'"
81
+ ),
82
+ (
83
+ "CREATE INDEX goatbeat_execution_result "
84
+ "ON goatbeat_execution (backend_result_id) WHERE backend_result_id IS NOT NULL"
85
+ ),
86
+ """
87
+ CREATE TABLE goatbeat_early_completion (
88
+ result_id text PRIMARY KEY,
89
+ status text NOT NULL,
90
+ completed_at timestamptz NOT NULL
91
+ )
92
+ """,
93
+ """
94
+ CREATE TABLE goatbeat_lease (
95
+ partition integer PRIMARY KEY,
96
+ holder_id text,
97
+ fencing_token bigint NOT NULL,
98
+ expires_at timestamptz NOT NULL
99
+ )
100
+ """,
101
+ (
102
+ "INSERT INTO goatbeat_lease (partition, holder_id, fencing_token, expires_at) "
103
+ "VALUES (0, NULL, 0, '-infinity')"
104
+ ),
105
+ )
106
+
107
+ MIGRATIONS: Final[tuple[tuple[LiteralString, ...], ...]] = (_V1,)
108
+ SCHEMA_VERSION: Final = len(MIGRATIONS)
109
+
110
+
111
+ def migrate(provider: ConnectionProvider) -> int:
112
+ """Apply pending migrations; concurrent callers are serialized by an advisory lock."""
113
+ with provider.transaction() as cursor:
114
+ cursor.execute(_MIGRATION_LOCK)
115
+ cursor.execute(_VERSION_TABLE)
116
+ cursor.execute("SELECT COALESCE(MAX(version), 0) FROM goatbeat_schema_version")
117
+ row = cursor.fetchone()
118
+ current = int(str(row[0])) if row is not None else 0
119
+ if current > SCHEMA_VERSION:
120
+ raise ConfigurationError(
121
+ f"the database schema is at version {current} but this release only knows "
122
+ f"{SCHEMA_VERSION}: upgrade goatbeat-postgres before running it"
123
+ )
124
+ for version, statements in enumerate(MIGRATIONS[current:], start=current + 1):
125
+ for statement in statements:
126
+ cursor.execute(statement)
127
+ cursor.execute("INSERT INTO goatbeat_schema_version (version) VALUES (%s)", [version])
128
+ return SCHEMA_VERSION
129
+
130
+
131
+ def applied_version(provider: ConnectionProvider) -> int:
132
+ with provider.transaction() as cursor:
133
+ cursor.execute("SELECT to_regclass('goatbeat_schema_version') IS NOT NULL")
134
+ exists = cursor.fetchone()
135
+ if exists is None or not exists[0]:
136
+ return 0
137
+ cursor.execute("SELECT COALESCE(MAX(version), 0) FROM goatbeat_schema_version")
138
+ row = cursor.fetchone()
139
+ return int(str(row[0])) if row is not None else 0
140
+
141
+
142
+ def database_clock(provider: ConnectionProvider) -> Instant:
143
+ """PostgreSQL's clock, which decides lease expiry; compared to local clocks by doctor."""
144
+ with provider.transaction() as cursor:
145
+ cursor.execute("SELECT clock_timestamp()")
146
+ row = cursor.fetchone()
147
+ if row is None or not isinstance(row[0], datetime):
148
+ raise StoreUnavailableError("the database returned no clock reading")
149
+ return Instant.of(row[0])
@@ -0,0 +1,602 @@
1
+ import json
2
+ import logging
3
+ import time
4
+ from collections.abc import Sequence
5
+ from dataclasses import replace
6
+ from datetime import datetime, timedelta
7
+ from typing import Final, LiteralString
8
+
9
+ from goatbeat.core.errors import ConfigurationError, LeaseLostError, StoreUnavailableError
10
+ from goatbeat.core.execution import (
11
+ EARLY_COMPLETION_RETENTION,
12
+ Execution,
13
+ ExecutionStatus,
14
+ Lease,
15
+ LeaseInfo,
16
+ SubmissionReceipt,
17
+ idempotency_key_for,
18
+ )
19
+ from goatbeat.core.ids import (
20
+ FencingToken,
21
+ IdempotencyKey,
22
+ ScheduleId,
23
+ SchedulerId,
24
+ new_schedule_id,
25
+ )
26
+ from goatbeat.core.misfire import OccurrencePlan
27
+ from goatbeat.core.observability import LOGGER_NAME, Event, log_event
28
+ from goatbeat.core.schedule import Schedule, ScheduleDefinition, ScheduleSource
29
+ from goatbeat.core.serialization import (
30
+ DefinitionColumns,
31
+ definition_from_columns,
32
+ definition_to_columns,
33
+ )
34
+ from goatbeat.core.submission import PendingSubmission
35
+ from goatbeat.core.time import Clock, Instant
36
+ from goatbeat.postgres.notify import CHANNEL, Notifier
37
+ from goatbeat.postgres.provider import ConnectionProvider, Cursor
38
+
39
+ _PARTITION: Final = 0
40
+ _MICROSECOND: Final = timedelta(microseconds=1)
41
+ # NULL means "use the database clock"; only tests pass an explicit instant.
42
+ _NOW: Final = "COALESCE(%s::timestamptz, now())"
43
+
44
+ _SCHEDULE_COLUMNS: Final = (
45
+ "id, name, source, task_path, task_args::text, task_kwargs::text, trigger_kind, "
46
+ "interval_us, start_at, cron, timezone, misfire_policy, misfire_grace_us, max_catch_up, "
47
+ "queue, priority, backend, enabled, next_run_at"
48
+ )
49
+ _SCHEDULE_FIELD_COUNT: Final = 19
50
+ _EXECUTION_COLUMNS: Final = (
51
+ "schedule_id, scheduled_at, idempotency_key, status, scheduler_id, fencing_token, "
52
+ "claimed_at, missed_count, retry_count, dispatched_at, completed_at, next_retry_at, "
53
+ "backend_result_id, error"
54
+ )
55
+ _SELECT_SCHEDULE: Final = "SELECT " + _SCHEDULE_COLUMNS + " FROM goatbeat_schedule "
56
+ DUE_SCHEDULES_QUERY: Final = (
57
+ _SELECT_SCHEDULE + "WHERE enabled AND next_run_at <= %s ORDER BY next_run_at LIMIT %s"
58
+ )
59
+ _SELECT_EXECUTION: Final = "SELECT " + _EXECUTION_COLUMNS + " FROM goatbeat_execution "
60
+ _INSERT_SCHEDULE: Final = (
61
+ "INSERT INTO goatbeat_schedule (id, name, source, task_path, task_args, task_kwargs, "
62
+ "trigger_kind, interval_us, start_at, cron, timezone, misfire_policy, misfire_grace_us, "
63
+ "max_catch_up, queue, priority, backend, enabled, next_run_at) VALUES "
64
+ "(%s, %s, %s, %s, %s::jsonb, %s::jsonb, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
65
+ )
66
+ _UPDATE_SCHEDULE: Final = (
67
+ "UPDATE goatbeat_schedule SET name = %s, source = %s, task_path = %s, "
68
+ "task_args = %s::jsonb, task_kwargs = %s::jsonb, trigger_kind = %s, interval_us = %s, "
69
+ "start_at = %s, cron = %s, timezone = %s, misfire_policy = %s, misfire_grace_us = %s, "
70
+ "max_catch_up = %s, queue = %s, priority = %s, backend = %s, enabled = %s, "
71
+ "next_run_at = %s, updated_at = now() WHERE id = %s"
72
+ )
73
+ _INSERT_EXECUTION: Final = (
74
+ "INSERT INTO goatbeat_execution (" + _EXECUTION_COLUMNS + ") VALUES "
75
+ "(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) "
76
+ "ON CONFLICT (idempotency_key) DO NOTHING"
77
+ )
78
+ DEFAULT_LOCK_TIMEOUT: Final = timedelta(seconds=1)
79
+ DEFAULT_FENCE_IDLE_TIMEOUT: Final = timedelta(seconds=10)
80
+ # set_config(..., true) is SET LOCAL with bindable values: it ends with the transaction.
81
+ _SET_LOCK_TIMEOUT: Final = "SELECT set_config('lock_timeout', %s, true)"
82
+ _SET_IDLE_TIMEOUT: Final = "SELECT set_config('idle_in_transaction_session_timeout', %s, true)"
83
+ # Row locks cannot guard names that do not exist yet: concurrent first starts would both insert.
84
+ _LOCK_SCHEDULE_WRITES: Final = "SELECT pg_advisory_xact_lock(hashtext('goatbeat.schedules'))"
85
+ # A completion and the dispatch record of the same task result must not interleave.
86
+ _LOCK_RESULT: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))"
87
+ _NOTIFY: Final = "SELECT pg_notify('" + CHANNEL + "', '')"
88
+
89
+
90
+ class PostgresStore:
91
+ """Implements ScheduleStore and LeaseManager. Lease time is PostgreSQL's ``now()``."""
92
+
93
+ def __init__(
94
+ self,
95
+ provider: ConnectionProvider,
96
+ *,
97
+ notifier: Notifier | None = None,
98
+ lease_clock: Clock | None = None,
99
+ lock_timeout: timedelta = DEFAULT_LOCK_TIMEOUT,
100
+ fence_idle_timeout: timedelta = DEFAULT_FENCE_IDLE_TIMEOUT,
101
+ logger: logging.Logger | None = None,
102
+ ) -> None:
103
+ if lock_timeout <= timedelta(0) or fence_idle_timeout <= timedelta(0):
104
+ raise ConfigurationError("lock_timeout and fence_idle_timeout must be positive")
105
+ self._provider = provider
106
+ self._lock_timeout = _milliseconds(lock_timeout)
107
+ self._fence_idle_timeout = _milliseconds(fence_idle_timeout)
108
+ self._notifier = notifier
109
+ self._lease_clock = lease_clock
110
+ self._logger = logger or logging.getLogger(LOGGER_NAME)
111
+
112
+ def _now(self) -> datetime | None:
113
+ return None if self._lease_clock is None else self._lease_clock.now().datetime
114
+
115
+ # LeaseManager
116
+
117
+ def acquire(self, holder: SchedulerId, ttl: timedelta) -> Lease | None:
118
+ now = self._now()
119
+ with self._provider.transaction() as cursor:
120
+ # A frozen leader holding the fence must not stall every standby indefinitely.
121
+ cursor.execute(_SET_LOCK_TIMEOUT, [self._lock_timeout])
122
+ cursor.execute(
123
+ "UPDATE goatbeat_lease SET "
124
+ "fencing_token = fencing_token + (holder_id IS DISTINCT FROM %s)::int, "
125
+ "holder_id = %s, expires_at = " + _NOW + " + %s "
126
+ "WHERE partition = %s AND (expires_at <= " + _NOW + " OR holder_id = %s) "
127
+ "RETURNING fencing_token",
128
+ [holder.value, holder.value, now, ttl, _PARTITION, now, holder.value],
129
+ )
130
+ row = cursor.fetchone()
131
+ return None if row is None else Lease(holder, FencingToken(_integer(row[0])))
132
+
133
+ def release(self, lease: Lease) -> None:
134
+ with self._provider.transaction() as cursor:
135
+ cursor.execute(_SET_LOCK_TIMEOUT, [self._lock_timeout])
136
+ cursor.execute(
137
+ "UPDATE goatbeat_lease SET expires_at = " + _NOW + " "
138
+ "WHERE partition = %s AND holder_id = %s AND fencing_token = %s",
139
+ [self._now(), _PARTITION, lease.holder.value, lease.token.value],
140
+ )
141
+
142
+ def current(self) -> LeaseInfo | None:
143
+ with self._provider.transaction() as cursor:
144
+ cursor.execute(
145
+ "SELECT holder_id, fencing_token, expires_at FROM goatbeat_lease "
146
+ "WHERE partition = %s AND holder_id IS NOT NULL",
147
+ [_PARTITION],
148
+ )
149
+ row = cursor.fetchone()
150
+ if row is None:
151
+ return None
152
+ return LeaseInfo(
153
+ SchedulerId(_text(row[0])), FencingToken(_integer(row[1])), _instant(row[2])
154
+ )
155
+
156
+ # ScheduleStore
157
+
158
+ def sync_code_schedules(
159
+ self, definitions: Sequence[ScheduleDefinition], now: Instant, *, prune: bool = False
160
+ ) -> None:
161
+ with self._provider.transaction() as cursor:
162
+ cursor.execute(_LOCK_SCHEDULE_WRITES)
163
+ for definition in sorted(definitions, key=lambda item: item.name):
164
+ existing = self._schedule_by_name(cursor, definition.name, lock=True)
165
+ if existing is None:
166
+ self._insert(cursor, definition, ScheduleSource.CODE, now)
167
+ else:
168
+ kept = replace(definition, enabled=existing.definition.enabled)
169
+ self._update(cursor, existing, kept, now)
170
+ if prune:
171
+ cursor.execute(
172
+ "UPDATE goatbeat_schedule SET enabled = false, updated_at = now() "
173
+ "WHERE source = %s AND enabled AND NOT (name = ANY(%s::text[]))",
174
+ [ScheduleSource.CODE.value, [d.name for d in definitions]],
175
+ )
176
+ cursor.execute(_NOTIFY)
177
+
178
+ def save_schedule(self, definition: ScheduleDefinition, now: Instant) -> Schedule:
179
+ with self._provider.transaction() as cursor:
180
+ cursor.execute(_LOCK_SCHEDULE_WRITES)
181
+ existing = self._schedule_by_name(cursor, definition.name, lock=True)
182
+ if existing is not None and existing.source is ScheduleSource.CODE:
183
+ raise ConfigurationError(f"schedule {definition.name!r} is declared in code")
184
+ if existing is None:
185
+ saved = self._insert(cursor, definition, ScheduleSource.STORE, now)
186
+ else:
187
+ saved = self._update(cursor, existing, definition, now)
188
+ cursor.execute(_NOTIFY)
189
+ return saved
190
+
191
+ def get_schedule(self, name: str) -> Schedule | None:
192
+ with self._provider.transaction() as cursor:
193
+ return self._schedule_by_name(cursor, name, lock=False)
194
+
195
+ def list_schedules(self) -> Sequence[Schedule]:
196
+ with self._provider.transaction() as cursor:
197
+ cursor.execute(_SELECT_SCHEDULE + "ORDER BY name")
198
+ return self._decodable(cursor.fetchall())
199
+
200
+ def set_enabled(self, name: str, *, enabled: bool, now: Instant) -> None:
201
+ with self._provider.transaction() as cursor:
202
+ schedule = self._schedule_by_name(cursor, name, lock=True)
203
+ if schedule is None:
204
+ raise ConfigurationError(f"unknown schedule {name!r}")
205
+ if schedule.definition.enabled == enabled:
206
+ return
207
+ next_run_at = schedule.next_run_at
208
+ if enabled:
209
+ next_run_at = schedule.definition.trigger.next_after(now)
210
+ cursor.execute(
211
+ "UPDATE goatbeat_schedule SET enabled = %s, next_run_at = %s, "
212
+ "updated_at = now() WHERE id = %s",
213
+ [enabled, next_run_at.datetime, schedule.id.value],
214
+ )
215
+ cursor.execute(_NOTIFY)
216
+
217
+ def due_schedules(self, now: Instant, limit: int) -> Sequence[Schedule]:
218
+ with self._provider.transaction() as cursor:
219
+ cursor.execute(DUE_SCHEDULES_QUERY, [now.datetime, limit])
220
+ return self._decodable(cursor.fetchall())
221
+
222
+ def claim(
223
+ self, lease: Lease, schedule: Schedule, plan: OccurrencePlan, now: Instant
224
+ ) -> Sequence[Execution]:
225
+ with self._provider.transaction() as cursor:
226
+ return self.claim_in(cursor, lease, schedule, plan, now)
227
+
228
+ def claim_in(
229
+ self, cursor: Cursor, lease: Lease, schedule: Schedule, plan: OccurrencePlan, now: Instant
230
+ ) -> Sequence[Execution]:
231
+ """Claim inside a transaction the caller owns (Django's ``atomic`` for instance)."""
232
+ self._ensure_fence(cursor, lease)
233
+ self._lock_schedule(cursor, schedule)
234
+ created: list[Execution] = []
235
+ for claim in plan.claims:
236
+ execution = Execution(
237
+ schedule_id=schedule.id,
238
+ scheduled_at=claim.scheduled_at,
239
+ idempotency_key=idempotency_key_for(schedule.id, claim.scheduled_at),
240
+ status=claim.status,
241
+ scheduler_id=lease.holder,
242
+ fencing_token=lease.token,
243
+ claimed_at=now,
244
+ missed_count=claim.missed_count,
245
+ )
246
+ cursor.execute(_INSERT_EXECUTION, _execution_values(execution))
247
+ if cursor.rowcount == 1:
248
+ created.append(execution)
249
+ cursor.execute(
250
+ "UPDATE goatbeat_schedule SET next_run_at = %s WHERE id = %s AND next_run_at = %s",
251
+ [plan.next_run_at.datetime, schedule.id.value, schedule.next_run_at.datetime],
252
+ )
253
+ return created
254
+
255
+ def request_run(
256
+ self, schedule: Schedule, key: IdempotencyKey, requested_by: SchedulerId, now: Instant
257
+ ) -> Execution:
258
+ execution = Execution(
259
+ schedule_id=schedule.id,
260
+ scheduled_at=now,
261
+ idempotency_key=key,
262
+ status=ExecutionStatus.REQUESTED,
263
+ scheduler_id=requested_by,
264
+ fencing_token=FencingToken(0),
265
+ claimed_at=now,
266
+ )
267
+ with self._provider.transaction() as cursor:
268
+ self._lock_schedule(cursor, schedule)
269
+ cursor.execute(_INSERT_EXECUTION, _execution_values(execution))
270
+ cursor.execute(_SELECT_EXECUTION + "WHERE idempotency_key = %s", [key.value])
271
+ row = cursor.fetchone()
272
+ cursor.execute(_NOTIFY)
273
+ if row is None:
274
+ raise StoreUnavailableError(f"requested run {key.value} was not recorded")
275
+ return _execution(row)
276
+
277
+ def pending_submissions(
278
+ self, now: Instant, stale_before: Instant, limit: int
279
+ ) -> Sequence[PendingSubmission]:
280
+ with self._provider.transaction() as cursor:
281
+ cursor.execute(
282
+ "SELECT e.schedule_id, e.scheduled_at, e.idempotency_key, e.status, "
283
+ "e.scheduler_id, e.fencing_token, e.claimed_at, e.missed_count, e.retry_count, "
284
+ "e.dispatched_at, e.completed_at, e.next_retry_at, e.backend_result_id, e.error, "
285
+ "s.id, s.name, s.source, s.task_path, s.task_args::text, s.task_kwargs::text, "
286
+ "s.trigger_kind, s.interval_us, s.start_at, s.cron, s.timezone, "
287
+ "s.misfire_policy, s.misfire_grace_us, s.max_catch_up, s.queue, s.priority, "
288
+ "s.backend, s.enabled, s.next_run_at "
289
+ "FROM goatbeat_execution e JOIN goatbeat_schedule s ON s.id = e.schedule_id "
290
+ "WHERE e.status = %s OR (e.status = %s AND e.claimed_at <= %s) "
291
+ "OR (e.status = %s AND e.next_retry_at <= %s) "
292
+ "ORDER BY e.scheduled_at LIMIT %s",
293
+ [
294
+ ExecutionStatus.REQUESTED.value,
295
+ ExecutionStatus.CLAIMED.value,
296
+ stale_before.datetime,
297
+ ExecutionStatus.SUBMIT_FAILED.value,
298
+ now.datetime,
299
+ limit,
300
+ ],
301
+ )
302
+ rows = cursor.fetchall()
303
+ pending: list[PendingSubmission] = []
304
+ split = len(_EXECUTION_COLUMNS.split(","))
305
+ for row in rows:
306
+ schedule = self._decode(row[split:])
307
+ if schedule is not None:
308
+ pending.append(PendingSubmission(_execution(row[:split]), schedule))
309
+ return pending
310
+
311
+ def mark_dispatched(self, key: IdempotencyKey, at: Instant, receipt: SubmissionReceipt) -> None:
312
+ status, completed_at = receipt.status, None if receipt.outcome is None else at.datetime
313
+ with self._provider.transaction() as cursor:
314
+ if receipt.result_id is not None:
315
+ cursor.execute(_LOCK_RESULT, [receipt.result_id])
316
+ cursor.execute(
317
+ "DELETE FROM goatbeat_early_completion WHERE result_id = %s "
318
+ "RETURNING status, completed_at",
319
+ [receipt.result_id],
320
+ )
321
+ early = cursor.fetchone()
322
+ if early is not None:
323
+ status, completed_at = (
324
+ ExecutionStatus(_text(early[0])),
325
+ _instant(early[1]).datetime,
326
+ )
327
+ cursor.execute(
328
+ "UPDATE goatbeat_execution SET status = %s, dispatched_at = %s, "
329
+ "completed_at = %s, backend_result_id = %s, next_retry_at = NULL, error = NULL "
330
+ "WHERE idempotency_key = %s AND status NOT IN ('completed', 'task_failed')",
331
+ [status.value, at.datetime, completed_at, receipt.result_id, key.value],
332
+ )
333
+
334
+ def mark_submit_failed(self, key: IdempotencyKey, error: str, retry_at: Instant | None) -> None:
335
+ status = ExecutionStatus.FAILED if retry_at is None else ExecutionStatus.SUBMIT_FAILED
336
+ with self._provider.transaction() as cursor:
337
+ cursor.execute(
338
+ "UPDATE goatbeat_execution SET status = %s, retry_count = retry_count + 1, "
339
+ "next_retry_at = %s, error = %s WHERE idempotency_key = %s",
340
+ [status.value, None if retry_at is None else retry_at.datetime, error, key.value],
341
+ )
342
+
343
+ def record_completion(
344
+ self, result_id: str, task_path: str, *, succeeded: bool, at: Instant
345
+ ) -> bool:
346
+ status = ExecutionStatus.COMPLETED if succeeded else ExecutionStatus.TASK_FAILED
347
+ with self._provider.transaction() as cursor:
348
+ cursor.execute(_LOCK_RESULT, [result_id])
349
+ cursor.execute(
350
+ "UPDATE goatbeat_execution SET status = %s, completed_at = %s "
351
+ "WHERE backend_result_id = %s AND status = %s",
352
+ [status.value, at.datetime, result_id, ExecutionStatus.DISPATCHED.value],
353
+ )
354
+ if cursor.rowcount > 0:
355
+ return True
356
+ cursor.execute(
357
+ "DELETE FROM goatbeat_early_completion WHERE completed_at < %s",
358
+ [(at - EARLY_COMPLETION_RETENTION).datetime],
359
+ )
360
+ cursor.execute(
361
+ "INSERT INTO goatbeat_early_completion (result_id, status, completed_at) "
362
+ "SELECT %s, %s, %s WHERE EXISTS "
363
+ "(SELECT 1 FROM goatbeat_schedule WHERE task_path = %s) ON CONFLICT DO NOTHING",
364
+ [result_id, status.value, at.datetime, task_path],
365
+ )
366
+ return False
367
+
368
+ def next_wakeup(self) -> Instant | None:
369
+ with self._provider.transaction() as cursor:
370
+ cursor.execute(
371
+ "SELECT LEAST("
372
+ "(SELECT MIN(next_run_at) FROM goatbeat_schedule WHERE enabled), "
373
+ "(SELECT MIN(next_retry_at) FROM goatbeat_execution WHERE status = %s))",
374
+ [ExecutionStatus.SUBMIT_FAILED.value],
375
+ )
376
+ row = cursor.fetchone()
377
+ return None if row is None else _optional_instant(row[0])
378
+
379
+ def recent_executions(self, schedule_id: ScheduleId | None, limit: int) -> Sequence[Execution]:
380
+ with self._provider.transaction() as cursor:
381
+ if schedule_id is None:
382
+ cursor.execute(_SELECT_EXECUTION + "ORDER BY scheduled_at DESC LIMIT %s", [limit])
383
+ else:
384
+ cursor.execute(
385
+ _SELECT_EXECUTION + "WHERE schedule_id = %s ORDER BY scheduled_at DESC "
386
+ "LIMIT %s",
387
+ [schedule_id.value, limit],
388
+ )
389
+ return [_execution(row) for row in cursor.fetchall()]
390
+
391
+ def wait_for_change(self, timeout: timedelta) -> bool:
392
+ if self._notifier is None:
393
+ time.sleep(timeout.total_seconds())
394
+ return False
395
+ return self._notifier.wait(timeout)
396
+
397
+ # Internals
398
+
399
+ def _ensure_fence(self, cursor: Cursor, lease: Lease) -> None:
400
+ """FOR SHARE makes a concurrent takeover wait until this claim transaction ends.
401
+
402
+ The idle timeout lets PostgreSQL end that transaction if the leader freezes.
403
+ """
404
+ cursor.execute(_SET_IDLE_TIMEOUT, [self._fence_idle_timeout])
405
+ cursor.execute(
406
+ "SELECT 1 FROM goatbeat_lease WHERE partition = %s AND holder_id = %s "
407
+ "AND fencing_token = %s AND expires_at > " + _NOW + " FOR SHARE",
408
+ [_PARTITION, lease.holder.value, lease.token.value, self._now()],
409
+ )
410
+ if cursor.fetchone() is None:
411
+ raise LeaseLostError(f"{lease.holder.value} no longer holds token {lease.token.value}")
412
+
413
+ def _lock_schedule(self, cursor: Cursor, schedule: Schedule) -> None:
414
+ """KEY SHARE keeps the schedule from being deleted before the executions commit."""
415
+ cursor.execute(
416
+ "SELECT 1 FROM goatbeat_schedule WHERE id = %s FOR KEY SHARE", [schedule.id.value]
417
+ )
418
+ if cursor.fetchone() is None:
419
+ raise ConfigurationError(f"schedule {schedule.name!r} no longer exists")
420
+
421
+ def _schedule_by_name(self, cursor: Cursor, name: str, *, lock: bool) -> Schedule | None:
422
+ query: LiteralString = _SELECT_SCHEDULE + "WHERE name = %s"
423
+ cursor.execute(query + " FOR UPDATE" if lock else query, [name])
424
+ row = cursor.fetchone()
425
+ return None if row is None else _schedule(row)
426
+
427
+ def _decode(self, row: Sequence[object]) -> Schedule | None:
428
+ try:
429
+ return _schedule(row)
430
+ except (ConfigurationError, ValueError) as error:
431
+ log_event(
432
+ self._logger,
433
+ Event.SCHEDULE_INVALID,
434
+ logging.ERROR,
435
+ schedule=str(row[1]),
436
+ error=str(error),
437
+ )
438
+ return None
439
+
440
+ def _decodable(self, rows: Sequence[Sequence[object]]) -> list[Schedule]:
441
+ """Skip rows this version cannot decode, so one bad row never blocks the others."""
442
+ return [schedule for row in rows if (schedule := self._decode(row)) is not None]
443
+
444
+ def _insert(
445
+ self, cursor: Cursor, definition: ScheduleDefinition, source: ScheduleSource, now: Instant
446
+ ) -> Schedule:
447
+ schedule = Schedule(
448
+ id=new_schedule_id(),
449
+ definition=definition,
450
+ source=source,
451
+ next_run_at=definition.trigger.next_after(now),
452
+ )
453
+ cursor.execute(_INSERT_SCHEDULE, _schedule_values(schedule))
454
+ return schedule
455
+
456
+ def _update(
457
+ self, cursor: Cursor, existing: Schedule, definition: ScheduleDefinition, now: Instant
458
+ ) -> Schedule:
459
+ next_run_at = existing.next_run_at
460
+ if definition.trigger != existing.definition.trigger:
461
+ next_run_at = definition.trigger.next_after(now)
462
+ updated = Schedule(existing.id, definition, existing.source, next_run_at)
463
+ values = _schedule_values(updated)
464
+ cursor.execute(_UPDATE_SCHEDULE, [*values[1:], values[0]])
465
+ return updated
466
+
467
+
468
+ def _schedule_values(schedule: Schedule) -> list[object]:
469
+ columns = definition_to_columns(schedule.definition)
470
+ interval = columns["interval"]
471
+ start_at = columns["start_at"]
472
+ return [
473
+ schedule.id.value,
474
+ columns["name"],
475
+ schedule.source.value,
476
+ columns["task_path"],
477
+ columns["task_args"],
478
+ columns["task_kwargs"],
479
+ columns["trigger_kind"],
480
+ None if interval is None else interval // _MICROSECOND,
481
+ None if start_at is None else start_at.datetime,
482
+ columns["cron"],
483
+ columns["timezone"],
484
+ columns["misfire_policy"],
485
+ columns["misfire_grace"] // _MICROSECOND,
486
+ columns["max_catch_up"],
487
+ columns["queue"],
488
+ columns["priority"],
489
+ columns["backend"],
490
+ columns["enabled"],
491
+ schedule.next_run_at.datetime,
492
+ ]
493
+
494
+
495
+ def _schedule(row: Sequence[object]) -> Schedule:
496
+ if len(row) != _SCHEDULE_FIELD_COUNT:
497
+ raise ValueError(f"expected {_SCHEDULE_FIELD_COUNT} schedule columns, got {len(row)}")
498
+ interval_us = _optional_integer(row[7])
499
+ columns = DefinitionColumns(
500
+ name=_text(row[1]),
501
+ task_path=_text(row[3]),
502
+ task_args=_json_text(row[4]),
503
+ task_kwargs=_json_text(row[5]),
504
+ trigger_kind=_text(row[6]),
505
+ interval=None if interval_us is None else interval_us * _MICROSECOND,
506
+ start_at=_optional_instant(row[8]),
507
+ cron=_optional_text(row[9]),
508
+ timezone=_optional_text(row[10]),
509
+ misfire_policy=_text(row[11]),
510
+ misfire_grace=_integer(row[12]) * _MICROSECOND,
511
+ max_catch_up=_integer(row[13]),
512
+ queue=_optional_text(row[14]),
513
+ priority=_optional_integer(row[15]),
514
+ backend=_optional_text(row[16]),
515
+ enabled=bool(row[17]),
516
+ )
517
+ return Schedule(
518
+ id=ScheduleId(_text(row[0])),
519
+ definition=definition_from_columns(columns),
520
+ source=ScheduleSource(_text(row[2])),
521
+ next_run_at=_instant(row[18]),
522
+ )
523
+
524
+
525
+ def _execution_values(execution: Execution) -> list[object]:
526
+ return [
527
+ execution.schedule_id.value,
528
+ execution.scheduled_at.datetime,
529
+ execution.idempotency_key.value,
530
+ execution.status.value,
531
+ execution.scheduler_id.value,
532
+ execution.fencing_token.value,
533
+ execution.claimed_at.datetime,
534
+ execution.missed_count,
535
+ execution.retry_count,
536
+ _optional_datetime(execution.dispatched_at),
537
+ _optional_datetime(execution.completed_at),
538
+ _optional_datetime(execution.next_retry_at),
539
+ execution.backend_result_id,
540
+ execution.error,
541
+ ]
542
+
543
+
544
+ def _execution(row: Sequence[object]) -> Execution:
545
+ return Execution(
546
+ schedule_id=ScheduleId(_text(row[0])),
547
+ scheduled_at=_instant(row[1]),
548
+ idempotency_key=IdempotencyKey(_text(row[2])),
549
+ status=ExecutionStatus(_text(row[3])),
550
+ scheduler_id=SchedulerId(_text(row[4])),
551
+ fencing_token=FencingToken(_integer(row[5])),
552
+ claimed_at=_instant(row[6]),
553
+ missed_count=_integer(row[7]),
554
+ retry_count=_integer(row[8]),
555
+ dispatched_at=_optional_instant(row[9]),
556
+ completed_at=_optional_instant(row[10]),
557
+ next_retry_at=_optional_instant(row[11]),
558
+ backend_result_id=_optional_text(row[12]),
559
+ error=_optional_text(row[13]),
560
+ )
561
+
562
+
563
+ def _text(value: object) -> str:
564
+ if not isinstance(value, str):
565
+ raise ValueError(f"expected text, got {value!r}")
566
+ return value
567
+
568
+
569
+ def _optional_text(value: object) -> str | None:
570
+ return None if value is None else _text(value)
571
+
572
+
573
+ def _json_text(value: object) -> str:
574
+ return value if isinstance(value, str) else json.dumps(value)
575
+
576
+
577
+ def _integer(value: object) -> int:
578
+ if isinstance(value, bool) or not isinstance(value, int):
579
+ raise ValueError(f"expected an integer, got {value!r}")
580
+ return value
581
+
582
+
583
+ def _optional_integer(value: object) -> int | None:
584
+ return None if value is None else _integer(value)
585
+
586
+
587
+ def _instant(value: object) -> Instant:
588
+ if not isinstance(value, datetime):
589
+ raise ValueError(f"expected a timestamp, got {value!r}")
590
+ return Instant.of(value)
591
+
592
+
593
+ def _optional_instant(value: object) -> Instant | None:
594
+ return None if value is None else _instant(value)
595
+
596
+
597
+ def _optional_datetime(value: Instant | None) -> datetime | None:
598
+ return None if value is None else value.datetime
599
+
600
+
601
+ def _milliseconds(value: timedelta) -> str:
602
+ return f"{max(1, round(value.total_seconds() * 1000))}ms"
@@ -0,0 +1,27 @@
1
+ import os
2
+ import secrets
3
+ from collections.abc import Generator
4
+ from contextlib import contextmanager
5
+ from typing import Final
6
+
7
+ import psycopg
8
+ from psycopg import sql
9
+ from psycopg.conninfo import make_conninfo
10
+
11
+ TEST_DATABASE_URL_VARIABLE: Final = "GOATBEAT_TEST_DATABASE_URL"
12
+
13
+
14
+ def database_url_from_env() -> str | None:
15
+ return os.environ.get(TEST_DATABASE_URL_VARIABLE) or None
16
+
17
+
18
+ @contextmanager
19
+ def isolated_schema(url: str) -> Generator[str]:
20
+ schema = f"goatbeat_test_{secrets.token_hex(6)}"
21
+ with psycopg.connect(url, autocommit=True) as connection:
22
+ connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
23
+ try:
24
+ yield make_conninfo(url, options=f"-c search_path={schema}")
25
+ finally:
26
+ with psycopg.connect(url, autocommit=True) as connection:
27
+ connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))