feat: add column if not exists (#5393)

* feat: add column if not exists

* chore: address reviews

src/operator/src/expr_factory.rs

@@ -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>>>()?,

src/operator/src/statement/ddl.rs

@@ -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(),

src/query/src/datafusion/planner.rs

@@ -70,7 +70,6 @@ impl DfContextProviderAdapter {
         );
 
         let tables = resolve_tables(table_names, &mut table_provider).await?;
-
         Ok(Self {
             engine_state,
             session_state,

src/sql/src/parsers/alter_parser.rs

@@ -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";

src/sql/src/statements/alter.rs

@@ -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 {

tests/cases/standalone/common/alter/add_col.result

@@ -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;
 
 +--------+----------------------+-----+------+---------+---------------+

tests/cases/standalone/common/alter/add_col.sql

@@ -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;