risingwave_connector/source/cdc/external/
mock_external_table.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::atomic::{AtomicUsize, Ordering};
16
17use futures::stream::BoxStream;
18use futures_async_stream::try_stream;
19use risingwave_common::catalog::Field;
20use risingwave_common::row::OwnedRow;
21use risingwave_common::types::{DataType, ScalarImpl};
22use risingwave_common::util::sort_util::{OrderType, cmp_datum};
23
24use crate::error::{ConnectorError, ConnectorResult};
25use crate::source::CdcTableSnapshotSplit;
26use crate::source::cdc::external::{
27    CdcOffset, CdcOffsetParseFunc, CdcTableSnapshotSplitOption, ExternalTableReader, MySqlOffset,
28    SchemaTableName,
29};
30#[derive(Debug)]
31pub struct MockExternalTableReader {
32    binlog_watermarks: Vec<MySqlOffset>,
33    snapshot_cnt: AtomicUsize,
34    snapshot_errors_remaining: AtomicUsize,
35    cdc_offset_idx: AtomicUsize,
36    parallel_backfill_snapshots: Vec<OwnedRow>,
37}
38
39impl MockExternalTableReader {
40    pub fn new() -> Self {
41        let binlog_file = String::from("1.binlog");
42        // mock binlog watermarks for backfill
43        // initial low watermark: 1.binlog, pos=2 and expected behaviors:
44        // - ignore events before (1.binlog, pos=2);
45        // - apply events in the range of (1.binlog, pos=2, 1.binlog, pos=4) to the snapshot
46        let binlog_watermarks = vec![
47            MySqlOffset::new(binlog_file.clone(), 2), // binlog low watermark
48            MySqlOffset::new(binlog_file.clone(), 4),
49            MySqlOffset::new(binlog_file.clone(), 6),
50            MySqlOffset::new(binlog_file.clone(), 8),
51            MySqlOffset::new(binlog_file, 10),
52        ];
53        let parallel_backfill_snapshots = vec![
54            OwnedRow::new(vec![
55                Some(ScalarImpl::Int64(1)),
56                Some(ScalarImpl::Float64(1.0001.into())),
57            ]),
58            OwnedRow::new(vec![
59                Some(ScalarImpl::Int64(1)),
60                Some(ScalarImpl::Float64(11.00.into())),
61            ]),
62            OwnedRow::new(vec![
63                Some(ScalarImpl::Int64(2)),
64                Some(ScalarImpl::Float64(22.00.into())),
65            ]),
66            OwnedRow::new(vec![
67                Some(ScalarImpl::Int64(5)),
68                Some(ScalarImpl::Float64(1.0005.into())),
69            ]),
70            OwnedRow::new(vec![
71                Some(ScalarImpl::Int64(6)),
72                Some(ScalarImpl::Float64(1.0006.into())),
73            ]),
74            OwnedRow::new(vec![
75                Some(ScalarImpl::Int64(900)),
76                Some(ScalarImpl::Float64(900.1.into())),
77            ]),
78            OwnedRow::new(vec![
79                Some(ScalarImpl::Int64(8)),
80                Some(ScalarImpl::Float64(1.0008.into())),
81            ]),
82            OwnedRow::new(vec![
83                Some(ScalarImpl::Int64(400)),
84                Some(ScalarImpl::Float64(400.1.into())),
85            ]),
86        ];
87        Self {
88            binlog_watermarks,
89            snapshot_cnt: AtomicUsize::new(0),
90            snapshot_errors_remaining: AtomicUsize::new(0),
91            parallel_backfill_snapshots,
92            cdc_offset_idx: AtomicUsize::new(0),
93        }
94    }
95
96    pub fn new_with_snapshot_errors(snapshot_errors: usize) -> Self {
97        Self {
98            snapshot_errors_remaining: AtomicUsize::new(snapshot_errors),
99            ..Self::new()
100        }
101    }
102
103    fn take_snapshot_error(&self) -> ConnectorResult<()> {
104        if self
105            .snapshot_errors_remaining
106            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |remaining| {
107                remaining.checked_sub(1)
108            })
109            .is_ok()
110        {
111            return Err(anyhow::anyhow!("mock snapshot read error").into());
112        }
113        Ok(())
114    }
115
116    pub fn get_normalized_table_name(_table_name: &SchemaTableName) -> String {
117        "`mock_table`".to_owned()
118    }
119
120    pub fn get_cdc_offset_parser() -> CdcOffsetParseFunc {
121        Box::new(move |offset| {
122            Ok(CdcOffset::MySql(MySqlOffset::parse_debezium_offset(
123                offset,
124            )?))
125        })
126    }
127
128    /// The snapshot will emit to downstream all in once, because it is too small.
129    /// After that we will emit the buffered upstream chunks all in one.
130    #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
131    async fn snapshot_read_inner(&self) {
132        self.take_snapshot_error()?;
133
134        let snap_idx = self.snapshot_cnt.fetch_add(1, Ordering::Relaxed);
135        println!("snapshot read: idx {}", snap_idx);
136
137        let snap0 = vec![
138            OwnedRow::new(vec![
139                Some(ScalarImpl::Int64(1)),
140                Some(ScalarImpl::Float64(1.0001.into())),
141            ]),
142            OwnedRow::new(vec![
143                Some(ScalarImpl::Int64(1)),
144                Some(ScalarImpl::Float64(11.00.into())),
145            ]),
146            OwnedRow::new(vec![
147                Some(ScalarImpl::Int64(2)),
148                Some(ScalarImpl::Float64(22.00.into())),
149            ]),
150            OwnedRow::new(vec![
151                Some(ScalarImpl::Int64(5)),
152                Some(ScalarImpl::Float64(1.0005.into())),
153            ]),
154            OwnedRow::new(vec![
155                Some(ScalarImpl::Int64(6)),
156                Some(ScalarImpl::Float64(1.0006.into())),
157            ]),
158            OwnedRow::new(vec![
159                Some(ScalarImpl::Int64(8)),
160                Some(ScalarImpl::Float64(1.0008.into())),
161            ]),
162        ];
163
164        let snapshots = [snap0];
165        if snap_idx >= snapshots.len() {
166            return Ok(());
167        }
168
169        for row in &snapshots[snap_idx] {
170            yield row.clone();
171        }
172    }
173
174    #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
175    async fn split_snapshot_read_inner(&self, left: OwnedRow, right: OwnedRow) {
176        self.take_snapshot_error()?;
177
178        for row in &self.parallel_backfill_snapshots {
179            if (left[0].is_none()
180                || cmp_datum(&row[0], &left[0], OrderType::ascending_nulls_first()).is_ge())
181                && (right[0].is_none()
182                    || cmp_datum(&row[0], &right[0], OrderType::ascending_nulls_first()).is_lt())
183            {
184                yield row.clone();
185            }
186        }
187    }
188}
189
190impl ExternalTableReader for MockExternalTableReader {
191    async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
192        let idx = self.cdc_offset_idx.fetch_add(1, Ordering::Relaxed);
193        if idx < self.binlog_watermarks.len() {
194            Ok(CdcOffset::MySql(self.binlog_watermarks[idx].clone()))
195        } else {
196            Ok(CdcOffset::MySql(MySqlOffset {
197                filename: "1.binlog".to_owned(),
198                position: u64::MAX,
199            }))
200        }
201    }
202
203    fn snapshot_read(
204        &self,
205        _table_name: SchemaTableName,
206        _start_pk: Option<OwnedRow>,
207        _primary_keys: Vec<String>,
208        _limit: u32,
209    ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
210        self.snapshot_read_inner()
211    }
212
213    fn get_parallel_cdc_splits(
214        &self,
215        _options: CdcTableSnapshotSplitOption,
216    ) -> BoxStream<'_, ConnectorResult<CdcTableSnapshotSplit>> {
217        unreachable!()
218    }
219
220    fn split_snapshot_read(
221        &self,
222        _table_name: SchemaTableName,
223        left: OwnedRow,
224        right: OwnedRow,
225        split_columns: Vec<Field>,
226    ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
227        assert_eq!(split_columns.len(), 1);
228        assert_eq!(split_columns[0].data_type, DataType::Int64);
229        self.split_snapshot_read_inner(left, right)
230    }
231}
risingwave_connector/source/cdc/external/mock_external_table.rs

risingwave_connector/source/cdc/external/
mock_external_table.rs
