servers/postgres/copy_in/
options.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//! Format-neutral collection of COPY `WITH (...)` options.
16//!
17//! The option list is **open**: every `KEY value` entry the client writes
18//! is collected here, regardless of whether any format understands it.
19//! Each format codec then takes the options it knows ([`CopyOptionSet::take_string`]
20//! and friends) and rejects everything else via
21//! [`CopyOptionSet::reject_leftovers`]. This way a format can introduce
22//! its own options — say a line-protocol format accepting
23//! `WITH (FORMAT influxdb_line, precision 'ns')` — without touching the
24//! statement parser or any other codec.
25
26use pgwire::error::{PgWireError, PgWireResult};
27
28use crate::postgres::copy_in::bad_copy_data;
29use crate::postgres::types::PgErrorCode;
30
31/// The value of one COPY option.
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub(crate) enum CopyOptionValue {
34    /// A single-quoted string literal.
35    String(String),
36    /// A numeric literal, kept as written.
37    Number(String),
38    /// A bare identifier.
39    Ident(String),
40    /// `TRUE` / `FALSE`.
41    Boolean(bool),
42    /// The option appears without a value.
43    Flag,
44}
45
46impl CopyOptionValue {
47    /// The value as a descriptive word for error messages.
48    pub(crate) fn kind(&self) -> &'static str {
49        match self {
50            CopyOptionValue::String(_) => "a string value",
51            CopyOptionValue::Number(_) => "a numeric value",
52            CopyOptionValue::Ident(_) => "an identifier value",
53            CopyOptionValue::Boolean(_) => "a boolean value",
54            CopyOptionValue::Flag => "no value",
55        }
56    }
57}
58
59/// The options of a `COPY ... FROM STDIN WITH (...)` statement.
60///
61/// Entries keep statement order; keys are lowercased for lookup and
62/// displayed uppercased in error messages.
63#[derive(Debug, Default)]
64pub(crate) struct CopyOptionSet {
65    entries: Vec<(String, CopyOptionValue)>,
66}
67
68impl CopyOptionSet {
69    /// Adds an option; duplicate keys are rejected like PostgreSQL does.
70    pub(crate) fn insert(&mut self, key: String, value: CopyOptionValue) -> PgWireResult<()> {
71        if self.entries.iter().any(|(k, _)| *k == key) {
72            return Err(bad_copy_data(format!(
73                "conflicting or redundant {} option",
74                key.to_uppercase()
75            )));
76        }
77        self.entries.push((key, value));
78        Ok(())
79    }
80
81    /// Takes the raw value of `key`; for codec-specific options whose
82    /// shape this module does not prescribe.
83    pub(crate) fn take_value(&mut self, key: &str) -> Option<CopyOptionValue> {
84        let key = key.to_lowercase();
85        let index = self.entries.iter().position(|(k, _)| *k == key)?;
86        Some(self.entries.remove(index).1)
87    }
88
89    /// Takes the `FORMAT` option, lowercased. Accepts identifier or string
90    /// values.
91    pub(crate) fn take_format(&mut self) -> PgWireResult<Option<String>> {
92        let Some(value) = self.take_value("format") else {
93            return Ok(None);
94        };
95        match value {
96            CopyOptionValue::Ident(name) | CopyOptionValue::String(name) => {
97                Ok(Some(name.to_lowercase()))
98            }
99            other => Err(bad_copy_data(format!(
100                "FORMAT requires an identifier value, got {}",
101                other.kind()
102            ))),
103        }
104    }
105
106    /// Takes a string-valued option.
107    pub(crate) fn take_string(&mut self, key: &str) -> PgWireResult<Option<String>> {
108        let Some(value) = self.take_value(key) else {
109            return Ok(None);
110        };
111        match value {
112            CopyOptionValue::String(s) => Ok(Some(s)),
113            other => Err(bad_copy_data(format!(
114                "COPY option {} requires a string value, got {}",
115                key.to_uppercase(),
116                other.kind()
117            ))),
118        }
119    }
120
121    /// Takes a single one-byte character option. A tab is a legal (and
122    /// common) CSV delimiter; newlines and other control characters are
123    /// not, matching PostgreSQL's rules.
124    pub(crate) fn take_single_byte(&mut self, key: &str) -> PgWireResult<Option<u8>> {
125        let Some(value) = self.take_string(key)? else {
126            return Ok(None);
127        };
128        let mut chars = value.chars();
129        let (c, rest) = (chars.next(), chars.as_str());
130        match c {
131            Some('\t') if rest.is_empty() => Ok(Some(b'\t')),
132            Some(c) if rest.is_empty() && c.is_ascii() && !c.is_ascii_control() => {
133                Ok(Some(c as u8))
134            }
135            _ => Err(bad_copy_data(format!(
136                "COPY option {} must be a single one-byte character other than a newline or carriage return",
137                key.to_uppercase()
138            ))),
139        }
140    }
141
142    /// Takes a boolean option; a valueless option means `TRUE`.
143    pub(crate) fn take_flag(&mut self, key: &str) -> PgWireResult<Option<bool>> {
144        let Some(value) = self.take_value(key) else {
145            return Ok(None);
146        };
147        match value {
148            CopyOptionValue::Flag => Ok(Some(true)),
149            CopyOptionValue::Boolean(b) => Ok(Some(b)),
150            other => Err(bad_copy_data(format!(
151                "COPY option {} requires a boolean value, got {}",
152                key.to_uppercase(),
153                other.kind()
154            ))),
155        }
156    }
157
158    /// Rejects every option not consumed by the format, with one message
159    /// naming them all. Codecs call this after their `take_*` calls so a
160    /// typo cannot be silently ignored.
161    pub(crate) fn reject_leftovers(self, format: &str) -> PgWireResult<()> {
162        if self.entries.is_empty() {
163            return Ok(());
164        }
165        let unknown = self
166            .entries
167            .iter()
168            .map(|(key, _)| key.to_uppercase())
169            .collect::<Vec<_>>()
170            .join(", ");
171        Err(PgWireError::UserError(Box::new(
172            PgErrorCode::Ec0A000.to_err_info(format!(
173                "COPY option(s) {unknown} are not supported by format {format}"
174            )),
175        )))
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    fn new_set(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet {
184        let mut set = CopyOptionSet::default();
185        for (key, value) in entries {
186            set.insert((*key).to_string(), value.clone()).unwrap();
187        }
188        set
189    }
190
191    #[test]
192    fn test_insert_rejects_duplicates() {
193        let mut set = CopyOptionSet::default();
194        set.insert("format".to_string(), CopyOptionValue::Ident("csv".into()))
195            .unwrap();
196        let err = set
197            .insert("format".to_string(), CopyOptionValue::Ident("text".into()))
198            .unwrap_err();
199        assert!(err.to_string().contains("redundant FORMAT"));
200    }
201
202    #[test]
203    fn test_take_single_byte_tab_allowed() {
204        // A literal tab is a legal delimiter (PostgreSQL permits it).
205        let mut set = new_set(&[("delimiter", CopyOptionValue::String("\t".into()))]);
206        assert_eq!(set.take_single_byte("DELIMITER").unwrap(), Some(b'\t'));
207
208        // Newlines and other control characters stay rejected.
209        for bad in ["\n", "\r", "\0"] {
210            let mut set = new_set(&[("delimiter", CopyOptionValue::String(bad.into()))]);
211            assert!(
212                set.take_single_byte("DELIMITER").is_err(),
213                "{bad:?} must be rejected"
214            );
215        }
216    }
217
218    #[test]
219    fn test_take_string_and_single_byte() {
220        let mut set = new_set(&[
221            ("null", CopyOptionValue::String("nil".into())),
222            ("delimiter", CopyOptionValue::String("|".into())),
223        ]);
224        assert_eq!(set.take_string("null").unwrap().as_deref(), Some("nil"));
225        assert_eq!(set.take_single_byte("delimiter").unwrap(), Some(b'|'));
226
227        let mut set = new_set(&[("delimiter", CopyOptionValue::Number("42".into()))]);
228        let err = set.take_single_byte("delimiter").unwrap_err();
229        assert!(err.to_string().contains("requires a string value"));
230
231        let mut set = new_set(&[("delimiter", CopyOptionValue::String("ab".into()))]);
232        let err = set.take_single_byte("delimiter").unwrap_err();
233        assert!(err.to_string().contains("single one-byte"));
234    }
235
236    #[test]
237    fn test_take_flag() {
238        let mut set = new_set(&[
239            ("header", CopyOptionValue::Flag),
240            ("frozen", CopyOptionValue::Boolean(false)),
241        ]);
242        assert_eq!(set.take_flag("header").unwrap(), Some(true));
243        assert_eq!(set.take_flag("frozen").unwrap(), Some(false));
244        assert_eq!(set.take_flag("absent").unwrap(), None);
245    }
246
247    #[test]
248    fn test_take_format_normalizes() {
249        let mut set = new_set(&[("format", CopyOptionValue::Ident("CSV".into()))]);
250        assert_eq!(set.take_format().unwrap().as_deref(), Some("csv"));
251    }
252
253    #[test]
254    fn test_reject_leftovers_lists_all() {
255        let mut set = new_set(&[
256            ("precision", CopyOptionValue::String("ns".into())),
257            ("weird", CopyOptionValue::Flag),
258        ]);
259        set.take_string("precision").unwrap();
260        let err = set.reject_leftovers("text").unwrap_err();
261        assert!(
262            err.to_string().contains("WEIRD"),
263            "leftover option must be listed: {err}"
264        );
265        assert!(
266            err.to_string().contains("not supported by format text"),
267            "{err}"
268        );
269        // The consumed option is not listed.
270        assert!(!err.to_string().contains("PRECISION"));
271    }
272
273    #[test]
274    fn test_raw_value_escape_hatch() {
275        let mut set = new_set(&[("precision", CopyOptionValue::Ident("ns".into()))]);
276        assert_eq!(
277            set.take_value("precision"),
278            Some(CopyOptionValue::Ident("ns".into()))
279        );
280    }
281}
servers/postgres/copy_in/options.rs

servers/postgres/copy_in/
options.rs
