286 lines
12 KiB
Python
286 lines
12 KiB
Python
import json
|
|
import os
|
|
import sqlite3
|
|
from collections.abc import Iterable, Mapping
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PushDelivery:
|
|
session_id: str
|
|
subscription: dict
|
|
thread_revisions: tuple[tuple[int, str], ...]
|
|
digest_revisions: tuple[tuple[int, str], ...] = ()
|
|
|
|
@property
|
|
def thread_ids(self) -> tuple[int, ...]:
|
|
return tuple(thread_id for thread_id, _revision in self.thread_revisions)
|
|
|
|
@property
|
|
def digest_ids(self) -> tuple[int, ...]:
|
|
return tuple(thread_id for thread_id, _revision in self.digest_revisions)
|
|
|
|
|
|
def _revisions(
|
|
values: Mapping[int, str] | Iterable[int | tuple[int, str]],
|
|
) -> tuple[tuple[int, str], ...]:
|
|
if isinstance(values, Mapping):
|
|
items = list(values.items())
|
|
else:
|
|
items = []
|
|
for value in values:
|
|
if isinstance(value, tuple):
|
|
items.append(value)
|
|
else:
|
|
items.append((value, ""))
|
|
normalized = {}
|
|
for thread_id, revision in items:
|
|
thread_id = int(thread_id)
|
|
if thread_id > 0:
|
|
normalized[thread_id] = str(revision or "")
|
|
return tuple(sorted(normalized.items()))
|
|
|
|
|
|
class PushSubscriptionStore:
|
|
"""Durable, device-bound Web Push subscriptions and delivery deduplication."""
|
|
|
|
def __init__(self, path: str | Path):
|
|
self.path = Path(path)
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
os.chmod(self.path.parent, 0o700)
|
|
with self._connect() as connection:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS push_subscriptions (
|
|
session_id TEXT PRIMARY KEY,
|
|
endpoint TEXT NOT NULL UNIQUE,
|
|
subscription_json TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS push_deliveries (
|
|
session_id TEXT NOT NULL,
|
|
thread_id INTEGER NOT NULL,
|
|
revision TEXT NOT NULL DEFAULT '',
|
|
PRIMARY KEY (session_id, thread_id),
|
|
FOREIGN KEY (session_id) REFERENCES push_subscriptions(session_id)
|
|
ON DELETE CASCADE
|
|
);
|
|
CREATE TABLE IF NOT EXISTS push_dispatch_lease (
|
|
singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
|
|
owner TEXT NOT NULL,
|
|
expires_at REAL NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS push_digest_pending (
|
|
session_id TEXT NOT NULL,
|
|
thread_id INTEGER NOT NULL,
|
|
revision TEXT NOT NULL DEFAULT '',
|
|
PRIMARY KEY (session_id, thread_id),
|
|
FOREIGN KEY (session_id) REFERENCES push_subscriptions(session_id)
|
|
ON DELETE CASCADE
|
|
);
|
|
"""
|
|
)
|
|
delivery_columns = {
|
|
row[1] for row in connection.execute("PRAGMA table_info(push_deliveries)")
|
|
}
|
|
if "revision" not in delivery_columns:
|
|
connection.execute(
|
|
"ALTER TABLE push_deliveries ADD COLUMN revision TEXT NOT NULL DEFAULT '*'"
|
|
)
|
|
digest_columns = {
|
|
row[1]
|
|
for row in connection.execute("PRAGMA table_info(push_digest_pending)")
|
|
}
|
|
if "revision" not in digest_columns:
|
|
connection.execute(
|
|
"ALTER TABLE push_digest_pending ADD COLUMN revision TEXT NOT NULL DEFAULT '*'"
|
|
)
|
|
os.chmod(self.path, 0o600)
|
|
|
|
def _connect(self):
|
|
connection = sqlite3.connect(self.path, timeout=2)
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
return connection
|
|
|
|
def acquire_dispatch_lease(
|
|
self, owner: str, *, now: float, lease_seconds: float
|
|
) -> bool:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
current = connection.execute(
|
|
"SELECT owner, expires_at FROM push_dispatch_lease WHERE singleton = 1"
|
|
).fetchone()
|
|
if current is not None and current[0] != owner and current[1] > now:
|
|
return False
|
|
connection.execute(
|
|
"""INSERT INTO push_dispatch_lease(singleton, owner, expires_at)
|
|
VALUES (1, ?, ?)
|
|
ON CONFLICT(singleton) DO UPDATE SET
|
|
owner = excluded.owner, expires_at = excluded.expires_at""",
|
|
(owner, now + lease_seconds),
|
|
)
|
|
return True
|
|
|
|
def release_dispatch_lease(self, owner: str) -> bool:
|
|
with self._connect() as connection:
|
|
result = connection.execute(
|
|
"DELETE FROM push_dispatch_lease WHERE singleton = 1 AND owner = ?",
|
|
(owner,),
|
|
)
|
|
return result.rowcount == 1
|
|
|
|
def upsert(self, session_id: str, subscription: dict) -> None:
|
|
endpoint = subscription["endpoint"]
|
|
encoded = json.dumps(subscription, separators=(",", ":"), sort_keys=True)
|
|
with self._connect() as connection:
|
|
connection.execute("DELETE FROM push_subscriptions WHERE endpoint = ?", (endpoint,))
|
|
connection.execute("DELETE FROM push_subscriptions WHERE session_id = ?", (session_id,))
|
|
connection.execute(
|
|
"INSERT INTO push_subscriptions(session_id, endpoint, subscription_json) VALUES (?, ?, ?)",
|
|
(session_id, endpoint, encoded),
|
|
)
|
|
|
|
def delete_session(self, session_id: str) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute("DELETE FROM push_subscriptions WHERE session_id = ?", (session_id,))
|
|
|
|
def delete_all(self) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute("DELETE FROM push_subscriptions")
|
|
|
|
def is_subscribed(self, session_id: str) -> bool:
|
|
with self._connect() as connection:
|
|
return connection.execute(
|
|
"SELECT 1 FROM push_subscriptions WHERE session_id = ?", (session_id,)
|
|
).fetchone() is not None
|
|
|
|
def claim_unseen(
|
|
self, thread_revisions: Mapping[int, str] | Iterable[int | tuple[int, str]]
|
|
) -> list[PushDelivery]:
|
|
candidates = _revisions(thread_revisions)
|
|
if not candidates:
|
|
return []
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT session_id, subscription_json FROM push_subscriptions ORDER BY session_id"
|
|
).fetchall()
|
|
deliveries = []
|
|
for session_id, encoded in rows:
|
|
delivered = {
|
|
row[0]: row[1]
|
|
for row in connection.execute(
|
|
"SELECT thread_id, revision FROM push_deliveries WHERE session_id = ?",
|
|
(session_id,),
|
|
)
|
|
}
|
|
for thread_id, revision in candidates:
|
|
if delivered.get(thread_id) == "*":
|
|
connection.execute(
|
|
"""UPDATE push_deliveries SET revision = ?
|
|
WHERE session_id = ? AND thread_id = ?""",
|
|
(revision, session_id, thread_id),
|
|
)
|
|
delivered[thread_id] = revision
|
|
unseen = tuple(
|
|
candidate
|
|
for candidate in candidates
|
|
if (
|
|
candidate[0] not in delivered
|
|
or (
|
|
bool(candidate[1])
|
|
and (
|
|
not delivered[candidate[0]]
|
|
or candidate[1] > delivered[candidate[0]]
|
|
)
|
|
)
|
|
)
|
|
)
|
|
if unseen:
|
|
pending = {
|
|
row[0]: row[1]
|
|
for row in connection.execute(
|
|
"SELECT thread_id, revision FROM push_digest_pending WHERE session_id = ?",
|
|
(session_id,),
|
|
)
|
|
}
|
|
digest_revisions = []
|
|
for thread_id, revision in unseen:
|
|
if pending.get(thread_id) in {revision, "*"}:
|
|
digest_revisions.append((thread_id, revision))
|
|
if pending[thread_id] == "*":
|
|
connection.execute(
|
|
"""UPDATE push_digest_pending SET revision = ?
|
|
WHERE session_id = ? AND thread_id = ?""",
|
|
(revision, session_id, thread_id),
|
|
)
|
|
deliveries.append(
|
|
PushDelivery(
|
|
session_id,
|
|
json.loads(encoded),
|
|
unseen,
|
|
tuple(digest_revisions),
|
|
)
|
|
)
|
|
return deliveries
|
|
|
|
def reconcile_unread(self, thread_ids: Iterable[int]) -> None:
|
|
"""Prune per-device checkpoints that are absent from a complete snapshot."""
|
|
unread_ids = tuple(sorted({int(value) for value in thread_ids if int(value) > 0}))
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
connection.execute(
|
|
"CREATE TEMP TABLE current_unread(thread_id INTEGER PRIMARY KEY)"
|
|
)
|
|
connection.executemany(
|
|
"INSERT INTO current_unread(thread_id) VALUES (?)",
|
|
((thread_id,) for thread_id in unread_ids),
|
|
)
|
|
connection.execute(
|
|
"""DELETE FROM push_deliveries
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM current_unread
|
|
WHERE current_unread.thread_id = push_deliveries.thread_id
|
|
)"""
|
|
)
|
|
connection.execute(
|
|
"""DELETE FROM push_digest_pending
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM current_unread
|
|
WHERE current_unread.thread_id = push_digest_pending.thread_id
|
|
)"""
|
|
)
|
|
|
|
def mark_digest_pending(
|
|
self,
|
|
session_id: str,
|
|
thread_revisions: Mapping[int, str] | Iterable[int | tuple[int, str]],
|
|
) -> None:
|
|
values = _revisions(thread_revisions)
|
|
with self._connect() as connection:
|
|
connection.executemany(
|
|
"""INSERT INTO push_digest_pending(session_id, thread_id, revision)
|
|
VALUES (?, ?, ?)
|
|
ON CONFLICT(session_id, thread_id) DO UPDATE SET
|
|
revision = excluded.revision""",
|
|
((session_id, thread_id, revision) for thread_id, revision in values),
|
|
)
|
|
|
|
def mark_delivered(
|
|
self,
|
|
session_id: str,
|
|
thread_revisions: Mapping[int, str] | Iterable[int | tuple[int, str]],
|
|
) -> None:
|
|
values = _revisions(thread_revisions)
|
|
with self._connect() as connection:
|
|
connection.executemany(
|
|
"""INSERT INTO push_deliveries(session_id, thread_id, revision)
|
|
VALUES (?, ?, ?)
|
|
ON CONFLICT(session_id, thread_id) DO UPDATE SET
|
|
revision = excluded.revision""",
|
|
((session_id, thread_id, revision) for thread_id, revision in values),
|
|
)
|
|
connection.executemany(
|
|
"DELETE FROM push_digest_pending WHERE session_id = ? AND thread_id = ?",
|
|
((session_id, thread_id) for thread_id, _revision in values),
|
|
)
|