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
62pub 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 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 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 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 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 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 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 let (changes_tx, changes_rx) = mpsc::channel(20480);
764
765 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 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 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 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 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(); 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 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 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 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 }
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 tx.prepare_cached(
1444 &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 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 let tx = self.conn.transaction()?;
1492
1493 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 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 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_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 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 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 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 for table in tables {
1674 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 }
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 t => {
1789 warn!("ignoring {t:?}");
1790 None
1791 }
1792 },
1793 _ => {
1794 unreachable!()
1797 }
1798 };
1799 if let Some(ref joins) = from.joins {
1800 for join in joins.iter() {
1801 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 t => {
1824 warn!("ignoring JOIN's non-SelectTable::Table: {t:?}");
1825 continue;
1826 }
1827 };
1828 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 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 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 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 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 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 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 _ => {}
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 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 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 {
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 {
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 {
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 {
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 assert_eq!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty));
2846
2847 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 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 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 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 {
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 {
3090 let res = subs.restore(id, &subscriptions_path, &schema, &pool, tripwire.clone());
3091 assert!(res.is_ok());
3092 }
3093
3094 {
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}