Skip to main content

assay_workflow/store/
sqlite.rs

1use anyhow::Result;
2use sqlx::SqlitePool;
3use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
4
5use crate::store::{
6    NamespaceRecord, NamespaceStats, QueueStats, RetryEvent, WorkflowStore, retry_denial,
7};
8use crate::types::*;
9
10const 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 = ? AND status = 'FAILED' ORDER BY seq DESC LIMIT 1";
11const RETRY_ACTIVITY_UPDATE: &str = "UPDATE workflow.activities SET status = 'PENDING', result = NULL, error = NULL, attempt = 1, claimed_by = NULL, scheduled_at = ?, started_at = NULL, completed_at = NULL, last_heartbeat = NULL WHERE id = ? 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";
12
13/// Workflow-module DDL. v0.1.2 schema-qualifies every table to the
14/// `workflow` schema, which on SQLite is an attached database (one
15/// `workflow.db` file per data dir, attached on connect). On PG the
16/// same DDL targets the `workflow` schema.
17///
18/// `engine.events` and `engine.lock` (engine-core infrastructure) live
19/// in the `engine` attachment; engine-core DDL is owned by
20/// `assay_domain::engine::SqliteEngineSchema`. We still bootstrap them
21/// here for v0.1.2 because the workflow store is the embedder for the
22/// `engine.events` notification outbox and for the SQLite single-instance
23/// lock — both pre-date the engine-core schema and stay co-located on
24/// SQLite to keep `SqliteStore::new(url)` self-sufficient for tests.
25const SCHEMA: &str = r#"
26CREATE TABLE IF NOT EXISTS workflow.namespaces (
27    name            TEXT PRIMARY KEY,
28    created_at      REAL NOT NULL
29);
30
31INSERT OR IGNORE INTO workflow.namespaces (name, created_at)
32    VALUES ('main', strftime('%s', 'now'));
33
34CREATE TABLE IF NOT EXISTS workflow.workflows (
35    id              TEXT PRIMARY KEY,
36    namespace       TEXT NOT NULL DEFAULT 'main',
37    run_id          TEXT NOT NULL,
38    workflow_type   TEXT NOT NULL,
39    task_queue      TEXT NOT NULL DEFAULT 'main',
40    status          TEXT NOT NULL DEFAULT 'PENDING',
41    input           TEXT,
42    result          TEXT,
43    error           TEXT,
44    parent_id       TEXT,
45    claimed_by      TEXT,
46    search_attributes TEXT,
47    archived_at     REAL,
48    archive_uri     TEXT,
49    -- Workflow-task dispatch (Phase 9): a workflow is "dispatchable" when
50    -- it has new events a worker needs to replay against. Set true on
51    -- start, on activity completion, on timer fire, on signal arrival.
52    -- Cleared when a worker claims the dispatch lease.
53    needs_dispatch  INTEGER NOT NULL DEFAULT 0,
54    dispatch_claimed_by    TEXT,
55    dispatch_last_heartbeat REAL,
56    created_at      REAL NOT NULL,
57    updated_at      REAL NOT NULL,
58    completed_at    REAL
59);
60CREATE INDEX IF NOT EXISTS workflow.idx_wf_status_queue ON workflows(status, task_queue);
61CREATE INDEX IF NOT EXISTS workflow.idx_wf_namespace ON workflows(namespace);
62CREATE INDEX IF NOT EXISTS workflow.idx_wf_dispatch ON workflows(task_queue, needs_dispatch, dispatch_claimed_by);
63
64CREATE TABLE IF NOT EXISTS workflow.events (
65    id              INTEGER PRIMARY KEY AUTOINCREMENT,
66    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
67    seq             INTEGER NOT NULL,
68    event_type      TEXT NOT NULL,
69    payload         TEXT,
70    timestamp       REAL NOT NULL
71);
72CREATE INDEX IF NOT EXISTS workflow.idx_wf_events_lookup ON events(workflow_id, seq);
73
74CREATE TABLE IF NOT EXISTS workflow.activities (
75    id              INTEGER PRIMARY KEY AUTOINCREMENT,
76    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
77    seq             INTEGER NOT NULL,
78    name            TEXT NOT NULL,
79    task_queue      TEXT NOT NULL DEFAULT 'main',
80    input           TEXT,
81    status          TEXT NOT NULL DEFAULT 'PENDING',
82    result          TEXT,
83    error           TEXT,
84    attempt         INTEGER NOT NULL DEFAULT 1,
85    max_attempts    INTEGER NOT NULL DEFAULT 3,
86    initial_interval_secs   REAL NOT NULL DEFAULT 1,
87    backoff_coefficient     REAL NOT NULL DEFAULT 2,
88    start_to_close_secs     REAL NOT NULL DEFAULT 300,
89    heartbeat_timeout_secs  REAL,
90    claimed_by      TEXT,
91    scheduled_at    REAL NOT NULL,
92    started_at      REAL,
93    completed_at    REAL,
94    last_heartbeat  REAL,
95    UNIQUE (workflow_id, seq)
96);
97CREATE INDEX IF NOT EXISTS workflow.idx_wf_act_pending ON activities(task_queue, status, scheduled_at);
98
99CREATE TABLE IF NOT EXISTS workflow.timers (
100    id              INTEGER PRIMARY KEY AUTOINCREMENT,
101    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
102    seq             INTEGER NOT NULL,
103    fire_at         REAL NOT NULL,
104    fired           INTEGER NOT NULL DEFAULT 0,
105    UNIQUE (workflow_id, seq)
106);
107CREATE INDEX IF NOT EXISTS workflow.idx_wf_timers_due ON timers(fire_at);
108
109CREATE TABLE IF NOT EXISTS workflow.signals (
110    id              INTEGER PRIMARY KEY AUTOINCREMENT,
111    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
112    name            TEXT NOT NULL,
113    payload         TEXT,
114    consumed        INTEGER NOT NULL DEFAULT 0,
115    received_at     REAL NOT NULL
116);
117CREATE INDEX IF NOT EXISTS workflow.idx_wf_signals_lookup ON signals(workflow_id, name, consumed);
118
119CREATE TABLE IF NOT EXISTS workflow.schedules (
120    name            TEXT NOT NULL,
121    namespace       TEXT NOT NULL DEFAULT 'main',
122    workflow_type   TEXT NOT NULL,
123    cron_expr       TEXT NOT NULL,
124    timezone        TEXT NOT NULL DEFAULT 'UTC',
125    input           TEXT,
126    task_queue      TEXT NOT NULL DEFAULT 'main',
127    overlap_policy  TEXT NOT NULL DEFAULT 'skip',
128    paused          INTEGER NOT NULL DEFAULT 0,
129    last_run_at     REAL,
130    next_run_at     REAL,
131    last_workflow_id TEXT,
132    created_at      REAL NOT NULL,
133    PRIMARY KEY (namespace, name)
134);
135
136CREATE TABLE IF NOT EXISTS workflow.workers (
137    id              TEXT PRIMARY KEY,
138    namespace       TEXT NOT NULL DEFAULT 'main',
139    identity        TEXT NOT NULL,
140    task_queue      TEXT NOT NULL,
141    workflows       TEXT,
142    activities      TEXT,
143    max_concurrent_workflows  INTEGER NOT NULL DEFAULT 10,
144    max_concurrent_activities INTEGER NOT NULL DEFAULT 10,
145    active_tasks    INTEGER NOT NULL DEFAULT 0,
146    last_heartbeat  REAL NOT NULL,
147    registered_at   REAL NOT NULL
148);
149
150CREATE TABLE IF NOT EXISTS workflow.snapshots (
151    workflow_id     TEXT NOT NULL REFERENCES workflows(id),
152    event_seq       INTEGER NOT NULL,
153    state_json      TEXT NOT NULL,
154    created_at      REAL NOT NULL,
155    PRIMARY KEY (workflow_id, event_seq)
156);
157
158-- workflow.api_keys retired in plan-15 slice 3 (auth tokens come from
159-- the auth module).
160DROP TABLE IF EXISTS workflow.api_keys;
161
162CREATE TABLE IF NOT EXISTS engine.lock (
163    id              INTEGER PRIMARY KEY CHECK (id = 1),
164    instance_id     TEXT NOT NULL,
165    started_at      REAL NOT NULL,
166    last_heartbeat  REAL NOT NULL
167);
168
169CREATE TABLE IF NOT EXISTS engine.events (
170    id              INTEGER PRIMARY KEY AUTOINCREMENT,
171    ts              REAL NOT NULL DEFAULT (CAST(strftime('%s','now') AS REAL)),
172    namespace       TEXT NOT NULL,
173    subsystem       TEXT NOT NULL,
174    kind            TEXT NOT NULL,
175    payload         TEXT NOT NULL DEFAULT '{}'
176);
177CREATE INDEX IF NOT EXISTS engine.idx_engine_events_ns_id ON events(namespace, id);
178CREATE INDEX IF NOT EXISTS engine.idx_engine_events_ts_prune ON events(ts);
179"#;
180
181/// Stale lock timeout — if the lock holder hasn't heartbeated in this
182/// many seconds, assume it's dead and allow takeover.
183const LOCK_STALE_SECS: f64 = 60.0;
184/// How often to refresh the lock heartbeat.
185const LOCK_HEARTBEAT_SECS: u64 = 15;
186
187/// `Clone` is derived because the underlying `SqlitePool` is itself
188/// `Clone` (it's `Arc<PoolInner>` internally) — cloning the store hands
189/// back a new wrapper around the same connection pool. The
190/// `instance_id` is per-store identity (heartbeat row tag), shared
191/// across clones so all clones look like the same instance to
192/// `engine.lock`.
193#[derive(Clone)]
194pub struct SqliteStore {
195    pool: SqlitePool,
196    instance_id: String,
197}
198
199/// Build a fresh [`SqlitePool`] with `engine` + `workflow` ATTACHed to
200/// in-memory shared-cache databases (one alias per pool). Each connection
201/// in the pool inherits the same ATTACHed databases via `after_connect`.
202///
203/// This is the test-friendly path for `SqliteStore::new(url)` callers that
204/// pass `sqlite::memory:` or any path-based URL — every connection sees
205/// the same `engine.*` / `workflow.*` data because the shared-cache URI
206/// pins the in-memory DB to a process-global name.
207///
208/// Production embedders (the engine binary) build their own pool with
209/// file-backed ATTACHes (`<data_dir>/engine.db`, `<data_dir>/workflow.db`)
210/// and call [`SqliteStore::from_attached_pool`] instead.
211async fn build_default_pool(url: &str) -> Result<SqlitePool> {
212    use std::str::FromStr;
213    use std::sync::atomic::{AtomicU64, Ordering};
214
215    static SEQ: AtomicU64 = AtomicU64::new(0);
216    let suffix = format!(
217        "{}_{}",
218        std::process::id(),
219        SEQ.fetch_add(1, Ordering::Relaxed)
220    );
221    let engine_alias = format!("file:assay_engine_{suffix}?mode=memory&cache=shared");
222    let workflow_alias = format!("file:assay_workflow_{suffix}?mode=memory&cache=shared");
223
224    let opts = SqliteConnectOptions::from_str(url)?.create_if_missing(true);
225
226    let pool = SqlitePoolOptions::new()
227        .max_connections(1)
228        .after_connect(move |conn, _meta| {
229            let engine_alias = engine_alias.clone();
230            let workflow_alias = workflow_alias.clone();
231            Box::pin(async move {
232                use sqlx::Executor;
233                conn.execute(format!("ATTACH DATABASE '{engine_alias}' AS engine").as_str())
234                    .await?;
235                conn.execute(format!("ATTACH DATABASE '{workflow_alias}' AS workflow").as_str())
236                    .await?;
237                Ok(())
238            })
239        })
240        .connect_with(opts)
241        .await?;
242    Ok(pool)
243}
244
245impl SqliteStore {
246    /// Open a SqliteStore at `url`. Provisions an in-memory `engine` +
247    /// `workflow` ATTACH automatically — convenient for tests and
248    /// embedders that don't need persistent module isolation. Production
249    /// deployments use [`SqliteStore::from_attached_pool`] with the
250    /// engine-controlled pool that ATTACHes to `<data_dir>/*.db` files.
251    pub async fn new(url: &str) -> Result<Self> {
252        let pool = build_default_pool(url).await?;
253        Self::from_attached_pool(pool).await
254    }
255
256    /// Construct from an externally-managed pool that already has the
257    /// `engine` and `workflow` databases ATTACHed. The engine binary
258    /// uses this — its pool's `after_connect` hook ATTACHes the
259    /// per-module file paths from `[backend].data_dir`.
260    pub async fn from_attached_pool(pool: SqlitePool) -> Result<Self> {
261        let instance_id = format!("assay-{:016x}", {
262            use std::collections::hash_map::DefaultHasher;
263            use std::hash::{Hash, Hasher};
264            let mut h = DefaultHasher::new();
265            std::time::SystemTime::now().hash(&mut h);
266            std::process::id().hash(&mut h);
267            h.finish()
268        });
269        let store = Self { pool, instance_id };
270        store.migrate().await?;
271        Ok(store)
272    }
273
274    /// Backward-compat alias for [`SqliteStore::from_attached_pool`].
275    /// Older call sites passed a bare pool from `SqlitePool::connect()`;
276    /// after v0.1.2 the pool must already have the engine + workflow
277    /// databases attached. The implementation is identical, kept under
278    /// the legacy name so external embedders don't break on upgrade.
279    pub async fn from_pool(pool: SqlitePool) -> Result<Self> {
280        Self::from_attached_pool(pool).await
281    }
282
283    /// Expose the underlying pool (used by the engine to build an
284    /// `SqliteEngineEventBus` that shares the same connection).
285    pub fn pool(&self) -> &SqlitePool {
286        &self.pool
287    }
288
289    /// Acquire the single-instance engine lock.
290    /// Returns an error if another instance is already running.
291    pub async fn acquire_engine_lock(&self) -> Result<()> {
292        let now = timestamp_now();
293
294        // Try to insert the lock
295        let result = sqlx::query(
296            "INSERT INTO engine.lock (id, instance_id, started_at, last_heartbeat) VALUES (1, ?, ?, ?)",
297        )
298        .bind(&self.instance_id)
299        .bind(now)
300        .bind(now)
301        .execute(&self.pool)
302        .await;
303
304        match result {
305            Ok(_) => Ok(()),
306            Err(_) => {
307                // Lock exists — check if it's stale
308                let row: Option<(String, f64)> = sqlx::query_as(
309                    "SELECT instance_id, last_heartbeat FROM engine.lock WHERE id = 1",
310                )
311                .fetch_optional(&self.pool)
312                .await?;
313
314                if let Some((existing_id, last_hb)) = row {
315                    if now - last_hb > LOCK_STALE_SECS {
316                        // Stale lock — take over
317                        sqlx::query(
318                            "UPDATE engine.lock SET instance_id = ?, started_at = ?, last_heartbeat = ? WHERE id = 1",
319                        )
320                        .bind(&self.instance_id)
321                        .bind(now)
322                        .bind(now)
323                        .execute(&self.pool)
324                        .await?;
325                        tracing::warn!(
326                            "Took over stale engine lock from {existing_id} (last heartbeat {:.0}s ago)",
327                            now - last_hb
328                        );
329                        Ok(())
330                    } else {
331                        let age = now - last_hb;
332                        anyhow::bail!(
333                            "Another assay engine instance is already running (id: {existing_id}, \
334                             last heartbeat {age:.0}s ago).\n\n\
335                             SQLite only supports a single engine instance. For multi-instance \
336                             deployment (Kubernetes, Docker Swarm), use PostgreSQL:\n\n\
337                             \x20 assay serve --backend postgres://user:pass@host:5432/dbname"
338                        );
339                    }
340                } else {
341                    anyhow::bail!("Unexpected engine lock state");
342                }
343            }
344        }
345    }
346
347    /// Refresh the engine lock heartbeat. Called periodically by the engine.
348    pub async fn refresh_engine_lock(&self) -> Result<()> {
349        sqlx::query("UPDATE engine.lock SET last_heartbeat = ? WHERE id = 1 AND instance_id = ?")
350            .bind(timestamp_now())
351            .bind(&self.instance_id)
352            .execute(&self.pool)
353            .await?;
354        Ok(())
355    }
356
357    /// Release the engine lock on shutdown.
358    pub async fn release_engine_lock(&self) -> Result<()> {
359        sqlx::query("DELETE FROM engine.lock WHERE id = 1 AND instance_id = ?")
360            .bind(&self.instance_id)
361            .execute(&self.pool)
362            .await?;
363        Ok(())
364    }
365
366    /// Start background task to keep the lock alive.
367    pub fn spawn_lock_heartbeat(self: &std::sync::Arc<Self>) {
368        let store = std::sync::Arc::clone(self);
369        tokio::spawn(async move {
370            let mut tick =
371                tokio::time::interval(std::time::Duration::from_secs(LOCK_HEARTBEAT_SECS));
372            loop {
373                tick.tick().await;
374                if let Err(e) = store.refresh_engine_lock().await {
375                    tracing::error!("Engine lock heartbeat failed: {e}");
376                }
377            }
378        });
379    }
380
381    /// Apply the baseline schema. SCHEMA's `CREATE TABLE IF NOT EXISTS`
382    /// statements are the source of truth — pre-1.0 we don't carry
383    /// `ALTER TABLE ADD COLUMN` history. For additive migrations later,
384    /// chain a `Self::add_column_if_missing(&self.pool, "<table>",
385    /// "<column>", "<type_def>")` call here before returning.
386    async fn migrate(&self) -> Result<()> {
387        for statement in SCHEMA.split(';') {
388            let trimmed = statement.trim();
389            if !trimmed.is_empty() {
390                sqlx::query(trimmed).execute(&self.pool).await?;
391            }
392        }
393        // Future additive migrations go here; see doc-comment above.
394        Ok(())
395    }
396
397    /// Add a column to an existing table if it's not already there.
398    ///
399    /// SQLite (unlike Postgres) doesn't support `ADD COLUMN IF NOT EXISTS`,
400    /// so we check via `pragma_table_info` before issuing the ALTER. Each
401    /// call is idempotent across startups.
402    ///
403    /// Currently unused — kept as the documented pattern for the first
404    /// additive migration after v0.11.3. Remove `#[allow(dead_code)]` when
405    /// a caller is added.
406    #[allow(dead_code)]
407    async fn add_column_if_missing(
408        pool: &SqlitePool,
409        table: &str,
410        column: &str,
411        type_def: &str,
412    ) -> Result<()> {
413        let exists: Option<(String,)> =
414            sqlx::query_as("SELECT name FROM pragma_table_info(?) WHERE name = ?")
415                .bind(table)
416                .bind(column)
417                .fetch_optional(pool)
418                .await?;
419        if exists.is_none() {
420            let sql = format!("ALTER TABLE {table} ADD COLUMN {column} {type_def}");
421            sqlx::query(&sql).execute(pool).await?;
422        }
423        Ok(())
424    }
425}
426
427impl WorkflowStore for SqliteStore {
428    // ── Namespaces ─────────────────────────────────────────
429
430    async fn create_namespace(&self, name: &str) -> Result<()> {
431        sqlx::query("INSERT INTO workflow.namespaces (name, created_at) VALUES (?, ?)")
432            .bind(name)
433            .bind(timestamp_now())
434            .execute(&self.pool)
435            .await?;
436        Ok(())
437    }
438
439    async fn list_namespaces(&self) -> Result<Vec<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)| NamespaceRecord { name, created_at })
448            .collect())
449    }
450
451    async fn delete_namespace(&self, name: &str) -> Result<bool> {
452        // Mirror PG: 'main' is always available, can't be deleted.
453        let res = sqlx::query("DELETE FROM workflow.namespaces WHERE name = ? AND name != 'main'")
454            .bind(name)
455            .execute(&self.pool)
456            .await?;
457        Ok(res.rows_affected() > 0)
458    }
459
460    async fn get_namespace_stats(&self, namespace: &str) -> Result<NamespaceStats> {
461        let total: (i64,) =
462            sqlx::query_as("SELECT COUNT(*) FROM workflow.workflows WHERE namespace = ?")
463                .bind(namespace)
464                .fetch_one(&self.pool)
465                .await?;
466        let running: (i64,) = sqlx::query_as(
467            "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = ? AND status = 'RUNNING'",
468        )
469        .bind(namespace)
470        .fetch_one(&self.pool)
471        .await?;
472        let pending: (i64,) = sqlx::query_as(
473            "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = ? AND status = 'PENDING'",
474        )
475        .bind(namespace)
476        .fetch_one(&self.pool)
477        .await?;
478        let completed: (i64,) = sqlx::query_as(
479            "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = ? AND status = 'COMPLETED'",
480        )
481        .bind(namespace)
482        .fetch_one(&self.pool)
483        .await?;
484        let failed: (i64,) = sqlx::query_as(
485            "SELECT COUNT(*) FROM workflow.workflows WHERE namespace = ? AND status = 'FAILED'",
486        )
487        .bind(namespace)
488        .fetch_one(&self.pool)
489        .await?;
490        let schedules: (i64,) =
491            sqlx::query_as("SELECT COUNT(*) FROM workflow.schedules WHERE namespace = ?")
492                .bind(namespace)
493                .fetch_one(&self.pool)
494                .await?;
495        let workers: (i64,) =
496            sqlx::query_as("SELECT COUNT(*) FROM workflow.workers WHERE namespace = ?")
497                .bind(namespace)
498                .fetch_one(&self.pool)
499                .await?;
500
501        Ok(NamespaceStats {
502            namespace: namespace.to_string(),
503            total_workflows: total.0,
504            running: running.0,
505            pending: pending.0,
506            completed: completed.0,
507            failed: failed.0,
508            schedules: schedules.0,
509            workers: workers.0,
510        })
511    }
512
513    // ── Workflows ──────────────────────────────────────────
514
515    async fn create_workflow(&self, wf: &WorkflowRecord) -> Result<()> {
516        sqlx::query(
517            "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)
518             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
519        )
520        .bind(&wf.id)
521        .bind(&wf.namespace)
522        .bind(&wf.run_id)
523        .bind(&wf.workflow_type)
524        .bind(&wf.task_queue)
525        .bind(&wf.status)
526        .bind(&wf.input)
527        .bind(&wf.result)
528        .bind(&wf.error)
529        .bind(&wf.parent_id)
530        .bind(&wf.claimed_by)
531        .bind(&wf.search_attributes)
532        .bind(wf.archived_at)
533        .bind(&wf.archive_uri)
534        .bind(wf.created_at)
535        .bind(wf.updated_at)
536        .bind(wf.completed_at)
537        .execute(&self.pool)
538        .await?;
539        Ok(())
540    }
541
542    async fn get_workflow(&self, id: &str) -> Result<Option<WorkflowRecord>> {
543        let row = sqlx::query_as::<_, SqliteWorkflowRow>(
544            "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 = ?",
545        )
546        .bind(id)
547        .fetch_optional(&self.pool)
548        .await?;
549        Ok(row.map(Into::into))
550    }
551
552    async fn list_workflows(
553        &self,
554        namespace: &str,
555        status: Option<WorkflowStatus>,
556        workflow_type: Option<&str>,
557        search_attrs_filter: Option<&str>,
558        limit: i64,
559        offset: i64,
560    ) -> Result<Vec<WorkflowRecord>> {
561        let status_str = status.map(|s| s.to_string());
562
563        // Parse search filter into (key, value) pairs. Each pair adds a
564        // `json_extract(search_attributes, '$.key') = value` predicate so
565        // matches require every filter key to be present in the stored
566        // attributes. Invalid/empty JSON → no filter (all pass).
567        let filter_pairs: Vec<(String, serde_json::Value)> = search_attrs_filter
568            .and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
569            .and_then(|v| v.as_object().cloned())
570            .map(|m| m.into_iter().collect())
571            .unwrap_or_default();
572
573        let mut sql = String::from(
574            "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
575             FROM workflow.workflows
576             WHERE namespace = ?
577               AND (? IS NULL OR status = ?)
578               AND (? IS NULL OR workflow_type = ?)",
579        );
580        for _ in &filter_pairs {
581            sql.push_str(" AND json_extract(search_attributes, '$.' || ?) = ?");
582        }
583        sql.push_str(" ORDER BY created_at DESC LIMIT ? OFFSET ?");
584
585        let mut q = sqlx::query_as::<_, SqliteWorkflowRow>(&sql)
586            .bind(namespace)
587            .bind(&status_str)
588            .bind(&status_str)
589            .bind(workflow_type)
590            .bind(workflow_type);
591        for (key, value) in &filter_pairs {
592            q = q.bind(key.clone());
593            // Bind the JSON value as its string/number representation.
594            // json_extract on a stored JSON string returns its "natural"
595            // SQLite type (text for strings, numeric for numbers), so we
596            // match by the same type.
597            match value {
598                serde_json::Value::String(s) => q = q.bind(s.clone()),
599                serde_json::Value::Number(n) => {
600                    if let Some(i) = n.as_i64() {
601                        q = q.bind(i);
602                    } else if let Some(f) = n.as_f64() {
603                        q = q.bind(f);
604                    } else {
605                        q = q.bind(n.to_string());
606                    }
607                }
608                serde_json::Value::Bool(b) => q = q.bind(*b as i64),
609                _ => q = q.bind(value.to_string()),
610            }
611        }
612        let rows = q.bind(limit).bind(offset).fetch_all(&self.pool).await?;
613        Ok(rows.into_iter().map(Into::into).collect())
614    }
615
616    async fn update_workflow_status(
617        &self,
618        id: &str,
619        status: WorkflowStatus,
620        result: Option<&str>,
621        error: Option<&str>,
622    ) -> Result<()> {
623        let now = timestamp_now();
624        let completed_at = if status.is_terminal() {
625            Some(now)
626        } else {
627            None
628        };
629        sqlx::query(
630            "UPDATE workflow.workflows SET status = ?, result = COALESCE(?, result), error = COALESCE(?, error), updated_at = ?, completed_at = COALESCE(?, completed_at) WHERE id = ?",
631        )
632        .bind(status.to_string())
633        .bind(result)
634        .bind(error)
635        .bind(now)
636        .bind(completed_at)
637        .bind(id)
638        .execute(&self.pool)
639        .await?;
640        Ok(())
641    }
642
643    async fn claim_workflow(&self, id: &str, worker_id: &str) -> Result<bool> {
644        let res = sqlx::query(
645            "UPDATE workflow.workflows SET claimed_by = ?, status = 'RUNNING', updated_at = ? WHERE id = ? AND claimed_by IS NULL",
646        )
647        .bind(worker_id)
648        .bind(timestamp_now())
649        .bind(id)
650        .execute(&self.pool)
651        .await?;
652        Ok(res.rows_affected() > 0)
653    }
654
655    async fn mark_workflow_dispatchable(&self, workflow_id: &str) -> Result<()> {
656        sqlx::query("UPDATE workflow.workflows SET needs_dispatch = 1 WHERE id = ?")
657            .bind(workflow_id)
658            .execute(&self.pool)
659            .await?;
660        Ok(())
661    }
662
663    async fn claim_workflow_task(
664        &self,
665        task_queue: &str,
666        worker_id: &str,
667    ) -> Result<Option<WorkflowRecord>> {
668        let now = timestamp_now();
669        // Atomic: pick the oldest dispatchable + unclaimed workflow on the queue
670        let row = sqlx::query_as::<_, SqliteWorkflowRow>(
671            "UPDATE workflow.workflows
672             SET dispatch_claimed_by = ?, dispatch_last_heartbeat = ?, needs_dispatch = 0
673             WHERE id = (
674                SELECT id FROM workflow.workflows
675                WHERE task_queue = ?
676                  AND needs_dispatch = 1
677                  AND dispatch_claimed_by IS NULL
678                  AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
679                ORDER BY updated_at ASC
680                LIMIT 1
681             )
682             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",
683        )
684        .bind(worker_id)
685        .bind(now)
686        .bind(task_queue)
687        .fetch_optional(&self.pool)
688        .await?;
689        Ok(row.map(Into::into))
690    }
691
692    async fn release_workflow_task(&self, workflow_id: &str, worker_id: &str) -> Result<()> {
693        sqlx::query(
694            "UPDATE workflow.workflows
695             SET dispatch_claimed_by = NULL, dispatch_last_heartbeat = NULL
696             WHERE id = ? AND dispatch_claimed_by = ?",
697        )
698        .bind(workflow_id)
699        .bind(worker_id)
700        .execute(&self.pool)
701        .await?;
702        Ok(())
703    }
704
705    async fn release_stale_dispatch_leases(&self, now: f64, timeout_secs: f64) -> Result<u64> {
706        // Re-arm needs_dispatch so the work goes back into the pool. Don't
707        // touch workflows that have reached a terminal state — those should
708        // never be re-dispatched.
709        let res = sqlx::query(
710            "UPDATE workflow.workflows
711             SET dispatch_claimed_by = NULL,
712                 dispatch_last_heartbeat = NULL,
713                 needs_dispatch = 1
714             WHERE dispatch_claimed_by IS NOT NULL
715               AND (? - dispatch_last_heartbeat) > ?
716               AND status NOT IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')",
717        )
718        .bind(now)
719        .bind(timeout_secs)
720        .execute(&self.pool)
721        .await?;
722        Ok(res.rows_affected())
723    }
724
725    // ── Events ─────────────────────────────────────────────
726
727    async fn append_event(&self, ev: &WorkflowEvent) -> Result<i64> {
728        let res = sqlx::query(
729            "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, timestamp) VALUES (?, ?, ?, ?, ?)",
730        )
731        .bind(&ev.workflow_id)
732        .bind(ev.seq)
733        .bind(&ev.event_type)
734        .bind(&ev.payload)
735        .bind(ev.timestamp)
736        .execute(&self.pool)
737        .await?;
738        Ok(res.last_insert_rowid())
739    }
740
741    async fn list_events(&self, workflow_id: &str) -> Result<Vec<WorkflowEvent>> {
742        let rows = sqlx::query_as::<_, SqliteEventRow>(
743            "SELECT id, workflow_id, seq, event_type, payload, timestamp FROM workflow.events WHERE workflow_id = ? ORDER BY seq ASC",
744        )
745        .bind(workflow_id)
746        .fetch_all(&self.pool)
747        .await?;
748        Ok(rows.into_iter().map(Into::into).collect())
749    }
750
751    async fn list_events_page(
752        &self,
753        workflow_id: &str,
754        cursor: Option<i32>,
755        limit: i64,
756        descending: bool,
757    ) -> Result<Vec<WorkflowEvent>> {
758        let limit = limit.clamp(0, 1_000);
759        if limit == 0 {
760            return Ok(Vec::new());
761        }
762        let rows = if descending {
763            sqlx::query_as::<_, SqliteEventRow>(
764                "SELECT id, workflow_id, seq, event_type, payload, timestamp
765                 FROM workflow.events
766                 WHERE workflow_id = ? AND (? IS NULL OR seq < ?)
767                 ORDER BY seq DESC LIMIT ?",
768            )
769            .bind(workflow_id)
770            .bind(cursor)
771            .bind(cursor)
772            .bind(limit)
773            .fetch_all(&self.pool)
774            .await?
775        } else {
776            sqlx::query_as::<_, SqliteEventRow>(
777                "SELECT id, workflow_id, seq, event_type, payload, timestamp
778                 FROM workflow.events
779                 WHERE workflow_id = ? AND (? IS NULL OR seq > ?)
780                 ORDER BY seq ASC LIMIT ?",
781            )
782            .bind(workflow_id)
783            .bind(cursor)
784            .bind(cursor)
785            .bind(limit)
786            .fetch_all(&self.pool)
787            .await?
788        };
789        Ok(rows.into_iter().map(Into::into).collect())
790    }
791
792    async fn get_event_count(&self, workflow_id: &str) -> Result<i64> {
793        let row: (i64,) =
794            sqlx::query_as("SELECT COUNT(*) FROM workflow.events WHERE workflow_id = ?")
795                .bind(workflow_id)
796                .fetch_one(&self.pool)
797                .await?;
798        Ok(row.0)
799    }
800
801    // ── Activities ──────────────────────────────────────────
802
803    async fn create_activity(&self, act: &WorkflowActivity) -> Result<i64> {
804        let res = sqlx::query(
805            "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)
806             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
807        )
808        .bind(&act.workflow_id)
809        .bind(act.seq)
810        .bind(&act.name)
811        .bind(&act.task_queue)
812        .bind(&act.input)
813        .bind(&act.status)
814        .bind(act.attempt)
815        .bind(act.max_attempts)
816        .bind(act.initial_interval_secs)
817        .bind(act.backoff_coefficient)
818        .bind(act.start_to_close_secs)
819        .bind(act.heartbeat_timeout_secs)
820        .bind(act.scheduled_at)
821        .execute(&self.pool)
822        .await?;
823        Ok(res.last_insert_rowid())
824    }
825
826    async fn get_activity(&self, id: i64) -> Result<Option<WorkflowActivity>> {
827        let row = sqlx::query_as::<_, SqliteActivityRow>(
828            "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
829             FROM workflow.activities WHERE id = ?",
830        )
831        .bind(id)
832        .fetch_optional(&self.pool)
833        .await?;
834        Ok(row.map(Into::into))
835    }
836
837    async fn get_activity_by_workflow_seq(
838        &self,
839        workflow_id: &str,
840        seq: i32,
841    ) -> Result<Option<WorkflowActivity>> {
842        let row = sqlx::query_as::<_, SqliteActivityRow>(
843            "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
844             FROM workflow.activities WHERE workflow_id = ? AND seq = ?",
845        )
846        .bind(workflow_id)
847        .bind(seq)
848        .fetch_optional(&self.pool)
849        .await?;
850        Ok(row.map(Into::into))
851    }
852
853    async fn claim_activity(
854        &self,
855        task_queue: &str,
856        worker_id: &str,
857    ) -> Result<Option<WorkflowActivity>> {
858        let now = timestamp_now();
859        let row = sqlx::query_as::<_, SqliteActivityRow>(
860            "UPDATE workflow.activities SET status = 'RUNNING', claimed_by = ?, started_at = ?
861             WHERE id = (
862                SELECT id FROM workflow.activities
863                WHERE task_queue = ? AND status = 'PENDING'
864                ORDER BY scheduled_at ASC
865                LIMIT 1
866             )
867             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",
868        )
869        .bind(worker_id)
870        .bind(now)
871        .bind(task_queue)
872        .fetch_optional(&self.pool)
873        .await?;
874        Ok(row.map(Into::into))
875    }
876
877    async fn requeue_activity_for_retry(
878        &self,
879        id: i64,
880        next_attempt: i32,
881        next_scheduled_at: f64,
882    ) -> Result<()> {
883        sqlx::query(
884            "UPDATE workflow.activities
885             SET status = 'PENDING', attempt = ?, scheduled_at = ?,
886                 claimed_by = NULL, started_at = NULL, last_heartbeat = NULL,
887                 error = NULL
888             WHERE id = ?",
889        )
890        .bind(next_attempt)
891        .bind(next_scheduled_at)
892        .bind(id)
893        .execute(&self.pool)
894        .await?;
895        Ok(())
896    }
897
898    async fn retry_failed_activity(
899        &self,
900        workflow_id: &str,
901        requested_by: &str,
902        reason: &str,
903        requested_at: f64,
904    ) -> Result<RetryFailedActivityResult> {
905        let mut tx = self.pool.begin().await?;
906        let workflow: Option<(String, Option<String>, Option<f64>)> = sqlx::query_as(
907            "SELECT status, parent_id, archived_at FROM workflow.workflows WHERE id = ?",
908        )
909        .bind(workflow_id)
910        .fetch_optional(&mut *tx)
911        .await?;
912        let Some((status, parent_id, archived_at)) = workflow else {
913            return Ok(RetryFailedActivityResult::NotFound);
914        };
915        if let Some(denial) = retry_denial(status, parent_id, archived_at) {
916            return Ok(denial);
917        }
918
919        let failed = sqlx::query_as::<_, SqliteActivityRow>(RETRY_ACTIVITY_SELECT)
920            .bind(workflow_id)
921            .fetch_optional(&mut *tx)
922            .await?;
923        let Some(failed) = failed else {
924            return Ok(RetryFailedActivityResult::NoFailedActivity);
925        };
926        let failed_event_seq: (i32,) = sqlx::query_as(
927            "SELECT seq FROM workflow.events
928             WHERE workflow_id = ? AND event_type = 'ActivityFailed'
929             ORDER BY seq DESC LIMIT 1",
930        )
931        .bind(workflow_id)
932        .fetch_one(&mut *tx)
933        .await?;
934        let invalidated =
935            sqlx::query("DELETE FROM workflow.activities WHERE workflow_id = ? AND seq > ?")
936                .bind(workflow_id)
937                .bind(failed.seq)
938                .execute(&mut *tx)
939                .await?
940                .rows_affected();
941        let activity = sqlx::query_as::<_, SqliteActivityRow>(RETRY_ACTIVITY_UPDATE)
942            .bind(requested_at)
943            .bind(failed.id)
944            .fetch_one(&mut *tx)
945            .await?;
946        sqlx::query(
947            "UPDATE workflow.workflows
948             SET status = 'WAITING', result = NULL, error = NULL, completed_at = NULL,
949                 updated_at = ?, needs_dispatch = 0, dispatch_claimed_by = NULL,
950                 dispatch_last_heartbeat = NULL
951             WHERE id = ?",
952        )
953        .bind(requested_at)
954        .bind(workflow_id)
955        .execute(&mut *tx)
956        .await?;
957        let event_seq: (i32,) = sqlx::query_as(
958            "SELECT COALESCE(MAX(seq), 0) + 1 FROM workflow.events WHERE workflow_id = ?",
959        )
960        .bind(workflow_id)
961        .fetch_one(&mut *tx)
962        .await?;
963        let payload = RetryEvent {
964            activity_id: failed.id,
965            activity_seq: failed.seq,
966            activity_name: &failed.name,
967            failed_event_seq: failed_event_seq.0,
968            requested_by,
969            reason,
970            invalidated_activities: invalidated,
971        }
972        .payload();
973        sqlx::query(
974            "INSERT INTO workflow.events (workflow_id, seq, event_type, payload, timestamp)
975             VALUES (?, ?, 'ActivityRetryRequested', ?, ?)",
976        )
977        .bind(workflow_id)
978        .bind(event_seq.0)
979        .bind(payload.to_string())
980        .bind(requested_at)
981        .execute(&mut *tx)
982        .await?;
983        tx.commit().await?;
984        Ok(RetryFailedActivityResult::Retried(Box::new(
985            RetriedActivity {
986                activity: activity.into(),
987                invalidated_activities: invalidated,
988            },
989        )))
990    }
991
992    async fn complete_activity(
993        &self,
994        id: i64,
995        result: Option<&str>,
996        error: Option<&str>,
997        failed: bool,
998    ) -> Result<()> {
999        let status = if failed { "FAILED" } else { "COMPLETED" };
1000        sqlx::query(
1001            "UPDATE workflow.activities SET status = ?, result = ?, error = ?, completed_at = ? WHERE id = ?",
1002        )
1003        .bind(status)
1004        .bind(result)
1005        .bind(error)
1006        .bind(timestamp_now())
1007        .bind(id)
1008        .execute(&self.pool)
1009        .await?;
1010        Ok(())
1011    }
1012
1013    async fn heartbeat_activity(&self, id: i64, _details: Option<&str>) -> Result<()> {
1014        sqlx::query("UPDATE workflow.activities SET last_heartbeat = ? WHERE id = ?")
1015            .bind(timestamp_now())
1016            .bind(id)
1017            .execute(&self.pool)
1018            .await?;
1019        Ok(())
1020    }
1021
1022    async fn get_timed_out_activities(&self, now: f64) -> Result<Vec<WorkflowActivity>> {
1023        let rows = sqlx::query_as::<_, SqliteActivityRow>(
1024            "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
1025             FROM workflow.activities
1026             WHERE status = 'RUNNING'
1027               AND heartbeat_timeout_secs IS NOT NULL
1028               AND (? - COALESCE(last_heartbeat, started_at)) > heartbeat_timeout_secs",
1029        )
1030        .bind(now)
1031        .fetch_all(&self.pool)
1032        .await?;
1033        Ok(rows.into_iter().map(Into::into).collect())
1034    }
1035
1036    // ── Timers ──────────────────────────────────────────────
1037
1038    async fn create_timer(&self, timer: &WorkflowTimer) -> Result<i64> {
1039        // Idempotent: INSERT OR IGNORE on UNIQUE (workflow_id, seq).
1040        // If the row already existed, last_insert_rowid() is 0 — fall back to SELECT.
1041        let res = sqlx::query(
1042            "INSERT OR IGNORE INTO workflow.timers (workflow_id, seq, fire_at, fired) VALUES (?, ?, ?, 0)",
1043        )
1044        .bind(&timer.workflow_id)
1045        .bind(timer.seq)
1046        .bind(timer.fire_at)
1047        .execute(&self.pool)
1048        .await?;
1049
1050        let id = res.last_insert_rowid();
1051        if id != 0 {
1052            return Ok(id);
1053        }
1054
1055        // Row already existed — return its id.
1056        let (existing_id,): (i64,) =
1057            sqlx::query_as("SELECT id FROM workflow.timers WHERE workflow_id = ? AND seq = ?")
1058                .bind(&timer.workflow_id)
1059                .bind(timer.seq)
1060                .fetch_one(&self.pool)
1061                .await?;
1062        Ok(existing_id)
1063    }
1064
1065    async fn cancel_pending_activities(&self, workflow_id: &str) -> Result<u64> {
1066        let res = sqlx::query(
1067            "UPDATE workflow.activities SET status = 'CANCELLED', completed_at = ?
1068             WHERE workflow_id = ? AND status = 'PENDING'",
1069        )
1070        .bind(timestamp_now())
1071        .bind(workflow_id)
1072        .execute(&self.pool)
1073        .await?;
1074        Ok(res.rows_affected())
1075    }
1076
1077    async fn cancel_pending_timers(&self, workflow_id: &str) -> Result<u64> {
1078        let res = sqlx::query(
1079            "UPDATE workflow.timers SET fired = 1
1080             WHERE workflow_id = ? AND fired = 0",
1081        )
1082        .bind(workflow_id)
1083        .execute(&self.pool)
1084        .await?;
1085        Ok(res.rows_affected())
1086    }
1087
1088    async fn get_timer_by_workflow_seq(
1089        &self,
1090        workflow_id: &str,
1091        seq: i32,
1092    ) -> Result<Option<WorkflowTimer>> {
1093        let row = sqlx::query_as::<_, SqliteTimerRow>(
1094            "SELECT id, workflow_id, seq, fire_at, fired
1095             FROM workflow.timers WHERE workflow_id = ? AND seq = ?",
1096        )
1097        .bind(workflow_id)
1098        .bind(seq)
1099        .fetch_optional(&self.pool)
1100        .await?;
1101        Ok(row.map(Into::into))
1102    }
1103
1104    async fn fire_due_timers(&self, now: f64) -> Result<Vec<WorkflowTimer>> {
1105        let rows = sqlx::query_as::<_, SqliteTimerRow>(
1106            "UPDATE workflow.timers SET fired = 1
1107             WHERE fired = 0 AND fire_at <= ?
1108             RETURNING id, workflow_id, seq, fire_at, fired",
1109        )
1110        .bind(now)
1111        .fetch_all(&self.pool)
1112        .await?;
1113        Ok(rows.into_iter().map(Into::into).collect())
1114    }
1115
1116    // ── Signals ─────────────────────────────────────────────
1117
1118    async fn send_signal(&self, sig: &WorkflowSignal) -> Result<i64> {
1119        let res = sqlx::query(
1120            "INSERT INTO workflow.signals (workflow_id, name, payload, consumed, received_at) VALUES (?, ?, ?, 0, ?)",
1121        )
1122        .bind(&sig.workflow_id)
1123        .bind(&sig.name)
1124        .bind(&sig.payload)
1125        .bind(sig.received_at)
1126        .execute(&self.pool)
1127        .await?;
1128        Ok(res.last_insert_rowid())
1129    }
1130
1131    async fn consume_signals(&self, workflow_id: &str, name: &str) -> Result<Vec<WorkflowSignal>> {
1132        let rows = sqlx::query_as::<_, SqliteSignalRow>(
1133            "UPDATE workflow.signals SET consumed = 1
1134             WHERE workflow_id = ? AND name = ? AND consumed = 0
1135             RETURNING id, workflow_id, name, payload, consumed, received_at",
1136        )
1137        .bind(workflow_id)
1138        .bind(name)
1139        .fetch_all(&self.pool)
1140        .await?;
1141        Ok(rows.into_iter().map(Into::into).collect())
1142    }
1143
1144    // ── Schedules ───────────────────────────────────────────
1145
1146    async fn create_schedule(&self, sched: &WorkflowSchedule) -> Result<()> {
1147        sqlx::query(
1148            "INSERT INTO workflow.schedules (name, namespace, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at)
1149             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
1150        )
1151        .bind(&sched.name)
1152        .bind(&sched.namespace)
1153        .bind(&sched.workflow_type)
1154        .bind(&sched.cron_expr)
1155        .bind(&sched.timezone)
1156        .bind(&sched.input)
1157        .bind(&sched.task_queue)
1158        .bind(&sched.overlap_policy)
1159        .bind(sched.paused)
1160        .bind(sched.last_run_at)
1161        .bind(sched.next_run_at)
1162        .bind(&sched.last_workflow_id)
1163        .bind(sched.created_at)
1164        .execute(&self.pool)
1165        .await?;
1166        Ok(())
1167    }
1168
1169    async fn get_schedule(&self, namespace: &str, name: &str) -> Result<Option<WorkflowSchedule>> {
1170        let row = sqlx::query_as::<_, SqliteScheduleRow>(
1171            "SELECT name, namespace, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at
1172             FROM workflow.schedules WHERE namespace = ? AND name = ?",
1173        )
1174        .bind(namespace)
1175        .bind(name)
1176        .fetch_optional(&self.pool)
1177        .await?;
1178        Ok(row.map(Into::into))
1179    }
1180
1181    async fn list_schedules(&self, namespace: &str) -> Result<Vec<WorkflowSchedule>> {
1182        let rows = sqlx::query_as::<_, SqliteScheduleRow>(
1183            "SELECT name, namespace, workflow_type, cron_expr, timezone, input, task_queue, overlap_policy, paused, last_run_at, next_run_at, last_workflow_id, created_at
1184             FROM workflow.schedules WHERE namespace = ? ORDER BY name",
1185        )
1186        .bind(namespace)
1187        .fetch_all(&self.pool)
1188        .await?;
1189        Ok(rows.into_iter().map(Into::into).collect())
1190    }
1191
1192    async fn update_schedule_last_run(
1193        &self,
1194        namespace: &str,
1195        name: &str,
1196        last_run_at: f64,
1197        next_run_at: f64,
1198        workflow_id: &str,
1199    ) -> Result<()> {
1200        sqlx::query(
1201            "UPDATE workflow.schedules SET last_run_at = ?, next_run_at = ?, last_workflow_id = ? WHERE namespace = ? AND name = ?",
1202        )
1203        .bind(last_run_at)
1204        .bind(next_run_at)
1205        .bind(workflow_id)
1206        .bind(namespace)
1207        .bind(name)
1208        .execute(&self.pool)
1209        .await?;
1210        Ok(())
1211    }
1212
1213    async fn delete_schedule(&self, namespace: &str, name: &str) -> Result<bool> {
1214        let res = sqlx::query("DELETE FROM workflow.schedules WHERE namespace = ? AND name = ?")
1215            .bind(namespace)
1216            .bind(name)
1217            .execute(&self.pool)
1218            .await?;
1219        Ok(res.rows_affected() > 0)
1220    }
1221
1222    async fn list_archivable_workflows(
1223        &self,
1224        cutoff: f64,
1225        limit: i64,
1226    ) -> Result<Vec<WorkflowRecord>> {
1227        let rows = sqlx::query_as::<_, SqliteWorkflowRow>(
1228            "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
1229             FROM workflow.workflows
1230             WHERE status IN ('COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT')
1231               AND completed_at IS NOT NULL
1232               AND completed_at < ?
1233               AND archived_at IS NULL
1234             ORDER BY completed_at ASC
1235             LIMIT ?",
1236        )
1237        .bind(cutoff)
1238        .bind(limit)
1239        .fetch_all(&self.pool)
1240        .await?;
1241        Ok(rows.into_iter().map(Into::into).collect())
1242    }
1243
1244    async fn mark_archived_and_purge(
1245        &self,
1246        workflow_id: &str,
1247        archive_uri: &str,
1248        archived_at: f64,
1249    ) -> Result<()> {
1250        let mut tx = self.pool.begin().await?;
1251        sqlx::query("DELETE FROM workflow.events WHERE workflow_id = ?")
1252            .bind(workflow_id)
1253            .execute(&mut *tx)
1254            .await?;
1255        sqlx::query("DELETE FROM workflow.activities WHERE workflow_id = ?")
1256            .bind(workflow_id)
1257            .execute(&mut *tx)
1258            .await?;
1259        sqlx::query("DELETE FROM workflow.timers WHERE workflow_id = ?")
1260            .bind(workflow_id)
1261            .execute(&mut *tx)
1262            .await?;
1263        sqlx::query("DELETE FROM workflow.signals WHERE workflow_id = ?")
1264            .bind(workflow_id)
1265            .execute(&mut *tx)
1266            .await?;
1267        sqlx::query("DELETE FROM workflow.snapshots WHERE workflow_id = ?")
1268            .bind(workflow_id)
1269            .execute(&mut *tx)
1270            .await?;
1271        sqlx::query("UPDATE workflow.workflows SET archived_at = ?, archive_uri = ? WHERE id = ?")
1272            .bind(archived_at)
1273            .bind(archive_uri)
1274            .bind(workflow_id)
1275            .execute(&mut *tx)
1276            .await?;
1277        tx.commit().await?;
1278        Ok(())
1279    }
1280
1281    async fn upsert_search_attributes(&self, workflow_id: &str, patch_json: &str) -> Result<()> {
1282        // Merge at the application layer so we don't depend on SQLite's
1283        // `json_patch`, which is only available with the json1 extension.
1284        let current: Option<(Option<String>,)> =
1285            sqlx::query_as("SELECT search_attributes FROM workflow.workflows WHERE id = ?")
1286                .bind(workflow_id)
1287                .fetch_optional(&self.pool)
1288                .await?;
1289        let merged = merge_search_attrs(current.and_then(|(s,)| s).as_deref(), patch_json)?;
1290        sqlx::query("UPDATE workflow.workflows SET search_attributes = ? WHERE id = ?")
1291            .bind(merged)
1292            .bind(workflow_id)
1293            .execute(&self.pool)
1294            .await?;
1295        Ok(())
1296    }
1297
1298    async fn update_schedule(
1299        &self,
1300        namespace: &str,
1301        name: &str,
1302        patch: &SchedulePatch,
1303    ) -> Result<Option<WorkflowSchedule>> {
1304        // Build the UPDATE dynamically so unchanged fields aren't touched
1305        // and NULL from `serde_json::Value::Null` round-trips cleanly.
1306        let mut sets: Vec<&'static str> = Vec::new();
1307        if patch.cron_expr.is_some() {
1308            sets.push("cron_expr = ?");
1309        }
1310        if patch.timezone.is_some() {
1311            sets.push("timezone = ?");
1312        }
1313        if patch.input.is_some() {
1314            sets.push("input = ?");
1315        }
1316        if patch.task_queue.is_some() {
1317            sets.push("task_queue = ?");
1318        }
1319        if patch.overlap_policy.is_some() {
1320            sets.push("overlap_policy = ?");
1321        }
1322        // Updating last_run_at/next_run_at is internal only (update_schedule_last_run).
1323        if sets.is_empty() {
1324            return self.get_schedule(namespace, name).await;
1325        }
1326
1327        let sql = format!(
1328            "UPDATE workflow.schedules SET {} WHERE namespace = ? AND name = ?",
1329            sets.join(", ")
1330        );
1331        let mut q = sqlx::query(&sql);
1332        if let Some(ref v) = patch.cron_expr {
1333            q = q.bind(v);
1334        }
1335        if let Some(ref v) = patch.timezone {
1336            q = q.bind(v);
1337        }
1338        if let Some(ref v) = patch.input {
1339            q = q.bind(v.to_string());
1340        }
1341        if let Some(ref v) = patch.task_queue {
1342            q = q.bind(v);
1343        }
1344        if let Some(ref v) = patch.overlap_policy {
1345            q = q.bind(v);
1346        }
1347        let res = q.bind(namespace).bind(name).execute(&self.pool).await?;
1348        if res.rows_affected() == 0 {
1349            return Ok(None);
1350        }
1351        self.get_schedule(namespace, name).await
1352    }
1353
1354    async fn set_schedule_paused(
1355        &self,
1356        namespace: &str,
1357        name: &str,
1358        paused: bool,
1359    ) -> Result<Option<WorkflowSchedule>> {
1360        let res = sqlx::query(
1361            "UPDATE workflow.schedules SET paused = ? WHERE namespace = ? AND name = ?",
1362        )
1363        .bind(paused)
1364        .bind(namespace)
1365        .bind(name)
1366        .execute(&self.pool)
1367        .await?;
1368        if res.rows_affected() == 0 {
1369            return Ok(None);
1370        }
1371        self.get_schedule(namespace, name).await
1372    }
1373
1374    // ── Workers ─────────────────────────────────────────────
1375
1376    async fn register_worker(&self, w: &WorkflowWorker) -> Result<()> {
1377        sqlx::query(
1378            "INSERT OR REPLACE INTO workflow.workers (id, namespace, identity, task_queue, workflows, activities, max_concurrent_workflows, max_concurrent_activities, active_tasks, last_heartbeat, registered_at)
1379             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
1380        )
1381        .bind(&w.id)
1382        .bind(&w.namespace)
1383        .bind(&w.identity)
1384        .bind(&w.task_queue)
1385        .bind(&w.workflows)
1386        .bind(&w.activities)
1387        .bind(w.max_concurrent_workflows)
1388        .bind(w.max_concurrent_activities)
1389        .bind(w.active_tasks)
1390        .bind(w.last_heartbeat)
1391        .bind(w.registered_at)
1392        .execute(&self.pool)
1393        .await?;
1394        Ok(())
1395    }
1396
1397    async fn heartbeat_worker(&self, id: &str, now: f64) -> Result<()> {
1398        sqlx::query("UPDATE workflow.workers SET last_heartbeat = ? WHERE id = ?")
1399            .bind(now)
1400            .bind(id)
1401            .execute(&self.pool)
1402            .await?;
1403        Ok(())
1404    }
1405
1406    async fn list_workers(&self, namespace: &str) -> Result<Vec<WorkflowWorker>> {
1407        let rows = sqlx::query_as::<_, SqliteWorkerRow>(
1408            "SELECT id, namespace, identity, task_queue, workflows, activities, max_concurrent_workflows, max_concurrent_activities, active_tasks, last_heartbeat, registered_at
1409             FROM workflow.workers WHERE namespace = ? ORDER BY registered_at",
1410        )
1411        .bind(namespace)
1412        .fetch_all(&self.pool)
1413        .await?;
1414        Ok(rows.into_iter().map(Into::into).collect())
1415    }
1416
1417    async fn remove_dead_workers(&self, cutoff: f64) -> Result<Vec<String>> {
1418        let rows: Vec<(String,)> =
1419            sqlx::query_as("SELECT id FROM workflow.workers WHERE last_heartbeat < ?")
1420                .bind(cutoff)
1421                .fetch_all(&self.pool)
1422                .await?;
1423        let ids: Vec<String> = rows.into_iter().map(|r| r.0).collect();
1424        if !ids.is_empty() {
1425            sqlx::query("DELETE FROM workflow.workers WHERE last_heartbeat < ?")
1426                .bind(cutoff)
1427                .execute(&self.pool)
1428                .await?;
1429        }
1430        Ok(ids)
1431    }
1432
1433    // ── Child Workflows ─────────────────────────────────────
1434
1435    async fn list_child_workflows(&self, parent_id: &str) -> Result<Vec<WorkflowRecord>> {
1436        let rows = sqlx::query_as::<_, SqliteWorkflowRow>(
1437            "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
1438             FROM workflow.workflows WHERE parent_id = ? ORDER BY created_at ASC",
1439        )
1440        .bind(parent_id)
1441        .fetch_all(&self.pool)
1442        .await?;
1443        Ok(rows.into_iter().map(Into::into).collect())
1444    }
1445
1446    // ── Snapshots ───────────────────────────────────────────
1447
1448    async fn create_snapshot(
1449        &self,
1450        workflow_id: &str,
1451        event_seq: i32,
1452        state_json: &str,
1453    ) -> Result<()> {
1454        sqlx::query(
1455            "INSERT OR REPLACE INTO workflow.snapshots (workflow_id, event_seq, state_json, created_at)
1456             VALUES (?, ?, ?, ?)",
1457        )
1458        .bind(workflow_id)
1459        .bind(event_seq)
1460        .bind(state_json)
1461        .bind(timestamp_now())
1462        .execute(&self.pool)
1463        .await?;
1464        Ok(())
1465    }
1466
1467    async fn get_latest_snapshot(&self, workflow_id: &str) -> Result<Option<WorkflowSnapshot>> {
1468        let row = sqlx::query_as::<_, (String, i32, String, f64)>(
1469            "SELECT workflow_id, event_seq, state_json, created_at
1470             FROM workflow.snapshots WHERE workflow_id = ?
1471             ORDER BY event_seq DESC LIMIT 1",
1472        )
1473        .bind(workflow_id)
1474        .fetch_optional(&self.pool)
1475        .await?;
1476
1477        Ok(row.map(
1478            |(workflow_id, event_seq, state_json, created_at)| WorkflowSnapshot {
1479                workflow_id,
1480                event_seq,
1481                state_json,
1482                created_at,
1483            },
1484        ))
1485    }
1486
1487    // ── Queue Stats ─────────────────────────────────────────
1488
1489    async fn get_queue_stats(&self, namespace: &str) -> Result<Vec<QueueStats>> {
1490        // Gather activity stats per queue for workflows in this namespace
1491        let rows = sqlx::query_as::<_, (String, i64, i64)>(
1492            "SELECT a.task_queue,
1493                    SUM(CASE WHEN a.status = 'PENDING' THEN 1 ELSE 0 END),
1494                    SUM(CASE WHEN a.status = 'RUNNING' THEN 1 ELSE 0 END)
1495             FROM workflow.activities a
1496             INNER JOIN workflow.workflows w ON w.id = a.workflow_id
1497             WHERE w.namespace = ?
1498             GROUP BY a.task_queue",
1499        )
1500        .bind(namespace)
1501        .fetch_all(&self.pool)
1502        .await?;
1503
1504        let mut stats: Vec<QueueStats> = rows
1505            .into_iter()
1506            .map(|(queue, pending, running)| QueueStats {
1507                queue,
1508                pending_activities: pending,
1509                running_activities: running,
1510                workers: 0,
1511            })
1512            .collect();
1513
1514        // Gather worker counts per queue in this namespace
1515        let worker_rows = sqlx::query_as::<_, (String, i64)>(
1516            "SELECT task_queue, COUNT(*) FROM workflow.workers WHERE namespace = ? GROUP BY task_queue",
1517        )
1518        .bind(namespace)
1519        .fetch_all(&self.pool)
1520        .await?;
1521
1522        for (queue, count) in worker_rows {
1523            if let Some(s) = stats.iter_mut().find(|s| s.queue == queue) {
1524                s.workers = count;
1525            } else {
1526                stats.push(QueueStats {
1527                    queue,
1528                    pending_activities: 0,
1529                    running_activities: 0,
1530                    workers: count,
1531                });
1532            }
1533        }
1534
1535        stats.sort_by(|a, b| a.queue.cmp(&b.queue));
1536        Ok(stats)
1537    }
1538
1539    // ── Leader Election ─────────────────────────────────────
1540
1541    async fn try_acquire_scheduler_lock(&self) -> Result<bool> {
1542        // SQLite is single-instance — always the leader.
1543        // Also refresh the engine lock heartbeat on each scheduler tick.
1544        self.refresh_engine_lock().await.ok();
1545        Ok(true)
1546    }
1547}
1548
1549fn timestamp_now() -> f64 {
1550    std::time::SystemTime::now()
1551        .duration_since(std::time::UNIX_EPOCH)
1552        .unwrap()
1553        .as_secs_f64()
1554}
1555
1556/// Merge a JSON-object patch into a (possibly-null) current JSON object,
1557/// returning the serialised result. Shared by SQLite and Postgres stores.
1558pub(crate) fn merge_search_attrs(current: Option<&str>, patch_json: &str) -> Result<String> {
1559    let mut current_map: serde_json::Map<String, serde_json::Value> = current
1560        .and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
1561        .and_then(|v| v.as_object().cloned())
1562        .unwrap_or_default();
1563    let patch: serde_json::Value = serde_json::from_str(patch_json)
1564        .map_err(|e| anyhow::anyhow!("invalid search_attributes patch: {e}"))?;
1565    let patch_obj = patch
1566        .as_object()
1567        .ok_or_else(|| anyhow::anyhow!("search_attributes patch must be a JSON object"))?;
1568    for (k, v) in patch_obj {
1569        current_map.insert(k.clone(), v.clone());
1570    }
1571    Ok(serde_json::Value::Object(current_map).to_string())
1572}
1573
1574// ── SQLite row types (sqlx::FromRow) ────────────────────────
1575
1576#[derive(sqlx::FromRow)]
1577struct SqliteWorkflowRow {
1578    id: String,
1579    namespace: String,
1580    run_id: String,
1581    workflow_type: String,
1582    task_queue: String,
1583    status: String,
1584    input: Option<String>,
1585    result: Option<String>,
1586    error: Option<String>,
1587    parent_id: Option<String>,
1588    claimed_by: Option<String>,
1589    search_attributes: Option<String>,
1590    archived_at: Option<f64>,
1591    archive_uri: Option<String>,
1592    created_at: f64,
1593    updated_at: f64,
1594    completed_at: Option<f64>,
1595}
1596
1597impl From<SqliteWorkflowRow> for WorkflowRecord {
1598    fn from(r: SqliteWorkflowRow) -> Self {
1599        Self {
1600            id: r.id,
1601            namespace: r.namespace,
1602            run_id: r.run_id,
1603            workflow_type: r.workflow_type,
1604            task_queue: r.task_queue,
1605            status: r.status,
1606            input: r.input,
1607            result: r.result,
1608            error: r.error,
1609            parent_id: r.parent_id,
1610            claimed_by: r.claimed_by,
1611            search_attributes: r.search_attributes,
1612            archived_at: r.archived_at,
1613            archive_uri: r.archive_uri,
1614            created_at: r.created_at,
1615            updated_at: r.updated_at,
1616            completed_at: r.completed_at,
1617        }
1618    }
1619}
1620
1621#[derive(sqlx::FromRow)]
1622struct SqliteEventRow {
1623    id: i64,
1624    workflow_id: String,
1625    seq: i32,
1626    event_type: String,
1627    payload: Option<String>,
1628    timestamp: f64,
1629}
1630
1631impl From<SqliteEventRow> for WorkflowEvent {
1632    fn from(r: SqliteEventRow) -> Self {
1633        Self {
1634            id: Some(r.id),
1635            workflow_id: r.workflow_id,
1636            seq: r.seq,
1637            event_type: r.event_type,
1638            payload: r.payload,
1639            timestamp: r.timestamp,
1640        }
1641    }
1642}
1643
1644#[derive(sqlx::FromRow)]
1645struct SqliteActivityRow {
1646    id: i64,
1647    workflow_id: String,
1648    seq: i32,
1649    name: String,
1650    task_queue: String,
1651    input: Option<String>,
1652    status: String,
1653    result: Option<String>,
1654    error: Option<String>,
1655    attempt: i32,
1656    max_attempts: i32,
1657    initial_interval_secs: f64,
1658    backoff_coefficient: f64,
1659    start_to_close_secs: f64,
1660    heartbeat_timeout_secs: Option<f64>,
1661    claimed_by: Option<String>,
1662    scheduled_at: f64,
1663    started_at: Option<f64>,
1664    completed_at: Option<f64>,
1665    last_heartbeat: Option<f64>,
1666}
1667
1668impl From<SqliteActivityRow> for WorkflowActivity {
1669    fn from(r: SqliteActivityRow) -> Self {
1670        Self {
1671            id: Some(r.id),
1672            workflow_id: r.workflow_id,
1673            seq: r.seq,
1674            name: r.name,
1675            task_queue: r.task_queue,
1676            input: r.input,
1677            status: r.status,
1678            result: r.result,
1679            error: r.error,
1680            attempt: r.attempt,
1681            max_attempts: r.max_attempts,
1682            initial_interval_secs: r.initial_interval_secs,
1683            backoff_coefficient: r.backoff_coefficient,
1684            start_to_close_secs: r.start_to_close_secs,
1685            heartbeat_timeout_secs: r.heartbeat_timeout_secs,
1686            claimed_by: r.claimed_by,
1687            scheduled_at: r.scheduled_at,
1688            started_at: r.started_at,
1689            completed_at: r.completed_at,
1690            last_heartbeat: r.last_heartbeat,
1691        }
1692    }
1693}
1694
1695#[derive(sqlx::FromRow)]
1696struct SqliteTimerRow {
1697    id: i64,
1698    workflow_id: String,
1699    seq: i32,
1700    fire_at: f64,
1701    fired: bool,
1702}
1703
1704impl From<SqliteTimerRow> for WorkflowTimer {
1705    fn from(r: SqliteTimerRow) -> Self {
1706        Self {
1707            id: Some(r.id),
1708            workflow_id: r.workflow_id,
1709            seq: r.seq,
1710            fire_at: r.fire_at,
1711            fired: r.fired,
1712        }
1713    }
1714}
1715
1716#[derive(sqlx::FromRow)]
1717struct SqliteSignalRow {
1718    id: i64,
1719    workflow_id: String,
1720    name: String,
1721    payload: Option<String>,
1722    consumed: bool,
1723    received_at: f64,
1724}
1725
1726impl From<SqliteSignalRow> for WorkflowSignal {
1727    fn from(r: SqliteSignalRow) -> Self {
1728        Self {
1729            id: Some(r.id),
1730            workflow_id: r.workflow_id,
1731            name: r.name,
1732            payload: r.payload,
1733            consumed: r.consumed,
1734            received_at: r.received_at,
1735        }
1736    }
1737}
1738
1739#[derive(sqlx::FromRow)]
1740struct SqliteScheduleRow {
1741    name: String,
1742    namespace: String,
1743    workflow_type: String,
1744    cron_expr: String,
1745    timezone: String,
1746    input: Option<String>,
1747    task_queue: String,
1748    overlap_policy: String,
1749    paused: bool,
1750    last_run_at: Option<f64>,
1751    next_run_at: Option<f64>,
1752    last_workflow_id: Option<String>,
1753    created_at: f64,
1754}
1755
1756impl From<SqliteScheduleRow> for WorkflowSchedule {
1757    fn from(r: SqliteScheduleRow) -> Self {
1758        Self {
1759            name: r.name,
1760            namespace: r.namespace,
1761            workflow_type: r.workflow_type,
1762            cron_expr: r.cron_expr,
1763            timezone: r.timezone,
1764            input: r.input,
1765            task_queue: r.task_queue,
1766            overlap_policy: r.overlap_policy,
1767            paused: r.paused,
1768            last_run_at: r.last_run_at,
1769            next_run_at: r.next_run_at,
1770            last_workflow_id: r.last_workflow_id,
1771            created_at: r.created_at,
1772        }
1773    }
1774}
1775
1776#[derive(sqlx::FromRow)]
1777struct SqliteWorkerRow {
1778    id: String,
1779    namespace: String,
1780    identity: String,
1781    task_queue: String,
1782    workflows: Option<String>,
1783    activities: Option<String>,
1784    max_concurrent_workflows: i32,
1785    max_concurrent_activities: i32,
1786    active_tasks: i32,
1787    last_heartbeat: f64,
1788    registered_at: f64,
1789}
1790
1791impl From<SqliteWorkerRow> for WorkflowWorker {
1792    fn from(r: SqliteWorkerRow) -> Self {
1793        Self {
1794            id: r.id,
1795            namespace: r.namespace,
1796            identity: r.identity,
1797            task_queue: r.task_queue,
1798            workflows: r.workflows,
1799            activities: r.activities,
1800            max_concurrent_workflows: r.max_concurrent_workflows,
1801            max_concurrent_activities: r.max_concurrent_activities,
1802            active_tasks: r.active_tasks,
1803            last_heartbeat: r.last_heartbeat,
1804            registered_at: r.registered_at,
1805        }
1806    }
1807}