Skip to main content

klukai_types/
pubsub.rs

1#![allow(clippy::result_large_err)]
2
3use std::{
4    cmp,
5    collections::{BTreeMap, HashMap, HashSet},
6    sync::Arc,
7    time::{Duration, Instant},
8};
9
10use crate::api::{ChangeId, ColumnName, ColumnType, RowId, SqliteValue, SqliteValueRef, TableName};
11use crate::spawn::spawn_counted;
12use crate::sqlite_pool::{self, RusqlitePool};
13use crate::tripwire::{Outcome, PreemptibleFutureExt, Tripwire};
14use bytes::{Buf, BufMut};
15use camino::{Utf8Path, Utf8PathBuf};
16use compact_str::{ToCompactString, format_compact};
17use enquote::unquote;
18use fallible_iterator::FallibleIterator;
19use indexmap::{IndexMap, IndexSet};
20use metrics::{counter, histogram};
21use parking_lot::{Condvar, Mutex, RwLock};
22use rusqlite::{
23    Connection, OptionalExtension, params_from_iter,
24    types::{FromSqlError, ValueRef},
25};
26use sqlite3_parser::{
27    ast::{
28        As, Cmd, Expr, FromClause, JoinConstraint, JoinOperator, JoinType, JoinedSelectTable, Name,
29        OneSelect, Operator, QualifiedName, ResultColumn, Select, SelectTable, Stmt,
30    },
31    lexer::sql::Parser,
32};
33use tokio::{
34    sync::{AcquireError, mpsc, watch},
35    task::block_in_place,
36};
37use tokio_util::sync::{CancellationToken, DropGuard, WaitForCancellationFuture};
38use tracing::{debug, error, info, trace, warn};
39use uuid::Uuid;
40
41use crate::{
42    agent::SplitPool,
43    api::QueryEvent,
44    change::Change,
45    schema::{Schema, Table},
46    sqlite::CrConn,
47    updates::HandleMetrics,
48};
49
50pub use crate::api::sqlite::ChangeType;
51use crate::updates::{Handle, Manager};
52
53#[derive(Debug, Default, Clone)]
54pub struct SubsManager(Arc<RwLock<InnerSubsManager>>);
55
56#[derive(Debug, Default)]
57struct InnerSubsManager {
58    handles: BTreeMap<Uuid, MatcherHandle>,
59    queries: HashMap<String, Uuid>,
60}
61
62// tools to bootstrap a new subscriber or notifier
63pub struct MatcherCreated {
64    pub evt_rx: mpsc::Receiver<QueryEvent>,
65}
66
67const SUB_EVENT_CHANNEL_CAP: usize = 512;
68
69impl Manager<MatcherHandle> for SubsManager {
70    fn trait_type(&self) -> String {
71        "subs".to_string()
72    }
73
74    fn get(&self, id: &Uuid) -> Option<MatcherHandle> {
75        self.0.read().get(id)
76    }
77
78    fn remove(&self, id: &Uuid) -> Option<MatcherHandle> {
79        let mut inner = self.0.write();
80        inner.remove(id)
81    }
82
83    fn get_handles(&self) -> BTreeMap<Uuid, MatcherHandle> {
84        self.0.read().handles.clone()
85    }
86}
87
88impl SubsManager {
89    pub fn get(&self, id: &Uuid) -> Option<MatcherHandle> {
90        self.0.read().get(id)
91    }
92
93    pub fn get_by_query(&self, sql: &str) -> Option<MatcherHandle> {
94        self.0.read().get_by_query(sql)
95    }
96
97    pub fn get_by_hash(&self, hash: &str) -> Option<MatcherHandle> {
98        self.0.read().get_by_hash(hash)
99    }
100
101    pub fn get_handles(&self) -> BTreeMap<Uuid, MatcherHandle> {
102        self.0.read().handles.clone()
103    }
104
105    pub async fn drop_handles(&self) {
106        let handles = {
107            let mut inner = self.0.write();
108            std::mem::take(&mut inner.handles)
109        };
110        for (_, handle) in handles.iter() {
111            handle.cleanup().await;
112        }
113    }
114
115    pub fn get_or_insert(
116        &self,
117        sql: &str,
118        subs_path: &Utf8Path,
119        schema: &Schema,
120        pool: &SplitPool,
121        tripwire: Tripwire,
122    ) -> Result<(MatcherHandle, Option<MatcherCreated>), MatcherError> {
123        if let Some(handle) = self.get_by_query(sql) {
124            return Ok((handle, None));
125        }
126
127        let mut inner = self.0.write();
128        if let Some(handle) = inner.get_by_query(sql) {
129            return Ok((handle, None));
130        }
131
132        let id = Uuid::new_v4();
133        let (evt_tx, evt_rx) = mpsc::channel(SUB_EVENT_CHANNEL_CAP);
134
135        let handle_res = Matcher::create(
136            id,
137            subs_path.to_path_buf(),
138            schema,
139            pool.client_dedicated()?,
140            evt_tx,
141            sql,
142            tripwire,
143        );
144
145        let handle = match handle_res {
146            Ok(handle) => handle,
147            Err(e) => {
148                error!(sub_id = %id, "could not create subscription: {e}");
149                if let Err(e) = Matcher::cleanup(id, Matcher::sub_path(subs_path, id)) {
150                    error!("could not cleanup subscription: {e}");
151                }
152
153                return Err(e);
154            }
155        };
156
157        inner.handles.insert(id, handle.clone());
158        inner.queries.insert(sql.to_owned(), id);
159
160        Ok((handle, Some(MatcherCreated { evt_rx })))
161    }
162
163    #[allow(clippy::too_many_arguments)]
164    pub fn restore(
165        &self,
166        id: Uuid,
167        subs_path: &Utf8Path,
168        schema: &Schema,
169        pool: &SplitPool,
170        tripwire: Tripwire,
171    ) -> Result<(MatcherHandle, MatcherCreated), MatcherError> {
172        let mut inner = self.0.write();
173
174        if inner.handles.contains_key(&id) {
175            return Err(MatcherError::CannotRestoreExisting);
176        }
177
178        let (evt_tx, evt_rx) = mpsc::channel(SUB_EVENT_CHANNEL_CAP);
179
180        let handle = Matcher::restore(
181            id,
182            subs_path.to_path_buf(),
183            schema,
184            pool.client_dedicated()?,
185            evt_tx,
186            tripwire,
187        )?;
188
189        inner.handles.insert(id, handle.clone());
190        inner.queries.insert(handle.inner.sql.clone(), id);
191
192        Ok((handle, MatcherCreated { evt_rx }))
193    }
194
195    pub fn remove(&self, id: &Uuid) -> Option<MatcherHandle> {
196        let mut inner = self.0.write();
197        inner.remove(id)
198    }
199}
200
201#[derive(Debug)]
202pub struct MatchableChange<'a> {
203    pub table: &'a TableName,
204    pub pk: &'a [u8],
205    pub column: &'a ColumnName,
206    pub cl: i64,
207}
208
209impl<'a> From<&'a Change> for MatchableChange<'a> {
210    fn from(value: &'a Change) -> Self {
211        MatchableChange {
212            table: &value.table,
213            pk: &value.pk,
214            column: &value.cid,
215            cl: value.cl,
216        }
217    }
218}
219
220impl InnerSubsManager {
221    fn get(&self, id: &Uuid) -> Option<MatcherHandle> {
222        self.handles.get(id).cloned()
223    }
224
225    fn get_by_query(&self, sql: &str) -> Option<MatcherHandle> {
226        self.queries
227            .get(sql)
228            .and_then(|id| self.handles.get(id).cloned())
229    }
230
231    pub fn get_by_hash(&self, hash: &str) -> Option<MatcherHandle> {
232        self.handles
233            .values()
234            .find(|x| x.inner.hash == hash)
235            .cloned()
236    }
237
238    fn remove(&mut self, id: &Uuid) -> Option<MatcherHandle> {
239        let handle = self.handles.remove(id)?;
240        self.queries.remove(&handle.inner.sql);
241        Some(handle)
242    }
243}
244
245#[derive(Debug, Clone, Copy)]
246pub enum MatcherState {
247    Created,
248    Running,
249}
250
251impl MatcherState {
252    fn is_running(&self) -> bool {
253        matches!(self, MatcherState::Running)
254    }
255}
256
257#[derive(Clone, Debug)]
258pub struct MatcherHandle {
259    inner: Arc<InnerMatcherHandle>,
260    state: StateLock,
261}
262
263#[derive(Debug)]
264struct InnerMatcherHandle {
265    id: Uuid,
266    sql: String,
267    hash: String,
268    pool: sqlite_pool::RusqlitePool,
269    parsed: ParsedSelect,
270    col_names: Vec<ColumnName>,
271    cancel: CancellationToken,
272    changes_tx: mpsc::Sender<MatchCandidates>,
273    last_change_rx: watch::Receiver<ChangeId>,
274    // some state from the matcher so we can take a look later
275    subs_path: String,
276    cached_statements: HashMap<String, MatcherStmt>,
277    metrics: HashMap<String, HandleMetrics>,
278}
279
280pub type MatchCandidates = IndexMap<TableName, IndexMap<Vec<u8>, i64>>;
281
282impl Handle for MatcherHandle {
283    fn id(&self) -> Uuid {
284        self.inner.id
285    }
286
287    fn cancelled(&self) -> WaitForCancellationFuture<'_> {
288        self.inner.cancel.cancelled()
289    }
290
291    fn changes_tx(&self) -> mpsc::Sender<MatchCandidates> {
292        self.inner.changes_tx.clone()
293    }
294
295    fn cleanup(&self) -> impl std::future::Future<Output = ()> + Send {
296        let cancel = self.inner.cancel.clone();
297        let id = self.inner.id;
298
299        async move {
300            cancel.cancel();
301            info!(sub_id = %id, "Canceled subscription");
302        }
303    }
304
305    fn filter_matchable_change(
306        &self,
307        candidates: &mut MatchCandidates,
308        change: MatchableChange,
309    ) -> bool {
310        trace!("filtering change {change:?}");
311        // don't double process the same pk
312        if candidates
313            .get(change.table)
314            .map(|pks| pks.contains_key(change.pk))
315            .unwrap_or_default()
316        {
317            trace!("already contained key");
318            return false;
319        }
320
321        // don't consider changes that don't have both the table + col in the matcher query
322        if !self
323            .inner
324            .parsed
325            .table_columns
326            .get(change.table.as_str())
327            .map(|cols| change.column.is_crsql_sentinel() || cols.contains(change.column.as_str()))
328            .unwrap_or_default()
329        {
330            trace!("could not match against parsed query table and columns");
331            return false;
332        }
333
334        if let Some(v) = candidates.get_mut(change.table) {
335            v.insert(change.pk.to_vec(), change.cl).is_none()
336        } else {
337            candidates.insert(
338                change.table.clone(),
339                [(change.pk.to_vec(), change.cl)].into(),
340            );
341            true
342        }
343    }
344
345    fn get_counter(&self, table: &str) -> &HandleMetrics {
346        self.inner.metrics.get(table).unwrap_or_else(|| {
347            panic!(
348                "metrics counter for table '{}' missing. subs hash {}!",
349                self.inner.hash, table
350            )
351        })
352    }
353}
354
355impl MatcherHandle {
356    pub fn sql(&self) -> &String {
357        &self.inner.sql
358    }
359
360    pub fn hash(&self) -> &str {
361        &self.inner.hash
362    }
363
364    pub fn parsed_columns(&self) -> &[ResultColumn] {
365        &self.inner.parsed.columns
366    }
367
368    pub fn col_names(&self) -> &[ColumnName] {
369        &self.inner.col_names
370    }
371
372    pub fn subs_path(&self) -> &String {
373        &self.inner.subs_path
374    }
375
376    pub fn cached_stmts(&self) -> &HashMap<String, MatcherStmt> {
377        &self.inner.cached_statements
378    }
379
380    pub fn pool(&self) -> &RusqlitePool {
381        &self.inner.pool
382    }
383
384    fn wait_for_running_state(&self) {
385        let (lock, cvar) = &*self.state;
386        let mut state = lock.lock();
387        while !state.is_running() {
388            cvar.wait(&mut state);
389        }
390    }
391
392    pub fn max_change_id(&self, conn: &Connection) -> rusqlite::Result<ChangeId> {
393        self.wait_for_running_state();
394        let mut prepped = conn.prepare_cached("SELECT COALESCE(MAX(id), 0) FROM changes")?;
395        prepped.query_row([], |row| row.get(0))
396    }
397
398    pub fn last_change_id_sent(&self) -> ChangeId {
399        *self.inner.last_change_rx.borrow()
400    }
401
402    pub fn max_row_id(&self, conn: &Connection) -> rusqlite::Result<RowId> {
403        self.wait_for_running_state();
404        let mut prepped =
405            conn.prepare_cached("SELECT COALESCE(MAX(__corro_rowid), 0) FROM query")?;
406        prepped.query_row([], |row| row.get(0))
407    }
408
409    pub fn changes_since(
410        &self,
411        since: ChangeId,
412        conn: &Connection,
413        tx: mpsc::Sender<QueryEvent>,
414    ) -> rusqlite::Result<ChangeId> {
415        self.wait_for_running_state();
416
417        let mut prepped = conn.prepare_cached("SELECT MIN(id) FROM changes")?;
418        let min_change_id: u64 = prepped.query_row([], |row| row.get(0))?;
419
420        // return error if we've cleared changes after the received change id
421        if since.0 + 1 < min_change_id {
422            return Err(rusqlite::Error::ModuleError(format!(
423                "subscription already deleted older changes, min change id: {min_change_id}",
424            )));
425        }
426
427        let mut query_cols = vec![];
428        for i in 0..(self.parsed_columns().len()) {
429            query_cols.push(format!("col_{i}"));
430        }
431        let mut prepped = conn.prepare_cached(&format!(
432            "SELECT id, type, __corro_rowid, {} FROM changes WHERE id > ? ORDER BY id ASC",
433            query_cols.join(",")
434        ))?;
435
436        let col_count = prepped.column_count();
437
438        let mut max_change_id = since;
439
440        let mut rows = prepped.query([since])?;
441
442        loop {
443            let row = match rows.next()? {
444                Some(row) => row,
445                None => break,
446            };
447
448            let change_id: ChangeId = row.get(0)?;
449            if change_id.0 > max_change_id.0 {
450                max_change_id = change_id;
451            }
452
453            if let Err(e) = tx.blocking_send(QueryEvent::Change(
454                row.get(1)?,
455                row.get(2)?,
456                (3..col_count)
457                    .map(|i| row.get::<_, SqliteValue>(i))
458                    .collect::<rusqlite::Result<Vec<_>>>()?,
459                change_id,
460            )) {
461                error!("could not send change to channel: {e}");
462                break;
463            }
464        }
465
466        Ok(max_change_id)
467    }
468
469    pub fn all_rows(
470        &self,
471        conn: &Connection,
472        tx: mpsc::Sender<QueryEvent>,
473    ) -> Result<ChangeId, MatcherError> {
474        self.wait_for_running_state();
475        let mut query_cols = vec![];
476        for i in 0..(self.parsed_columns().len()) {
477            query_cols.push(format!("col_{i}"));
478        }
479        let mut prepped = conn.prepare_cached(&format!(
480            "SELECT __corro_rowid, {} FROM query",
481            query_cols.join(",")
482        ))?;
483
484        let col_count = prepped.column_count();
485
486        tx.blocking_send(QueryEvent::Columns(self.col_names().to_vec()))
487            .map_err(|_| MatcherError::EventReceiverClosed)?;
488
489        let start = Instant::now();
490        let mut rows = prepped.query([])?;
491        let elapsed = start.elapsed();
492
493        let mut count = 0;
494
495        loop {
496            let row = match rows.next()? {
497                Some(row) => row,
498                None => break,
499            };
500
501            tx.blocking_send(QueryEvent::Row(
502                row.get(0)?,
503                (1..col_count)
504                    .map(|i| row.get::<_, SqliteValue>(i))
505                    .collect::<rusqlite::Result<Vec<_>>>()?,
506            ))
507            .map_err(|_| MatcherError::EventReceiverClosed)?;
508            count += 1;
509        }
510
511        trace!("sent {count} rows");
512
513        let max_change_id = conn
514            .prepare("SELECT COALESCE(MAX(id),0) FROM changes")?
515            .query_row([], |row| row.get(0))?;
516
517        tx.blocking_send(QueryEvent::EndOfQuery {
518            time: elapsed.as_secs_f64(),
519            change_id: Some(max_change_id),
520        })
521        .map_err(|_| MatcherError::EventReceiverClosed)?;
522
523        Ok(max_change_id)
524    }
525}
526
527type StateLock = Arc<(Mutex<MatcherState>, Condvar)>;
528
529pub struct Matcher {
530    pub id: Uuid,
531    pub hash: String,
532    pub query: Stmt,
533    pub cached_statements: HashMap<String, MatcherStmt>,
534    pub pks: IndexMap<String, Vec<String>>,
535    pub parsed: ParsedSelect,
536    pub evt_tx: mpsc::Sender<QueryEvent>,
537    pub col_names: Vec<ColumnName>,
538    pub last_rowid: u64,
539    conn: Connection,
540    base_path: Utf8PathBuf,
541    cancel: CancellationToken,
542    state: StateLock,
543    last_change_tx: watch::Sender<ChangeId>,
544    changes_rx: mpsc::Receiver<MatchCandidates>,
545}
546
547#[derive(Debug, Clone)]
548pub struct MatcherStmt {
549    new_query: String,
550    temp_query: String,
551}
552
553impl MatcherStmt {
554    pub fn new_query(&self) -> &String {
555        &self.new_query
556    }
557}
558
559const CHANGE_ID_COL: &str = "id";
560const CHANGE_TYPE_COL: &str = "type";
561
562pub const QUERY_TABLE_NAME: &str = "query";
563
564pub const SUB_DB_PATH: &str = "sub.sqlite";
565
566impl Matcher {
567    fn new(
568        id: Uuid,
569        subs_path: Utf8PathBuf,
570        schema: &Schema,
571        state_conn: &Connection,
572        evt_tx: mpsc::Sender<QueryEvent>,
573        sql: &str,
574    ) -> Result<(Matcher, MatcherHandle), MatcherError> {
575        let sub_path = Self::sub_path(subs_path.as_path(), id);
576        let sql_hash = hex::encode(seahash::hash(sql.as_bytes()).to_be_bytes());
577
578        info!(%sql_hash, sub_id = %id, "Initializing subscription at {sub_path}");
579
580        std::fs::create_dir_all(&sub_path)?;
581
582        let sub_db_path = sub_path.join(SUB_DB_PATH);
583
584        let col_names: Vec<ColumnName> = {
585            state_conn
586                .prepare(sql)?
587                .column_names()
588                .into_iter()
589                .map(|s| ColumnName(s.to_compact_string()))
590                .collect()
591        };
592
593        let conn = Connection::open(&sub_db_path)?;
594        conn.execute_batch(
595            r#"
596                PRAGMA journal_mode = WAL;
597                PRAGMA synchronous = NORMAL;
598                PRAGMA temp_store = memory;
599                PRAGMA cache_size = -32000; -- 32MB
600                PRAGMA mmap_size = 536870912; -- 512MB
601            "#,
602        )?;
603
604        let mut parser = Parser::new(sql.as_bytes());
605
606        let (mut stmt, parsed) = match parser.next()?.ok_or(MatcherError::StatementRequired)? {
607            Cmd::Stmt(stmt) => {
608                let parsed = match stmt {
609                    Stmt::Select(ref select) => extract_select_columns(select, schema)?,
610                    _ => return Err(MatcherError::UnsupportedStatement),
611                };
612
613                (stmt, parsed)
614            }
615            _ => return Err(MatcherError::StatementRequired),
616        };
617
618        if parsed.table_columns.is_empty() {
619            return Err(MatcherError::TableRequired);
620        }
621
622        let mut statements = HashMap::new();
623
624        let mut pks = IndexMap::default();
625
626        match &mut stmt {
627            Stmt::Select(select) => match &mut select.body.select {
628                OneSelect::Select { columns, .. } => {
629                    let mut new_cols = parsed
630                        .table_columns
631                        .iter()
632                        .filter_map(|(tbl_name, _cols)| {
633                            schema.tables.get(tbl_name).map(|table| {
634                                let tbl_name = parsed
635                                    .aliases
636                                    .iter()
637                                    .find_map(|(alias, actual)| {
638                                        (actual == tbl_name).then_some(alias)
639                                    })
640                                    .unwrap_or(tbl_name);
641                                table
642                                    .pk
643                                    .iter()
644                                    .map(|pk| {
645                                        let alias = format!("__corro_pk_{tbl_name}_{pk}");
646                                        let entry: &mut Vec<String> =
647                                            pks.entry(table.name.clone()).or_default();
648                                        entry.push(alias.clone());
649
650                                        ResultColumn::Expr(
651                                            Expr::Qualified(
652                                                Name(tbl_name.clone().into()),
653                                                Name(pk.clone().into()),
654                                            ),
655                                            Some(As::As(Name(alias.into()))),
656                                        )
657                                    })
658                                    .collect::<Vec<_>>()
659                            })
660                        })
661                        .flatten()
662                        .collect::<Vec<_>>();
663
664                    new_cols.append(&mut parsed.columns.clone());
665                    *columns = new_cols;
666                }
667                _ => unreachable!(),
668            },
669            _ => unreachable!(),
670        }
671
672        for (idx, (tbl_name, _cols)) in parsed.table_columns.iter().enumerate() {
673            let expr = table_to_expr(
674                &parsed.aliases,
675                schema
676                    .tables
677                    .get(tbl_name)
678                    .expect("this should not happen, missing table in schema"),
679                tbl_name,
680            )?;
681
682            let mut stmt = stmt.clone();
683
684            if let Stmt::Select(select) = &mut stmt
685                && let OneSelect::Select {
686                    where_clause, from, ..
687                } = &mut select.body.select
688            {
689                *where_clause = if let Some(prev) = where_clause.take() {
690                    Some(Box::new(Expr::Binary(
691                        Box::new(expr),
692                        Operator::And,
693                        Box::new(Expr::parenthesized(*prev)),
694                    )))
695                } else {
696                    Some(Box::new(expr))
697                };
698
699                match from {
700                    Some(FromClause {
701                        joins: Some(joins), ..
702                    }) if idx > 0 => {
703                        // Replace LEFT JOIN with INNER join if the target is the joined table
704                        if let Some(joined_table) = joins.get_mut(idx - 1)
705                            && let JoinOperator::TypedJoin(Some(join_type)) = &joined_table.operator
706                            && join_type.contains(JoinType::LEFT)
707                        {
708                            joined_table.operator = JoinOperator::TypedJoin(Some(JoinType::INNER));
709                        };
710
711                        // Remove all custom INDEXED BY clauses for the table as the most efficient
712                        // way is to query it by the primary keys
713                        if let Some(JoinedSelectTable {
714                            table: SelectTable::Table(_, _, indexed @ Some(_)),
715                            ..
716                        }) = joins.get_mut(idx - 1)
717                        {
718                            *indexed = None
719                        };
720                    }
721                    _ => (),
722                };
723            }
724
725            let mut new_query = Cmd::Stmt(stmt).to_string();
726            new_query.pop();
727
728            let mut all_cols = pks.values().flatten().cloned().collect::<Vec<String>>();
729            for i in 0..(parsed.columns.len()) {
730                all_cols.push(format!("col_{i}"));
731            }
732
733            let temp_query = format!(
734                "SELECT {} FROM query WHERE ({}) IN temp_{tbl_name}",
735                all_cols.join(","),
736                pks.get(tbl_name)
737                    .cloned()
738                    .ok_or(MatcherError::MissingPrimaryKeys)?
739                    .into_iter()
740                    .map(|pk| format!("coalesce({pk}, \"\")"))
741                    .collect::<Vec<_>>()
742                    .join(","),
743            );
744
745            info!(%sql_hash, sub_id = %id, "modified query for table '{tbl_name}': {new_query}");
746
747            statements.insert(
748                tbl_name.clone(),
749                MatcherStmt {
750                    new_query,
751                    temp_query,
752                },
753            );
754        }
755
756        let cancel = CancellationToken::new();
757
758        let state = Arc::new((Mutex::new(MatcherState::Created), Condvar::new()));
759
760        let (last_change_tx, last_change_rx) = watch::channel(ChangeId(0));
761
762        // big channel to not miss anything
763        let (changes_tx, changes_rx) = mpsc::channel(20480);
764
765        // metrics counters
766        let mut counter_map = HashMap::new();
767        for table in parsed.table_columns.keys() {
768            counter_map.insert(table.clone(), HandleMetrics{
769                matched_count: counter!("corro.subs.changes.matched.count", "sql_hash" => sql_hash.clone(), "table" => table.to_string()),
770            });
771        }
772
773        let handle = MatcherHandle {
774            inner: Arc::new(InnerMatcherHandle {
775                id,
776                sql: sql.to_owned(),
777                hash: sql_hash.clone(),
778                pool: sqlite_pool::Config::new(sub_db_path.into_std_path_buf())
779                    .max_size(5)
780                    .read_only()
781                    .create_pool()
782                    .expect("could not build pool, this can't fail because we specified a runtime"),
783                parsed: parsed.clone(),
784                col_names: col_names.clone(),
785                cancel: cancel.clone(),
786                last_change_rx,
787                changes_tx,
788                cached_statements: statements.clone(),
789                subs_path: sub_path.to_string(),
790                metrics: counter_map,
791            }),
792            state: state.clone(),
793        };
794
795        let matcher = Self {
796            id,
797            hash: sql_hash,
798            query: stmt,
799            cached_statements: statements,
800            pks,
801            parsed,
802            evt_tx,
803            col_names,
804            last_rowid: 0,
805            conn,
806            base_path: sub_path,
807            cancel,
808            state,
809            last_change_tx,
810            changes_rx,
811        };
812
813        Ok((matcher, handle))
814    }
815
816    pub fn cleanup(id: Uuid, sub_path: Utf8PathBuf) -> rusqlite::Result<()> {
817        info!(sub_id = %id, "Attempting to cleanup... {}", sub_path);
818
819        block_in_place(|| {
820            if let Err(e) = std::fs::remove_dir_all(&sub_path) {
821                error!(sub_id = %id, "could not delete subscription base path {} due to: {e}", sub_path);
822            }
823
824            Ok(())
825        })
826    }
827
828    pub fn sub_path(subs_path: &Utf8Path, id: Uuid) -> Utf8PathBuf {
829        subs_path.join(id.as_simple().to_string())
830    }
831
832    pub fn sub_db_path(subs_path: &Utf8Path, id: Uuid) -> Utf8PathBuf {
833        Self::sub_path(subs_path, id).join(SUB_DB_PATH)
834    }
835
836    #[allow(clippy::too_many_arguments)]
837    pub fn restore(
838        id: Uuid,
839        subs_path: Utf8PathBuf,
840        schema: &Schema,
841        state_conn: CrConn,
842        evt_tx: mpsc::Sender<QueryEvent>,
843        tripwire: Tripwire,
844    ) -> Result<MatcherHandle, MatcherError> {
845        let sql: String = block_in_place(|| {
846            let conn = Connection::open(Matcher::sub_db_path(&subs_path, id))?;
847            let state: Option<String> = conn
848                .query_row("SELECT value FROM meta WHERE key = 'state'", [], |row| {
849                    row.get(0)
850                })
851                .optional()?;
852            if !matches!(state.as_deref(), Some("completed")) {
853                return Err(MatcherError::NotRunning);
854            }
855
856            let sql: Option<String> = conn
857                .query_row("SELECT value FROM meta WHERE key = 'sql'", [], |row| {
858                    row.get(0)
859                })
860                .optional()?;
861
862            match sql {
863                Some(sql) => Ok(sql),
864                None => Err(MatcherError::MissingSql),
865            }
866        })?;
867
868        let (matcher, handle) = Self::new(id, subs_path, schema, &state_conn, evt_tx, &sql)?;
869
870        spawn_counted(matcher.run_restore(state_conn, tripwire));
871
872        Ok(handle)
873    }
874
875    #[allow(clippy::too_many_arguments)]
876    pub fn create(
877        id: Uuid,
878        subs_path: Utf8PathBuf,
879        schema: &Schema,
880        state_conn: CrConn,
881        evt_tx: mpsc::Sender<QueryEvent>,
882        sql: &str,
883        tripwire: Tripwire,
884    ) -> Result<MatcherHandle, MatcherError> {
885        let (mut matcher, handle) = Self::new(id, subs_path, schema, &state_conn, evt_tx, sql)?;
886
887        let pk_cols = matcher
888            .pks
889            .values()
890            .flatten()
891            .cloned()
892            .collect::<Vec<String>>();
893
894        let mut all_cols = pk_cols.clone();
895
896        let mut query_cols = vec![];
897
898        for i in 0..(matcher.parsed.columns.len()) {
899            let col_name = format!("col_{i}");
900            all_cols.push(col_name.clone());
901            query_cols.push(col_name);
902        }
903
904        block_in_place(|| {
905            let tx = matcher.conn.transaction()?;
906
907            info!(sub_id = %id, "Creating subscription database schema");
908            let create_temp_table = format!(
909                r#"
910                CREATE TABLE query (__corro_rowid INTEGER PRIMARY KEY AUTOINCREMENT, {columns});
911
912                CREATE UNIQUE INDEX index_{id}_pk ON query ({pks_coalesced});
913
914                CREATE TABLE changes (
915                    {CHANGE_ID_COL} INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
916                    __corro_rowid INTEGER NOT NULL,
917                    {CHANGE_TYPE_COL} INTEGER NOT NULL,
918                    {actual_columns}
919                );
920
921                CREATE TABLE meta (
922                    key TEXT PRIMARY KEY NOT NULL,
923                    value
924                ) WITHOUT ROWID;
925
926                CREATE TABLE columns (
927                    "table" TEXT NOT NULL,
928                    cid TEXT,
929
930                    PRIMARY KEY ("table", cid)
931                );
932            "#,
933                columns = all_cols.join(","),
934                id = id.as_simple(),
935                pks_coalesced = pk_cols
936                    .iter()
937                    .map(|pk| format!("coalesce({pk}, \"\")"))
938                    .collect::<Vec<_>>()
939                    .join(","),
940                actual_columns = query_cols.join(","),
941            );
942
943            tx.execute_batch(&create_temp_table)?;
944            trace!("created sub tables");
945
946            for (table, pks) in matcher.pks.iter() {
947                tx.execute(
948                    &format!(
949                        "CREATE INDEX index_{id}_{table}_pk ON query ({pks})",
950                        id = id.as_simple(),
951                        table = table,
952                        pks = pks.to_vec().join(","),
953                    ),
954                    [],
955                )?;
956            }
957            trace!("created query indexes");
958
959            for (table, columns) in matcher.parsed.table_columns.iter() {
960                tx.execute(
961                    r#"INSERT INTO columns ("table", cid) VALUES (?, '-1')"#,
962                    [table.as_str()],
963                )?;
964                for column in columns.iter() {
965                    trace!("inserting sub column {} => {}", table, column);
966                    tx.execute(
967                        r#"INSERT INTO columns ("table", cid) VALUES (?, ?)"#,
968                        [table.as_str(), column.as_str()],
969                    )?;
970                }
971            }
972            trace!("inserted sub columns");
973
974            tx.execute("INSERT INTO meta (key, value) VALUES ('sql', ?)", [sql])?;
975            tx.execute(
976                "INSERT INTO meta (key, value) VALUES ('state', 'created')",
977                [],
978            )?;
979
980            tx.commit()?;
981            trace!("committed subscription");
982
983            Ok::<_, MatcherError>(())
984        })?;
985
986        spawn_counted(matcher.run(state_conn, tripwire));
987
988        Ok(handle)
989    }
990
991    async fn run_restore(mut self, mut state_conn: CrConn, tripwire: Tripwire) {
992        info!(sub_id = %self.id, "Restoring subscription");
993        let init_res = block_in_place(|| {
994            self.last_rowid = self
995                .conn
996                .query_row(
997                    "SELECT COALESCE(MAX(__corro_rowid), 0) FROM query",
998                    (),
999                    |row| row.get(0),
1000                )
1001                .optional()?
1002                .unwrap_or_default();
1003
1004            let max_change_id = self
1005                .conn
1006                .prepare_cached("SELECT COALESCE(MAX(id), 0) FROM changes")?
1007                .query_row([], |row| row.get(0))?;
1008
1009            _ = self.last_change_tx.send(max_change_id);
1010
1011            Ok::<_, MatcherError>(())
1012        });
1013
1014        if let Err(e) = init_res {
1015            error!(sub_id = %self.id, "could not re-initialize subscription: {e}");
1016            return;
1017        }
1018
1019        if let Err(e) = block_in_place(|| self.setup(&mut state_conn)) {
1020            error!(sub_id = %self.id, "could not setup connection: {e}");
1021            return;
1022        }
1023
1024        if let Err(e) = self.set_status("running") {
1025            error!(sub_id = %self.id, "could not set status: {e}");
1026            return;
1027        }
1028
1029        self.cmd_loop(state_conn, tripwire).await
1030    }
1031
1032    fn set_status(&self, status: &str) -> Result<(), MatcherError> {
1033        self.conn.execute(
1034            "INSERT OR REPLACE INTO meta (key, value) VALUES ('state', ?)",
1035            [status],
1036        )?;
1037        Ok(())
1038    }
1039
1040    fn setup(&self, state_conn: &mut Connection) -> Result<(), MatcherError> {
1041        info!(sub_id = %self.id, "Attaching __corro_sub to state db");
1042        if let Err(e) = state_conn.execute_batch(&format!(
1043            "ATTACH DATABASE {} AS __corro_sub",
1044            enquote::enquote('\'', self.base_path.join(SUB_DB_PATH).as_str()),
1045        )) {
1046            error!(sub_id = %self.id, "could not ATTACH sub db as __corro_sub on state db: {e}");
1047            _ = self.evt_tx.try_send(QueryEvent::Error(format_compact!(
1048                "could not ATTACH subscription db: {e}"
1049            )));
1050            return Err(e.into());
1051        }
1052
1053        info!(sub_id = %self.id, "Attached __corro_sub to state db");
1054
1055        for (tbl_name, pks) in &self.pks {
1056            if let Ok(plan) = dump_query_plan(
1057                state_conn,
1058                &self
1059                    .cached_statements
1060                    .get(tbl_name)
1061                    .ok_or(MatcherError::StatementRequired)?
1062                    .new_query,
1063                tbl_name,
1064                pks,
1065            ) {
1066                info!(sub_id = %self.id, sql_hash = %self.hash, "query plan for table '{tbl_name}':\n{plan}");
1067            }
1068        }
1069
1070        Ok(())
1071    }
1072
1073    async fn cmd_loop(mut self, mut state_conn: CrConn, mut tripwire: Tripwire) {
1074        const PROCESS_CHANGES_THRESHOLD: usize = 1000;
1075        const PROCESSING_WARN_THRESHOLD: Duration = Duration::from_secs(5);
1076        const PROCESS_BUFFER_DEADLINE: Duration = Duration::from_millis(600);
1077
1078        info!(sub_id = %self.id, "Starting loop to run the subscription");
1079        {
1080            let (lock, cvar) = &*self.state;
1081            let mut state = lock.lock();
1082            *state = MatcherState::Running;
1083            cvar.notify_all();
1084            info!(sub_id = %self.id, "Notified condvar that the subscription is 'running'");
1085        }
1086        trace!("set state!");
1087
1088        let mut buf = MatchCandidates::new();
1089        let mut buf_count = 0;
1090
1091        let mut purge_changes_interval = tokio::time::interval(Duration::from_secs(300));
1092
1093        // max duration of aggregating candidates
1094        let process_changes_deadline = tokio::time::sleep(PROCESS_BUFFER_DEADLINE);
1095        tokio::pin!(process_changes_deadline);
1096
1097        loop {
1098            enum Branch {
1099                NewCandidates(MatchCandidates),
1100                PurgeOldChanges,
1101            }
1102
1103            // trace!("looping...");
1104
1105            let branch = tokio::select! {
1106                biased;
1107                _ = self.cancel.cancelled() => {
1108                    info!(sub_id = %self.id, "Acknowledged subscription cancellation, breaking loop.");
1109                    if let Err(e) = self.set_status("cancelled") {
1110                        error!(sub_id = %self.id, "could not set status during cancellation: {e}");
1111                    }
1112                    info!(sub_id = %self.id, "attempting to cleanup");
1113                    if let Err(e) = Self::cleanup(self.id, self.base_path.clone()) {
1114                        error!("could not handle cleanup: {e}");
1115                    }
1116                    return;
1117                }
1118                Some(candidates) = self.changes_rx.recv() => {
1119                    for (table, pks) in  candidates {
1120                        let buffed = buf.entry(table).or_default();
1121                        for (pk, cl) in pks {
1122                            if buffed.insert(pk, cl).is_none() {
1123                                buf_count += 1;
1124                            }
1125                        }
1126                    }
1127
1128                    if buf_count >= PROCESS_CHANGES_THRESHOLD {
1129                        buf_count = 0;
1130                        Branch::NewCandidates(std::mem::take(&mut buf))
1131                    } else {
1132                        continue;
1133                    }
1134                },
1135                _ = process_changes_deadline.as_mut() => {
1136                    process_changes_deadline
1137                        .as_mut()
1138                        .reset((Instant::now() + PROCESS_BUFFER_DEADLINE).into());
1139                    if buf_count == 0 {
1140                        continue;
1141                    }
1142                    Branch::NewCandidates(std::mem::take(&mut buf))
1143                },
1144                _ = &mut tripwire => {
1145                    info!(sub_id = %self.id, "tripped cmd_loop, returning");
1146                    // just return!
1147                    break;
1148                }
1149                _ = purge_changes_interval.tick() => Branch::PurgeOldChanges,
1150                else => {
1151                    return;
1152                }
1153            };
1154
1155            match branch {
1156                Branch::NewCandidates(candidates) => {
1157                    let start = Instant::now();
1158                    if let Err(e) = block_in_place(|| {
1159                        self.handle_candidates(&mut state_conn, candidates, false)
1160                    }) {
1161                        if !matches!(e, MatcherError::EventReceiverClosed) {
1162                            error!(sub_id = %self.id, "could not handle change: {e}");
1163                        }
1164                        info!(sub_id = %self.id, "attempting to cleanup");
1165                        if let Err(e) = Self::cleanup(self.id, self.base_path.clone()) {
1166                            error!("could not handle cleanup: {e}");
1167                        }
1168                        return;
1169                    }
1170                    let elapsed = start.elapsed();
1171
1172                    histogram!("corro.subs.changes.processing.duration.seconds", "sql_hash" => self.hash.clone()).record(elapsed);
1173
1174                    if elapsed >= PROCESSING_WARN_THRESHOLD {
1175                        warn!(sub_id = %self.id, "processed {buf_count} changes (very slowly) for subscription in {elapsed:?}");
1176                    } else {
1177                        debug!(sub_id = %self.id, "processed {buf_count} changes for subscription in {elapsed:?}");
1178                    }
1179                    buf_count = 0;
1180
1181                    // reset the deadline
1182                    process_changes_deadline
1183                        .as_mut()
1184                        .reset((Instant::now() + PROCESS_BUFFER_DEADLINE).into());
1185                }
1186                Branch::PurgeOldChanges => {
1187                    let start = Instant::now();
1188                    let res = block_in_place(|| {
1189                        let tx = self.conn.transaction()?;
1190
1191                        let deleted = tx
1192                            .prepare_cached(
1193                                "DELETE FROM changes WHERE id < (SELECT COALESCE(MAX(id),0) - 500 FROM changes)"
1194                            )?
1195                            .execute([])?;
1196
1197                        tx.commit().map(|_| deleted)
1198                    });
1199
1200                    match res {
1201                        Ok(deleted) => {
1202                            info!(sub_id = %self.id, "Deleted {deleted} old changes row in {:?}", start.elapsed());
1203                        }
1204                        Err(e) => {
1205                            error!(sub_id = %self.id, "could not delete old changes: {e}");
1206                        }
1207                    }
1208                }
1209            }
1210        }
1211
1212        info!(sub_id = %self.id, "draining changes channel");
1213        while let Some(candidates) = self.changes_rx.recv().await {
1214            for (table, pks) in candidates {
1215                let buffed = buf.entry(table).or_default();
1216                for (pk, cl) in pks {
1217                    buffed.insert(pk, cl);
1218                }
1219            }
1220        }
1221
1222        if !buf.is_empty() {
1223            info!(sub_id = %self.id, "handling buffered candidates");
1224            let start = Instant::now();
1225            if let Err(e) = block_in_place(|| self.handle_candidates(&mut state_conn, buf, true)) {
1226                error!(sub_id = %self.id, "could not handle final buffered candidates: {e}");
1227                if let Err(e) = Self::cleanup(self.id, self.base_path.clone()) {
1228                    error!(sub_id = %self.id, "could not handle cleanup: {e}");
1229                }
1230                return;
1231            }
1232            let elapsed = start.elapsed();
1233            info!(sub_id = %self.id, "handled final buffered candidates in {elapsed:?}");
1234        }
1235
1236        if let Err(e) = self.set_status("completed") {
1237            error!(sub_id = %self.id, "could not set status: {e}");
1238        };
1239
1240        info!(sub_id = %self.id, "matcher loop is done");
1241    }
1242
1243    async fn run(mut self, mut state_conn: CrConn, tripwire: Tripwire) {
1244        info!(sub_id = %self.id, "Running initial query");
1245        if let Err(e) = self
1246            .evt_tx
1247            .send(QueryEvent::Columns(self.col_names.clone()))
1248            .await
1249        {
1250            error!(sub_id = %self.id, "could not send back columns, probably means no receivers! {e}");
1251            return;
1252        }
1253
1254        let mut query_cols = vec![];
1255        for i in 0..(self.parsed.columns.len()) {
1256            query_cols.push(format!("col_{i}"));
1257        }
1258
1259        let res = block_in_place(|| {
1260            let tx = self.conn.transaction()?;
1261
1262            let mut stmt_str = Cmd::Stmt(self.query.clone()).to_string();
1263            stmt_str.pop(); // remove trailing `;`
1264
1265            let mut all_cols = self
1266                .pks
1267                .values()
1268                .flatten()
1269                .cloned()
1270                .collect::<Vec<String>>();
1271
1272            for i in 0..(self.parsed.columns.len()) {
1273                let col_name = format!("col_{i}");
1274                all_cols.push(col_name.clone());
1275            }
1276
1277            let mut last_rowid = 0;
1278
1279            // ensure drop and recreate
1280            tx.execute_batch(&format!(
1281                "DROP TABLE IF EXISTS state_rows;
1282                    CREATE TEMP TABLE state_rows ({})",
1283                all_cols.join(",")
1284            ))?;
1285
1286            info!(sub_id = %self.id, "Starting state conn read transaction for initial query");
1287            // this is read transaction up until the end!
1288            let state_tx = state_conn.transaction()?;
1289
1290            let elapsed = {
1291                debug!("select stmt: {stmt_str:?}");
1292
1293                let mut select = state_tx.prepare(&stmt_str)?;
1294                let start = Instant::now();
1295                let mut select_rows = {
1296                    let _guard = interrupt_deadline_guard(&state_tx, Duration::from_secs(15));
1297                    select.query(())?
1298                };
1299                let elapsed = start.elapsed();
1300                info!(sub_id = %self.id, "Initial query done in {elapsed:?}");
1301
1302                let insert_into = format!(
1303                    "INSERT INTO query ({}) VALUES ({}) RETURNING __corro_rowid,{}",
1304                    all_cols.join(","),
1305                    all_cols
1306                        .iter()
1307                        .map(|_| "?".to_string())
1308                        .collect::<Vec<_>>()
1309                        .join(","),
1310                    query_cols.join(","),
1311                );
1312                trace!("insert stmt: {insert_into:?}");
1313
1314                {
1315                    let mut insert = tx.prepare(&insert_into)?;
1316
1317                    loop {
1318                        match select_rows.next() {
1319                            Ok(Some(row)) => {
1320                                for i in 0..all_cols.len() {
1321                                    insert.raw_bind_parameter(
1322                                        i + 1,
1323                                        SqliteValueRef::from(row.get_ref(i)?),
1324                                    )?;
1325                                }
1326
1327                                let mut rows = insert.raw_query();
1328
1329                                let row = match rows.next()? {
1330                                    Some(row) => row,
1331                                    None => continue,
1332                                };
1333
1334                                let rowid = row.get(0)?;
1335                                let cells = (1..=query_cols.len())
1336                                    .map(|i| row.get::<_, SqliteValue>(i))
1337                                    .collect::<rusqlite::Result<Vec<_>>>()?;
1338
1339                                if let Err(e) = self
1340                                    .evt_tx
1341                                    .blocking_send(QueryEvent::Row(RowId(rowid), cells))
1342                                {
1343                                    error!(sub_id = %self.id, "could not send back row: {e}");
1344                                    return Err(MatcherError::EventReceiverClosed);
1345                                }
1346
1347                                last_rowid = cmp::max(rowid, last_rowid);
1348                            }
1349                            Ok(None) => {
1350                                info!(sub_id = %self.id, "Done iterating through rows for initial query");
1351                                // done!
1352                                break;
1353                            }
1354                            Err(e) => {
1355                                return Err(e.into());
1356                            }
1357                        }
1358                    }
1359                }
1360
1361                tx.execute_batch("DROP TABLE IF EXISTS state_rows;")?;
1362
1363                tx.execute(
1364                    "INSERT OR REPLACE INTO meta (key, value) VALUES ('state', 'running')",
1365                    [],
1366                )?;
1367
1368                tx.commit()?;
1369
1370                elapsed
1371            };
1372
1373            self.last_rowid = last_rowid;
1374
1375            Ok::<_, MatcherError>(elapsed)
1376        });
1377
1378        match res {
1379            Ok(elapsed) => {
1380                trace!(
1381                    "done w/ block_in_place for initial query, elapsed: {:?}",
1382                    elapsed
1383                );
1384                if let Err(e) = self
1385                    .evt_tx
1386                    .send(QueryEvent::EndOfQuery {
1387                        time: elapsed.as_secs_f64(),
1388                        change_id: Some(ChangeId(0)),
1389                    })
1390                    .await
1391                {
1392                    error!(sub_id = %self.id, "could not return end of query event: {e}");
1393                    return;
1394                }
1395                // db_version
1396            }
1397            Err(e) => {
1398                warn!(sub_id = %self.id, "could not complete initial query: {e}");
1399                _ = self
1400                    .evt_tx
1401                    .send(QueryEvent::Error(e.to_compact_string()))
1402                    .await;
1403
1404                return;
1405            }
1406        };
1407
1408        if let Err(e) = block_in_place(|| self.setup(&mut state_conn)) {
1409            error!(sub_id = %self.id, "could not setup connection: {e}");
1410            return;
1411        }
1412
1413        self.cmd_loop(state_conn, tripwire).await
1414    }
1415
1416    fn handle_candidates(
1417        &mut self,
1418        state_conn: &mut Connection,
1419        candidates: MatchCandidates,
1420        skip_send: bool,
1421    ) -> Result<(), MatcherError> {
1422        let mut tables = IndexSet::new();
1423
1424        if candidates.is_empty() {
1425            return Ok(());
1426        }
1427
1428        trace!(
1429            "got some candidates! {:?}",
1430            candidates.keys().collect::<Vec<_>>()
1431        );
1432
1433        let tx = self.conn.transaction()?;
1434        for (table, pks) in candidates {
1435            let pks = pks
1436                .iter()
1437                .map(|(pk, _)| unpack_columns(pk))
1438                .collect::<Result<Vec<Vec<SqliteValueRef>>, _>>()?;
1439
1440            let tmp_table_name = format!("temp_{table}");
1441            if tables.insert(table.clone()) {
1442                // create a temporary table to mix and match the data
1443                tx.prepare_cached(
1444                    // TODO: cache the statement's string somewhere, it's always the same!
1445                    // NOTE: this can't be an actual CREATE TEMP TABLE or it won't be visible
1446                    //       from the state db
1447                    &format!(
1448                        "CREATE TABLE IF NOT EXISTS {tmp_table_name} ({})",
1449                        self.pks
1450                            .get(table.as_str())
1451                            .ok_or_else(|| MatcherError::MissingPrimaryKeys)?
1452                            .to_vec()
1453                            .join(",")
1454                    ),
1455                )?
1456                .execute(())?;
1457            }
1458
1459            for pks in pks {
1460                tx.prepare_cached(&format!(
1461                    "INSERT INTO {tmp_table_name} VALUES ({})",
1462                    (0..pks.len())
1463                        .map(|_i| "coalesce(?, \"\")")
1464                        .collect::<Vec<_>>()
1465                        .join(",")
1466                ))?
1467                .execute(params_from_iter(pks))?;
1468            }
1469        }
1470        // commit so it is visible to the other db!
1471        tx.commit()?;
1472        trace!("committed temp pk tables");
1473
1474        let mut query_cols = vec![];
1475
1476        let pk_cols = self
1477            .pks
1478            .values()
1479            .flatten()
1480            .cloned()
1481            .collect::<Vec<String>>();
1482
1483        let mut all_cols = pk_cols.clone();
1484        for i in 0..(self.parsed.columns.len()) {
1485            let col_name = format!("col_{i}");
1486            all_cols.push(col_name.clone());
1487            query_cols.push(col_name);
1488        }
1489
1490        // start a new tx
1491        let tx = self.conn.transaction()?;
1492
1493        // ensure drop and recreate
1494        tx.execute_batch(&format!(
1495            "CREATE TEMP TABLE IF NOT EXISTS state_results ({})",
1496            all_cols.join(",")
1497        ))?;
1498
1499        let mut new_last_rowid = self.last_rowid;
1500
1501        {
1502            // read-only!
1503            let state_tx = state_conn.transaction()?;
1504            let mut tmp_insert_prepped = tx.prepare_cached(&format!(
1505                "INSERT INTO state_results VALUES ({})",
1506                (0..all_cols.len())
1507                    .map(|_| "?")
1508                    .collect::<Vec<_>>()
1509                    .join(",")
1510            ))?;
1511
1512            for table in tables.iter() {
1513                let start = Instant::now();
1514                let stmt = match self.cached_statements.get(table.as_str()) {
1515                    Some(stmt) => stmt,
1516                    None => {
1517                        warn!(sub_id = %self.id, "no statements pre-computed for table {table}");
1518                        continue;
1519                    }
1520                };
1521
1522                trace!("SELECT SQL: {}", stmt.new_query);
1523
1524                let mut prepped = state_tx.prepare_cached(&stmt.new_query)?;
1525
1526                let col_count = prepped.column_count();
1527
1528                let mut rows = prepped.raw_query();
1529
1530                // insert each row in the other DB...
1531                while let Some(row) = rows.next()? {
1532                    for i in 0..col_count {
1533                        tmp_insert_prepped
1534                            .raw_bind_parameter(i + 1, SqliteValueRef::from(row.get_ref(i)?))?;
1535                    }
1536                    tmp_insert_prepped.raw_execute()?;
1537                }
1538
1539                let coalesced_pks = pk_cols
1540                    .iter()
1541                    .map(|pk| format!("coalesce({pk},\"\")"))
1542                    .collect::<Vec<_>>()
1543                    .join(",");
1544                let sql = format!(
1545                    "INSERT INTO query ({insert_cols})
1546                        SELECT * FROM (
1547                            SELECT * FROM state_results
1548                            EXCEPT
1549                            {query_query}
1550                        ) WHERE 1
1551                        ON CONFLICT({conflict_clause})
1552                            DO UPDATE SET
1553                                {excluded}
1554                            WHERE {excluded_not_same}
1555                        RETURNING __corro_rowid,{return_cols}",
1556                    // insert into
1557                    insert_cols = all_cols.join(","),
1558                    query_query = stmt.temp_query,
1559                    conflict_clause = coalesced_pks,
1560                    excluded = (0..(self.parsed.columns.len()))
1561                        .map(|i| format!("col_{i} = excluded.col_{i}"))
1562                        .collect::<Vec<_>>()
1563                        .join(","),
1564                    excluded_not_same = (0..(self.parsed.columns.len()))
1565                        .map(|i| format!("col_{i} IS NOT excluded.col_{i}"))
1566                        .collect::<Vec<_>>()
1567                        .join(" OR "),
1568                    return_cols = query_cols.join(",")
1569                );
1570
1571                trace!("INSERT SQL: {sql}");
1572
1573                let insert_prepped = tx.prepare_cached(&sql)?;
1574
1575                let sql = format!(
1576                    "
1577                    DELETE FROM query WHERE ({pks}) IN (SELECT {select_pks} FROM (
1578                        {query_query}
1579                        EXCEPT
1580                        SELECT * FROM state_results
1581                    )) RETURNING __corro_rowid,{return_cols}
1582                ",
1583                    // delete from
1584                    pks = coalesced_pks,
1585                    select_pks = coalesced_pks,
1586                    query_query = stmt.temp_query,
1587                    return_cols = query_cols.join(",")
1588                );
1589
1590                trace!("DELETE SQL: {sql}");
1591
1592                let delete_prepped = tx.prepare_cached(&sql)?;
1593
1594                let mut change_insert_stmt = tx.prepare_cached(&format!(
1595                    "INSERT INTO changes (__corro_rowid, {CHANGE_TYPE_COL}, {}) VALUES (?, ?, {}) RETURNING {CHANGE_ID_COL}",
1596                    query_cols.join(","),
1597                    (0..query_cols.len())
1598                        .map(|_i| "?")
1599                        .collect::<Vec<_>>()
1600                        .join(",")
1601                ))?;
1602
1603                for (change_type, mut prepped) in [
1604                    (None, insert_prepped),
1605                    (Some(ChangeType::Delete), delete_prepped),
1606                ] {
1607                    let col_count = prepped.column_count();
1608
1609                    let mut rows = prepped.raw_query();
1610
1611                    while let Ok(Some(row)) = rows.next() {
1612                        let rowid: RowId = row.get(0)?;
1613
1614                        let change_type = change_type.unwrap_or({
1615                            if rowid.0 > self.last_rowid {
1616                                ChangeType::Insert
1617                            } else {
1618                                ChangeType::Update
1619                            }
1620                        });
1621
1622                        new_last_rowid = cmp::max(new_last_rowid, rowid.0);
1623
1624                        match (1..col_count)
1625                            .map(|i| row.get::<_, SqliteValue>(i))
1626                            .collect::<rusqlite::Result<Vec<_>>>()
1627                        {
1628                            Ok(cells) => {
1629                                change_insert_stmt.raw_bind_parameter(1, rowid)?;
1630                                change_insert_stmt.raw_bind_parameter(2, change_type)?;
1631                                for (i, cell) in cells.iter().enumerate() {
1632                                    // increment index by 3 because that's where we're starting...
1633                                    change_insert_stmt.raw_bind_parameter(i + 3, cell)?;
1634                                }
1635
1636                                let mut change_rows = change_insert_stmt.raw_query();
1637
1638                                let change_id: ChangeId = change_rows
1639                                    .next()?
1640                                    .ok_or(MatcherError::NoChangeInserted)?
1641                                    .get(0)?;
1642
1643                                trace!("got change id: {change_id}");
1644
1645                                if !skip_send
1646                                    && let Err(e) = self.evt_tx.blocking_send(QueryEvent::Change(
1647                                        change_type,
1648                                        rowid,
1649                                        cells,
1650                                        change_id,
1651                                    ))
1652                                {
1653                                    warn!("could not send back row to matcher sub sender: {e}");
1654                                    return Err(MatcherError::EventReceiverClosed);
1655                                }
1656                                _ = self.last_change_tx.send(change_id);
1657                            }
1658                            Err(e) => {
1659                                error!("could not deserialize row's cells: {e}");
1660                                return Ok(());
1661                            }
1662                        }
1663                    }
1664                }
1665                // clean that up
1666                tx.execute_batch("DELETE FROM state_results")?;
1667
1668                let elapsed = start.elapsed();
1669                histogram!("corro.subs.changes.processing.table.duration.seconds", "sql_hash" => self.hash.clone(), "table" => table.0.to_string()).record(elapsed);
1670            }
1671
1672            // clean up temporary tables immediately
1673            for table in tables {
1674                // TODO: reduce mistakes by computing this table name once
1675                tx.prepare_cached(&format!("DELETE FROM temp_{table}",))?
1676                    .execute(())?;
1677                trace!("cleaned up temp_{table}");
1678            }
1679        }
1680
1681        tx.commit()?;
1682
1683        trace!("committed!");
1684
1685        self.last_rowid = new_last_rowid;
1686
1687        Ok(())
1688    }
1689}
1690
1691fn dump_query_plan(
1692    conn: &mut Connection,
1693    query: &str,
1694    table_name: &str,
1695    pks: &[String],
1696) -> Result<String, MatcherError> {
1697    let tx = conn.transaction()?;
1698
1699    tx.prepare_cached(&format!(
1700        "CREATE TABLE IF NOT EXISTS __corro_sub.temp_{table_name} ({})",
1701        pks.join(",")
1702    ))?
1703    .execute(())?;
1704
1705    let mut prepped = tx.prepare(&format!("EXPLAIN QUERY PLAN {query}"))?;
1706    let mut rows = prepped.query(())?;
1707
1708    let mut output = String::new();
1709    while let Some(row) = rows.next()? {
1710        output.push_str(&row.get::<_, String>("detail")?);
1711        output.push('\n');
1712    }
1713
1714    Ok(output)
1715}
1716
1717fn interrupt_deadline_guard(conn: &Connection, dur: Duration) -> DropGuard {
1718    let int_handle = conn.get_interrupt_handle();
1719    let cancel = CancellationToken::new();
1720    tokio::spawn({
1721        let cancel = cancel.clone();
1722        async move {
1723            match tokio::time::sleep(dur)
1724                .preemptible(cancel.cancelled())
1725                .await
1726            {
1727                Outcome::Completed(_) => {
1728                    warn!("subscription query deadline reached, interrupting!");
1729                    int_handle.interrupt();
1730                    // no need to send any query event, it should bubble up properly and if not
1731                    // then it means the conn was not interrupted in time which is also fine
1732                }
1733                Outcome::Preempted(_) => {
1734                    debug!("deadline was canceled, not interrupting query");
1735                }
1736            }
1737        }
1738    });
1739    cancel.drop_guard()
1740}
1741
1742#[derive(Debug, Default, Clone)]
1743pub struct ParsedSelect {
1744    table_columns: IndexMap<String, HashSet<String>>,
1745    aliases: HashMap<String, String>,
1746    pub columns: Vec<ResultColumn>,
1747    children: Vec<ParsedSelect>,
1748}
1749
1750fn extract_select_columns(select: &Select, schema: &Schema) -> Result<ParsedSelect, MatcherError> {
1751    let mut parsed = ParsedSelect::default();
1752
1753    if let OneSelect::Select {
1754        ref from,
1755        ref columns,
1756        ref where_clause,
1757        ..
1758    } = select.body.select
1759    {
1760        let from_table = match from {
1761            Some(from) => {
1762                let from_table = match &from.select {
1763                    Some(table) => match table.as_ref() {
1764                        SelectTable::Table(name, alias, _) => {
1765                            if schema.tables.contains_key(name.name.0.as_ref()) {
1766                                if let Some(As::As(alias) | As::Elided(alias)) = alias {
1767                                    parsed
1768                                        .aliases
1769                                        .insert(alias.0.to_string(), name.name.0.to_string());
1770                                } else if let Some(ref alias) = name.alias {
1771                                    parsed
1772                                        .aliases
1773                                        .insert(alias.0.to_string(), name.name.0.to_string());
1774                                }
1775                                parsed
1776                                    .table_columns
1777                                    .entry(name.name.0.to_string())
1778                                    .or_default();
1779                                Some(&name.name)
1780                            } else {
1781                                return Err(MatcherError::TableNotFound(name.name.0.to_string()));
1782                            }
1783                        }
1784                        // TODO: add support for:
1785                        // TableCall(QualifiedName, Option<Vec<Expr>>, Option<As>),
1786                        // Select(Select, Option<As>),
1787                        // Sub(FromClause, Option<As>),
1788                        t => {
1789                            warn!("ignoring {t:?}");
1790                            None
1791                        }
1792                    },
1793                    _ => {
1794                        // according to the sqlite3-parser docs, this can't really happen
1795                        // ignore!
1796                        unreachable!()
1797                    }
1798                };
1799                if let Some(ref joins) = from.joins {
1800                    for join in joins.iter() {
1801                        // let mut tbl_name = None;
1802                        let tbl_name = match &join.table {
1803                            SelectTable::Table(name, alias, _) => {
1804                                if let Some(As::As(alias) | As::Elided(alias)) = alias {
1805                                    parsed
1806                                        .aliases
1807                                        .insert(alias.0.to_string(), name.name.0.to_string());
1808                                } else if let Some(ref alias) = name.alias {
1809                                    parsed
1810                                        .aliases
1811                                        .insert(alias.0.to_string(), name.name.0.to_string());
1812                                }
1813                                parsed
1814                                    .table_columns
1815                                    .entry(name.name.0.to_string())
1816                                    .or_default();
1817                                &name.name
1818                            }
1819                            // TODO: add support for:
1820                            // TableCall(QualifiedName, Option<Vec<Expr>>, Option<As>),
1821                            // Select(Select, Option<As>),
1822                            // Sub(FromClause, Option<As>),
1823                            t => {
1824                                warn!("ignoring JOIN's non-SelectTable::Table:  {t:?}");
1825                                continue;
1826                            }
1827                        };
1828                        // ON or USING
1829                        if let Some(constraint) = &join.constraint {
1830                            match constraint {
1831                                JoinConstraint::On(expr) => {
1832                                    extract_expr_columns(expr, schema, &mut parsed)?;
1833                                }
1834                                JoinConstraint::Using(names) => {
1835                                    let entry = parsed
1836                                        .table_columns
1837                                        .entry(tbl_name.0.to_string())
1838                                        .or_default();
1839                                    for name in names.iter() {
1840                                        insert_col(entry, schema, &tbl_name.0, &name.0);
1841                                    }
1842                                }
1843                            }
1844                        }
1845                    }
1846                }
1847                if let Some(expr) = where_clause {
1848                    extract_expr_columns(expr, schema, &mut parsed)?;
1849                }
1850                from_table
1851            }
1852            _ => None,
1853        };
1854
1855        extract_columns(columns.as_slice(), from_table, schema, &mut parsed)?;
1856    }
1857
1858    Ok(parsed)
1859}
1860
1861fn insert_col(set: &mut HashSet<String>, schema: &Schema, tbl_name: &str, name: &str) {
1862    let table = schema.tables.get(tbl_name);
1863    if let Some(generated) =
1864        table.and_then(|tbl| tbl.columns.get(name).and_then(|col| col.generated.as_ref()))
1865    {
1866        // recursively check for generated columns
1867        for name in generated.from.iter() {
1868            insert_col(set, schema, tbl_name, name);
1869        }
1870    } else {
1871        set.insert(name.to_owned());
1872    }
1873}
1874
1875fn extract_expr_columns(
1876    expr: &Expr,
1877    schema: &Schema,
1878    parsed: &mut ParsedSelect,
1879) -> Result<(), MatcherError> {
1880    match expr {
1881        // simplest case
1882        Expr::Qualified(tblname, colname) => {
1883            let resolved_name = parsed
1884                .aliases
1885                .get(tblname.0.as_ref())
1886                .map(|s| s.as_str())
1887                .unwrap_or(tblname.0.as_ref());
1888            // println!("adding column: {resolved_name} => {colname:?}");
1889            insert_col(
1890                parsed
1891                    .table_columns
1892                    .entry(resolved_name.to_string())
1893                    .or_default(),
1894                schema,
1895                resolved_name,
1896                &colname.0,
1897            );
1898        }
1899        // simplest case but also mentioning the schema
1900        Expr::DoublyQualified(schema_name, tblname, colname)
1901            if schema_name.0.as_ref() == "main" =>
1902        {
1903            let resolved_name = parsed
1904                .aliases
1905                .get(tblname.0.as_ref())
1906                .map(|s| s.as_str())
1907                .unwrap_or(tblname.0.as_ref());
1908            // println!("adding column: {resolved_name} => {colname:?}");
1909            insert_col(
1910                parsed
1911                    .table_columns
1912                    .entry(resolved_name.to_string())
1913                    .or_default(),
1914                schema,
1915                resolved_name,
1916                &colname.0,
1917            );
1918        }
1919
1920        Expr::Name(colname) => {
1921            let check_col_name = unquote(&colname.0).ok().unwrap_or(colname.0.to_string());
1922
1923            let mut found = None;
1924            for tbl in parsed.table_columns.keys() {
1925                if let Some(tbl) = schema.tables.get(tbl)
1926                    && tbl.columns.contains_key(&check_col_name)
1927                {
1928                    if found.is_some() {
1929                        return Err(MatcherError::QualificationRequired {
1930                            col_name: check_col_name,
1931                        });
1932                    }
1933                    found = Some(tbl.name.as_str());
1934                }
1935            }
1936
1937            if let Some(found) = found {
1938                insert_col(
1939                    parsed.table_columns.entry(found.to_owned()).or_default(),
1940                    schema,
1941                    found,
1942                    &check_col_name,
1943                );
1944            } else {
1945                return Err(MatcherError::TableForColumnNotFound {
1946                    col_name: check_col_name,
1947                });
1948            }
1949        }
1950
1951        Expr::Id(colname) => {
1952            let check_col_name = unquote(&colname.0).ok().unwrap_or(colname.0.to_string());
1953
1954            let mut found = None;
1955            for tbl in parsed.table_columns.keys() {
1956                if let Some(tbl) = schema.tables.get(tbl)
1957                    && tbl.columns.contains_key(&check_col_name)
1958                {
1959                    if found.is_some() {
1960                        return Err(MatcherError::QualificationRequired {
1961                            col_name: check_col_name,
1962                        });
1963                    }
1964                    found = Some(tbl.name.as_str());
1965                }
1966            }
1967
1968            if let Some(found) = found {
1969                insert_col(
1970                    parsed.table_columns.entry(found.to_owned()).or_default(),
1971                    schema,
1972                    found,
1973                    &colname.0,
1974                );
1975            } else {
1976                if colname.0.starts_with('"') {
1977                    return Ok(());
1978                }
1979                return Err(MatcherError::TableForColumnNotFound {
1980                    col_name: colname.0.to_string(),
1981                });
1982            }
1983        }
1984
1985        Expr::Between { lhs, .. } => extract_expr_columns(lhs, schema, parsed)?,
1986        Expr::Binary(lhs, _, rhs) => {
1987            extract_expr_columns(lhs, schema, parsed)?;
1988            extract_expr_columns(rhs, schema, parsed)?;
1989        }
1990        Expr::Case {
1991            base,
1992            when_then_pairs,
1993            else_expr,
1994        } => {
1995            if let Some(expr) = base {
1996                extract_expr_columns(expr, schema, parsed)?;
1997            }
1998            for (when_expr, _then_expr) in when_then_pairs.iter() {
1999                // NOTE: should we also parse the then expr?
2000                extract_expr_columns(when_expr, schema, parsed)?;
2001            }
2002            if let Some(expr) = else_expr {
2003                extract_expr_columns(expr, schema, parsed)?;
2004            }
2005        }
2006        Expr::Cast { expr, .. } => extract_expr_columns(expr, schema, parsed)?,
2007        Expr::Collate(expr, _) => extract_expr_columns(expr, schema, parsed)?,
2008        Expr::Exists(select) => {
2009            parsed
2010                .children
2011                .push(extract_select_columns(select, schema)?);
2012        }
2013        Expr::FunctionCall { args, .. } => {
2014            if let Some(args) = args {
2015                for expr in args.iter() {
2016                    extract_expr_columns(expr, schema, parsed)?;
2017                }
2018            }
2019        }
2020        Expr::InList { lhs, rhs, .. } => {
2021            extract_expr_columns(lhs, schema, parsed)?;
2022            if let Some(rhs) = rhs {
2023                for expr in rhs.iter() {
2024                    extract_expr_columns(expr, schema, parsed)?;
2025                }
2026            }
2027        }
2028        Expr::InSelect { lhs, rhs, .. } => {
2029            extract_expr_columns(lhs, schema, parsed)?;
2030            parsed.children.push(extract_select_columns(rhs, schema)?);
2031        }
2032        expr @ Expr::InTable { .. } => {
2033            return Err(MatcherError::UnsupportedExpr { expr: expr.clone() });
2034        }
2035        Expr::IsNull(expr) => {
2036            extract_expr_columns(expr, schema, parsed)?;
2037        }
2038        Expr::Like { lhs, rhs, .. } => {
2039            extract_expr_columns(lhs, schema, parsed)?;
2040            extract_expr_columns(rhs, schema, parsed)?;
2041        }
2042
2043        Expr::NotNull(expr) => {
2044            extract_expr_columns(expr, schema, parsed)?;
2045        }
2046        Expr::Parenthesized(parens) => {
2047            for expr in parens.iter() {
2048                extract_expr_columns(expr, schema, parsed)?;
2049            }
2050        }
2051        Expr::Subquery(select) => {
2052            parsed
2053                .children
2054                .push(extract_select_columns(select, schema)?);
2055        }
2056        Expr::Unary(_, expr) => {
2057            extract_expr_columns(expr, schema, parsed)?;
2058        }
2059
2060        // no column names in there...
2061        // Expr::FunctionCallStar { name, filter_over } => todo!(),
2062        // Expr::Id(_) => todo!(),
2063        // Expr::Literal(_) => todo!(),
2064        // Expr::Raise(_, _) => todo!(),
2065        // Expr::Variable(_) => todo!(),
2066        _ => {}
2067    }
2068
2069    Ok(())
2070}
2071
2072fn extract_columns(
2073    columns: &[ResultColumn],
2074    from: Option<&Name>,
2075    schema: &Schema,
2076    parsed: &mut ParsedSelect,
2077) -> Result<(), MatcherError> {
2078    let mut i = 0;
2079    for col in columns.iter() {
2080        match col {
2081            ResultColumn::Expr(expr, _) => {
2082                // println!("extracting col: {expr:?} (as: {maybe_as:?})");
2083                extract_expr_columns(expr, schema, parsed)?;
2084                parsed.columns.push(ResultColumn::Expr(
2085                    expr.clone(),
2086                    Some(As::As(Name(format!("col_{i}").into()))),
2087                ));
2088                i += 1;
2089            }
2090            ResultColumn::Star => {
2091                if let Some(tbl_name) = from {
2092                    if let Some(table) = schema.tables.get(tbl_name.0.as_ref()) {
2093                        let entry = parsed.table_columns.entry(table.name.clone()).or_default();
2094                        for col in table.columns.keys() {
2095                            entry.insert(col.clone());
2096                            parsed.columns.push(ResultColumn::Expr(
2097                                Expr::Name(Name(col.clone().into())),
2098                                Some(As::As(Name(format!("col_{i}").into()))),
2099                            ));
2100                            i += 1;
2101                        }
2102                    } else {
2103                        return Err(MatcherError::TableStarNotFound {
2104                            tbl_name: tbl_name.0.to_string(),
2105                        });
2106                    }
2107                } else {
2108                    unreachable!()
2109                }
2110            }
2111            ResultColumn::TableStar(tbl_name) => {
2112                let name = parsed
2113                    .aliases
2114                    .get(tbl_name.0.as_ref())
2115                    .map(|s| s.as_str())
2116                    .unwrap_or(tbl_name.0.as_ref());
2117                if let Some(table) = schema.tables.get(name) {
2118                    let entry = parsed.table_columns.entry(table.name.clone()).or_default();
2119                    for col in table.columns.keys() {
2120                        entry.insert(col.clone());
2121                        parsed.columns.push(ResultColumn::Expr(
2122                            Expr::Qualified(tbl_name.clone(), Name(col.clone().into())),
2123                            Some(As::As(Name(format!("col_{i}").into()))),
2124                        ));
2125                        i += 1;
2126                    }
2127                } else {
2128                    return Err(MatcherError::TableStarNotFound {
2129                        tbl_name: name.to_string(),
2130                    });
2131                }
2132            }
2133        }
2134    }
2135    Ok(())
2136}
2137
2138fn table_to_expr(
2139    aliases: &HashMap<String, String>,
2140    tbl: &Table,
2141    table: &str,
2142) -> Result<Expr, MatcherError> {
2143    let tbl_name = aliases
2144        .iter()
2145        .find_map(|(alias, actual)| (actual == table).then_some(alias))
2146        .cloned()
2147        .unwrap_or_else(|| table.to_owned());
2148
2149    let expr = Expr::in_table(
2150        Expr::Parenthesized(
2151            tbl.pk
2152                .iter()
2153                .map(|pk| {
2154                    Expr::Qualified(Name(tbl_name.clone().into()), Name(pk.to_owned().into()))
2155                })
2156                .collect(),
2157        ),
2158        false,
2159        QualifiedName::fullname(
2160            Name("__corro_sub".into()),
2161            Name(format!("temp_{table}").into()),
2162        ),
2163        None,
2164    );
2165
2166    Ok(expr)
2167}
2168
2169#[derive(Debug, thiserror::Error)]
2170pub enum MatcherError {
2171    #[error(transparent)]
2172    Lexer(#[from] sqlite3_parser::lexer::sql::Error),
2173    #[error("one statement is required for matching")]
2174    StatementRequired,
2175    #[error("unsupported statement")]
2176    UnsupportedStatement,
2177    #[error("at least 1 table is required in FROM / JOIN clause")]
2178    TableRequired,
2179    #[error(transparent)]
2180    Sqlite(#[from] rusqlite::Error),
2181    #[error("table not found in schema: {0}")]
2182    TableNotFound(String),
2183    #[error("no primary key for table: {0}")]
2184    NoPrimaryKey(String),
2185    #[error("aggregate missing primary key {0}.{1}")]
2186    AggPrimaryKeyMissing(String, String),
2187    #[error("JOIN .. ON expression is not supported for join on table '{table}': {expr:?}")]
2188    JoinOnExprUnsupported { table: String, expr: Box<Expr> },
2189    #[error("expression is not supported: {expr:?}")]
2190    UnsupportedExpr { expr: Expr },
2191    #[error("could not find table for {tbl_name}.* in corrosion's schema")]
2192    TableStarNotFound { tbl_name: String },
2193    #[error("<tbl>.{col_name} qualification required for ambiguous column name")]
2194    QualificationRequired { col_name: String },
2195    #[error("could not find table for column {col_name}")]
2196    TableForColumnNotFound { col_name: String },
2197    #[error("missing primary keys, this shouldn't happen")]
2198    MissingPrimaryKeys,
2199    #[error("change queue has been closed or is full")]
2200    ChangeQueueClosedOrFull,
2201    #[error("no change was inserted, this is not supposed to happen")]
2202    NoChangeInserted,
2203    #[error("change receiver is closed")]
2204    EventReceiverClosed,
2205    #[error(transparent)]
2206    Unpack(#[from] UnpackError),
2207    #[error("did not insert subscription")]
2208    InsertSub,
2209    #[error(transparent)]
2210    FromSql(#[from] FromSqlError),
2211    #[error(transparent)]
2212    Io(#[from] std::io::Error),
2213    #[error("could not encode changeset for queue: {0}")]
2214    ChangesetEncode(#[from] speedy::Error),
2215    #[error("queue is full")]
2216    QueueFull,
2217    #[error("cannot restore existing subscription")]
2218    CannotRestoreExisting,
2219    #[error("could not acquire write permit")]
2220    WritePermitAcquire(#[from] AcquireError),
2221    #[error("subscription is not running")]
2222    NotRunning,
2223    #[error("subscription restore is missing SQL query")]
2224    MissingSql,
2225}
2226
2227impl MatcherError {
2228    pub fn is_event_recv_closed(&self) -> bool {
2229        matches!(self, MatcherError::EventReceiverClosed)
2230    }
2231}
2232
2233#[derive(Debug, thiserror::Error)]
2234#[allow(clippy::large_enum_variant)]
2235pub enum NormalizeStatementError {
2236    #[error(transparent)]
2237    Parse(#[from] sqlite3_parser::lexer::sql::Error),
2238    #[error("unexpected statement: {0}")]
2239    UnexpectedStatement(Box<Cmd>),
2240    #[error("only 1 statement is supported")]
2241    Multiple,
2242    #[error("at least 1 statement is required")]
2243    NoStatement,
2244}
2245
2246pub fn normalize_sql(sql: &str) -> Result<String, NormalizeStatementError> {
2247    let mut parser = Parser::new(sql.as_bytes());
2248
2249    let stmt = match parser.next()? {
2250        Some(Cmd::Stmt(stmt)) => stmt,
2251        Some(cmd) => {
2252            return Err(NormalizeStatementError::UnexpectedStatement(Box::new(cmd)));
2253        }
2254        None => {
2255            return Err(NormalizeStatementError::NoStatement);
2256        }
2257    };
2258
2259    if parser.next()?.is_some() {
2260        return Err(NormalizeStatementError::Multiple);
2261    }
2262
2263    Ok(Cmd::Stmt(stmt).to_string())
2264}
2265
2266#[derive(Debug, thiserror::Error)]
2267pub enum PackError {
2268    #[error("abort")]
2269    Abort,
2270}
2271
2272pub fn pack_columns(args: &[SqliteValue]) -> Result<Vec<u8>, PackError> {
2273    let mut buf = vec![];
2274    /*
2275     * Format:
2276     * [num_columns:u8,...[(type(0-3),num_bytes?(3-7)):u8, length?:i32, ...bytes:u8[]]]
2277     *
2278     * The byte used for column type also encodes the number of bytes used for the integer.
2279     * e.g.: (type(0-3),num_bytes?(3-7)):u8
2280     * first 3 bits are type
2281     * last 5 encode how long the following integer, if there is a following integer, is. 1, 2, 3, ... 8 bytes.
2282     *
2283     * Not packing an integer into the minimal number of bytes required is rather wasteful.
2284     * E.g., the number `0` would take 8 bytes rather than 1 byte.
2285     */
2286    let len_result: Result<u8, _> = args.len().try_into();
2287    if let Ok(len) = len_result {
2288        buf.put_u8(len);
2289        for value in args {
2290            match value {
2291                SqliteValue::Null => {
2292                    buf.put_u8(ColumnType::Null as u8);
2293                }
2294                SqliteValue::Integer(val) => {
2295                    let num_bytes_for_int = num_bytes_needed_i64(*val);
2296                    let type_byte = num_bytes_for_int << 3 | (ColumnType::Integer as u8);
2297                    buf.put_u8(type_byte);
2298                    buf.put_int(*val, num_bytes_for_int as usize);
2299                }
2300                SqliteValue::Real(v) => {
2301                    buf.put_u8(ColumnType::Float as u8);
2302                    buf.put_f64(v.0);
2303                }
2304                SqliteValue::Text(value) => {
2305                    let len = value.len() as i32;
2306                    let num_bytes_for_len = num_bytes_needed_i32(len);
2307                    let type_byte = num_bytes_for_len << 3 | (ColumnType::Text as u8);
2308                    buf.put_u8(type_byte);
2309                    buf.put_int(len as i64, num_bytes_for_len as usize);
2310                    buf.put_slice(value.as_bytes());
2311                }
2312                SqliteValue::Blob(value) => {
2313                    let len = value.len() as i32;
2314                    let num_bytes_for_len = num_bytes_needed_i32(len);
2315                    let type_byte = num_bytes_for_len << 3 | (ColumnType::Blob as u8);
2316                    buf.put_u8(type_byte);
2317                    buf.put_int(len as i64, num_bytes_for_len as usize);
2318                    buf.put_slice(value);
2319                }
2320            }
2321        }
2322        Ok(buf)
2323    } else {
2324        Err(PackError::Abort)
2325    }
2326}
2327
2328fn num_bytes_needed_i64(val: i64) -> u8 {
2329    if val & 0xFF00000000000000u64 as i64 != 0 {
2330        8
2331    } else if val & 0x00FF000000000000 != 0 {
2332        7
2333    } else if val & 0x0000FF0000000000 != 0 {
2334        6
2335    } else if val & 0x000000FF00000000 != 0 {
2336        5
2337    } else {
2338        num_bytes_needed_i32(val as i32)
2339    }
2340}
2341
2342fn num_bytes_needed_i32(val: i32) -> u8 {
2343    if val & 0xFF000000u32 as i32 != 0 {
2344        4
2345    } else if val & 0x00FF0000 != 0 {
2346        3
2347    } else if val & 0x0000FF00 != 0 {
2348        2
2349    } else if val * 0x000000FF != 0 {
2350        1
2351    } else {
2352        0
2353    }
2354}
2355
2356#[derive(Debug, thiserror::Error)]
2357pub enum UnpackError {
2358    #[error("abort")]
2359    Abort,
2360    #[error("misuse")]
2361    Misuse,
2362}
2363
2364pub fn unpack_columns(mut buf: &[u8]) -> Result<Vec<SqliteValueRef<'_>>, UnpackError> {
2365    let mut ret = vec![];
2366    let num_columns = buf.get_u8();
2367
2368    for _i in 0..num_columns {
2369        if !buf.has_remaining() {
2370            return Err(UnpackError::Abort);
2371        }
2372        let column_type_and_maybe_intlen = buf.get_u8();
2373        let column_type = ColumnType::from_u8(column_type_and_maybe_intlen & 0x07);
2374        let intlen = (column_type_and_maybe_intlen >> 3) as usize;
2375
2376        match column_type {
2377            Some(ColumnType::Blob) => {
2378                if buf.remaining() < intlen {
2379                    return Err(UnpackError::Abort);
2380                }
2381                let len = buf.get_uint(intlen) as usize;
2382                if buf.remaining() < len {
2383                    return Err(UnpackError::Abort);
2384                }
2385                ret.push(SqliteValueRef(ValueRef::Blob(&buf[0..len])));
2386                buf.advance(len);
2387            }
2388            Some(ColumnType::Float) => {
2389                if buf.remaining() < 8 {
2390                    return Err(UnpackError::Abort);
2391                }
2392                ret.push(SqliteValueRef(ValueRef::Real(buf.get_f64())));
2393            }
2394            Some(ColumnType::Integer) => {
2395                if buf.remaining() < intlen {
2396                    return Err(UnpackError::Abort);
2397                }
2398                let unsigned = buf.get_uint(intlen);
2399                ret.push(SqliteValueRef(ValueRef::Integer(unsigned as i64)));
2400            }
2401            Some(ColumnType::Null) => {
2402                ret.push(SqliteValueRef(ValueRef::Null));
2403            }
2404            Some(ColumnType::Text) => {
2405                if buf.remaining() < intlen {
2406                    return Err(UnpackError::Abort);
2407                }
2408                let len = buf.get_uint(intlen) as usize;
2409                if buf.remaining() < len {
2410                    return Err(UnpackError::Abort);
2411                }
2412                ret.push(SqliteValueRef(ValueRef::Text(&buf[0..len])));
2413                buf.advance(len);
2414            }
2415            None => return Err(UnpackError::Misuse),
2416        }
2417    }
2418
2419    Ok(ret)
2420}
2421
2422#[cfg(test)]
2423mod tests {
2424    use std::net::Ipv4Addr;
2425
2426    use crate::spawn::wait_for_all_pending_handles;
2427    use camino::Utf8PathBuf;
2428    use rusqlite::{Connection, params};
2429    use tokio::sync::Semaphore;
2430
2431    use crate::{
2432        actor::ActorId,
2433        agent::migrate,
2434        base::CrsqlDbVersion,
2435        change::row_to_change,
2436        schema::{apply_schema, parse_sql},
2437        sqlite::{CrConn, rusqlite_to_crsqlite, setup_conn},
2438    };
2439    use klukai_tests::tempdir::TempDir;
2440
2441    use super::*;
2442
2443    #[test]
2444    fn test_pack_unpack() {
2445        let neg_one: i64 = -1;
2446        let big_neg: i64 = -2500000;
2447        let i8_max: i64 = i8::MAX as i64;
2448        let i16_max: i64 = i16::MAX as i64;
2449        let columns = vec![
2450            vec![1_i64.into(), SqliteValue::Null],
2451            vec![
2452                neg_one.into(),
2453                "abcdefghijklmnopqrstuvwxyz1234567890?".into(),
2454            ],
2455            vec![i64::MIN.into(), 1.0.into()],
2456            vec![i64::MAX.into(), vec![1, 2, 3].into()],
2457            vec![big_neg.into(), f64::MAX.into()],
2458            vec![10156800_i64.into(), f64::MIN.into()],
2459            vec![10000000_i64.into(), "".into()],
2460            vec![i8_max.into(), "a".into()],
2461            vec![i16_max.into(), vec![0].into()],
2462        ];
2463        for cols in columns.clone() {
2464            let packed = pack_columns(&cols).unwrap();
2465            let unpacked: Vec<_> = unpack_columns(&packed)
2466                .unwrap()
2467                .into_iter()
2468                .map(|v| v.to_owned())
2469                .collect();
2470            assert_eq!(cols, unpacked);
2471        }
2472
2473        let conn = Connection::open_in_memory().unwrap();
2474        let cr_conn = rusqlite_to_crsqlite(conn).unwrap();
2475        cr_conn
2476            .execute_batch("CREATE TABLE foo (pk INTEGER NOT NULL PRIMARY KEY, col_1);")
2477            .unwrap();
2478
2479        for cols in columns {
2480            cr_conn
2481                .execute(
2482                    "INSERT INTO foo (pk, col_1) VALUES (?, ?);",
2483                    params![cols[0], cols[1]],
2484                )
2485                .unwrap();
2486            let packed = cr_conn
2487                .query_row(
2488                    "SELECT crsql_pack_columns(pk, col_1) FROM foo where pk = ?",
2489                    params![cols[0]],
2490                    |row| row.get::<_, Vec<u8>>(0),
2491                )
2492                .unwrap();
2493            let unpacked: Vec<_> = unpack_columns(&packed)
2494                .unwrap()
2495                .into_iter()
2496                .map(|v| v.to_owned())
2497                .collect();
2498            assert_eq!(cols, unpacked);
2499        }
2500    }
2501
2502    #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
2503    async fn test_matcher() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
2504        _ = tracing_subscriber::fmt::try_init();
2505        let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple();
2506        let schema_sql = "CREATE TABLE sw (pk TEXT NOT NULL PRIMARY KEY, sandwich TEXT);";
2507        let mut schema = parse_sql(schema_sql)?;
2508
2509        let sql = "SELECT sandwich FROM sw WHERE pk=\"mad\"";
2510
2511        let subs = SubsManager::default();
2512
2513        let tmpdir = TempDir::new(tempfile::tempdir()?);
2514        let db_path = tmpdir.path().join("test.db");
2515        println!("db_path: {}", db_path.display());
2516        let subscriptions_path: Utf8PathBuf =
2517            tmpdir.path().join("subs").display().to_string().into();
2518
2519        let pool = SplitPool::create(db_path, Arc::new(Semaphore::new(1))).await?;
2520        let clock = Arc::new(uhlc::HLC::default());
2521
2522        {
2523            let mut conn = pool.write_priority().await?;
2524            setup_conn(&conn)?;
2525            migrate(clock, &mut conn)?;
2526            let tx = conn.transaction()?;
2527            apply_schema(&tx, &Schema::default(), &mut schema)?;
2528            tx.commit()?;
2529        }
2530
2531        {
2532            let (handle, maybe_created) = subs.get_or_insert(
2533                sql,
2534                subscriptions_path.as_path(),
2535                &schema,
2536                &pool,
2537                tripwire.clone(),
2538            )?;
2539
2540            assert!(maybe_created.is_some());
2541
2542            handle.cleanup().await;
2543            subs.remove(&handle.id());
2544        }
2545
2546        tripwire_tx.send(()).await.ok();
2547        tripwire_worker.await;
2548        wait_for_all_pending_handles().await;
2549
2550        Ok(())
2551    }
2552
2553    #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
2554    async fn test_diff() {
2555        _ = tracing_subscriber::fmt::try_init();
2556        let sql = "SELECT json_object(
2557            'targets', json_array(cs.address||':'||cs.port),
2558            'labels',  json_object(
2559              '__metrics_path__', JSON_EXTRACT(cs.meta, '$.path'),
2560              'app',            cs.app_name,
2561              'vm_account_id',  cs.organization_id,
2562              'instance',       cs.instance_id
2563            )
2564          )
2565          FROM consul_services cs
2566            LEFT JOIN machines m                   ON m.id = cs.instance_id
2567            LEFT JOIN machine_versions mv          ON m.id = mv.machine_id  AND m.machine_version_id = mv.id
2568            LEFT JOIN machine_version_statuses mvs ON m.id = mvs.machine_id AND m.machine_version_id = mvs.id
2569          WHERE cs.node = 'test-hostname'
2570            AND (mvs.status IS NULL OR mvs.status = 'started')
2571            AND cs.name == 'app-prometheus'";
2572
2573        let schema_sql = "
2574          CREATE TABLE consul_services (
2575              node TEXT NOT NULL,
2576              id TEXT NOT NULL,
2577              name TEXT NOT NULL DEFAULT '',
2578              tags TEXT NOT NULL DEFAULT '[]',
2579              meta TEXT NOT NULL DEFAULT '{}',
2580              port INTEGER NOT NULL DEFAULT 0,
2581              address TEXT NOT NULL DEFAULT '',
2582              updated_at INTEGER NOT NULL DEFAULT 0,
2583              app_id INTEGER AS (CAST(JSON_EXTRACT(meta, '$.app_id') AS INTEGER)), network_id INTEGER AS (
2584                  CAST(JSON_EXTRACT(meta, '$.network_id') AS INTEGER)
2585              ), app_name TEXT AS (JSON_EXTRACT(meta, '$.app_name')), instance_id TEXT AS (
2586                  COALESCE(
2587                      JSON_EXTRACT(meta, '$.machine_id'),
2588                      SUBSTR(JSON_EXTRACT(meta, '$.alloc_id'), 1, 8),
2589                      CASE
2590                          WHEN INSTR(id, '_nomad-task-') = 1 THEN SUBSTR(id, 13, 8)
2591                          ELSE NULL
2592                      END
2593                  )
2594              ), organization_id INTEGER AS (
2595                  CAST(
2596                      JSON_EXTRACT(meta, '$.organization_id') AS INTEGER
2597                  )
2598              ), protocol TEXT
2599          AS (JSON_EXTRACT(meta, '$.protocol')),
2600              PRIMARY KEY (node, id)
2601          );
2602
2603          CREATE TABLE machines (
2604              id TEXT NOT NULL PRIMARY KEY,
2605              node TEXT NOT NULL DEFAULT '',
2606              name TEXT NOT NULL DEFAULT '',
2607              machine_version_id TEXT NOT NULL DEFAULT '',
2608              app_id INTEGER NOT NULL DEFAULT 0,
2609              organization_id INTEGER NOT NULL DEFAULT 0,
2610              network_id INTEGER NOT NULL DEFAULT 0,
2611              updated_at INTEGER NOT NULL DEFAULT 0
2612          );
2613
2614          CREATE TABLE machine_versions (
2615              machine_id TEXT NOT NULL,
2616              id TEXT NOT NULL DEFAULT '',
2617              config TEXT NOT NULL DEFAULT '{}',
2618              updated_at INTEGER NOT NULL DEFAULT 0,
2619              PRIMARY KEY (machine_id, id)
2620          );
2621
2622          CREATE TABLE machine_version_statuses (
2623              machine_id TEXT NOT NULL,
2624              id TEXT NOT NULL,
2625              status TEXT NOT NULL DEFAULT '',
2626              updated_at INTEGER NOT NULL DEFAULT 0,
2627              PRIMARY KEY (machine_id, id)
2628          );
2629          ";
2630
2631        let mut schema = parse_sql(schema_sql).unwrap();
2632        let subs = SubsManager::default();
2633
2634        let tmpdir = TempDir::new(tempfile::tempdir().unwrap());
2635        let db_path = tmpdir.path().join("test.db");
2636        let subscriptions_path: Utf8PathBuf =
2637            tmpdir.path().join("subs").display().to_string().into();
2638
2639        let pool = SplitPool::create(&db_path, Arc::new(Semaphore::new(1)))
2640            .await
2641            .unwrap();
2642        let mut conn = pool.write_priority().await.unwrap();
2643        let clock = Arc::new(uhlc::HLC::default());
2644        {
2645            setup_conn(&conn).unwrap();
2646            migrate(clock.clone(), &mut conn).unwrap();
2647            let tx = conn.transaction().unwrap();
2648            apply_schema(&tx, &Schema::default(), &mut schema).unwrap();
2649            tx.commit().unwrap();
2650        }
2651
2652        // let's seed some data in there
2653        {
2654            let tx = conn.transaction().unwrap();
2655            tx.execute_batch(r#"
2656                INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', 'service-1', 'app-prometheus', '127.0.0.1', 1, '{"path": "/1", "machine_id": "m-1"}');
2657
2658                INSERT INTO machines (id, machine_version_id) VALUES ('m-1', 'mv-1');
2659
2660                INSERT INTO machine_versions (machine_id, id) VALUES ('m-1', 'mv-1');
2661
2662                INSERT INTO machine_version_statuses (machine_id, id, status) VALUES ('m-1', 'mv-1', 'started');
2663
2664                INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', 'service-2', 'not-app-prometheus', '127.0.0.1', 1, '{"path": "/1", "machine_id": "m-2"}');
2665
2666                INSERT INTO machines (id, machine_version_id) VALUES ('m-2', 'mv-2');
2667
2668                INSERT INTO machine_versions (machine_id, id) VALUES ('m-2', 'mv-2');
2669
2670                INSERT INTO machine_version_statuses (machine_id, id, status) VALUES ('m-2', 'mv-2', 'started');
2671                    "#).unwrap();
2672            tx.commit().unwrap();
2673        }
2674
2675        {
2676            let mut conn2 = CrConn::init(
2677                rusqlite::Connection::open(tmpdir.path().join("test2.db"))
2678                    .expect("could not open conn"),
2679            )
2680            .expect("could not init crsql");
2681
2682            setup_conn(&conn2).unwrap();
2683            migrate(clock.clone(), &mut conn2).unwrap();
2684
2685            {
2686                let tx = conn2.transaction().unwrap();
2687                apply_schema(&tx, &Schema::default(), &mut schema).unwrap();
2688                tx.commit().unwrap();
2689            }
2690
2691            let changes = {
2692                let mut prepped = conn.prepare_cached(r#"SELECT "table", pk, cid, val, col_version, db_version, seq, site_id, cl FROM crsql_changes WHERE db_version = ? ORDER BY seq ASC"#).unwrap();
2693                let rows = prepped.query_map([1], row_to_change).unwrap();
2694
2695                let mut changes = vec![];
2696
2697                for row in rows {
2698                    changes.push(row.unwrap());
2699                }
2700                changes
2701            };
2702
2703            let tx = conn2.transaction().unwrap();
2704
2705            for change in changes {
2706                debug!("change: {change:?}");
2707                tx.prepare_cached(
2708                    r#"
2709                    INSERT INTO crsql_changes
2710                        ("table", pk, cid, val, col_version, db_version, site_id, cl, seq)
2711                    VALUES
2712                        (?,       ?,  ?,   ?,   ?,           ?,          ?,       ?,  ?)
2713                "#,
2714                )
2715                .unwrap()
2716                .execute(params![
2717                    change.table.as_str(),
2718                    change.pk,
2719                    change.cid.as_str(),
2720                    &change.val,
2721                    change.col_version,
2722                    change.db_version,
2723                    &change.site_id,
2724                    change.cl,
2725                    change.seq
2726                ])
2727                .unwrap();
2728            }
2729        }
2730
2731        let matcher_conn =
2732            CrConn::init(rusqlite::Connection::open(&db_path).expect("could not open conn"))
2733                .expect("could not init crconn");
2734
2735        setup_conn(&matcher_conn).unwrap();
2736
2737        let mut last_change_id = None;
2738
2739        let id = {
2740            let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple();
2741
2742            let (matcher, maybe_created) = subs
2743                .get_or_insert(
2744                    sql,
2745                    subscriptions_path.as_path(),
2746                    &schema,
2747                    &pool,
2748                    tripwire.clone(),
2749                )
2750                .unwrap();
2751
2752            let mut rx = maybe_created.unwrap().evt_rx;
2753
2754            println!("matcher created w/ id: {}", matcher.id().as_simple());
2755            println!("parsed: {:?}", matcher.inner.parsed);
2756
2757            assert!(matches!(rx.recv().await.unwrap(), QueryEvent::Columns(_)));
2758
2759            let cells = vec![SqliteValue::Text("{\"targets\":[\"127.0.0.1:1\"],\"labels\":{\"__metrics_path__\":\"/1\",\"app\":null,\"vm_account_id\":null,\"instance\":\"m-1\"}}".into())];
2760
2761            assert_eq!(rx.recv().await.unwrap(), QueryEvent::Row(RowId(1), cells));
2762            println!("received a row");
2763            assert!(matches!(
2764                rx.recv().await.unwrap(),
2765                QueryEvent::EndOfQuery { .. }
2766            ));
2767            println!("received end of query");
2768
2769            // insert the second row
2770            {
2771                let tx = conn.transaction().unwrap();
2772                tx.execute_batch(r#"
2773                INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', 'service-3', 'app-prometheus', '127.0.0.1', 1, '{"path": "/1", "machine_id": "m-3"}');
2774
2775                INSERT INTO machines (id, machine_version_id) VALUES ('m-3', 'mv-3');
2776
2777                INSERT INTO machine_versions (machine_id, id) VALUES ('m-3', 'mv-3');
2778
2779                INSERT INTO machine_version_statuses (machine_id, id, status) VALUES ('m-3', 'mv-3', 'started');
2780            "#).unwrap();
2781                tx.commit().unwrap();
2782            }
2783
2784            println!("processing change...");
2785            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(2)).unwrap();
2786            println!("processed changes");
2787
2788            let cells = vec![SqliteValue::Text("{\"targets\":[\"127.0.0.1:1\"],\"labels\":{\"__metrics_path__\":\"/1\",\"app\":null,\"vm_account_id\":null,\"instance\":\"m-3\"}}".into())];
2789
2790            assert_eq!(
2791                rx.recv().await.unwrap(),
2792                QueryEvent::Change(ChangeType::Insert, RowId(2), cells, ChangeId(1))
2793            );
2794
2795            println!("received change");
2796
2797            // delete the first row
2798            {
2799                let tx = conn.transaction().unwrap();
2800                tx.execute_batch(r#"
2801                        DELETE FROM consul_services where node = 'test-hostname' AND id = 'service-1';
2802                    "#).unwrap();
2803                tx.commit().unwrap();
2804            }
2805
2806            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(3)).unwrap();
2807
2808            let cells = vec![SqliteValue::Text("{\"targets\":[\"127.0.0.1:1\"],\"labels\":{\"__metrics_path__\":\"/1\",\"app\":null,\"vm_account_id\":null,\"instance\":\"m-1\"}}".into())];
2809
2810            println!("waiting for a change (A)");
2811
2812            assert_eq!(
2813                rx.recv().await.unwrap(),
2814                QueryEvent::Change(ChangeType::Delete, RowId(1), cells, ChangeId(2))
2815            );
2816
2817            println!("got change (A)");
2818
2819            // update the second row
2820            {
2821                let tx = conn.transaction().unwrap();
2822                let n = tx.execute(r#"
2823                        UPDATE consul_services SET address = '127.0.0.2' WHERE node = 'test-hostname' AND id = 'service-3';
2824                    "#, ()).unwrap();
2825                assert_eq!(n, 1);
2826                tx.commit().unwrap();
2827            }
2828
2829            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(4)).unwrap();
2830
2831            let cells = vec![SqliteValue::Text("{\"targets\":[\"127.0.0.2:1\"],\"labels\":{\"__metrics_path__\":\"/1\",\"app\":null,\"vm_account_id\":null,\"instance\":\"m-3\"}}".into())];
2832
2833            println!("waiting for a change (B)");
2834
2835            assert_eq!(
2836                rx.recv().await.unwrap(),
2837                QueryEvent::Change(ChangeType::Update, RowId(2), cells, ChangeId(3))
2838            );
2839
2840            println!("got change (B)");
2841
2842            // process the same and check that it doesn't produce a change again!
2843            // db_v_tx.send(db_version).unwrap();
2844
2845            assert_eq!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty));
2846
2847            // lots of operations
2848
2849            let range = 4u32..1000u32;
2850
2851            {
2852                let tx = conn.transaction().unwrap();
2853
2854                for n in range.clone() {
2855                    let svc_id = format!("service-{n}");
2856                    let ip = Ipv4Addr::from(n).to_string();
2857                    let port = n;
2858                    let machine_id = format!("m-{n}");
2859                    let mv = format!("mv-{n}");
2860                    let meta =
2861                        serde_json::json!({"path": format!("/path-{n}"), "machine_id": machine_id});
2862                    tx.execute("
2863                                    INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', ?, 'app-prometheus', ?, ?, ?);
2864                                    ", params![svc_id, ip, port, meta]).unwrap();
2865                    tx.execute(
2866                        "
2867                        INSERT INTO machines (id, machine_version_id) VALUES (?, ?);
2868                        ",
2869                        params![machine_id, mv],
2870                    )
2871                    .unwrap();
2872
2873                    tx.execute(
2874                        "
2875                        INSERT INTO machine_versions (machine_id, id) VALUES (?, ?);
2876                        ",
2877                        params![machine_id, mv],
2878                    )
2879                    .unwrap();
2880
2881                    tx.execute("
2882                        INSERT INTO machine_version_statuses (machine_id, id, status) VALUES (?, ?, 'started');", params![machine_id, mv]).unwrap();
2883                }
2884
2885                tx.commit().unwrap();
2886            }
2887
2888            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(5)).unwrap();
2889
2890            let start = Instant::now();
2891            for _ in range {
2892                if let QueryEvent::Change(change_type, _, _, change_id) = rx.recv().await.unwrap() {
2893                    assert_eq!(change_type, ChangeType::Insert);
2894                    last_change_id = Some(change_id);
2895                }
2896            }
2897
2898            assert_eq!(last_change_id, Some(ChangeId(999)));
2899
2900            println!("took: {:?}", start.elapsed());
2901            {
2902                let conn = rusqlite::Connection::open(
2903                    subscriptions_path
2904                        .join(matcher.id().as_simple().to_string())
2905                        .join("sub.sqlite"),
2906                )
2907                .unwrap();
2908                let mut prepped = conn.prepare("SELECT * FROM changes").unwrap();
2909                let cols = prepped
2910                    .column_names()
2911                    .iter()
2912                    .map(|s| s.to_string())
2913                    .collect::<Vec<_>>();
2914                let mut rows = prepped.query([]).unwrap();
2915                loop {
2916                    let row = rows.next().unwrap();
2917                    let row = match row {
2918                        None => break,
2919                        Some(row) => row,
2920                    };
2921
2922                    let s = cols
2923                        .iter()
2924                        .enumerate()
2925                        .map(|(idx, col)| {
2926                            format!("{col} = {}", row.get::<_, SqliteValue>(idx).unwrap())
2927                        })
2928                        .collect::<Vec<_>>();
2929                    println!("{}", s.join(", "));
2930                }
2931            }
2932
2933            subs.remove(&matcher.id());
2934
2935            // trip the wire so subs loop exits cleanly
2936            tripwire_tx.send(()).await.ok();
2937            tripwire_worker.await;
2938            matcher.id()
2939        };
2940
2941        tokio::time::sleep(Duration::from_secs(1)).await;
2942
2943        let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple();
2944        // restore subscription
2945        let matcher_id = {
2946            let (matcher, created) = subs
2947                .restore(id, &subscriptions_path, &schema, &pool, tripwire.clone())
2948                .unwrap();
2949            let mut rx = created.evt_rx;
2950
2951            let matcher_id = matcher.id().as_simple().to_string();
2952            println!("matcher restored w/ id: {matcher_id}");
2953
2954            let (catch_up_tx, mut catch_up_rx) = mpsc::channel(1);
2955
2956            tokio::spawn({
2957                let matcher = matcher.clone();
2958                async move {
2959                    let mut conn = matcher.pool().get().await.unwrap();
2960                    block_in_place(|| {
2961                        let conn_tx = conn.transaction()?;
2962
2963                        println!(
2964                            "max change id: {}",
2965                            matcher.max_change_id(&conn_tx).unwrap()
2966                        );
2967                        matcher.all_rows(&conn_tx, catch_up_tx)?;
2968                        Ok::<_, MatcherError>(())
2969                    })
2970                }
2971            });
2972
2973            assert!(matches!(
2974                catch_up_rx.recv().await.unwrap(),
2975                QueryEvent::Columns(_)
2976            ));
2977
2978            let mut rows_count = 0;
2979            let mut eoq_change_id = None;
2980
2981            while eoq_change_id.is_none() {
2982                match catch_up_rx.recv().await.unwrap() {
2983                    QueryEvent::EndOfQuery { change_id, .. } => eoq_change_id = Some(change_id),
2984                    QueryEvent::Row(_, _) => rows_count += 1,
2985                    QueryEvent::Change(_, _, _, _) => {
2986                        panic!("received a change intertwined w/ rows");
2987                    }
2988                    _ => {}
2989                }
2990            }
2991
2992            assert_eq!(rows_count, 997);
2993            assert_eq!(eoq_change_id, Some(Some(ChangeId(999))));
2994
2995            {
2996                let tx = conn.transaction().unwrap();
2997                tx.execute_batch(
2998                    r#"
2999                        DELETE FROM consul_services where node = 'test-hostname' AND id = 'service-3';
3000                    "#,
3001                )
3002                .unwrap();
3003                tx.commit().unwrap();
3004            }
3005            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(6)).unwrap();
3006            println!("waiting for a change (A)");
3007
3008            let cells = vec![SqliteValue::Text("{\"targets\":[\"127.0.0.2:1\"],\"labels\":{\"__metrics_path__\":\"/1\",\"app\":null,\"vm_account_id\":null,\"instance\":\"m-3\"}}".into())];
3009
3010            assert_eq!(
3011                rx.recv().await.unwrap(),
3012                QueryEvent::Change(ChangeType::Delete, RowId(2), cells, ChangeId(1000))
3013            );
3014
3015            assert!(rx.try_recv().is_err());
3016
3017            {
3018                let tx = conn.transaction().unwrap();
3019                for n in 1001..=1005 {
3020                    tx.execute_batch(format!(r#"
3021                        INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', 'service-{n}', 'app-prometheus', '127.0.0.{n}', 1, '{{"path": "/1", "machine_id": "m-{n}"}}');
3022
3023                        INSERT INTO machines (id, machine_version_id) VALUES ('m-{n}', 'mv-{n}');
3024
3025                        INSERT INTO machine_versions (machine_id, id) VALUES ('m-{n}', 'mv-{n}');
3026
3027                        INSERT INTO machine_version_statuses (machine_id, id, status) VALUES ('m-{n}', 'mv-{n}', 'started');
3028                    "#).as_str()).unwrap();
3029                }
3030                tx.commit().unwrap();
3031            }
3032            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(7)).unwrap();
3033
3034            // initiate shutdown
3035            tripwire_tx.send(()).await.ok();
3036            tripwire_worker.await;
3037
3038            {
3039                let tx = conn.transaction().unwrap();
3040                for n in 1006..=1010 {
3041                    tx.execute_batch(format!(r#"
3042                        INSERT INTO consul_services (node, id, name, address, port, meta) VALUES ('test-hostname', 'service-{n}', 'app-prometheus', '127.0.0.{n}', 1, '{{"path": "/1", "machine_id": "m-{n}"}}');
3043
3044                        INSERT INTO machines (id, machine_version_id) VALUES ('m-{n}', 'mv-{n}');
3045
3046                        INSERT INTO machine_versions (machine_id, id) VALUES ('m-{n}', 'mv-{n}');
3047
3048                        INSERT INTO machine_version_statuses (machine_id, id, status) VALUES ('m-{n}', 'mv-{n}', 'started');
3049                    "#).as_str()).unwrap();
3050                }
3051                tx.commit().unwrap();
3052            }
3053            filter_changes_from_db(&matcher, &conn, None, CrsqlDbVersion(8)).unwrap();
3054
3055            matcher_id
3056        };
3057
3058        subs.drop_handles().await;
3059        wait_for_all_pending_handles().await;
3060        info!("tripwire worker finished");
3061
3062        // check that changes sent after tripping the wire were persisted to db
3063
3064        {
3065            let conn = rusqlite::Connection::open(
3066                subscriptions_path
3067                    .join(matcher_id.clone())
3068                    .join("sub.sqlite"),
3069            )
3070            .unwrap();
3071
3072            let mut prepped = conn.prepare("SELECT max(id) FROM changes").unwrap();
3073            let max_id = prepped.query_row([], |row| row.get::<_, i64>(0)).unwrap();
3074            assert_eq!(max_id, 1010);
3075
3076            for n in 1001..=1010 {
3077                let mut prepped = conn
3078                    .prepare(
3079                        format!("SELECT 1 from query where __corro_pk_cs_id = 'service-{n}'")
3080                            .as_str(),
3081                    )
3082                    .unwrap();
3083                let max_db_version = prepped.query_row([], |row| row.get::<_, i64>(0)).unwrap();
3084                assert_eq!(max_db_version, 1);
3085            }
3086        }
3087
3088        // a restore should start ok if we shutdown properly
3089        {
3090            let res = subs.restore(id, &subscriptions_path, &schema, &pool, tripwire.clone());
3091            assert!(res.is_ok());
3092        }
3093
3094        // restore should fail if we don't shutdown properly
3095        {
3096            let conn =
3097                rusqlite::Connection::open(subscriptions_path.join(matcher_id).join("sub.sqlite"))
3098                    .unwrap();
3099            conn.execute(
3100                "INSERT OR REPLACE INTO meta (key, value) VALUES ('state', 'running')",
3101                (),
3102            )
3103            .unwrap();
3104
3105            let res = subs.restore(id, &subscriptions_path, &schema, &pool, tripwire.clone());
3106            assert!(res.is_err());
3107        }
3108    }
3109
3110    fn filter_changes_from_db(
3111        matcher: &MatcherHandle,
3112        state_conn: &Connection,
3113        actor_id: Option<ActorId>,
3114        db_version: CrsqlDbVersion,
3115    ) -> rusqlite::Result<()> {
3116        let mut candidates = MatchCandidates::new();
3117
3118        let mut prepped = state_conn.prepare_cached(r#"SELECT "table", pk, cid, val, col_version, db_version, seq, site_id, cl  FROM crsql_changes WHERE site_id = COALESCE(?, crsql_site_id()) AND db_version = ? ORDER BY seq ASC"#)?;
3119        let rows = prepped
3120            .query_map(params![actor_id, db_version], row_to_change)
3121            .unwrap();
3122
3123        for row in rows {
3124            let change = row?;
3125            matcher.filter_matchable_change(&mut candidates, (&change).into());
3126        }
3127
3128        if let Err(e) = matcher.inner.changes_tx.try_send(candidates) {
3129            error!(sub_id = %matcher.inner.id, "could not send candidates to matcher: {e}");
3130        }
3131        Ok(())
3132    }
3133}