sqlparser/ast/helpers/
stmt_data_loading.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! AST types specific to loading and unloading syntax, like one available in Snowflake which
19//! contains: STAGE ddl operations, PUT upload or COPY INTO
20//! See [this page](https://docs.snowflake.com/en/sql-reference/commands-data-loading) for more details.
21
22#[cfg(not(feature = "std"))]
23use alloc::string::String;
24use core::fmt;
25
26#[cfg(feature = "serde")]
27use serde::{Deserialize, Serialize};
28
29use crate::ast::helpers::key_value_options::KeyValueOptions;
30use crate::ast::{Ident, ObjectName, SelectItem};
31#[cfg(feature = "visitor")]
32use yachtsql_sqlparser_derive::{Visit, VisitMut};
33
34#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
35#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
36#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
37pub struct StageParamsObject {
38    pub url: Option<String>,
39    pub encryption: KeyValueOptions,
40    pub endpoint: Option<String>,
41    pub storage_integration: Option<String>,
42    pub credentials: KeyValueOptions,
43}
44
45/// This enum enables support for both standard SQL select item expressions
46/// and Snowflake-specific ones for data loading.
47#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
48#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
49#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
50pub enum StageLoadSelectItemKind {
51    SelectItem(SelectItem),
52    StageLoadSelectItem(StageLoadSelectItem),
53}
54
55impl fmt::Display for StageLoadSelectItemKind {
56    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57        match &self {
58            StageLoadSelectItemKind::SelectItem(item) => write!(f, "{item}"),
59            StageLoadSelectItemKind::StageLoadSelectItem(item) => write!(f, "{item}"),
60        }
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
65#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
66#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
67pub struct StageLoadSelectItem {
68    pub alias: Option<Ident>,
69    pub file_col_num: i32,
70    pub element: Option<Ident>,
71    pub item_as: Option<Ident>,
72}
73
74impl fmt::Display for StageParamsObject {
75    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
76        let url = &self.url.as_ref();
77        let storage_integration = &self.storage_integration.as_ref();
78        let endpoint = &self.endpoint.as_ref();
79
80        if url.is_some() {
81            write!(f, " URL='{}'", url.unwrap())?;
82        }
83        if storage_integration.is_some() {
84            write!(f, " STORAGE_INTEGRATION={}", storage_integration.unwrap())?;
85        }
86        if endpoint.is_some() {
87            write!(f, " ENDPOINT='{}'", endpoint.unwrap())?;
88        }
89        if !self.credentials.options.is_empty() {
90            write!(f, " CREDENTIALS=({})", self.credentials)?;
91        }
92        if !self.encryption.options.is_empty() {
93            write!(f, " ENCRYPTION=({})", self.encryption)?;
94        }
95
96        Ok(())
97    }
98}
99
100impl fmt::Display for StageLoadSelectItem {
101    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
102        if self.alias.is_some() {
103            write!(f, "{}.", self.alias.as_ref().unwrap())?;
104        }
105        write!(f, "${}", self.file_col_num)?;
106        if self.element.is_some() {
107            write!(f, ":{}", self.element.as_ref().unwrap())?;
108        }
109        if self.item_as.is_some() {
110            write!(f, " AS {}", self.item_as.as_ref().unwrap())?;
111        }
112        Ok(())
113    }
114}
115
116#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
117#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
118#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
119pub struct FileStagingCommand {
120    #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
121    pub stage: ObjectName,
122    pub pattern: Option<String>,
123}
124
125impl fmt::Display for FileStagingCommand {
126    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
127        write!(f, "{}", self.stage)?;
128        if let Some(pattern) = self.pattern.as_ref() {
129            write!(f, " PATTERN='{pattern}'")?;
130        }
131        Ok(())
132    }
133}