Skip to content

Commit

Permalink
feat: support INSERT INTO [TABLE] FUNCTION of Clickhouse (#1633)
Browse files Browse the repository at this point in the history
Co-authored-by: Kermit <[email protected]>
Co-authored-by: Ifeanyi Ubah <[email protected]>
  • Loading branch information
3 people authored Jan 10, 2025
1 parent c54ba4d commit b09514e
Show file tree
Hide file tree
Showing 11 changed files with 147 additions and 53 deletions.
13 changes: 6 additions & 7 deletions src/ast/dml.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use super::{
FileFormat, FromTable, HiveDistributionStyle, HiveFormat, HiveIOFormat, HiveRowFormat, Ident,
InsertAliases, MysqlInsertPriority, ObjectName, OnCommit, OnInsert, OneOrManyWithParens,
OrderByExpr, Query, RowAccessPolicy, SelectItem, SqlOption, SqliteOnConflict, TableEngine,
TableWithJoins, Tag, WrappedCollection,
TableObject, TableWithJoins, Tag, WrappedCollection,
};

/// CREATE INDEX statement.
Expand Down Expand Up @@ -470,8 +470,7 @@ pub struct Insert {
/// INTO - optional keyword
pub into: bool,
/// TABLE
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
pub table_name: ObjectName,
pub table: TableObject,
/// table_name as foo (for PostgreSQL)
pub table_alias: Option<Ident>,
/// COLUMNS
Expand All @@ -488,7 +487,7 @@ pub struct Insert {
/// Columns defined after PARTITION
pub after_columns: Vec<Ident>,
/// whether the insert has the table keyword (Hive)
pub table: bool,
pub has_table_keyword: bool,
pub on: Option<OnInsert>,
/// RETURNING
pub returning: Option<Vec<SelectItem>>,
Expand All @@ -503,9 +502,9 @@ pub struct Insert {
impl Display for Insert {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let table_name = if let Some(alias) = &self.table_alias {
format!("{0} AS {alias}", self.table_name)
format!("{0} AS {alias}", self.table)
} else {
self.table_name.to_string()
self.table.to_string()
};

if let Some(on_conflict) = self.or {
Expand All @@ -531,7 +530,7 @@ impl Display for Insert {
ignore = if self.ignore { " IGNORE" } else { "" },
over = if self.overwrite { " OVERWRITE" } else { "" },
int = if self.into { " INTO" } else { "" },
tbl = if self.table { " TABLE" } else { "" },
tbl = if self.has_table_keyword { " TABLE" } else { "" },
)?;
}
if !self.columns.is_empty() {
Expand Down
30 changes: 30 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7907,6 +7907,36 @@ impl fmt::Display for RenameTable {
}
}

/// Represents the referenced table in an `INSERT INTO` statement
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum TableObject {
/// Table specified by name.
/// Example:
/// ```sql
/// INSERT INTO my_table
/// ```
TableName(#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] ObjectName),

/// Table specified as a function.
/// Example:
/// ```sql
/// INSERT INTO TABLE FUNCTION remote('localhost', default.simple_table)
/// ```
/// [Clickhouse](https://clickhouse.com/docs/en/sql-reference/table-functions)
TableFunction(Function),
}

impl fmt::Display for TableObject {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Self::TableName(table_name) => write!(f, "{table_name}"),
Self::TableFunction(func) => write!(f, "FUNCTION {}", func),
}
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
22 changes: 17 additions & 5 deletions src/ast/spans.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,9 @@ use super::{
OrderBy, OrderByExpr, Partition, PivotValueSource, ProjectionSelect, Query, ReferentialAction,
RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem,
SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef,
TableConstraint, TableFactor, TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use,
Value, Values, ViewColumnDef, WildcardAdditionalOptions, With, WithFill,
TableConstraint, TableFactor, TableObject, TableOptionsClustered, TableWithJoins,
UpdateTableFromKind, Use, Value, Values, ViewColumnDef, WildcardAdditionalOptions, With,
WithFill,
};

/// Given an iterator of spans, return the [Span::union] of all spans.
Expand Down Expand Up @@ -1141,14 +1142,14 @@ impl Spanned for Insert {
or: _, // enum, sqlite specific
ignore: _, // bool
into: _, // bool
table_name,
table,
table_alias,
columns,
overwrite: _, // bool
source,
partitioned,
after_columns,
table: _, // bool
has_table_keyword: _, // bool
on,
returning,
replace_into: _, // bool
Expand All @@ -1158,7 +1159,7 @@ impl Spanned for Insert {
} = self;

union_spans(
core::iter::once(table_name.span())
core::iter::once(table.span())
.chain(table_alias.as_ref().map(|i| i.span))
.chain(columns.iter().map(|i| i.span))
.chain(source.as_ref().map(|q| q.span()))
Expand Down Expand Up @@ -2121,6 +2122,17 @@ impl Spanned for UpdateTableFromKind {
}
}

impl Spanned for TableObject {
fn span(&self) -> Span {
match self {
TableObject::TableName(ObjectName(segments)) => {
union_spans(segments.iter().map(|i| i.span))
}
TableObject::TableFunction(func) => func.span(),
}
}
}

#[cfg(test)]
pub mod tests {
use crate::dialect::{Dialect, GenericDialect, SnowflakeDialect};
Expand Down
4 changes: 4 additions & 0 deletions src/dialect/clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,4 +50,8 @@ impl Dialect for ClickHouseDialect {
fn supports_limit_comma(&self) -> bool {
true
}

fn supports_insert_table_function(&self) -> bool {
true
}
}
5 changes: 5 additions & 0 deletions src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -792,6 +792,11 @@ pub trait Dialect: Debug + Any {
fn supports_insert_set(&self) -> bool {
false
}

/// Does the dialect support table function in insertion?
fn supports_insert_table_function(&self) -> bool {
false
}
}

/// This represents the operators for which precedence must be defined
Expand Down
18 changes: 15 additions & 3 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8937,6 +8937,18 @@ impl<'a> Parser<'a> {
}
}

/// Parse a table object for insetion
/// e.g. `some_database.some_table` or `FUNCTION some_table_func(...)`
pub fn parse_table_object(&mut self) -> Result<TableObject, ParserError> {
if self.dialect.supports_insert_table_function() && self.parse_keyword(Keyword::FUNCTION) {
let fn_name = self.parse_object_name(false)?;
self.parse_function_call(fn_name)
.map(TableObject::TableFunction)
} else {
self.parse_object_name(false).map(TableObject::TableName)
}
}

/// Parse a possibly qualified, possibly quoted identifier, optionally allowing for wildcards,
/// e.g. *, *.*, `foo`.*, or "foo"."bar"
fn parse_object_name_with_wildcards(
Expand Down Expand Up @@ -12010,7 +12022,7 @@ impl<'a> Parser<'a> {
} else {
// Hive lets you put table here regardless
let table = self.parse_keyword(Keyword::TABLE);
let table_name = self.parse_object_name(false)?;
let table_object = self.parse_table_object()?;

let table_alias =
if dialect_of!(self is PostgreSqlDialect) && self.parse_keyword(Keyword::AS) {
Expand Down Expand Up @@ -12118,7 +12130,7 @@ impl<'a> Parser<'a> {

Ok(Statement::Insert(Insert {
or,
table_name,
table: table_object,
table_alias,
ignore,
into,
Expand All @@ -12128,7 +12140,7 @@ impl<'a> Parser<'a> {
after_columns,
source,
assignments,
table,
has_table_keyword: table,
on,
returning,
replace_into,
Expand Down
1 change: 0 additions & 1 deletion src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,6 @@ impl TestedDialects {
pub fn one_statement_parses_to(&self, sql: &str, canonical: &str) -> Statement {
let mut statements = self.parse_sql_statements(sql).expect(sql);
assert_eq!(statements.len(), 1);

if !canonical.is_empty() && sql != canonical {
assert_eq!(self.parse_sql_statements(canonical).unwrap(), statements);
}
Expand Down
6 changes: 6 additions & 0 deletions tests/sqlparser_clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,12 @@ fn parse_create_table() {
);
}

#[test]
fn parse_insert_into_function() {
clickhouse().verified_stmt(r#"INSERT INTO TABLE FUNCTION remote('localhost', default.simple_table) VALUES (100, 'inserted via remote()')"#);
clickhouse().verified_stmt(r#"INSERT INTO FUNCTION remote('localhost', default.simple_table) VALUES (100, 'inserted via remote()')"#);
}

#[test]
fn parse_alter_table_attach_and_detach_partition() {
for operation in &["ATTACH", "DETACH"] {
Expand Down
23 changes: 16 additions & 7 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ fn parse_insert_values() {
) {
match verified_stmt(sql) {
Statement::Insert(Insert {
table_name,
table: table_name,
columns,
source: Some(source),
..
Expand Down Expand Up @@ -149,7 +149,7 @@ fn parse_insert_default_values() {
partitioned,
returning,
source,
table_name,
table: table_name,
..
}) => {
assert_eq!(columns, vec![]);
Expand All @@ -158,7 +158,10 @@ fn parse_insert_default_values() {
assert_eq!(partitioned, None);
assert_eq!(returning, None);
assert_eq!(source, None);
assert_eq!(table_name, ObjectName(vec!["test_table".into()]));
assert_eq!(
table_name,
TableObject::TableName(ObjectName(vec!["test_table".into()]))
);
}
_ => unreachable!(),
}
Expand All @@ -174,7 +177,7 @@ fn parse_insert_default_values() {
partitioned,
returning,
source,
table_name,
table: table_name,
..
}) => {
assert_eq!(after_columns, vec![]);
Expand All @@ -183,7 +186,10 @@ fn parse_insert_default_values() {
assert_eq!(partitioned, None);
assert!(returning.is_some());
assert_eq!(source, None);
assert_eq!(table_name, ObjectName(vec!["test_table".into()]));
assert_eq!(
table_name,
TableObject::TableName(ObjectName(vec!["test_table".into()]))
);
}
_ => unreachable!(),
}
Expand All @@ -199,7 +205,7 @@ fn parse_insert_default_values() {
partitioned,
returning,
source,
table_name,
table: table_name,
..
}) => {
assert_eq!(after_columns, vec![]);
Expand All @@ -208,7 +214,10 @@ fn parse_insert_default_values() {
assert_eq!(partitioned, None);
assert_eq!(returning, None);
assert_eq!(source, None);
assert_eq!(table_name, ObjectName(vec!["test_table".into()]));
assert_eq!(
table_name,
TableObject::TableName(ObjectName(vec!["test_table".into()]))
);
}
_ => unreachable!(),
}
Expand Down
Loading

0 comments on commit b09514e

Please sign in to comment.