Skip to main content

assay_workflow/store/
postgres.rs

1use anyhow::Result;
2use sqlx::PgPool;
3
4use crate::store::WorkflowStore;
5use crate::types::*;
6
7const SCHEMA: &str = r#"
8CREATE TABLE IF NOT EXISTS namespaces (
9    name            TEXT PRIMARY KEY,
10    created_at      DOUBLE PRECISION NOT NULL
11);
12INSERT INTO namespaces (name, created_at)
13    VALUES ('main', EXTRACT(EPOCH FROM NOW()))
14    ON CONFLICT DO NOTHING;
15
16CREATE TABLE IF NOT EXISTS workflows (
17    id              TEXT PRIMARY KEY,
18    namespace       TEXT NOT NULL DEFAULT 'main',
19    run_id          TEXT NOT NULL,
20    workflow_type   TEXT NOT NULL,
21    task_queue      TEXT NOT NULL DEFAULT 'main',
22    status          TEXT NOT NULL DEFAULT 'PENDING',
23    input           TEXT,
24    result          TEXT,
25    error           TEXT,
26    parent_id       TEXT,
27    claimed_by      TEXT,
28    search_attributes TEXT,
29    archived_at     DOUBLE PRECISION,
30    archive_uri     TEXT,
31    -- Workflow-task dispatch (Phase 9): see sqlite.rs for the full comment.
32    needs_dispatch  BOOLEAN NOT NULL DEFAULT FALSE,
33    dispatch_claimed_by    TEXT,
34    dispatch_last_heartbeat DOUBLE PRECISION,
35    created_at      DOUBLE PRECISION NOT NULL,
36    updated_at      DOUBLE PRECISION NOT NULL,
37    completed_at    DOUBLE PRECISION
38);
39CREATE INDEX IF NOT EXISTS idx_wf_status_queue ON workflows(status, task_queue);
40CREATE INDEX IF NOT EXISTS idx_wf_namespace ON workflows(namespace);
41CREATE INDEX IF NOT EXISTS idx_wf_dispatch ON workflows(task_queue, needs_dispatch, dispatch_claimed_by);
42
43CREATE TABLE IF NOT EXISTS workflow_events (
44    id              BIGSERIAL PRIMARY KEY,
45    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
46    seq             INTEGER NOT NULL,
47    event_type      TEXT NOT NULL,
48    payload         TEXT,
49    timestamp       DOUBLE PRECISION NOT NULL
50);
51CREATE INDEX IF NOT EXISTS idx_wf_events_lookup ON workflow_events(workflow_id, seq);
52
53CREATE TABLE IF NOT EXISTS workflow_activities (
54    id              BIGSERIAL PRIMARY KEY,
55    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
56    seq             INTEGER NOT NULL,
57    name            TEXT NOT NULL,
58    task_queue      TEXT NOT NULL DEFAULT 'main',
59    input           TEXT,
60    status          TEXT NOT NULL DEFAULT 'PENDING',
61    result          TEXT,
62    error           TEXT,
63    attempt         INTEGER NOT NULL DEFAULT 1,
64    max_attempts    INTEGER NOT NULL DEFAULT 3,
65    initial_interval_secs   DOUBLE PRECISION NOT NULL DEFAULT 1,
66    backoff_coefficient     DOUBLE PRECISION NOT NULL DEFAULT 2,
67    start_to_close_secs     DOUBLE PRECISION NOT NULL DEFAULT 300,
68    heartbeat_timeout_secs  DOUBLE PRECISION,
69    claimed_by      TEXT,
70    scheduled_at    DOUBLE PRECISION NOT NULL,
71    started_at      DOUBLE PRECISION,
72    completed_at    DOUBLE PRECISION,
73    last_heartbeat  DOUBLE PRECISION,
74    UNIQUE (workflow_id, seq)
75);
76CREATE INDEX IF NOT EXISTS idx_wf_act_pending ON workflow_activities(task_queue, status, scheduled_at);
77
78CREATE TABLE IF NOT EXISTS workflow_timers (
79    id              BIGSERIAL PRIMARY KEY,
80    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
81    seq             INTEGER NOT NULL,
82    fire_at         DOUBLE PRECISION NOT NULL,
83    fired           BOOLEAN NOT NULL DEFAULT FALSE,
84    UNIQUE (workflow_id, seq)
85);
86CREATE INDEX IF NOT EXISTS idx_wf_timers_due ON workflow_timers(fire_at) WHERE fired = FALSE;
87
88CREATE TABLE IF NOT EXISTS workflow_signals (
89    id              BIGSERIAL PRIMARY KEY,
90    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
91    name            TEXT NOT NULL,
92    payload         TEXT,
93    consumed        BOOLEAN NOT NULL DEFAULT FALSE,
94    received_at     DOUBLE PRECISION NOT NULL
95);
96CREATE INDEX IF NOT EXISTS idx_wf_signals_lookup ON workflow_signals(workflow_id, name, consumed);
97
98CREATE TABLE IF NOT EXISTS workflow_schedules (
99    namespace       TEXT NOT NULL DEFAULT 'main',
100    name            TEXT NOT NULL,
101    workflow_type   TEXT NOT NULL,
102    cron_expr       TEXT NOT NULL,
103    timezone        TEXT NOT NULL DEFAULT 'UTC',
104    input           TEXT,
105    task_queue      TEXT NOT NULL DEFAULT 'main',
106    overlap_policy  TEXT NOT NULL DEFAULT 'skip',
107    paused          BOOLEAN NOT NULL DEFAULT FALSE,
108    last_run_at     DOUBLE PRECISION,
109    next_run_at     DOUBLE PRECISION,
110    last_workflow_id TEXT,
111    created_at      DOUBLE PRECISION NOT NULL,
112    PRIMARY KEY (namespace, name)
113);
114
115CREATE TABLE IF NOT EXISTS workflow_workers (
116    id              TEXT PRIMARY KEY,
117    namespace       TEXT NOT NULL DEFAULT 'main',
118    identity        TEXT NOT NULL,
119    task_queue      TEXT NOT NULL,
120    workflows       TEXT,
121    activities      TEXT,
122    max_concurrent_workflows  INTEGER NOT NULL DEFAULT 10,
123    max_concurrent_activities INTEGER NOT NULL DEFAULT 10,
124    active_tasks    INTEGER NOT NULL DEFAULT 0,
125    last_heartbeat  DOUBLE PRECISION NOT NULL,
126    registered_at   DOUBLE PRECISION NOT NULL
127);
128
129CREATE TABLE IF NOT EXISTS workflow_snapshots (
130    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
131    event_seq       INTEGER NOT NULL,
132    state_json      TEXT NOT NULL,
133    created_at      DOUBLE PRECISION NOT NULL,
134    PRIMARY KEY (workflow_id, event_seq)
135);
136
137CREATE TABLE IF NOT EXISTS api_keys (
138    key_hash        TEXT PRIMARY KEY,
139    prefix          TEXT NOT NULL,
140    label           TEXT,
141    created_at      DOUBLE PRECISION NOT NULL
142);
143CREATE INDEX IF NOT EXISTS idx_api_keys_prefix ON api_keys(prefix);
144
145CREATE TABLE IF NOT EXISTS engine_events (
146    id              BIGSERIAL PRIMARY KEY,
147    ts              DOUBLE PRECISION NOT NULL DEFAULT EXTRACT(EPOCH FROM NOW()),
148    namespace       TEXT NOT NULL,
149    subsystem       TEXT NOT NULL,
150    kind            TEXT NOT NULL,
151    payload         JSONB NOT NULL DEFAULT '{}'::jsonb
152);
153CREATE INDEX IF NOT EXISTS idx_engine_events_ns_id ON engine_events(namespace, id);
154CREATE INDEX IF NOT EXISTS idx_engine_events_ts_prune ON engine_events(ts);
155
156"#;
157
158/// Split a Postgres DDL script into individual statements ready for `sqlx::query`.
159///
160/// Drops pure-comment lines (those starting with `--` after optional whitespace)
161/// *before* splitting on `;`. Without this step, a semicolon inside a line comment
162/// (e.g. `-- Idempotent across startups; fresh installs pick the column up`) would
163/// split the surrounding comment into fragments — one of which is naked prose that
164/// Postgres tries to parse as SQL and rejects with `syntax error at or near "<word>"`.
165///
166/// The filter only drops *pure-comment* lines (leading whitespace then `--`), leaving
167/// `--`-after-code untouched. That keeps string literals safe (could legally contain
168/// `--`) and is conservative enough to remain correct if the SCHEMA grows more prose.
169fn sanitise_schema(schema: &str) -> Vec<String> {
170    let without_comments: String = schema
171        .lines()
172        .filter(|line| !line.trim_start().starts_with("--"))
173        .collect::<Vec<_>>()
174        .join("\n");
175
176    without_comments
177        .split(';')
178        .map(|s| s.trim().to_string())
179        .filter(|s| !s.is_empty())
180        .collect()
181}
182
183pub struct PostgresStore {
184    pool: PgPool,
185}
186
187impl PostgresStore {
188    pub async fn new(url: &str) -> Result<Self> {
189        let pool = PgPool::connect(url).await?;
190        Self::from_pool(pool).await
191    }
192
193    /// Build a store from an existing pool. Runs migrations on the target
194    /// database. Useful when the engine owns the pool (shared with other
195    /// modules) and hands a clone to the workflow module, or for tests that
196    /// point many stores at different databases in the same Postgres server.
197    pub async fn from_pool(pool: PgPool) -> Result<Self> {
198        let store = Self { pool };
199        store.migrate().await?;
200        Ok(store)
201    }
202
203    /// Expose the underlying pool (used by the engine to build a
204    /// `PgEngineEventBus` that shares the same connection pool).
205    pub fn pool(&self) -> &PgPool {
206        &self.pool
207    }
208
209    async fn migrate(&self) -> Result<()> {
210        // Apply the base schema (tables + indexes) statement-by-statement.
211        for statement in sanitise_schema(SCHEMA) {
212            sqlx::query(&statement).execute(&self.pool).await?;
213        }
214        // Drop the v0.13.0 LISTEN/NOTIFY triggers if they still exist on
215        // the target database. The Rust-managed CDC outbox in
216        // assay_domain::events is the replacement; leaving stale
217        // triggers in place would double-publish NOTIFYs with channels
218        // no one listens to.
219        sqlx::raw_sql(
220            r#"
221            DROP TRIGGER IF EXISTS workflow_runnable_notify ON workflows;
222            DROP TRIGGER IF EXISTS workflow_task_notify ON workflow_activities;
223            DROP FUNCTION IF EXISTS assay_notify_runnable();
224            DROP FUNCTION IF EXISTS assay_notify_task();
225            "#,
226        )
227        .execute(&self.pool)
228        .await?;
229        Ok(())
230    }
231
232    /// Try to acquire pg_advisory_lock for leader election.
233    /// Returns true if this instance is the leader (scheduler should run).
234    pub async fn try_acquire_leader_lock(&self) -> Result<bool> {
235        let row: (bool,) =
236            sqlx::query_as("SELECT pg_try_advisory_lock(1)")
237                .fetch_one(&self.pool)
238                .await?;
239        Ok(row.0)
240    }
241}
242
243impl WorkflowStore for PostgresStore {
244    // ── Namespaces ─────────────────────────────────────────
245
246    async fn create_namespace(&self, name: &str) -> Result<()> {
247        sqlx::query("INSERT INTO namespaces (name, created_at) VALUES ($1, EXTRACT(EPOCH FROM NOW()))")
248            .bind(name)
249            .execute(&self.pool)
250            .await?;
251        Ok(())
252    }
253
254    async fn list_namespaces(&self) -> Result<Vec<crate::store::NamespaceRecord>> {
255        let rows = sqlx::query_as::<_, (String, f64)>(
256            "SELECT name, created_at FROM namespaces ORDER BY name",
257        )
258        .fetch_all(&self.pool)
259        .await?;
260        Ok(rows
261            .into_iter()
262            .map(|(name, created_at)| crate::store::NamespaceRecord { name, created_at })
263            .collect())
264    }
265
266    async fn delete_namespace(&self, name: &str) -> Result<bool> {
267        let res = sqlx::query("DELETE FROM namespaces WHERE name = $1 AND name != 'main'")
268            .bind(name)
269            .execute(&self.pool)
270            .await?;
271        Ok(res.rows_affected() > 0)
272    }
273
274    async fn get_namespace_stats(&self, namespace: &str) -> Result<crate::store::NamespaceStats> {
275        let total: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM workflows WHERE namespace = $1")
276            .bind(namespace)
277            .fetch_one(&self.pool)
278            .await?;
279        let running: (i64,) = sqlx::query_as(
280            "SELECT COUNT(*) FROM workflows WHERE namespace = $1 AND status = 'RUNNING'",
281        )
282        .bind(namespace)
283        .fetch_one(&self.pool)
284        .await?;
285        let pending: (i64,) = sqlx::query_as(
286            "SELECT COUNT(*) FROM workflows WHERE namespace = $1 AND status = 'PENDING'",
287        )
288        .bind(namespace)
289        .fetch_one(&self.pool)
290        .await?;
291        let completed: (i64,) = sqlx::query_as(
292            "SELECT COUNT(*) FROM workflows WHERE namespace = $1 AND status = 'COMPLETED'",
293        )
294        .bind(namespace)
295        .fetch_one(&self.pool)
296        .await?;
297        let failed: (i64,) = sqlx::query_as(
298            "SELECT COUNT(*) FROM workflows WHERE namespace = $1 AND status = 'FAILED'",
299        )
300        .bind(namespace)
301        .fetch_one(&self.pool)
302        .await?;
303        let schedules: (i64,) =
304            sqlx::query_as("SELECT COUNT(*) FROM workflow_schedules WHERE namespace = $1")
305                .bind(namespace)
306                .fetch_one(&self.pool)
307                .await?;
308        let workers: (i64,) =
309            sqlx::query_as("SELECT COUNT(*) FROM workflow_workers WHERE namespace = $1")
310                .bind(namespace)
311                .fetch_one(&self.pool)
312                .await?;
313
314        Ok(crate::store::NamespaceStats {
315            namespace: namespace.to_string(),
316            total_workflows: total.0,
317            running: running.0,
318            pending: pending.0,
319            completed: completed.0,
320            failed: failed.0,
321            schedules: schedules.0,
322            workers: workers.0,
323        })
324    }
325
326    // ── Workflows ──────────────────────────────────────────
327
328    async fn create_workflow(&self, wf: &WorkflowRecord) -> Result<()> {
329        sqlx::query(
330            "INSERT INTO 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)
331             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)",
332        )
333        .bind(&wf.id)
334        .bind(&wf.namespace)
335        .bind(&wf.run_id)
336        .bind(&wf.workflow_type)
337        .bind(&wf.task_queue)
338        .bind(&wf.status)
339        .bind(&wf.input)
340        .bind(&wf.result)
341        .bind(&wf.error)
342        .bind(&wf.parent_id)
343        .bind(&wf.claimed_by)
344        .bind(&wf.search_attributes)
345        .bind(wf.archived_at)
346        .bind(&wf.archive_uri)
347        .bind(wf.created_at)
348        .bind(wf.updated_at)
349        .bind(wf.completed_at)
350        .execute(&self.pool)
351        .await?;
352        Ok(())
353    }
354
355    async fn get_workflow(&self, id: &str) -> Result<Option<WorkflowRecord>> {
356        let row = sqlx::query_as::<_, PgWorkflowRow>(
357            "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 workflows WHERE id = $1",
358        )
359        .bind(id)
360        .fetch_optional(&self.pool)
361        .await?;
362        Ok(row.map(Into::into))
363    }
364
365    async fn list_workflows(
366        &self,
367        namespace: &str,
368        status: Option<WorkflowStatus>,
369        workflow_type: Option<&str>,
370        search_attrs_filter: Option<&str>,
371        limit: i64,
372        offset: i64,
373    ) -> Result<Vec<WorkflowRecord>> {
374        let status_str = status.map(|s| s.to_string());
375
376        let filter_pairs: Vec<(String, serde_json::Value)> = search_attrs_filter
377            .and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
378            .and_then(|v| v.as_object().cloned())
379            .map(|m| m.into_iter().collect())
380            .unwrap_or_default();
381
382        let mut sql = String::from(
383            "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
384             FROM workflows
385             WHERE namespace = $1
386               AND ($2::TEXT IS NULL OR status = $2)
387               AND ($3::TEXT IS NULL OR workflow_type = $3)",
388        );
389        // Bind placeholders for the filter follow $3; next index is 4.
390        let mut idx = 4usize;
391        for _ in &filter_pairs {
392            sql.push_str(&format!(
393                " AND (search_attributes::jsonb)->>${} = ${}",
394                idx,
395                idx + 1
396            ));
397            idx += 2;
398        }
399        sql.push_str(&format!(" ORDER BY created_at DESC LIMIT ${} OFFSET ${}", idx, idx + 1));
400
401        let mut q = sqlx::query_as::<_, PgWorkflowRow>(&sql)
402            .bind(namespace)
403            .bind(&status_str)
404            .bind(workflow_type);
405        for (key, value) in &filter_pairs {
406            q = q.bind(key.clone());
407            // JSONB ->> always returns TEXT; compare by stringified value.
408            let as_text = match value {
409                serde_json::Value::String(s) => s.clone(),
410                other => other.to_string(),
411            };
412            q = q.bind(as_text);
413        }
414        let rows = q.bind(limit).bind(offset).fetch_all(&self.pool).await?;
415        Ok(rows.into_iter().map(Into::into).collect())
416    }
417
418    async fn update_workflow_status(
419        &self,
420        id: &str,
421        status: WorkflowStatus,
422        result: Option<&str>,
423        error: Option<&str>,
424    ) -> Result<()> {
425        let now = timestamp_now();
426        let completed_at = if status.is_terminal() { Some(now) } else { None };
427        sqlx::query(
428            "UPDATE workflows SET status = $1, result = COALESCE($2, result), error = COALESCE($3, error), updated_at = $4, completed_at = COALESCE($5, completed_at) WHERE id = $6",
429        )
430        .bind(status.to_string())
431        .bind(result)
432        .bind(error)
433        .bind(now)
434        .bind(completed_at)
435        .bind(id)
436        .execute(&self.pool)
437        .await?;
438        Ok(())
439    }
440
441    async fn claim_workflow(&self, id: &str, worker_id: &str) -> Result<bool> {
442        let res = sqlx::query(
443            "UPDATE workflows SET claimed_by = $1, status = 'RUNNING', updated_at = $2 WHERE id = $3 AND claimed_by IS NULL",
444        )
445        .bind(worker_id)
446        .bind(timestamp_now())
447        .bind(id)
448        .execute(&self.pool)
449        .await?;
450        Ok(res.rows_affected() > 0)
451    }
452
453    async fn mark_workflow_dispatchable(&self, workflow_id: &str) -> Result<()> {
454        sqlx::query("UPDATE workflows SET needs_dispatch = TRUE WHERE id = $1")
455            .bind(workflow_id)
456            .execute(&self.pool)
457            .await?;
458        Ok(())
459    }
460
461    async fn claim_workflow_task(
462        &self,
463        task_queue: &str,
464        worker_id: &str,
465    ) -> Result<Option<WorkflowRecord>> {
466        let now = timestamp_now();
467        // Atomic claim with FOR UPDATE SKIP LOCKED so multiple engine
468        // replicas don't fight over the same workflow task.
469        let row = sqlx::query_as::<_, PgWorkflowRow>(
470            "UPDATE workflows
471             SET dispatch_claimed_by = $1, dispatch_last_heartbeat = $2, needs_dispatch = FALSE
472             WHERE id = (
473                SELECT id FROM workflows
474                WHERE task_queue = $3
475                  AND needs_dispatch = TRUE
476                  AND dispatch_claimed_by IS NULL
477                  AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
478                ORDER BY updated_at ASC
479                FOR UPDATE SKIP LOCKED
480                LIMIT 1
481             )
482             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",
483        )
484        .bind(worker_id)
485        .bind(now)
486        .bind(task_queue)
487        .fetch_optional(&self.pool)
488        .await?;
489        Ok(row.map(Into::into))
490    }
491
492    async fn release_workflow_task(&self, workflow_id: &str, worker_id: &str) -> Result<()> {
493        sqlx::query(
494            "UPDATE workflows
495             SET dispatch_claimed_by = NULL, dispatch_last_heartbeat = NULL
496             WHERE id = $1 AND dispatch_claimed_by = $2",
497        )
498        .bind(workflow_id)
499        .bind(worker_id)
500        .execute(&self.pool)
501        .await?;
502        Ok(())
503    }
504
505    async fn release_stale_dispatch_leases(
506        &self,
507        now: f64,
508        timeout_secs: f64,
509    ) -> Result<u64> {
510        let res = sqlx::query(
511            "UPDATE workflows
512             SET dispatch_claimed_by = NULL,
513                 dispatch_last_heartbeat = NULL,
514                 needs_dispatch = TRUE
515             WHERE dispatch_claimed_by IS NOT NULL
516               AND ($1 - dispatch_last_heartbeat) > $2
517               AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')",
518        )
519        .bind(now)
520        .bind(timeout_secs)
521        .execute(&self.pool)
522        .await?;
523        Ok(res.rows_affected())
524    }
525
526    // ── Events ─────────────────────────────────────────────
527
528    async fn append_event(&self, ev: &WorkflowEvent) -> Result<i64> {
529        let row: (i64,) = sqlx::query_as(
530            "INSERT INTO workflow_events (workflow_id, seq, event_type, payload, timestamp) VALUES ($1, $2, $3, $4, $5) RETURNING id",
531        )
532        .bind(&ev.workflow_id)
533        .bind(ev.seq)
534        .bind(&ev.event_type)
535        .bind(&ev.payload)
536        .bind(ev.timestamp)
537        .fetch_one(&self.pool)
538        .await?;
539        Ok(row.0)
540    }
541
542    async fn list_events(&self, workflow_id: &str) -> Result<Vec<WorkflowEvent>> {
543        let rows = sqlx::query_as::<_, PgEventRow>(
544            "SELECT id, workflow_id, seq, event_type, payload, timestamp FROM workflow_events WHERE workflow_id = $1 ORDER BY seq ASC",
545        )
546        .bind(workflow_id)
547        .fetch_all(&self.pool)
548        .await?;
549        Ok(rows.into_iter().map(Into::into).collect())
550    }
551
552    async fn get_event_count(&self, workflow_id: &str) -> Result<i64> {
553        let row: (i64,) =
554            sqlx::query_as("SELECT COUNT(*) FROM workflow_events WHERE workflow_id = $1")
555                .bind(workflow_id)
556                .fetch_one(&self.pool)
557                .await?;
558        Ok(row.0)
559    }
560
561    // ── Activities ──────────────────────────────────────────
562
563    async fn create_activity(&self, act: &WorkflowActivity) -> Result<i64> {
564        let row: (i64,) = sqlx::query_as(
565            "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)
566             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) RETURNING id",
567        )
568        .bind(&act.workflow_id)
569        .bind(act.seq)
570        .bind(&act.name)
571        .bind(&act.task_queue)
572        .bind(&act.input)
573        .bind(&act.status)
574        .bind(act.attempt)
575        .bind(act.max_attempts)
576        .bind(act.initial_interval_secs)
577        .bind(act.backoff_coefficient)
578        .bind(act.start_to_close_secs)
579        .bind(act.heartbeat_timeout_secs)
580        .bind(act.scheduled_at)
581        .fetch_one(&self.pool)
582        .await?;
583        Ok(row.0)
584    }
585
586    async fn get_activity(&self, id: i64) -> Result<Option<WorkflowActivity>> {
587        let row = sqlx::query_as::<_, PgActivityRow>(
588            "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
589             FROM workflow_activities WHERE id = $1",
590        )
591        .bind(id)
592        .fetch_optional(&self.pool)
593        .await?;
594        Ok(row.map(Into::into))
595    }
596
597    async fn get_activity_by_workflow_seq(
598        &self,
599        workflow_id: &str,
600        seq: i32,
601    ) -> Result<Option<WorkflowActivity>> {
602        let row = sqlx::query_as::<_, PgActivityRow>(
603            "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
604             FROM workflow_activities WHERE workflow_id = $1 AND seq = $2",
605        )
606        .bind(workflow_id)
607        .bind(seq)
608        .fetch_optional(&self.pool)
609        .await?;
610        Ok(row.map(Into::into))
611    }
612
613    async fn claim_activity(
614        &self,
615        task_queue: &str,
616        worker_id: &str,
617    ) -> Result<Option<WorkflowActivity>> {
618        let now = timestamp_now();
619        // Atomic claim using FOR UPDATE SKIP LOCKED — prevents contention
620        // between multiple assay serve instances claiming the same activity
621        let row = sqlx::query_as::<_, PgActivityRow>(
622            "UPDATE workflow_activities SET status = 'RUNNING', claimed_by = $1, started_at = $2
623             WHERE id = (
624                SELECT id FROM workflow_activities
625                WHERE task_queue = $3 AND status = 'PENDING'
626                ORDER BY scheduled_at ASC
627                FOR UPDATE SKIP LOCKED
628                LIMIT 1
629             )
630             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",
631        )
632        .bind(worker_id)
633        .bind(now)
634        .bind(task_queue)
635        .fetch_optional(&self.pool)
636        .await?;
637        Ok(row.map(Into::into))
638    }
639
640    async fn requeue_activity_for_retry(
641        &self,
642        id: i64,
643        next_attempt: i32,
644        next_scheduled_at: f64,
645    ) -> Result<()> {
646        sqlx::query(
647            "UPDATE workflow_activities
648             SET status = 'PENDING', attempt = $1, scheduled_at = $2,
649                 claimed_by = NULL, started_at = NULL, last_heartbeat = NULL,
650                 error = NULL
651             WHERE id = $3",
652        )
653        .bind(next_attempt)
654        .bind(next_scheduled_at)
655        .bind(id)
656        .execute(&self.pool)
657        .await?;
658        Ok(())
659    }
660
661    async fn complete_activity(
662        &self,
663        id: i64,
664        result: Option<&str>,
665        error: Option<&str>,
666        failed: bool,
667    ) -> Result<()> {
668        let status = if failed { "FAILED" } else { "COMPLETED" };
669        sqlx::query(
670            "UPDATE workflow_activities SET status = $1, result = $2, error = $3, completed_at = $4 WHERE id = $5",
671        )
672        .bind(status)
673        .bind(result)
674        .bind(error)
675        .bind(timestamp_now())
676        .bind(id)
677        .execute(&self.pool)
678        .await?;
679        Ok(())
680    }
681
682    async fn heartbeat_activity(&self, id: i64, _details: Option<&str>) -> Result<()> {
683        sqlx::query("UPDATE workflow_activities SET last_heartbeat = $1 WHERE id = $2")
684            .bind(timestamp_now())
685            .bind(id)
686            .execute(&self.pool)
687            .await?;
688        Ok(())
689    }
690
691    async fn get_timed_out_activities(&self, now: f64) -> Result<Vec<WorkflowActivity>> {
692        let rows = sqlx::query_as::<_, PgActivityRow>(
693            "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
694             FROM workflow_activities
695             WHERE status = 'RUNNING'
696               AND heartbeat_timeout_secs IS NOT NULL
697               AND last_heartbeat IS NOT NULL
698               AND ($1 - last_heartbeat) > heartbeat_timeout_secs",
699        )
700        .bind(now)
701        .fetch_all(&self.pool)
702        .await?;
703        Ok(rows.into_iter().map(Into::into).collect())
704    }
705
706    // ── Timers ──────────────────────────────────────────────
707
708    async fn create_timer(&self, timer: &WorkflowTimer) -> Result<i64> {
709        // Idempotent: ON CONFLICT (workflow_id, seq) DO NOTHING.
710        // If a row already exists, RETURNING produces no rows — fall back to SELECT.
711        let inserted: Option<(i64,)> = sqlx::query_as(
712            "INSERT INTO workflow_timers (workflow_id, seq, fire_at, fired)
713             VALUES ($1, $2, $3, FALSE)
714             ON CONFLICT (workflow_id, seq) DO NOTHING
715             RETURNING id",
716        )
717        .bind(&timer.workflow_id)
718        .bind(timer.seq)
719        .bind(timer.fire_at)
720        .fetch_optional(&self.pool)
721        .await?;
722
723        if let Some((id,)) = inserted {
724            return Ok(id);
725        }
726
727        // Row already existed — return its id.
728        let (id,): (i64,) = sqlx::query_as(
729            "SELECT id FROM workflow_timers WHERE workflow_id = $1 AND seq = $2",
730        )
731        .bind(&timer.workflow_id)
732        .bind(timer.seq)
733        .fetch_one(&self.pool)
734        .await?;
735        Ok(id)
736    }
737
738    async fn cancel_pending_activities(&self, workflow_id: &str) -> Result<u64> {
739        let res = sqlx::query(
740            "UPDATE workflow_activities SET status = 'CANCELLED', completed_at = $1
741             WHERE workflow_id = $2 AND status = 'PENDING'",
742        )
743        .bind(timestamp_now())
744        .bind(workflow_id)
745        .execute(&self.pool)
746        .await?;
747        Ok(res.rows_affected())
748    }
749
750    async fn cancel_pending_timers(&self, workflow_id: &str) -> Result<u64> {
751        let res = sqlx::query(
752            "UPDATE workflow_timers SET fired = TRUE
753             WHERE workflow_id = $1 AND fired = FALSE",
754        )
755        .bind(workflow_id)
756        .execute(&self.pool)
757        .await?;
758        Ok(res.rows_affected())
759    }
760
761    async fn get_timer_by_workflow_seq(
762        &self,
763        workflow_id: &str,
764        seq: i32,
765    ) -> Result<Option<WorkflowTimer>> {
766        let row = sqlx::query_as::<_, PgTimerRow>(
767            "SELECT id, workflow_id, seq, fire_at, fired
768             FROM workflow_timers WHERE workflow_id = $1 AND seq = $2",
769        )
770        .bind(workflow_id)
771        .bind(seq)
772        .fetch_optional(&self.pool)
773        .await?;
774        Ok(row.map(Into::into))
775    }
776
777    async fn fire_due_timers(&self, now: f64) -> Result<Vec<WorkflowTimer>> {
778        let rows = sqlx::query_as::<_, PgTimerRow>(
779            "UPDATE workflow_timers SET fired = TRUE
780             WHERE fired = FALSE AND fire_at <= $1
781             RETURNING id, workflow_id, seq, fire_at, fired",
782        )
783        .bind(now)
784        .fetch_all(&self.pool)
785        .await?;
786        Ok(rows.into_iter().map(Into::into).collect())
787    }
788
789    // ── Signals ─────────────────────────────────────────────
790
791    async fn send_signal(&self, sig: &WorkflowSignal) -> Result<i64> {
792        let row: (i64,) = sqlx::query_as(
793            "INSERT INTO workflow_signals (workflow_id, name, payload, consumed, received_at) VALUES ($1, $2, $3, FALSE, $4) RETURNING id",
794        )
795        .bind(&sig.workflow_id)
796        .bind(&sig.name)
797        .bind(&sig.payload)
798        .bind(sig.received_at)
799        .fetch_one(&self.pool)
800        .await?;
801        Ok(row.0)
802    }
803
804    async fn consume_signals(
805        &self,
806        workflow_id: &str,
807        name: &str,
808    ) -> Result<Vec<WorkflowSignal>> {
809        let rows = sqlx::query_as::<_, PgSignalRow>(
810            "UPDATE workflow_signals SET consumed = TRUE
811             WHERE workflow_id = $1 AND name = $2 AND consumed = FALSE
812             RETURNING id, workflow_id, name, payload, consumed, received_at",
813        )
814        .bind(workflow_id)
815        .bind(name)
816        .fetch_all(&self.pool)
817        .await?;
818        Ok(rows.into_iter().map(Into::into).collect())
819    }
820
821    // ── Schedules ───────────────────────────────────────────
822
823    async fn create_schedule(&self, sched: &WorkflowSchedule) -> Result<()> {
824        sqlx::query(
825            "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)
826             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)",
827        )
828        .bind(&sched.namespace)
829        .bind(&sched.name)
830        .bind(&sched.workflow_type)
831        .bind(&sched.cron_expr)
832        .bind(&sched.timezone)
833        .bind(&sched.input)
834        .bind(&sched.task_queue)
835        .bind(&sched.overlap_policy)
836        .bind(sched.paused)
837        .bind(sched.last_run_at)
838        .bind(sched.next_run_at)
839        .bind(&sched.last_workflow_id)
840        .bind(sched.created_at)
841        .execute(&self.pool)
842        .await?;
843        Ok(())
844    }
845
846    async fn get_schedule(&self, namespace: &str, name: &str) -> Result<Option<WorkflowSchedule>> {
847        let row = sqlx::query_as::<_, PgScheduleRow>(
848            "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",
849        )
850        .bind(namespace)
851        .bind(name)
852        .fetch_optional(&self.pool)
853        .await?;
854        Ok(row.map(Into::into))
855    }
856
857    async fn list_schedules(&self, namespace: &str) -> Result<Vec<WorkflowSchedule>> {
858        let rows = sqlx::query_as::<_, PgScheduleRow>(
859            "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",
860        )
861        .bind(namespace)
862        .fetch_all(&self.pool)
863        .await?;
864        Ok(rows.into_iter().map(Into::into).collect())
865    }
866
867    async fn update_schedule_last_run(
868        &self,
869        namespace: &str,
870        name: &str,
871        last_run_at: f64,
872        next_run_at: f64,
873        workflow_id: &str,
874    ) -> Result<()> {
875        sqlx::query(
876            "UPDATE workflow_schedules SET last_run_at = $1, next_run_at = $2, last_workflow_id = $3 WHERE namespace = $4 AND name = $5",
877        )
878        .bind(last_run_at)
879        .bind(next_run_at)
880        .bind(workflow_id)
881        .bind(namespace)
882        .bind(name)
883        .execute(&self.pool)
884        .await?;
885        Ok(())
886    }
887
888    async fn delete_schedule(&self, namespace: &str, name: &str) -> Result<bool> {
889        let res = sqlx::query("DELETE FROM workflow_schedules WHERE namespace = $1 AND name = $2")
890            .bind(namespace)
891            .bind(name)
892            .execute(&self.pool)
893            .await?;
894        Ok(res.rows_affected() > 0)
895    }
896
897    async fn list_archivable_workflows(
898        &self,
899        cutoff: f64,
900        limit: i64,
901    ) -> Result<Vec<WorkflowRecord>> {
902        let rows = sqlx::query_as::<_, PgWorkflowRow>(
903            "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
904             FROM workflows
905             WHERE status IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
906               AND completed_at IS NOT NULL
907               AND completed_at < $1
908               AND archived_at IS NULL
909             ORDER BY completed_at ASC
910             LIMIT $2",
911        )
912        .bind(cutoff)
913        .bind(limit)
914        .fetch_all(&self.pool)
915        .await?;
916        Ok(rows.into_iter().map(Into::into).collect())
917    }
918
919    async fn mark_archived_and_purge(
920        &self,
921        workflow_id: &str,
922        archive_uri: &str,
923        archived_at: f64,
924    ) -> Result<()> {
925        let mut tx = self.pool.begin().await?;
926        sqlx::query("DELETE FROM workflow_events WHERE workflow_id = $1")
927            .bind(workflow_id)
928            .execute(&mut *tx)
929            .await?;
930        sqlx::query("DELETE FROM workflow_activities WHERE workflow_id = $1")
931            .bind(workflow_id)
932            .execute(&mut *tx)
933            .await?;
934        sqlx::query("DELETE FROM workflow_timers WHERE workflow_id = $1")
935            .bind(workflow_id)
936            .execute(&mut *tx)
937            .await?;
938        sqlx::query("DELETE FROM workflow_signals WHERE workflow_id = $1")
939            .bind(workflow_id)
940            .execute(&mut *tx)
941            .await?;
942        sqlx::query("DELETE FROM workflow_snapshots WHERE workflow_id = $1")
943            .bind(workflow_id)
944            .execute(&mut *tx)
945            .await?;
946        sqlx::query(
947            "UPDATE workflows SET archived_at = $1, archive_uri = $2 WHERE id = $3",
948        )
949        .bind(archived_at)
950        .bind(archive_uri)
951        .bind(workflow_id)
952        .execute(&mut *tx)
953        .await?;
954        tx.commit().await?;
955        Ok(())
956    }
957
958    async fn upsert_search_attributes(
959        &self,
960        workflow_id: &str,
961        patch_json: &str,
962    ) -> Result<()> {
963        let current: Option<(Option<String>,)> =
964            sqlx::query_as("SELECT search_attributes FROM workflows WHERE id = $1")
965                .bind(workflow_id)
966                .fetch_optional(&self.pool)
967                .await?;
968        let merged = crate::store::sqlite::merge_search_attrs(
969            current.and_then(|(s,)| s).as_deref(),
970            patch_json,
971        )?;
972        sqlx::query("UPDATE workflows SET search_attributes = $1 WHERE id = $2")
973            .bind(merged)
974            .bind(workflow_id)
975            .execute(&self.pool)
976            .await?;
977        Ok(())
978    }
979
980    async fn update_schedule(
981        &self,
982        namespace: &str,
983        name: &str,
984        patch: &SchedulePatch,
985    ) -> Result<Option<WorkflowSchedule>> {
986        let mut sets: Vec<String> = Vec::new();
987        let mut idx = 1usize;
988        if patch.cron_expr.is_some() {
989            sets.push(format!("cron_expr = ${idx}"));
990            idx += 1;
991        }
992        if patch.timezone.is_some() {
993            sets.push(format!("timezone = ${idx}"));
994            idx += 1;
995        }
996        if patch.input.is_some() {
997            sets.push(format!("input = ${idx}"));
998            idx += 1;
999        }
1000        if patch.task_queue.is_some() {
1001            sets.push(format!("task_queue = ${idx}"));
1002            idx += 1;
1003        }
1004        if patch.overlap_policy.is_some() {
1005            sets.push(format!("overlap_policy = ${idx}"));
1006            idx += 1;
1007        }
1008        if sets.is_empty() {
1009            return self.get_schedule(namespace, name).await;
1010        }
1011        let sql = format!(
1012            "UPDATE workflow_schedules SET {} WHERE namespace = ${} AND name = ${}",
1013            sets.join(", "),
1014            idx,
1015            idx + 1
1016        );
1017        let mut q = sqlx::query(&sql);
1018        if let Some(ref v) = patch.cron_expr {
1019            q = q.bind(v);
1020        }
1021        if let Some(ref v) = patch.timezone {
1022            q = q.bind(v);
1023        }
1024        if let Some(ref v) = patch.input {
1025            q = q.bind(v.to_string());
1026        }
1027        if let Some(ref v) = patch.task_queue {
1028            q = q.bind(v);
1029        }
1030        if let Some(ref v) = patch.overlap_policy {
1031            q = q.bind(v);
1032        }
1033        let res = q
1034            .bind(namespace)
1035            .bind(name)
1036            .execute(&self.pool)
1037            .await?;
1038        if res.rows_affected() == 0 {
1039            return Ok(None);
1040        }
1041        self.get_schedule(namespace, name).await
1042    }
1043
1044    async fn set_schedule_paused(
1045        &self,
1046        namespace: &str,
1047        name: &str,
1048        paused: bool,
1049    ) -> Result<Option<WorkflowSchedule>> {
1050        let res = sqlx::query(
1051            "UPDATE workflow_schedules SET paused = $1 WHERE namespace = $2 AND name = $3",
1052        )
1053        .bind(paused)
1054        .bind(namespace)
1055        .bind(name)
1056        .execute(&self.pool)
1057        .await?;
1058        if res.rows_affected() == 0 {
1059            return Ok(None);
1060        }
1061        self.get_schedule(namespace, name).await
1062    }
1063
1064    // ── Workers ─────────────────────────────────────────────
1065
1066    async fn register_worker(&self, w: &WorkflowWorker) -> Result<()> {
1067        sqlx::query(
1068            "INSERT INTO workflow_workers (id, namespace, identity, task_queue, workflows, activities, max_concurrent_workflows, max_concurrent_activities, active_tasks, last_heartbeat, registered_at)
1069             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
1070             ON CONFLICT (id) DO UPDATE SET last_heartbeat = EXCLUDED.last_heartbeat, identity = EXCLUDED.identity",
1071        )
1072        .bind(&w.id)
1073        .bind(&w.namespace)
1074        .bind(&w.identity)
1075        .bind(&w.task_queue)
1076        .bind(&w.workflows)
1077        .bind(&w.activities)
1078        .bind(w.max_concurrent_workflows)
1079        .bind(w.max_concurrent_activities)
1080        .bind(w.active_tasks)
1081        .bind(w.last_heartbeat)
1082        .bind(w.registered_at)
1083        .execute(&self.pool)
1084        .await?;
1085        Ok(())
1086    }
1087
1088    async fn heartbeat_worker(&self, id: &str, now: f64) -> Result<()> {
1089        sqlx::query("UPDATE workflow_workers SET last_heartbeat = $1 WHERE id = $2")
1090            .bind(now)
1091            .bind(id)
1092            .execute(&self.pool)
1093            .await?;
1094        Ok(())
1095    }
1096
1097    async fn list_workers(&self, namespace: &str) -> Result<Vec<WorkflowWorker>> {
1098        let rows = sqlx::query_as::<_, PgWorkerRow>(
1099            "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",
1100        )
1101        .bind(namespace)
1102        .fetch_all(&self.pool)
1103        .await?;
1104        Ok(rows.into_iter().map(Into::into).collect())
1105    }
1106
1107    async fn remove_dead_workers(&self, cutoff: f64) -> Result<Vec<String>> {
1108        let rows: Vec<(String,)> =
1109            sqlx::query_as("SELECT id FROM workflow_workers WHERE last_heartbeat < $1")
1110                .bind(cutoff)
1111                .fetch_all(&self.pool)
1112                .await?;
1113        let ids: Vec<String> = rows.into_iter().map(|r| r.0).collect();
1114        if !ids.is_empty() {
1115            sqlx::query("DELETE FROM workflow_workers WHERE last_heartbeat < $1")
1116                .bind(cutoff)
1117                .execute(&self.pool)
1118                .await?;
1119        }
1120        Ok(ids)
1121    }
1122
1123    // ── API Keys ────────────────────────────────────────────
1124
1125    async fn create_api_key(
1126        &self,
1127        key_hash: &str,
1128        prefix: &str,
1129        label: Option<&str>,
1130        created_at: f64,
1131    ) -> Result<()> {
1132        sqlx::query("INSERT INTO api_keys (key_hash, prefix, label, created_at) VALUES ($1, $2, $3, $4)")
1133            .bind(key_hash)
1134            .bind(prefix)
1135            .bind(label)
1136            .bind(created_at)
1137            .execute(&self.pool)
1138            .await?;
1139        Ok(())
1140    }
1141
1142    async fn validate_api_key(&self, key_hash: &str) -> Result<bool> {
1143        let row: Option<(i64,)> =
1144            sqlx::query_as("SELECT 1::BIGINT FROM api_keys WHERE key_hash = $1")
1145                .bind(key_hash)
1146                .fetch_optional(&self.pool)
1147                .await?;
1148        Ok(row.is_some())
1149    }
1150
1151    async fn list_api_keys(&self) -> Result<Vec<crate::store::ApiKeyRecord>> {
1152        let rows = sqlx::query_as::<_, (String, Option<String>, f64)>(
1153            "SELECT prefix, label, created_at FROM api_keys ORDER BY created_at DESC",
1154        )
1155        .fetch_all(&self.pool)
1156        .await?;
1157        Ok(rows
1158            .into_iter()
1159            .map(|(prefix, label, created_at)| crate::store::ApiKeyRecord {
1160                prefix,
1161                label,
1162                created_at,
1163            })
1164            .collect())
1165    }
1166
1167    async fn revoke_api_key(&self, prefix: &str) -> Result<bool> {
1168        let res = sqlx::query("DELETE FROM api_keys WHERE prefix = $1")
1169            .bind(prefix)
1170            .execute(&self.pool)
1171            .await?;
1172        Ok(res.rows_affected() > 0)
1173    }
1174
1175    async fn api_keys_empty(&self) -> Result<bool> {
1176        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM api_keys")
1177            .fetch_one(&self.pool)
1178            .await?;
1179        Ok(row.0 == 0)
1180    }
1181
1182    async fn get_api_key_by_label(
1183        &self,
1184        label: &str,
1185    ) -> Result<Option<crate::store::ApiKeyRecord>> {
1186        let row: Option<(String, Option<String>, f64)> = sqlx::query_as(
1187            "SELECT prefix, label, created_at FROM api_keys WHERE label = $1 LIMIT 1",
1188        )
1189        .bind(label)
1190        .fetch_optional(&self.pool)
1191        .await?;
1192        Ok(row.map(|(prefix, label, created_at)| crate::store::ApiKeyRecord {
1193            prefix,
1194            label,
1195            created_at,
1196        }))
1197    }
1198
1199    // ── Child Workflows ─────────────────────────────────────
1200
1201    async fn list_child_workflows(&self, parent_id: &str) -> Result<Vec<WorkflowRecord>> {
1202        let rows = sqlx::query_as::<_, PgWorkflowRow>(
1203            "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
1204             FROM workflows WHERE parent_id = $1 ORDER BY created_at ASC",
1205        )
1206        .bind(parent_id)
1207        .fetch_all(&self.pool)
1208        .await?;
1209        Ok(rows.into_iter().map(Into::into).collect())
1210    }
1211
1212    // ── Snapshots ───────────────────────────────────────────
1213
1214    async fn create_snapshot(
1215        &self,
1216        workflow_id: &str,
1217        event_seq: i32,
1218        state_json: &str,
1219    ) -> Result<()> {
1220        sqlx::query(
1221            "INSERT INTO workflow_snapshots (workflow_id, event_seq, state_json, created_at)
1222             VALUES ($1, $2, $3, $4)
1223             ON CONFLICT (workflow_id, event_seq) DO UPDATE SET state_json = EXCLUDED.state_json, created_at = EXCLUDED.created_at",
1224        )
1225        .bind(workflow_id)
1226        .bind(event_seq)
1227        .bind(state_json)
1228        .bind(timestamp_now())
1229        .execute(&self.pool)
1230        .await?;
1231        Ok(())
1232    }
1233
1234    async fn get_latest_snapshot(
1235        &self,
1236        workflow_id: &str,
1237    ) -> Result<Option<WorkflowSnapshot>> {
1238        let row = sqlx::query_as::<_, (String, i32, String, f64)>(
1239            "SELECT workflow_id, event_seq, state_json, created_at
1240             FROM workflow_snapshots WHERE workflow_id = $1
1241             ORDER BY event_seq DESC LIMIT 1",
1242        )
1243        .bind(workflow_id)
1244        .fetch_optional(&self.pool)
1245        .await?;
1246
1247        Ok(row.map(|(workflow_id, event_seq, state_json, created_at)| WorkflowSnapshot {
1248            workflow_id,
1249            event_seq,
1250            state_json,
1251            created_at,
1252        }))
1253    }
1254
1255    // ── Queue Stats ─────────────────────────────────────────
1256
1257    async fn get_queue_stats(&self, namespace: &str) -> Result<Vec<crate::store::QueueStats>> {
1258        let rows = sqlx::query_as::<_, (String, i64, i64, i64)>(
1259            "SELECT
1260                a.task_queue AS queue,
1261                SUM(CASE WHEN a.status = 'PENDING' THEN 1 ELSE 0 END) AS pending,
1262                SUM(CASE WHEN a.status = 'RUNNING' THEN 1 ELSE 0 END) AS running,
1263                (SELECT COUNT(*) FROM workflow_workers w WHERE w.task_queue = a.task_queue AND w.namespace = $1) AS workers
1264             FROM workflow_activities a
1265             JOIN workflows wf ON a.workflow_id = wf.id AND wf.namespace = $1
1266             GROUP BY a.task_queue",
1267        )
1268        .bind(namespace)
1269        .fetch_all(&self.pool)
1270        .await?;
1271
1272        Ok(rows
1273            .into_iter()
1274            .map(|(queue, pending, running, workers)| crate::store::QueueStats {
1275                queue,
1276                pending_activities: pending,
1277                running_activities: running,
1278                workers,
1279            })
1280            .collect())
1281    }
1282
1283    // ── Leader Election ─────────────────────────────────────
1284
1285    async fn try_acquire_scheduler_lock(&self) -> Result<bool> {
1286        // pg_try_advisory_lock is session-scoped — only one connection
1287        // in the pool will hold the lock. In a multi-replica Kubernetes
1288        // deployment, only one pod's connection wins.
1289        let row: (bool,) =
1290            sqlx::query_as("SELECT pg_try_advisory_lock(42)")
1291                .fetch_one(&self.pool)
1292                .await?;
1293        Ok(row.0)
1294    }
1295
1296}
1297
1298fn timestamp_now() -> f64 {
1299    std::time::SystemTime::now()
1300        .duration_since(std::time::UNIX_EPOCH)
1301        .unwrap()
1302        .as_secs_f64()
1303}
1304
1305// ── Postgres row types (sqlx::FromRow) ──────────────────────
1306
1307#[derive(sqlx::FromRow)]
1308struct PgWorkflowRow {
1309    id: String,
1310    namespace: String,
1311    run_id: String,
1312    workflow_type: String,
1313    task_queue: String,
1314    status: String,
1315    input: Option<String>,
1316    result: Option<String>,
1317    error: Option<String>,
1318    parent_id: Option<String>,
1319    claimed_by: Option<String>,
1320    search_attributes: Option<String>,
1321    archived_at: Option<f64>,
1322    archive_uri: Option<String>,
1323    created_at: f64,
1324    updated_at: f64,
1325    completed_at: Option<f64>,
1326}
1327
1328impl From<PgWorkflowRow> for WorkflowRecord {
1329    fn from(r: PgWorkflowRow) -> Self {
1330        Self {
1331            id: r.id,
1332            namespace: r.namespace,
1333            run_id: r.run_id,
1334            workflow_type: r.workflow_type,
1335            task_queue: r.task_queue,
1336            status: r.status,
1337            input: r.input,
1338            result: r.result,
1339            error: r.error,
1340            parent_id: r.parent_id,
1341            claimed_by: r.claimed_by,
1342            search_attributes: r.search_attributes,
1343            archived_at: r.archived_at,
1344            archive_uri: r.archive_uri,
1345            created_at: r.created_at,
1346            updated_at: r.updated_at,
1347            completed_at: r.completed_at,
1348        }
1349    }
1350}
1351
1352#[derive(sqlx::FromRow)]
1353struct PgEventRow {
1354    id: i64,
1355    workflow_id: String,
1356    seq: i32,
1357    event_type: String,
1358    payload: Option<String>,
1359    timestamp: f64,
1360}
1361
1362impl From<PgEventRow> for WorkflowEvent {
1363    fn from(r: PgEventRow) -> Self {
1364        Self {
1365            id: Some(r.id),
1366            workflow_id: r.workflow_id,
1367            seq: r.seq,
1368            event_type: r.event_type,
1369            payload: r.payload,
1370            timestamp: r.timestamp,
1371        }
1372    }
1373}
1374
1375#[derive(sqlx::FromRow)]
1376struct PgActivityRow {
1377    id: i64,
1378    workflow_id: String,
1379    seq: i32,
1380    name: String,
1381    task_queue: String,
1382    input: Option<String>,
1383    status: String,
1384    result: Option<String>,
1385    error: Option<String>,
1386    attempt: i32,
1387    max_attempts: i32,
1388    initial_interval_secs: f64,
1389    backoff_coefficient: f64,
1390    start_to_close_secs: f64,
1391    heartbeat_timeout_secs: Option<f64>,
1392    claimed_by: Option<String>,
1393    scheduled_at: f64,
1394    started_at: Option<f64>,
1395    completed_at: Option<f64>,
1396    last_heartbeat: Option<f64>,
1397}
1398
1399impl From<PgActivityRow> for WorkflowActivity {
1400    fn from(r: PgActivityRow) -> Self {
1401        Self {
1402            id: Some(r.id),
1403            workflow_id: r.workflow_id,
1404            seq: r.seq,
1405            name: r.name,
1406            task_queue: r.task_queue,
1407            input: r.input,
1408            status: r.status,
1409            result: r.result,
1410            error: r.error,
1411            attempt: r.attempt,
1412            max_attempts: r.max_attempts,
1413            initial_interval_secs: r.initial_interval_secs,
1414            backoff_coefficient: r.backoff_coefficient,
1415            start_to_close_secs: r.start_to_close_secs,
1416            heartbeat_timeout_secs: r.heartbeat_timeout_secs,
1417            claimed_by: r.claimed_by,
1418            scheduled_at: r.scheduled_at,
1419            started_at: r.started_at,
1420            completed_at: r.completed_at,
1421            last_heartbeat: r.last_heartbeat,
1422        }
1423    }
1424}
1425
1426#[derive(sqlx::FromRow)]
1427struct PgTimerRow {
1428    id: i64,
1429    workflow_id: String,
1430    seq: i32,
1431    fire_at: f64,
1432    fired: bool,
1433}
1434
1435impl From<PgTimerRow> for WorkflowTimer {
1436    fn from(r: PgTimerRow) -> Self {
1437        Self {
1438            id: Some(r.id),
1439            workflow_id: r.workflow_id,
1440            seq: r.seq,
1441            fire_at: r.fire_at,
1442            fired: r.fired,
1443        }
1444    }
1445}
1446
1447#[derive(sqlx::FromRow)]
1448struct PgSignalRow {
1449    id: i64,
1450    workflow_id: String,
1451    name: String,
1452    payload: Option<String>,
1453    consumed: bool,
1454    received_at: f64,
1455}
1456
1457impl From<PgSignalRow> for WorkflowSignal {
1458    fn from(r: PgSignalRow) -> Self {
1459        Self {
1460            id: Some(r.id),
1461            workflow_id: r.workflow_id,
1462            name: r.name,
1463            payload: r.payload,
1464            consumed: r.consumed,
1465            received_at: r.received_at,
1466        }
1467    }
1468}
1469
1470#[derive(sqlx::FromRow)]
1471struct PgScheduleRow {
1472    namespace: String,
1473    name: String,
1474    workflow_type: String,
1475    cron_expr: String,
1476    timezone: String,
1477    input: Option<String>,
1478    task_queue: String,
1479    overlap_policy: String,
1480    paused: bool,
1481    last_run_at: Option<f64>,
1482    next_run_at: Option<f64>,
1483    last_workflow_id: Option<String>,
1484    created_at: f64,
1485}
1486
1487impl From<PgScheduleRow> for WorkflowSchedule {
1488    fn from(r: PgScheduleRow) -> Self {
1489        Self {
1490            namespace: r.namespace,
1491            name: r.name,
1492            workflow_type: r.workflow_type,
1493            cron_expr: r.cron_expr,
1494            timezone: r.timezone,
1495            input: r.input,
1496            task_queue: r.task_queue,
1497            overlap_policy: r.overlap_policy,
1498            paused: r.paused,
1499            last_run_at: r.last_run_at,
1500            next_run_at: r.next_run_at,
1501            last_workflow_id: r.last_workflow_id,
1502            created_at: r.created_at,
1503        }
1504    }
1505}
1506
1507#[derive(sqlx::FromRow)]
1508struct PgWorkerRow {
1509    id: String,
1510    namespace: String,
1511    identity: String,
1512    task_queue: String,
1513    workflows: Option<String>,
1514    activities: Option<String>,
1515    max_concurrent_workflows: i32,
1516    max_concurrent_activities: i32,
1517    active_tasks: i32,
1518    last_heartbeat: f64,
1519    registered_at: f64,
1520}
1521
1522impl From<PgWorkerRow> for WorkflowWorker {
1523    fn from(r: PgWorkerRow) -> Self {
1524        Self {
1525            id: r.id,
1526            namespace: r.namespace,
1527            identity: r.identity,
1528            task_queue: r.task_queue,
1529            workflows: r.workflows,
1530            activities: r.activities,
1531            max_concurrent_workflows: r.max_concurrent_workflows,
1532            max_concurrent_activities: r.max_concurrent_activities,
1533            active_tasks: r.active_tasks,
1534            last_heartbeat: r.last_heartbeat,
1535            registered_at: r.registered_at,
1536        }
1537    }
1538}
1539
1540#[cfg(test)]
1541mod tests {
1542    use super::*;
1543
1544    #[test]
1545    fn sanitise_schema_keeps_statements_intact() {
1546        let input = "CREATE TABLE foo (x INT);\nCREATE INDEX idx_foo ON foo(x);\n";
1547        let out = sanitise_schema(input);
1548        assert_eq!(out.len(), 2);
1549        assert!(out[0].starts_with("CREATE TABLE foo"));
1550        assert!(out[1].starts_with("CREATE INDEX idx_foo"));
1551    }
1552
1553    #[test]
1554    fn sanitise_schema_drops_pure_comment_lines() {
1555        let input = "-- header comment\nCREATE TABLE foo (x INT);\n-- trailing comment\n";
1556        let out = sanitise_schema(input);
1557        assert_eq!(out.len(), 1);
1558        assert!(out[0].starts_with("CREATE TABLE foo"));
1559    }
1560
1561    #[test]
1562    fn sanitise_schema_ignores_semicolons_inside_comment_prose() {
1563        // Regression: the exact shape that broke v0.11.3–v0.11.5 in production.
1564        // `-- foo; bar` used to split into "foo" and " bar" fragments, the second
1565        // of which was executed as SQL and rejected with `syntax error at or near "bar"`.
1566        let input = "\
1567CREATE TABLE foo (x INT);
1568-- Idempotent across startups; fresh installs pick the column up from the
1569-- CREATE TABLE above so the ADD is a no-op.
1570";
1571        let out = sanitise_schema(input);
1572        assert_eq!(
1573            out.len(),
1574            1,
1575            "expected 1 real statement, got {}: {:?}",
1576            out.len(),
1577            out
1578        );
1579        assert!(out[0].starts_with("CREATE TABLE foo"));
1580    }
1581
1582    #[test]
1583    fn sanitise_schema_drops_indented_comment_lines() {
1584        let input = "  -- indented comment\n\tCREATE TABLE foo (x INT);\n";
1585        let out = sanitise_schema(input);
1586        assert_eq!(out.len(), 1);
1587        assert!(out[0].contains("CREATE TABLE foo"));
1588    }
1589
1590    #[test]
1591    fn sanitise_schema_real_constant_produces_only_ddl() {
1592        // The real SCHEMA constant must not produce any statement whose first
1593        // token isn't a recognised SQL keyword. A prose fragment leaking in
1594        // (e.g. "fresh installs...") means the filter regressed.
1595        for stmt in sanitise_schema(SCHEMA) {
1596            let first_word = stmt
1597                .split_whitespace()
1598                .next()
1599                .expect("non-empty statement")
1600                .to_uppercase();
1601            assert!(
1602                matches!(
1603                    first_word.as_str(),
1604                    "CREATE" | "INSERT" | "UPDATE" | "DROP" | "ALTER" | "WITH"
1605                ),
1606                "SCHEMA produced non-DDL statement starting with {first_word:?}: {stmt:?}"
1607            );
1608        }
1609    }
1610}