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