elefant_tools/storage/sql_file/
mod.rs

1use crate::chunk_reader::{ChunkResult, StringChunkReader};
2use crate::helpers::IMPORT_PREFIX;
3use crate::models::PostgresSchema;
4use crate::models::PostgresTable;
5use crate::models::SimplifiedDataType;
6use crate::quoting::{AttemptedKeywordUsage, IdentifierQuoter, Quotable};
7use crate::storage::data_format::DataFormat;
8use crate::storage::table_data::TableData;
9use crate::storage::{BaseCopyTarget, CopyDestination, CopyTransaction};
10use crate::{
11    AsyncCleanup, ColumnIdentity, CopyDestinationFactory, ParallelCopyDestinationNotAvailable,
12    PostgresClientWrapper, Result, SequentialOrParallel, SupportedParallelism, TableDataReader,
13};
14use itertools::Itertools;
15use std::fmt::{Display, Formatter};
16use std::sync::Arc;
17use std::vec;
18use tokio::fs::File;
19use tokio::io::{
20    AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufWriter,
21};
22use tracing::instrument;
23use uuid::Uuid;
24
25#[cfg(test)]
26mod tests;
27
28/// Options that control how the SQL file is generated.
29pub struct SqlFileOptions {
30    /// How many rows are inserted per insert statement.
31    pub max_rows_per_insert: usize,
32    /// The string that separates chunks of commands in the file.
33    pub chunk_separator: String,
34    /// How many DDL commands to generate per chunk at most.
35    pub max_commands_per_chunk: usize,
36    /// How to generate statements for inserting data. See the specific option values
37    /// in [SqlDataMode] for more information.
38    pub data_mode: SqlDataMode,
39}
40
41/// How to generate statements for inserting data.
42#[allow(clippy::tabs_in_doc_comments)]
43#[derive(Debug, Clone, Eq, PartialEq)]
44pub enum SqlDataMode {
45    /// Generate insert statements. A bit slower on import, but might work across many
46    /// database systems.
47    /// Example:
48    /// ```sql
49    /// insert into public.store (store_id, manager_staff_id, address_id, last_update) values
50    /// (1, 1, 1, E'2006-02-15 09:57:12'),
51    /// (2, 2, 2, E'2006-02-15 09:57:12');
52    /// ```
53    InsertStatements,
54    /// Generate copy statements. Much faster on import, but might not work across many
55    /// database systems.
56    /// Example:
57    /// ```sql
58    /// copy public.store (store_id, manager_staff_id, address_id, last_update) from stdin with (format text, header false);
59    /// 1	1	1	2006-02-15 09:57:12
60    /// 2	2	2	2006-02-15 09:57:12
61    /// \.
62    /// ```
63    CopyStatements,
64}
65
66impl Display for SqlDataMode {
67    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
68        match self {
69            SqlDataMode::InsertStatements => write!(f, "InsertStatements"),
70            SqlDataMode::CopyStatements => write!(f, "CopyStatements"),
71        }
72    }
73}
74
75impl From<String> for SqlDataMode {
76    fn from(value: String) -> Self {
77        match value.as_str() {
78            "InsertStatements" => SqlDataMode::InsertStatements,
79            "CopyStatements" => SqlDataMode::CopyStatements,
80            _ => panic!("Invalid value for SqlDataMode"),
81        }
82    }
83}
84
85impl Default for SqlFileOptions {
86    fn default() -> Self {
87        Self {
88            max_rows_per_insert: 1000,
89            chunk_separator: Uuid::new_v4().to_string(),
90            max_commands_per_chunk: 10,
91            data_mode: SqlDataMode::InsertStatements,
92        }
93    }
94}
95
96/// A file to output sql to
97pub struct SqlFile<F: AsyncWrite + Unpin + Send + Sync> {
98    /// The underlying file, though it can be anything that implements `AsyncWrite`
99    file: F,
100    /// If 'nothing' has been written to the chunk yet.
101    is_empty: bool,
102    /// The options that control how the file is generated.
103    options: SqlFileOptions,
104    /// The quoter to use for escaping identifiers.
105    quoter: Arc<IdentifierQuoter>,
106    /// The number of commands written to the current chunk.
107    current_command_count: usize,
108    /// The string that separates chunks of commands in the file.
109    chunk_separator: Vec<u8>,
110}
111
112impl SqlFile<BufWriter<File>> {
113    /// Create a new `SqlFile` from a file path.
114    /// This automatically creates a new file and returns a `SqlFile` that writes to it.
115    #[instrument(skip_all)]
116    pub async fn new_file(
117        path: &str,
118        identifier_quoter: Arc<IdentifierQuoter>,
119        options: SqlFileOptions,
120    ) -> Result<Self> {
121        let file = File::create(path).await?;
122
123        let file = BufWriter::new(file);
124
125        SqlFile::new(file, identifier_quoter, options).await
126    }
127}
128
129static CHUNK_SEPARATOR_PREFIX: &str = "-- chunk-separator-";
130
131impl<F: AsyncWrite + Unpin + Send + Sync> SqlFile<F> {
132    /// Create a new `SqlFile` from a file-like object. This does not do any additional buffering
133    /// so it's recommended to use a `BufWriter` or similar.
134    pub async fn new(
135        mut file: F,
136        identifier_quoter: Arc<IdentifierQuoter>,
137        options: SqlFileOptions,
138    ) -> Result<Self> {
139        let chunk_separator =
140            format!("{}{} --", CHUNK_SEPARATOR_PREFIX, options.chunk_separator).into_bytes();
141
142        file.write_all(&chunk_separator).await?;
143        file.write_all(IMPORT_PREFIX.as_bytes()).await?;
144
145        Ok(SqlFile {
146            file,
147            is_empty: true,
148            options,
149            quoter: identifier_quoter,
150            current_command_count: 0,
151            chunk_separator,
152        })
153    }
154}
155
156impl<F: AsyncWrite + Unpin + Send + Sync> BaseCopyTarget for SqlFile<F> {
157    async fn supported_data_format(&self) -> Result<Vec<DataFormat>> {
158        Ok(vec![DataFormat::Text])
159    }
160}
161
162impl<'a, F: AsyncWrite + Unpin + Send + Sync + 'a> CopyDestinationFactory<'a> for SqlFile<F> {
163    type SequentialDestination = &'a mut SqlFile<F>;
164    type ParallelDestination = ParallelCopyDestinationNotAvailable;
165
166    async fn create_destination(
167        &'a mut self,
168    ) -> Result<SequentialOrParallel<Self::SequentialDestination, Self::ParallelDestination>> {
169        Ok(SequentialOrParallel::Sequential(self))
170    }
171
172    async fn create_sequential_destination(&'a mut self) -> Result<Self::SequentialDestination> {
173        Ok(self)
174    }
175
176    fn supported_parallelism(&self) -> SupportedParallelism {
177        SupportedParallelism::Sequential
178    }
179}
180
181impl<F: AsyncWrite + Unpin + Send + Sync> CopyTransaction for &mut SqlFile<F> {
182    #[instrument(skip_all)]
183    async fn apply_statement(&mut self, statement: &str) -> Result<()> {
184        self.write_statement(statement).await
185    }
186
187    async fn commit(self) -> Result<()> {
188        Ok(())
189    }
190}
191
192impl<F: AsyncWrite + Unpin + Send + Sync> CopyDestination for &mut SqlFile<F> {
193    type Transaction<'a>
194        = &'a mut SqlFile<F>
195    where
196        Self: 'a;
197
198    #[instrument(skip_all)]
199    async fn apply_data<R: TableDataReader, C: AsyncCleanup>(
200        &mut self,
201        schema: &PostgresSchema,
202        table: &PostgresTable,
203        mut data: TableData<R, C>,
204    ) -> Result<()> {
205        let file = &mut self.file;
206        if self.current_command_count > 0 {
207            file.write_all(b"\n").await?;
208            self.current_command_count = 0;
209        }
210
211        if self.options.data_mode == SqlDataMode::InsertStatements {
212            self.write_data_stream_to_insert_statements(&mut data.data, schema, table)
213                .await?;
214        } else {
215            self.write_data_stream_to_copy_statements(&mut data.data, schema, table)
216                .await?;
217        }
218
219        Ok(())
220    }
221
222    #[instrument(skip_all)]
223    async fn apply_non_transactional_statement(&mut self, statement: &str) -> Result<()> {
224        self.write_statement(statement).await
225    }
226
227    async fn begin_transaction(&mut self) -> Result<&mut SqlFile<F>> {
228        Ok(self)
229    }
230
231    fn get_identifier_quoter(&self) -> Arc<IdentifierQuoter> {
232        self.quoter.clone()
233    }
234
235    async fn finish(&mut self) -> Result<()> {
236        self.file.flush().await?;
237        Ok(())
238    }
239}
240
241impl<F: AsyncWrite + Unpin + Send + Sync> SqlFile<F> {
242    /// Writes a single DDL statement to the file, handling chunk separators.
243    #[instrument(skip_all)]
244    async fn write_statement(&mut self, statement: &str) -> Result<()> {
245        if self
246            .current_command_count
247            .is_multiple_of(self.options.max_commands_per_chunk)
248        {
249            if !self.is_empty {
250                self.file.write_all(b"\n\n").await?;
251            }
252
253            self.file.write_all(&self.chunk_separator).await?;
254            self.file.write_all(b"\n").await?;
255            self.is_empty = true;
256        }
257
258        if self.is_empty {
259            self.file.write_all(statement.as_bytes()).await?;
260            self.is_empty = false;
261        } else {
262            self.file.write_all(b"\n\n").await?;
263            self.file.write_all(statement.as_bytes()).await?;
264        }
265
266        self.current_command_count += 1;
267
268        Ok(())
269    }
270
271    /// Writes the data stream to the file as insert statements.
272    #[instrument(skip_all)]
273    async fn write_data_stream_to_insert_statements<R: TableDataReader>(
274        &mut self,
275        reader: &mut R,
276        schema: &PostgresSchema,
277        table: &PostgresTable,
278    ) -> Result<()> {
279        let file = &mut self.file;
280
281        let column_types = table
282            .get_writable_columns()
283            .map(|c| c.get_simplified_data_type())
284            .collect_vec();
285
286        let mut count = 0;
287        while let Some(bytes) = reader.read_chunk().await? {
288            if count == 0 {
289                file.write_all(b"\n").await?;
290                file.write_all(&self.chunk_separator).await?;
291                file.write_all(b"\n").await?;
292            }
293
294            if count % self.options.max_rows_per_insert == 0 {
295                if count > 0 {
296                    file.write_all(b";\n").await?;
297                    file.write_all(&self.chunk_separator).await?;
298                    file.write_all(b"\n").await?;
299                }
300
301                file.write_all(b"insert into ").await?;
302                file.write_all(
303                    schema
304                        .name
305                        .quote(&self.quoter, AttemptedKeywordUsage::TypeOrFunctionName)
306                        .as_bytes(),
307                )
308                .await?;
309                file.write_all(b".").await?;
310                file.write_all(
311                    table
312                        .name
313                        .quote(&self.quoter, AttemptedKeywordUsage::TypeOrFunctionName)
314                        .as_bytes(),
315                )
316                .await?;
317                file.write_all(b" (").await?;
318                for (index, column) in table.get_writable_columns().enumerate() {
319                    if index != 0 {
320                        file.write_all(b", ").await?;
321                    }
322                    file.write_all(column.name.as_bytes()).await?;
323                }
324                file.write_all(b")").await?;
325
326                if table
327                    .columns
328                    .iter()
329                    .any(|c| c.identity == Some(ColumnIdentity::GeneratedAlways))
330                {
331                    file.write_all(b" overriding system value").await?;
332                }
333
334                file.write_all(b" values").await?;
335
336                file.write_all(b"\n").await?;
337                count = 0;
338            } else {
339                file.write_all(b",\n").await?;
340            }
341            count += 1;
342
343            write_row(file, &column_types, bytes).await?;
344        }
345
346        if count > 0 {
347            file.write_all(b";\n").await?;
348        }
349
350        file.flush().await?;
351
352        Ok(())
353    }
354
355    /// Writes the data stream to the file as copy statements.
356    #[instrument(skip_all)]
357    async fn write_data_stream_to_copy_statements<R: TableDataReader>(
358        &mut self,
359        reader: &mut R,
360        schema: &PostgresSchema,
361        table: &PostgresTable,
362    ) -> Result<()> {
363        let file = &mut self.file;
364
365        let mut count = 0;
366        while let Some(bytes) = reader.read_chunk().await? {
367            if count == 0 {
368                file.write_all(b"\n").await?;
369                file.write_all(&self.chunk_separator).await?;
370                file.write_all(b"\n").await?;
371
372                let copy_command =
373                    table.get_copy_in_command(schema, &DataFormat::Text, &self.quoter);
374                file.write_all(copy_command.as_bytes()).await?;
375
376                file.write_all(b"\n").await?;
377                file.write_all(&self.chunk_separator).await?;
378                file.write_all(b"\n").await?;
379            }
380
381            file.write_all(bytes).await?;
382            count += 1;
383        }
384
385        if count > 0 {
386            file.write_all(b"\\.\n").await?;
387            file.flush().await?;
388        }
389
390        Ok(())
391    }
392}
393
394/// Writes a single insert row
395async fn write_row<F: AsyncWrite + Unpin + Send + Sync>(
396    file: &mut F,
397    column_types: &[SimplifiedDataType],
398    bytes: &[u8],
399) -> Result<()> {
400    let without_line_break = &bytes[..bytes.len() - 1];
401    let column_bytes = without_line_break.split(|b| *b == b'\t');
402
403    let cols = column_bytes.zip(column_types.iter());
404    file.write_all(b"(").await?;
405    for (index, (bytes, col_data_type)) in cols.enumerate() {
406        if index != 0 {
407            file.write_all(b", ").await?;
408        }
409
410        write_column(file, bytes, col_data_type).await?;
411    }
412    file.write_all(b")").await?;
413
414    Ok(())
415}
416
417/// Writes a single column in an insert row
418async fn write_column<F: AsyncWrite + Unpin + Send + Sync>(
419    content: &mut F,
420    bytes: &[u8],
421    col_data_type: &SimplifiedDataType,
422) -> Result<()> {
423    if bytes == [b'\\', b'N'] {
424        content.write_all(b"null").await?;
425        return Ok(());
426    }
427
428    match col_data_type {
429        SimplifiedDataType::Number => {
430            write_number_column(content, bytes).await?;
431        }
432        SimplifiedDataType::Text => {
433            write_text_column(content, bytes).await?;
434        }
435        SimplifiedDataType::Bool => {
436            write_bool_column(content, bytes).await?;
437        }
438    }
439
440    Ok(())
441}
442
443/// Writes a `bool` column
444async fn write_bool_column<F: AsyncWrite + Unpin + Send + Sync>(
445    content: &mut F,
446    bytes: &[u8],
447) -> Result<()> {
448    let value = bytes[0] == b't';
449    content.write_all(format!("{value}").as_bytes()).await?;
450    Ok(())
451}
452
453/// Writes a generic `text` column
454async fn write_text_column<F: AsyncWrite + Unpin + Send + Sync>(
455    content: &mut F,
456    bytes: &[u8],
457) -> Result<()> {
458    content.write_all(b"E'").await?;
459
460    if bytes.contains(&b'\'') {
461        let s = std::str::from_utf8(bytes).unwrap();
462        let s = s.replace('\'', "''");
463        content.write_all(s.as_bytes()).await?;
464    } else {
465        content.write_all(bytes).await?;
466    }
467    content.write_all(b"'").await?;
468
469    Ok(())
470}
471
472/// Writes a `number` column
473async fn write_number_column<F: AsyncWrite + Unpin + Send + Sync>(
474    content: &mut F,
475    bytes: &[u8],
476) -> Result<()> {
477    match bytes[..] {
478        [b'N', b'a', b'N']
479        | [b'I', b'n', b'f', b'i', b'n', b'i', b't', b'y']
480        | [b'-', b'I', b'n', b'f', b'i', b'n', b'i', b't', b'y'] => {
481            content.write_all(b"'").await?;
482            content.write_all(bytes).await?;
483            content.write_all(b"'").await?;
484        }
485        _ => {
486            content.write_all(bytes).await?;
487        }
488    }
489
490    Ok(())
491}
492
493/// Applies the provided sql file context to the provided connection.
494/// If the sql file was generated by using the [SqlFile] struct,
495/// this function is quite memory efficient. If not the entire file
496/// will be read into memory before being executed in a single transaction.
497#[instrument(skip_all)]
498pub async fn apply_sql_file<F: AsyncBufRead + Unpin + Send + Sync>(
499    content: &mut F,
500    target_connection: &PostgresClientWrapper,
501) -> Result<()> {
502    let mut sql_chunk = String::with_capacity(10000);
503
504    let read = content.read_line(&mut sql_chunk).await?;
505
506    if read == 0 {
507        return Ok(());
508    }
509
510    if sql_chunk.starts_with(CHUNK_SEPARATOR_PREFIX) {
511        let separator = sql_chunk.clone();
512
513        loop {
514            sql_chunk.clear();
515
516            let read = content
517                .read_lines_until_separator_line(&separator, &mut sql_chunk)
518                .await?;
519            match read {
520                ChunkResult::Chunk(_) => {
521                    if sql_chunk.starts_with("copy ")
522                        && sql_chunk.ends_with(" from stdin with (format text, header false);\n")
523                    {
524                        let mut client = target_connection.pool().get_client().await?;
525                        let mut copy_writer = client.copy_in(&*sql_chunk, &[]).await?;
526
527                        loop {
528                            sql_chunk.clear();
529                            let read = content.read_line(&mut sql_chunk).await?;
530                            if read == 0 {
531                                break;
532                            }
533                            if sql_chunk.starts_with("\\.") {
534                                break;
535                            }
536
537                            copy_writer.write(sql_chunk.as_bytes()).await?;
538                        }
539
540                        copy_writer.end().await?;
541                    } else {
542                        target_connection.execute_non_query(&sql_chunk).await?;
543                    }
544                }
545                ChunkResult::End(read) => {
546                    if read > 0 {
547                        target_connection.execute_non_query(&sql_chunk).await?;
548                    }
549                    break;
550                }
551            }
552        }
553    } else {
554        content.read_to_string(&mut sql_chunk).await?;
555        target_connection.execute_non_query(&sql_chunk).await?;
556    }
557
558    Ok(())
559}
560
561/// Applies the provided sql string to the provided connection. See [apply_sql_file] for more information.
562pub async fn apply_sql_string(
563    file_content: &str,
564    target_connection: &PostgresClientWrapper,
565) -> Result<()> {
566    let mut bytes = file_content.as_bytes();
567    apply_sql_file(&mut bytes, target_connection).await
568}
elefant_tools/storage/sql_file/mod.rs

elefant_tools/storage/sql_file/
mod.rs
