1use anyhow::Result;
2use sqlx::PgPool;
3
4use crate::store::{
5 NOT_A_SETTLEMENT, RetryEvent, WorkflowStore, payload_activity_id, retry_denial, settle_outcome,
6};
7use crate::types::*;
8
9const RETRY_ACTIVITY_SELECT: &str = "SELECT id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat FROM workflow.activities WHERE workflow_id = $1 AND status = 'FAILED' ORDER BY seq DESC LIMIT 1 FOR UPDATE";
10const UNSETTLED_ACTIVITY_SELECT: &str = "SELECT a.id, a.workflow_id, a.seq, a.name, a.task_queue, a.input, a.status, a.result, a.error, a.attempt, a.max_attempts, a.initial_interval_secs, a.backoff_coefficient, a.start_to_close_secs, a.heartbeat_timeout_secs, a.claimed_by, a.scheduled_at, a.started_at, a.completed_at, a.last_heartbeat FROM workflow.activities a JOIN workflow.workflows w ON w.id = a.workflow_id WHERE a.status IN ('COMPLETED', 'FAILED') AND w.status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT') AND w.archived_at IS NULL AND NOT EXISTS (SELECT 1 FROM workflow.events e WHERE e.workflow_id = a.workflow_id AND e.activity_id = a.id AND e.event_type IN ('ActivityCompleted', 'ActivityFailed')) ORDER BY a.completed_at ASC LIMIT $1";
14const RETRY_ACTIVITY_UPDATE: &str = "UPDATE workflow.activities SET status = 'PENDING', result = NULL, error = NULL, attempt = 1, claimed_by = NULL, scheduled_at = $1, started_at = NULL, completed_at = NULL, last_heartbeat = NULL WHERE id = $2 RETURNING id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat";
15
16const SCHEMA: &str = r#"
22CREATE SCHEMA IF NOT EXISTS workflow;
23CREATE SCHEMA IF NOT EXISTS engine;
24
25CREATE TABLE IF NOT EXISTS workflow.namespaces (
26 name TEXT PRIMARY KEY,
27 created_at DOUBLE PRECISION NOT NULL
28);
29INSERT INTO workflow.namespaces (name, created_at)
30 VALUES ('main', EXTRACT(EPOCH FROM NOW()))
31 ON CONFLICT DO NOTHING;
32
33CREATE TABLE IF NOT EXISTS workflow.workflows (
34 id TEXT PRIMARY KEY,
35 namespace TEXT NOT NULL DEFAULT 'main',
36 run_id TEXT NOT NULL,
37 workflow_type TEXT NOT NULL,
38 task_queue TEXT NOT NULL DEFAULT 'main',
39 status TEXT NOT NULL DEFAULT 'PENDING',
40 input TEXT,
41 result TEXT,
42 error TEXT,
43 parent_id TEXT,
44 claimed_by TEXT,
45 search_attributes TEXT,
46 archived_at DOUBLE PRECISION,
47 archive_uri TEXT,
48 -- Workflow-task dispatch (Phase 9): see sqlite.rs for the full comment.
49 needs_dispatch BOOLEAN NOT NULL DEFAULT FALSE,
50 dispatch_claimed_by TEXT,
51 dispatch_last_heartbeat DOUBLE PRECISION,
52 created_at DOUBLE PRECISION NOT NULL,
53 updated_at DOUBLE PRECISION NOT NULL,
54 completed_at DOUBLE PRECISION
55);
56CREATE INDEX IF NOT EXISTS idx_wf_status_queue ON workflow.workflows(status, task_queue);
57CREATE INDEX IF NOT EXISTS idx_wf_namespace ON workflow.workflows(namespace);
58CREATE INDEX IF NOT EXISTS idx_wf_dispatch ON workflow.workflows(task_queue, needs_dispatch, dispatch_claimed_by);
59
60CREATE TABLE IF NOT EXISTS workflow.events (
61 id BIGSERIAL PRIMARY KEY,
62 workflow_id TEXT NOT NULL REFERENCES workflow.workflows(id),
63 seq INTEGER NOT NULL,
64 event_type TEXT NOT NULL,
65 payload TEXT,
66 -- Set on ActivityCompleted / ActivityFailed only. Answers "did this
67 -- activity's terminal event land" without parsing payload JSON, which
68 -- is what the settle transaction and the reconciler both need.
69 activity_id BIGINT,
70 timestamp DOUBLE PRECISION NOT NULL
71);
72CREATE INDEX IF NOT EXISTS idx_wf_events_lookup ON workflow.events(workflow_id, seq);
73
74CREATE TABLE IF NOT EXISTS workflow.activities (
75 id BIGSERIAL PRIMARY KEY,
76 workflow_id TEXT NOT NULL REFERENCES workflow.workflows(id),
77 seq INTEGER NOT NULL,
78 name TEXT NOT NULL,
79 task_queue TEXT NOT NULL DEFAULT 'main',
80 input TEXT,
81 status TEXT NOT NULL DEFAULT 'PENDING',
82 result TEXT,
83 error TEXT,
84 attempt INTEGER NOT NULL DEFAULT 1,
85 max_attempts INTEGER NOT NULL DEFAULT 3,
86 initial_interval_secs DOUBLE PRECISION NOT NULL DEFAULT 1,
87 backoff_coefficient DOUBLE PRECISION NOT NULL DEFAULT 2,
88 start_to_close_secs DOUBLE PRECISION NOT NULL DEFAULT 300,
89 heartbeat_timeout_secs DOUBLE PRECISION,
90 claimed_by TEXT,
91 scheduled_at DOUBLE PRECISION NOT NULL,
92 started_at DOUBLE PRECISION,
93 completed_at DOUBLE PRECISION,
94 last_heartbeat DOUBLE PRECISION,
95 UNIQUE (workflow_id, seq)
96);
97CREATE INDEX IF NOT EXISTS idx_wf_act_pending ON workflow.activities(task_queue, status, scheduled_at);
98
99CREATE TABLE IF NOT EXISTS workflow.timers (
100 id BIGSERIAL PRIMARY KEY,
101 workflow_id TEXT NOT NULL REFERENCES workflow.workflows(id),
102 seq INTEGER NOT NULL,
103 fire_at DOUBLE PRECISION NOT NULL,
104 fired BOOLEAN NOT NULL DEFAULT FALSE,
105 UNIQUE (workflow_id, seq)
106);
107CREATE INDEX IF NOT EXISTS idx_wf_timers_due ON workflow.timers(fire_at) WHERE fired = FALSE;
108
109CREATE TABLE IF NOT EXISTS workflow.signals (
110 id BIGSERIAL PRIMARY KEY,
111 workflow_id TEXT NOT NULL REFERENCES workflow.workflows(id),
112 name TEXT NOT NULL,
113 payload TEXT,
114 consumed BOOLEAN NOT NULL DEFAULT FALSE,
115 received_at DOUBLE PRECISION NOT NULL
116);
117CREATE INDEX IF NOT EXISTS idx_wf_signals_lookup ON workflow.signals(workflow_id, name, consumed);
118
119CREATE TABLE IF NOT EXISTS workflow.schedules (
120 namespace TEXT NOT NULL DEFAULT 'main',
121 name TEXT NOT NULL,
122 workflow_type TEXT NOT NULL,
123 cron_expr TEXT NOT NULL,
124 timezone TEXT NOT NULL DEFAULT 'UTC',
125 input TEXT,
126 task_queue TEXT NOT NULL DEFAULT 'main',
127 overlap_policy TEXT NOT NULL DEFAULT 'skip',
128 paused BOOLEAN NOT NULL DEFAULT FALSE,
129 last_run_at DOUBLE PRECISION,
130 next_run_at DOUBLE PRECISION,
131 last_workflow_id TEXT,
132 created_at DOUBLE PRECISION NOT NULL,
133 PRIMARY KEY (namespace, name)
134);
135
136CREATE TABLE IF NOT EXISTS workflow.workers (
137 id TEXT PRIMARY KEY,
138 namespace TEXT NOT NULL DEFAULT 'main',
139 identity TEXT NOT NULL,
140 task_queue TEXT NOT NULL,
141 workflows TEXT,
142 activities TEXT,
143 max_concurrent_workflows INTEGER NOT NULL DEFAULT 10,
144 max_concurrent_activities INTEGER NOT NULL DEFAULT 10,
145 active_tasks INTEGER NOT NULL DEFAULT 0,
146 last_heartbeat DOUBLE PRECISION NOT NULL,
147 registered_at DOUBLE PRECISION NOT NULL
148);
149
150CREATE TABLE IF NOT EXISTS workflow.snapshots (
151 workflow_id TEXT NOT NULL REFERENCES workflow.workflows(id),
152 event_seq INTEGER NOT NULL,
153 state_json TEXT NOT NULL,
154 created_at DOUBLE PRECISION NOT NULL,
155 PRIMARY KEY (workflow_id, event_seq)
156);
157
158-- Plan-15 slice 3: workflow.api_keys retired in favour of the auth
159-- module (sessions / JWT / Zanzibar tuples). Table is dropped on
160-- migration; nothing here re-creates it.
161DROP TABLE IF EXISTS workflow.api_keys CASCADE;
162
163CREATE TABLE IF NOT EXISTS engine.events (
164 id BIGSERIAL PRIMARY KEY,
165 ts DOUBLE PRECISION NOT NULL DEFAULT EXTRACT(EPOCH FROM NOW()),
166 namespace TEXT NOT NULL,
167 subsystem TEXT NOT NULL,
168 kind TEXT NOT NULL,
169 payload JSONB NOT NULL DEFAULT '{}'::jsonb
170);
171CREATE INDEX IF NOT EXISTS idx_engine_events_ns_id ON engine.events(namespace, id);
172CREATE INDEX IF NOT EXISTS idx_engine_events_ts_prune ON engine.events(ts);
173
174"#;
175
176const V0_13_2_RELOCATION_SQL: &str = r#"
187DO $$
188DECLARE
189 has_old BOOLEAN;
190BEGIN
191 -- Each table: if the legacy public.<old> exists, drop the empty
192 -- schema-qualified twin (created above by SCHEMA) and move the
193 -- legacy table into its new home.
194
195 -- workflows
196 SELECT to_regclass('public.workflows') IS NOT NULL INTO has_old;
197 IF has_old THEN
198 DROP TABLE IF EXISTS workflow.workflows CASCADE;
199 ALTER TABLE public.workflows SET SCHEMA workflow;
200 END IF;
201
202 -- workflow_events → workflow.events
203 SELECT to_regclass('public.workflow_events') IS NOT NULL INTO has_old;
204 IF has_old THEN
205 DROP TABLE IF EXISTS workflow.events CASCADE;
206 ALTER TABLE public.workflow_events SET SCHEMA workflow;
207 ALTER TABLE workflow.workflow_events RENAME TO events;
208 END IF;
209
210 -- workflow_activities → workflow.activities
211 SELECT to_regclass('public.workflow_activities') IS NOT NULL INTO has_old;
212 IF has_old THEN
213 DROP TABLE IF EXISTS workflow.activities CASCADE;
214 ALTER TABLE public.workflow_activities SET SCHEMA workflow;
215 ALTER TABLE workflow.workflow_activities RENAME TO activities;
216 END IF;
217
218 -- workflow_timers → workflow.timers
219 SELECT to_regclass('public.workflow_timers') IS NOT NULL INTO has_old;
220 IF has_old THEN
221 DROP TABLE IF EXISTS workflow.timers CASCADE;
222 ALTER TABLE public.workflow_timers SET SCHEMA workflow;
223 ALTER TABLE workflow.workflow_timers RENAME TO timers;
224 END IF;
225
226 -- workflow_signals → workflow.signals
227 SELECT to_regclass('public.workflow_signals') IS NOT NULL INTO has_old;
228 IF has_old THEN
229 DROP TABLE IF EXISTS workflow.signals CASCADE;
230 ALTER TABLE public.workflow_signals SET SCHEMA workflow;
231 ALTER TABLE workflow.workflow_signals RENAME TO signals;
232 END IF;
233
234 -- workflow_snapshots → workflow.snapshots
235 SELECT to_regclass('public.workflow_snapshots') IS NOT NULL INTO has_old;
236 IF has_old THEN
237 DROP TABLE IF EXISTS workflow.snapshots CASCADE;
238 ALTER TABLE public.workflow_snapshots SET SCHEMA workflow;
239 ALTER TABLE workflow.workflow_snapshots RENAME TO snapshots;
240 END IF;
241
242 -- workflow_schedules → workflow.schedules
243 SELECT to_regclass('public.workflow_schedules') IS NOT NULL INTO has_old;
244 IF has_old THEN
245 DROP TABLE IF EXISTS workflow.schedules CASCADE;
246 ALTER TABLE public.workflow_schedules SET SCHEMA workflow;
247 ALTER TABLE workflow.workflow_schedules RENAME TO schedules;
248 END IF;
249
250 -- workflow_workers → workflow.workers
251 SELECT to_regclass('public.workflow_workers') IS NOT NULL INTO has_old;
252 IF has_old THEN
253 DROP TABLE IF EXISTS workflow.workers CASCADE;
254 ALTER TABLE public.workflow_workers SET SCHEMA workflow;
255 ALTER TABLE workflow.workflow_workers RENAME TO workers;
256 END IF;
257
258 -- namespaces → workflow.namespaces
259 SELECT to_regclass('public.namespaces') IS NOT NULL INTO has_old;
260 IF has_old THEN
261 DROP TABLE IF EXISTS workflow.namespaces CASCADE;
262 ALTER TABLE public.namespaces SET SCHEMA workflow;
263 END IF;
264
265 -- public.api_keys: retired in plan-15 slice 3 (workflow REST API
266 -- auth moved to the auth module — see CHANGELOG). Drop any
267 -- orphaned legacy table so an upgraded v0.13.1 install doesn't
268 -- carry it forward.
269 SELECT to_regclass('public.api_keys') IS NOT NULL INTO has_old;
270 IF has_old THEN
271 DROP TABLE public.api_keys CASCADE;
272 END IF;
273
274 -- engine_events → engine.events (notification outbox; preserves the
275 -- v0.13.1 publish-on-commit guarantee since the new INSERT into
276 -- engine.events sits in the same transaction as the pg_notify call).
277 SELECT to_regclass('public.engine_events') IS NOT NULL INTO has_old;
278 IF has_old THEN
279 DROP TABLE IF EXISTS engine.events CASCADE;
280 ALTER TABLE public.engine_events SET SCHEMA engine;
281 ALTER TABLE engine.engine_events RENAME TO events;
282 END IF;
283END
284$$;
285"#;
286
287fn sanitise_schema(schema: &str) -> Vec<String> {
299 let without_comments: String = schema
300 .lines()
301 .filter(|line| !line.trim_start().starts_with("--"))
302 .collect::<Vec<_>>()
303 .join("\n");
304
305 without_comments
306 .split(';')
307 .map(|s| s.trim().to_string())
308 .filter(|s| !s.is_empty())
309 .collect()
310}
311
312#[derive(Clone)]
318pub struct PostgresStore {
319 pool: PgPool,
320}
321
322impl PostgresStore {
323 pub async fn new(url: &str) -> Result<Self> {
324 let pool = PgPool::connect(url).await?;
325 Self::from_pool(pool).await
326 }
327
328 pub async fn from_pool(pool: PgPool) -> Result<Self> {
333 let store = Self { pool };
334 store.migrate().await?;
335 Ok(store)
336 }
337
338 pub fn pool(&self) -> &PgPool {
341 &self.pool
342 }
343
344 async fn migrate(&self) -> Result<()> {
345 for statement in sanitise_schema(SCHEMA) {
352 sqlx::query(&statement).execute(&self.pool).await?;
353 }
354 sqlx::raw_sql(V0_13_2_RELOCATION_SQL)
357 .execute(&self.pool)
358 .await?;
359 sqlx::raw_sql(
367 r#"
368 DROP TRIGGER IF EXISTS workflow_runnable_notify ON workflow.workflows;
369 DROP TRIGGER IF EXISTS workflow_task_notify ON workflow.activities;
370 DROP FUNCTION IF EXISTS assay_notify_runnable();
371 DROP FUNCTION IF EXISTS assay_notify_task();
372 "#,
373 )
374 .execute(&self.pool)
375 .await?;
376 sqlx::query("ALTER TABLE workflow.events ADD COLUMN IF NOT EXISTS activity_id BIGINT")
377 .execute(&self.pool)
378 .await?;
379 sqlx::query(
380 "CREATE INDEX IF NOT EXISTS idx_wf_events_activity ON workflow.events(activity_id)",
381 )
382 .execute(&self.pool)
383 .await?;
384 self.backfill_event_activity_ids().await?;
385 Ok(())
386 }
387
388 async fn backfill_event_activity_ids(&self) -> Result<()> {
394 const BATCH: i64 = 500;
395 loop {
396 let rows: Vec<(i64, Option<String>)> = sqlx::query_as(
397 "SELECT id, payload FROM workflow.events
398 WHERE activity_id IS NULL
399 AND event_type IN ('ActivityCompleted', 'ActivityFailed')
400 LIMIT $1",
401 )
402 .bind(BATCH)
403 .fetch_all(&self.pool)
404 .await?;
405 if rows.is_empty() {
406 return Ok(());
407 }
408 let batch_len = rows.len() as i64;
409 for (id, payload) in rows {
410 sqlx::query("UPDATE workflow.events SET activity_id = $1 WHERE id = $2")
411 .bind(payload_activity_id(payload.as_deref()))
412 .bind(id)
413 .execute(&self.pool)
414 .await?;
415 }
416 if batch_len < BATCH {
417 return Ok(());
418 }
419 }
420 }
421
422 pub async fn try_acquire_leader_lock(&self) -> Result<bool> {
425 let row: (bool,) = sqlx::query_as("SELECT pg_try_advisory_lock(1)")
426 .fetch_one(&self.pool)
427 .await?;
428 Ok(row.0)
429 }
430}
431
432impl WorkflowStore for PostgresStore {
433 async fn create_namespace(&self, name: &str) -> Result<()> {
436 sqlx::query("INSERT INTO workflow.namespaces (name, created_at) VALUES ($1, EXTRACT(EPOCH FROM NOW()))")
437 .bind(name)
438 .execute(&self.pool)
439 .await?;
440 Ok(())
441 }
442
443 async fn list_namespaces(&self) -> Result<Vec<crate::store::NamespaceRecord>> {
444 let rows = sqlx::query_as::<_, (String, f64)>(
445 "SELECT name, created_at FROM workflow.namespaces ORDER BY name",
446 )
447 .fetch_all(&self.pool)
448 .await?;
449 Ok(rows
450 .into_iter()
451 .map(|(name, created_at)| crate::store::NamespaceRecord { name, created_at })
452 .collect())
453 }
454
455 async fn delete_namespace(&self, name: &str) -> Result<bool> {
456 let res = sqlx::query("DELETE FROM workflow.namespaces WHERE name = $1 AND name != 'main'")
457 .bind(name)
458 .execute(&self.pool)
459 .await?;
460 Ok(res.rows_affected() > 0)
461 }
462
463 async fn get_namespace_stats(&self, namespace: &str) -> Result<crate::store::NamespaceStats> {
464 let total: (i64,) =
465 sqlx::query_as("SELECT COUNT(*) FROM workflow.workflows WHERE namespace = $1")
466 .bind(namespace)
467 .fetch_one(&self.pool)
468 .await?;
469 let running: (i64,) = sqlx::query_as(
470 "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = $1 AND status = 'RUNNING'",
471 )
472 .bind(namespace)
473 .fetch_one(&self.pool)
474 .await?;
475 let pending: (i64,) = sqlx::query_as(
476 "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = $1 AND status = 'PENDING'",
477 )
478 .bind(namespace)
479 .fetch_one(&self.pool)
480 .await?;
481 let completed: (i64,) = sqlx::query_as(
482 "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = $1 AND status = 'COMPLETED'",
483 )
484 .bind(namespace)
485 .fetch_one(&self.pool)
486 .await?;
487 let failed: (i64,) = sqlx::query_as(
488 "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = $1 AND status = 'FAILED'",
489 )
490 .bind(namespace)
491 .fetch_one(&self.pool)
492 .await?;
493 let schedules: (i64,) =
494 sqlx::query_as("SELECT COUNT(*) FROM workflow.schedules WHERE namespace = $1")
495 .bind(namespace)
496 .fetch_one(&self.pool)
497 .await?;
498 let workers: (i64,) =
499 sqlx::query_as("SELECT COUNT(*) FROM workflow.workers WHERE namespace = $1")
500 .bind(namespace)
501 .fetch_one(&self.pool)
502 .await?;
503
504 Ok(crate::store::NamespaceStats {
505 namespace: namespace.to_string(),
506 total_workflows: total.0,
507 running: running.0,
508 pending: pending.0,
509 completed: completed.0,
510 failed: failed.0,
511 schedules: schedules.0,
512 workers: workers.0,
513 })
514 }
515
516 async fn create_workflow(&self, wf: &WorkflowRecord) -> Result<()> {
519 sqlx::query(
520 "INSERT INTO workflow.workflows (id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at)
521 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)",
522 )
523 .bind(&wf.id)
524 .bind(&wf.namespace)
525 .bind(&wf.run_id)
526 .bind(&wf.workflow_type)
527 .bind(&wf.task_queue)
528 .bind(&wf.status)
529 .bind(&wf.input)
530 .bind(&wf.result)
531 .bind(&wf.error)
532 .bind(&wf.parent_id)
533 .bind(&wf.claimed_by)
534 .bind(&wf.search_attributes)
535 .bind(wf.archived_at)
536 .bind(&wf.archive_uri)
537 .bind(wf.created_at)
538 .bind(wf.updated_at)
539 .bind(wf.completed_at)
540 .execute(&self.pool)
541 .await?;
542 Ok(())
543 }
544
545 async fn get_workflow(&self, id: &str) -> Result<Option<WorkflowRecord>> {
546 let row = sqlx::query_as::<_, PgWorkflowRow>(
547 "SELECT id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at FROM workflow.workflows WHERE id = $1",
548 )
549 .bind(id)
550 .fetch_optional(&self.pool)
551 .await?;
552 Ok(row.map(Into::into))
553 }
554
555 async fn list_workflows(
556 &self,
557 namespace: &str,
558 status: Option<WorkflowStatus>,
559 workflow_type: Option<&str>,
560 search_attrs_filter: Option<&str>,
561 limit: i64,
562 offset: i64,
563 ) -> Result<Vec<WorkflowRecord>> {
564 let status_str = status.map(|s| s.to_string());
565
566 let filter_pairs: Vec<(String, serde_json::Value)> = search_attrs_filter
567 .and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
568 .and_then(|v| v.as_object().cloned())
569 .map(|m| m.into_iter().collect())
570 .unwrap_or_default();
571
572 let mut sql = String::from(
573 "SELECT id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at
574 FROM workflow.workflows
575 WHERE namespace = $1
576 AND ($2::TEXT IS NULL OR status = $2)
577 AND ($3::TEXT IS NULL OR workflow_type = $3)",
578 );
579 let mut idx = 4usize;
581 for _ in &filter_pairs {
582 sql.push_str(&format!(
583 " AND (search_attributes::jsonb)->>${} = ${}",
584 idx,
585 idx + 1
586 ));
587 idx += 2;
588 }
589 sql.push_str(&format!(
590 " ORDER BY created_at DESC LIMIT ${} OFFSET ${}",
591 idx,
592 idx + 1
593 ));
594
595 let mut q = sqlx::query_as::<_, PgWorkflowRow>(&sql)
596 .bind(namespace)
597 .bind(&status_str)
598 .bind(workflow_type);
599 for (key, value) in &filter_pairs {
600 q = q.bind(key.clone());
601 let as_text = match value {
603 serde_json::Value::String(s) => s.clone(),
604 other => other.to_string(),
605 };
606 q = q.bind(as_text);
607 }
608 let rows = q.bind(limit).bind(offset).fetch_all(&self.pool).await?;
609 Ok(rows.into_iter().map(Into::into).collect())
610 }
611
612 async fn update_workflow_status(
613 &self,
614 id: &str,
615 status: WorkflowStatus,
616 result: Option<&str>,
617 error: Option<&str>,
618 ) -> Result<()> {
619 let now = timestamp_now();
620 let completed_at = if status.is_terminal() {
621 Some(now)
622 } else {
623 None
624 };
625 sqlx::query(
626 "UPDATE workflow.workflows SET status = $1, result = COALESCE($2, result), error = COALESCE($3, error), updated_at = $4, completed_at = COALESCE($5, completed_at) WHERE id = $6",
627 )
628 .bind(status.to_string())
629 .bind(result)
630 .bind(error)
631 .bind(now)
632 .bind(completed_at)
633 .bind(id)
634 .execute(&self.pool)
635 .await?;
636 Ok(())
637 }
638
639 async fn claim_workflow(&self, id: &str, worker_id: &str) -> Result<bool> {
640 let res = sqlx::query(
641 "UPDATE workflow.workflows SET claimed_by = $1, status = 'RUNNING', updated_at = $2 WHERE id = $3 AND claimed_by IS NULL",
642 )
643 .bind(worker_id)
644 .bind(timestamp_now())
645 .bind(id)
646 .execute(&self.pool)
647 .await?;
648 Ok(res.rows_affected() > 0)
649 }
650
651 async fn mark_workflow_dispatchable(&self, workflow_id: &str) -> Result<()> {
652 sqlx::query("UPDATE workflow.workflows SET needs_dispatch = TRUE WHERE id = $1")
653 .bind(workflow_id)
654 .execute(&self.pool)
655 .await?;
656 Ok(())
657 }
658
659 async fn claim_workflow_task(
660 &self,
661 task_queue: &str,
662 worker_id: &str,
663 ) -> Result<Option<WorkflowRecord>> {
664 let now = timestamp_now();
665 let row = sqlx::query_as::<_, PgWorkflowRow>(
668 "UPDATE workflow.workflows
669 SET dispatch_claimed_by = $1, dispatch_last_heartbeat = $2, needs_dispatch = FALSE
670 WHERE id = (
671 SELECT id FROM workflow.workflows
672 WHERE task_queue = $3
673 AND needs_dispatch = TRUE
674 AND dispatch_claimed_by IS NULL
675 AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
676 ORDER BY updated_at ASC
677 FOR UPDATE SKIP LOCKED
678 LIMIT 1
679 )
680 RETURNING id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at",
681 )
682 .bind(worker_id)
683 .bind(now)
684 .bind(task_queue)
685 .fetch_optional(&self.pool)
686 .await?;
687 Ok(row.map(Into::into))
688 }
689
690 async fn release_workflow_task(&self, workflow_id: &str, worker_id: &str) -> Result<()> {
691 sqlx::query(
692 "UPDATE workflow.workflows
693 SET dispatch_claimed_by = NULL, dispatch_last_heartbeat = NULL
694 WHERE id = $1 AND dispatch_claimed_by = $2",
695 )
696 .bind(workflow_id)
697 .bind(worker_id)
698 .execute(&self.pool)
699 .await?;
700 Ok(())
701 }
702
703 async fn release_stale_dispatch_leases(&self, now: f64, timeout_secs: f64) -> Result<u64> {
704 let res = sqlx::query(
705 "UPDATE workflow.workflows
706 SET dispatch_claimed_by = NULL,
707 dispatch_last_heartbeat = NULL,
708 needs_dispatch = TRUE
709 WHERE dispatch_claimed_by IS NOT NULL
710 AND ($1 - dispatch_last_heartbeat) > $2
711 AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')",
712 )
713 .bind(now)
714 .bind(timeout_secs)
715 .execute(&self.pool)
716 .await?;
717 Ok(res.rows_affected())
718 }
719
720 async fn append_event(&self, ev: &WorkflowEvent) -> Result<i64> {
723 let row: (i64,) = sqlx::query_as(
724 "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, timestamp) VALUES ($1, $2, $3, $4, $5) RETURNING id",
725 )
726 .bind(&ev.workflow_id)
727 .bind(ev.seq)
728 .bind(&ev.event_type)
729 .bind(&ev.payload)
730 .bind(ev.timestamp)
731 .fetch_one(&self.pool)
732 .await?;
733 Ok(row.0)
734 }
735
736 async fn list_events(&self, workflow_id: &str) -> Result<Vec<WorkflowEvent>> {
737 let rows = sqlx::query_as::<_, PgEventRow>(
738 "SELECT id, workflow_id, seq, event_type, payload, timestamp FROM workflow.events WHERE workflow_id = $1 ORDER BY seq ASC",
739 )
740 .bind(workflow_id)
741 .fetch_all(&self.pool)
742 .await?;
743 Ok(rows.into_iter().map(Into::into).collect())
744 }
745
746 async fn list_events_page(
747 &self,
748 workflow_id: &str,
749 cursor: Option<i32>,
750 limit: i64,
751 descending: bool,
752 ) -> Result<Vec<WorkflowEvent>> {
753 let limit = limit.clamp(0, 1_000);
754 if limit == 0 {
755 return Ok(Vec::new());
756 }
757 let rows = if descending {
758 sqlx::query_as::<_, PgEventRow>(
759 "SELECT id, workflow_id, seq, event_type, payload, timestamp
760 FROM workflow.events
761 WHERE workflow_id = $1 AND ($2::INTEGER IS NULL OR seq < $2)
762 ORDER BY seq DESC LIMIT $3",
763 )
764 .bind(workflow_id)
765 .bind(cursor)
766 .bind(limit)
767 .fetch_all(&self.pool)
768 .await?
769 } else {
770 sqlx::query_as::<_, PgEventRow>(
771 "SELECT id, workflow_id, seq, event_type, payload, timestamp
772 FROM workflow.events
773 WHERE workflow_id = $1 AND ($2::INTEGER IS NULL OR seq > $2)
774 ORDER BY seq ASC LIMIT $3",
775 )
776 .bind(workflow_id)
777 .bind(cursor)
778 .bind(limit)
779 .fetch_all(&self.pool)
780 .await?
781 };
782 Ok(rows.into_iter().map(Into::into).collect())
783 }
784
785 async fn get_event_count(&self, workflow_id: &str) -> Result<i64> {
786 let row: (i64,) =
787 sqlx::query_as("SELECT COUNT(*) FROM workflow.events WHERE workflow_id = $1")
788 .bind(workflow_id)
789 .fetch_one(&self.pool)
790 .await?;
791 Ok(row.0)
792 }
793
794 async fn create_activity(&self, act: &WorkflowActivity) -> Result<i64> {
797 let row: (i64,) = sqlx::query_as(
798 "INSERT INTO workflow.activities (workflow_id, seq, name, task_queue, input, status, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, scheduled_at)
799 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) RETURNING id",
800 )
801 .bind(&act.workflow_id)
802 .bind(act.seq)
803 .bind(&act.name)
804 .bind(&act.task_queue)
805 .bind(&act.input)
806 .bind(&act.status)
807 .bind(act.attempt)
808 .bind(act.max_attempts)
809 .bind(act.initial_interval_secs)
810 .bind(act.backoff_coefficient)
811 .bind(act.start_to_close_secs)
812 .bind(act.heartbeat_timeout_secs)
813 .bind(act.scheduled_at)
814 .fetch_one(&self.pool)
815 .await?;
816 Ok(row.0)
817 }
818
819 async fn get_activity(&self, id: i64) -> Result<Option<WorkflowActivity>> {
820 let row = sqlx::query_as::<_, PgActivityRow>(
821 "SELECT id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat
822 FROM workflow.activities WHERE id = $1",
823 )
824 .bind(id)
825 .fetch_optional(&self.pool)
826 .await?;
827 Ok(row.map(Into::into))
828 }
829
830 async fn get_activity_by_workflow_seq(
831 &self,
832 workflow_id: &str,
833 seq: i32,
834 ) -> Result<Option<WorkflowActivity>> {
835 let row = sqlx::query_as::<_, PgActivityRow>(
836 "SELECT id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat
837 FROM workflow.activities WHERE workflow_id = $1 AND seq = $2",
838 )
839 .bind(workflow_id)
840 .bind(seq)
841 .fetch_optional(&self.pool)
842 .await?;
843 Ok(row.map(Into::into))
844 }
845
846 async fn claim_activity(
847 &self,
848 task_queue: &str,
849 worker_id: &str,
850 ) -> Result<Option<WorkflowActivity>> {
851 let now = timestamp_now();
852 let row = sqlx::query_as::<_, PgActivityRow>(
855 "UPDATE workflow.activities SET status = 'RUNNING', claimed_by = $1, started_at = $2
856 WHERE id = (
857 SELECT id FROM workflow.activities
858 WHERE task_queue = $3 AND status = 'PENDING'
859 ORDER BY scheduled_at ASC
860 FOR UPDATE SKIP LOCKED
861 LIMIT 1
862 )
863 RETURNING id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat",
864 )
865 .bind(worker_id)
866 .bind(now)
867 .bind(task_queue)
868 .fetch_optional(&self.pool)
869 .await?;
870 Ok(row.map(Into::into))
871 }
872
873 async fn requeue_activity_for_retry(
874 &self,
875 id: i64,
876 next_attempt: i32,
877 next_scheduled_at: f64,
878 ) -> Result<()> {
879 sqlx::query(
880 "UPDATE workflow.activities
881 SET status = 'PENDING', attempt = $1, scheduled_at = $2,
882 claimed_by = NULL, started_at = NULL, last_heartbeat = NULL,
883 error = NULL
884 WHERE id = $3",
885 )
886 .bind(next_attempt)
887 .bind(next_scheduled_at)
888 .bind(id)
889 .execute(&self.pool)
890 .await?;
891 Ok(())
892 }
893
894 async fn retry_failed_activity(
895 &self,
896 workflow_id: &str,
897 requested_by: &str,
898 reason: &str,
899 requested_at: f64,
900 ) -> Result<RetryFailedActivityResult> {
901 let mut tx = self.pool.begin().await?;
902 let workflow: Option<(String, Option<String>, Option<f64>)> = sqlx::query_as(
903 "SELECT status, parent_id, archived_at FROM workflow.workflows WHERE id = $1 FOR UPDATE",
904 )
905 .bind(workflow_id)
906 .fetch_optional(&mut *tx)
907 .await?;
908 let Some((status, parent_id, archived_at)) = workflow else {
909 return Ok(RetryFailedActivityResult::NotFound);
910 };
911 if let Some(denial) = retry_denial(status, parent_id, archived_at) {
912 return Ok(denial);
913 }
914
915 let failed = sqlx::query_as::<_, PgActivityRow>(RETRY_ACTIVITY_SELECT)
916 .bind(workflow_id)
917 .fetch_optional(&mut *tx)
918 .await?;
919 let Some(failed) = failed else {
920 return Ok(RetryFailedActivityResult::NoFailedActivity);
921 };
922 let failed_event_seq: (i32,) = sqlx::query_as(
923 "SELECT seq FROM workflow.events
924 WHERE workflow_id = $1 AND event_type = 'ActivityFailed'
925 ORDER BY seq DESC LIMIT 1",
926 )
927 .bind(workflow_id)
928 .fetch_one(&mut *tx)
929 .await?;
930 let invalidated =
931 sqlx::query("DELETE FROM workflow.activities WHERE workflow_id = $1 AND seq > $2")
932 .bind(workflow_id)
933 .bind(failed.seq)
934 .execute(&mut *tx)
935 .await?
936 .rows_affected();
937 let activity = sqlx::query_as::<_, PgActivityRow>(RETRY_ACTIVITY_UPDATE)
938 .bind(requested_at)
939 .bind(failed.id)
940 .fetch_one(&mut *tx)
941 .await?;
942 sqlx::query("UPDATE workflow.events SET activity_id = $1 WHERE activity_id = $2")
945 .bind(NOT_A_SETTLEMENT)
946 .bind(failed.id)
947 .execute(&mut *tx)
948 .await?;
949 sqlx::query(
950 "UPDATE workflow.workflows
951 SET status = 'WAITING', result = NULL, error = NULL, completed_at = NULL,
952 updated_at = $1, needs_dispatch = FALSE, dispatch_claimed_by = NULL,
953 dispatch_last_heartbeat = NULL
954 WHERE id = $2",
955 )
956 .bind(requested_at)
957 .bind(workflow_id)
958 .execute(&mut *tx)
959 .await?;
960 let event_seq: (i32,) = sqlx::query_as(
961 "SELECT COALESCE(MAX(seq), 0) + 1 FROM workflow.events WHERE workflow_id = $1",
962 )
963 .bind(workflow_id)
964 .fetch_one(&mut *tx)
965 .await?;
966 let payload = RetryEvent {
967 activity_id: failed.id,
968 activity_seq: failed.seq,
969 activity_name: &failed.name,
970 failed_event_seq: failed_event_seq.0,
971 requested_by,
972 reason,
973 invalidated_activities: invalidated,
974 }
975 .payload();
976 sqlx::query(
977 "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, timestamp)
978 VALUES ($1, $2, 'ActivityRetryRequested', $3, $4)",
979 )
980 .bind(workflow_id)
981 .bind(event_seq.0)
982 .bind(payload.to_string())
983 .bind(requested_at)
984 .execute(&mut *tx)
985 .await?;
986 tx.commit().await?;
987 Ok(RetryFailedActivityResult::Retried(Box::new(
988 RetriedActivity {
989 activity: activity.into(),
990 invalidated_activities: invalidated,
991 },
992 )))
993 }
994
995 async fn complete_activity(
996 &self,
997 id: i64,
998 result: Option<&str>,
999 error: Option<&str>,
1000 failed: bool,
1001 ) -> Result<()> {
1002 let status = if failed { "FAILED" } else { "COMPLETED" };
1003 sqlx::query(
1004 "UPDATE workflow.activities SET status = $1, result = $2, error = $3, completed_at = $4 WHERE id = $5",
1005 )
1006 .bind(status)
1007 .bind(result)
1008 .bind(error)
1009 .bind(timestamp_now())
1010 .bind(id)
1011 .execute(&self.pool)
1012 .await?;
1013 Ok(())
1014 }
1015
1016 async fn settle_activity(&self, settlement: &ActivitySettlement<'_>) -> Result<SettleOutcome> {
1017 let mut tx = self.pool.begin().await?;
1018 let current: Option<(String,)> =
1019 sqlx::query_as("SELECT status FROM workflow.activities WHERE id = $1 FOR UPDATE")
1020 .bind(settlement.activity_id)
1021 .fetch_optional(&mut *tx)
1022 .await?;
1023 let Some((status,)) = current else {
1024 return Ok(SettleOutcome::Unknown);
1025 };
1026 let settled = matches!(status.as_str(), "COMPLETED" | "FAILED");
1027 let event_id: Option<(i64,)> = sqlx::query_as(
1028 "SELECT id FROM workflow.events
1029 WHERE workflow_id = $1 AND activity_id = $2
1030 AND event_type IN ('ActivityCompleted', 'ActivityFailed')
1031 LIMIT 1",
1032 )
1033 .bind(settlement.workflow_id)
1034 .bind(settlement.activity_id)
1035 .fetch_optional(&mut *tx)
1036 .await?;
1037
1038 if !settled {
1039 sqlx::query(
1040 "UPDATE workflow.activities
1041 SET status = $1, result = $2, error = $3, completed_at = $4
1042 WHERE id = $5",
1043 )
1044 .bind(if settlement.failed {
1045 "FAILED"
1046 } else {
1047 "COMPLETED"
1048 })
1049 .bind(settlement.result)
1050 .bind(settlement.error)
1051 .bind(settlement.now)
1052 .bind(settlement.activity_id)
1053 .execute(&mut *tx)
1054 .await?;
1055 }
1056 if !settled || event_id.is_none() {
1060 let seq: (i32,) = sqlx::query_as(
1061 "SELECT COALESCE(MAX(seq), 0) + 1 FROM workflow.events WHERE workflow_id = $1",
1062 )
1063 .bind(settlement.workflow_id)
1064 .fetch_one(&mut *tx)
1065 .await?;
1066 sqlx::query(
1067 "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, activity_id, timestamp)
1068 VALUES ($1, $2, $3, $4, $5, $6)",
1069 )
1070 .bind(settlement.workflow_id)
1071 .bind(seq.0)
1072 .bind(settlement.event_type)
1073 .bind(settlement.payload)
1074 .bind(settlement.activity_id)
1075 .bind(settlement.now)
1076 .execute(&mut *tx)
1077 .await?;
1078 }
1079 sqlx::query("UPDATE workflow.workflows SET needs_dispatch = TRUE WHERE id = $1")
1080 .bind(settlement.workflow_id)
1081 .execute(&mut *tx)
1082 .await?;
1083 tx.commit().await?;
1084 Ok(settle_outcome(settled, event_id.is_some()))
1085 }
1086
1087 async fn list_unsettled_activities(&self, limit: i64) -> Result<Vec<WorkflowActivity>> {
1088 let rows = sqlx::query_as::<_, PgActivityRow>(UNSETTLED_ACTIVITY_SELECT)
1089 .bind(limit)
1090 .fetch_all(&self.pool)
1091 .await?;
1092 Ok(rows.into_iter().map(Into::into).collect())
1093 }
1094
1095 async fn heartbeat_activity(&self, id: i64, _details: Option<&str>) -> Result<()> {
1096 sqlx::query("UPDATE workflow.activities SET last_heartbeat = $1 WHERE id = $2")
1097 .bind(timestamp_now())
1098 .bind(id)
1099 .execute(&self.pool)
1100 .await?;
1101 Ok(())
1102 }
1103
1104 async fn get_timed_out_activities(&self, now: f64) -> Result<Vec<WorkflowActivity>> {
1105 let rows = sqlx::query_as::<_, PgActivityRow>(
1106 "SELECT id, workflow_id, seq, name, task_queue, input, status, result, error, attempt, max_attempts, initial_interval_secs, backoff_coefficient, start_to_close_secs, heartbeat_timeout_secs, claimed_by, scheduled_at, started_at, completed_at, last_heartbeat
1107 FROM workflow.activities
1108 WHERE status = 'RUNNING'
1109 AND heartbeat_timeout_secs IS NOT NULL
1110 AND ($1 - COALESCE(last_heartbeat, started_at)) > heartbeat_timeout_secs",
1111 )
1112 .bind(now)
1113 .fetch_all(&self.pool)
1114 .await?;
1115 Ok(rows.into_iter().map(Into::into).collect())
1116 }
1117
1118 async fn create_timer(&self, timer: &WorkflowTimer) -> Result<i64> {
1121 let inserted: Option<(i64,)> = sqlx::query_as(
1124 "INSERT INTO workflow.timers (workflow_id, seq, fire_at, fired)
1125 VALUES ($1, $2, $3, FALSE)
1126 ON CONFLICT (workflow_id, seq) DO NOTHING
1127 RETURNING id",
1128 )
1129 .bind(&timer.workflow_id)
1130 .bind(timer.seq)
1131 .bind(timer.fire_at)
1132 .fetch_optional(&self.pool)
1133 .await?;
1134
1135 if let Some((id,)) = inserted {
1136 return Ok(id);
1137 }
1138
1139 let (id,): (i64,) =
1141 sqlx::query_as("SELECT id FROM workflow.timers WHERE workflow_id = $1 AND seq = $2")
1142 .bind(&timer.workflow_id)
1143 .bind(timer.seq)
1144 .fetch_one(&self.pool)
1145 .await?;
1146 Ok(id)
1147 }
1148
1149 async fn cancel_pending_activities(&self, workflow_id: &str) -> Result<u64> {
1150 let res = sqlx::query(
1151 "UPDATE workflow.activities SET status = 'CANCELLED', completed_at = $1
1152 WHERE workflow_id = $2 AND status = 'PENDING'",
1153 )
1154 .bind(timestamp_now())
1155 .bind(workflow_id)
1156 .execute(&self.pool)
1157 .await?;
1158 Ok(res.rows_affected())
1159 }
1160
1161 async fn cancel_pending_timers(&self, workflow_id: &str) -> Result<u64> {
1162 let res = sqlx::query(
1163 "UPDATE workflow.timers SET fired = TRUE
1164 WHERE workflow_id = $1 AND fired = FALSE",
1165 )
1166 .bind(workflow_id)
1167 .execute(&self.pool)
1168 .await?;
1169 Ok(res.rows_affected())
1170 }
1171
1172 async fn get_timer_by_workflow_seq(
1173 &self,
1174 workflow_id: &str,
1175 seq: i32,
1176 ) -> Result<Option<WorkflowTimer>> {
1177 let row = sqlx::query_as::<_, PgTimerRow>(
1178 "SELECT id, workflow_id, seq, fire_at, fired
1179 FROM workflow.timers WHERE workflow_id = $1 AND seq = $2",
1180 )
1181 .bind(workflow_id)
1182 .bind(seq)
1183 .fetch_optional(&self.pool)
1184 .await?;
1185 Ok(row.map(Into::into))
1186 }
1187
1188 async fn fire_due_timers(&self, now: f64) -> Result<Vec<WorkflowTimer>> {
1189 let rows = sqlx::query_as::<_, PgTimerRow>(
1190 "UPDATE workflow.timers SET fired = TRUE
1191 WHERE fired = FALSE AND fire_at <= $1
1192 RETURNING id, workflow_id, seq, fire_at, fired",
1193 )
1194 .bind(now)
1195 .fetch_all(&self.pool)
1196 .await?;
1197 Ok(rows.into_iter().map(Into::into).collect())
1198 }
1199
1200 async fn send_signal(&self, sig: &WorkflowSignal) -> Result<i64> {
1203 let row: (i64,) = sqlx::query_as(
1204 "INSERT INTO workflow.signals (workflow_id, name, payload, consumed, received_at) VALUES ($1, $2, $3, FALSE, $4) RETURNING id",
1205 )
1206 .bind(&sig.workflow_id)
1207 .bind(&sig.name)
1208 .bind(&sig.payload)
1209 .bind(sig.received_at)
1210 .fetch_one(&self.pool)
1211 .await?;
1212 Ok(row.0)
1213 }
1214
1215 async fn deliver_signal(&self, sig: &WorkflowSignal, payload_json: &str) -> Result<i64> {
1216 let mut tx = self.pool.begin().await?;
1217 let row: (i64,) = sqlx::query_as(
1218 "INSERT INTO workflow.signals (workflow_id, name, payload, consumed, received_at) VALUES ($1, $2, $3, FALSE, $4) RETURNING id",
1219 )
1220 .bind(&sig.workflow_id)
1221 .bind(&sig.name)
1222 .bind(&sig.payload)
1223 .bind(sig.received_at)
1224 .fetch_one(&mut *tx)
1225 .await?;
1226 let seq: (i32,) = sqlx::query_as(
1227 "SELECT COALESCE(MAX(seq), 0) + 1 FROM workflow.events WHERE workflow_id = $1",
1228 )
1229 .bind(&sig.workflow_id)
1230 .fetch_one(&mut *tx)
1231 .await?;
1232 sqlx::query(
1233 "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, timestamp)
1234 VALUES ($1, $2, 'SignalReceived', $3, $4)",
1235 )
1236 .bind(&sig.workflow_id)
1237 .bind(seq.0)
1238 .bind(payload_json)
1239 .bind(sig.received_at)
1240 .execute(&mut *tx)
1241 .await?;
1242 sqlx::query("UPDATE workflow.workflows SET needs_dispatch = TRUE WHERE id = $1")
1243 .bind(&sig.workflow_id)
1244 .execute(&mut *tx)
1245 .await?;
1246 tx.commit().await?;
1247 Ok(row.0)
1248 }
1249
1250 async fn consume_signals(&self, workflow_id: &str, name: &str) -> Result<Vec<WorkflowSignal>> {
1251 let rows = sqlx::query_as::<_, PgSignalRow>(
1252 "UPDATE workflow.signals SET consumed = TRUE
1253 WHERE workflow_id = $1 AND name = $2 AND consumed = FALSE
1254 RETURNING id, workflow_id, name, payload, consumed, received_at",
1255 )
1256 .bind(workflow_id)
1257 .bind(name)
1258 .fetch_all(&self.pool)
1259 .await?;
1260 Ok(rows.into_iter().map(Into::into).collect())
1261 }
1262
1263 async fn create_schedule(&self, sched: &WorkflowSchedule) -> Result<()> {
1266 sqlx::query(
1267 "INSERT INTO workflow.schedules (namespace, name, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at)
1268 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)",
1269 )
1270 .bind(&sched.namespace)
1271 .bind(&sched.name)
1272 .bind(&sched.workflow_type)
1273 .bind(&sched.cron_expr)
1274 .bind(&sched.timezone)
1275 .bind(&sched.input)
1276 .bind(&sched.task_queue)
1277 .bind(&sched.overlap_policy)
1278 .bind(sched.paused)
1279 .bind(sched.last_run_at)
1280 .bind(crate::scheduler::seed_next_run(sched))
1281 .bind(&sched.last_workflow_id)
1282 .bind(sched.created_at)
1283 .execute(&self.pool)
1284 .await?;
1285 Ok(())
1286 }
1287
1288 async fn get_schedule(&self, namespace: &str, name: &str) -> Result<Option<WorkflowSchedule>> {
1289 let row = sqlx::query_as::<_, PgScheduleRow>(
1290 "SELECT namespace, name, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at FROM workflow.schedules WHERE namespace = $1 AND name = $2",
1291 )
1292 .bind(namespace)
1293 .bind(name)
1294 .fetch_optional(&self.pool)
1295 .await?;
1296 Ok(row.map(Into::into))
1297 }
1298
1299 async fn list_schedules(&self, namespace: &str) -> Result<Vec<WorkflowSchedule>> {
1300 let rows = sqlx::query_as::<_, PgScheduleRow>(
1301 "SELECT namespace, name, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at FROM workflow.schedules WHERE namespace = $1 ORDER BY name",
1302 )
1303 .bind(namespace)
1304 .fetch_all(&self.pool)
1305 .await?;
1306 Ok(rows.into_iter().map(Into::into).collect())
1307 }
1308
1309 async fn update_schedule_last_run(
1310 &self,
1311 namespace: &str,
1312 name: &str,
1313 last_run_at: f64,
1314 next_run_at: f64,
1315 workflow_id: &str,
1316 ) -> Result<()> {
1317 sqlx::query(
1318 "UPDATE workflow.schedules SET last_run_at = $1, next_run_at = $2, last_workflow_id = $3 WHERE namespace = $4 AND name = $5",
1319 )
1320 .bind(last_run_at)
1321 .bind(next_run_at)
1322 .bind(workflow_id)
1323 .bind(namespace)
1324 .bind(name)
1325 .execute(&self.pool)
1326 .await?;
1327 Ok(())
1328 }
1329
1330 async fn delete_schedule(&self, namespace: &str, name: &str) -> Result<bool> {
1331 let res = sqlx::query("DELETE FROM workflow.schedules WHERE namespace = $1 AND name = $2")
1332 .bind(namespace)
1333 .bind(name)
1334 .execute(&self.pool)
1335 .await?;
1336 Ok(res.rows_affected() > 0)
1337 }
1338
1339 async fn list_archivable_workflows(
1340 &self,
1341 cutoff: f64,
1342 limit: i64,
1343 ) -> Result<Vec<WorkflowRecord>> {
1344 let rows = sqlx::query_as::<_, PgWorkflowRow>(
1345 "SELECT id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at
1346 FROM workflow.workflows
1347 WHERE status IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
1348 AND completed_at IS NOT NULL
1349 AND completed_at < $1
1350 AND archived_at IS NULL
1351 ORDER BY completed_at ASC
1352 LIMIT $2",
1353 )
1354 .bind(cutoff)
1355 .bind(limit)
1356 .fetch_all(&self.pool)
1357 .await?;
1358 Ok(rows.into_iter().map(Into::into).collect())
1359 }
1360
1361 async fn mark_archived_and_purge(
1362 &self,
1363 workflow_id: &str,
1364 archive_uri: &str,
1365 archived_at: f64,
1366 ) -> Result<()> {
1367 let mut tx = self.pool.begin().await?;
1368 sqlx::query("DELETE FROM workflow.events WHERE workflow_id = $1")
1369 .bind(workflow_id)
1370 .execute(&mut *tx)
1371 .await?;
1372 sqlx::query("DELETE FROM workflow.activities WHERE workflow_id = $1")
1373 .bind(workflow_id)
1374 .execute(&mut *tx)
1375 .await?;
1376 sqlx::query("DELETE FROM workflow.timers WHERE workflow_id = $1")
1377 .bind(workflow_id)
1378 .execute(&mut *tx)
1379 .await?;
1380 sqlx::query("DELETE FROM workflow.signals WHERE workflow_id = $1")
1381 .bind(workflow_id)
1382 .execute(&mut *tx)
1383 .await?;
1384 sqlx::query("DELETE FROM workflow.snapshots WHERE workflow_id = $1")
1385 .bind(workflow_id)
1386 .execute(&mut *tx)
1387 .await?;
1388 sqlx::query(
1389 "UPDATE workflow.workflows SET archived_at = $1, archive_uri = $2 WHERE id = $3",
1390 )
1391 .bind(archived_at)
1392 .bind(archive_uri)
1393 .bind(workflow_id)
1394 .execute(&mut *tx)
1395 .await?;
1396 tx.commit().await?;
1397 Ok(())
1398 }
1399
1400 async fn upsert_search_attributes(&self, workflow_id: &str, patch_json: &str) -> Result<()> {
1401 let current: Option<(Option<String>,)> =
1402 sqlx::query_as("SELECT search_attributes FROM workflow.workflows WHERE id = $1")
1403 .bind(workflow_id)
1404 .fetch_optional(&self.pool)
1405 .await?;
1406 let merged = crate::store::sqlite::merge_search_attrs(
1407 current.and_then(|(s,)| s).as_deref(),
1408 patch_json,
1409 )?;
1410 sqlx::query("UPDATE workflow.workflows SET search_attributes = $1 WHERE id = $2")
1411 .bind(merged)
1412 .bind(workflow_id)
1413 .execute(&self.pool)
1414 .await?;
1415 Ok(())
1416 }
1417
1418 async fn update_schedule(
1419 &self,
1420 namespace: &str,
1421 name: &str,
1422 patch: &SchedulePatch,
1423 ) -> Result<Option<WorkflowSchedule>> {
1424 let mut sets: Vec<String> = Vec::new();
1425 let mut idx = 1usize;
1426 if patch.cron_expr.is_some() {
1427 sets.push(format!("cron_expr = ${idx}"));
1428 idx += 1;
1429 }
1430 if patch.timezone.is_some() {
1431 sets.push(format!("timezone = ${idx}"));
1432 idx += 1;
1433 }
1434 if patch.input.is_some() {
1435 sets.push(format!("input = ${idx}"));
1436 idx += 1;
1437 }
1438 if patch.task_queue.is_some() {
1439 sets.push(format!("task_queue = ${idx}"));
1440 idx += 1;
1441 }
1442 if patch.overlap_policy.is_some() {
1443 sets.push(format!("overlap_policy = ${idx}"));
1444 idx += 1;
1445 }
1446 if sets.is_empty() {
1447 return self.get_schedule(namespace, name).await;
1448 }
1449 let sql = format!(
1450 "UPDATE workflow.schedules SET {} WHERE namespace = ${} AND name = ${}",
1451 sets.join(", "),
1452 idx,
1453 idx + 1
1454 );
1455 let mut q = sqlx::query(&sql);
1456 if let Some(ref v) = patch.cron_expr {
1457 q = q.bind(v);
1458 }
1459 if let Some(ref v) = patch.timezone {
1460 q = q.bind(v);
1461 }
1462 if let Some(ref v) = patch.input {
1463 q = q.bind(v.to_string());
1464 }
1465 if let Some(ref v) = patch.task_queue {
1466 q = q.bind(v);
1467 }
1468 if let Some(ref v) = patch.overlap_policy {
1469 q = q.bind(v);
1470 }
1471 let res = q.bind(namespace).bind(name).execute(&self.pool).await?;
1472 if res.rows_affected() == 0 {
1473 return Ok(None);
1474 }
1475 self.get_schedule(namespace, name).await
1476 }
1477
1478 async fn set_schedule_paused(
1479 &self,
1480 namespace: &str,
1481 name: &str,
1482 paused: bool,
1483 ) -> Result<Option<WorkflowSchedule>> {
1484 let res = sqlx::query(
1485 "UPDATE workflow.schedules SET paused = $1 WHERE namespace = $2 AND name = $3",
1486 )
1487 .bind(paused)
1488 .bind(namespace)
1489 .bind(name)
1490 .execute(&self.pool)
1491 .await?;
1492 if res.rows_affected() == 0 {
1493 return Ok(None);
1494 }
1495 self.get_schedule(namespace, name).await
1496 }
1497
1498 async fn register_worker(&self, w: &WorkflowWorker) -> Result<()> {
1501 sqlx::query(
1502 "INSERT INTO workflow.workers (id, namespace, identity, task_queue, workflows, activities, max_concurrent_workflows, max_concurrent_activities, active_tasks, last_heartbeat, registered_at)
1503 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
1504 ON CONFLICT (id) DO UPDATE SET last_heartbeat = EXCLUDED.last_heartbeat, identity = EXCLUDED.identity",
1505 )
1506 .bind(&w.id)
1507 .bind(&w.namespace)
1508 .bind(&w.identity)
1509 .bind(&w.task_queue)
1510 .bind(&w.workflows)
1511 .bind(&w.activities)
1512 .bind(w.max_concurrent_workflows)
1513 .bind(w.max_concurrent_activities)
1514 .bind(w.active_tasks)
1515 .bind(w.last_heartbeat)
1516 .bind(w.registered_at)
1517 .execute(&self.pool)
1518 .await?;
1519 Ok(())
1520 }
1521
1522 async fn heartbeat_worker(&self, id: &str, now: f64) -> Result<bool> {
1523 let res = sqlx::query("UPDATE workflow.workers SET last_heartbeat = $1 WHERE id = $2")
1524 .bind(now)
1525 .bind(id)
1526 .execute(&self.pool)
1527 .await?;
1528 Ok(res.rows_affected() > 0)
1529 }
1530
1531 async fn list_workers(&self, namespace: &str) -> Result<Vec<WorkflowWorker>> {
1532 let rows = sqlx::query_as::<_, PgWorkerRow>(
1533 "SELECT id, namespace, identity, task_queue, workflows, activities, max_concurrent_workflows, max_concurrent_activities, active_tasks, last_heartbeat, registered_at FROM workflow.workers WHERE namespace = $1 ORDER BY registered_at",
1534 )
1535 .bind(namespace)
1536 .fetch_all(&self.pool)
1537 .await?;
1538 Ok(rows.into_iter().map(Into::into).collect())
1539 }
1540
1541 async fn remove_dead_workers(&self, cutoff: f64) -> Result<Vec<String>> {
1542 let rows: Vec<(String,)> =
1543 sqlx::query_as("SELECT id FROM workflow.workers WHERE last_heartbeat < $1")
1544 .bind(cutoff)
1545 .fetch_all(&self.pool)
1546 .await?;
1547 let ids: Vec<String> = rows.into_iter().map(|r| r.0).collect();
1548 if !ids.is_empty() {
1549 sqlx::query("DELETE FROM workflow.workers WHERE last_heartbeat < $1")
1550 .bind(cutoff)
1551 .execute(&self.pool)
1552 .await?;
1553 }
1554 Ok(ids)
1555 }
1556
1557 async fn list_child_workflows(&self, parent_id: &str) -> Result<Vec<WorkflowRecord>> {
1560 let rows = sqlx::query_as::<_, PgWorkflowRow>(
1561 "SELECT id, namespace, run_id, workflow_type, task_queue, status, input, result, error, parent_id, claimed_by, search_attributes, archived_at, archive_uri, created_at, updated_at, completed_at
1562 FROM workflow.workflows WHERE parent_id = $1 ORDER BY created_at ASC",
1563 )
1564 .bind(parent_id)
1565 .fetch_all(&self.pool)
1566 .await?;
1567 Ok(rows.into_iter().map(Into::into).collect())
1568 }
1569
1570 async fn create_snapshot(
1573 &self,
1574 workflow_id: &str,
1575 event_seq: i32,
1576 state_json: &str,
1577 ) -> Result<()> {
1578 sqlx::query(
1579 "INSERT INTO workflow.snapshots (workflow_id, event_seq, state_json, created_at)
1580 VALUES ($1, $2, $3, $4)
1581 ON CONFLICT (workflow_id, event_seq) DO UPDATE SET state_json = EXCLUDED.state_json, created_at = EXCLUDED.created_at",
1582 )
1583 .bind(workflow_id)
1584 .bind(event_seq)
1585 .bind(state_json)
1586 .bind(timestamp_now())
1587 .execute(&self.pool)
1588 .await?;
1589 Ok(())
1590 }
1591
1592 async fn get_latest_snapshot(&self, workflow_id: &str) -> Result<Option<WorkflowSnapshot>> {
1593 let row = sqlx::query_as::<_, (String, i32, String, f64)>(
1594 "SELECT workflow_id, event_seq, state_json, created_at
1595 FROM workflow.snapshots WHERE workflow_id = $1
1596 ORDER BY event_seq DESC LIMIT 1",
1597 )
1598 .bind(workflow_id)
1599 .fetch_optional(&self.pool)
1600 .await?;
1601
1602 Ok(row.map(
1603 |(workflow_id, event_seq, state_json, created_at)| WorkflowSnapshot {
1604 workflow_id,
1605 event_seq,
1606 state_json,
1607 created_at,
1608 },
1609 ))
1610 }
1611
1612 async fn get_queue_stats(&self, namespace: &str) -> Result<Vec<crate::store::QueueStats>> {
1615 let rows = sqlx::query_as::<_, (String, i64, i64, i64)>(
1616 "SELECT
1617 a.task_queue AS queue,
1618 SUM(CASE WHEN a.status = 'PENDING' THEN 1 ELSE 0 END) AS pending,
1619 SUM(CASE WHEN a.status = 'RUNNING' THEN 1 ELSE 0 END) AS running,
1620 (SELECT COUNT(*) FROM workflow.workers w WHERE w.task_queue = a.task_queue AND w.namespace = $1) AS workers
1621 FROM workflow.activities a
1622 JOIN workflow.workflows wf ON a.workflow_id = wf.id AND wf.namespace = $1
1623 GROUP BY a.task_queue",
1624 )
1625 .bind(namespace)
1626 .fetch_all(&self.pool)
1627 .await?;
1628
1629 Ok(rows
1630 .into_iter()
1631 .map(
1632 |(queue, pending, running, workers)| crate::store::QueueStats {
1633 queue,
1634 pending_activities: pending,
1635 running_activities: running,
1636 workers,
1637 },
1638 )
1639 .collect())
1640 }
1641
1642 async fn try_acquire_scheduler_lock(&self) -> Result<bool> {
1645 let row: (bool,) = sqlx::query_as("SELECT pg_try_advisory_lock(42)")
1649 .fetch_one(&self.pool)
1650 .await?;
1651 Ok(row.0)
1652 }
1653}
1654
1655fn timestamp_now() -> f64 {
1656 std::time::SystemTime::now()
1657 .duration_since(std::time::UNIX_EPOCH)
1658 .unwrap()
1659 .as_secs_f64()
1660}
1661
1662#[derive(sqlx::FromRow)]
1665struct PgWorkflowRow {
1666 id: String,
1667 namespace: String,
1668 run_id: String,
1669 workflow_type: String,
1670 task_queue: String,
1671 status: String,
1672 input: Option<String>,
1673 result: Option<String>,
1674 error: Option<String>,
1675 parent_id: Option<String>,
1676 claimed_by: Option<String>,
1677 search_attributes: Option<String>,
1678 archived_at: Option<f64>,
1679 archive_uri: Option<String>,
1680 created_at: f64,
1681 updated_at: f64,
1682 completed_at: Option<f64>,
1683}
1684
1685impl From<PgWorkflowRow> for WorkflowRecord {
1686 fn from(r: PgWorkflowRow) -> Self {
1687 Self {
1688 id: r.id,
1689 namespace: r.namespace,
1690 run_id: r.run_id,
1691 workflow_type: r.workflow_type,
1692 task_queue: r.task_queue,
1693 status: r.status,
1694 input: r.input,
1695 result: r.result,
1696 error: r.error,
1697 parent_id: r.parent_id,
1698 claimed_by: r.claimed_by,
1699 search_attributes: r.search_attributes,
1700 archived_at: r.archived_at,
1701 archive_uri: r.archive_uri,
1702 created_at: r.created_at,
1703 updated_at: r.updated_at,
1704 completed_at: r.completed_at,
1705 }
1706 }
1707}
1708
1709#[derive(sqlx::FromRow)]
1710struct PgEventRow {
1711 id: i64,
1712 workflow_id: String,
1713 seq: i32,
1714 event_type: String,
1715 payload: Option<String>,
1716 timestamp: f64,
1717}
1718
1719impl From<PgEventRow> for WorkflowEvent {
1720 fn from(r: PgEventRow) -> Self {
1721 Self {
1722 id: Some(r.id),
1723 workflow_id: r.workflow_id,
1724 seq: r.seq,
1725 event_type: r.event_type,
1726 payload: r.payload,
1727 timestamp: r.timestamp,
1728 }
1729 }
1730}
1731
1732#[derive(sqlx::FromRow)]
1733struct PgActivityRow {
1734 id: i64,
1735 workflow_id: String,
1736 seq: i32,
1737 name: String,
1738 task_queue: String,
1739 input: Option<String>,
1740 status: String,
1741 result: Option<String>,
1742 error: Option<String>,
1743 attempt: i32,
1744 max_attempts: i32,
1745 initial_interval_secs: f64,
1746 backoff_coefficient: f64,
1747 start_to_close_secs: f64,
1748 heartbeat_timeout_secs: Option<f64>,
1749 claimed_by: Option<String>,
1750 scheduled_at: f64,
1751 started_at: Option<f64>,
1752 completed_at: Option<f64>,
1753 last_heartbeat: Option<f64>,
1754}
1755
1756impl From<PgActivityRow> for WorkflowActivity {
1757 fn from(r: PgActivityRow) -> Self {
1758 Self {
1759 id: Some(r.id),
1760 workflow_id: r.workflow_id,
1761 seq: r.seq,
1762 name: r.name,
1763 task_queue: r.task_queue,
1764 input: r.input,
1765 status: r.status,
1766 result: r.result,
1767 error: r.error,
1768 attempt: r.attempt,
1769 max_attempts: r.max_attempts,
1770 initial_interval_secs: r.initial_interval_secs,
1771 backoff_coefficient: r.backoff_coefficient,
1772 start_to_close_secs: r.start_to_close_secs,
1773 heartbeat_timeout_secs: r.heartbeat_timeout_secs,
1774 claimed_by: r.claimed_by,
1775 scheduled_at: r.scheduled_at,
1776 started_at: r.started_at,
1777 completed_at: r.completed_at,
1778 last_heartbeat: r.last_heartbeat,
1779 }
1780 }
1781}
1782
1783#[derive(sqlx::FromRow)]
1784struct PgTimerRow {
1785 id: i64,
1786 workflow_id: String,
1787 seq: i32,
1788 fire_at: f64,
1789 fired: bool,
1790}
1791
1792impl From<PgTimerRow> for WorkflowTimer {
1793 fn from(r: PgTimerRow) -> Self {
1794 Self {
1795 id: Some(r.id),
1796 workflow_id: r.workflow_id,
1797 seq: r.seq,
1798 fire_at: r.fire_at,
1799 fired: r.fired,
1800 }
1801 }
1802}
1803
1804#[derive(sqlx::FromRow)]
1805struct PgSignalRow {
1806 id: i64,
1807 workflow_id: String,
1808 name: String,
1809 payload: Option<String>,
1810 consumed: bool,
1811 received_at: f64,
1812}
1813
1814impl From<PgSignalRow> for WorkflowSignal {
1815 fn from(r: PgSignalRow) -> Self {
1816 Self {
1817 id: Some(r.id),
1818 workflow_id: r.workflow_id,
1819 name: r.name,
1820 payload: r.payload,
1821 consumed: r.consumed,
1822 received_at: r.received_at,
1823 }
1824 }
1825}
1826
1827#[derive(sqlx::FromRow)]
1828struct PgScheduleRow {
1829 namespace: String,
1830 name: String,
1831 workflow_type: String,
1832 cron_expr: String,
1833 timezone: String,
1834 input: Option<String>,
1835 task_queue: String,
1836 overlap_policy: String,
1837 paused: bool,
1838 last_run_at: Option<f64>,
1839 next_run_at: Option<f64>,
1840 last_workflow_id: Option<String>,
1841 created_at: f64,
1842}
1843
1844impl From<PgScheduleRow> for WorkflowSchedule {
1845 fn from(r: PgScheduleRow) -> Self {
1846 Self {
1847 namespace: r.namespace,
1848 name: r.name,
1849 workflow_type: r.workflow_type,
1850 cron_expr: r.cron_expr,
1851 timezone: r.timezone,
1852 input: r.input,
1853 task_queue: r.task_queue,
1854 overlap_policy: r.overlap_policy,
1855 paused: r.paused,
1856 last_run_at: r.last_run_at,
1857 next_run_at: r.next_run_at,
1858 last_workflow_id: r.last_workflow_id,
1859 created_at: r.created_at,
1860 }
1861 }
1862}
1863
1864#[derive(sqlx::FromRow)]
1865struct PgWorkerRow {
1866 id: String,
1867 namespace: String,
1868 identity: String,
1869 task_queue: String,
1870 workflows: Option<String>,
1871 activities: Option<String>,
1872 max_concurrent_workflows: i32,
1873 max_concurrent_activities: i32,
1874 active_tasks: i32,
1875 last_heartbeat: f64,
1876 registered_at: f64,
1877}
1878
1879impl From<PgWorkerRow> for WorkflowWorker {
1880 fn from(r: PgWorkerRow) -> Self {
1881 Self {
1882 id: r.id,
1883 namespace: r.namespace,
1884 identity: r.identity,
1885 task_queue: r.task_queue,
1886 workflows: r.workflows,
1887 activities: r.activities,
1888 max_concurrent_workflows: r.max_concurrent_workflows,
1889 max_concurrent_activities: r.max_concurrent_activities,
1890 active_tasks: r.active_tasks,
1891 last_heartbeat: r.last_heartbeat,
1892 registered_at: r.registered_at,
1893 }
1894 }
1895}
1896
1897#[cfg(test)]
1898mod tests {
1899 use super::*;
1900
1901 #[test]
1902 fn sanitise_schema_keeps_statements_intact() {
1903 let input = "CREATE TABLE foo (x INT);\nCREATE INDEX idx_foo ON foo(x);\n";
1904 let out = sanitise_schema(input);
1905 assert_eq!(out.len(), 2);
1906 assert!(out[0].starts_with("CREATE TABLE foo"));
1907 assert!(out[1].starts_with("CREATE INDEX idx_foo"));
1908 }
1909
1910 #[test]
1911 fn sanitise_schema_drops_pure_comment_lines() {
1912 let input = "-- header comment\nCREATE TABLE foo (x INT);\n-- trailing comment\n";
1913 let out = sanitise_schema(input);
1914 assert_eq!(out.len(), 1);
1915 assert!(out[0].starts_with("CREATE TABLE foo"));
1916 }
1917
1918 #[test]
1919 fn sanitise_schema_ignores_semicolons_inside_comment_prose() {
1920 let input = "\
1924CREATE TABLE foo (x INT);
1925-- Idempotent across startups; fresh installs pick the column up from the
1926-- CREATE TABLE above so the ADD is a no-op.
1927";
1928 let out = sanitise_schema(input);
1929 assert_eq!(
1930 out.len(),
1931 1,
1932 "expected 1 real statement, got {}: {:?}",
1933 out.len(),
1934 out
1935 );
1936 assert!(out[0].starts_with("CREATE TABLE foo"));
1937 }
1938
1939 #[test]
1940 fn sanitise_schema_drops_indented_comment_lines() {
1941 let input = " -- indented comment\n\tCREATE TABLE foo (x INT);\n";
1942 let out = sanitise_schema(input);
1943 assert_eq!(out.len(), 1);
1944 assert!(out[0].contains("CREATE TABLE foo"));
1945 }
1946
1947 #[test]
1948 fn sanitise_schema_real_constant_produces_only_ddl() {
1949 for stmt in sanitise_schema(SCHEMA) {
1953 let first_word = stmt
1954 .split_whitespace()
1955 .next()
1956 .expect("non-empty statement")
1957 .to_uppercase();
1958 assert!(
1959 matches!(
1960 first_word.as_str(),
1961 "CREATE" | "INSERT" | "UPDATE" | "DROP" | "ALTER" | "WITH"
1962 ),
1963 "SCHEMA produced non-DDL statement starting with {first_word:?}: {stmt:?}"
1964 );
1965 }
1966 }
1967}