chore: Add type check to alter parallelism statement. (#14323) (#14331)

Signed-off-by: Shanicky Chen <[email protected]> Co-authored-by: Shanicky Chen <[email protected]> Co-authored-by: August <[email protected]>
risingwavelabs · Jan 5, 2024 · 0b0075c · 0b0075c
1 parent d1fc91a
commit 0b0075c
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 3 deletions.
diff --git a/src/frontend/src/handler/alter_parallelism.rs b/src/frontend/src/handler/alter_parallelism.rs
@@ -22,7 +22,10 @@ use risingwave_sqlparser::keywords::Keyword;
 
 use super::{HandlerArgs, RwPgResponse};
 use crate::catalog::root_catalog::SchemaPath;
+use crate::catalog::table_catalog::TableType;
+use crate::catalog::CatalogError;
 use crate::Binder;
+
 pub async fn handle_alter_parallelism(
     handler_args: HandlerArgs,
     obj_name: ObjectName,
@@ -41,9 +44,33 @@ pub async fn handle_alter_parallelism(
         let reader = session.env().catalog_reader().read_guard();
 
         match stmt_type {
-            StatementType::ALTER_TABLE | StatementType::ALTER_MATERIALIZED_VIEW => {
+            StatementType::ALTER_TABLE
+            | StatementType::ALTER_MATERIALIZED_VIEW
+            | StatementType::ALTER_INDEX => {
                 let (table, schema_name) =
                     reader.get_table_by_name(db_name, schema_path, &real_table_name)?;
+
+                match (table.table_type(), stmt_type) {
+                    (TableType::Internal, _) => {
+                        // we treat internal table as NOT FOUND
+                        return Err(
+                            CatalogError::NotFound("table", table.name().to_string()).into()
+                        );
+                    }
+                    (TableType::Table, StatementType::ALTER_TABLE)
+                    | (TableType::MaterializedView, StatementType::ALTER_MATERIALIZED_VIEW)
+                    | (TableType::Index, StatementType::ALTER_INDEX) => {}
+                    _ => {
+                        return Err(ErrorCode::InvalidInputSyntax(format!(
+                            "cannot alter parallelism of {} {} by {}",
+                            table.table_type().to_prost().as_str_name(),
+                            table.name(),
+                            stmt_type,
+                        ))
+                        .into());
+                    }
+                }
+
                 session.check_privilege_for_drop_alter(schema_name, &**table)?;
                 table.id.table_id()
             }

diff --git a/src/frontend/src/handler/mod.rs b/src/frontend/src/handler/mod.rs
@@ -528,6 +528,18 @@ pub async fn handle(
             name,
             operation: AlterIndexOperation::RenameIndex { index_name },
         } => alter_rename::handle_rename_index(handler_args, name, index_name).await,
+        Statement::AlterIndex {
+            name,
+            operation: AlterIndexOperation::SetParallelism { parallelism },
+        } => {
+            alter_parallelism::handle_alter_parallelism(
+                handler_args,
+                name,
+                parallelism,
+                StatementType::ALTER_INDEX,
+            )
+            .await
+        }
         Statement::AlterView {
             materialized,
             name,

diff --git a/src/sqlparser/src/ast/ddl.rs b/src/sqlparser/src/ast/ddl.rs
@@ -92,7 +92,13 @@ pub enum AlterTableOperation {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum AlterIndexOperation {
-    RenameIndex { index_name: ObjectName },
+    RenameIndex {
+        index_name: ObjectName,
+    },
+    /// `SET PARALLELISM TO <parallelism>`
+    SetParallelism {
+        parallelism: SetVariableValue,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -251,6 +257,9 @@ impl fmt::Display for AlterIndexOperation {
             AlterIndexOperation::RenameIndex { index_name } => {
                 write!(f, "RENAME TO {index_name}")
             }
+            AlterIndexOperation::SetParallelism { parallelism } => {
+                write!(f, "SET PARALLELISM TO {}", parallelism)
+            }
         }
     }
 }

diff --git a/src/sqlparser/src/parser.rs b/src/sqlparser/src/parser.rs
@@ -2960,7 +2960,7 @@ impl Parser {
 
                 AlterTableOperation::SetParallelism { parallelism: value }
             } else {
-                return self.expected("SCHEMA after SET", self.peek_token());
+                return self.expected("SCHEMA/PARALLELISM after SET", self.peek_token());
             }
         } else if self.parse_keyword(Keyword::DROP) {
             let _ = self.parse_keyword(Keyword::COLUMN);
@@ -3024,6 +3024,23 @@ impl Parser {
             } else {
                 return self.expected("TO after RENAME", self.peek_token());
             }
+        } else if self.parse_keyword(Keyword::SET) {
+            if self.parse_keyword(Keyword::PARALLELISM) {
+                if self.expect_keyword(Keyword::TO).is_err()
+                    && self.expect_token(&Token::Eq).is_err()
+                {
+                    return self.expected(
+                        "TO or = after ALTER TABLE SET PARALLELISM",
+                        self.peek_token(),
+                    );
+                }
+
+                let value = self.parse_set_variable()?;
+
+                AlterIndexOperation::SetParallelism { parallelism: value }
+            } else {
+                return self.expected("PARALLELISM after SET", self.peek_token());
+            }
         } else {
             return self.expected("RENAME after ALTER INDEX", self.peek_token());
         };