drasi_bootstrap_sqlite/
lib.rs

1// Copyright 2025 The Drasi Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![allow(unexpected_cfgs)]
16
17use anyhow::Result;
18use async_trait::async_trait;
19use base64::Engine;
20use drasi_core::models::{
21    Element, ElementMetadata, ElementPropertyMap, ElementReference, ElementValue, SourceChange,
22};
23use drasi_lib::bootstrap::{
24    BootstrapContext, BootstrapProvider, BootstrapRequest, BootstrapResult,
25};
26use drasi_lib::channels::BootstrapEventSender;
27use log::{info, warn};
28use ordered_float::OrderedFloat;
29use rusqlite::types::ValueRef;
30use rusqlite::{Connection, OpenFlags};
31use std::sync::Arc;
32
33pub mod descriptor;
34
35/// Per-table key configuration for stable element IDs.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct TableKeyConfig {
38    pub table: String,
39    pub key_columns: Vec<String>,
40}
41
42/// SQLite bootstrap provider.
43#[derive(Clone)]
44pub struct SqliteBootstrapProvider {
45    path: Option<String>,
46    tables: Option<Vec<String>>,
47    table_keys: Vec<TableKeyConfig>,
48}
49
50impl SqliteBootstrapProvider {
51    pub fn builder() -> SqliteBootstrapBuilder {
52        SqliteBootstrapBuilder::new()
53    }
54}
55
56/// Builder for [`SqliteBootstrapProvider`].
57pub struct SqliteBootstrapBuilder {
58    path: Option<String>,
59    tables: Option<Vec<String>>,
60    table_keys: Vec<TableKeyConfig>,
61}
62
63impl SqliteBootstrapBuilder {
64    fn new() -> Self {
65        Self {
66            path: None,
67            tables: None,
68            table_keys: Vec::new(),
69        }
70    }
71
72    pub fn with_path(mut self, path: impl Into<String>) -> Self {
73        self.path = Some(path.into());
74        self
75    }
76
77    pub fn in_memory(mut self) -> Self {
78        self.path = None;
79        self
80    }
81
82    pub fn with_tables(mut self, tables: Vec<String>) -> Self {
83        self.tables = Some(tables);
84        self
85    }
86
87    pub fn with_table_keys(mut self, table_keys: Vec<TableKeyConfig>) -> Self {
88        self.table_keys = table_keys;
89        self
90    }
91
92    pub fn build(self) -> SqliteBootstrapProvider {
93        SqliteBootstrapProvider {
94            path: self.path,
95            tables: self.tables,
96            table_keys: self.table_keys,
97        }
98    }
99}
100
101#[async_trait]
102impl BootstrapProvider for SqliteBootstrapProvider {
103    async fn bootstrap(
104        &self,
105        request: BootstrapRequest,
106        context: &BootstrapContext,
107        event_tx: BootstrapEventSender,
108        _settings: Option<&drasi_lib::config::SourceSubscriptionSettings>,
109    ) -> Result<BootstrapResult> {
110        info!("Starting SQLite bootstrap for query '{}'", request.query_id);
111
112        let Some(path) = self.path.clone() else {
113            warn!("SQLite bootstrap skipped for in-memory database");
114            return Ok(BootstrapResult::default());
115        };
116
117        let configured_tables = self.tables.clone();
118        let table_keys = self.table_keys.clone();
119        let context = context.clone();
120        let query_id = request.query_id.clone();
121
122        let count = tokio::task::spawn_blocking(move || {
123            stream_sqlite_tables(
124                &path,
125                configured_tables.as_ref(),
126                &table_keys,
127                &request,
128                &context,
129                &event_tx,
130            )
131        })
132        .await??;
133
134        info!("SQLite bootstrap completed for query '{query_id}': {count} rows");
135        Ok(BootstrapResult {
136            event_count: count,
137            ..Default::default()
138        })
139    }
140}
141
142fn stream_sqlite_tables(
143    path: &str,
144    configured_tables: Option<&Vec<String>>,
145    table_keys: &[TableKeyConfig],
146    request: &BootstrapRequest,
147    context: &BootstrapContext,
148    event_tx: &BootstrapEventSender,
149) -> Result<usize> {
150    let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
151    let tables = resolve_tables(&conn, configured_tables)?;
152    let mut count = 0usize;
153
154    for table in tables {
155        if !request.node_labels.is_empty() && !request.node_labels.contains(&table) {
156            continue;
157        }
158
159        let key_columns = key_columns_for_table(&conn, table_keys, &table)?;
160        let query = format!("SELECT rowid, * FROM {}", quote_ident(&table));
161        let mut stmt = conn.prepare(&query)?;
162        let column_names = user_column_names(&stmt);
163        let mut rows = stmt.query([])?;
164
165        while let Some(row) = rows.next()? {
166            let (values, rowid) = map_query_row(row, &column_names)?;
167            let element_id = generate_element_id(&table, &values, &key_columns, Some(rowid));
168            let mut properties = ElementPropertyMap::new();
169            for (name, value) in values {
170                properties.insert(&name, value);
171            }
172
173            let labels: Arc<[Arc<str>]> = vec![Arc::<str>::from(table.as_str())].into();
174            let element = Element::Node {
175                metadata: ElementMetadata {
176                    reference: ElementReference::new(&context.source_id, &element_id),
177                    labels,
178                    effective_from: chrono::Utc::now().timestamp_millis() as u64,
179                },
180                properties,
181            };
182
183            event_tx
184                .blocking_send(drasi_lib::channels::BootstrapEvent {
185                    source_id: context.source_id.clone(),
186                    change: SourceChange::Insert { element },
187                    timestamp: chrono::Utc::now(),
188                    sequence: context.next_sequence(),
189                })
190                .map_err(|e| anyhow::anyhow!("Failed to send SQLite bootstrap event: {e}"))?;
191            count += 1;
192        }
193    }
194
195    Ok(count)
196}
197
198fn key_columns_for_table(
199    conn: &Connection,
200    table_keys: &[TableKeyConfig],
201    table: &str,
202) -> Result<Vec<String>> {
203    if let Some(cfg) = table_keys.iter().find(|item| item.table == table) {
204        return Ok(cfg.key_columns.clone());
205    }
206    detect_primary_key(conn, table)
207}
208
209fn resolve_tables(
210    conn: &Connection,
211    configured_tables: Option<&Vec<String>>,
212) -> Result<Vec<String>> {
213    if let Some(tables) = configured_tables {
214        return Ok(tables.clone());
215    }
216
217    let mut stmt = conn.prepare(
218        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'",
219    )?;
220    let tables = stmt
221        .query_map([], |row| row.get::<_, String>(0))?
222        .collect::<std::result::Result<Vec<_>, _>>()?;
223    Ok(tables)
224}
225
226fn user_column_names(stmt: &rusqlite::Statement<'_>) -> Vec<String> {
227    // First column is rowid, remaining are user columns
228    (1..stmt.column_count())
229        .map(|index| stmt.column_name(index).unwrap_or("").to_string())
230        .collect()
231}
232
233fn map_query_row(
234    row: &rusqlite::Row<'_>,
235    column_names: &[String],
236) -> Result<(Vec<(String, ElementValue)>, i64)> {
237    let rowid: i64 = row.get(0)?;
238    let mut values = Vec::with_capacity(column_names.len());
239    for (i, name) in column_names.iter().enumerate() {
240        let value_ref = row.get_ref(i + 1)?;
241        values.push((name.clone(), value_ref_to_element_value(value_ref)));
242    }
243    Ok((values, rowid))
244}
245
246fn read_table_rows(
247    conn: &Connection,
248    table: &str,
249) -> Result<Vec<(Vec<(String, ElementValue)>, i64)>> {
250    let query = format!("SELECT rowid, * FROM {}", quote_ident(table));
251    let mut stmt = conn.prepare(&query)?;
252    let column_names = user_column_names(&stmt);
253    let mut rows = stmt.query([])?;
254    let mut result = Vec::new();
255    while let Some(row) = rows.next()? {
256        result.push(map_query_row(row, &column_names)?);
257    }
258    Ok(result)
259}
260
261fn detect_primary_key(conn: &Connection, table: &str) -> Result<Vec<String>> {
262    let sql = format!("PRAGMA table_info({})", quote_ident(table));
263    let mut stmt = conn.prepare(&sql)?;
264    let mut key_pairs = stmt
265        .query_map([], |row| {
266            let name: String = row.get(1)?;
267            let pk: i64 = row.get(5)?;
268            Ok((pk, name))
269        })?
270        .collect::<std::result::Result<Vec<_>, _>>()?;
271
272    key_pairs.retain(|(pk, _)| *pk > 0);
273    key_pairs.sort_by_key(|(pk, _)| *pk);
274    Ok(key_pairs.into_iter().map(|(_, name)| name).collect())
275}
276
277fn value_ref_to_element_value(value_ref: ValueRef<'_>) -> ElementValue {
278    match value_ref {
279        ValueRef::Null => ElementValue::Null,
280        ValueRef::Integer(i) => ElementValue::Integer(i),
281        ValueRef::Real(f) => ElementValue::Float(OrderedFloat(f)),
282        ValueRef::Text(t) => ElementValue::String(Arc::from(String::from_utf8_lossy(t).as_ref())),
283        ValueRef::Blob(b) => ElementValue::String(Arc::from(
284            base64::engine::general_purpose::STANDARD.encode(b),
285        )),
286    }
287}
288
289fn generate_element_id(
290    table: &str,
291    values: &[(String, ElementValue)],
292    key_columns: &[String],
293    rowid: Option<i64>,
294) -> String {
295    if !key_columns.is_empty() {
296        let key_parts = key_columns
297            .iter()
298            .filter_map(|column| {
299                values
300                    .iter()
301                    .find(|(name, _)| name == column)
302                    .map(|(_, v)| v)
303            })
304            .map(value_to_id_fragment)
305            .collect::<Vec<_>>();
306
307        if !key_parts.is_empty() {
308            return format!("{table}:{}", key_parts.join(":"));
309        }
310    }
311
312    // Fall back to rowid, matching the source's CDC behavior
313    if let Some(id) = rowid {
314        return format!("{table}:{id}");
315    }
316
317    format!("{table}:unknown")
318}
319
320fn value_to_id_fragment(value: &ElementValue) -> String {
321    match value {
322        ElementValue::Null => "null".to_string(),
323        ElementValue::Bool(v) => v.to_string(),
324        ElementValue::Float(v) => v.to_string(),
325        ElementValue::Integer(v) => v.to_string(),
326        ElementValue::String(v) => v.to_string(),
327        ElementValue::LocalDateTime(v) => v.to_string(),
328        ElementValue::ZonedDateTime(v) => v.to_rfc3339(),
329        ElementValue::List(v) => format!("{v:?}").replace(':', "%3A"),
330        ElementValue::Object(v) => format!("{v:?}").replace(':', "%3A"),
331    }
332}
333
334fn quote_ident(identifier: &str) -> String {
335    format!("\"{}\"", identifier.replace('"', "\"\""))
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341
342    #[test]
343    fn generate_element_id_uses_key_columns_when_provided() {
344        let values = vec![
345            ("id".to_string(), ElementValue::Integer(42)),
346            ("name".to_string(), ElementValue::String(Arc::from("test"))),
347        ];
348        let keys = vec!["id".to_string()];
349        assert_eq!(
350            generate_element_id("sensors", &values, &keys, Some(1)),
351            "sensors:42"
352        );
353    }
354
355    #[test]
356    fn generate_element_id_uses_composite_keys() {
357        let values = vec![
358            ("tenant".to_string(), ElementValue::String(Arc::from("t1"))),
359            (
360                "event_id".to_string(),
361                ElementValue::String(Arc::from("e1")),
362            ),
363        ];
364        let keys = vec!["tenant".to_string(), "event_id".to_string()];
365        assert_eq!(
366            generate_element_id("events", &values, &keys, Some(99)),
367            "events:t1:e1"
368        );
369    }
370
371    #[test]
372    fn generate_element_id_falls_back_to_rowid_when_no_keys() {
373        let values = vec![
374            ("name".to_string(), ElementValue::String(Arc::from("test"))),
375            ("value".to_string(), ElementValue::Integer(100)),
376        ];
377        let keys: Vec<String> = vec![];
378        assert_eq!(
379            generate_element_id("data", &values, &keys, Some(7)),
380            "data:7"
381        );
382    }
383
384    #[test]
385    fn generate_element_id_returns_unknown_without_keys_or_rowid() {
386        let values = vec![("x".to_string(), ElementValue::Integer(1))];
387        let keys: Vec<String> = vec![];
388        assert_eq!(
389            generate_element_id("data", &values, &keys, None),
390            "data:unknown"
391        );
392    }
393
394    #[test]
395    fn stream_sqlite_tables_filters_labels_and_assigns_sequences() {
396        let path = std::env::temp_dir().join(format!(
397            "drasi-sqlite-stream-{}-{}.db",
398            std::process::id(),
399            std::time::SystemTime::now()
400                .duration_since(std::time::UNIX_EPOCH)
401                .expect("time")
402                .as_nanos()
403        ));
404        let path_str = path.to_string_lossy().to_string();
405        let conn = Connection::open(&path).expect("open");
406        conn.execute_batch(
407            "CREATE TABLE sensors (id INTEGER PRIMARY KEY, name TEXT);
408             INSERT INTO sensors VALUES (1, 'alpha');
409             INSERT INTO sensors VALUES (2, 'beta');
410             CREATE TABLE ignored (id INTEGER PRIMARY KEY, name TEXT);
411             INSERT INTO ignored VALUES (9, 'skip');",
412        )
413        .expect("setup");
414        drop(conn);
415
416        let (tx, mut rx) = tokio::sync::mpsc::channel(16);
417        let context = BootstrapContext::new_minimal("server".to_string(), "src".to_string());
418        let request = BootstrapRequest {
419            query_id: "q1".to_string(),
420            node_labels: vec!["sensors".to_string()],
421            relation_labels: Vec::new(),
422            request_id: "req-1".to_string(),
423        };
424        let tables = vec!["sensors".to_string(), "ignored".to_string()];
425        let count = stream_sqlite_tables(&path_str, Some(&tables), &[], &request, &context, &tx)
426            .expect("stream");
427        drop(tx);
428
429        assert_eq!(count, 2, "only sensors rows should be streamed");
430        let mut events = Vec::new();
431        while let Ok(event) = rx.try_recv() {
432            events.push(event);
433        }
434        let _ = std::fs::remove_file(&path);
435
436        assert_eq!(events.len(), 2);
437        assert_eq!(events[0].source_id, "src");
438        assert_eq!(events[0].sequence, 0);
439        assert_eq!(events[1].sequence, 1);
440
441        let names: Vec<String> = events
442            .iter()
443            .map(|event| match &event.change {
444                SourceChange::Insert { element } => match element {
445                    Element::Node { properties, .. } => match &properties["name"] {
446                        ElementValue::String(name) => name.to_string(),
447                        other => panic!("unexpected name value: {other:?}"),
448                    },
449                    other => panic!("expected node, got {other:?}"),
450                },
451                other => panic!("expected insert, got {other:?}"),
452            })
453            .collect();
454        assert_eq!(names, vec!["alpha".to_string(), "beta".to_string()]);
455    }
456
457    #[test]
458    fn read_table_rows_returns_rows_with_rowid() {
459        let conn = Connection::open_in_memory().expect("open");
460        conn.execute_batch("CREATE TABLE items (name TEXT, value INTEGER); INSERT INTO items VALUES ('a', 1); INSERT INTO items VALUES ('b', 2);")
461            .expect("setup");
462
463        let rows = read_table_rows(&conn, "items").expect("read");
464        assert_eq!(rows.len(), 2);
465
466        let (first_row, first_rowid) = &rows[0];
467        assert_eq!(*first_rowid, 1);
468        assert_eq!(first_row.len(), 2);
469        assert_eq!(first_row[0].0, "name");
470
471        let (second_row, second_rowid) = &rows[1];
472        assert_eq!(*second_rowid, 2);
473        assert_eq!(second_row.len(), 2);
474        let _ = second_row;
475    }
476
477    #[test]
478    fn detect_primary_key_finds_pk_columns() {
479        let conn = Connection::open_in_memory().expect("open");
480        conn.execute_batch("CREATE TABLE sensors (id INTEGER PRIMARY KEY, name TEXT)")
481            .expect("setup");
482
483        let pks = detect_primary_key(&conn, "sensors").expect("detect");
484        assert_eq!(pks, vec!["id".to_string()]);
485    }
486
487    #[test]
488    fn detect_primary_key_returns_empty_for_no_pk() {
489        let conn = Connection::open_in_memory().expect("open");
490        conn.execute_batch("CREATE TABLE data (x TEXT, y TEXT)")
491            .expect("setup");
492
493        let pks = detect_primary_key(&conn, "data").expect("detect");
494        assert!(pks.is_empty());
495    }
496
497    #[test]
498    fn element_id_matches_between_bootstrap_and_source_with_pk() {
499        // Verifies that bootstrap and source produce the same element ID
500        // when a table has a declared PK.
501        let conn = Connection::open_in_memory().expect("open");
502        conn.execute_batch("CREATE TABLE sensors (id INTEGER PRIMARY KEY, name TEXT); INSERT INTO sensors VALUES (42, 'test');")
503            .expect("setup");
504
505        let pks = detect_primary_key(&conn, "sensors").expect("detect pk");
506        let rows = read_table_rows(&conn, "sensors").expect("read");
507        let (row, rowid) = &rows[0];
508
509        let bootstrap_id = generate_element_id("sensors", row, &pks, Some(*rowid));
510        // Source would produce "sensors:42" via PK column "id"
511        assert_eq!(bootstrap_id, "sensors:42");
512    }
513
514    #[test]
515    fn element_id_matches_between_bootstrap_and_source_without_pk() {
516        // Verifies that bootstrap falls back to rowid just like source does
517        // when no PK is declared.
518        let conn = Connection::open_in_memory().expect("open");
519        conn.execute_batch(
520            "CREATE TABLE data (x TEXT, y TEXT); INSERT INTO data VALUES ('a', 'b');",
521        )
522        .expect("setup");
523
524        let pks = detect_primary_key(&conn, "data").expect("detect pk");
525        assert!(pks.is_empty());
526
527        let rows = read_table_rows(&conn, "data").expect("read");
528        let (row, rowid) = &rows[0];
529
530        let bootstrap_id = generate_element_id("data", row, &pks, Some(*rowid));
531        // Source would produce "data:1" via rowid fallback
532        assert_eq!(bootstrap_id, "data:1");
533    }
534}
535
536/// Dynamic plugin entry point.
537#[cfg(feature = "dynamic-plugin")]
538drasi_plugin_sdk::export_plugin!(
539    plugin_id = "sqlite-bootstrap",
540    core_version = env!("CARGO_PKG_VERSION"),
541    lib_version = env!("CARGO_PKG_VERSION"),
542    plugin_version = env!("CARGO_PKG_VERSION"),
543    source_descriptors = [],
544    reaction_descriptors = [],
545    bootstrap_descriptors = [descriptor::SqliteBootstrapDescriptor],
546);
drasi_bootstrap_sqlite/lib.rs

drasi_bootstrap_sqlite/
lib.rs
