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;
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    cdc_offset_idx: AtomicUsize,
35    parallel_backfill_snapshots: Vec<OwnedRow>,
36}
37
38impl MockExternalTableReader {
39    pub fn new() -> Self {
40        let binlog_file = String::from("1.binlog");
41        // mock binlog watermarks for backfill
42        // initial low watermark: 1.binlog, pos=2 and expected behaviors:
43        // - ignore events before (1.binlog, pos=2);
44        // - apply events in the range of (1.binlog, pos=2, 1.binlog, pos=4) to the snapshot
45        let binlog_watermarks = vec![
46            MySqlOffset::new(binlog_file.clone(), 2), // binlog low watermark
47            MySqlOffset::new(binlog_file.clone(), 4),
48            MySqlOffset::new(binlog_file.clone(), 6),
49            MySqlOffset::new(binlog_file.clone(), 8),
50            MySqlOffset::new(binlog_file, 10),
51        ];
52        let parallel_backfill_snapshots = vec![
53            OwnedRow::new(vec![
54                Some(ScalarImpl::Int64(1)),
55                Some(ScalarImpl::Float64(1.0001.into())),
56            ]),
57            OwnedRow::new(vec![
58                Some(ScalarImpl::Int64(1)),
59                Some(ScalarImpl::Float64(11.00.into())),
60            ]),
61            OwnedRow::new(vec![
62                Some(ScalarImpl::Int64(2)),
63                Some(ScalarImpl::Float64(22.00.into())),
64            ]),
65            OwnedRow::new(vec![
66                Some(ScalarImpl::Int64(5)),
67                Some(ScalarImpl::Float64(1.0005.into())),
68            ]),
69            OwnedRow::new(vec![
70                Some(ScalarImpl::Int64(6)),
71                Some(ScalarImpl::Float64(1.0006.into())),
72            ]),
73            OwnedRow::new(vec![
74                Some(ScalarImpl::Int64(900)),
75                Some(ScalarImpl::Float64(900.1.into())),
76            ]),
77            OwnedRow::new(vec![
78                Some(ScalarImpl::Int64(8)),
79                Some(ScalarImpl::Float64(1.0008.into())),
80            ]),
81            OwnedRow::new(vec![
82                Some(ScalarImpl::Int64(400)),
83                Some(ScalarImpl::Float64(400.1.into())),
84            ]),
85        ];
86        Self {
87            binlog_watermarks,
88            snapshot_cnt: AtomicUsize::new(0),
89            parallel_backfill_snapshots,
90            cdc_offset_idx: AtomicUsize::new(0),
91        }
92    }
93
94    pub fn get_normalized_table_name(_table_name: &SchemaTableName) -> String {
95        "`mock_table`".to_owned()
96    }
97
98    pub fn get_cdc_offset_parser() -> CdcOffsetParseFunc {
99        Box::new(move |offset| {
100            Ok(CdcOffset::MySql(MySqlOffset::parse_debezium_offset(
101                offset,
102            )?))
103        })
104    }
105
106    /// The snapshot will emit to downstream all in once, because it is too small.
107    /// After that we will emit the buffered upstream chunks all in one.
108    #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
109    async fn snapshot_read_inner(&self) {
110        let snap_idx = self
111            .snapshot_cnt
112            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
113        println!("snapshot read: idx {}", snap_idx);
114
115        let snap0 = vec![
116            OwnedRow::new(vec![
117                Some(ScalarImpl::Int64(1)),
118                Some(ScalarImpl::Float64(1.0001.into())),
119            ]),
120            OwnedRow::new(vec![
121                Some(ScalarImpl::Int64(1)),
122                Some(ScalarImpl::Float64(11.00.into())),
123            ]),
124            OwnedRow::new(vec![
125                Some(ScalarImpl::Int64(2)),
126                Some(ScalarImpl::Float64(22.00.into())),
127            ]),
128            OwnedRow::new(vec![
129                Some(ScalarImpl::Int64(5)),
130                Some(ScalarImpl::Float64(1.0005.into())),
131            ]),
132            OwnedRow::new(vec![
133                Some(ScalarImpl::Int64(6)),
134                Some(ScalarImpl::Float64(1.0006.into())),
135            ]),
136            OwnedRow::new(vec![
137                Some(ScalarImpl::Int64(8)),
138                Some(ScalarImpl::Float64(1.0008.into())),
139            ]),
140        ];
141
142        let snapshots = [snap0];
143        if snap_idx >= snapshots.len() {
144            return Ok(());
145        }
146
147        for row in &snapshots[snap_idx] {
148            yield row.clone();
149        }
150    }
151
152    #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
153    async fn split_snapshot_read_inner(&self, left: OwnedRow, right: OwnedRow) {
154        for row in &self.parallel_backfill_snapshots {
155            if (left[0].is_none()
156                || cmp_datum(&row[0], &left[0], OrderType::ascending_nulls_first()).is_ge())
157                && (right[0].is_none()
158                    || cmp_datum(&row[0], &right[0], OrderType::ascending_nulls_first()).is_lt())
159            {
160                yield row.clone();
161            }
162        }
163    }
164}
165
166impl ExternalTableReader for MockExternalTableReader {
167    async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
168        let idx = self
169            .cdc_offset_idx
170            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
171        if idx < self.binlog_watermarks.len() {
172            Ok(CdcOffset::MySql(self.binlog_watermarks[idx].clone()))
173        } else {
174            Ok(CdcOffset::MySql(MySqlOffset {
175                filename: "1.binlog".to_owned(),
176                position: u64::MAX,
177            }))
178        }
179    }
180
181    fn snapshot_read(
182        &self,
183        _table_name: SchemaTableName,
184        _start_pk: Option<OwnedRow>,
185        _primary_keys: Vec<String>,
186        _limit: u32,
187    ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
188        self.snapshot_read_inner()
189    }
190
191    fn get_parallel_cdc_splits(
192        &self,
193        _options: CdcTableSnapshotSplitOption,
194    ) -> BoxStream<'_, ConnectorResult<CdcTableSnapshotSplit>> {
195        unreachable!()
196    }
197
198    fn split_snapshot_read(
199        &self,
200        _table_name: SchemaTableName,
201        left: OwnedRow,
202        right: OwnedRow,
203        split_columns: Vec<Field>,
204    ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
205        assert_eq!(split_columns.len(), 1);
206        assert_eq!(split_columns[0].data_type, DataType::Int64);
207        self.split_snapshot_read_inner(left, right)
208    }
209}
risingwave_connector/source/cdc/external/mock_external_table.rs

risingwave_connector/source/cdc/external/
mock_external_table.rs
