35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278 | class RunStore:
"""Small SQLite metadata store. A fresh bounded connection is used for every operation."""
def __init__(self, path: Path, *, retention_days=30, max_rows=10000, clock=time.time):
self.path = path
self.retention_seconds = retention_days * 86400
self.max_rows = max_rows
self.clock = clock
@contextmanager
def connection(self, *, write=False):
connection = sqlite3.connect(str(self.path), timeout=2, isolation_level=None)
connection.row_factory = sqlite3.Row
try:
connection.execute("PRAGMA busy_timeout=2000")
if write:
connection.execute("BEGIN IMMEDIATE")
yield connection
if write:
connection.commit()
except BaseException:
if write:
connection.rollback()
raise
finally:
connection.close()
def initialize(self):
# The directory is provisioned by the operator, avoiding silently ephemeral storage.
# SQLite may report BUSY immediately while another worker changes journal mode.
# Retry only that bounded initialization step; never conceal schema or filesystem errors.
for attempt in range(20):
try:
with self.connection() as db:
db.execute("PRAGMA busy_timeout=100")
db.execute("PRAGMA journal_mode=WAL")
break
except sqlite3.OperationalError as error:
if (
error.sqlite_errorcode & 255 not in {sqlite3.SQLITE_BUSY, sqlite3.SQLITE_LOCKED}
or attempt == 19
):
raise
time.sleep(0.01)
with self.connection(write=True) as db:
version = db.execute("PRAGMA user_version").fetchone()[0]
if version not in {0, 1}:
raise ValueError("Unsupported run-history database schema version")
db.execute("""CREATE TABLE IF NOT EXISTS runs (
uid TEXT PRIMARY KEY, namespace TEXT NOT NULL, app_id TEXT, name TEXT NOT NULL,
driver_pod TEXT NOT NULL, image TEXT, spark_version TEXT, status TEXT NOT NULL,
created_at TEXT, first_seen REAL NOT NULL, last_seen REAL NOT NULL, disappeared_at REAL
)""")
db.execute("CREATE INDEX IF NOT EXISTS runs_listing ON runs(first_seen DESC, uid DESC)")
db.execute("CREATE INDEX IF NOT EXISTS runs_namespace ON runs(namespace)")
db.execute("""CREATE TABLE IF NOT EXISTS collection (
namespace TEXT PRIMARY KEY, last_attempt REAL, last_success REAL,
successes INTEGER NOT NULL DEFAULT 0, failures INTEGER NOT NULL DEFAULT 0, last_error TEXT
)""")
# Earlier version 1 databases lack the failure category; older readers ignore the column.
if "last_error" not in {row["name"] for row in db.execute("PRAGMA table_info(collection)")}:
db.execute("ALTER TABLE collection ADD COLUMN last_error TEXT")
db.execute("""CREATE TABLE IF NOT EXISTS collector_lease (
singleton INTEGER PRIMARY KEY CHECK(singleton=1), owner TEXT NOT NULL,
token INTEGER NOT NULL, expires REAL NOT NULL
)""")
db.execute("INSERT OR IGNORE INTO collector_lease VALUES(1, '', 0, 0)")
db.execute("PRAGMA user_version=1")
def acquire(self, owner: str, seconds: float) -> Lease | None:
with self.connection(write=True) as db:
row = db.execute("SELECT * FROM collector_lease WHERE singleton=1").fetchone()
now = self.clock()
if row["owner"] != owner and row["expires"] > now:
return None
token = row["token"] + 1
db.execute(
"UPDATE collector_lease SET owner=?,token=?,expires=? WHERE singleton=1",
(owner, token, now + seconds),
)
return Lease(owner, token)
def release(self, lease: Lease):
with self.connection(write=True) as db:
db.execute(
"UPDATE collector_lease SET owner='',expires=0 WHERE singleton=1 AND owner=? AND token=?",
(lease.owner, lease.token),
)
def _fence(self, db, lease):
row = db.execute("SELECT owner,token,expires FROM collector_lease WHERE singleton=1").fetchone()
if row["owner"] != lease.owner or row["token"] != lease.token or row["expires"] <= self.clock():
raise LostLease("The run-history collector lease expired or changed")
def snapshot(self, namespace: str, records: list[dict], lease: Lease):
"""Commit a complete namespaced snapshot, or roll it back in its entirety."""
with self.connection(write=True) as db:
self._fence(db, lease)
now = self.clock()
db.execute("CREATE TEMP TABLE seen(uid TEXT PRIMARY KEY)")
for record in records:
if record["namespace"] != namespace or record["status"] not in STATUSES:
raise ValueError("Invalid run metadata")
db.execute("INSERT INTO seen VALUES(?)", (record["uid"],))
db.execute(
"""INSERT INTO runs VALUES(?,?,?,?,?,?,?,?,?,?,?,NULL)
ON CONFLICT(uid) DO UPDATE SET
app_id=excluded.app_id,name=excluded.name,driver_pod=excluded.driver_pod,
image=excluded.image,spark_version=excluded.spark_version,
status=CASE WHEN runs.status IN ('Succeeded','Failed')
THEN runs.status ELSE excluded.status END,
last_seen=excluded.last_seen,disappeared_at=NULL""",
tuple(
record[key]
for key in (
"uid",
"namespace",
"app_id",
"name",
"driver_pod",
"image",
"spark_version",
"status",
"created_at",
)
)
+ (now, now),
)
db.execute(
"""UPDATE runs SET disappeared_at=? WHERE namespace=? AND disappeared_at IS NULL
AND uid NOT IN (SELECT uid FROM seen)""",
(now, namespace),
)
db.execute(
"""INSERT INTO collection(namespace,last_attempt,last_success,successes)
VALUES(?,?,?,1) ON CONFLICT(namespace)
DO UPDATE SET last_attempt=excluded.last_attempt,
last_success=excluded.last_success,successes=collection.successes+1,
last_error=NULL""",
(namespace, now, now),
)
self._prune(db, now)
# The transaction lock prevents takeover during the write; this second check also fences expiry.
self._fence(db, lease)
def _prune(self, db, now):
db.execute("DELETE FROM runs WHERE last_seen < ?", (now - self.retention_seconds,))
db.execute(
"""DELETE FROM runs WHERE uid IN (
SELECT uid FROM runs ORDER BY last_seen DESC,first_seen DESC,uid DESC LIMIT -1 OFFSET ?
)""",
(self.max_rows,),
)
def failure(self, namespace: str, lease: Lease, category: str = "unexpected"):
"""Count a failed collection; ``category`` is a bounded reason, never exception text."""
with self.connection(write=True) as db:
self._fence(db, lease)
db.execute(
"""INSERT INTO collection(namespace,last_attempt,failures,last_error) VALUES(?,?,1,?)
ON CONFLICT(namespace) DO UPDATE SET last_attempt=excluded.last_attempt,
failures=collection.failures+1,last_error=excluded.last_error""",
(namespace, self.clock(), category),
)
self._prune(db, self.clock())
self._fence(db, lease)
def list_runs(
self, *, namespace=None, status=None, presence="all", limit=50, cursor=None, allowed_namespaces=None
):
if (
status is not None
and status not in STATUSES
or presence not in {"all", "present", "missing"}
or not 1 <= limit <= 200
):
raise ValueError("Invalid run-list filter")
filters = [namespace, status, presence, list(allowed_namespaces) if allowed_namespaces else None]
conditions, params = [], []
if allowed_namespaces:
conditions.append("namespace IN (" + ",".join("?" for _ in allowed_namespaces) + ")")
params.extend(allowed_namespaces)
if namespace is not None:
conditions.append("namespace=?")
params.append(namespace)
if status is not None:
conditions.append("status=?")
params.append(status)
if presence != "all":
conditions.append("disappeared_at IS " + ("NULL" if presence == "present" else "NOT NULL"))
if cursor:
try:
data = json.loads(base64.b64decode(cursor, altchars=b"-_", validate=True))
first, uid = data["first"], data["uid"]
if (
data["filters"] != filters
or type(first) not in {int, float}
or not math.isfinite(first)
or not isinstance(uid, str)
):
raise ValueError
# first_seen is REAL; also bounds integers that SQLite cannot bind.
first = float(first)
# Python 3.11 raises RecursionError for JSON nested deeper than its recursion limit.
except (ValueError, KeyError, TypeError, OverflowError, RecursionError) as error:
raise ValueError("Invalid run-list cursor") from error
conditions.append("(first_seen < ? OR (first_seen = ? AND uid < ?))")
params.extend((first, first, uid))
where = " WHERE " + " AND ".join(conditions) if conditions else ""
with self.connection() as db:
rows = db.execute(
"SELECT * FROM runs" + where + " ORDER BY first_seen DESC,uid DESC LIMIT ?",
(*params, limit + 1),
).fetchall()
next_cursor = None
if len(rows) > limit:
last = rows[limit - 1]
next_cursor = base64.urlsafe_b64encode(
json.dumps(
{"first": last["first_seen"], "uid": last["uid"], "filters": filters},
separators=(",", ":"),
).encode()
).decode()
items = []
for row in rows[:limit]:
item = dict(row)
for field in ("first_seen", "last_seen", "disappeared_at"):
item[field] = timestamp(item[field])
items.append(item)
return {"schema_version": 1, "items": items, "next_cursor": next_cursor}
def health(self):
with self.connection() as db:
return {row["namespace"]: dict(row) for row in db.execute("SELECT * FROM collection")}
def counts(self):
with self.connection() as db:
return [
dict(row)
for row in db.execute(
"SELECT namespace,status,COUNT(*) AS count FROM runs "
"GROUP BY namespace,status ORDER BY namespace,status"
)
]
|