Skip to main content

assay_workflow/store/
sqlite.rs

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