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.
- goatbeat_postgres-0.1.0/LICENSE +21 -0
- goatbeat_postgres-0.1.0/PKG-INFO +29 -0
- goatbeat_postgres-0.1.0/README.md +4 -0
- goatbeat_postgres-0.1.0/pyproject.toml +49 -0
- goatbeat_postgres-0.1.0/pyproject.toml.orig +39 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/__init__.py +0 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/notify.py +74 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/provider.py +112 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/py.typed +0 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/schema.py +149 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/store.py +602 -0
- goatbeat_postgres-0.1.0/src/goatbeat/postgres/testing.py +27 -0
|
@@ -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,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 }
|
|
File without changes
|
|
@@ -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)))
|