Skip to main content

udb/control/
engine.rs

1// src/engine.rs — Migration FSM state machine types.
2//
3// The Engine struct tracks the FSM state, run lifecycle, and error history for a
4// single migration run.  This is a **pure data model** — no backend connections required.
5// The Go UDB service drives the state transitions; the Rust library validates them
6// and produces the run snapshot that the service persists to `migration_runtime_state`.
7//
8// APPLYING phase scope (spec §16.3.1, step 5):
9//   • Tier 1 (PostgreSQL)  — apply SQL DDL migration files via `sqlx`
10//   • Tier 3 (Qdrant)     — create / update vector collections and payload indexes
11//   • Tier 4 (MinIO)      — create buckets, set lifecycle policies
12//   Tier 2 (Redis) is managed by config/TTL only; no structural migration needed.
13//
14// Aligned with:
15//   - legacy_sql  legacycore/db/migrate.go      (FSM states + validTransitions) — SQL-only reference
16//   - UDB spec §16.1                        (four storage tiers)
17//   - UDB spec §16.3.1                      (FSM Lifecycle)
18//   - UDB spec §16.3.2                      (Double-Entry Ledger)
19
20use serde::{Deserialize, Serialize};
21use std::time::{SystemTime, UNIX_EPOCH};
22
23// ── Constants ─────────────────────────────────────────────────────────────────
24
25/// Maximum number of automatic recovery retries before the engine halts.
26/// Mirrors `maxRetries` in legacy_sql `migrate.go`.
27pub const MAX_RETRIES: u32 = 3;
28
29/// PostgreSQL advisory lock key used to prevent concurrent migration runs.
30/// Stable across restarts; chosen so as not to conflict with application locks.
31pub const PG_ADVISORY_LOCK_KEY: i64 = 0x7072_6973_6d76_6973; // "primvis" in hex
32
33// ── FSM States ────────────────────────────────────────────────────────────────
34
35/// The current FSM state of a migration run.
36///
37/// Each state corresponds to one phase of the double-entry migration lifecycle
38/// described in the UDB spec §16.3 and legacy_sql's migrate.go.
39#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
40#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
41pub enum FsmState {
42    /// Quiescent — no migration in progress.
43    #[default]
44    Idle,
45    /// Bootstrap: creating migration ledger tables, acquiring advisory lock.
46    Initialising,
47    /// Loading `.proto` files and building the desired AST.
48    LoadProtoState,
49    /// Comparing the parsed AST checksum against the `proto_schema_versions` ledger.
50    ProtoChecksumLint,
51    /// Diffing old manifest vs new to discover which schema changes are needed.
52    PlanProtoDiff,
53    /// Generating SQL delta files from the diff result.
54    GenerateSql,
55    /// Verifying checksums of all SQL files in the migration directory.
56    ChecksumLint,
57    /// Executing migration artefacts across all active backends:
58    /// SQL DDL → PostgreSQL, Collections/Indexes → Qdrant, Buckets → MinIO.
59    Applying,
60    /// Running the proto linter against the live PostgreSQL `information_schema`
61    /// and verifying Qdrant collections + MinIO bucket existence.
62    Linting,
63    /// Applying safe structural repairs: SQL DDL for Postgres, collection
64    /// recreation for Qdrant, bucket creation for MinIO.
65    AutoAltering,
66    /// Cross-checking the live topology across ALL active backends against the
67    /// desired proto AST (SQL tables, vector collections, object buckets).
68    Verifying,
69    /// Stepping back from `Error` state after a manual recovery action.
70    Recovering,
71    /// Migration run completed successfully — ledger entries recorded.
72    Completed,
73    /// The FSM encountered an unrecoverable error and has halted.
74    Error,
75}
76
77impl std::fmt::Display for FsmState {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        write!(f, "{}", self.as_str())
80    }
81}
82
83impl FsmState {
84    pub const ALL: [FsmState; 14] = [
85        FsmState::Idle,
86        FsmState::Initialising,
87        FsmState::LoadProtoState,
88        FsmState::ProtoChecksumLint,
89        FsmState::PlanProtoDiff,
90        FsmState::GenerateSql,
91        FsmState::ChecksumLint,
92        FsmState::Applying,
93        FsmState::Linting,
94        FsmState::AutoAltering,
95        FsmState::Verifying,
96        FsmState::Recovering,
97        FsmState::Completed,
98        FsmState::Error,
99    ];
100
101    pub const VARIANT_COUNT: usize = Self::ALL.len();
102
103    /// Returns the canonical uppercase string stored in DB records / JSON.
104    /// Mirrors the `MigrationState` constants in legacy_sql `migrate.go`.
105    pub fn as_str(&self) -> &'static str {
106        match self {
107            Self::Idle => "IDLE",
108            Self::Initialising => "INITIALISING",
109            Self::LoadProtoState => "LOAD_PROTO_STATE",
110            Self::ProtoChecksumLint => "PROTO_CHECKSUM_LINT",
111            Self::PlanProtoDiff => "PLAN_PROTO_DIFF",
112            Self::GenerateSql => "GENERATE_SQL",
113            Self::ChecksumLint => "CHECKSUM_LINT",
114            Self::Applying => "APPLYING",
115            Self::Linting => "LINTING",
116            Self::AutoAltering => "AUTO_ALTERING",
117            Self::Verifying => "VERIFYING",
118            Self::Recovering => "RECOVERING",
119            Self::Completed => "COMPLETED",
120            Self::Error => "ERROR",
121        }
122    }
123
124    /// Parse a state from its canonical uppercase string.
125    /// Returns `None` for unrecognised values.
126    #[allow(clippy::should_implement_trait)]
127    pub fn from_str(s: &str) -> Option<Self> {
128        Some(match s {
129            "IDLE" => Self::Idle,
130            "INITIALISING" => Self::Initialising,
131            "LOAD_PROTO_STATE" => Self::LoadProtoState,
132            "PROTO_CHECKSUM_LINT" => Self::ProtoChecksumLint,
133            "PLAN_PROTO_DIFF" => Self::PlanProtoDiff,
134            "GENERATE_SQL" => Self::GenerateSql,
135            "CHECKSUM_LINT" => Self::ChecksumLint,
136            "APPLYING" => Self::Applying,
137            "LINTING" => Self::Linting,
138            "AUTO_ALTERING" => Self::AutoAltering,
139            "VERIFYING" => Self::Verifying,
140            "RECOVERING" => Self::Recovering,
141            "COMPLETED" => Self::Completed,
142            "ERROR" => Self::Error,
143            _ => return None,
144        })
145    }
146
147    /// Returns all legal successor states from this state.
148    ///
149    /// Mirrors `validTransitions` in legacy_sql `migrate.go`.
150    /// The Go service calls this to guard every `e.transition(next)` call.
151    pub fn valid_transitions(&self) -> Vec<FsmState> {
152        use FsmState::*;
153        match self {
154            Idle => vec![Initialising, Recovering],
155            Initialising => vec![LoadProtoState, ChecksumLint, Error],
156            LoadProtoState => vec![ProtoChecksumLint, Error],
157            ProtoChecksumLint => vec![PlanProtoDiff, ChecksumLint, Error],
158            PlanProtoDiff => vec![GenerateSql, ChecksumLint, Error],
159            GenerateSql => vec![ChecksumLint, Error],
160            ChecksumLint => vec![Applying, Error],
161            Applying => vec![Verifying, Linting, Error],
162            Verifying => vec![Completed, AutoAltering, Error],
163            Linting => vec![AutoAltering, Completed, Error],
164            AutoAltering => vec![Completed, Error],
165            Recovering => vec![Idle, Error],
166            Completed => vec![Idle],
167            Error => vec![Recovering, Idle],
168        }
169    }
170
171    /// Returns `true` when transitioning to `next` is legal from this state.
172    pub fn can_transition_to(&self, next: &FsmState) -> bool {
173        self.valid_transitions().contains(next)
174    }
175
176    /// Returns `true` for terminal states where the FSM run has ended.
177    pub fn is_terminal(&self) -> bool {
178        matches!(self, FsmState::Completed | FsmState::Error)
179    }
180}
181
182// ── Run snapshot ──────────────────────────────────────────────────────────────
183
184/// A point-in-time snapshot of a migration run, persisted to
185/// `public.migration_runtime_state` by the Go UDB service.
186#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
187pub struct RuntimeStateSnapshot {
188    pub run_id: String,
189    pub state: FsmState,
190    /// The migration artefact currently being applied (SQL filename, collection
191    /// name, or bucket name depending on which backend tier is active).
192    pub active_file: String,
193    pub retry_count: u32,
194    pub updated_at_unix: u64,
195    /// Backend tiers active in the current run, e.g. `["postgres", "qdrant", "minio"]`.
196    pub active_backends: Vec<String>,
197}
198
199// ── Error record ──────────────────────────────────────────────────────────────
200
201/// A single migration error record, persisted to `public.migration_error_log`.
202#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
203pub struct EngineError {
204    /// Matches the `run_id` of the `Engine` that produced the error.
205    pub run_id: String,
206    /// The FSM state the engine was in when the error occurred.
207    pub fsm_state: String,
208    /// The artefact being applied when the error occurred (SQL file, collection
209    /// name, or bucket name; empty for non-apply errors).
210    pub filename: String,
211    /// Human-readable error message.
212    pub message: String,
213    /// Unix timestamp of the error.
214    pub failed_at_unix: u64,
215    /// `true` after `RecoverFromError()` has been called.
216    pub resolved: bool,
217}
218
219// ── Engine ────────────────────────────────────────────────────────────────────
220
221/// In-memory FSM driver used by the Go UDB service.
222///
223/// The Go service creates an `Engine` at startup, calls `transition()` to move
224/// through the states, and serializes `RuntimeStateSnapshot` to the migration
225/// ledger (PostgreSQL `public.migration_runtime_state`) after each transition.
226/// Blocked/error transitions are refused with a descriptive error.
227#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
228pub struct Engine {
229    /// Unique identifier for this run, formatted as `20060102T150405.000000000Z`.
230    pub run_id: String,
231    /// Current FSM state.
232    pub state: FsmState,
233    /// The migration file currently being applied (empty when not in `Applying`).
234    pub active_file: String,
235    /// Number of auto-recovery attempts in this run.
236    pub retry_count: u32,
237    /// Last recorded error, if any.
238    pub error: Option<EngineError>,
239    /// `true` once the migration ledger tables have been bootstrapped.
240    pub tracker_ready: bool,
241    /// Backend tiers active in this run, e.g. `["postgres", "qdrant", "minio"]`.
242    /// Set by the Go service during INITIALISING based on `UdbConfig.active_tiers()`.
243    pub active_backends: Vec<String>,
244}
245
246impl Engine {
247    /// Create a new `Engine` in `IDLE` state with the given run ID.
248    pub fn new(run_id: impl Into<String>) -> Self {
249        Self {
250            run_id: run_id.into(),
251            state: FsmState::Idle,
252            active_file: String::new(),
253            retry_count: 0,
254            error: None,
255            tracker_ready: false,
256            active_backends: vec!["postgres".to_string()], // Tier 1 always present
257        }
258    }
259
260    /// Create a new `Engine` with an auto-generated run ID based on current UTC time.
261    pub fn new_auto_id() -> Self {
262        let ts = SystemTime::now()
263            .duration_since(UNIX_EPOCH)
264            .map(|d| d.as_secs())
265            .unwrap_or_default();
266        Self::new(format!("{ts:020}"))
267    }
268
269    /// Attempt to transition to `next`.
270    ///
271    /// Returns `Ok(())` on success and updates `self.state`.
272    /// Returns `Err(String)` with a human-readable message for illegal transitions.
273    pub fn transition(&mut self, next: FsmState) -> Result<(), String> {
274        if self.state.can_transition_to(&next) {
275            self.state = next;
276            Ok(())
277        } else {
278            Err(format!(
279                "[fsm] illegal transition {} → {} (run={})",
280                self.state, next, self.run_id
281            ))
282        }
283    }
284
285    /// Record an error and unconditionally move to `Error` state.
286    ///
287    /// The previous state is captured in the `EngineError.fsm_state` field so
288    /// the Go service can log which phase failed.
289    pub fn fail(&mut self, filename: impl Into<String>, message: impl Into<String>) {
290        self.error = Some(EngineError {
291            run_id: self.run_id.clone(),
292            fsm_state: self.state.as_str().to_string(),
293            filename: filename.into(),
294            message: message.into(),
295            failed_at_unix: unix_now(),
296            resolved: false,
297        });
298        self.state = FsmState::Error;
299    }
300
301    /// Attempt an auto-recovery: increment retry counter and transition to `Recovering`.
302    ///
303    /// Returns `Err` when `MAX_RETRIES` has been exceeded.
304    pub fn recover(&mut self) -> Result<(), String> {
305        if self.retry_count >= MAX_RETRIES {
306            return Err(format!(
307                "[fsm] max retries ({MAX_RETRIES}) exceeded — manual intervention required (run={})",
308                self.run_id
309            ));
310        }
311        self.retry_count += 1;
312        self.transition(FsmState::Recovering)
313    }
314
315    /// Reset the engine back to `Idle` after a successful recovery.
316    pub fn reset_to_idle(&mut self) {
317        if let Some(err) = &mut self.error {
318            err.resolved = true;
319        }
320        self.state = FsmState::Idle;
321        self.active_file = String::new();
322    }
323
324    /// Produce a `RuntimeStateSnapshot` for persistence.
325    pub fn snapshot(&self) -> RuntimeStateSnapshot {
326        RuntimeStateSnapshot {
327            run_id: self.run_id.clone(),
328            state: self.state.clone(),
329            active_file: self.active_file.clone(),
330            retry_count: self.retry_count,
331            updated_at_unix: unix_now(),
332            active_backends: self.active_backends.clone(),
333        }
334    }
335}
336
337// ── Helpers ───────────────────────────────────────────────────────────────────
338
339fn unix_now() -> u64 {
340    SystemTime::now()
341        .duration_since(UNIX_EPOCH)
342        .map(|d| d.as_secs())
343        .unwrap_or_default()
344}
345
346// ── Tests ─────────────────────────────────────────────────────────────────────
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn idle_can_transition_to_initialising() {
354        let mut engine = Engine::new("test-run-1");
355        engine.transition(FsmState::Initialising).unwrap();
356        assert_eq!(engine.state, FsmState::Initialising);
357    }
358
359    #[test]
360    fn illegal_transition_returns_err() {
361        let mut engine = Engine::new("test-run-2");
362        let result = engine.transition(FsmState::Completed);
363        assert!(result.is_err(), "IDLE → COMPLETED must be illegal");
364    }
365
366    #[test]
367    fn fail_sets_error_state() {
368        let mut engine = Engine::new("test-run-3");
369        engine.transition(FsmState::Initialising).unwrap();
370        engine.fail("001_init.sql", "syntax error at line 42");
371        assert_eq!(engine.state, FsmState::Error);
372        assert!(engine.error.is_some());
373        let err = engine.error.as_ref().unwrap();
374        assert_eq!(err.fsm_state, "INITIALISING");
375        assert!(!err.resolved);
376    }
377
378    #[test]
379    fn recover_increments_retry_count() {
380        let mut engine = Engine::new("test-run-4");
381        engine.state = FsmState::Error; // simulate error
382        engine.recover().unwrap();
383        assert_eq!(engine.state, FsmState::Recovering);
384        assert_eq!(engine.retry_count, 1);
385    }
386
387    #[test]
388    fn max_retries_exceeded_returns_err() {
389        let mut engine = Engine::new("test-run-5");
390        engine.retry_count = MAX_RETRIES;
391        engine.state = FsmState::Error;
392        let result = engine.recover();
393        assert!(result.is_err());
394    }
395
396    #[test]
397    fn from_str_round_trip() {
398        assert_eq!(FsmState::VARIANT_COUNT, 14);
399        for state in FsmState::ALL {
400            let s = state.as_str();
401            assert_eq!(
402                FsmState::from_str(s).unwrap(),
403                state,
404                "round-trip failed for {s}"
405            );
406        }
407    }
408
409    #[test]
410    fn full_happy_path_transitions() {
411        let mut e = Engine::new("happy-run");
412        // Walk the standard migration path end-to-end
413        for next in [
414            FsmState::Initialising,
415            FsmState::LoadProtoState,
416            FsmState::ProtoChecksumLint,
417            FsmState::PlanProtoDiff,
418            FsmState::GenerateSql,
419            FsmState::ChecksumLint,
420            FsmState::Applying,
421            FsmState::Verifying,
422            FsmState::Completed,
423            FsmState::Idle,
424        ] {
425            e.transition(next).expect("unexpected illegal transition");
426        }
427    }
428}