Merge pull request #7044 from lichuang/share_sql

feat: add alter share tenants sql

common/ast/src/ast/statements/share.rs

@@ -17,6 +17,7 @@ use std::fmt::Formatter;
 
 use common_meta_app::share::ShareGrantObjectName;
 use common_meta_app::share::ShareGrantObjectPrivilege;
+use itertools::Itertools;
 
 use crate::ast::Identifier;
 
@@ -96,3 +97,33 @@ impl Display for RevokeShareObjectStmt<'_> {
         Ok(())
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AlterShareTenantsStmt<'a> {
+    pub share: Identifier<'a>,
+    pub if_exists: bool,
+    pub tenants: Vec<Identifier<'a>>,
+    pub is_add: bool,
+}
+
+impl Display for AlterShareTenantsStmt<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "ALTER SHARE ")?;
+        if self.if_exists {
+            write!(f, "IF EXISTS ")?;
+        }
+        write!(f, "{}", self.share)?;
+        if self.is_add {
+            write!(f, " ADD TENANTS = ")?;
+        } else {
+            write!(f, " REMOVE TENANTS = ")?;
+        }
+        write!(
+            f,
+            "{}",
+            self.tenants.iter().map(|v| v.to_string()).join(",")
+        )?;
+
+        Ok(())
+    }
+}

common/ast/src/ast/statements/statement.rs

@@ -163,6 +163,7 @@ pub enum Statement<'a> {
     DropShare(DropShareStmt<'a>),
     GrantShareObject(GrantShareObjectStmt<'a>),
     RevokeShareObject(RevokeShareObjectStmt<'a>),
+    AlterShareTenants(AlterShareTenantsStmt<'a>),
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -368,6 +369,7 @@ impl<'a> Display for Statement<'a> {
             Statement::DropShare(stmt) => write!(f, "{stmt}")?,
             Statement::GrantShareObject(stmt) => write!(f, "{stmt}")?,
             Statement::RevokeShareObject(stmt) => write!(f, "{stmt}")?,
+            Statement::AlterShareTenants(stmt) => write!(f, "{stmt}")?,
         }
         Ok(())
     }

common/ast/src/parser/statement.rs

@@ -794,6 +794,19 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
             })
         },
     );
+    let alter_share_tenants = map(
+        rule! {
+            ALTER ~ SHARE ~ (IF ~ EXISTS )? ~ #ident ~ #alter_add_share_accounts ~ TENANTS ~ Eq ~ #comma_separated_list1(ident)
+        },
+        |(_, _, opt_if_exists, share, is_add, _, _, tenants)| {
+            Statement::AlterShareTenants(AlterShareTenantsStmt {
+                share,
+                if_exists: opt_if_exists.is_some(),
+                is_add,
+                tenants,
+            })
+        },
+    );
 
     let statement_body = alt((
         rule!(
@@ -886,6 +899,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
             | #drop_share: "`DROP SHARE [IF EXISTS] <share_name>`"
             | #grant_share_object: "`GRANT { USAGE | SELECT | REFERENCE_USAGE } ON { DATABASE db | TABLE db.table } TO SHARE <share_name>`"
             | #revoke_share_object: "`REVOKE { USAGE | SELECT | REFERENCE_USAGE } ON { DATABASE db | TABLE db.table } FROM SHARE <share_name>`"
+            | #alter_share_tenants: "`ALTER SHARE [IF EXISTS] <share_name> { ADD | REMOVE } TENANTS = tenant [, tenant, ...]`"
         ),
     ));
 
@@ -1054,6 +1068,10 @@ pub fn priv_share_type(i: Input) -> IResult<ShareGrantObjectPrivilege> {
     ))(i)
 }
 
+pub fn alter_add_share_accounts(i: Input) -> IResult<bool> {
+    alt((value(true, rule! { ADD }), value(false, rule! { REMOVE })))(i)
+}
+
 pub fn grant_share_object_name(i: Input) -> IResult<ShareGrantObjectName> {
     let database = map(
         rule! {

common/ast/src/parser/token.rs

@@ -237,6 +237,8 @@ pub enum TokenKind {
     //    reserved list.
     #[token("ALL", ignore(ascii_case))]
     ALL,
+    #[token("ADD", ignore(ascii_case))]
+    ADD,
     #[token("ANY", ignore(ascii_case))]
     ANY,
     #[token("SOME", ignore(ascii_case))]
@@ -625,6 +627,8 @@ pub enum TokenKind {
     TEXT,
     #[token("TENANTSETTING", ignore(ascii_case))]
     TENANTSETTING,
+    #[token("TENANTS", ignore(ascii_case))]
+    TENANTS,
     #[token("THEN", ignore(ascii_case))]
     THEN,
     #[token("TIMESTAMP", ignore(ascii_case))]

common/ast/tests/it/parser.rs

@@ -265,6 +265,9 @@ fn test_statement() {
         r#"GRANT SELECT ON TABLE db1.tb1 TO SHARE a;"#,
         r#"REVOKE USAGE ON DATABASE db1 FROM SHARE a;"#,
         r#"REVOKE SELECT ON TABLE db1.tb1 FROM SHARE a;"#,
+        r#"ALTER SHARE a ADD TENANTS = b,c;"#,
+        r#"ALTER SHARE IF EXISTS a ADD TENANTS = b,c;"#,
+        r#"ALTER SHARE IF EXISTS a REMOVE TENANTS = b,c;"#,
     ];
 
     for case in cases {

common/ast/tests/it/testdata/statement.txt

@@ -6090,3 +6090,93 @@ RevokeShareObject(
 )
 
 
+---------- Input ----------
+ALTER SHARE a ADD TENANTS = b,c;
+---------- Output ---------
+ALTER SHARE a ADD TENANTS = b,c
+---------- AST ------------
+AlterShareTenants(
+    AlterShareTenantsStmt {
+        share: Identifier {
+            name: "a",
+            quote: None,
+            span: Ident(12..13),
+        },
+        if_exists: false,
+        tenants: [
+            Identifier {
+                name: "b",
+                quote: None,
+                span: Ident(28..29),
+            },
+            Identifier {
+                name: "c",
+                quote: None,
+                span: Ident(30..31),
+            },
+        ],
+        is_add: true,
+    },
+)
+
+
+---------- Input ----------
+ALTER SHARE IF EXISTS a ADD TENANTS = b,c;
+---------- Output ---------
+ALTER SHARE IF EXISTS a ADD TENANTS = b,c
+---------- AST ------------
+AlterShareTenants(
+    AlterShareTenantsStmt {
+        share: Identifier {
+            name: "a",
+            quote: None,
+            span: Ident(22..23),
+        },
+        if_exists: true,
+        tenants: [
+            Identifier {
+                name: "b",
+                quote: None,
+                span: Ident(38..39),
+            },
+            Identifier {
+                name: "c",
+                quote: None,
+                span: Ident(40..41),
+            },
+        ],
+        is_add: true,
+    },
+)
+
+
+---------- Input ----------
+ALTER SHARE IF EXISTS a REMOVE TENANTS = b,c;
+---------- Output ---------
+ALTER SHARE IF EXISTS a REMOVE TENANTS = b,c
+---------- AST ------------
+AlterShareTenants(
+    AlterShareTenantsStmt {
+        share: Identifier {
+            name: "a",
+            quote: None,
+            span: Ident(22..23),
+        },
+        if_exists: true,
+        tenants: [
+            Identifier {
+                name: "b",
+                quote: None,
+                span: Ident(41..42),
+            },
+            Identifier {
+                name: "c",
+                quote: None,
+                span: Ident(43..44),
+            },
+        ],
+        is_add: false,
+    },
+)
+
+

query/src/interpreters/interpreter_factory_v2.rs

@@ -267,6 +267,10 @@ impl InterpreterFactoryV2 {
                 ctx,
                 *p.clone(),
             )?)),
+            Plan::AlterShareTenants(p) => Ok(Arc::new(AlterShareTenantsInterpreter::try_create(
+                ctx,
+                *p.clone(),
+            )?)),
         }
     }
 }

query/src/interpreters/interpreter_share_alter_tenants.rs

@@ -0,0 +1,81 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use common_datavalues::chrono::Utc;
+use common_exception::Result;
+use common_meta_api::ShareApi;
+use common_meta_app::share::AddShareAccountsReq;
+use common_meta_app::share::RemoveShareAccountsReq;
+use common_meta_app::share::ShareNameIdent;
+use common_streams::DataBlockStream;
+use common_streams::SendableDataBlockStream;
+
+use crate::interpreters::Interpreter;
+use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
+use crate::sql::plans::share::AlterShareTenantsPlan;
+
+pub struct AlterShareTenantsInterpreter {
+    ctx: Arc<QueryContext>,
+    plan: AlterShareTenantsPlan,
+}
+
+impl AlterShareTenantsInterpreter {
+    pub fn try_create(ctx: Arc<QueryContext>, plan: AlterShareTenantsPlan) -> Result<Self> {
+        Ok(AlterShareTenantsInterpreter { ctx, plan })
+    }
+}
+
+#[async_trait::async_trait]
+impl Interpreter for AlterShareTenantsInterpreter {
+    fn name(&self) -> &str {
+        "AlterShareTenantsInterpreter"
+    }
+
+    async fn execute(&self) -> Result<SendableDataBlockStream> {
+        let tenant = self.ctx.get_tenant();
+        let user_mgr = self.ctx.get_user_manager();
+        let meta_api = user_mgr.get_meta_store_client();
+        if self.plan.is_add {
+            let req = AddShareAccountsReq {
+                share_name: ShareNameIdent {
+                    tenant,
+                    share_name: self.plan.share.clone(),
+                },
+                if_exists: self.plan.if_exists,
+                accounts: self.plan.accounts.clone(),
+                share_on: Utc::now(),
+            };
+            meta_api.add_share_accounts(req).await?;
+        } else {
+            let req = RemoveShareAccountsReq {
+                share_name: ShareNameIdent {
+                    tenant,
+                    share_name: self.plan.share.clone(),
+                },
+                if_exists: self.plan.if_exists,
+                accounts: self.plan.accounts.clone(),
+            };
+            meta_api.remove_share_accounts(req).await?;
+        }
+
+        Ok(Box::pin(DataBlockStream::create(
+            self.plan.schema(),
+            None,
+            vec![],
+        )))
+    }
+}

query/src/interpreters/mod.rs

@@ -49,6 +49,7 @@ mod interpreter_role_revoke;
 mod interpreter_select;
 mod interpreter_select_v2;
 mod interpreter_setting;
+mod interpreter_share_alter_tenants;
 mod interpreter_share_create;
 mod interpreter_share_drop;
 mod interpreter_share_grant_object;
@@ -133,6 +134,7 @@ pub use interpreter_role_revoke::RevokeRoleInterpreter;
 pub use interpreter_select::SelectInterpreter;
 pub use interpreter_select_v2::SelectInterpreterV2;
 pub use interpreter_setting::SettingInterpreter;
+pub use interpreter_share_alter_tenants::AlterShareTenantsInterpreter;
 pub use interpreter_share_create::CreateShareInterpreter;
 pub use interpreter_share_drop::DropShareInterpreter;
 pub use interpreter_share_grant_object::GrantShareObjectInterpreter;

query/src/sql/planner/binder/ddl/share.rs

@@ -14,10 +14,12 @@
 
 use common_ast::ast::*;
 use common_exception::Result;
+use itertools::Itertools;
 
 use crate::sessions::TableContext;
 use crate::sql::binder::Binder;
 use crate::sql::normalize_identifier;
+use crate::sql::plans::AlterShareTenantsPlan;
 use crate::sql::plans::CreateSharePlan;
 use crate::sql::plans::DropSharePlan;
 use crate::sql::plans::GrantShareObjectPlan;
@@ -101,4 +103,26 @@ impl<'a> Binder {
         };
         Ok(Plan::RevokeShareObject(Box::new(plan)))
     }
+
+    pub(in crate::sql::planner::binder) async fn bind_alter_share_accounts(
+        &mut self,
+        stmt: &AlterShareTenantsStmt<'a>,
+    ) -> Result<Plan> {
+        let AlterShareTenantsStmt {
+            share,
+            if_exists,
+            tenants,
+            is_add,
+        } = stmt;
+
+        let share = normalize_identifier(share, &self.name_resolution_ctx).name;
+
+        let plan = AlterShareTenantsPlan {
+            share,
+            if_exists: *if_exists,
+            is_add: *is_add,
+            accounts: tenants.iter().map(|v| v.to_string()).collect_vec(),
+        };
+        Ok(Plan::AlterShareTenants(Box::new(plan)))
+    }
 }

query/src/sql/planner/binder/mod.rs

@@ -312,6 +312,9 @@ impl<'a> Binder {
             Statement::RevokeShareObject(stmt) => {
                 self.bind_revoke_share_object(stmt).await?
             }
+            Statement::AlterShareTenants(stmt) => {
+                self.bind_alter_share_accounts(stmt).await?
+            }
         };
         Ok(plan)
     }

query/src/sql/planner/format/display_plan.rs

@@ -98,6 +98,7 @@ impl Plan {
             Plan::DropShare(p) => Ok(format!("{:?}", p)),
             Plan::GrantShareObject(p) => Ok(format!("{:?}", p)),
             Plan::RevokeShareObject(p) => Ok(format!("{:?}", p)),
+            Plan::AlterShareTenants(p) => Ok(format!("{:?}", p)),
         }
     }
 }