mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-10-11 22:42:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Parse Snowflake COPY INTO <location> (#1669)

Browse source
This commit is contained in:
Yoav Cohen 2025-02-05 20:23:27 +01:00 • committed by GitHub
parent 486b29ffab
commit 751dc5afce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 381 additions and 185 deletions
Download patch file Download diff file Expand all files Collapse all files

102
src/ast/mod.rs
View file

@ -2498,24 +2498,30 @@ pub enum Statement {
values: Vec<Option<String>>,
},
/// ```sql
/// COPY INTO
/// COPY INTO <table> | <location>
/// ```
/// See <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
/// See:
/// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
/// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
///
/// Copy Into syntax available for Snowflake is different than the one implemented in
/// Postgres. Although they share common prefix, it is reasonable to implement them
/// in different enums. This can be refactored later once custom dialects
/// are allowed to have custom Statements.
CopyIntoSnowflake {
kind: CopyIntoSnowflakeKind,
into: ObjectName,
from_stage: ObjectName,
from_stage_alias: Option<Ident>,
from_obj: Option<ObjectName>,
from_obj_alias: Option<Ident>,
stage_params: StageParamsObject,
from_transformations: Option<Vec<StageLoadSelectItem>>,
from_query: Option<Box<Query>>,
files: Option<Vec<String>>,
pattern: Option<String>,
file_format: DataLoadingOptions,
copy_options: DataLoadingOptions,
validation_mode: Option<String>,
partition: Option<Box<Expr>>,
},
/// ```sql
/// CLOSE
@ -5048,60 +5054,69 @@ impl fmt::Display for Statement {
Ok(())
}
Statement::CopyIntoSnowflake {
kind,
into,
from_stage,
from_stage_alias,
from_obj,
from_obj_alias,
stage_params,
from_transformations,
from_query,
files,
pattern,
file_format,
copy_options,
validation_mode,
partition,
} => {
write!(f, "COPY INTO {}", into)?;
if from_transformations.is_none() {
// Standard data load
write!(f, " FROM {}{}", from_stage, stage_params)?;
if from_stage_alias.as_ref().is_some() {
write!(f, " AS {}", from_stage_alias.as_ref().unwrap())?;
}
} else {
if let Some(from_transformations) = from_transformations {
// Data load with transformation
write!(
f,
" FROM (SELECT {} FROM {}{}",
display_separated(from_transformations.as_ref().unwrap(), ", "),
from_stage,
stage_params,
)?;
if from_stage_alias.as_ref().is_some() {
write!(f, " AS {}", from_stage_alias.as_ref().unwrap())?;
if let Some(from_stage) = from_obj {
write!(
f,
" FROM (SELECT {} FROM {}{}",
display_separated(from_transformations, ", "),
from_stage,
stage_params
)?;
}
if let Some(from_obj_alias) = from_obj_alias {
write!(f, " AS {}", from_obj_alias)?;
}
write!(f, ")")?;
} else if let Some(from_obj) = from_obj {
// Standard data load
write!(f, " FROM {}{}", from_obj, stage_params)?;
if let Some(from_obj_alias) = from_obj_alias {
write!(f, " AS {from_obj_alias}")?;
}
} else if let Some(from_query) = from_query {
// Data unload from query
write!(f, " FROM ({from_query})")?;
}
if files.is_some() {
write!(
f,
" FILES = ('{}')",
display_separated(files.as_ref().unwrap(), "', '")
)?;
if let Some(files) = files {
write!(f, " FILES = ('{}')", display_separated(files, "', '"))?;
}
if pattern.is_some() {
write!(f, " PATTERN = '{}'", pattern.as_ref().unwrap())?;
if let Some(pattern) = pattern {
write!(f, " PATTERN = '{}'", pattern)?;
}
if let Some(partition) = partition {
write!(f, " PARTITION BY {partition}")?;
}
if !file_format.options.is_empty() {
write!(f, " FILE_FORMAT=({})", file_format)?;
}
if !copy_options.options.is_empty() {
write!(f, " COPY_OPTIONS=({})", copy_options)?;
match kind {
CopyIntoSnowflakeKind::Table => {
write!(f, " COPY_OPTIONS=({})", copy_options)?
}
CopyIntoSnowflakeKind::Location => write!(f, " {copy_options}")?,
}
}
if validation_mode.is_some() {
write!(
f,
" VALIDATION_MODE = {}",
validation_mode.as_ref().unwrap()
)?;
if let Some(validation_mode) = validation_mode {
write!(f, " VALIDATION_MODE = {}", validation_mode)?;
}
Ok(())
}
@ -8543,6 +8558,19 @@ impl Display for StorageSerializationPolicy {
}
}
/// Variants of the Snowflake `COPY INTO` statement
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum CopyIntoSnowflakeKind {
/// Loads data from files to a table
/// See: <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
Table,
/// Unloads data from a table or query to external files
/// See: <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
Location,
}
#[cfg(test)]
mod tests {
use super::*;