1use std::sync::atomic::Ordering;
2
3use anyhow::Context;
4use thiserror::Error;
5use tracing::info;
6
7use crate::{
8 args::Args,
9 channel,
10 databases::{
11 table::{Column, ColumnType, Row, TableInfo},
12 traits::{DBReader, DBWriter},
13 },
14 progress::TableMigrationProgress,
15 retry::ExponentialRetry,
16};
17
18#[derive(Clone, Debug, PartialEq)]
19pub struct TableMigratorSettings {
20 queue_size: usize,
21 quiet: bool,
22 no_count: bool,
23 writer_workers: usize,
24 batch_write_size: usize,
25 batch_write_retries: usize,
26}
27
28impl From<&Args> for TableMigratorSettings {
29 fn from(args: &Args) -> Self {
30 return Self {
31 queue_size: args.queue_size,
32 quiet: args.quiet,
33 no_count: args.no_count,
34 writer_workers: args.writer_workers,
35 batch_write_size: args.batch_write_size,
36 batch_write_retries: args.batch_write_retries,
37 };
38 }
39}
40
41#[derive(Error, Debug)]
42enum MigratorError {
43 #[error("Stopeed because of the error in an another thread")]
44 Stopped,
45 #[error(transparent)]
46 Other(#[from] anyhow::Error),
47}
48
49pub struct TableMigrator {
50 reader: Box<dyn DBReader>,
51 writers: Vec<Box<dyn DBWriter>>,
52 tracker: TableMigrationProgress,
53 target_format: TableInfo,
54 sender: channel::Sender,
55 reciever: channel::Reciever,
56 stopped: std::sync::atomic::AtomicBool,
57 settings: TableMigratorSettings,
58}
59
60impl TableMigrator {
61 pub fn new(
62 mut reader: Box<dyn DBReader>,
63 mut writer: Box<dyn DBWriter>,
64 table: &str,
65 settings: TableMigratorSettings,
66 ) -> anyhow::Result<TableMigrator> {
67 info!("Collecting info about table \"{table}\"");
68 let reader_table_info = reader
69 .get_table_info(table, settings.no_count)
70 .context("Unable to get information about source table")?;
71 let writer_table_info = writer
72 .get_table_info(table, false)
73 .context("Unable to get information about destination table")?;
74 Self::check_table_compatibility(&reader_table_info, &writer_table_info)?;
75 info!("Table \"{table}\" has passed compatability checks");
76 let mut writers = Vec::new();
77 if settings.writer_workers > 1 {
78 for _ in 0..settings.writer_workers {
79 writers.push(writer.opt_clone()?);
80 }
81 } else {
82 writers.push(writer);
83 }
84 let tracker = TableMigrationProgress::new(table, reader_table_info.num_rows);
85 let (sender, reciever) = channel::create_channel(settings.queue_size);
86 return Ok(TableMigrator {
87 reader,
88 writers,
89 tracker,
90 target_format: writer_table_info,
91 sender,
92 reciever,
93 stopped: std::sync::atomic::AtomicBool::new(false),
94 settings,
95 });
96 }
97
98 fn check_table_compatibility(
99 reader_info: &TableInfo,
100 writer_info: &TableInfo,
101 ) -> anyhow::Result<()> {
102 if writer_info.num_rows != Some(0) {
103 return Err(anyhow::anyhow!(
104 "Destination table \"{}\" should be empty",
105 &writer_info.name
106 ));
107 }
108 let mut reader_columns: Vec<&Column> = reader_info.columns.iter().collect();
109 reader_columns.sort_by(|l, r| r.name.cmp(&l.name));
110 let mut writer_columns: Vec<&Column> = writer_info.columns.iter().collect();
111 writer_columns.sort_by(|l, r| r.name.cmp(&l.name));
112 let err = Err(anyhow::anyhow!(
113 "Incompatable set of columns for table \"{}\".\n\
114 In: {reader_columns:?}\n\
115 Out: {writer_columns:?}",
116 &writer_info.name
117 ));
118 if reader_columns.len() != writer_columns.len() {
119 return err;
120 }
121 for (reader_column, writer_column) in std::iter::zip(reader_columns, writer_columns) {
122 if reader_column.name != writer_column.name
123 || (!writer_column.nullable && reader_column.nullable)
124 {
125 return err;
126 }
127 let allowed = matches!(
128 (reader_column.column_type, writer_column.column_type),
129 (ColumnType::I16, ColumnType::I32 | ColumnType::I64)
130 | (ColumnType::I32, ColumnType::I64)
131 | (ColumnType::F32, ColumnType::F64)
132 );
133
134 if !allowed && reader_column.column_type != writer_column.column_type {
135 return err;
136 }
137 }
138 return Ok(());
139 }
140
141 fn start_reading(
142 mut reader: Box<dyn DBReader>,
143 sender: channel::Sender,
144 tracker: &TableMigrationProgress,
145 target_format: TableInfo,
146 stopped: &std::sync::atomic::AtomicBool,
147 ) -> Result<(), MigratorError> {
148 let iterator = reader.read_iter(target_format)?;
149 for result in iterator {
150 if stopped.load(Ordering::Relaxed) {
151 return Err(MigratorError::Stopped);
152 }
153 let row = result?;
154 sender.send(row).map_err(|_| MigratorError::Stopped)?;
155 tracker.inc_reader(1);
156 }
157 tracker.finish_reader();
158 return Ok(());
159 }
160
161 fn start_writing(
162 mut writer: Box<dyn DBWriter>,
163 reciever: channel::Reciever,
164 tracker: &TableMigrationProgress,
165 table: &TableInfo,
166 batch_size: usize,
167 retries: usize,
168 stopped: &std::sync::atomic::AtomicBool,
169 ) -> Result<(), MigratorError> {
170 let mut batch: Vec<Row> = Vec::with_capacity(batch_size);
171 while let Ok(row) = reciever.recv() {
172 if stopped.load(Ordering::Relaxed) {
173 return Err(MigratorError::Stopped);
174 }
175 batch.push(row);
176 if batch.len() == batch_size {
177 writer
178 .write_batch_with_retry(&batch, table, ExponentialRetry::new(retries))
179 .context("Writing error")?;
180 tracker.inc_writer(batch.len().try_into().unwrap());
181 batch.clear();
182 }
183 }
184 if !batch.is_empty() {
185 writer
186 .write_batch_with_retry(&batch, table, ExponentialRetry::new(retries))
187 .context("Writing error")?;
188 tracker.inc_writer(batch.len().try_into().unwrap());
189 }
190 return Ok(());
191 }
192
193 pub fn run(self) -> anyhow::Result<()> {
194 let process_result = |r: Result<(), MigratorError>| match r {
195 Ok(()) | Err(MigratorError::Stopped) => Ok(()),
196 Err(MigratorError::Other(e)) => {
197 self.stopped.store(true, Ordering::Relaxed);
198 Err(e)
199 }
200 };
201 info!("Start moving data of table \"{}\"", self.target_format.name);
202 return std::thread::scope(|s| {
203 let mut handles = Vec::new();
204 handles.push(s.spawn(|| {
205 return process_result(Self::start_reading(
206 self.reader,
207 self.sender,
208 &self.tracker,
209 self.target_format.clone(),
210 &self.stopped,
211 ));
212 }));
213 for writer in self.writers {
214 handles.push(s.spawn(|| {
215 return process_result(Self::start_writing(
216 writer,
217 self.reciever.clone(),
218 &self.tracker,
219 &self.target_format,
220 self.settings.batch_write_size,
221 self.settings.batch_write_retries,
222 &self.stopped,
223 ));
224 }));
225 }
226 for handle in handles {
228 handle.join().unwrap()?;
229 }
230 self.tracker.finish_writer();
231 return Ok(());
232 });
233 }
234}
235
236#[cfg(test)]
237mod tests {
238 use crate::databases::table::TableInfo;
239 use crate::databases::traits::{DBInfoProvider, ReaderIterator, WriterError};
240
241 use super::*;
242 use mockall::{mock, predicate::*};
243 use std::sync::atomic::AtomicBool;
244
245 mock! {
246 RowsIter {}
247
248 impl Iterator for RowsIter {
249 type Item = anyhow::Result<Row>;
250
251 fn next(&mut self) -> Option<<Self as Iterator>::Item>;
252 }
253 }
254
255 mock! {
256 DB {}
257
258 impl DBInfoProvider for DB {
259 fn get_tables(&mut self) -> anyhow::Result<Vec<String>>;
260 fn get_table_info(&mut self, table: &str, no_count: bool) -> anyhow::Result<TableInfo>;
261 }
262 impl DBReader for DB {
263 fn read_iter<'a>(&'a mut self, target_format: TableInfo) -> anyhow::Result<ReaderIterator<'a>>;
264 }
265 impl DBWriter for DB {
266 fn write_batch(&mut self, batch: &[Row], table: &TableInfo) -> Result<(), WriterError>;
267 fn recover(&mut self) -> anyhow::Result<()>;
268 }
269 }
270
271 impl Default for TableMigratorSettings {
272 fn default() -> Self {
273 Self {
274 queue_size: 10,
275 quiet: true,
276 no_count: false,
277 writer_workers: 1,
278 batch_write_size: 10,
279 batch_write_retries: 0,
280 }
281 }
282 }
283
284 const NUM_ROWS: u64 = 5;
285
286 const TABLE_NAME: &str = "test";
287
288 impl TableInfo {
289 fn default_in() -> Self {
290 Self {
291 name: TABLE_NAME.to_string(),
292 num_rows: None,
293 columns: Vec::new(),
294 }
295 }
296
297 fn default_out() -> Self {
298 Self {
299 name: TABLE_NAME.to_string(),
300 num_rows: Some(0),
301 columns: Vec::new(),
302 }
303 }
304 }
305
306 #[test]
307 fn test_reading() {
308 let mut db_mock = MockDB::new();
309 db_mock.expect_read_iter().returning(|_| {
310 let mut rows = MockRowsIter::new();
311 let mut count = 0;
312 rows.expect_next().returning(move || {
313 if count == NUM_ROWS {
314 return None;
315 }
316 count += 1;
317 return Some(Ok(Row::default()));
318 });
319 Ok(Box::new(rows))
320 });
321 let (sender, receiver) = channel::create_channel(10);
322 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
323 let stopped = AtomicBool::new(false);
324
325 let result = TableMigrator::start_reading(
326 Box::new(db_mock),
327 sender,
328 &tracker,
329 TableInfo::default_out(),
330 &stopped,
331 );
332 assert!(matches!(result, Ok(())));
333 assert_eq!(tracker.reader_processed(), NUM_ROWS);
334 assert_eq!(receiver.len() as u64, NUM_ROWS);
335 }
336
337 #[test]
338 fn test_reading_stops_on_signal() {
339 let mut db_mock = MockDB::new();
340 db_mock.expect_read_iter().returning(|_| {
341 let mut rows = MockRowsIter::new();
342 rows.expect_next().returning(|| Some(Ok(Row::default())));
343 Ok(Box::new(rows))
344 });
345 let (sender, _receiver) = channel::create_channel(10);
346 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
347 let stopped = AtomicBool::new(true);
348
349 let result = TableMigrator::start_reading(
350 Box::new(db_mock),
351 sender,
352 &tracker,
353 TableInfo::default_out(),
354 &stopped,
355 );
356 assert!(matches!(result, Err(MigratorError::Stopped)));
357 }
358
359 #[test]
360 fn test_reading_stops_on_dropped_writers() {
361 let mut db_mock = MockDB::new();
362 db_mock.expect_read_iter().returning(|_| {
363 let mut rows = MockRowsIter::new();
364 rows.expect_next().returning(|| Some(Ok(Row::default())));
365 Ok(Box::new(rows))
366 });
367 let (sender, receiver) = channel::create_channel(10);
368 drop(receiver);
369 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
370 let stopped = AtomicBool::new(true);
371
372 let result = TableMigrator::start_reading(
373 Box::new(db_mock),
374 sender,
375 &tracker,
376 TableInfo::default_out(),
377 &stopped,
378 );
379 assert!(matches!(result, Err(MigratorError::Stopped)));
380 }
381
382 #[test]
383 fn test_writing_one_batch() {
384 let mut db_mock = MockDB::new();
385 let batch_size = 10;
386 db_mock.expect_write_batch().times(1).returning(|rows, _| {
387 assert_eq!(rows.len() as u64, NUM_ROWS);
388 Ok(())
389 });
390 let (sender, receiver) = channel::create_channel(10);
391 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
392 for _ in 0..NUM_ROWS {
393 sender.send(Row::default()).unwrap();
394 }
395 drop(sender);
396 let stopped = AtomicBool::new(false);
397
398 let result = TableMigrator::start_writing(
399 Box::new(db_mock),
400 receiver,
401 &tracker,
402 &TableInfo::default_out(),
403 batch_size,
404 0,
405 &stopped,
406 );
407 assert!(matches!(result, Ok(())));
408 }
409
410 #[test]
411 fn test_writing_multiple_batches() {
412 let mut db_mock = MockDB::new();
413 let batch_size = 1;
414 db_mock
415 .expect_write_batch()
416 .times(NUM_ROWS as usize)
417 .returning(|rows, _| {
418 assert_eq!(rows.len(), 1);
419 Ok(())
420 });
421
422 let (sender, receiver) = channel::create_channel(10);
423 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
424 for _ in 0..NUM_ROWS {
425 sender.send(Row::default()).unwrap();
426 }
427 drop(sender);
428 let stopped = AtomicBool::new(false);
429
430 let result = TableMigrator::start_writing(
431 Box::new(db_mock),
432 receiver,
433 &tracker,
434 &TableInfo::default_out(),
435 batch_size,
436 0,
437 &stopped,
438 );
439 assert!(matches!(result, Ok(())));
440 }
441
442 #[test]
443 fn test_writing_stops_on_signal() {
444 let db_mock = MockDB::new();
445 let (sender, receiver) = channel::create_channel(10);
446 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
447 sender.send(Row::default()).unwrap();
448 let stopped = AtomicBool::new(true);
449
450 let result = TableMigrator::start_writing(
451 Box::new(db_mock),
452 receiver,
453 &tracker,
454 &TableInfo::default_out(),
455 10,
456 3,
457 &stopped,
458 );
459 assert!(matches!(result, Err(MigratorError::Stopped)));
460 }
461
462 #[test]
463 fn test_writing_stops_on_dropped_reader() {
464 let db_mock = MockDB::new();
465 let (sender, receiver) = channel::create_channel(10);
466 drop(sender);
467 let tracker = TableMigrationProgress::new(TABLE_NAME, None);
468 let stopped = AtomicBool::new(true);
469
470 let result = TableMigrator::start_writing(
471 Box::new(db_mock),
472 receiver,
473 &tracker,
474 &TableInfo::default_out(),
475 10,
476 3,
477 &stopped,
478 );
479 assert!(matches!(result, Ok(())));
480 }
481
482 #[test]
483 fn test_run_success() {
484 let mut reader_mock = MockDB::new();
485 let mut writer_mock = MockDB::new();
486
487 reader_mock
488 .expect_get_table_info()
489 .returning(|_, _| Ok(TableInfo::default_in()));
490
491 writer_mock
492 .expect_get_table_info()
493 .returning(|_, _| Ok(TableInfo::default_out()));
494
495 reader_mock.expect_read_iter().returning(|_| {
496 let mut rows = MockRowsIter::new();
497 let mut count = 0;
498 rows.expect_next().returning(move || {
499 if count == NUM_ROWS {
500 return None;
501 }
502 count += 1;
503 Some(Ok(Row::default()))
504 });
505 Ok(Box::new(rows))
506 });
507
508 writer_mock
509 .expect_write_batch()
510 .times(1)
511 .returning(|rows, _| {
512 assert_eq!(rows.len() as u64, NUM_ROWS);
513 Ok(())
514 });
515
516 let settings = TableMigratorSettings::default();
517 let migrator = TableMigrator::new(
518 Box::new(reader_mock),
519 Box::new(writer_mock),
520 TABLE_NAME,
521 settings,
522 )
523 .expect("Failed to create TableMigrator");
524
525 let result = migrator.run();
526 assert!(result.is_ok());
527 }
528
529 #[test]
530 fn test_run_reader_error() {
531 let mut reader_mock = MockDB::new();
532 let mut writer_mock = MockDB::new();
533
534 reader_mock
535 .expect_get_table_info()
536 .returning(|_, _| Ok(TableInfo::default_in()));
537
538 writer_mock
539 .expect_get_table_info()
540 .returning(|_, _| Ok(TableInfo::default_out()));
541
542 reader_mock.expect_read_iter().returning(|_| {
543 let mut rows = MockRowsIter::new();
544 rows.expect_next()
545 .returning(move || Some(Err(anyhow::anyhow!("Test error"))));
546 Ok(Box::new(rows))
547 });
548
549 let settings = TableMigratorSettings::default();
550 let migrator = TableMigrator::new(
551 Box::new(reader_mock),
552 Box::new(writer_mock),
553 TABLE_NAME,
554 settings,
555 )
556 .expect("Failed to create TableMigrator");
557
558 let result = migrator.run();
559 assert!(result.is_err());
560 let error = result.unwrap_err();
561
562 let root_cause = error.root_cause();
563 assert_eq!(format!("{}", root_cause), "Test error");
564 }
565
566 #[test]
567 fn test_run_writer_error() {
568 let mut reader_mock = MockDB::new();
569 let mut writer_mock = MockDB::new();
570
571 reader_mock
572 .expect_get_table_info()
573 .returning(|_, _| Ok(TableInfo::default_in()));
574
575 writer_mock
576 .expect_get_table_info()
577 .returning(|_, _| Ok(TableInfo::default_out()));
578
579 reader_mock.expect_read_iter().returning(|_| {
580 let mut rows = MockRowsIter::new();
581 let mut count = 0;
582 rows.expect_next().returning(move || {
583 if count == NUM_ROWS {
584 return None;
585 }
586 count += 1;
587 Some(Ok(Row::default()))
588 });
589 Ok(Box::new(rows))
590 });
591
592 writer_mock
593 .expect_write_batch()
594 .times(1)
595 .returning(|_, _| Err(WriterError::Unrecoverable(anyhow::anyhow!("Test error"))));
596
597 let settings = TableMigratorSettings::default();
598 let migrator = TableMigrator::new(
599 Box::new(reader_mock),
600 Box::new(writer_mock),
601 TABLE_NAME,
602 settings,
603 )
604 .expect("Failed to create TableMigrator");
605
606 let result = migrator.run();
607 assert!(result.is_err());
608 let error = result.unwrap_err();
609
610 let root_cause = error.root_cause();
611 assert_eq!(format!("{}", root_cause), "Test error");
612 }
613}