1use std::collections::{HashMap, VecDeque};
2use std::num::NonZeroU32;
3use std::sync::Arc;
4use std::time::Duration;
5
6use chrono::{DateTime, Utc};
7use rusqlite::types::Value;
8use tokio::sync::Mutex;
9use tokio::sync::Notify;
10use tokio::sync::mpsc;
11
12use eventuary_core::io::cursor::{CursorOrder, JsonCursorCodec};
13use eventuary_core::io::filter::{EventFilter, NamespacePattern, TopicPattern};
14use eventuary_core::io::reader::{
15 CoordinatedAcker, CoordinatedCursor, CoordinatedReader, CoordinatedReaderConfig,
16 CoordinatedStream, CoordinatedSubscription, PartitionAcker, PartitionedCoordAdapter,
17 PartitionedCursor,
18};
19use eventuary_core::io::stream::SpawnedStream;
20use eventuary_core::io::{Acker, Cursor, Filter, Message, Reader};
21use eventuary_core::partition::{HasPartition, Partition, PartitionGroup, PartitionSelection};
22use eventuary_core::{
23 Error, PartitionableSubscription, Result, SerializedEvent, SerializedPayload, StartFrom,
24 StartableSubscription, StopAt,
25};
26
27use crate::coordinator::SqlitePartitionCoordinator;
28
29use crate::database::SqliteConn;
30use crate::event_log::{SqliteEventLogSchema, SqliteEventLogSchemaConfig};
31use crate::relation::SqliteRelationName;
32
33#[derive(
34 Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, serde::Serialize, serde::Deserialize,
35)]
36pub struct SqliteCursor {
37 pub sequence: i64,
38 pub partition: Partition,
39}
40
41impl SqliteCursor {
42 pub fn new(sequence: i64, partition: Partition) -> Self {
43 Self {
44 sequence,
45 partition,
46 }
47 }
48
49 pub fn sequence(&self) -> i64 {
50 self.sequence
51 }
52
53 pub fn partition(&self) -> Partition {
54 self.partition
55 }
56}
57
58impl Cursor for SqliteCursor {
59 fn order_key(&self) -> CursorOrder {
60 CursorOrder::from_i64(self.sequence)
61 }
62}
63
64impl HasPartition for SqliteCursor {
65 fn partition(&self) -> Partition {
66 self.partition
67 }
68}
69
70impl SqliteCursor {
71 pub fn codec() -> Result<JsonCursorCodec<Self>> {
72 JsonCursorCodec::new("eventuary.sqlite.sqlite_cursor.v1")
73 }
74}
75
76#[derive(Debug, Clone)]
77pub struct SqliteSubscription {
78 pub start: StartFrom<SqliteCursor>,
79 pub stop_at: StopAt<SqliteCursor>,
80 pub filter: EventFilter,
81 pub partitions: PartitionSelection,
82 pub batch_size: Option<usize>,
83 pub limit: Option<usize>,
84}
85
86impl Default for SqliteSubscription {
87 fn default() -> Self {
88 Self {
89 start: StartFrom::Latest,
90 stop_at: StopAt::Never,
91 filter: EventFilter::default(),
92 batch_size: None,
93 limit: None,
94 partitions: PartitionSelection::default(),
95 }
96 }
97}
98
99impl StartableSubscription<SqliteCursor> for SqliteSubscription {
100 fn with_start(mut self, start: StartFrom<SqliteCursor>) -> Self {
101 self.start = start;
102 self
103 }
104}
105
106impl PartitionableSubscription<SqliteCursor> for SqliteSubscription {
107 fn with_partitions(mut self, group: PartitionGroup) -> Self {
114 self.partitions = PartitionSelection::Many(group);
115 self
116 }
117}
118
119#[derive(Debug, Clone)]
120pub struct SqliteReaderConfig {
121 pub events_relation: SqliteRelationName,
122 pub poll_interval: Duration,
123 pub default_batch_size: usize,
124}
125
126impl Default for SqliteReaderConfig {
127 fn default() -> Self {
128 Self {
129 events_relation: SqliteRelationName::new("events").expect("default events relation"),
130 poll_interval: Duration::from_millis(100),
131 default_batch_size: 100,
132 }
133 }
134}
135
136#[derive(Clone)]
137pub struct SqliteCursorAcker {
138 state: Arc<Mutex<CursorState>>,
139 notify: Arc<Notify>,
140 sequence: i64,
141}
142
143struct CursorState {
144 last_acked: i64,
145 pending_nack: bool,
146}
147
148impl Acker for SqliteCursorAcker {
149 async fn ack(&self) -> Result<()> {
150 let mut state = self.state.lock().await;
151 if self.sequence > state.last_acked {
152 state.last_acked = self.sequence;
153 }
154 state.pending_nack = false;
155 self.notify.notify_waiters();
156 Ok(())
157 }
158
159 async fn nack(&self) -> Result<()> {
160 let mut state = self.state.lock().await;
161 state.pending_nack = true;
162 self.notify.notify_waiters();
163 Ok(())
164 }
165}
166
167pub struct SqliteReader {
168 conn: SqliteConn,
169 config: SqliteReaderConfig,
170}
171
172impl Clone for SqliteReader {
173 fn clone(&self) -> Self {
174 Self {
175 conn: Arc::clone(&self.conn),
176 config: self.config.clone(),
177 }
178 }
179}
180
181impl SqliteReader {
182 pub fn connect(conn: SqliteConn, config: SqliteReaderConfig) -> Result<Self> {
183 Self::prepare_schema(&conn, &config)?;
184 Ok(Self::new(conn, config))
185 }
186
187 pub fn prepare_schema(conn: &SqliteConn, config: &SqliteReaderConfig) -> Result<()> {
188 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
189 SqliteEventLogSchema::prepare(
190 &guard,
191 &SqliteEventLogSchemaConfig {
192 events_relation: config.events_relation.clone(),
193 },
194 )
195 }
196
197 pub fn schema_sql(config: &SqliteReaderConfig) -> String {
198 SqliteEventLogSchema::schema_sql(&SqliteEventLogSchemaConfig {
199 events_relation: config.events_relation.clone(),
200 })
201 }
202
203 pub fn new(conn: SqliteConn, config: SqliteReaderConfig) -> Self {
204 Self { conn, config }
205 }
206}
207
208impl Reader for SqliteReader {
209 type Subscription = SqliteSubscription;
210 type Acker = SqliteCursorAcker;
211 type Cursor = SqliteCursor;
212 type Stream = SpawnedStream<SqliteCursorAcker, SqliteCursor>;
213
214 async fn read(&self, subscription: Self::Subscription) -> Result<Self::Stream> {
215 let conn = Arc::clone(&self.conn);
216 let events_relation = self.config.events_relation.render();
217 let poll_interval = self.config.poll_interval;
218 let batch_size = subscription
219 .batch_size
220 .unwrap_or(self.config.default_batch_size)
221 .clamp(1, 1000);
222 let filter = subscription.filter.clone();
223 let limit = subscription.limit;
224 let partitions = subscription.partitions.clone();
225 let (tx, rx) = mpsc::channel(64);
226
227 let (mut after_seq, lower_bound_ts) =
228 match resolve_initial_position(&conn, &events_relation, &subscription).await {
229 Ok(pos) => pos,
230 Err(e) => {
231 let _ = tx.send(Err(e)).await;
232 return Ok(SpawnedStream::from_receiver(rx));
233 }
234 };
235
236 let stop_seq = match resolve_stop_position(&conn, &events_relation, &subscription).await {
237 Ok(pos) => pos,
238 Err(e) => {
239 let _ = tx.send(Err(e)).await;
240 return Ok(SpawnedStream::from_receiver(rx));
241 }
242 };
243
244 let state = Arc::new(Mutex::new(CursorState {
245 last_acked: after_seq,
246 pending_nack: false,
247 }));
248 let notify = Arc::new(Notify::new());
249
250 let handle = tokio::spawn(async move {
251 let mut delivered = 0usize;
252 let mut buffer: VecDeque<(SerializedEvent, i64, Partition)> = VecDeque::new();
253 loop {
254 if buffer.is_empty() {
255 let fetched = match fetch_batch(
256 &conn,
257 FetchBatchParams {
258 events_relation: &events_relation,
259 after_seq,
260 stop_seq,
261 take: batch_size,
262 lower_bound_ts,
263 filter: &filter,
264 partitions: &partitions,
265 },
266 )
267 .await
268 {
269 Ok(b) => b,
270 Err(e) => {
271 let _ = tx.send(Err(e)).await;
272 return;
273 }
274 };
275 if fetched.is_empty() {
276 if stop_seq.is_some() {
277 return;
278 }
279 tokio::time::sleep(poll_interval).await;
280 continue;
281 }
282 buffer.extend(fetched);
283 }
284
285 while let Some((serialized, sequence, partition)) = buffer.front() {
286 let sequence = *sequence;
287 let partition = *partition;
288 let event = match serialized.to_event() {
289 Ok(e) => e,
290 Err(e) => {
291 let _ = tx
292 .send(Err(Error::Serialization(format!(
293 "decode event at sequence {sequence}: {e}"
294 ))))
295 .await;
296 return;
297 }
298 };
299 if !filter.matches(&event) {
300 buffer.pop_front();
301 after_seq = sequence;
302 continue;
303 }
304 if let Some(l) = limit
305 && delivered >= l
306 {
307 return;
308 }
309 let acker = SqliteCursorAcker {
310 state: Arc::clone(&state),
311 notify: Arc::clone(¬ify),
312 sequence,
313 };
314 let cursor = SqliteCursor {
315 sequence,
316 partition,
317 };
318 if tx
319 .send(Ok(Message::new(event, acker, cursor)))
320 .await
321 .is_err()
322 {
323 return;
324 }
325 delivered += 1;
326
327 loop {
328 {
329 let guard = state.lock().await;
330 if guard.last_acked >= sequence {
331 after_seq = sequence;
332 buffer.pop_front();
333 break;
334 }
335 if guard.pending_nack {
336 break;
337 }
338 if tx.is_closed() {
339 return;
340 }
341 }
342 notify.notified().await;
343 }
344 }
345 }
346 });
347
348 Ok(SpawnedStream::new(rx, handle))
349 }
350}
351
352async fn resolve_initial_position(
353 conn: &SqliteConn,
354 events_relation: &str,
355 subscription: &SqliteSubscription,
356) -> Result<(i64, Option<DateTime<Utc>>)> {
357 match subscription.start.clone() {
358 StartFrom::After(cursor) => Ok((cursor.sequence, None)),
359 StartFrom::Earliest => Ok((0, None)),
360 StartFrom::Latest => {
361 let conn = Arc::clone(conn);
362 let org = subscription
363 .filter
364 .organization
365 .as_ref()
366 .map(|o| o.as_str().to_owned());
367 let relation = events_relation.to_owned();
368 tokio::task::spawn_blocking(move || {
369 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
370 let seq: i64 = match org {
371 Some(o) => guard
372 .query_row(
373 &format!(
374 "SELECT COALESCE(MAX(sequence), 0) FROM {relation} WHERE organization = ?1"
375 ),
376 rusqlite::params![o],
377 |r| r.get(0),
378 )
379 .map_err(|e| Error::Store(e.to_string()))?,
380 None => guard
381 .query_row(
382 &format!("SELECT COALESCE(MAX(sequence), 0) FROM {relation}"),
383 [],
384 |r| r.get(0),
385 )
386 .map_err(|e| Error::Store(e.to_string()))?,
387 };
388 Ok::<(i64, Option<DateTime<Utc>>), Error>((seq, None))
389 })
390 .await
391 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
392 }
393 StartFrom::Timestamp(ts) => {
394 let conn = Arc::clone(conn);
395 let org = subscription
396 .filter
397 .organization
398 .as_ref()
399 .map(|o| o.as_str().to_owned());
400 let ts_str = ts.to_rfc3339();
401 let relation = events_relation.to_owned();
402 tokio::task::spawn_blocking(move || {
403 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
404 let seq: i64 = match org {
405 Some(o) => guard
406 .query_row(
407 &format!(
408 "SELECT COALESCE(MIN(sequence), 1) - 1 FROM {relation} \
409 WHERE organization = ?1 AND timestamp >= ?2"
410 ),
411 rusqlite::params![o, ts_str],
412 |r| r.get(0),
413 )
414 .map_err(|e| Error::Store(e.to_string()))?,
415 None => guard
416 .query_row(
417 &format!(
418 "SELECT COALESCE(MIN(sequence), 1) - 1 FROM {relation} \
419 WHERE timestamp >= ?1"
420 ),
421 rusqlite::params![ts_str],
422 |r| r.get(0),
423 )
424 .map_err(|e| Error::Store(e.to_string()))?,
425 };
426 Ok::<(i64, Option<DateTime<Utc>>), Error>((seq.max(0), Some(ts)))
427 })
428 .await
429 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
430 }
431 }
432}
433
434async fn resolve_stop_position(
435 conn: &SqliteConn,
436 events_relation: &str,
437 subscription: &SqliteSubscription,
438) -> Result<Option<i64>> {
439 match subscription.stop_at {
440 StopAt::Never => Ok(None),
441 StopAt::Cursor(cursor) => Ok(Some(cursor.sequence)),
442 StopAt::CurrentEnd => {
443 let conn = Arc::clone(conn);
444 let org = subscription
445 .filter
446 .organization
447 .as_ref()
448 .map(|o| o.as_str().to_owned());
449 let relation = events_relation.to_owned();
450 tokio::task::spawn_blocking(move || {
451 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
452 let seq: i64 = match org {
453 Some(o) => guard
454 .query_row(
455 &format!(
456 "SELECT COALESCE(MAX(sequence), 0) FROM {relation} WHERE organization = ?1"
457 ),
458 rusqlite::params![o],
459 |r| r.get(0),
460 )
461 .map_err(|e| Error::Store(e.to_string()))?,
462 None => guard
463 .query_row(
464 &format!("SELECT COALESCE(MAX(sequence), 0) FROM {relation}"),
465 [],
466 |r| r.get(0),
467 )
468 .map_err(|e| Error::Store(e.to_string()))?,
469 };
470 Ok::<Option<i64>, Error>(Some(seq))
471 })
472 .await
473 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
474 }
475 }
476}
477
478struct FetchBatchParams<'a> {
479 events_relation: &'a str,
480 after_seq: i64,
481 stop_seq: Option<i64>,
482 take: usize,
483 lower_bound_ts: Option<DateTime<Utc>>,
484 filter: &'a EventFilter,
485 partitions: &'a PartitionSelection,
486}
487
488async fn fetch_batch(
489 conn: &SqliteConn,
490 p: FetchBatchParams<'_>,
491) -> Result<Vec<(SerializedEvent, i64, Partition)>> {
492 let FetchBatchParams {
493 events_relation,
494 after_seq,
495 stop_seq,
496 take,
497 lower_bound_ts,
498 filter,
499 partitions,
500 } = p;
501 let conn = Arc::clone(conn);
502 let relation = events_relation.to_owned();
503 let org = filter.organization.as_ref().map(|o| o.as_str().to_owned());
504 let exact_topic: Option<String> = filter.topic.as_ref().map(|p| match p {
505 TopicPattern::Exact(t) => t.as_str().to_owned(),
506 });
507 let ns_prefix = filter.namespace.as_ref().and_then(|p| match p {
508 NamespacePattern::Prefix(ns) if !ns.is_root() => Some(ns.as_str().to_owned()),
509 _ => None,
510 });
511 let ts_str = lower_bound_ts.map(|t| t.to_rfc3339());
512 let partitions = partitions.clone();
513
514 tokio::task::spawn_blocking(move || {
515 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
516
517 let mut sql = format!(
518 "SELECT sequence, id, organization, namespace, topic, event_key, payload, content_type, metadata, \
519 timestamp, version, parent_id, correlation_id, causation_id, partition_id, partition_count \
520 FROM {relation} WHERE sequence > ?1"
521 );
522 let mut params: Vec<Value> = vec![Value::Integer(after_seq)];
523 let mut idx = 2usize;
524
525 if let Some(stop) = stop_seq {
526 sql.push_str(&format!(" AND sequence <= ?{idx}"));
527 params.push(Value::Integer(stop));
528 idx += 1;
529 }
530
531 if let Some(o) = &org {
532 sql.push_str(&format!(" AND organization = ?{idx}"));
533 params.push(Value::Text(o.clone()));
534 idx += 1;
535 }
536 if let Some(t) = &exact_topic {
537 sql.push_str(&format!(" AND topic = ?{idx}"));
538 params.push(Value::Text(t.clone()));
539 idx += 1;
540 }
541 if let Some(prefix) = &ns_prefix {
542 sql.push_str(&format!(
543 " AND (namespace = ?{idx} OR namespace LIKE ?{} || '/%')",
544 idx
545 ));
546 params.push(Value::Text(prefix.clone()));
547 idx += 1;
548 }
549 if let Some(ts) = &ts_str {
550 sql.push_str(&format!(" AND timestamp >= ?{idx}"));
551 params.push(Value::Text(ts.clone()));
552 idx += 1;
553 }
554 match &partitions {
555 PartitionSelection::All => {}
556 PartitionSelection::One(partition) => {
557 sql.push_str(&format!(" AND partition_count = ?{idx}"));
558 params.push(Value::Integer(partition.count() as i64));
559 idx += 1;
560 sql.push_str(&format!(" AND partition_id = ?{idx}"));
561 params.push(Value::Integer(partition.id() as i64));
562 idx += 1;
563 }
564 PartitionSelection::Many(group) => {
565 sql.push_str(&format!(" AND partition_count = ?{idx}"));
566 params.push(Value::Integer(group.count() as i64));
567 idx += 1;
568 let placeholders: Vec<String> = (0..group.partitions().len())
569 .map(|i| format!("?{}", idx + i))
570 .collect();
571 sql.push_str(&format!(
572 " AND partition_id IN ({})",
573 placeholders.join(",")
574 ));
575 for partition in group.partitions() {
576 params.push(Value::Integer(partition.id() as i64));
577 }
578 idx += group.partitions().len();
579 }
580 }
581 sql.push_str(&format!(" ORDER BY sequence ASC LIMIT ?{idx}"));
582 params.push(Value::Integer(take as i64));
583
584 let mut stmt = guard
585 .prepare(&sql)
586 .map_err(|e| Error::Store(e.to_string()))?;
587 let rows = stmt
588 .query_map(rusqlite::params_from_iter(params.iter()), |row| {
589 let sequence: i64 = row.get(0)?;
590 let id: String = row.get(1)?;
591 let organization: String = row.get(2)?;
592 let namespace: String = row.get(3)?;
593 let topic: String = row.get(4)?;
594 let key: String = row.get(5)?;
595 let payload_str: String = row.get(6)?;
596 let content_type: String = row.get(7)?;
597 let metadata_str: String = row.get(8)?;
598 let timestamp_str: String = row.get(9)?;
599 let version: i64 = row.get(10)?;
600 let parent_id: Option<String> = row.get(11)?;
601 let correlation_id: Option<String> = row.get(12)?;
602 let causation_id: Option<String> = row.get(13)?;
603 let partition_id: Option<i64> = row.get(14)?;
604 let partition_count: Option<i64> = row.get(15)?;
605 Ok((
606 sequence,
607 id,
608 organization,
609 namespace,
610 topic,
611 key,
612 payload_str,
613 content_type,
614 metadata_str,
615 timestamp_str,
616 version,
617 parent_id,
618 correlation_id,
619 causation_id,
620 partition_id,
621 partition_count,
622 ))
623 })
624 .map_err(|e| Error::Store(e.to_string()))?;
625
626 let mut out = Vec::new();
627 for row in rows {
628 let (
629 sequence,
630 id,
631 organization,
632 namespace,
633 topic,
634 key,
635 payload_str,
636 content_type,
637 metadata_str,
638 timestamp_str,
639 version,
640 parent_id,
641 correlation_id,
642 causation_id,
643 partition_id,
644 partition_count,
645 ) = row.map_err(|e| Error::Store(e.to_string()))?;
646
647 let payload: SerializedPayload = serde_json::from_str(&payload_str)
648 .map_err(|e| Error::Serialization(format!("decode payload: {e}")))?;
649 let _ = content_type;
650 let id = uuid::Uuid::parse_str(&id)
651 .map_err(|e| Error::Serialization(format!("decode id: {e}")))?;
652 let parent_id = parent_id
653 .as_deref()
654 .map(uuid::Uuid::parse_str)
655 .transpose()
656 .map_err(|e| Error::Serialization(format!("decode parent_id: {e}")))?;
657 let metadata: HashMap<String, String> = serde_json::from_str(&metadata_str)
658 .map_err(|e| Error::Serialization(format!("decode metadata: {e}")))?;
659 let timestamp = DateTime::parse_from_rfc3339(×tamp_str)
660 .map(|d| d.with_timezone(&Utc))
661 .map_err(|e| Error::Serialization(format!("decode timestamp: {e}")))?;
662 let partition = decode_partition(partition_id, partition_count)?;
663 out.push((
664 SerializedEvent {
665 id,
666 organization,
667 namespace,
668 topic,
669 payload,
670 metadata,
671 timestamp,
672 version: version as u64,
673 key,
674 parent_id,
675 correlation_id,
676 causation_id,
677 },
678 sequence,
679 partition,
680 ));
681 }
682 Ok(out)
683 })
684 .await
685 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
686}
687
688fn decode_partition(partition_id: Option<i64>, partition_count: Option<i64>) -> Result<Partition> {
689 match (partition_id, partition_count) {
690 (Some(id), Some(count)) => {
691 let id = u32::try_from(id)
692 .map_err(|_| Error::Store(format!("partition_id {id} exceeds u32::MAX")))?;
693 let count = u32::try_from(count)
694 .map_err(|_| Error::Store(format!("partition_count {count} exceeds u32::MAX")))?;
695 let count = NonZeroU32::new(count)
696 .ok_or_else(|| Error::Store("partition_count must be positive".to_owned()))?;
697 Partition::new(id, count).map_err(|e| Error::Store(format!("invalid partition: {e}")))
698 }
699 (None, None) => Ok(Partition::new(0, NonZeroU32::new(1).unwrap())
700 .expect("synthetic single partition is valid")),
701 _ => Err(Error::Store(
702 "event has incomplete partition columns — run partition backfill first".to_owned(),
703 )),
704 }
705}
706
707pub type SqlitePartitionedCursor = PartitionedCursor<SqliteCursor>;
724pub type SqliteCoordinatedReaderConfig = CoordinatedReaderConfig;
725pub type SqliteCoordinatedSubscription = CoordinatedSubscription<SqliteSubscription, SqliteCursor>;
726pub type SqliteCoordinatedReader = CoordinatedReader<SqliteReader, SqlitePartitionCoordinator>;
727pub type SqliteCoordinatedAcker =
732 CoordinatedAcker<SqliteCursorAcker, SqliteCursor, SqlitePartitionCoordinator>;
733pub type SqliteCoordinatedStreamAcker = CoordinatedAcker<
735 PartitionAcker<SqliteCursorAcker, SqliteCursor>,
736 PartitionedCursor<SqliteCursor>,
737 PartitionedCoordAdapter<SqlitePartitionCoordinator, SqliteCursor>,
738>;
739pub type SqliteCoordinatedCursor = CoordinatedCursor<PartitionedCursor<SqliteCursor>>;
740pub type SqliteCoordinatedStream = CoordinatedStream<
741 PartitionAcker<SqliteCursorAcker, SqliteCursor>,
742 PartitionedCursor<SqliteCursor>,
743 PartitionedCoordAdapter<SqlitePartitionCoordinator, SqliteCursor>,
744>;
745
746#[cfg(test)]
747mod tests {
748 use std::collections::HashMap;
749 use std::time::Duration;
750
751 use futures::StreamExt;
752 use tokio::time::timeout;
753
754 use super::*;
755 use crate::database::SqliteDatabase;
756 use crate::writer::{SqlitePartitioningConfig, SqliteWriter, SqliteWriterConfig};
757 use eventuary_core::io::cursor::{CursorCodec, CursorOrder};
758 use eventuary_core::io::{Cursor, CursorId, Reader, Writer};
759 use eventuary_core::partition::{
760 EventKeyPartitionKeyResolver, Fnv1a64PartitionHasher, PartitionHasher, PartitionKey,
761 };
762 use eventuary_core::{Event, PartitionableSubscription, Payload, StartFrom, StopAt};
763
764 fn test_partition() -> Partition {
765 Partition::new(0, NonZeroU32::new(1).unwrap()).unwrap()
766 }
767
768 #[test]
769 fn sqlite_subscription_with_partition_wraps_in_singleton_partition_group() {
770 let count = NonZeroU32::new(8).unwrap();
771 let partition = Partition::new(3, count).unwrap();
772 let sub = SqliteSubscription::default().with_partition(partition);
773 match sub.partitions {
774 PartitionSelection::Many(g) => {
775 assert_eq!(g.len(), 1);
776 assert_eq!(g.partitions()[0].id(), 3);
777 assert_eq!(g.count(), 8);
778 }
779 _ => panic!("expected PartitionSelection::Many(singleton)"),
780 }
781 }
782
783 #[test]
784 fn sqlite_subscription_with_partitions_sets_partition_selection_many() {
785 let count = NonZeroU32::new(8).unwrap();
786 let group = PartitionGroup::new(vec![
787 Partition::new(2, count).unwrap(),
788 Partition::new(5, count).unwrap(),
789 ])
790 .unwrap();
791 let sub = SqliteSubscription::default().with_partitions(group);
792 match sub.partitions {
793 PartitionSelection::Many(g) => {
794 assert_eq!(g.len(), 2);
795 let ids: Vec<u32> = g.partitions().iter().map(|p| p.id()).collect();
796 assert_eq!(ids, vec![2, 5]);
797 }
798 _ => panic!("expected PartitionSelection::Many"),
799 }
800 }
801
802 #[test]
803 fn sqlite_cursor_id_is_global() {
804 assert_eq!(
805 SqliteCursor::new(42, test_partition()).id(),
806 CursorId::global()
807 );
808 }
809
810 #[test]
811 fn sqlite_cursor_order_key_from_sequence() {
812 assert_eq!(
813 SqliteCursor::new(42, test_partition()).order_key(),
814 CursorOrder::from_i64(42)
815 );
816 assert!(
817 SqliteCursor::new(9, test_partition()).order_key()
818 < SqliteCursor::new(10, test_partition()).order_key()
819 );
820 }
821
822 #[test]
823 fn sqlite_cursor_codec_roundtrips() {
824 let codec = SqliteCursor::codec().unwrap();
825 let cursor = SqliteCursor::new(42, test_partition());
826 let encoded = codec.encode(&cursor).unwrap();
827 assert_eq!(encoded.kind().as_str(), "eventuary.sqlite.sqlite_cursor.v1");
828 assert_eq!(encoded.order(), &CursorOrder::from_i64(42));
829 assert_eq!(codec.decode(&encoded).unwrap(), cursor);
830 }
831
832 #[test]
833 fn sqlite_cursor_codec_preserves_typed_ord() {
834 let codec = SqliteCursor::codec().unwrap();
835 let lo = codec
836 .encode(&SqliteCursor::new(9, test_partition()))
837 .unwrap();
838 let hi = codec
839 .encode(&SqliteCursor::new(10, test_partition()))
840 .unwrap();
841 assert!(lo < hi);
842 }
843
844 const PARTITION_COUNT: u32 = 4;
845
846 fn event_with_key(key: &str) -> Event {
847 Event::builder(
848 "acme",
849 "/orders",
850 "order.placed",
851 key,
852 Payload::from_string("{}"),
853 )
854 .unwrap()
855 .build()
856 .unwrap()
857 }
858
859 fn partition_for_key(key: &str) -> u32 {
860 let k = PartitionKey::new(key).unwrap();
861 let hash = Fnv1a64PartitionHasher.hash(&k);
862 (hash.get() % PARTITION_COUNT as u64) as u32
863 }
864
865 fn fast_config() -> SqliteReaderConfig {
866 SqliteReaderConfig {
867 poll_interval: Duration::from_millis(10),
868 ..SqliteReaderConfig::default()
869 }
870 }
871
872 #[tokio::test]
873 async fn reader_default_all_returns_every_event() {
874 let db = SqliteDatabase::open_in_memory().unwrap();
875 let config = SqliteWriterConfig {
876 partitioning: SqlitePartitioningConfig::inline(
877 NonZeroU32::new(PARTITION_COUNT).unwrap(),
878 EventKeyPartitionKeyResolver::new(),
879 Fnv1a64PartitionHasher,
880 ),
881 ..SqliteWriterConfig::default()
882 };
883 SqliteWriter::prepare_schema(&db.conn(), &config).unwrap();
884 let writer = SqliteWriter::new_with_config(db.conn(), config);
885
886 let keys = ["k0", "k1", "k2", "k3", "k4", "k5", "k6", "k7"];
887 for key in &keys {
888 writer.write(&event_with_key(key)).await.unwrap();
889 }
890
891 let reader = SqliteReader::new(db.conn(), fast_config());
892 let subscription = SqliteSubscription {
893 start: StartFrom::Earliest,
894 stop_at: StopAt::CurrentEnd,
895 ..SqliteSubscription::default()
896 };
897 let mut stream = reader.read(subscription).await.unwrap();
898
899 let mut count = 0usize;
900 while let Ok(Some(Ok(msg))) = timeout(Duration::from_secs(5), stream.next()).await {
901 msg.acker().ack().await.unwrap();
902 count += 1;
903 }
904
905 assert_eq!(count, keys.len());
906 }
907
908 #[tokio::test]
909 async fn reader_one_filters_to_single_partition() {
910 let db = SqliteDatabase::open_in_memory().unwrap();
911 let config = SqliteWriterConfig {
912 partitioning: SqlitePartitioningConfig::inline(
913 NonZeroU32::new(PARTITION_COUNT).unwrap(),
914 EventKeyPartitionKeyResolver::new(),
915 Fnv1a64PartitionHasher,
916 ),
917 ..SqliteWriterConfig::default()
918 };
919 SqliteWriter::prepare_schema(&db.conn(), &config).unwrap();
920 let writer = SqliteWriter::new_with_config(db.conn(), config);
921
922 let keys = ["k0", "k1", "k2", "k3", "k4", "k5", "k6", "k7"];
923 for key in &keys {
924 writer.write(&event_with_key(key)).await.unwrap();
925 }
926
927 let partitions_by_id: HashMap<u32, Vec<&str>> =
928 keys.iter().fold(HashMap::new(), |mut acc, key| {
929 acc.entry(partition_for_key(key)).or_default().push(key);
930 acc
931 });
932
933 let (chosen_partition, expected_keys) = partitions_by_id
934 .iter()
935 .find(|(_, ks)| ks.len() >= 2)
936 .map(|(id, ks)| (*id, ks.clone()))
937 .expect("expected at least one partition with >=2 events");
938
939 let reader = SqliteReader::new(db.conn(), fast_config());
940 let subscription = SqliteSubscription {
941 start: StartFrom::Earliest,
942 stop_at: StopAt::CurrentEnd,
943 partitions: PartitionSelection::One(
944 Partition::new(chosen_partition, NonZeroU32::new(PARTITION_COUNT).unwrap())
945 .unwrap(),
946 ),
947 ..SqliteSubscription::default()
948 };
949 let mut stream = reader.read(subscription).await.unwrap();
950
951 let mut received_keys: Vec<String> = Vec::new();
952 while let Ok(Some(Ok(msg))) = timeout(Duration::from_secs(5), stream.next()).await {
953 let key = msg.event().key().as_str().to_owned();
954 msg.acker().ack().await.unwrap();
955 received_keys.push(key);
956 }
957
958 assert_eq!(received_keys.len(), expected_keys.len());
959 for key in &received_keys {
960 assert_eq!(
961 partition_for_key(key),
962 chosen_partition,
963 "event key {key} maps to wrong partition"
964 );
965 }
966 }
967
968 #[tokio::test]
969 async fn reader_many_filters_to_selected_partitions() {
970 let db = SqliteDatabase::open_in_memory().unwrap();
971 let config = SqliteWriterConfig {
972 partitioning: SqlitePartitioningConfig::inline(
973 NonZeroU32::new(PARTITION_COUNT).unwrap(),
974 EventKeyPartitionKeyResolver::new(),
975 Fnv1a64PartitionHasher,
976 ),
977 ..SqliteWriterConfig::default()
978 };
979 SqliteWriter::prepare_schema(&db.conn(), &config).unwrap();
980 let writer = SqliteWriter::new_with_config(db.conn(), config);
981
982 let keys = ["k0", "k1", "k2", "k3", "k4", "k5", "k6", "k7"];
983 for key in &keys {
984 writer.write(&event_with_key(key)).await.unwrap();
985 }
986
987 let count = NonZeroU32::new(PARTITION_COUNT).unwrap();
988 let mut populated: Vec<u32> = keys
989 .iter()
990 .map(|k| partition_for_key(k))
991 .collect::<std::collections::BTreeSet<_>>()
992 .into_iter()
993 .collect();
994 populated.truncate(2);
995 assert!(
996 populated.len() >= 2,
997 "fixture must populate at least 2 distinct partitions"
998 );
999
1000 let selected: std::collections::HashSet<u32> = populated.iter().copied().collect();
1001 let expected_len = keys
1002 .iter()
1003 .copied()
1004 .filter(|k| selected.contains(&partition_for_key(k)))
1005 .count();
1006
1007 let group = PartitionGroup::new(
1008 populated
1009 .iter()
1010 .map(|id| Partition::new(*id, count).unwrap())
1011 .collect(),
1012 )
1013 .unwrap();
1014
1015 let reader = SqliteReader::new(db.conn(), fast_config());
1016 let subscription = SqliteSubscription {
1017 start: StartFrom::Earliest,
1018 stop_at: StopAt::CurrentEnd,
1019 ..SqliteSubscription::default()
1020 }
1021 .with_partitions(group);
1022 let mut stream = reader.read(subscription).await.unwrap();
1023
1024 let mut received: Vec<String> = Vec::new();
1025 while let Ok(Some(Ok(msg))) = timeout(Duration::from_secs(5), stream.next()).await {
1026 let key = msg.event().key().as_str().to_owned();
1027 msg.acker().ack().await.unwrap();
1028 received.push(key);
1029 }
1030
1031 assert_eq!(received.len(), expected_len);
1032 for key in &received {
1033 assert!(
1034 selected.contains(&partition_for_key(key)),
1035 "event key {key} maps to partition outside selected group"
1036 );
1037 }
1038 }
1039}