1use anyhow::{ensure, Result};
2use rusqlite::{params, Connection, OptionalExtension};
3use serde::Serialize;
4
5use crate::db::ExtractionTaskKind;
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
8pub struct ExtractionReplayRange {
9 pub id: i64,
10 pub source_task_id: i64,
11 pub replay_task_id: Option<i64>,
12 pub task_kind: String,
13 pub project: String,
14 pub session_id: Option<String>,
15 pub from_event_id: i64,
16 pub to_event_id: i64,
17 pub status: String,
18 pub attempts: i64,
19 pub updated_at_epoch: i64,
20 pub last_error: Option<String>,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
24pub struct ExtractionReplayTaskEvidence {
25 pub id: i64,
26 pub status: String,
27 pub attempts: i64,
28 pub last_error: Option<String>,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
32pub struct ExtractionReplayRangeEvidence {
33 pub range: ExtractionReplayRange,
34 pub replay_task: Option<ExtractionReplayTaskEvidence>,
35}
36
37pub fn list_extraction_replay_ranges(
38 conn: &Connection,
39 project: Option<&str>,
40 limit: i64,
41) -> Result<Vec<ExtractionReplayRange>> {
42 let mut stmt = conn.prepare(
43 "SELECT r.id, r.source_task_id, r.replay_task_id, r.task_kind, p.project_path,
44 s.session_id, r.from_event_id, r.to_event_id, r.status, r.attempts,
45 r.updated_at_epoch, r.last_error
46 FROM extraction_replay_ranges r
47 JOIN projects p ON p.id = r.project_id
48 LEFT JOIN sessions s ON s.id = r.session_row_id
49 WHERE r.status IN ('pending', 'failed', 'requeued', 'quarantined')
50 AND (?1 IS NULL OR p.project_path = ?1)
51 ORDER BY r.updated_at_epoch DESC, r.id DESC
52 LIMIT ?2",
53 )?;
54 let rows = stmt.query_map(params![project, limit.max(1)], |row| {
55 Ok(ExtractionReplayRange {
56 id: row.get(0)?,
57 source_task_id: row.get(1)?,
58 replay_task_id: row.get(2)?,
59 task_kind: row.get(3)?,
60 project: row.get(4)?,
61 session_id: row.get(5)?,
62 from_event_id: row.get(6)?,
63 to_event_id: row.get(7)?,
64 status: row.get(8)?,
65 attempts: row.get(9)?,
66 updated_at_epoch: row.get(10)?,
67 last_error: row.get(11)?,
68 })
69 })?;
70 crate::db::query::collect_rows(rows)
71}
72
73pub fn get_extraction_replay_range_evidence(
74 conn: &Connection,
75 range_id: i64,
76) -> Result<ExtractionReplayRangeEvidence> {
77 ensure!(range_id > 0, "extraction replay range id must be positive");
78 conn.query_row(
79 "SELECT r.id, r.source_task_id, r.replay_task_id, r.task_kind, p.project_path,
80 s.session_id, r.from_event_id, r.to_event_id, r.status, r.attempts,
81 r.updated_at_epoch, r.last_error,
82 t.id, t.status, t.attempts, t.last_error
83 FROM extraction_replay_ranges r
84 JOIN projects p ON p.id = r.project_id
85 LEFT JOIN sessions s ON s.id = r.session_row_id
86 LEFT JOIN extraction_tasks t ON t.id = r.replay_task_id
87 WHERE r.id = ?1",
88 params![range_id],
89 |row| {
90 let replay_task_id = row.get::<_, Option<i64>>(12)?;
91 Ok(ExtractionReplayRangeEvidence {
92 range: ExtractionReplayRange {
93 id: row.get(0)?,
94 source_task_id: row.get(1)?,
95 replay_task_id: row.get(2)?,
96 task_kind: row.get(3)?,
97 project: row.get(4)?,
98 session_id: row.get(5)?,
99 from_event_id: row.get(6)?,
100 to_event_id: row.get(7)?,
101 status: row.get(8)?,
102 attempts: row.get(9)?,
103 updated_at_epoch: row.get(10)?,
104 last_error: row.get(11)?,
105 },
106 replay_task: if let Some(id) = replay_task_id {
107 Some(ExtractionReplayTaskEvidence {
108 id,
109 status: row.get(13)?,
110 attempts: row.get(14)?,
111 last_error: row.get(15)?,
112 })
113 } else {
114 None
115 },
116 })
117 },
118 )
119 .optional()?
120 .ok_or_else(|| anyhow::anyhow!("extraction replay range {range_id} does not exist"))
121}
122
123pub fn count_retryable_extraction_replay_ranges(
124 conn: &Connection,
125 project: Option<&str>,
126 limit: i64,
127) -> Result<i64> {
128 conn.query_row(
129 "SELECT COUNT(*)
130 FROM (
131 SELECT r.id
132 FROM extraction_replay_ranges r
133 JOIN projects p ON p.id = r.project_id
134 WHERE r.status IN ('pending', 'failed')
135 AND r.archived_at_epoch IS NULL
136 AND NOT EXISTS (
137 SELECT 1
138 FROM extraction_tasks t
139 WHERE t.replay_range_id = r.id
140 AND t.status IN ('pending', 'processing')
141 )
142 AND (?1 IS NULL OR p.project_path = ?1)
143 ORDER BY r.updated_at_epoch ASC, r.id ASC
144 LIMIT ?2
145 )",
146 params![project, limit.max(1)],
147 |row| row.get(0),
148 )
149 .map_err(Into::into)
150}
151
152fn query_retryable_replay_range_ids(
153 conn: &Connection,
154 project: Option<&str>,
155 range_id: Option<i64>,
156 acknowledge_quarantine: bool,
157 include_archived: bool,
158 limit: i64,
159) -> Result<Vec<i64>> {
160 let mut stmt = conn.prepare(
161 "SELECT r.id
162 FROM extraction_replay_ranges r
163 JOIN projects p ON p.id = r.project_id
164 WHERE (r.status IN ('pending', 'failed')
165 OR (?3 = 1 AND r.status = 'quarantined'))
166 AND (?4 = 1 OR r.archived_at_epoch IS NULL)
167 AND NOT EXISTS (
168 SELECT 1
169 FROM extraction_tasks t
170 WHERE t.replay_range_id = r.id
171 AND t.status IN ('pending', 'processing')
172 )
173 AND (?1 IS NULL OR p.project_path = ?1)
174 AND (?2 IS NULL OR r.id = ?2)
175 ORDER BY r.updated_at_epoch ASC, r.id ASC
176 LIMIT ?5",
177 )?;
178 let rows = stmt.query_map(
179 params![
180 project,
181 range_id,
182 acknowledge_quarantine,
183 include_archived,
184 limit.max(1)
185 ],
186 |row| row.get::<_, i64>(0),
187 )?;
188 crate::db::query::collect_rows(rows)
189}
190
191pub fn ensure_extraction_replay_range_retryable(
192 conn: &Connection,
193 range_id: i64,
194 acknowledge_quarantine: bool,
195 include_archived: bool,
196) -> Result<()> {
197 ensure!(range_id > 0, "extraction replay range id must be positive");
198 let range_ids = query_retryable_replay_range_ids(
199 conn,
200 None,
201 Some(range_id),
202 acknowledge_quarantine,
203 include_archived,
204 1,
205 )?;
206 ensure!(
207 range_ids == [range_id],
208 "extraction replay range {range_id} is not retryable"
209 );
210 Ok(())
211}
212
213pub fn retry_extraction_replay_range(
214 conn: &Connection,
215 range_id: i64,
216 acknowledge_quarantine: bool,
217) -> Result<()> {
218 let tx = conn.unchecked_transaction()?;
219 ensure_extraction_replay_range_retryable(&tx, range_id, acknowledge_quarantine, false)?;
220 enqueue_replay_extraction_task(&tx, range_id, acknowledge_quarantine)?;
221 tx.commit()?;
222 Ok(())
223}
224
225pub(crate) fn retry_and_claim_extraction_replay_range(
226 conn: &mut Connection,
227 range_id: i64,
228 acknowledge_quarantine: bool,
229 include_archived: bool,
230 lease_owner: &str,
231 lease_secs: i64,
232) -> Result<crate::db::ExtractionTask> {
233 ensure!(
234 crate::db::is_exact_replay_worker_owner(lease_owner),
235 "exact replay recovery requires an exact replay worker owner"
236 );
237 let tx = conn.transaction()?;
238 ensure_extraction_replay_range_retryable(
239 &tx,
240 range_id,
241 acknowledge_quarantine,
242 include_archived,
243 )?;
244 let task_id = enqueue_replay_extraction_task(&tx, range_id, acknowledge_quarantine)?;
245 let task = crate::db::claim_extraction_task_by_id_in_transaction(
246 &tx,
247 task_id,
248 lease_owner,
249 lease_secs,
250 )?
251 .ok_or_else(|| {
252 anyhow::anyhow!(
253 "extraction replay task {task_id} is not pending and retry-ready for exact claim"
254 )
255 })?;
256 tx.commit()?;
257 Ok(task)
258}
259
260pub fn quarantine_extraction_replay_range(conn: &Connection, range_id: i64) -> Result<()> {
261 let tx = conn.unchecked_transaction()?;
262 ensure_extraction_replay_range_retryable(&tx, range_id, false, false)?;
263 let now = chrono::Utc::now().timestamp();
264 tx.execute(
265 "UPDATE extraction_replay_ranges
266 SET status = 'quarantined', updated_at_epoch = ?1
267 WHERE id = ?2",
268 params![now, range_id],
269 )?;
270 clear_terminal_failures_for_quiesced_range(&tx, range_id, now)?;
271 tx.commit()?;
272 Ok(())
273}
274
275pub fn retry_extraction_replay_ranges(
276 conn: &Connection,
277 project: Option<&str>,
278 limit: i64,
279) -> Result<usize> {
280 let tx = conn.unchecked_transaction()?;
281 let range_ids = query_retryable_replay_range_ids(&tx, project, None, false, false, limit)?;
282 for range_id in &range_ids {
283 enqueue_replay_extraction_task(&tx, *range_id, false)?;
284 }
285 tx.commit()?;
286 Ok(range_ids.len())
287}
288
289pub fn quarantine_extraction_replay_ranges(
290 conn: &Connection,
291 project: Option<&str>,
292 limit: i64,
293) -> Result<usize> {
294 let tx = conn.unchecked_transaction()?;
295 let range_ids = query_retryable_replay_range_ids(&tx, project, None, false, false, limit)?;
296 let now = chrono::Utc::now().timestamp();
297 for range_id in &range_ids {
298 tx.execute(
299 "UPDATE extraction_replay_ranges
300 SET status = 'quarantined', updated_at_epoch = ?1
301 WHERE id = ?2",
302 params![now, range_id],
303 )?;
304 clear_terminal_failures_for_quiesced_range(&tx, *range_id, now)?;
305 }
306 tx.commit()?;
307 Ok(range_ids.len())
308}
309
310pub(crate) fn enqueue_replay_extraction_task(
311 conn: &Connection,
312 range_id: i64,
313 acknowledge_quarantine: bool,
314) -> Result<i64> {
315 let (task_kind, host_id, workspace_id, project_id, session_row_id, from_event_id, to_event_id) =
316 conn.query_row(
317 "SELECT task_kind, host_id, workspace_id, project_id, session_row_id,
318 from_event_id, to_event_id
319 FROM extraction_replay_ranges
320 WHERE id = ?1
321 AND (status IN ('pending', 'failed')
322 OR (?2 = 1 AND status = 'quarantined'))",
323 params![range_id, acknowledge_quarantine],
324 |row| {
325 Ok((
326 row.get::<_, String>(0)?,
327 row.get::<_, i64>(1)?,
328 row.get::<_, i64>(2)?,
329 row.get::<_, i64>(3)?,
330 row.get::<_, Option<i64>>(4)?,
331 row.get::<_, i64>(5)?,
332 row.get::<_, i64>(6)?,
333 ))
334 },
335 )?;
336 let session_row_id = session_row_id.ok_or_else(|| {
337 anyhow::anyhow!("extraction replay range {range_id} is missing session_row_id")
338 })?;
339 let task_kind_value = ExtractionTaskKind::from_db(&task_kind)?;
340 let now = chrono::Utc::now().timestamp();
341 let idempotency_key =
342 format!("{host_id}:{project_id}:{session_row_id}:{task_kind}:replay-range:{range_id}");
343 conn.execute(
344 "INSERT INTO extraction_tasks
345 (task_kind, host_id, workspace_id, project_id, session_row_id, priority, status,
346 idempotency_key, cursor_event_id, high_watermark_event_id, attempts,
347 next_retry_epoch, lease_owner, lease_expires_epoch, last_error, created_at_epoch,
348 updated_at_epoch, replay_range_id)
349 VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'pending', ?7, ?8, ?9, 0, NULL, NULL, NULL, NULL,
350 ?10, ?10, ?11)
351 ON CONFLICT(idempotency_key) DO UPDATE SET
352 status = CASE
353 WHEN extraction_tasks.status IN ('done', 'failed') THEN 'pending'
354 ELSE extraction_tasks.status
355 END,
356 cursor_event_id = excluded.cursor_event_id,
357 high_watermark_event_id = excluded.high_watermark_event_id,
358 attempts = CASE
359 WHEN extraction_tasks.status IN ('done', 'failed') THEN 0
360 ELSE extraction_tasks.attempts
361 END,
362 next_retry_epoch = CASE
363 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
364 ELSE extraction_tasks.next_retry_epoch
365 END,
366 last_error = CASE
367 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
368 ELSE extraction_tasks.last_error
369 END,
370 failure_class = CASE
371 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
372 ELSE extraction_tasks.failure_class
373 END,
374 failed_at_epoch = CASE
375 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
376 ELSE extraction_tasks.failed_at_epoch
377 END,
378 archived_at_epoch = CASE
379 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
380 ELSE extraction_tasks.archived_at_epoch
381 END,
382 replay_range_id = excluded.replay_range_id,
383 updated_at_epoch = excluded.updated_at_epoch",
384 params![
385 task_kind,
386 host_id,
387 workspace_id,
388 project_id,
389 session_row_id,
390 task_kind_value.priority(),
391 idempotency_key,
392 from_event_id - 1,
393 to_event_id,
394 now,
395 range_id
396 ],
397 )?;
398 let replay_task_id: i64 = conn.query_row(
399 "SELECT id FROM extraction_tasks WHERE idempotency_key = ?1",
400 params![idempotency_key],
401 |row| row.get(0),
402 )?;
403 conn.execute(
404 "UPDATE extraction_replay_ranges
405 SET status = 'requeued',
406 replay_task_id = ?1,
407 attempts = attempts + 1,
408 failure_class = NULL,
409 failed_at_epoch = NULL,
410 archived_at_epoch = NULL,
411 updated_at_epoch = ?2
412 WHERE id = ?3",
413 params![replay_task_id, now, range_id],
414 )?;
415 Ok(replay_task_id)
416}
417
418pub(crate) fn archive_exact_replay_range_after_task_failure(
419 conn: &Connection,
420 range_id: i64,
421 replay_task_id: i64,
422 error: &str,
423 now: i64,
424) -> Result<()> {
425 let updated = conn.execute(
426 "UPDATE extraction_replay_ranges
427 SET status = 'quarantined',
428 replay_task_id = ?1,
429 last_error = ?2,
430 failure_class = ?3,
431 failed_at_epoch = COALESCE(failed_at_epoch, ?4),
432 archived_at_epoch = ?4,
433 updated_at_epoch = ?4
434 WHERE id = ?5
435 AND EXISTS (
436 SELECT 1
437 FROM extraction_tasks t
438 WHERE t.id = ?1 AND t.replay_range_id = extraction_replay_ranges.id
439 )",
440 params![
441 replay_task_id,
442 crate::db::truncate_str(error, 2000),
443 crate::db::classify_failure(error).as_str(),
444 now,
445 range_id
446 ],
447 )?;
448 ensure!(
449 updated == 1,
450 "exact replay range {range_id} is not linked to replay task {replay_task_id}"
451 );
452 Ok(())
453}
454
455#[allow(clippy::too_many_arguments)]
456pub(crate) fn record_exhausted_replay_range(
457 conn: &Connection,
458 source_task_id: i64,
459 task_kind: &str,
460 host_id: i64,
461 workspace_id: i64,
462 project_id: i64,
463 session_row_id: Option<i64>,
464 from_event_id: i64,
465 to_event_id: i64,
466 _attempts: i64,
467 err: &str,
468 now: i64,
469) -> Result<i64> {
470 conn.execute(
471 "INSERT INTO extraction_replay_ranges
472 (source_task_id, task_kind, host_id, workspace_id, project_id, session_row_id,
473 from_event_id, to_event_id, status, attempts, last_error, failure_class,
474 failed_at_epoch, archived_at_epoch, created_at_epoch, updated_at_epoch)
475 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, 'pending', 0, ?9, ?10, ?11, NULL, ?11, ?11)
476 ON CONFLICT(source_task_id, from_event_id, to_event_id) DO UPDATE SET
477 status = CASE
478 WHEN extraction_replay_ranges.status = 'quarantined' THEN 'quarantined'
479 ELSE 'pending'
480 END,
481 last_error = excluded.last_error,
482 failure_class = excluded.failure_class,
483 failed_at_epoch = COALESCE(extraction_replay_ranges.failed_at_epoch, excluded.failed_at_epoch),
484 archived_at_epoch = NULL,
485 updated_at_epoch = excluded.updated_at_epoch",
486 params![
487 source_task_id,
488 task_kind,
489 host_id,
490 workspace_id,
491 project_id,
492 session_row_id,
493 from_event_id,
494 to_event_id,
495 crate::db::truncate_str(err, 2000),
496 crate::db::classify_failure(err).as_str(),
497 now
498 ],
499 )?;
500 conn.query_row(
501 "SELECT id
502 FROM extraction_replay_ranges
503 WHERE source_task_id = ?1 AND from_event_id = ?2 AND to_event_id = ?3",
504 params![source_task_id, from_event_id, to_event_id],
505 |row| row.get(0),
506 )
507 .map_err(Into::into)
508}
509
510pub(crate) fn mark_replay_range_replayed_if_done(
511 conn: &Connection,
512 task_id: i64,
513 now: i64,
514) -> Result<()> {
515 let updated = conn.execute(
516 "UPDATE extraction_replay_ranges
517 SET status = 'replayed',
518 replay_task_id = COALESCE(replay_task_id, ?1),
519 last_error = NULL,
520 failure_class = NULL,
521 failed_at_epoch = NULL,
522 archived_at_epoch = NULL,
523 updated_at_epoch = ?2
524 WHERE id = (
525 SELECT replay_range_id FROM extraction_tasks
526 WHERE id = ?1 AND status = 'done' AND replay_range_id IS NOT NULL
527 )
528 AND status != 'quarantined'
529 AND NOT EXISTS (
530 SELECT 1 FROM extraction_tasks t
531 WHERE t.replay_range_id = extraction_replay_ranges.id
532 AND t.status != 'done'
533 )",
534 params![task_id, now],
535 )?;
536 if updated > 0 {
537 let range_id = conn.query_row(
538 "SELECT replay_range_id
539 FROM extraction_tasks
540 WHERE id = ?1",
541 params![task_id],
542 |row| row.get::<_, i64>(0),
543 )?;
544 clear_terminal_failures_for_quiesced_range(conn, range_id, now)?;
545 }
546 Ok(())
547}
548
549fn clear_terminal_failures_for_quiesced_range(
550 conn: &Connection,
551 range_id: i64,
552 now: i64,
553) -> Result<()> {
554 conn.execute(
555 "UPDATE extraction_tasks
556 SET status = 'done',
557 attempts = 0,
558 lease_owner = NULL,
559 lease_expires_epoch = NULL,
560 next_retry_epoch = NULL,
561 last_error = NULL,
562 failure_class = NULL,
563 failed_at_epoch = NULL,
564 archived_at_epoch = NULL,
565 updated_at_epoch = ?2
566 WHERE replay_range_id = ?1
567 AND status = 'failed'",
568 params![range_id, now],
569 )?;
570 conn.execute(
571 "UPDATE extraction_tasks
572 SET status = 'done',
573 attempts = 0,
574 lease_owner = NULL,
575 lease_expires_epoch = NULL,
576 next_retry_epoch = NULL,
577 last_error = NULL,
578 failure_class = NULL,
579 failed_at_epoch = NULL,
580 archived_at_epoch = NULL,
581 updated_at_epoch = ?2
582 WHERE id = (
583 SELECT source_task_id
584 FROM extraction_replay_ranges
585 WHERE id = ?1
586 )
587 AND status = 'failed'
588 AND NOT EXISTS (
589 SELECT 1
590 FROM extraction_replay_ranges r
591 WHERE r.source_task_id = extraction_tasks.id
592 AND r.status NOT IN ('replayed', 'quarantined')
593 )",
594 params![range_id, now],
595 )?;
596 Ok(())
597}
598
599pub(crate) fn mark_replay_range_failed(
600 conn: &Connection,
601 task_id: i64,
602 now: i64,
603 err: &str,
604) -> Result<()> {
605 conn.execute(
606 "UPDATE extraction_replay_ranges
607 SET status = 'failed',
608 replay_task_id = COALESCE(replay_task_id, ?1),
609 attempts = COALESCE((SELECT attempts FROM extraction_tasks WHERE id = ?1), attempts),
610 last_error = ?2,
611 failure_class = ?3,
612 failed_at_epoch = COALESCE(failed_at_epoch, ?4),
613 archived_at_epoch = NULL,
614 updated_at_epoch = ?4
615 WHERE id = (
616 SELECT replay_range_id FROM extraction_tasks
617 WHERE id = ?1 AND replay_range_id IS NOT NULL
618 )",
619 params![
620 task_id,
621 crate::db::truncate_str(err, 2000),
622 crate::db::classify_failure(err).as_str(),
623 now
624 ],
625 )?;
626 Ok(())
627}