feat: support alter add multiple columns (#5262)

* feat: support alter add multiple columns

* fix: address review

* chore: add column format
This commit is contained in:
Niwaka
2025-01-10 15:14:17 +09:00
committed by GitHub
parent 24ea9cf215
commit 50583815de
9 changed files with 204 additions and 132 deletions

View File

@@ -494,20 +494,24 @@ pub(crate) fn to_alter_table_expr(
}
.fail();
}
AlterTableOperation::AddColumn {
column_def,
location,
} => AlterTableKind::AddColumns(AddColumns {
add_columns: vec![AddColumn {
column_def: Some(
sql_column_def_to_grpc_column_def(&column_def, Some(&query_ctx.timezone()))
.map_err(BoxedError::new)
.context(ExternalSnafu)?,
),
location: location.as_ref().map(From::from),
// TODO(yingwen): We don't support `IF NOT EXISTS` for `ADD COLUMN` yet.
add_if_not_exists: false,
}],
AlterTableOperation::AddColumns { add_columns } => AlterTableKind::AddColumns(AddColumns {
add_columns: add_columns
.into_iter()
.map(|add_column| {
let column_def = sql_column_def_to_grpc_column_def(
&add_column.column_def,
Some(&query_ctx.timezone()),
)
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
Ok(AddColumn {
column_def: Some(column_def),
location: add_column.location.as_ref().map(From::from),
// TODO(yingwen): We don't support `IF NOT EXISTS` for `ADD COLUMN` yet.
add_if_not_exists: false,
})
})
.collect::<Result<Vec<AddColumn>>>()?,
}),
AlterTableOperation::ModifyColumnType {
column_name,

View File

@@ -283,8 +283,8 @@ impl ParserContext<'_> {
/// but with constant argument "false".
/// Because the argument is always "false" for us (it's introduced by BigQuery),
/// we don't want to write it again and again.
pub(crate) fn parse_identifier(&mut self) -> std::result::Result<Ident, ParserError> {
self.parser.parse_identifier(false)
pub(crate) fn parse_identifier(parser: &mut Parser) -> std::result::Result<Ident, ParserError> {
parser.parse_identifier(false)
}
}

View File

@@ -26,7 +26,8 @@ use crate::error::{self, InvalidColumnOptionSnafu, Result, SetFulltextOptionSnaf
use crate::parser::ParserContext;
use crate::parsers::utils::validate_column_fulltext_create_option;
use crate::statements::alter::{
AlterDatabase, AlterDatabaseOperation, AlterTable, AlterTableOperation, KeyValueOption,
AddColumn, AlterDatabase, AlterDatabaseOperation, AlterTable, AlterTableOperation,
KeyValueOption,
};
use crate::statements::statement::Statement;
use crate::util::parse_option_string;
@@ -120,7 +121,8 @@ impl ParserContext<'_> {
.expect_keyword(Keyword::COLUMN)
.context(error::SyntaxSnafu)?;
let name = Self::canonicalize_identifier(
self.parse_identifier().context(error::SyntaxSnafu)?,
Self::parse_identifier(&mut self.parser)
.context(error::SyntaxSnafu)?,
);
AlterTableOperation::DropColumn { name }
}
@@ -185,30 +187,12 @@ impl ParserContext<'_> {
{
Ok(AlterTableOperation::AddConstraint(constraint))
} else {
let _ = self.parser.parse_keyword(Keyword::COLUMN);
let mut column_def = self.parser.parse_column_def().context(error::SyntaxSnafu)?;
column_def.name = Self::canonicalize_identifier(column_def.name);
let location = if self.parser.parse_keyword(Keyword::FIRST) {
Some(AddColumnLocation::First)
} else if let Token::Word(word) = self.parser.peek_token().token {
if word.value.eq_ignore_ascii_case("AFTER") {
let _ = self.parser.next_token();
let name = Self::canonicalize_identifier(
self.parse_identifier().context(error::SyntaxSnafu)?,
);
Some(AddColumnLocation::After {
column_name: name.value,
})
} else {
None
}
} else {
None
};
Ok(AlterTableOperation::AddColumn {
column_def,
location,
})
self.parser.prev_token();
let add_columns = self
.parser
.parse_comma_separated(parse_add_columns)
.context(error::SyntaxSnafu)?;
Ok(AlterTableOperation::AddColumns { add_columns })
}
}
@@ -304,6 +288,33 @@ fn parse_string_options(parser: &mut Parser) -> std::result::Result<(String, Str
Ok((name, value))
}
fn parse_add_columns(parser: &mut Parser) -> std::result::Result<AddColumn, ParserError> {
parser.expect_keyword(Keyword::ADD)?;
let _ = parser.parse_keyword(Keyword::COLUMN);
let mut column_def = parser.parse_column_def()?;
column_def.name = ParserContext::canonicalize_identifier(column_def.name);
let location = if parser.parse_keyword(Keyword::FIRST) {
Some(AddColumnLocation::First)
} else if let Token::Word(word) = parser.peek_token().token {
if word.value.eq_ignore_ascii_case("AFTER") {
let _ = parser.next_token();
let name =
ParserContext::canonicalize_identifier(ParserContext::parse_identifier(parser)?);
Some(AddColumnLocation::After {
column_name: name.value,
})
} else {
None
}
} else {
None
};
Ok(AddColumn {
column_def,
location,
})
}
/// Parses a comma separated list of string literals.
fn parse_string_option_names(parser: &mut Parser) -> std::result::Result<String, ParserError> {
parser.parse_literal_string()
@@ -315,7 +326,7 @@ mod tests {
use common_error::ext::ErrorExt;
use datatypes::schema::{FulltextAnalyzer, FulltextOptions};
use sqlparser::ast::{ColumnOption, DataType};
use sqlparser::ast::{ColumnDef, ColumnOption, ColumnOptionDef, DataType};
use super::*;
use crate::dialect::GreptimeDbDialect;
@@ -397,19 +408,18 @@ mod tests {
assert_eq!("my_metric_1", alter_table.table_name().0[0].value);
let alter_operation = alter_table.alter_operation();
assert_matches!(alter_operation, AlterTableOperation::AddColumn { .. });
assert_matches!(alter_operation, AlterTableOperation::AddColumns { .. });
match alter_operation {
AlterTableOperation::AddColumn {
column_def,
location,
} => {
assert_eq!("tagk_i", column_def.name.value);
assert_eq!(DataType::String(None), column_def.data_type);
assert!(column_def
AlterTableOperation::AddColumns { add_columns } => {
assert_eq!(add_columns.len(), 1);
assert_eq!("tagk_i", add_columns[0].column_def.name.value);
assert_eq!(DataType::String(None), add_columns[0].column_def.data_type);
assert!(add_columns[0]
.column_def
.options
.iter()
.any(|o| matches!(o.option, ColumnOption::Null)));
assert_eq!(&None, location);
assert_eq!(&None, &add_columns[0].location);
}
_ => unreachable!(),
}
@@ -433,19 +443,17 @@ mod tests {
assert_eq!("my_metric_1", alter_table.table_name().0[0].value);
let alter_operation = alter_table.alter_operation();
assert_matches!(alter_operation, AlterTableOperation::AddColumn { .. });
assert_matches!(alter_operation, AlterTableOperation::AddColumns { .. });
match alter_operation {
AlterTableOperation::AddColumn {
column_def,
location,
} => {
assert_eq!("tagk_i", column_def.name.value);
assert_eq!(DataType::String(None), column_def.data_type);
assert!(column_def
AlterTableOperation::AddColumns { add_columns } => {
assert_eq!("tagk_i", add_columns[0].column_def.name.value);
assert_eq!(DataType::String(None), add_columns[0].column_def.data_type);
assert!(add_columns[0]
.column_def
.options
.iter()
.any(|o| matches!(o.option, ColumnOption::Null)));
assert_eq!(&Some(AddColumnLocation::First), location);
assert_eq!(&Some(AddColumnLocation::First), &add_columns[0].location);
}
_ => unreachable!(),
}
@@ -456,7 +464,8 @@ mod tests {
#[test]
fn test_parse_alter_add_column_with_after() {
let sql = "ALTER TABLE my_metric_1 ADD tagk_i STRING Null AFTER ts;";
let sql =
"ALTER TABLE my_metric_1 ADD tagk_i STRING Null AFTER ts, add column tagl_i String;";
let mut result =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
@@ -469,24 +478,42 @@ mod tests {
assert_eq!("my_metric_1", alter_table.table_name().0[0].value);
let alter_operation = alter_table.alter_operation();
assert_matches!(alter_operation, AlterTableOperation::AddColumn { .. });
assert_matches!(alter_operation, AlterTableOperation::AddColumns { .. });
match alter_operation {
AlterTableOperation::AddColumn {
column_def,
location,
} => {
assert_eq!("tagk_i", column_def.name.value);
assert_eq!(DataType::String(None), column_def.data_type);
assert!(column_def
.options
AlterTableOperation::AddColumns { add_columns } => {
let expecteds: Vec<(Option<AddColumnLocation>, ColumnDef)> = vec![
(
Some(AddColumnLocation::After {
column_name: "ts".to_string(),
}),
ColumnDef {
name: Ident::new("tagk_i"),
data_type: DataType::String(None),
options: vec![ColumnOptionDef {
name: None,
option: ColumnOption::Null,
}],
collation: None,
},
),
(
None,
ColumnDef {
name: Ident::new("tagl_i"),
data_type: DataType::String(None),
options: vec![],
collation: None,
},
),
];
for (add_column, expected) in add_columns
.iter()
.any(|o| matches!(o.option, ColumnOption::Null)));
assert_eq!(
&Some(AddColumnLocation::After {
column_name: "ts".to_string()
}),
location
);
.zip(expecteds)
.collect::<Vec<(&AddColumn, (Option<AddColumnLocation>, ColumnDef))>>()
{
assert_eq!(add_column.column_def, expected.1);
assert_eq!(&expected.0, &add_column.location);
}
}
_ => unreachable!(),
}

View File

@@ -291,12 +291,14 @@ impl ParserContext<'_> {
Token::Word(w) => match w.keyword {
Keyword::LIKE => {
self.parser.next_token();
Ok(ShowKind::Like(self.parse_identifier().with_context(
|_| error::UnexpectedSnafu {
expected: "LIKE",
actual: self.peek_token_as_string(),
},
)?))
Ok(ShowKind::Like(
Self::parse_identifier(&mut self.parser).with_context(|_| {
error::UnexpectedSnafu {
expected: "LIKE",
actual: self.peek_token_as_string(),
}
})?,
))
}
Keyword::WHERE => {
self.parser.next_token();
@@ -435,12 +437,12 @@ impl ParserContext<'_> {
))),
Token::Word(w) => match w.keyword {
Keyword::LIKE => Ok(Statement::ShowDatabases(ShowDatabases::new(
ShowKind::Like(self.parse_identifier().with_context(|_| {
error::UnexpectedSnafu {
ShowKind::Like(Self::parse_identifier(&mut self.parser).with_context(
|_| error::UnexpectedSnafu {
expected: "LIKE",
actual: tok.to_string(),
}
})?),
},
)?),
full,
))),
Keyword::WHERE => Ok(Statement::ShowDatabases(ShowDatabases::new(

View File

@@ -62,10 +62,7 @@ pub enum AlterTableOperation {
/// `ADD <table_constraint>`
AddConstraint(TableConstraint),
/// `ADD [ COLUMN ] <column_def> [location]`
AddColumn {
column_def: ColumnDef,
location: Option<AddColumnLocation>,
},
AddColumns { add_columns: Vec<AddColumn> },
/// `MODIFY <column_name> [target_type]`
ModifyColumnType {
column_name: Ident,
@@ -88,19 +85,32 @@ pub enum AlterTableOperation {
UnsetColumnFulltext { column_name: Ident },
}
#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
pub struct AddColumn {
pub column_def: ColumnDef,
pub location: Option<AddColumnLocation>,
}
impl Display for AddColumn {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(location) = &self.location {
write!(f, "{} {location}", self.column_def)
} else {
write!(f, "{}", self.column_def)
}
}
}
impl Display for AlterTableOperation {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AlterTableOperation::AddConstraint(constraint) => write!(f, r#"ADD {constraint}"#),
AlterTableOperation::AddColumn {
column_def,
location,
} => {
if let Some(location) = location {
write!(f, r#"ADD COLUMN {column_def} {location}"#)
} else {
write!(f, r#"ADD COLUMN {column_def}"#)
}
AlterTableOperation::AddColumns { add_columns } => {
let columns = add_columns
.iter()
.map(|add_column| format!("ADD COLUMN {add_column}"))
.join(", ");
write!(f, "{columns}")
}
AlterTableOperation::DropColumn { name } => write!(f, r#"DROP COLUMN {name}"#),
AlterTableOperation::RenameTable { new_table_name } => {
@@ -275,7 +285,8 @@ ALTER DATABASE db UNSET 'a','c'"#,
}
}
let sql = r"alter table monitor add column app string default 'shop' primary key;";
let sql =
r"alter table monitor add column app string default 'shop' primary key, add foo INT;";
let stmts =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
@@ -287,7 +298,7 @@ ALTER DATABASE db UNSET 'a','c'"#,
let new_sql = format!("\n{}", set);
assert_eq!(
r#"
ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY"#,
ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY, ADD COLUMN foo INT"#,
&new_sql
);
}

View File

@@ -57,10 +57,12 @@ impl TransformRule for TypeAliasTransformRule {
alter_table.alter_operation_mut()
{
replace_type_alias(target_type)
} else if let AlterTableOperation::AddColumn { column_def, .. } =
} else if let AlterTableOperation::AddColumns { add_columns, .. } =
alter_table.alter_operation_mut()
{
replace_type_alias(&mut column_def.data_type);
for add_column in add_columns {
replace_type_alias(&mut add_column.column_def.data_type);
}
}
}
_ => {}

View File

@@ -1347,14 +1347,17 @@ async fn test_alter_table(instance: Arc<dyn MockInstance>) {
));
// Add column
let output = execute_sql(&instance, "alter table demo add my_tag string null")
.await
.data;
let output = execute_sql(
&instance,
"alter table demo add my_tag1 string null, add my_tag2 integer",
)
.await
.data;
assert!(matches!(output, OutputData::AffectedRows(0)));
let output = execute_sql(
&instance,
"insert into demo(host, cpu, memory, ts, my_tag) values ('host2', 2.2, 200, 2000, 'hello')",
"insert into demo(host, cpu, memory, ts, my_tag1, my_tag2) values ('host2', 2.2, 200, 2000, 'hello', 123)",
)
.await
.data;
@@ -1371,13 +1374,13 @@ async fn test_alter_table(instance: Arc<dyn MockInstance>) {
.await
.data;
let expected = "\
+-------+-----+--------+---------------------+--------+
| host | cpu | memory | ts | my_tag |
+-------+-----+--------+---------------------+--------+
| host1 | 1.1 | 100.0 | 1970-01-01T00:00:01 | |
| host2 | 2.2 | 200.0 | 1970-01-01T00:00:02 | hello |
| host3 | 3.3 | 300.0 | 1970-01-01T00:00:03 | |
+-------+-----+--------+---------------------+--------+";
+-------+-----+--------+---------------------+---------+---------+
| host | cpu | memory | ts | my_tag1 | my_tag2 |
+-------+-----+--------+---------------------+---------+---------+
| host1 | 1.1 | 100.0 | 1970-01-01T00:00:01 | | |
| host2 | 2.2 | 200.0 | 1970-01-01T00:00:02 | hello | 123 |
| host3 | 3.3 | 300.0 | 1970-01-01T00:00:03 | | |
+-------+-----+--------+---------------------+---------+---------+";
check_output_stream(output, expected).await;
// Drop a column
@@ -1390,19 +1393,19 @@ async fn test_alter_table(instance: Arc<dyn MockInstance>) {
.await
.data;
let expected = "\
+-------+-----+---------------------+--------+
| host | cpu | ts | my_tag |
+-------+-----+---------------------+--------+
| host1 | 1.1 | 1970-01-01T00:00:01 | |
| host2 | 2.2 | 1970-01-01T00:00:02 | hello |
| host3 | 3.3 | 1970-01-01T00:00:03 | |
+-------+-----+---------------------+--------+";
+-------+-----+---------------------+---------+---------+
| host | cpu | ts | my_tag1 | my_tag2 |
+-------+-----+---------------------+---------+---------+
| host1 | 1.1 | 1970-01-01T00:00:01 | | |
| host2 | 2.2 | 1970-01-01T00:00:02 | hello | 123 |
| host3 | 3.3 | 1970-01-01T00:00:03 | | |
+-------+-----+---------------------+---------+---------+";
check_output_stream(output, expected).await;
// insert a new row
let output = execute_sql(
&instance,
"insert into demo(host, cpu, ts, my_tag) values ('host4', 400, 4000, 'world')",
"insert into demo(host, cpu, ts, my_tag1, my_tag2) values ('host4', 400, 4000, 'world', 321)",
)
.await
.data;
@@ -1412,14 +1415,14 @@ async fn test_alter_table(instance: Arc<dyn MockInstance>) {
.await
.data;
let expected = "\
+-------+-------+---------------------+--------+
| host | cpu | ts | my_tag |
+-------+-------+---------------------+--------+
| host1 | 1.1 | 1970-01-01T00:00:01 | |
| host2 | 2.2 | 1970-01-01T00:00:02 | hello |
| host3 | 3.3 | 1970-01-01T00:00:03 | |
| host4 | 400.0 | 1970-01-01T00:00:04 | world |
+-------+-------+---------------------+--------+";
+-------+-------+---------------------+---------+---------+
| host | cpu | ts | my_tag1 | my_tag2 |
+-------+-------+---------------------+---------+---------+
| host1 | 1.1 | 1970-01-01T00:00:01 | | |
| host2 | 2.2 | 1970-01-01T00:00:02 | hello | 123 |
| host3 | 3.3 | 1970-01-01T00:00:03 | | |
| host4 | 400.0 | 1970-01-01T00:00:04 | world | 321 |
+-------+-------+---------------------+---------+---------+";
check_output_stream(output, expected).await;
}

View File

@@ -104,6 +104,25 @@ DESC TABLE test;
| IdC | String | PRI | YES | idc | TAG |
+--------+----------------------+-----+------+---------+---------------+
ALTER TABLE test ADD COLUMN "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "bar" STRING default 'bar' PRIMARY KEY;
Affected Rows: 0
DESC TABLE test;
+--------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+--------+----------------------+-----+------+---------+---------------+
| i | Int32 | | YES | | FIELD |
| j | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| k | Int32 | | YES | | FIELD |
| host | String | PRI | YES | | TAG |
| idc | String | PRI | YES | idc | TAG |
| IdC | String | PRI | YES | idc | TAG |
| foo | String | PRI | YES | foo | TAG |
| bar | String | PRI | YES | bar | TAG |
+--------+----------------------+-----+------+---------+---------------+
DROP TABLE test;
Affected Rows: 0

View File

@@ -26,4 +26,8 @@ ALTER TABLE test ADD COLUMN "IdC" STRING default 'idc' PRIMARY KEY;
DESC TABLE test;
ALTER TABLE test ADD COLUMN "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "bar" STRING default 'bar' PRIMARY KEY;
DESC TABLE test;
DROP TABLE test;