sql_splitter/splitter/
mod.rs

1use crate::parser::{determine_buffer_size, ContentFilter, Parser, SqlDialect, StatementType};
2use crate::writer::WriterPool;
3use ahash::AHashSet;
4use std::fs::File;
5use std::io::Read;
6use std::path::{Path, PathBuf};
7
8pub struct Stats {
9    pub statements_processed: u64,
10    pub tables_found: usize,
11    pub bytes_processed: u64,
12    pub table_names: Vec<String>,
13}
14
15#[derive(Default)]
16pub struct SplitterConfig {
17    pub dialect: SqlDialect,
18    pub dry_run: bool,
19    pub table_filter: Option<AHashSet<String>>,
20    pub progress_fn: Option<Box<dyn Fn(u64)>>,
21    pub content_filter: ContentFilter,
22}
23
24/// Compression format detected from file extension
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum Compression {
27    None,
28    Gzip,
29    Bzip2,
30    Xz,
31    Zstd,
32}
33
34impl Compression {
35    /// Detect compression format from file extension
36    pub fn from_path(path: &Path) -> Self {
37        let ext = path
38            .extension()
39            .and_then(|e| e.to_str())
40            .map(|e| e.to_lowercase());
41
42        match ext.as_deref() {
43            Some("gz" | "gzip") => Compression::Gzip,
44            Some("bz2" | "bzip2") => Compression::Bzip2,
45            Some("xz" | "lzma") => Compression::Xz,
46            Some("zst" | "zstd") => Compression::Zstd,
47            _ => Compression::None,
48        }
49    }
50
51    /// Wrap a reader with the appropriate decompressor
52    pub fn wrap_reader<'a>(&self, reader: Box<dyn Read + 'a>) -> Box<dyn Read + 'a> {
53        match self {
54            Compression::None => reader,
55            Compression::Gzip => Box::new(flate2::read::GzDecoder::new(reader)),
56            Compression::Bzip2 => Box::new(bzip2::read::BzDecoder::new(reader)),
57            Compression::Xz => Box::new(xz2::read::XzDecoder::new(reader)),
58            Compression::Zstd => Box::new(zstd::stream::read::Decoder::new(reader).unwrap()),
59        }
60    }
61}
62
63impl std::fmt::Display for Compression {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        match self {
66            Compression::None => write!(f, "none"),
67            Compression::Gzip => write!(f, "gzip"),
68            Compression::Bzip2 => write!(f, "bzip2"),
69            Compression::Xz => write!(f, "xz"),
70            Compression::Zstd => write!(f, "zstd"),
71        }
72    }
73}
74
75pub struct Splitter {
76    input_file: PathBuf,
77    output_dir: PathBuf,
78    config: SplitterConfig,
79}
80
81impl Splitter {
82    pub fn new(input_file: PathBuf, output_dir: PathBuf) -> Self {
83        Self {
84            input_file,
85            output_dir,
86            config: SplitterConfig::default(),
87        }
88    }
89
90    pub fn with_dialect(mut self, dialect: SqlDialect) -> Self {
91        self.config.dialect = dialect;
92        self
93    }
94
95    pub fn with_dry_run(mut self, dry_run: bool) -> Self {
96        self.config.dry_run = dry_run;
97        self
98    }
99
100    pub fn with_table_filter(mut self, tables: Vec<String>) -> Self {
101        if !tables.is_empty() {
102            self.config.table_filter = Some(tables.into_iter().collect());
103        }
104        self
105    }
106
107    pub fn with_progress<F: Fn(u64) + 'static>(mut self, f: F) -> Self {
108        self.config.progress_fn = Some(Box::new(f));
109        self
110    }
111
112    pub fn with_content_filter(mut self, filter: ContentFilter) -> Self {
113        self.config.content_filter = filter;
114        self
115    }
116
117    pub fn split(self) -> anyhow::Result<Stats> {
118        let file = File::open(&self.input_file)?;
119        let file_size = file.metadata()?.len();
120        let buffer_size = determine_buffer_size(file_size);
121        let dialect = self.config.dialect;
122        let content_filter = self.config.content_filter;
123
124        // Detect and apply decompression
125        let compression = Compression::from_path(&self.input_file);
126
127        let reader: Box<dyn Read> = if self.config.progress_fn.is_some() {
128            let progress_reader = ProgressReader::new(file, self.config.progress_fn.unwrap());
129            compression.wrap_reader(Box::new(progress_reader))
130        } else {
131            compression.wrap_reader(Box::new(file))
132        };
133
134        let mut parser = Parser::with_dialect(reader, buffer_size, dialect);
135
136        let mut writer_pool = WriterPool::new(self.output_dir.clone());
137        if !self.config.dry_run {
138            writer_pool.ensure_output_dir()?;
139        }
140
141        let mut tables_seen: AHashSet<String> = AHashSet::new();
142        let mut stats = Stats {
143            statements_processed: 0,
144            tables_found: 0,
145            bytes_processed: 0,
146            table_names: Vec::new(),
147        };
148
149        // Track the last COPY table for PostgreSQL COPY data blocks
150        let mut last_copy_table: Option<String> = None;
151
152        while let Some(stmt) = parser.read_statement()? {
153            let (stmt_type, mut table_name) =
154                Parser::<&[u8]>::parse_statement_with_dialect(&stmt, dialect);
155
156            // Track COPY statements for data association
157            if stmt_type == StatementType::Copy {
158                last_copy_table = Some(table_name.clone());
159            }
160
161            // Handle PostgreSQL COPY data blocks - associate with last COPY table
162            let is_copy_data = if stmt_type == StatementType::Unknown && last_copy_table.is_some() {
163                // Check if this looks like COPY data (ends with \.\n)
164                if stmt.ends_with(b"\\.\n") || stmt.ends_with(b"\\.\r\n") {
165                    table_name = last_copy_table.take().unwrap();
166                    true
167                } else {
168                    false
169                }
170            } else {
171                false
172            };
173
174            if !is_copy_data && (stmt_type == StatementType::Unknown || table_name.is_empty()) {
175                continue;
176            }
177
178            // Apply content filter (schema-only or data-only)
179            match content_filter {
180                ContentFilter::SchemaOnly => {
181                    if !stmt_type.is_schema() {
182                        continue;
183                    }
184                }
185                ContentFilter::DataOnly => {
186                    // For data-only, include INSERT, COPY, and COPY data blocks
187                    if !stmt_type.is_data() && !is_copy_data {
188                        continue;
189                    }
190                }
191                ContentFilter::All => {}
192            }
193
194            if let Some(ref filter) = self.config.table_filter {
195                if !filter.contains(&table_name) {
196                    continue;
197                }
198            }
199
200            if !tables_seen.contains(&table_name) {
201                tables_seen.insert(table_name.clone());
202                stats.tables_found += 1;
203                stats.table_names.push(table_name.clone());
204            }
205
206            if !self.config.dry_run {
207                writer_pool.write_statement(&table_name, &stmt)?;
208            }
209
210            stats.statements_processed += 1;
211            stats.bytes_processed += stmt.len() as u64;
212        }
213
214        if !self.config.dry_run {
215            writer_pool.close_all()?;
216        }
217
218        Ok(stats)
219    }
220}
221
222struct ProgressReader<R: Read> {
223    reader: R,
224    callback: Box<dyn Fn(u64)>,
225    bytes_read: u64,
226}
227
228impl<R: Read> ProgressReader<R> {
229    fn new(reader: R, callback: Box<dyn Fn(u64)>) -> Self {
230        Self {
231            reader,
232            callback,
233            bytes_read: 0,
234        }
235    }
236}
237
238impl<R: Read> Read for ProgressReader<R> {
239    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
240        let n = self.reader.read(buf)?;
241        self.bytes_read += n as u64;
242        (self.callback)(self.bytes_read);
243        Ok(n)
244    }
245}
sql_splitter/splitter/mod.rs

sql_splitter/splitter/
mod.rs
