feat: add column if not exists (#5393)

* feat: add column if not exists

* chore: address reviews
This commit is contained in:
Niwaka
2025-01-21 02:38:25 +00:00
committed by GitHub
parent 2ab235ec9d
commit 0163ce8df9
7 changed files with 110 additions and 8 deletions
+1 -2
View File
@@ -507,8 +507,7 @@ pub(crate) fn to_alter_table_expr(
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,
add_if_not_exists: add_column.add_if_not_exists,
})
})
.collect::<Result<Vec<AddColumn>>>()?,
+27 -5
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use api::helper::ColumnDataTypeWrapper;
@@ -885,12 +885,17 @@ impl StatementExecutor {
Ok(Output::new_with_affected_rows(0))
}
/// Verifies an alter and returns whether it is necessary to perform the alter.
///
/// # Returns
///
/// Returns true if the alter need to be porformed; otherwise, it returns false.
fn verify_alter(
&self,
table_id: TableId,
table_info: Arc<TableInfo>,
expr: AlterTableExpr,
) -> Result<()> {
) -> Result<bool> {
let request: AlterTableRequest = common_grpc_expr::alter_expr_to_request(table_id, expr)
.context(AlterExprToRequestSnafu)?;
@@ -907,6 +912,21 @@ impl StatementExecutor {
violated: format!("Invalid table name: {}", new_table_name)
}
);
} else if let AlterKind::AddColumns { columns } = alter_kind {
// If all the columns are marked as add_if_not_exists and they already exist in the table,
// there is no need to perform the alter.
let column_names: HashSet<_> = table_info
.meta
.schema
.column_schemas()
.iter()
.map(|schema| &schema.name)
.collect();
if columns.iter().all(|column| {
column_names.contains(&column.column_schema.name) && column.add_if_not_exists
}) {
return Ok(false);
}
}
let _ = table_info
@@ -916,7 +936,7 @@ impl StatementExecutor {
.build()
.context(error::BuildTableMetaSnafu { table_name })?;
Ok(())
Ok(true)
}
#[tracing::instrument(skip_all)]
@@ -971,8 +991,10 @@ impl StatementExecutor {
})?;
let table_id = table.table_info().ident.table_id;
self.verify_alter(table_id, table.table_info(), expr.clone())?;
let need_alter = self.verify_alter(table_id, table.table_info(), expr.clone())?;
if !need_alter {
return Ok(Output::new_with_affected_rows(0));
}
info!(
"Table info before alter is {:?}, expr: {:?}",
table.table_info(),
-1
View File
@@ -70,7 +70,6 @@ impl DfContextProviderAdapter {
);
let tables = resolve_tables(table_names, &mut table_provider).await?;
Ok(Self {
engine_state,
session_state,
+63
View File
@@ -341,6 +341,7 @@ fn parse_string_options(parser: &mut Parser) -> std::result::Result<(String, Str
fn parse_add_columns(parser: &mut Parser) -> std::result::Result<AddColumn, ParserError> {
parser.expect_keyword(Keyword::ADD)?;
let _ = parser.parse_keyword(Keyword::COLUMN);
let add_if_not_exists = parser.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
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) {
@@ -362,6 +363,7 @@ fn parse_add_columns(parser: &mut Parser) -> std::result::Result<AddColumn, Pars
Ok(AddColumn {
column_def,
location,
add_if_not_exists,
})
}
@@ -572,6 +574,67 @@ mod tests {
}
}
#[test]
fn test_parse_add_column_if_not_exists() {
let sql = "ALTER TABLE test ADD COLUMN IF NOT EXISTS a INTEGER, ADD COLUMN b STRING, ADD COLUMN IF NOT EXISTS c INT;";
let mut result =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
assert_eq!(result.len(), 1);
let statement = result.remove(0);
assert_matches!(statement, Statement::AlterTable { .. });
match statement {
Statement::AlterTable(alter) => {
assert_eq!(alter.table_name.0[0].value, "test");
assert_matches!(
alter.alter_operation,
AlterTableOperation::AddColumns { .. }
);
match alter.alter_operation {
AlterTableOperation::AddColumns { add_columns } => {
let expected = vec![
AddColumn {
column_def: ColumnDef {
name: Ident::new("a"),
data_type: DataType::Integer(None),
collation: None,
options: vec![],
},
location: None,
add_if_not_exists: true,
},
AddColumn {
column_def: ColumnDef {
name: Ident::new("b"),
data_type: DataType::String(None),
collation: None,
options: vec![],
},
location: None,
add_if_not_exists: false,
},
AddColumn {
column_def: ColumnDef {
name: Ident::new("c"),
data_type: DataType::Int(None),
collation: None,
options: vec![],
},
location: None,
add_if_not_exists: true,
},
];
for (idx, add_column) in add_columns.into_iter().enumerate() {
assert_eq!(add_column, expected[idx]);
}
}
_ => unreachable!(),
}
}
_ => unreachable!(),
}
}
#[test]
fn test_parse_alter_drop_column() {
let sql = "ALTER TABLE my_metric_1 DROP a";
+1
View File
@@ -118,6 +118,7 @@ pub enum UnsetIndexOperation {
pub struct AddColumn {
pub column_def: ColumnDef,
pub location: Option<AddColumnLocation>,
pub add_if_not_exists: bool,
}
impl Display for AddColumn {
@@ -19,6 +19,10 @@ ALTER TABLE test ADD COLUMN k INTEGER;
Affected Rows: 0
ALTER TABLE test ADD COLUMN IF NOT EXISTS k INTEGER;
Affected Rows: 0
SELECT * FROM test;
+---+-------------------------+---+
@@ -108,6 +112,14 @@ ALTER TABLE test ADD COLUMN "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "
Affected Rows: 0
ALTER TABLE test ADD COLUMN IF NOT EXISTS "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "bar" STRING default 'bar' PRIMARY KEY;
Error: 4003(TableColumnExists), Column bar already exists in table test
ALTER TABLE test ADD COLUMN IF NOT EXISTS "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN IF NOT EXISTS "bar" STRING default 'bar' PRIMARY KEY;
Affected Rows: 0
DESC TABLE test;
+--------+----------------------+-----+------+---------+---------------+
@@ -6,6 +6,8 @@ INSERT INTO test VALUES (1, 1), (2, 2);
ALTER TABLE test ADD COLUMN k INTEGER;
ALTER TABLE test ADD COLUMN IF NOT EXISTS k INTEGER;
SELECT * FROM test;
DESC TABLE test;
@@ -28,6 +30,10 @@ DESC TABLE test;
ALTER TABLE test ADD COLUMN "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "bar" STRING default 'bar' PRIMARY KEY;
ALTER TABLE test ADD COLUMN IF NOT EXISTS "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN "bar" STRING default 'bar' PRIMARY KEY;
ALTER TABLE test ADD COLUMN IF NOT EXISTS "foo" STRING default 'foo' PRIMARY KEY, ADD COLUMN IF NOT EXISTS "bar" STRING default 'bar' PRIMARY KEY;
DESC TABLE test;
DROP TABLE test;