exarrow_rs/query/
import.rs

1//! Import query builder for generating IMPORT SQL statements.
2//!
3//! This module provides a builder pattern for constructing Exasol IMPORT statements
4//! that import data from HTTP endpoints into database tables.
5//!
6//! # Example
7//!
8
9use super::clauses;
10
11/// Row separator options for CSV import.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
13pub enum RowSeparator {
14    /// Line feed (Unix-style)
15    #[default]
16    LF,
17    /// Carriage return (old Mac-style)
18    CR,
19    /// Carriage return + line feed (Windows-style)
20    CRLF,
21}
22
23impl RowSeparator {
24    /// Convert to SQL string representation.
25    pub fn to_sql(&self) -> &'static str {
26        match self {
27            RowSeparator::LF => "LF",
28            RowSeparator::CR => "CR",
29            RowSeparator::CRLF => "CRLF",
30        }
31    }
32}
33
34/// Compression options for import files.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
36pub enum Compression {
37    /// No compression
38    #[default]
39    None,
40    /// Gzip compression (.gz extension)
41    Gzip,
42    /// Bzip2 compression (.bz2 extension)
43    Bzip2,
44}
45
46impl Compression {
47    /// Get the file extension for this compression type.
48    pub fn extension(&self) -> &'static str {
49        match self {
50            Compression::None => ".csv",
51            Compression::Gzip => ".csv.gz",
52            Compression::Bzip2 => ".csv.bz2",
53        }
54    }
55}
56
57/// File format for an IMPORT statement.
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
59pub enum ImportFormat {
60    /// CSV format. Format options (encoding, separators, etc.) apply.
61    #[default]
62    Csv,
63    /// Native Parquet format. The server detects the schema from the file;
64    /// CSV format options and compression suffixes do not apply.
65    Parquet,
66}
67
68/// Trim mode options for CSV import.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
70pub enum TrimMode {
71    /// No trimming
72    #[default]
73    None,
74    /// Trim leading whitespace
75    LTrim,
76    /// Trim trailing whitespace
77    RTrim,
78    /// Trim both leading and trailing whitespace
79    Trim,
80}
81
82impl TrimMode {
83    /// Convert to SQL string representation.
84    pub fn to_sql(&self) -> Option<&'static str> {
85        match self {
86            TrimMode::None => None,
87            TrimMode::LTrim => Some("LTRIM"),
88            TrimMode::RTrim => Some("RTRIM"),
89            TrimMode::Trim => Some("TRIM"),
90        }
91    }
92}
93
94/// Entry for a single file in a multi-file IMPORT statement.
95///
96/// Used with `ImportQuery::with_files()` to build IMPORT statements
97/// that reference multiple FILE clauses for parallel import.
98#[derive(Debug, Clone)]
99pub struct ImportFileEntry {
100    /// Internal address from EXA handshake (format: "host:port")
101    pub address: String,
102    /// File name for this entry (e.g., "001.csv", "002.csv")
103    pub file_name: String,
104    /// Optional public key fingerprint for TLS
105    pub public_key: Option<String>,
106}
107
108impl ImportFileEntry {
109    /// Create a new import file entry.
110    ///
111    /// # Arguments
112    ///
113    /// * `address` - Internal address from EXA handshake
114    /// * `file_name` - File name for this entry
115    /// * `public_key` - Optional public key fingerprint for TLS
116    pub fn new(address: String, file_name: String, public_key: Option<String>) -> Self {
117        Self {
118            address,
119            file_name,
120            public_key,
121        }
122    }
123}
124
125/// Builder for constructing Exasol IMPORT SQL statements.
126///
127/// The ImportQuery builder allows you to configure all aspects of an IMPORT statement
128/// including the target table, columns, CSV format options, and error handling.
129///
130/// # Multi-file Import
131///
132/// For parallel imports, use `with_files()` instead of `at_address()` and `file_name()`:
133///
134/// ```rust
135/// use exarrow_rs::query::import::{ImportQuery, ImportFileEntry};
136///
137/// let entries = vec![
138///     ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
139///     ImportFileEntry::new("10.0.0.6:8563".to_string(), "002.csv".to_string(), None),
140/// ];
141///
142/// let query = ImportQuery::new("my_table")
143///     .with_files(entries)
144///     .build();
145/// ```
146#[derive(Debug, Clone)]
147pub struct ImportQuery {
148    /// Target table name (required)
149    table: String,
150    /// Target schema (optional)
151    schema: Option<String>,
152    /// Columns to import into (optional, imports all if not specified)
153    columns: Option<Vec<String>>,
154    /// HTTP address for the CSV source (e.g., "192.168.1.1:8080") - single file mode
155    address: Option<String>,
156    /// SHA-256 fingerprint for TLS public key verification - single file mode
157    public_key: Option<String>,
158    /// File name (default "001.csv") - single file mode
159    file_name: String,
160    /// Multiple file entries for parallel import
161    file_entries: Option<Vec<ImportFileEntry>>,
162    /// Column separator character (default ',')
163    column_separator: char,
164    /// Column delimiter character for quoting (default '"')
165    column_delimiter: char,
166    /// Row separator (default LF)
167    row_separator: RowSeparator,
168    /// Character encoding (default "UTF-8")
169    encoding: String,
170    /// Number of header rows to skip (default 0)
171    skip: u32,
172    /// Custom NULL value representation
173    null_value: Option<String>,
174    /// Trim mode for values
175    trim: TrimMode,
176    /// Compression type
177    compression: Compression,
178    /// Maximum number of invalid rows before failure
179    reject_limit: Option<u32>,
180    /// File format (CSV or Parquet)
181    format: ImportFormat,
182}
183
184impl ImportQuery {
185    /// Create a new ImportQuery builder for the specified table.
186    ///
187    /// # Arguments
188    /// * `table` - The name of the target table to import into
189    ///
190    pub fn new(table: &str) -> Self {
191        Self {
192            table: table.to_string(),
193            schema: None,
194            columns: None,
195            address: None,
196            public_key: None,
197            file_name: "001.csv".to_string(),
198            file_entries: None,
199            column_separator: ',',
200            column_delimiter: '"',
201            row_separator: RowSeparator::default(),
202            encoding: "UTF-8".to_string(),
203            skip: 0,
204            null_value: None,
205            trim: TrimMode::default(),
206            compression: Compression::default(),
207            reject_limit: None,
208            format: ImportFormat::default(),
209        }
210    }
211
212    /// Set the target schema for the import.
213    ///
214    /// # Arguments
215    /// * `schema` - The schema name
216    pub fn schema(mut self, schema: &str) -> Self {
217        self.schema = Some(schema.to_string());
218        self
219    }
220
221    /// Set the columns to import into.
222    ///
223    /// If not specified, all columns in the table will be used.
224    ///
225    /// # Arguments
226    /// * `cols` - List of column names
227    pub fn columns(mut self, cols: Vec<&str>) -> Self {
228        self.columns = Some(cols.into_iter().map(String::from).collect());
229        self
230    }
231
232    /// Set the HTTP address to import from.
233    ///
234    /// # Arguments
235    /// * `addr` - HTTP address in "host:port" format
236    pub fn at_address(mut self, addr: &str) -> Self {
237        self.address = Some(addr.to_string());
238        self
239    }
240
241    /// Set the public key fingerprint for TLS verification.
242    ///
243    /// When set, HTTPS will be used and the PUBLIC KEY clause will be added.
244    ///
245    /// # Arguments
246    /// * `fingerprint` - SHA-256 fingerprint of the server's public key
247    pub fn with_public_key(mut self, fingerprint: &str) -> Self {
248        self.public_key = Some(fingerprint.to_string());
249        self
250    }
251
252    /// Set the file name for the import.
253    ///
254    /// # Arguments
255    /// * `name` - File name (default "001.csv")
256    pub fn file_name(mut self, name: &str) -> Self {
257        self.file_name = name.to_string();
258        self
259    }
260
261    /// Set the column separator character.
262    ///
263    /// # Arguments
264    /// * `sep` - Separator character (default ',')
265    pub fn column_separator(mut self, sep: char) -> Self {
266        self.column_separator = sep;
267        self
268    }
269
270    /// Set the column delimiter character for quoting.
271    ///
272    /// # Arguments
273    /// * `delim` - Delimiter character (default '"')
274    pub fn column_delimiter(mut self, delim: char) -> Self {
275        self.column_delimiter = delim;
276        self
277    }
278
279    /// Set the row separator.
280    ///
281    /// # Arguments
282    /// * `sep` - Row separator (default LF)
283    pub fn row_separator(mut self, sep: RowSeparator) -> Self {
284        self.row_separator = sep;
285        self
286    }
287
288    /// Set the character encoding.
289    ///
290    /// # Arguments
291    /// * `enc` - Encoding name (default "UTF-8")
292    pub fn encoding(mut self, enc: &str) -> Self {
293        self.encoding = enc.to_string();
294        self
295    }
296
297    /// Set the number of header rows to skip.
298    ///
299    /// # Arguments
300    /// * `rows` - Number of rows to skip (default 0)
301    pub fn skip(mut self, rows: u32) -> Self {
302        self.skip = rows;
303        self
304    }
305
306    /// Set a custom NULL value representation.
307    ///
308    /// # Arguments
309    /// * `val` - String representing NULL values in the CSV
310    pub fn null_value(mut self, val: &str) -> Self {
311        self.null_value = Some(val.to_string());
312        self
313    }
314
315    /// Set the trim mode for imported values.
316    ///
317    /// # Arguments
318    /// * `trim` - Trim mode to apply
319    pub fn trim(mut self, trim: TrimMode) -> Self {
320        self.trim = trim;
321        self
322    }
323
324    /// Set the compression type for the import file.
325    ///
326    /// This affects the file extension in the generated SQL.
327    ///
328    /// # Arguments
329    /// * `compression` - Compression type
330    pub fn compressed(mut self, compression: Compression) -> Self {
331        self.compression = compression;
332        self
333    }
334
335    /// Set the reject limit for error handling.
336    ///
337    /// # Arguments
338    /// * `limit` - Maximum number of invalid rows before the import fails
339    pub fn reject_limit(mut self, limit: u32) -> Self {
340        self.reject_limit = Some(limit);
341        self
342    }
343
344    /// Set the file format for the import.
345    ///
346    /// When set to `ImportFormat::Parquet`, the generated SQL emits
347    /// `FROM PARQUET`, the file extension becomes `.parquet` regardless of
348    /// the configured compression, and the CSV format-options clause is
349    /// omitted (the server reads the schema from the Parquet file).
350    ///
351    /// # Arguments
352    /// * `format` - File format
353    pub fn with_format(mut self, format: ImportFormat) -> Self {
354        self.format = format;
355        self
356    }
357
358    /// Set multiple file entries for parallel import.
359    ///
360    /// This method enables multi-file IMPORT statements where each file
361    /// has its own AT/FILE clause with a unique internal address.
362    ///
363    /// # Arguments
364    /// * `entries` - Vector of file entries with addresses and file names
365    ///
366    /// # Note
367    ///
368    /// When using `with_files()`, the `at_address()` and `file_name()` methods
369    /// are ignored and should not be called.
370    pub fn with_files(mut self, entries: Vec<ImportFileEntry>) -> Self {
371        self.file_entries = Some(entries);
372        self
373    }
374
375    /// Build the complete IMPORT SQL statement.
376    ///
377    /// # Returns
378    /// The generated IMPORT SQL statement as a string.
379    ///
380    pub fn build(&self) -> String {
381        let mut sql = String::with_capacity(512);
382        sql.push_str(&self.target_clause());
383        sql.push_str(&self.source_clause());
384        // Format options apply ONLY to CSV. For PARQUET the server reads the
385        // schema from the file directly, so emitting any of these clauses
386        // (ENCODING, COLUMN SEPARATOR, COLUMN DELIMITER, ROW SEPARATOR, SKIP,
387        // NULL, TRIM, REJECT LIMIT) would either be rejected or ignored.
388        if !self.is_parquet() {
389            sql.push_str(&self.csv_option_clauses());
390        }
391        sql
392    }
393
394    /// `IMPORT INTO [schema.]table [(col, ...)]`
395    fn target_clause(&self) -> String {
396        let mut clause = String::from("IMPORT INTO ");
397        if let Some(ref schema) = self.schema {
398            clause.push_str(schema);
399            clause.push('.');
400        }
401        clause.push_str(&self.table);
402
403        if let Some(ref cols) = self.columns {
404            clause.push_str(" (");
405            clause.push_str(&cols.join(", "));
406            clause.push(')');
407        }
408        clause
409    }
410
411    /// `FROM <FORMAT> AT '<url>' [PUBLIC KEY '...'] FILE '...'`, once per file.
412    ///
413    /// Multi-file mode keeps every `AT`/`FILE` pair on the `FROM` line so the
414    /// server can read the files in parallel; single-file mode puts `FILE` on
415    /// its own line. Only that layout differs between the two modes.
416    fn source_clause(&self) -> String {
417        let mut clause = String::from("\nFROM ");
418        clause.push_str(if self.is_parquet() { "PARQUET" } else { "CSV" });
419
420        match self.file_entries {
421            Some(ref entries) => {
422                for entry in entries {
423                    clause.push_str(" AT ");
424                    clause.push_str(&self.endpoint_of(&entry.address, entry.public_key.as_deref()));
425                    clause.push_str(" FILE '");
426                    clause.push_str(&self.file_name_for(&entry.file_name));
427                    clause.push('\'');
428                }
429            }
430            None => {
431                clause.push_str(" AT ");
432                clause.push_str(&self.endpoint_of(
433                    self.address.as_deref().unwrap_or_default(),
434                    self.public_key.as_deref(),
435                ));
436                clause.push_str("\nFILE '");
437                clause.push_str(&self.file_name_for(&self.file_name));
438                clause.push('\'');
439            }
440        }
441        clause
442    }
443
444    /// Render one quoted source URL plus its optional `PUBLIC KEY` clause.
445    ///
446    /// The `;MaxConcurrentReads=1` option is appended INSIDE the quoted URL for
447    /// native Parquet imports, mirroring the JDBC reference driver's
448    /// on-the-wire shape.
449    fn endpoint_of(&self, address: &str, public_key: Option<&str>) -> String {
450        let url_suffix = if self.is_parquet() {
451            ";MaxConcurrentReads=1"
452        } else {
453            ""
454        };
455        clauses::quoted_endpoint(address, url_suffix, public_key)
456    }
457
458    /// The CSV-only clauses, in the order Exasol expects them.
459    fn csv_option_clauses(&self) -> String {
460        let mut clause = clauses::CsvFraming {
461            encoding: &self.encoding,
462            column_separator: self.column_separator,
463            column_delimiter: self.column_delimiter,
464            row_separator: self.row_separator.to_sql(),
465        }
466        .to_sql();
467
468        if self.skip > 0 {
469            clause.push_str(&format!("\nSKIP = {}", self.skip));
470        }
471        clause.push_str(&clauses::null_clause(self.null_value.as_deref()));
472        if let Some(trim_sql) = self.trim.to_sql() {
473            clause.push_str(&format!("\nTRIM = '{}'", trim_sql));
474        }
475        if let Some(limit) = self.reject_limit {
476            clause.push_str(&format!("\nREJECT LIMIT {}", limit));
477        }
478        clause
479    }
480
481    fn is_parquet(&self) -> bool {
482        matches!(self.format, ImportFormat::Parquet)
483    }
484
485    /// Get a file name with the appropriate extension for the configured format.
486    ///
487    /// For `ImportFormat::Parquet`, the extension is always `.parquet` and
488    /// any configured compression suffix is bypassed. For `ImportFormat::Csv`,
489    /// the extension reflects the configured `Compression`.
490    fn file_name_for(&self, base_name: &str) -> String {
491        let base = strip_known_extensions(base_name);
492        match self.format {
493            ImportFormat::Parquet => format!("{}.parquet", base),
494            ImportFormat::Csv => format!("{}{}", base, self.compression.extension()),
495        }
496    }
497}
498
499fn strip_known_extensions(name: &str) -> &str {
500    let mut current = name;
501    loop {
502        let stripped = current
503            .strip_suffix(".gz")
504            .or_else(|| current.strip_suffix(".bz2"))
505            .or_else(|| current.strip_suffix(".csv"))
506            .or_else(|| current.strip_suffix(".parquet"));
507        match stripped {
508            Some(next) if next.len() < current.len() => current = next,
509            _ => return current,
510        }
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517
518    #[test]
519    fn test_basic_import_statement() {
520        let sql = ImportQuery::new("users")
521            .at_address("192.168.1.1:8080")
522            .build();
523
524        assert!(sql.contains("IMPORT INTO users"));
525        assert!(sql.contains("FROM CSV AT 'http://192.168.1.1:8080'"));
526        assert!(sql.contains("FILE '001.csv'"));
527        assert!(sql.contains("ENCODING = 'UTF-8'"));
528        assert!(sql.contains("COLUMN SEPARATOR = ','"));
529        assert!(sql.contains("COLUMN DELIMITER = '\"'"));
530        assert!(sql.contains("ROW SEPARATOR = 'LF'"));
531    }
532
533    #[test]
534    fn test_import_with_schema() {
535        let sql = ImportQuery::new("users")
536            .schema("myschema")
537            .at_address("192.168.1.1:8080")
538            .build();
539
540        assert!(sql.contains("IMPORT INTO myschema.users"));
541    }
542
543    #[test]
544    fn test_import_with_columns() {
545        let sql = ImportQuery::new("users")
546            .columns(vec!["id", "name", "email"])
547            .at_address("192.168.1.1:8080")
548            .build();
549
550        assert!(sql.contains("IMPORT INTO users (id, name, email)"));
551    }
552
553    #[test]
554    fn test_import_with_all_format_options() {
555        let sql = ImportQuery::new("data")
556            .at_address("10.0.0.1:9000")
557            .column_separator(';')
558            .column_delimiter('\'')
559            .row_separator(RowSeparator::CRLF)
560            .encoding("ISO-8859-1")
561            .skip(2)
562            .null_value("NULL")
563            .trim(TrimMode::Trim)
564            .reject_limit(100)
565            .build();
566
567        assert!(sql.contains("COLUMN SEPARATOR = ';'"));
568        assert!(sql.contains("COLUMN DELIMITER = '''"));
569        assert!(sql.contains("ROW SEPARATOR = 'CRLF'"));
570        assert!(sql.contains("ENCODING = 'ISO-8859-1'"));
571        assert!(sql.contains("SKIP = 2"));
572        assert!(sql.contains("NULL = 'NULL'"));
573        assert!(sql.contains("TRIM = 'TRIM'"));
574        assert!(sql.contains("REJECT LIMIT 100"));
575    }
576
577    #[test]
578    fn test_import_with_use_tls() {
579        let fingerprint = "SHA256:abc123def456";
580        let sql = ImportQuery::new("secure_data")
581            .at_address("192.168.1.1:8443")
582            .with_public_key(fingerprint)
583            .build();
584
585        assert!(sql.contains("FROM CSV AT 'https://192.168.1.1:8443'"));
586        assert!(sql.contains(&format!("PUBLIC KEY '{}'", fingerprint)));
587    }
588
589    #[test]
590    fn test_import_with_gzip_compression() {
591        let sql = ImportQuery::new("compressed_data")
592            .at_address("192.168.1.1:8080")
593            .compressed(Compression::Gzip)
594            .build();
595
596        assert!(sql.contains("FILE '001.csv.gz'"));
597    }
598
599    #[test]
600    fn test_import_with_bzip2_compression() {
601        let sql = ImportQuery::new("compressed_data")
602            .at_address("192.168.1.1:8080")
603            .compressed(Compression::Bzip2)
604            .build();
605
606        assert!(sql.contains("FILE '001.csv.bz2'"));
607    }
608
609    #[test]
610    fn test_import_custom_file_name() {
611        let sql = ImportQuery::new("data")
612            .at_address("192.168.1.1:8080")
613            .file_name("custom_file")
614            .build();
615
616        assert!(sql.contains("FILE 'custom_file.csv'"));
617    }
618
619    #[test]
620    fn test_import_custom_file_name_with_compression() {
621        let sql = ImportQuery::new("data")
622            .at_address("192.168.1.1:8080")
623            .file_name("custom_file")
624            .compressed(Compression::Gzip)
625            .build();
626
627        assert!(sql.contains("FILE 'custom_file.csv.gz'"));
628    }
629
630    #[test]
631    fn test_row_separator_to_sql() {
632        assert_eq!(RowSeparator::LF.to_sql(), "LF");
633        assert_eq!(RowSeparator::CR.to_sql(), "CR");
634        assert_eq!(RowSeparator::CRLF.to_sql(), "CRLF");
635    }
636
637    #[test]
638    fn test_compression_extension() {
639        assert_eq!(Compression::None.extension(), ".csv");
640        assert_eq!(Compression::Gzip.extension(), ".csv.gz");
641        assert_eq!(Compression::Bzip2.extension(), ".csv.bz2");
642    }
643
644    #[test]
645    fn test_trim_mode_to_sql() {
646        assert_eq!(TrimMode::None.to_sql(), None);
647        assert_eq!(TrimMode::LTrim.to_sql(), Some("LTRIM"));
648        assert_eq!(TrimMode::RTrim.to_sql(), Some("RTRIM"));
649        assert_eq!(TrimMode::Trim.to_sql(), Some("TRIM"));
650    }
651
652    #[test]
653    fn test_defaults() {
654        assert_eq!(RowSeparator::default(), RowSeparator::LF);
655        assert_eq!(Compression::default(), Compression::None);
656        assert_eq!(TrimMode::default(), TrimMode::None);
657    }
658
659    #[test]
660    fn test_import_file_entry() {
661        let entry = ImportFileEntry::new(
662            "10.0.0.5:8563".to_string(),
663            "001.csv".to_string(),
664            Some("sha256//abc123".to_string()),
665        );
666
667        assert_eq!(entry.address, "10.0.0.5:8563");
668        assert_eq!(entry.file_name, "001.csv");
669        assert_eq!(entry.public_key, Some("sha256//abc123".to_string()));
670    }
671
672    #[test]
673    fn test_multi_file_import_basic() {
674        let entries = vec![
675            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
676            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.csv".to_string(), None),
677        ];
678
679        let sql = ImportQuery::new("my_table").with_files(entries).build();
680
681        assert!(sql.contains("IMPORT INTO my_table"));
682        assert!(sql.contains("FROM CSV"));
683        assert!(sql.contains("AT 'http://10.0.0.5:8563' FILE '001.csv'"));
684        assert!(sql.contains("AT 'http://10.0.0.6:8564' FILE '002.csv'"));
685    }
686
687    #[test]
688    fn test_multi_file_import_with_tls() {
689        let entries = vec![
690            ImportFileEntry::new(
691                "10.0.0.5:8563".to_string(),
692                "001.csv".to_string(),
693                Some("sha256//fingerprint1".to_string()),
694            ),
695            ImportFileEntry::new(
696                "10.0.0.6:8564".to_string(),
697                "002.csv".to_string(),
698                Some("sha256//fingerprint2".to_string()),
699            ),
700        ];
701
702        let sql = ImportQuery::new("secure_table").with_files(entries).build();
703
704        assert!(sql.contains("AT 'https://10.0.0.5:8563' PUBLIC KEY 'sha256//fingerprint1'"));
705        assert!(sql.contains("AT 'https://10.0.0.6:8564' PUBLIC KEY 'sha256//fingerprint2'"));
706    }
707
708    #[test]
709    fn test_multi_file_import_with_compression() {
710        let entries = vec![
711            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
712            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.csv".to_string(), None),
713        ];
714
715        let sql = ImportQuery::new("compressed_table")
716            .with_files(entries)
717            .compressed(Compression::Gzip)
718            .build();
719
720        assert!(sql.contains("FILE '001.csv.gz'"));
721        assert!(sql.contains("FILE '002.csv.gz'"));
722    }
723
724    #[test]
725    fn test_multi_file_import_three_files() {
726        let entries = vec![
727            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
728            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.csv".to_string(), None),
729            ImportFileEntry::new("10.0.0.7:8565".to_string(), "003.csv".to_string(), None),
730        ];
731
732        let sql = ImportQuery::new("data").with_files(entries).build();
733
734        assert!(sql.contains("AT 'http://10.0.0.5:8563' FILE '001.csv'"));
735        assert!(sql.contains("AT 'http://10.0.0.6:8564' FILE '002.csv'"));
736        assert!(sql.contains("AT 'http://10.0.0.7:8565' FILE '003.csv'"));
737    }
738
739    #[test]
740    fn test_multi_file_import_with_schema_and_columns() {
741        let entries = vec![
742            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
743            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.csv".to_string(), None),
744        ];
745
746        let sql = ImportQuery::new("my_table")
747            .schema("my_schema")
748            .columns(vec!["id", "name", "value"])
749            .with_files(entries)
750            .build();
751
752        assert!(sql.contains("IMPORT INTO my_schema.my_table (id, name, value)"));
753    }
754
755    #[test]
756    fn test_multi_file_import_with_all_options() {
757        let entries = vec![
758            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
759            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.csv".to_string(), None),
760        ];
761
762        let sql = ImportQuery::new("data")
763            .with_files(entries)
764            .encoding("ISO-8859-1")
765            .column_separator(';')
766            .skip(1)
767            .null_value("NULL")
768            .reject_limit(100)
769            .build();
770
771        assert!(sql.contains("ENCODING = 'ISO-8859-1'"));
772        assert!(sql.contains("COLUMN SEPARATOR = ';'"));
773        assert!(sql.contains("SKIP = 1"));
774        assert!(sql.contains("NULL = 'NULL'"));
775        assert!(sql.contains("REJECT LIMIT 100"));
776    }
777
778    #[test]
779    fn test_import_no_skip_when_zero() {
780        let sql = ImportQuery::new("data")
781            .at_address("192.168.1.1:8080")
782            .build();
783
784        // SKIP should not appear when it's 0
785        assert!(!sql.contains("SKIP"));
786    }
787
788    #[test]
789    fn test_import_skip_header_row() {
790        let sql = ImportQuery::new("data")
791            .at_address("192.168.1.1:8080")
792            .skip(1)
793            .build();
794
795        assert!(sql.contains("SKIP = 1"));
796    }
797
798    #[test]
799    fn test_complete_import_statement_format() {
800        let sql = ImportQuery::new("employees")
801            .schema("hr")
802            .columns(vec!["id", "first_name", "last_name", "department"])
803            .at_address("10.20.30.40:8080")
804            .with_public_key("SHA256:fingerprint123")
805            .skip(1)
806            .reject_limit(10)
807            .build();
808
809        // Verify the complete structure
810        let expected_parts = [
811            "IMPORT INTO hr.employees (id, first_name, last_name, department)",
812            "FROM CSV AT 'https://10.20.30.40:8080' PUBLIC KEY 'SHA256:fingerprint123'",
813            "FILE '001.csv'",
814            "ENCODING = 'UTF-8'",
815            "COLUMN SEPARATOR = ','",
816            "COLUMN DELIMITER = '\"'",
817            "ROW SEPARATOR = 'LF'",
818            "SKIP = 1",
819            "REJECT LIMIT 10",
820        ];
821
822        for part in expected_parts {
823            assert!(sql.contains(part), "SQL should contain: {}", part);
824        }
825    }
826
827    #[test]
828    fn test_import_query_format_parquet_single_file() {
829        let sql = ImportQuery::new("my_table")
830            .at_address("10.0.0.5:8563")
831            .with_format(ImportFormat::Parquet)
832            .build();
833        assert!(sql.contains("FROM PARQUET AT 'http://10.0.0.5:8563;MaxConcurrentReads=1'"));
834        assert!(sql.contains("FILE '001.parquet'"));
835        assert!(!sql.contains("ENCODING"));
836    }
837
838    #[test]
839    fn test_import_query_format_parquet_multi_file() {
840        let entries = vec![
841            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.parquet".to_string(), None),
842            ImportFileEntry::new("10.0.0.6:8564".to_string(), "002.parquet".to_string(), None),
843        ];
844        let sql = ImportQuery::new("my_table")
845            .with_files(entries)
846            .with_format(ImportFormat::Parquet)
847            .build();
848        assert!(sql.contains("FROM PARQUET"));
849        assert!(sql.contains("AT 'http://10.0.0.5:8563;MaxConcurrentReads=1'"));
850        assert!(sql.contains("AT 'http://10.0.0.6:8564;MaxConcurrentReads=1'"));
851        assert!(sql.contains("FILE '001.parquet'"));
852        assert!(sql.contains("FILE '002.parquet'"));
853        assert!(!sql.contains("MULTIPLE LOCAL FILES"));
854    }
855
856    #[test]
857    fn test_import_query_format_parquet_with_tls_emits_public_key() {
858        let sql = ImportQuery::new("t")
859            .at_address("10.0.0.5:8563")
860            .with_public_key("sha256//abc")
861            .with_format(ImportFormat::Parquet)
862            .build();
863        assert!(sql
864            .contains("AT 'https://10.0.0.5:8563;MaxConcurrentReads=1' PUBLIC KEY 'sha256//abc'"));
865    }
866
867    #[test]
868    fn test_import_query_format_parquet_omits_csv_format_options() {
869        let sql = ImportQuery::new("t")
870            .at_address("addr:1234")
871            .with_format(ImportFormat::Parquet)
872            .skip(2)
873            .null_value("NULL")
874            .trim(TrimMode::Trim)
875            .reject_limit(100)
876            .build();
877        for keyword in &[
878            "ENCODING",
879            "COLUMN SEPARATOR",
880            "COLUMN DELIMITER",
881            "ROW SEPARATOR",
882            "SKIP",
883            "NULL",
884            "TRIM",
885            "REJECT LIMIT",
886        ] {
887            assert!(
888                !sql.contains(keyword),
889                "SQL should NOT contain {} for Parquet format",
890                keyword
891            );
892        }
893    }
894
895    #[test]
896    fn test_import_query_parquet_ignores_compression() {
897        let sql = ImportQuery::new("t")
898            .at_address("addr:1234")
899            .with_format(ImportFormat::Parquet)
900            .compressed(Compression::Gzip)
901            .build();
902        assert!(sql.contains("FILE '001.parquet'"));
903        assert!(!sql.contains(".gz"));
904    }
905
906    #[test]
907    fn test_import_query_format_csv_unchanged() {
908        let sql = ImportQuery::new("t").at_address("192.168.1.1:8080").build();
909        assert!(sql.contains("FROM CSV AT 'http://192.168.1.1:8080'"));
910        assert!(sql.contains("ENCODING = 'UTF-8'"));
911        assert!(sql.contains("COLUMN SEPARATOR = ','"));
912    }
913
914    // =========================================================================
915    // Whole-statement characterization: exact text, exact clause order
916    // =========================================================================
917
918    #[test]
919    fn test_build_renders_the_full_default_csv_statement_verbatim() {
920        let sql = ImportQuery::new("users")
921            .at_address("192.168.1.1:8080")
922            .build();
923
924        assert_eq!(
925            sql,
926            "IMPORT INTO users\n\
927             FROM CSV AT 'http://192.168.1.1:8080'\n\
928             FILE '001.csv'\n\
929             ENCODING = 'UTF-8'\n\
930             COLUMN SEPARATOR = ','\n\
931             COLUMN DELIMITER = '\"'\n\
932             ROW SEPARATOR = 'LF'"
933        );
934    }
935
936    #[test]
937    fn test_build_renders_every_csv_option_in_order() {
938        let sql = ImportQuery::new("data")
939            .schema("hr")
940            .columns(vec!["a", "b"])
941            .at_address("10.0.0.1:9000")
942            .with_public_key("SHA256:fp")
943            .column_separator(';')
944            .column_delimiter('|')
945            .row_separator(RowSeparator::CRLF)
946            .encoding("ISO-8859-1")
947            .skip(2)
948            .null_value("\\N")
949            .trim(TrimMode::Trim)
950            .reject_limit(100)
951            .compressed(Compression::Gzip)
952            .build();
953
954        assert_eq!(
955            sql,
956            "IMPORT INTO hr.data (a, b)\n\
957             FROM CSV AT 'https://10.0.0.1:9000' PUBLIC KEY 'SHA256:fp'\n\
958             FILE '001.csv.gz'\n\
959             ENCODING = 'ISO-8859-1'\n\
960             COLUMN SEPARATOR = ';'\n\
961             COLUMN DELIMITER = '|'\n\
962             ROW SEPARATOR = 'CRLF'\n\
963             SKIP = 2\n\
964             NULL = '\\N'\n\
965             TRIM = 'TRIM'\n\
966             REJECT LIMIT 100"
967        );
968    }
969
970    #[test]
971    fn test_build_renders_multi_file_entries_on_one_from_line() {
972        let entries = vec![
973            ImportFileEntry::new("10.0.0.5:8563".to_string(), "001.csv".to_string(), None),
974            ImportFileEntry::new(
975                "10.0.0.6:8564".to_string(),
976                "002.csv".to_string(),
977                Some("sha256//fp2".to_string()),
978            ),
979        ];
980
981        let sql = ImportQuery::new("my_table").with_files(entries).build();
982
983        assert_eq!(
984            sql,
985            "IMPORT INTO my_table\n\
986             FROM CSV AT 'http://10.0.0.5:8563' FILE '001.csv' \
987             AT 'https://10.0.0.6:8564' PUBLIC KEY 'sha256//fp2' FILE '002.csv'\n\
988             ENCODING = 'UTF-8'\n\
989             COLUMN SEPARATOR = ','\n\
990             COLUMN DELIMITER = '\"'\n\
991             ROW SEPARATOR = 'LF'"
992        );
993    }
994
995    #[test]
996    fn test_build_renders_the_full_parquet_statement_verbatim() {
997        let sql = ImportQuery::new("my_table")
998            .at_address("10.0.0.5:8563")
999            .with_format(ImportFormat::Parquet)
1000            .build();
1001
1002        assert_eq!(
1003            sql,
1004            "IMPORT INTO my_table\n\
1005             FROM PARQUET AT 'http://10.0.0.5:8563;MaxConcurrentReads=1'\n\
1006             FILE '001.parquet'"
1007        );
1008    }
1009
1010    #[test]
1011    fn test_build_without_address_still_renders_a_scheme_only_url() {
1012        let sql = ImportQuery::new("t").build();
1013
1014        assert!(sql.contains("FROM CSV AT 'http://'"), "got: {}", sql);
1015    }
1016
1017    #[test]
1018    fn test_strip_known_extensions_unwinds_stacked_suffixes() {
1019        assert_eq!(strip_known_extensions("001"), "001");
1020        assert_eq!(strip_known_extensions("001.csv"), "001");
1021        assert_eq!(strip_known_extensions("001.csv.gz"), "001");
1022        assert_eq!(strip_known_extensions("001.csv.bz2"), "001");
1023        assert_eq!(strip_known_extensions("001.parquet"), "001");
1024        assert_eq!(strip_known_extensions("archive.tar"), "archive.tar");
1025    }
1026}
exarrow_rs/query/import.rs

exarrow_rs/query/
import.rs
