Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added ability to parse STRUCT and MAP fields as well as nested arrays #966

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions src/ast/data_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ use serde::{Deserialize, Serialize};
#[cfg(feature = "visitor")]
use sqlparser_derive::{Visit, VisitMut};

use crate::ast::ObjectName;
use crate::ast::{display_comma_separated, ObjectName};
use crate::ast::ddl::StructField;

use super::value::escape_single_quote_string;

Expand Down Expand Up @@ -198,10 +199,15 @@ pub enum DataType {
Custom(ObjectName, Vec<String>),
/// Arrays
Array(Option<Box<DataType>>),
BracketArray(Option<Box<DataType>>),
/// Enums
Enum(Vec<String>),
/// Set
Set(Vec<String>),
/// Map
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you please provide a link to what SQL dialect supports this syntax?

Map(Box<DataType>, Box<DataType>),
/// Struct
Struct(Vec<StructField>)
}

impl fmt::Display for DataType {
Expand Down Expand Up @@ -320,11 +326,17 @@ impl fmt::Display for DataType {
DataType::Bytea => write!(f, "BYTEA"),
DataType::Array(ty) => {
if let Some(t) = &ty {
write!(f, "{t}[]")
write!(f, "ARRAY<{t}>")
} else {
write!(f, "ARRAY")
}
}
DataType::BracketArray(ty) => {
if let Some(t) = &ty {
write!(f, "{t}[]")?
}
Ok(())
}
DataType::Custom(ty, modifiers) => {
if modifiers.is_empty() {
write!(f, "{ty}")
Expand Down Expand Up @@ -352,6 +364,14 @@ impl fmt::Display for DataType {
}
write!(f, ")")
}
DataType::Map(key, value) => {
write!(f, "MAP<{}>", display_comma_separated(&[key, value]))
}
DataType::Struct(vals) => {
write!(f, "STRUCT<")?;
write!(f, "{}", display_comma_separated(vals))?;
write!(f, ">")
}
}
}
}
Expand Down
19 changes: 19 additions & 0 deletions src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -678,6 +678,25 @@ impl fmt::Display for ColumnOption {
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct StructField {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you please add docstrings that explain what this struct is for (with an example SQL snippet)

pub(crate) name: Ident,
pub(crate) data_type: DataType,
pub(crate) options: Option<ColumnOption>,
}

impl fmt::Display for StructField {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}: {}", self.name, self.data_type)?;
if let Some(option) = self.options.as_ref() {
write!(f, "{option}")?;
}
Ok(())
}
}

/// `GeneratedAs`s are modifiers that follow a column option in a `generated`.
/// 'ExpStored' is PostgreSQL specific
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
Expand Down
2 changes: 1 addition & 1 deletion src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ pub use self::data_type::{
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
pub use self::ddl::{
AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, ProcedureParam, ReferentialAction,
ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, ProcedureParam, ReferentialAction, StructField,
TableConstraint, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
};
pub use self::operator::{BinaryOperator, UnaryOperator};
Expand Down
2 changes: 2 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,7 @@ define_keywords!(
LOWER,
MACRO,
MANAGEDLOCATION,
MAP,
MATCH,
MATCHED,
MATERIALIZED,
Expand Down Expand Up @@ -575,6 +576,7 @@ define_keywords!(
STORED,
STRICT,
STRING,
STRUCT,
SUBMULTISET,
SUBSTRING,
SUBSTRING_REGEX,
Expand Down
94 changes: 84 additions & 10 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use alloc::{
vec::Vec,
};
use core::fmt;
use std::ops::Rem;

use log::debug;

Expand Down Expand Up @@ -114,6 +115,7 @@ mod recursion {
Self { remaining_depth }
}
}

impl Drop for DepthGuard {
fn drop(&mut self) {
self.remaining_depth.fetch_add(1, Ordering::SeqCst);
Expand Down Expand Up @@ -257,6 +259,7 @@ pub struct Parser<'a> {
options: ParserOptions,
/// ensure the stack does not overflow by limiting recursion depth
recursion_counter: RecursionCounter,
max_depth: usize,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please document what this field is foe.

Also, could you please explain why you didn't use the counter in RecusionCounter?

}

impl<'a> Parser<'a> {
Expand All @@ -282,6 +285,7 @@ impl<'a> Parser<'a> {
dialect,
recursion_counter: RecursionCounter::new(DEFAULT_REMAINING_DEPTH),
options: ParserOptions::default(),
max_depth: 1,
}
}

Expand Down Expand Up @@ -2181,7 +2185,7 @@ impl<'a> Parser<'a> {
token => {
return token
.cloned()
.unwrap_or_else(|| TokenWithLocation::wrap(Token::EOF))
.unwrap_or_else(|| TokenWithLocation::wrap(Token::EOF));
}
}
}
Expand Down Expand Up @@ -4648,6 +4652,13 @@ impl<'a> Parser<'a> {

/// Parse a SQL datatype (in the context of a CREATE TABLE statement for example)
pub fn parse_data_type(&mut self) -> Result<DataType, ParserError> {
self.parse_data_type_with_depth(1)
}

pub fn parse_data_type_with_depth(&mut self, depth: usize) -> Result<DataType, ParserError> {
if depth > self.max_depth {
self.max_depth = depth - 1;
}
let next_token = self.next_token();
let mut data = match next_token.token {
Token::Word(w) => match w.keyword {
Expand Down Expand Up @@ -4837,10 +4848,57 @@ impl<'a> Parser<'a> {
// that ends with > will fail due to "C++" problem - >> is parsed as
// Token::ShiftRight
self.expect_token(&Token::Lt)?;
let inside_type = self.parse_data_type()?;

let inside_type = self.parse_data_type_with_depth(depth + 1)?;
dbg!(depth, self.max_depth);

if depth <= 1 {
dbg!("First Level");
if (depth == 1 && self.max_depth == depth)
|| (self.peek_previous_token()? == &Token::ShiftRight
&& self.max_depth.rem(2) != 0)
{
self.expect_token(&Token::Gt)?;
}
} else if depth.rem(2) == 0 && depth != self.max_depth {
} else {
dbg!("Else Level");
self.expect_token(&Token::ShiftRight)?;
}

if dialect_of!(self is PostgreSqlDialect) {
Ok(DataType::BracketArray(Some(Box::new(inside_type))))
} else {
Ok(DataType::Array(Some(Box::new(inside_type))))
}
}
}
Keyword::MAP => {
self.expect_token(&Token::Lt)?;
let key = self.parse_data_type_with_depth(depth + 1)?;
let tok = self.consume_token(&Token::Comma);
debug!("Tok: {tok}");
let value = self.parse_data_type_with_depth(depth + 1)?;
let tok = self.peek_token().token;
debug!("Next Tok: {tok}");
if tok == Token::ShiftRight {
self.expect_token(&Token::ShiftRight)?;
} else if tok == Token::Gt {
self.expect_token(&Token::Gt)?;
Ok(DataType::Array(Some(Box::new(inside_type))))
}
Ok(DataType::Map(Box::new(key), Box::new(value)))
}
Keyword::STRUCT => {
self.expect_token(&Token::Lt)?;
let fields = self.parse_comma_separated(Parser::parse_struct_fields)?;
let tok = self.peek_token().token;
debug!("Next Tok: {tok}");
if tok == Token::ShiftRight {
self.expect_token(&Token::ShiftRight)?;
} else if tok == Token::Gt {
self.expect_token(&Token::Gt)?;
}
Ok(DataType::Struct(fields))
}
_ => {
self.prev_token();
Expand All @@ -4859,11 +4917,27 @@ impl<'a> Parser<'a> {
// Keyword::ARRAY syntax from above
while self.consume_token(&Token::LBracket) {
self.expect_token(&Token::RBracket)?;
data = DataType::Array(Some(Box::new(data)))
data = DataType::BracketArray(Some(Box::new(data)))
}
Ok(data)
}

pub fn peek_previous_token(&mut self) -> Result<&TokenWithLocation, ParserError> {
Ok(&self.tokens[self.index - 1])
}

pub fn parse_struct_fields(&mut self) -> Result<StructField, ParserError> {
let name = self.parse_identifier()?;
self.expect_token(&Token::Colon)?;
let data_type = self.parse_data_type()?;
let options = self.parse_optional_column_option()?;
Ok(StructField {
name,
data_type,
options,
})
}

pub fn parse_string_values(&mut self) -> Result<Vec<String>, ParserError> {
self.expect_token(&Token::LParen)?;
let mut values = Vec::new();
Expand Down Expand Up @@ -5028,12 +5102,12 @@ impl<'a> Parser<'a> {
Token::EOF => {
return Err(ParserError::ParserError(
"Empty input when parsing identifier".to_string(),
))?
))?;
}
token => {
return Err(ParserError::ParserError(format!(
"Unexpected token in identifier: {token}"
)))?
)))?;
}
};

Expand All @@ -5046,19 +5120,19 @@ impl<'a> Parser<'a> {
Token::EOF => {
return Err(ParserError::ParserError(
"Trailing period in identifier".to_string(),
))?
))?;
}
token => {
return Err(ParserError::ParserError(format!(
"Unexpected token following period in identifier: {token}"
)))?
)))?;
}
},
Token::EOF => break,
token => {
return Err(ParserError::ParserError(format!(
"Unexpected token in identifier: {token}"
)))?
)))?;
}
}
}
Expand Down Expand Up @@ -6031,7 +6105,7 @@ impl<'a> Parser<'a> {
_ => {
return Err(ParserError::ParserError(format!(
"expected OUTER, SEMI, ANTI or JOIN after {kw:?}"
)))
)));
}
}
}
Expand Down
36 changes: 28 additions & 8 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2416,17 +2416,41 @@ fn parse_create_table() {
.contains("Expected constraint details after CONSTRAINT <name>"));
}

#[test]
fn parse_double_greater_than_array() {
let supported_dialects = TestedDialects {
dialects: vec![Box::new(HiveDialect {})],
options: None,
};
let levels = &[
"CREATE TABLE t (a ARRAY<INT>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<INT>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<INT>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>>>>>, n INT)",
"CREATE TABLE t (a ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<ARRAY<INT>>>>>>>>>>, n INT)",
];
for q in levels {
let statements = supported_dialects.parse_sql_statements(q).unwrap();
println!("{}", statements[0]);
}
}

#[test]
fn parse_create_table_hive_array() {
// Parsing [] type arrays does not work in MsSql since [ is used in is_delimited_identifier_start
let dialects = TestedDialects {
dialects: vec![Box::new(PostgreSqlDialect {}), Box::new(HiveDialect {})],
dialects: vec![Box::new(HiveDialect {})],
options: None,
};
let sql = "CREATE TABLE IF NOT EXISTS something (name int, val array<int>)";
match dialects.one_statement_parses_to(
sql,
"CREATE TABLE IF NOT EXISTS something (name INT, val INT[])",
"CREATE TABLE IF NOT EXISTS something (name INT, val ARRAY<INT>)",
) {
Statement::CreateTable {
if_not_exists,
Expand Down Expand Up @@ -2457,13 +2481,9 @@ fn parse_create_table_hive_array() {
_ => unreachable!(),
}

// SnowflakeDialect using array diffrent
// SnowflakeDialect using array different
let dialects = TestedDialects {
dialects: vec![
Box::new(PostgreSqlDialect {}),
Box::new(HiveDialect {}),
Box::new(MySqlDialect {}),
],
dialects: vec![Box::new(HiveDialect {}), Box::new(MySqlDialect {})],
options: None,
};
let sql = "CREATE TABLE IF NOT EXISTS something (name int, val array<int)";
Expand Down
4 changes: 2 additions & 2 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1741,7 +1741,7 @@ fn parse_array_index_expr() {
})],
named: true,
})),
data_type: DataType::Array(Some(Box::new(DataType::Array(Some(Box::new(
data_type: DataType::BracketArray(Some(Box::new(DataType::BracketArray(Some(Box::new(
DataType::Int(None)
))))))
}))),
Expand All @@ -1753,7 +1753,7 @@ fn parse_array_index_expr() {
let sql = "SELECT ARRAY[]";
let select = pg_and_generic().verified_only_select(sql);
assert_eq!(
&Expr::Array(sqlparser::ast::Array {
&Expr::Array(Array {
elem: vec![],
named: true
}),
Expand Down