iddqueue 0.13.0rc1__py3-none-any.whl
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.
- iddqueue/__init__.py +15 -0
- iddqueue/broker.py +667 -0
- iddqueue/cancellation.py +51 -0
- iddqueue/cancellation.sql +6 -0
- iddqueue/cli.py +474 -0
- iddqueue/control.py +87 -0
- iddqueue/control.sql +6 -0
- iddqueue/coordination.sql +11 -0
- iddqueue/deduplication.sql +12 -0
- iddqueue/failures.py +18 -0
- iddqueue/history.py +70 -0
- iddqueue/history.sql +13 -0
- iddqueue/metrics.py +68 -0
- iddqueue/rate_limits.py +122 -0
- iddqueue/results.py +118 -0
- iddqueue/scheduler.py +75 -0
- iddqueue/scheduler.sql +9 -0
- iddqueue/schema.py +55 -0
- iddqueue/schema.sql +31 -0
- iddqueue/utils.py +206 -0
- iddqueue-0.13.0rc1.dist-info/METADATA +730 -0
- iddqueue-0.13.0rc1.dist-info/RECORD +25 -0
- iddqueue-0.13.0rc1.dist-info/WHEEL +4 -0
- iddqueue-0.13.0rc1.dist-info/entry_points.txt +3 -0
- iddqueue-0.13.0rc1.dist-info/licenses/LICENSE +19 -0
iddqueue/history.py
ADDED
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""Optional per-execution diagnostics, independent of queue/result retention."""
|
|
2
|
+
from uuid import uuid4
|
|
3
|
+
|
|
4
|
+
from dramatiq.middleware import Middleware
|
|
5
|
+
from psycopg import sql
|
|
6
|
+
|
|
7
|
+
from .utils import transaction
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def attempt_table(schema, prefix):
|
|
11
|
+
return sql.Identifier(schema, prefix + "attempts")
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class AttemptHistory(Middleware):
|
|
15
|
+
def before_process_message(self, broker, message):
|
|
16
|
+
attempt_id = uuid4()
|
|
17
|
+
with transaction(broker.pool) as curs:
|
|
18
|
+
curs.execute(sql.SQL(
|
|
19
|
+
"INSERT INTO {} (attempt_id, message_id, queue_name, actor_name) VALUES (%s, %s, %s, %s)"
|
|
20
|
+
).format(attempt_table(broker.queries.schema, broker.queries.prefix)),
|
|
21
|
+
(attempt_id, message.message_id, message.queue_name, message.actor_name))
|
|
22
|
+
message._pg_attempt_id = attempt_id
|
|
23
|
+
|
|
24
|
+
def after_process_message(self, broker, message, *, result=None, exception=None):
|
|
25
|
+
attempt_id = getattr(message, "_pg_attempt_id", None)
|
|
26
|
+
if attempt_id is None:
|
|
27
|
+
return
|
|
28
|
+
with transaction(broker.pool) as curs:
|
|
29
|
+
curs.execute(sql.SQL(
|
|
30
|
+
"UPDATE {} SET finished_at = clock_timestamp(), outcome = %s, error_type = %s, "
|
|
31
|
+
"error_text = %s WHERE attempt_id = %s"
|
|
32
|
+
).format(attempt_table(broker.queries.schema, broker.queries.prefix)),
|
|
33
|
+
("successful" if exception is None else "failed",
|
|
34
|
+
type(exception).__name__ if exception is not None else None,
|
|
35
|
+
str(exception)[:2000] if exception is not None else None, attempt_id))
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def list_attempts(pool, message_id, *, schema="dramatiq", prefix="", limit=50, after=None):
|
|
39
|
+
if not 1 <= limit <= 1000:
|
|
40
|
+
raise ValueError("limit must be between 1 and 1000")
|
|
41
|
+
with transaction(pool) as curs:
|
|
42
|
+
curs.execute(sql.SQL(
|
|
43
|
+
"SELECT attempt_id, message_id, queue_name, actor_name, started_at, finished_at, "
|
|
44
|
+
"COALESCE(outcome, 'incomplete'), error_type, error_text FROM {} "
|
|
45
|
+
"WHERE message_id = %s AND (%s::uuid IS NULL OR attempt_id > %s) "
|
|
46
|
+
"ORDER BY attempt_id LIMIT %s"
|
|
47
|
+
).format(attempt_table(schema, prefix)), (message_id, after, after, limit + 1))
|
|
48
|
+
rows = curs.fetchall()
|
|
49
|
+
names = ("attempt_id", "message_id", "queue", "actor", "started_at", "finished_at",
|
|
50
|
+
"outcome", "error_type", "error_text")
|
|
51
|
+
items = []
|
|
52
|
+
for row in rows[:limit]:
|
|
53
|
+
item = dict(zip(names, row))
|
|
54
|
+
item["duration_ms"] = (row[5] - row[4]).total_seconds() * 1000 if row[5] else None
|
|
55
|
+
for key in ("attempt_id", "message_id"):
|
|
56
|
+
item[key] = str(item[key])
|
|
57
|
+
for key in ("started_at", "finished_at"):
|
|
58
|
+
item[key] = item[key].isoformat() if item[key] else None
|
|
59
|
+
items.append(item)
|
|
60
|
+
return dict(items=items, next_after=items[-1]["attempt_id"] if len(rows) > limit else None)
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def purge_attempts(pool, maxage="30 days", *, schema="dramatiq", prefix=""):
|
|
64
|
+
with transaction(pool) as curs:
|
|
65
|
+
curs.execute("SELECT %s::interval > interval '0'", (maxage,))
|
|
66
|
+
if not curs.fetchone()[0]:
|
|
67
|
+
raise ValueError("maxage must be positive")
|
|
68
|
+
curs.execute(sql.SQL("DELETE FROM {} WHERE started_at < clock_timestamp() - %s::interval")
|
|
69
|
+
.format(attempt_table(schema, prefix)), (maxage,))
|
|
70
|
+
return curs.rowcount
|
iddqueue/history.sql
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
CREATE TABLE IF NOT EXISTS :"schema".:"attempts" (
|
|
2
|
+
attempt_id UUID PRIMARY KEY,
|
|
3
|
+
message_id UUID NOT NULL,
|
|
4
|
+
queue_name TEXT NOT NULL,
|
|
5
|
+
actor_name TEXT NOT NULL,
|
|
6
|
+
started_at TIMESTAMPTZ NOT NULL DEFAULT clock_timestamp(),
|
|
7
|
+
finished_at TIMESTAMPTZ,
|
|
8
|
+
outcome TEXT CHECK (outcome IN ('successful', 'failed')),
|
|
9
|
+
error_type TEXT,
|
|
10
|
+
error_text TEXT CHECK (length(error_text) <= 2000)
|
|
11
|
+
);
|
|
12
|
+
CREATE INDEX IF NOT EXISTS :"attempts_message" ON :"schema".:"attempts" (message_id, attempt_id);
|
|
13
|
+
CREATE INDEX IF NOT EXISTS :"attempts_age" ON :"schema".:"attempts" (started_at);
|
iddqueue/metrics.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Queue snapshots; Prometheus is imported only when collecting metrics."""
|
|
2
|
+
|
|
3
|
+
from psycopg import sql
|
|
4
|
+
|
|
5
|
+
from .utils import transaction
|
|
6
|
+
|
|
7
|
+
STATES = ("queued", "consumed", "done", "rejected", "cancelled")
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def statistics_query(schema="dramatiq", prefix=""):
|
|
11
|
+
return sql.SQL("""
|
|
12
|
+
SELECT queue_name,
|
|
13
|
+
count(*) FILTER (WHERE state = 'queued'),
|
|
14
|
+
count(*) FILTER (WHERE state = 'consumed'),
|
|
15
|
+
count(*) FILTER (WHERE state = 'done'),
|
|
16
|
+
count(*) FILTER (WHERE state = 'rejected'),
|
|
17
|
+
count(*) FILTER (WHERE state = 'cancelled'),
|
|
18
|
+
count(*) FILTER (WHERE state = 'queued' AND ready_at <= now()),
|
|
19
|
+
count(*) FILTER (WHERE state IN ('queued', 'consumed') AND ready_at > now()),
|
|
20
|
+
coalesce(greatest(0, extract(epoch FROM now() - min(ready_at)
|
|
21
|
+
FILTER (WHERE state = 'queued' AND ready_at <= now()))), 0)
|
|
22
|
+
FROM (
|
|
23
|
+
SELECT *, greatest(mtime,
|
|
24
|
+
to_timestamp(coalesce((message->'options'->>'eta')::double precision / 1000, 0))) AS ready_at
|
|
25
|
+
FROM {} WHERE (%s::text IS NULL OR queue_name = %s)
|
|
26
|
+
) AS messages GROUP BY queue_name ORDER BY queue_name
|
|
27
|
+
""").format(sql.Identifier(schema, prefix + "queue"))
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def queue_statistics(pool, *, schema="dramatiq", prefix="", queue=None):
|
|
31
|
+
with transaction(pool) as cursor:
|
|
32
|
+
cursor.execute(statistics_query(schema, prefix), (queue, queue))
|
|
33
|
+
rows = cursor.fetchall()
|
|
34
|
+
snapshots = [dict(queue=row[0], counts=dict(zip(STATES, row[1:6])),
|
|
35
|
+
ready=row[6], scheduled=row[7], oldest_ready_seconds=float(row[8]))
|
|
36
|
+
for row in rows]
|
|
37
|
+
if not snapshots and queue is not None:
|
|
38
|
+
snapshots.append(dict(queue=queue, counts=dict.fromkeys(STATES, 0),
|
|
39
|
+
ready=0, scheduled=0, oldest_ready_seconds=0.0))
|
|
40
|
+
return snapshots
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class PostgresQueueCollector:
|
|
44
|
+
"""Register in an exporter process; the caller owns the pool."""
|
|
45
|
+
|
|
46
|
+
def __init__(self, pool, *, schema="dramatiq", prefix="", queue=None):
|
|
47
|
+
self.pool = pool
|
|
48
|
+
self.options = dict(schema=schema, prefix=prefix, queue=queue)
|
|
49
|
+
|
|
50
|
+
def describe(self):
|
|
51
|
+
# Avoid database access during registry registration.
|
|
52
|
+
return iter(())
|
|
53
|
+
|
|
54
|
+
def collect(self):
|
|
55
|
+
from prometheus_client.core import GaugeMetricFamily
|
|
56
|
+
|
|
57
|
+
counts = GaugeMetricFamily("iddqueue_queue_messages", "Stored messages by state", labels=["queue", "state"])
|
|
58
|
+
ready = GaugeMetricFamily("iddqueue_queue_ready", "Ready queued messages", labels=["queue"])
|
|
59
|
+
scheduled = GaugeMetricFamily("iddqueue_queue_scheduled", "Future delayed messages, including prefetched", labels=["queue"])
|
|
60
|
+
age = GaugeMetricFamily("iddqueue_queue_oldest_ready_seconds", "Age of oldest ready queued message", labels=["queue"])
|
|
61
|
+
for snapshot in queue_statistics(self.pool, **self.options):
|
|
62
|
+
queue = snapshot["queue"]
|
|
63
|
+
for state, count in snapshot["counts"].items():
|
|
64
|
+
counts.add_metric([queue, state], count)
|
|
65
|
+
ready.add_metric([queue], snapshot["ready"])
|
|
66
|
+
scheduled.add_metric([queue], snapshot["scheduled"])
|
|
67
|
+
age.add_metric([queue], snapshot["oldest_ready_seconds"])
|
|
68
|
+
yield from (counts, ready, scheduled, age)
|
iddqueue/rate_limits.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
"""PostgreSQL storage for Dramatiq's standard limiters and barriers."""
|
|
2
|
+
|
|
3
|
+
import hashlib
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
from dramatiq.rate_limits.backend import RateLimiterBackend
|
|
7
|
+
from psycopg import sql
|
|
8
|
+
|
|
9
|
+
from .utils import make_pool, storage_namespace, transaction, wait_for_notifies
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class PostgresRateLimiterBackend(RateLimiterBackend):
|
|
13
|
+
def __init__(self, *, url=None, pool=None, schema="dramatiq", prefix=""):
|
|
14
|
+
if pool is not None and url:
|
|
15
|
+
raise ValueError("You can't set both pool and URL!")
|
|
16
|
+
self._owns_pool = pool is None
|
|
17
|
+
self.pool = make_pool(url or "") if pool is None else pool
|
|
18
|
+
self.table = sql.Identifier(schema, prefix + "coordination")
|
|
19
|
+
self.namespace = storage_namespace(schema, prefix)
|
|
20
|
+
|
|
21
|
+
def close(self):
|
|
22
|
+
if self._owns_pool:
|
|
23
|
+
self.pool.close()
|
|
24
|
+
|
|
25
|
+
def _digest(self, key):
|
|
26
|
+
return hashlib.sha256((self.namespace + key).encode()).digest()
|
|
27
|
+
|
|
28
|
+
def _lock(self, curs, keys):
|
|
29
|
+
# Every counter mutation uses the same locks, including missing rows.
|
|
30
|
+
locks = {int.from_bytes(self._digest(key)[:8], "big", signed=True) for key in keys}
|
|
31
|
+
for lock in sorted(locks):
|
|
32
|
+
curs.execute("SELECT pg_advisory_xact_lock(%s)", (lock,))
|
|
33
|
+
|
|
34
|
+
def _values(self, curs, keys):
|
|
35
|
+
curs.execute(sql.SQL(
|
|
36
|
+
"SELECT key, value FROM {} WHERE key = ANY(%s) "
|
|
37
|
+
"AND expires_at > clock_timestamp()"
|
|
38
|
+
).format(self.table), (list(keys),))
|
|
39
|
+
return dict(curs.fetchall())
|
|
40
|
+
|
|
41
|
+
def _put(self, curs, key, value, ttl):
|
|
42
|
+
curs.execute(sql.SQL(
|
|
43
|
+
"INSERT INTO {} AS target (key, value, expires_at) "
|
|
44
|
+
"VALUES (%s, %s, clock_timestamp() + %s * interval '1 millisecond') "
|
|
45
|
+
"ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, "
|
|
46
|
+
"expires_at = EXCLUDED.expires_at"
|
|
47
|
+
).format(self.table), (key, value, ttl))
|
|
48
|
+
|
|
49
|
+
def add(self, key, value, ttl):
|
|
50
|
+
with transaction(self.pool) as curs:
|
|
51
|
+
self._lock(curs, [key])
|
|
52
|
+
if key in self._values(curs, [key]):
|
|
53
|
+
return False
|
|
54
|
+
self._put(curs, key, value, ttl)
|
|
55
|
+
return True
|
|
56
|
+
|
|
57
|
+
def _change(self, key, amount, bound, ttl, *, decrement=False):
|
|
58
|
+
with transaction(self.pool) as curs:
|
|
59
|
+
self._lock(curs, [key])
|
|
60
|
+
value = self._values(curs, [key]).get(key, 0) + amount
|
|
61
|
+
if (value < bound) if decrement else (value > bound):
|
|
62
|
+
return False
|
|
63
|
+
self._put(curs, key, value, ttl)
|
|
64
|
+
return True
|
|
65
|
+
|
|
66
|
+
def incr(self, key, amount, maximum, ttl):
|
|
67
|
+
return self._change(key, amount, maximum, ttl)
|
|
68
|
+
|
|
69
|
+
def decr(self, key, amount, minimum, ttl):
|
|
70
|
+
return self._change(key, -amount, minimum, ttl, decrement=True)
|
|
71
|
+
|
|
72
|
+
def incr_and_sum(self, key, keys, amount, maximum, ttl):
|
|
73
|
+
while True:
|
|
74
|
+
with transaction(self.pool) as curs:
|
|
75
|
+
window_keys = list(keys())
|
|
76
|
+
self._lock(curs, [key, *window_keys])
|
|
77
|
+
# The window may advance while this transaction waits for locks.
|
|
78
|
+
# Restart to lock its new keys in the same global order.
|
|
79
|
+
if list(keys()) != window_keys:
|
|
80
|
+
continue
|
|
81
|
+
values = self._values(curs, [key, *window_keys])
|
|
82
|
+
value = values.get(key, 0) + amount
|
|
83
|
+
if value > maximum or amount + sum(values.get(k, 0) for k in window_keys) > maximum:
|
|
84
|
+
return False
|
|
85
|
+
self._put(curs, key, value, ttl)
|
|
86
|
+
return True
|
|
87
|
+
|
|
88
|
+
def _channel(self, key):
|
|
89
|
+
return "dpg.c." + self._digest(key).hex()[:48]
|
|
90
|
+
|
|
91
|
+
def wait(self, key, timeout):
|
|
92
|
+
deadline = None if timeout is None else time.monotonic() + timeout / 1000
|
|
93
|
+
with transaction(self.pool, listen=self._channel(key)) as curs:
|
|
94
|
+
while True:
|
|
95
|
+
curs.execute(sql.SQL(
|
|
96
|
+
"SELECT 1 FROM {} WHERE key = %s "
|
|
97
|
+
"AND event_expires_at > clock_timestamp()"
|
|
98
|
+
).format(self.table), (key,))
|
|
99
|
+
if curs.fetchone() is not None:
|
|
100
|
+
return True
|
|
101
|
+
remaining = None if deadline is None else deadline - time.monotonic()
|
|
102
|
+
if remaining is not None and remaining <= 0:
|
|
103
|
+
return False
|
|
104
|
+
wait_for_notifies(curs.connection, timeout=remaining)
|
|
105
|
+
|
|
106
|
+
def wait_notify(self, key, ttl):
|
|
107
|
+
with transaction(self.pool) as curs:
|
|
108
|
+
curs.execute(sql.SQL(
|
|
109
|
+
"INSERT INTO {} (key, event_expires_at) "
|
|
110
|
+
"VALUES (%s, clock_timestamp() + %s * interval '1 millisecond') "
|
|
111
|
+
"ON CONFLICT (key) DO UPDATE SET event_expires_at = EXCLUDED.event_expires_at"
|
|
112
|
+
).format(self.table), (key, ttl))
|
|
113
|
+
curs.execute("SELECT pg_notify(%s, '')", (self._channel(key),))
|
|
114
|
+
|
|
115
|
+
def purge(self):
|
|
116
|
+
"""Remove expired counters and events; schedule periodically."""
|
|
117
|
+
with transaction(self.pool) as curs:
|
|
118
|
+
curs.execute(sql.SQL(
|
|
119
|
+
"DELETE FROM {} WHERE expires_at <= clock_timestamp() "
|
|
120
|
+
"AND event_expires_at <= clock_timestamp()"
|
|
121
|
+
).format(self.table))
|
|
122
|
+
return curs.rowcount
|
iddqueue/results.py
ADDED
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
#
|
|
2
|
+
# R E S U L T S
|
|
3
|
+
#
|
|
4
|
+
# Implements a result backend using Postgres. See
|
|
5
|
+
# https://dramatiq.io/cookbook.html#results.
|
|
6
|
+
#
|
|
7
|
+
|
|
8
|
+
import logging
|
|
9
|
+
import time
|
|
10
|
+
from textwrap import dedent
|
|
11
|
+
|
|
12
|
+
from dramatiq.results import ResultBackend, ResultMissing, ResultTimeout
|
|
13
|
+
from psycopg.types.json import Jsonb
|
|
14
|
+
|
|
15
|
+
from .cancellation import ResultCancelled
|
|
16
|
+
from .utils import (
|
|
17
|
+
QueryManager,
|
|
18
|
+
make_pool,
|
|
19
|
+
retry_pg,
|
|
20
|
+
tidy4json,
|
|
21
|
+
transaction,
|
|
22
|
+
wait_for_notifies,
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
logger = logging.getLogger(__name__)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class PostgresBackend(ResultBackend):
|
|
29
|
+
def __init__(self, *, url=None, pool=None, schema=None, prefix=None, **kw):
|
|
30
|
+
if kw.get("use_namespace_prefix_keys"):
|
|
31
|
+
raise ValueError("use_namespace_prefix_keys is unsupported for UUID results; use schema/prefix isolation")
|
|
32
|
+
super().__init__(**kw)
|
|
33
|
+
|
|
34
|
+
if pool is not None and url:
|
|
35
|
+
raise ValueError("You can't set both pool and URL!")
|
|
36
|
+
self._owns_pool = pool is None
|
|
37
|
+
self.pool = make_pool(url or "") if pool is None else pool
|
|
38
|
+
|
|
39
|
+
self.queries = QueryManager(QUERIES.queries, schema or "dramatiq", prefix or "")
|
|
40
|
+
|
|
41
|
+
def close(self):
|
|
42
|
+
if self._owns_pool:
|
|
43
|
+
self.pool.close()
|
|
44
|
+
|
|
45
|
+
def build_message_key(self, message):
|
|
46
|
+
# Just use message_id, it's UNIQUE in table.
|
|
47
|
+
return str(message.message_id)
|
|
48
|
+
|
|
49
|
+
@retry_pg
|
|
50
|
+
def get_result(self, message, *, block=False, timeout=None):
|
|
51
|
+
key = self.build_message_key(message)
|
|
52
|
+
|
|
53
|
+
timeout = 300_000 if timeout is None else timeout
|
|
54
|
+
deadline = time.monotonic() + timeout / 1000
|
|
55
|
+
channel = self.queries.channel(key, "results")
|
|
56
|
+
with transaction(self.pool, listen=channel) as curs:
|
|
57
|
+
while True:
|
|
58
|
+
curs.execute(self.queries.GET, (key,))
|
|
59
|
+
row = curs.fetchone()
|
|
60
|
+
if row is not None:
|
|
61
|
+
if row[1] == "cancelled":
|
|
62
|
+
raise ResultCancelled(str(message))
|
|
63
|
+
return self.unwrap_result(row[0])
|
|
64
|
+
if not block:
|
|
65
|
+
raise ResultMissing(message)
|
|
66
|
+
remaining = deadline - time.monotonic()
|
|
67
|
+
if remaining <= 0:
|
|
68
|
+
raise ResultTimeout(message)
|
|
69
|
+
wait_for_notifies(curs.connection, timeout=remaining)
|
|
70
|
+
|
|
71
|
+
@retry_pg
|
|
72
|
+
def _store(self, key, result, ttl):
|
|
73
|
+
with transaction(self.pool) as curs:
|
|
74
|
+
logger.debug("Storing result for %s.", key)
|
|
75
|
+
curs.execute(
|
|
76
|
+
self.queries.STORE,
|
|
77
|
+
(
|
|
78
|
+
key,
|
|
79
|
+
Jsonb(tidy4json(result)),
|
|
80
|
+
f"{ttl} ms",
|
|
81
|
+
self.queries.channel(key, "results"),
|
|
82
|
+
),
|
|
83
|
+
)
|
|
84
|
+
if 0 == curs.rowcount:
|
|
85
|
+
raise Exception(f"Can't store result of message {key}.")
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
QUERIES = QueryManager(
|
|
89
|
+
dict(
|
|
90
|
+
GET=dedent(
|
|
91
|
+
"""\
|
|
92
|
+
SELECT result, state::text
|
|
93
|
+
FROM {schema}.{tablename}
|
|
94
|
+
WHERE message_id = %s AND (state::text = 'cancelled' OR
|
|
95
|
+
(result IS NOT NULL AND result_ttl > NOW()));
|
|
96
|
+
"""
|
|
97
|
+
),
|
|
98
|
+
STORE=dedent(
|
|
99
|
+
"""\
|
|
100
|
+
WITH stored AS (
|
|
101
|
+
INSERT INTO {schema}.{tablename}
|
|
102
|
+
(queue_name, message_id, "state", result, result_ttl)
|
|
103
|
+
VALUES ('__RQ__', %s, 'done',
|
|
104
|
+
%s, NOW() + %s::interval)
|
|
105
|
+
ON CONFLICT (message_id)
|
|
106
|
+
DO UPDATE SET mtime = NOW(),
|
|
107
|
+
result = EXCLUDED.result,
|
|
108
|
+
result_ttl = EXCLUDED.result_ttl
|
|
109
|
+
WHERE {schema}.{tablename}.state <> 'cancelled'
|
|
110
|
+
RETURNING queue_name, message_id, result
|
|
111
|
+
)
|
|
112
|
+
SELECT
|
|
113
|
+
pg_notify(%s, message_id::text)
|
|
114
|
+
FROM stored;
|
|
115
|
+
"""
|
|
116
|
+
),
|
|
117
|
+
)
|
|
118
|
+
)
|
iddqueue/scheduler.py
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
"""Fixed intervals; PostgreSQL owns time and the occurrence commit boundary."""
|
|
2
|
+
from datetime import timezone
|
|
3
|
+
from uuid import uuid4, uuid5
|
|
4
|
+
|
|
5
|
+
from dramatiq import Message
|
|
6
|
+
from psycopg import sql
|
|
7
|
+
from psycopg.types.json import Jsonb
|
|
8
|
+
|
|
9
|
+
from .utils import tidy4json, transaction
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class PostgresScheduler:
|
|
13
|
+
def __init__(self, broker):
|
|
14
|
+
self.broker = broker
|
|
15
|
+
self.table = sql.Identifier(broker.queries.schema, broker.queries.prefix + "schedules")
|
|
16
|
+
|
|
17
|
+
def create(self, name, message, *, interval_ms, start_at=None):
|
|
18
|
+
if not isinstance(name, str) or not name:
|
|
19
|
+
raise ValueError("schedule name must be nonempty")
|
|
20
|
+
if type(interval_ms) is not int or not 0 < interval_ms <= 2**63 - 1:
|
|
21
|
+
raise ValueError("interval_ms must be positive bigint milliseconds")
|
|
22
|
+
if start_at is not None and (start_at.tzinfo is None or start_at.utcoffset() is None):
|
|
23
|
+
raise ValueError("start_at must include a timezone")
|
|
24
|
+
schedule_id = uuid4()
|
|
25
|
+
with transaction(self.broker.pool) as curs:
|
|
26
|
+
curs.execute(sql.SQL(
|
|
27
|
+
"INSERT INTO {} (schedule_id, name, message, interval_ms, next_run) "
|
|
28
|
+
"VALUES (%s, %s, %s, %s, COALESCE(%s::timestamptz, clock_timestamp()))"
|
|
29
|
+
).format(self.table), (schedule_id, name, Jsonb(tidy4json(message)), interval_ms, start_at))
|
|
30
|
+
return str(schedule_id)
|
|
31
|
+
|
|
32
|
+
def list(self):
|
|
33
|
+
with transaction(self.broker.pool) as curs:
|
|
34
|
+
curs.execute(sql.SQL(
|
|
35
|
+
"SELECT schedule_id, name, message->>'actor_name', message->>'queue_name', "
|
|
36
|
+
"interval_ms, next_run, enabled FROM {} ORDER BY name"
|
|
37
|
+
).format(self.table))
|
|
38
|
+
rows = curs.fetchall()
|
|
39
|
+
return [dict(schedule_id=str(i), name=n, actor=a, queue=q, interval_ms=interval,
|
|
40
|
+
next_run=due.astimezone(timezone.utc).isoformat(), enabled=enabled)
|
|
41
|
+
for i, n, a, q, interval, due, enabled in rows]
|
|
42
|
+
|
|
43
|
+
def disable(self, name):
|
|
44
|
+
with transaction(self.broker.pool) as curs:
|
|
45
|
+
curs.execute(sql.SQL("UPDATE {} SET enabled = FALSE WHERE name = %s").format(self.table), (name,))
|
|
46
|
+
return bool(curs.rowcount)
|
|
47
|
+
|
|
48
|
+
def tick(self, *, limit=100):
|
|
49
|
+
"""Publish due rows once, coalescing missed intervals; no automatic retry."""
|
|
50
|
+
if type(limit) is not int or not 1 <= limit <= 1000:
|
|
51
|
+
raise ValueError("limit must be between 1 and 1000")
|
|
52
|
+
published = []
|
|
53
|
+
with transaction(self.broker.pool) as curs:
|
|
54
|
+
curs.execute("SELECT clock_timestamp()")
|
|
55
|
+
now = curs.fetchone()[0]
|
|
56
|
+
curs.execute(sql.SQL(
|
|
57
|
+
"SELECT schedule_id, message, next_run FROM {} WHERE enabled AND next_run <= %s "
|
|
58
|
+
"ORDER BY next_run, schedule_id LIMIT %s FOR UPDATE SKIP LOCKED"
|
|
59
|
+
).format(self.table), (now, limit))
|
|
60
|
+
for schedule_id, payload, due in curs.fetchall():
|
|
61
|
+
occurrence = due.astimezone(timezone.utc).isoformat()
|
|
62
|
+
message = Message(**payload).copy(
|
|
63
|
+
message_id=str(uuid5(schedule_id, occurrence)), message_timestamp=int(now.timestamp() * 1000))
|
|
64
|
+
message = self.broker.enqueue_in_transaction(
|
|
65
|
+
message, connection=curs.connection,
|
|
66
|
+
deduplication_key=f"iddqueue:schedule:{schedule_id}:{occurrence}",
|
|
67
|
+
deduplication_ttl=7 * 24 * 60 * 60 * 1000)
|
|
68
|
+
curs.execute(sql.SQL(
|
|
69
|
+
"WITH tick AS (SELECT clock_timestamp() AS now) UPDATE {} "
|
|
70
|
+
"SET next_run = next_run + (floor(extract(epoch FROM (tick.now - next_run)) "
|
|
71
|
+
"* 1000 / interval_ms) + 1) * interval_ms * interval '1 millisecond' "
|
|
72
|
+
"FROM tick WHERE schedule_id = %s"
|
|
73
|
+
).format(self.table), (schedule_id,))
|
|
74
|
+
published.append(message)
|
|
75
|
+
return published
|
iddqueue/scheduler.sql
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
CREATE TABLE IF NOT EXISTS :"schema".:"schedules" (
|
|
2
|
+
schedule_id UUID PRIMARY KEY,
|
|
3
|
+
name TEXT NOT NULL UNIQUE CHECK (name <> ''),
|
|
4
|
+
message JSONB NOT NULL,
|
|
5
|
+
interval_ms BIGINT NOT NULL CHECK (interval_ms > 0),
|
|
6
|
+
next_run TIMESTAMPTZ NOT NULL,
|
|
7
|
+
enabled BOOLEAN NOT NULL DEFAULT TRUE
|
|
8
|
+
);
|
|
9
|
+
CREATE INDEX IF NOT EXISTS :"schedules_due" ON :"schema".:"schedules" (next_run) WHERE enabled;
|
iddqueue/schema.py
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
import os.path
|
|
2
|
+
|
|
3
|
+
from .utils import quote_ident
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def process_psql_lines(raw_lines, schema, prefix):
|
|
7
|
+
schema = quote_ident(schema)
|
|
8
|
+
tablename = quote_ident(prefix + "queue")
|
|
9
|
+
statename = quote_ident(prefix + "state")
|
|
10
|
+
|
|
11
|
+
for line in raw_lines:
|
|
12
|
+
if line.startswith("\\"):
|
|
13
|
+
continue
|
|
14
|
+
yield (
|
|
15
|
+
line.replace(':"schema"', schema)
|
|
16
|
+
.replace(':"state"', statename)
|
|
17
|
+
.replace(':"queue"', tablename)
|
|
18
|
+
.replace(':"schedules"', quote_ident(prefix + "schedules"))
|
|
19
|
+
.replace(':"schedules_due"', quote_ident(prefix + "schedules_due"))
|
|
20
|
+
.replace(':"attempts"', quote_ident(prefix + "attempts"))
|
|
21
|
+
.replace(':"attempts_message"', quote_ident(prefix + "attempts_message"))
|
|
22
|
+
.replace(':"attempts_age"', quote_ident(prefix + "attempts_age"))
|
|
23
|
+
.replace(':"coordination"', quote_ident(prefix + "coordination"))
|
|
24
|
+
.replace(':"queue_control"', quote_ident(prefix + "queue_control"))
|
|
25
|
+
.replace(':"deduplication"', quote_ident(prefix + "deduplication"))
|
|
26
|
+
.replace(':"deduplication_expiry"', quote_ident(prefix + "deduplication_expiry"))
|
|
27
|
+
)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def generate_init_sql(schema="dramatiq", prefix=""):
|
|
31
|
+
"""Returns SQL for schema initialisation
|
|
32
|
+
|
|
33
|
+
Interpolate schema and prefix and return a single SQL string for execution
|
|
34
|
+
on a PostgreSQL connection.
|
|
35
|
+
"""
|
|
36
|
+
|
|
37
|
+
path = os.path.dirname(__file__) + "/schema.sql"
|
|
38
|
+
with open(path) as fo:
|
|
39
|
+
return "\n".join(process_psql_lines(fo, schema, prefix)) + "\n" + generate_upgrade_sql(schema, prefix)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def generate_coordination_sql(schema="dramatiq", prefix=""):
|
|
43
|
+
"""Idempotent upgrade for existing databases; leaves queue data intact."""
|
|
44
|
+
path = os.path.join(os.path.dirname(__file__), "coordination.sql")
|
|
45
|
+
with open(path) as fo:
|
|
46
|
+
return "\n".join(process_psql_lines(fo, schema, prefix))
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def generate_upgrade_sql(schema="dramatiq", prefix=""):
|
|
50
|
+
"""Add optional storage without changing existing queue data."""
|
|
51
|
+
parts = [generate_coordination_sql(schema, prefix)]
|
|
52
|
+
for name in ("deduplication.sql", "control.sql", "cancellation.sql", "history.sql", "scheduler.sql"):
|
|
53
|
+
with open(os.path.join(os.path.dirname(__file__), name)) as fo:
|
|
54
|
+
parts.append("\n".join(process_psql_lines(fo, schema, prefix)))
|
|
55
|
+
return "\n".join(parts)
|
iddqueue/schema.sql
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
\set ON_ERROR_STOP on
|
|
2
|
+
\set schema 'dramatiq'
|
|
3
|
+
\set state 'state'
|
|
4
|
+
\set queue 'queue'
|
|
5
|
+
|
|
6
|
+
CREATE SCHEMA IF NOT EXISTS :"schema";
|
|
7
|
+
|
|
8
|
+
CREATE TYPE :"schema".:"state" AS ENUM (
|
|
9
|
+
'queued',
|
|
10
|
+
'consumed',
|
|
11
|
+
'rejected',
|
|
12
|
+
'cancelled',
|
|
13
|
+
'done'
|
|
14
|
+
);
|
|
15
|
+
|
|
16
|
+
CREATE TABLE :"schema".:"queue"(
|
|
17
|
+
message_id uuid PRIMARY KEY,
|
|
18
|
+
started BOOLEAN NOT NULL DEFAULT FALSE,
|
|
19
|
+
cancel_requested BOOLEAN NOT NULL DEFAULT FALSE,
|
|
20
|
+
queue_name TEXT NOT NULL DEFAULT 'default',
|
|
21
|
+
"state" :"schema".:"state",
|
|
22
|
+
mtime TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
|
23
|
+
-- message as encoded by dramatiq.
|
|
24
|
+
message JSONB,
|
|
25
|
+
"result" JSONB,
|
|
26
|
+
result_ttl TIMESTAMP WITH TIME ZONE
|
|
27
|
+
);
|
|
28
|
+
|
|
29
|
+
-- Index state and mtime together to speed up deletion. This can also speed up
|
|
30
|
+
-- statistics when VACUUM ANALYZE is recent enough.
|
|
31
|
+
CREATE INDEX ON :"schema".:"queue"("state", mtime);
|