servers/postgres/copy_in/
csv.rs

1// Copyright 2023 Greptime Team
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
15//! The `csv` format — currently the only implemented COPY FROM STDIN
16//! format.
17//!
18//! Parsing and typed conversion both reuse arrow-csv's incremental
19//! [`Decoder`](arrow::csv::reader::Decoder) — the same engine backing
20//! GreptimeDB's file-based `COPY ... FROM 'file.csv'` (see
21//! `common-datasource`'s `csv` file format) — so both COPY routes accept
22//! identical data. The decoder splits records (quoting, escaped quotes,
23//! custom delimiters, `\n`/`\r\n` separators, optional headers) and
24//! converts fields against the targeted columns' schema in one pass.
25//!
26//! Consequences of sharing the engine: JSON and Binary columns are
27//! rejected, timestamps must be date/time strings (offset-free ones
28//! parse in the session timezone), and an empty field — quoted or not —
29//! is NULL, exactly like file-based COPY.
30
31use std::sync::Arc;
32
33use api::helper::vectors_to_rows;
34use arrow::csv::ReaderBuilder;
35use arrow::csv::reader::Decoder as CsvDecoder;
36use arrow_schema::{ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema};
37use datatypes::data_type::DataType;
38use datatypes::prelude::ConcreteDataType;
39use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema};
40use datatypes::vectors::{BooleanVector, Helper, StringVector, Vector};
41use pgwire::error::{PgWireError, PgWireResult};
42
43use crate::postgres::copy_in::codec::{CopyInCodec, ParsedRow, RecordParser};
44use crate::postgres::copy_in::options::CopyOptionSet;
45use crate::postgres::copy_in::{COPY_IN_FLUSH_ROWS, bad_copy_data, invalid_text_repr, unsupported};
46
47/// The validated `csv` format configuration.
48#[derive(Debug)]
49pub(super) struct CsvCodec {
50    delimiter: u8,
51    quote: u8,
52    escape: u8,
53    header: bool,
54}
55
56/// Builds the CSV codec from the parsed options.
57pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult<Arc<dyn CopyInCodec>> {
58    let delimiter = options.take_single_byte("DELIMITER")?.unwrap_or(b',');
59    let quote = options.take_single_byte("QUOTE")?.unwrap_or(b'"');
60    // In PostgreSQL the default escape character for CSV is the quote
61    // character itself (doubled quotes inside quoted fields).
62    let escape = options.take_single_byte("ESCAPE")?.unwrap_or(quote);
63    // An empty string is the default NULL marker (an empty field is NULL);
64    // a custom NULL string is not supported because arrow-csv's null
65    // matching does not distinguish quoted from unquoted fields.
66    if let Some(null) = options.take_string("NULL")?
67        && !null.is_empty()
68    {
69        return Err(unsupported(
70            "custom NULL string is not supported by format csv; an empty field is NULL",
71        ));
72    }
73    let header = options.take_flag("HEADER")?.unwrap_or(false);
74    options.reject_leftovers("csv")?;
75    if delimiter == quote {
76        return Err(bad_copy_data("DELIMITER must not be equal to QUOTE"));
77    }
78    Ok(Arc::new(CsvCodec {
79        delimiter,
80        quote,
81        escape,
82        header,
83    }))
84}
85
86impl CopyInCodec for CsvCodec {
87    fn name(&self) -> &'static str {
88        "csv"
89    }
90
91    /// The column types arrow-csv can decode against.
92    fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> {
93        let supported = matches!(
94            data_type,
95            ConcreteDataType::Boolean(_)
96                | ConcreteDataType::Int8(_)
97                | ConcreteDataType::Int16(_)
98                | ConcreteDataType::Int32(_)
99                | ConcreteDataType::Int64(_)
100                | ConcreteDataType::UInt8(_)
101                | ConcreteDataType::UInt16(_)
102                | ConcreteDataType::UInt32(_)
103                | ConcreteDataType::UInt64(_)
104                | ConcreteDataType::Float32(_)
105                | ConcreteDataType::Float64(_)
106                | ConcreteDataType::String(_)
107                | ConcreteDataType::Date(_)
108                | ConcreteDataType::Timestamp(_)
109                | ConcreteDataType::Decimal128(_)
110        );
111        if !supported {
112            return Err(unsupported(format!(
113                "column type \"{}\" is not supported by format csv",
114                data_type.name()
115            )));
116        }
117        Ok(())
118    }
119
120    fn create_parser(
121        &self,
122        columns: &[ColumnSchema],
123        timezone: &str,
124    ) -> PgWireResult<Box<dyn RecordParser + Send>> {
125        // Decode schema: the greptime schema's arrow projection, except
126        // Boolean columns decode as Utf8 — arrow-csv only accepts
127        // `true`/`false`, while PostgreSQL COPY IN accepts its full set
128        // of boolean spellings (`t`, `f`, `yes`, `on`, ...), which COPY
129        // OUT emits — and Timestamp fields carry the SESSION timezone so
130        // offset-free strings parse like SQL INSERT does instead of as
131        // UTC. The substitutions are applied in `flush_batch` and by the
132        // decoder's tz-aware parsing respectively; explicit offsets in
133        // the input always win.
134        let boolean_columns: Vec<usize> = columns
135            .iter()
136            .enumerate()
137            .filter(|(_, c)| matches!(c.data_type, ConcreteDataType::Boolean(_)))
138            .map(|(i, _)| i)
139            .collect();
140        let arrow_schema = GreptimeSchema::try_new(columns.to_vec())
141            .map_err(|e| unsupported(format!("invalid columns for format csv: {e}")))?
142            .arrow_schema()
143            .clone();
144        let timezone = timezone.to_string();
145        let fields = arrow_schema
146            .fields()
147            .iter()
148            .zip(columns)
149            .map(|(field, column)| match &column.data_type {
150                ConcreteDataType::Boolean(_) => {
151                    Field::new(field.name(), ArrowDataType::Utf8, field.is_nullable())
152                }
153                ConcreteDataType::Timestamp(_) => match field.data_type() {
154                    ArrowDataType::Timestamp(unit, None) => Field::new(
155                        field.name(),
156                        ArrowDataType::Timestamp(*unit, Some(Arc::from(timezone.as_str()))),
157                        field.is_nullable(),
158                    ),
159                    _ => field.as_ref().clone(),
160                },
161                _ => field.as_ref().clone(),
162            })
163            .collect::<Vec<_>>();
164        let mut builder = ReaderBuilder::new(Arc::new(ArrowSchema::new(fields)))
165            .with_header(self.header)
166            .with_delimiter(self.delimiter)
167            .with_quote(self.quote)
168            .with_batch_size(COPY_IN_FLUSH_ROWS);
169        // An escape character equal to the quote is the doubled-quote
170        // default; csv-core handles that through quoting alone.
171        if self.escape != self.quote {
172            builder = builder.with_escape(self.escape);
173        }
174        Ok(Box::new(CsvRecordParser {
175            decoder: builder.build_decoder(),
176            boolean_columns,
177        }))
178    }
179}
180
181/// Incremental parser backed by arrow-csv's decoder.
182struct CsvRecordParser {
183    decoder: CsvDecoder,
184    /// Indices of Boolean columns, decoded as Utf8 and converted by
185    /// [`flush_batch`] to accept PostgreSQL's boolean spellings.
186    boolean_columns: Vec<usize>,
187}
188
189impl RecordParser for CsvRecordParser {
190    fn feed(&mut self, data: &[u8]) -> PgWireResult<Vec<ParsedRow>> {
191        let mut buf = data;
192        let mut rows = Vec::new();
193        while !buf.is_empty() {
194            let consumed = self.decoder.decode(buf).map_err(decode_error)?;
195            if consumed == 0 {
196                while let Some(batch) = self.decoder.flush().map_err(decode_error)? {
197                    rows.extend(self.flush_batch(&batch)?);
198                }
199            } else {
200                buf = &buf[consumed..];
201            }
202        }
203        Ok(rows)
204    }
205
206    fn finish(&mut self) -> PgWireResult<Vec<ParsedRow>> {
207        // An empty decode signals EOF, finalizing a trailing record that
208        // lacks its record separator.
209        self.decoder.decode(&[]).map_err(decode_error)?;
210        let mut rows = Vec::new();
211        while let Some(batch) = self.decoder.flush().map_err(decode_error)? {
212            rows.extend(self.flush_batch(&batch)?);
213        }
214        Ok(rows)
215    }
216}
217
218impl CsvRecordParser {
219    fn flush_batch(
220        &mut self,
221        batch: &arrow::record_batch::RecordBatch,
222    ) -> PgWireResult<Vec<ParsedRow>> {
223        let mut vectors = Helper::try_into_vectors(batch.columns())
224            .map_err(|e| bad_copy_data(format!("failed to convert decoded CSV batch: {e}")))?;
225        for &index in &self.boolean_columns {
226            let Some(vector) = vectors[index].as_any().downcast_ref::<StringVector>() else {
227                return Err(bad_copy_data(
228                    "expected boolean column to decode as a string vector",
229                ));
230            };
231            let mut values = Vec::with_capacity(batch.num_rows());
232            for row_index in 0..batch.num_rows() {
233                values.push(match vector.get(row_index) {
234                    datatypes::value::Value::Null => None,
235                    datatypes::value::Value::String(s) => Some(parse_pg_bool(s.as_utf8())?),
236                    other => {
237                        return Err(bad_copy_data(format!(
238                            "expected string values in a boolean column, got {other:?}"
239                        )));
240                    }
241                });
242            }
243            vectors[index] = Arc::new(BooleanVector::from(values));
244        }
245        Ok(vectors_to_rows(vectors.iter(), batch.num_rows()))
246    }
247}
248
249/// Parses PostgreSQL's boolean spellings: `true`/`false`, `t`/`f`,
250/// `yes`/`no`, `on`/`off` and `1`/`0`, case-insensitive with surrounding
251/// whitespace allowed (see the BOOLEAN type's input grammar). COPY OUT
252/// emits `t`/`f`, so accepting them keeps both COPY directions
253/// round-trippable.
254fn parse_pg_bool(input: &str) -> PgWireResult<bool> {
255    match input.trim().to_ascii_lowercase().as_str() {
256        "true" | "t" | "yes" | "y" | "on" | "1" => Ok(true),
257        "false" | "f" | "no" | "n" | "off" | "0" => Ok(false),
258        _ => Err(invalid_text_repr(format!(
259            "invalid input syntax for type boolean: \"{input}\""
260        ))),
261    }
262}
263
264fn decode_error(error: ArrowError) -> PgWireError {
265    match &error {
266        ArrowError::CsvError(message) => {
267            bad_copy_data(format!("failed to parse CSV data: {message}"))
268        }
269        ArrowError::ParseError(message) | ArrowError::CastError(message) => {
270            invalid_text_repr(format!("failed to parse CSV data: {message}"))
271        }
272        _ => bad_copy_data(format!("failed to parse CSV data: {error}")),
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use datatypes::schema::ColumnSchema;
279
280    use super::*;
281    use crate::postgres::copy_in::options::CopyOptionValue;
282
283    fn test_columns() -> Vec<ColumnSchema> {
284        vec![
285            ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
286            ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true),
287            ColumnSchema::new(
288                "ts",
289                ConcreteDataType::timestamp_millisecond_datatype(),
290                false,
291            ),
292        ]
293    }
294
295    fn parser() -> Box<dyn RecordParser + Send> {
296        build_csv_codec(CopyOptionSet::default())
297            .unwrap()
298            .create_parser(&test_columns(), "UTC")
299            .unwrap()
300    }
301
302    fn string_at(row: &ParsedRow, index: usize) -> Option<String> {
303        row.values[index].value_data.as_ref().map(|v| match v {
304            api::v1::value::ValueData::StringValue(s) => s.clone(),
305            other => panic!("expected string, got {other:?}"),
306        })
307    }
308
309    fn rows_with(mut set: CopyOptionSet) -> Box<dyn RecordParser + Send> {
310        set.take_format().unwrap();
311        build_csv_codec(set)
312            .unwrap()
313            .create_parser(&test_columns(), "UTC")
314            .unwrap()
315    }
316
317    /// Feeds chunks and then the EOF signal, returning all rows: with the
318    /// strict decoder protocol, feeds only emit full
319    /// `COPY_IN_FLUSH_ROWS` batches; a trailing partial batch surfaces at
320    /// `finish`.
321    fn feed_all(parser: &mut Box<dyn RecordParser + Send>, chunks: &[&[u8]]) -> Vec<ParsedRow> {
322        let mut rows = Vec::new();
323        for chunk in chunks {
324            rows.extend(parser.feed(chunk).unwrap());
325        }
326        rows.extend(parser.finish().unwrap());
327        rows
328    }
329
330    #[test]
331    fn test_build_defaults() {
332        let codec = build_csv_codec(CopyOptionSet::default()).unwrap();
333        assert_eq!(codec.name(), "csv");
334    }
335
336    #[test]
337    fn test_build_rejects_bad_options() {
338        let mut set = CopyOptionSet::default();
339        set.insert("delimiter".into(), CopyOptionValue::String("\"".into()))
340            .unwrap();
341        let err = build_csv_codec(set).unwrap_err();
342        assert!(
343            err.to_string()
344                .contains("DELIMITER must not be equal to QUOTE")
345        );
346
347        let mut set = CopyOptionSet::default();
348        set.insert("null".into(), CopyOptionValue::String("NULL".into()))
349            .unwrap();
350        assert!(
351            build_csv_codec(set)
352                .unwrap_err()
353                .to_string()
354                .contains("custom NULL string")
355        );
356
357        let mut set = CopyOptionSet::default();
358        set.take_format().unwrap();
359        set.insert("precision".into(), CopyOptionValue::Ident("ns".into()))
360            .unwrap();
361        assert!(
362            build_csv_codec(set)
363                .unwrap_err()
364                .to_string()
365                .contains("not supported by format csv")
366        );
367    }
368
369    #[test]
370    fn test_check_column_type() {
371        let codec = build_csv_codec(CopyOptionSet::default()).unwrap();
372        assert!(
373            codec
374                .check_column_type(&ConcreteDataType::json_datatype())
375                .is_err()
376        );
377        assert!(
378            codec
379                .check_column_type(&ConcreteDataType::binary_datatype())
380                .is_err()
381        );
382        assert!(
383            codec
384                .check_column_type(&ConcreteDataType::string_datatype())
385                .is_ok()
386        );
387    }
388
389    #[test]
390    fn test_csv_parse_records() {
391        let mut parser = parser();
392        let rows = feed_all(
393            &mut parser,
394            &[b"host1,1.5,2023-11-14 22:13:20\n\
395                ,2.5,2023-11-14 22:13:21\n\
396                \"quoted, comma\",3.5,2023-11-14 22:13:22\n"],
397        );
398        assert_eq!(rows.len(), 3);
399        assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
400        // An empty field is NULL.
401        assert_eq!(string_at(&rows[1], 0), None);
402        assert_eq!(string_at(&rows[2], 0), Some("quoted, comma".into()));
403    }
404
405    #[test]
406    fn test_csv_escaped_quotes() {
407        let mut parser = parser();
408        let rows = feed_all(
409            &mut parser,
410            &[b"\"say \"\"hi\"\"\",1.0,2023-11-14 22:13:20\n"],
411        );
412        assert_eq!(string_at(&rows[0], 0), Some("say \"hi\"".into()));
413    }
414
415    #[test]
416    fn test_csv_custom_escape() {
417        let mut set = CopyOptionSet::default();
418        set.insert("escape".into(), CopyOptionValue::String("\\".into()))
419            .unwrap();
420        let mut parser = rows_with(set);
421        let rows = feed_all(&mut parser, &[b"\"a\\\"b\",1.0,2023-11-14 22:13:20\n"]);
422        assert_eq!(string_at(&rows[0], 0), Some("a\"b".into()));
423    }
424
425    #[test]
426    fn test_csv_quoted_newline_and_chunk_boundaries() {
427        let mut parser = parser();
428        // Split inside a quoted field containing a newline, and a CRLF
429        // record separator split across chunks. No batch is full, so all
430        // feeds return nothing and both rows surface at the EOF signal.
431        assert!(parser.feed(b"\"multi").unwrap().is_empty());
432        assert!(
433            parser
434                .feed(b"line\nfield\",1.0,2023-11-14 22:13:20")
435                .unwrap()
436                .is_empty()
437        );
438        assert!(
439            parser
440                .feed(b"\r\nhost2,2.0,2023-11-14 22:13:21\n")
441                .unwrap()
442                .is_empty()
443        );
444        let rows = parser.finish().unwrap();
445        assert_eq!(rows.len(), 2);
446        assert_eq!(string_at(&rows[0], 0), Some("multiline\nfield".into()));
447        assert_eq!(string_at(&rows[1], 0), Some("host2".into()));
448    }
449
450    #[test]
451    fn test_csv_header() {
452        let mut set = CopyOptionSet::default();
453        set.insert("header".into(), CopyOptionValue::Flag).unwrap();
454        let mut parser = rows_with(set);
455        let rows = feed_all(
456            &mut parser,
457            &[b"host,val,ts\nhost1,1.0,2023-11-14 22:13:20\n"],
458        );
459        assert_eq!(rows.len(), 1);
460        assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
461    }
462
463    #[test]
464    fn test_csv_trailing_record_without_newline() {
465        let mut parser = parser();
466        let fed = parser
467            .feed(b"host1,1.0,2023-11-14 22:13:20\nhost2,2.0,2023-11-14 22:13:21")
468            .unwrap();
469        // While the trailing record is pending, flush cannot emit the
470        // completed one either; both surface at the EOF signal.
471        assert!(fed.is_empty());
472        let rows = parser.finish().unwrap();
473        assert_eq!(rows.len(), 2);
474        assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
475        assert_eq!(string_at(&rows[1], 0), Some("host2".into()));
476    }
477
478    #[test]
479    fn test_csv_finish_at_record_boundary() {
480        let mut parser = parser();
481        // The complete record stays buffered (no full batch); the feed
482        // emits nothing and `finish` surfaces it.
483        assert!(
484            parser
485                .feed(b"host1,1.0,2023-11-14 22:13:20\n")
486                .unwrap()
487                .is_empty()
488        );
489        let rows = parser.finish().unwrap();
490        assert_eq!(rows.len(), 1);
491        assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
492    }
493
494    /// Regression test for chunk-boundary corruption: CopyData chunking is
495    /// arbitrary (records may split mid-field), and an earlier feed loop
496    /// flushed partial batches at record-boundary chunk ends, which
497    /// arrow-csv does not permit mid-stream — fields then parse
498    /// ("00001h,ost-001u,..."). Feed the same payload at many chunk sizes
499    /// and verify EVERY cell.
500    #[test]
501    fn test_csv_chunked_feeds_are_alignment_safe() {
502        // (int32 id, timestamp) rows of near-equal length.
503        let sweep_columns = vec![
504            ColumnSchema::new("i", ConcreteDataType::int32_datatype(), true),
505            ColumnSchema::new(
506                "ts",
507                ConcreteDataType::timestamp_millisecond_datatype(),
508                false,
509            ),
510        ];
511        // A payload wider than one decoder batch so batch-full flushes
512        // interleave with chunk boundaries too.
513        let total = 2 * COPY_IN_FLUSH_ROWS + 137;
514        let payload: Vec<u8> = (0..total)
515            .map(|i| {
516                format!(
517                    "{:06},2024-01-01 00:00:{:02}.{:03}\n",
518                    i,
519                    (i / 1000) % 60,
520                    i % 1000
521                )
522            })
523            .collect::<String>()
524            .into_bytes();
525
526        let check = |chunk: usize| {
527            let mut parser = build_csv_codec(CopyOptionSet::default())
528                .unwrap()
529                .create_parser(&sweep_columns, "UTC")
530                .unwrap();
531            let mut ids: Vec<i32> = Vec::new();
532            for c in payload.chunks(chunk) {
533                for row in parser.feed(c).unwrap() {
534                    ids.push(match row.values[0].value_data.as_ref().unwrap() {
535                        api::v1::value::ValueData::I32Value(v) => *v,
536                        other => panic!("expected i32, got {other:?}"),
537                    });
538                }
539            }
540            for row in parser.finish().unwrap() {
541                ids.push(match row.values[0].value_data.as_ref().unwrap() {
542                    api::v1::value::ValueData::I32Value(v) => *v,
543                    other => panic!("expected i32, got {other:?}"),
544                });
545            }
546            assert_eq!(ids.len(), total, "chunk={chunk}");
547            for (i, v) in ids.iter().enumerate() {
548                assert_eq!(*v, i as i32, "chunk={chunk} row {i} corrupted");
549            }
550        };
551
552        // Small chunk sizes hit every possible boundary alignment,
553        // including multiple decode calls per record.
554        for chunk in [1usize, 2, 3, 7, 8, 13, 31, 64, 97, 128, 1024] {
555            check(chunk);
556        }
557        // Client-natural sizes with a shifted start to vary where
558        // boundaries land relative to records.
559        for chunk in [4096usize, 8192, 65531] {
560            check(chunk);
561            let mut shifted = payload.clone();
562            shifted.drain(..37);
563            let mut parser = build_csv_codec(CopyOptionSet::default())
564                .unwrap()
565                .create_parser(&sweep_columns, "UTC")
566                .unwrap();
567            let mut count = 0usize;
568            for c in shifted.chunks(chunk) {
569                count += parser.feed(c).unwrap().len();
570            }
571            count += parser.finish().unwrap().len();
572            assert_eq!(count, total - 1, "chunk={chunk} shifted");
573        }
574    }
575
576    #[test]
577    fn test_csv_boolean_spellings() {
578        // PostgreSQL COPY OUT emits `t`/`f`; COPY IN accepts the full
579        // boolean input grammar.
580        let columns = vec![
581            ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true),
582            ColumnSchema::new(
583                "ts",
584                ConcreteDataType::timestamp_millisecond_datatype(),
585                false,
586            ),
587        ];
588        let mut parser = build_csv_codec(CopyOptionSet::default())
589            .unwrap()
590            .create_parser(&columns, "UTC")
591            .unwrap();
592        let rows = feed_all(
593            &mut parser,
594            &[b"t,2024-01-01 00:00:00\n\
595                 f,2024-01-01 00:00:01\n\
596                 TRUE,2024-01-01 00:00:02\n\
597                 no,2024-01-01 00:00:03\n\
598                 on,2024-01-01 00:00:04\n\
599                 0,2024-01-01 00:00:05\n\
600                 ,2024-01-01 00:00:06\n"],
601        );
602        let expected = [
603            Some(true),
604            Some(false),
605            Some(true),
606            Some(false),
607            Some(true),
608            Some(false),
609            None,
610        ];
611        for (row, want) in rows.iter().zip(expected) {
612            assert_eq!(
613                row.values[0].value_data,
614                want.map(api::v1::value::ValueData::BoolValue),
615                "{:?}",
616                row.values[0].value_data
617            );
618        }
619    }
620
621    #[test]
622    fn test_csv_boolean_rejects_garbage() {
623        let columns = vec![
624            ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true),
625            ColumnSchema::new(
626                "ts",
627                ConcreteDataType::timestamp_millisecond_datatype(),
628                false,
629            ),
630        ];
631        let mut parser = build_csv_codec(CopyOptionSet::default())
632            .unwrap()
633            .create_parser(&columns, "UTC")
634            .unwrap();
635        parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap();
636        let err = match parser.finish() {
637            Ok(_) => panic!("expected parse error"),
638            Err(e) => e,
639        };
640        assert!(
641            err.to_string()
642                .contains("invalid input syntax for type boolean"),
643            "{err}"
644        );
645    }
646
647    #[test]
648    fn test_csv_timestamp_session_timezone() {
649        // Offset-free timestamp strings must parse in the session
650        // timezone (parity with SQL INSERT); explicit offsets win.
651        let columns = vec![
652            ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
653            ColumnSchema::new(
654                "ts",
655                ConcreteDataType::timestamp_millisecond_datatype(),
656                false,
657            ),
658        ];
659        let mut parser = build_csv_codec(CopyOptionSet::default())
660            .unwrap()
661            .create_parser(&columns, "Asia/Shanghai")
662            .unwrap();
663        let rows = feed_all(
664            &mut parser,
665            &[b"host1,2024-01-01 00:00:00\n\
666                host2,2024-01-01T00:00:00+08:00\n\
667                host3,2024-01-01T00:00:00Z\n"],
668        );
669        let millis = |row: &ParsedRow| match row.values[1].value_data.as_ref().unwrap() {
670            api::v1::value::ValueData::TimestampMillisecondValue(v) => *v,
671            other => panic!("expected ts, got {other:?}"),
672        };
673        // 2024-01-01T00:00:00+08:00 == 2023-12-31T16:00:00Z
674        assert_eq!(millis(&rows[0]), 1704038400000);
675        // Explicit +08:00 offset: same instant as the naive parse.
676        assert_eq!(millis(&rows[1]), 1704038400000);
677        // Explicit Z: parsed as UTC regardless of the session zone.
678        assert_eq!(millis(&rows[2]), 1704067200000);
679    }
680
681    #[test]
682    fn test_csv_errors() {
683        let mut p = parser();
684        let err = match p.feed(b"host1,1.0\n") {
685            Ok(_) => panic!("expected parse error"),
686            Err(e) => e,
687        };
688        assert!(
689            err.to_string().contains("incorrect number of fields"),
690            "{err}"
691        );
692
693        // Value conversion runs at flush time, so the bad double
694        // surfaces at the EOF signal, not at feed time.
695        let mut p = parser();
696        p.feed(b"host1,not_a_double,2023-11-14 22:13:20\n").unwrap();
697        let err = match p.finish() {
698            Ok(_) => panic!("expected parse error"),
699            Err(e) => e,
700        };
701        assert!(err.to_string().contains("failed to parse"), "{err}");
702    }
703}
servers/postgres/copy_in/csv.rs

servers/postgres/copy_in/
csv.rs
