surreal_sync_runtime/pipeline/apply/
row_chunk.rs

1//! Finite row/relation chunk streams for full-sync / keyset table scans.
2//!
3//! Unlike per-call [`crate::pipeline::write_rows`] / [`crate::pipeline::write_relations`],
4//! a [`RowChunkDriver`] / [`RelationChunkDriver`] fed to
5//! [`crate::pipeline::run_source_runtime`] keeps one apply window alive across chunks so
6//! the next read can overlap prior-chunk transform/sink when
7//! `max_in_flight > 1`.
8
9use anyhow::Result;
10use async_trait::async_trait;
11use surreal_sync_core::{Change, Relation, RelationChange, Row};
12
13use super::event::PositionedEvent;
14use super::source_driver::{CheckpointPolicy, SourceDriver};
15
16/// Produces successive row chunks until the table scan is exhausted.
17///
18/// Return `Ok(None)` (or an empty `Vec`) when there are no more rows.
19#[async_trait]
20pub trait RowChunkSource: Send {
21    async fn next_chunk(&mut self) -> Result<Option<Vec<Row>>>;
22}
23
24/// Long-lived [`SourceDriver`] over a [`RowChunkSource`] (CSV-like full-sync pattern).
25///
26/// Each poll loads one chunk, converts rows to upsert changes, and returns them
27/// as positioned events. [`run_source_runtime`](crate::pipeline::run_source_runtime) may
28/// poll the next chunk while earlier chunks are still transforming or sinking.
29pub struct RowChunkDriver<C> {
30    source: C,
31    next_index: u64,
32    sunk_count: u64,
33    finished: bool,
34}
35
36impl<C> RowChunkDriver<C> {
37    /// Wrap a chunk source. `next_index` seeds [`Row::index`] / positions.
38    pub fn new(source: C) -> Self {
39        Self {
40            source,
41            next_index: 0,
42            sunk_count: 0,
43            finished: false,
44        }
45    }
46
47    /// Rows successfully sunk (pre-transform input count via [`SourceDriver::note_sunk_events`]).
48    pub fn sunk_count(&self) -> u64 {
49        self.sunk_count
50    }
51}
52
53#[async_trait]
54impl<C> SourceDriver for RowChunkDriver<C>
55where
56    C: RowChunkSource,
57{
58    type Position = u64;
59
60    async fn poll_work(&mut self) -> Result<Vec<PositionedEvent<Self::Position>>> {
61        if self.finished {
62            return Ok(Vec::new());
63        }
64
65        match self.source.next_chunk().await? {
66            None => {
67                self.finished = true;
68                Ok(Vec::new())
69            }
70            Some(rows) if rows.is_empty() => {
71                self.finished = true;
72                Ok(Vec::new())
73            }
74            Some(mut rows) => {
75                let mut events = Vec::with_capacity(rows.len());
76                for row in rows.drain(..) {
77                    let mut row = row;
78                    row.index = self.next_index;
79                    let pos = self.next_index;
80                    self.next_index = self.next_index.saturating_add(1);
81                    let change = Change::update(row.table, row.id, row.fields);
82                    events.push(PositionedEvent::change(change, pos));
83                }
84                Ok(events)
85            }
86        }
87    }
88
89    async fn advance_watermark(&mut self, _position: Self::Position) -> Result<()> {
90        // File / keyset scans have no durable mid-run cursor.
91        Ok(())
92    }
93
94    fn is_finished(&self) -> bool {
95        self.finished
96    }
97
98    fn checkpoint_policy(&self) -> CheckpointPolicy {
99        CheckpointPolicy::AdvanceOnly
100    }
101
102    fn note_sunk_events(&mut self, count: u64) {
103        self.sunk_count = self.sunk_count.saturating_add(count);
104    }
105}
106
107/// Produces successive relation chunks until the scan is exhausted.
108#[async_trait]
109pub trait RelationChunkSource: Send {
110    async fn next_chunk(&mut self) -> Result<Option<Vec<Relation>>>;
111}
112
113/// Long-lived [`SourceDriver`] over a [`RelationChunkSource`].
114///
115/// Same windowing model as [`RowChunkDriver`], emitting relation upserts.
116pub struct RelationChunkDriver<C> {
117    source: C,
118    next_index: u64,
119    sunk_count: u64,
120    finished: bool,
121}
122
123impl<C> RelationChunkDriver<C> {
124    /// Wrap a relation chunk source.
125    pub fn new(source: C) -> Self {
126        Self {
127            source,
128            next_index: 0,
129            sunk_count: 0,
130            finished: false,
131        }
132    }
133
134    /// Relations successfully sunk (pre-transform input count).
135    pub fn sunk_count(&self) -> u64 {
136        self.sunk_count
137    }
138}
139
140#[async_trait]
141impl<C> SourceDriver for RelationChunkDriver<C>
142where
143    C: RelationChunkSource,
144{
145    type Position = u64;
146
147    async fn poll_work(&mut self) -> Result<Vec<PositionedEvent<Self::Position>>> {
148        if self.finished {
149            return Ok(Vec::new());
150        }
151
152        match self.source.next_chunk().await? {
153            None => {
154                self.finished = true;
155                Ok(Vec::new())
156            }
157            Some(rels) if rels.is_empty() => {
158                self.finished = true;
159                Ok(Vec::new())
160            }
161            Some(mut rels) => {
162                let mut events = Vec::with_capacity(rels.len());
163                for relation in rels.drain(..) {
164                    let pos = self.next_index;
165                    self.next_index = self.next_index.saturating_add(1);
166                    let change = RelationChange::update(relation);
167                    events.push(PositionedEvent::relation_change(change, pos));
168                }
169                Ok(events)
170            }
171        }
172    }
173
174    async fn advance_watermark(&mut self, _position: Self::Position) -> Result<()> {
175        Ok(())
176    }
177
178    fn is_finished(&self) -> bool {
179        self.finished
180    }
181
182    fn checkpoint_policy(&self) -> CheckpointPolicy {
183        CheckpointPolicy::AdvanceOnly
184    }
185
186    fn note_sunk_events(&mut self, count: u64) {
187        self.sunk_count = self.sunk_count.saturating_add(count);
188    }
189}
surreal_sync_runtime/pipeline/apply/row_chunk.rs

surreal_sync_runtime/pipeline/apply/
row_chunk.rs
