add test

databendlabs · May 22, 2024 · 9ae328a · 9ae328a
1 parent a8f8bc9
commit 9ae328a
Show file tree

Hide file tree

Showing 16 changed files with 359 additions and 134 deletions.
diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs
@@ -204,6 +204,7 @@ pub enum Statement {
         principal: Option<PrincipalIdentity>,
         show_options: Option<ShowOptions>,
     },
+    ShowObjectPrivileges(ShowObjectPrivilegesStmt),
     Revoke(RevokeStmt),
 
     // UDF
@@ -619,6 +620,7 @@ impl Display for Statement {
                     write!(f, " {show_options}")?;
                 }
             }
+            Statement::ShowObjectPrivileges(stmt) => write!(f, "{stmt}")?,
             Statement::Revoke(stmt) => write!(f, "{stmt}")?,
             Statement::CreateUDF(stmt) => write!(f, "{stmt}")?,
             Statement::DropUDF {

diff --git a/src/query/ast/src/ast/statements/user.rs b/src/query/ast/src/ast/statements/user.rs
@@ -22,6 +22,7 @@ use crate::ast::write_comma_separated_list;
 use crate::ast::AuthType;
 use crate::ast::CreateOption;
 use crate::ast::PrincipalIdentity;
+use crate::ast::ShowOptions;
 use crate::ast::UserIdentity;
 use crate::ast::UserPrivilegeType;
 
@@ -135,6 +136,50 @@ impl Display for RevokeStmt {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct ShowObjectPrivilegesStmt {
+    pub object: GrantObjectName,
+    pub show_option: Option<ShowOptions>,
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum GrantObjectName {
+    Database(#[drive(skip)] String),
+    Table(#[drive(skip)] Option<String>, #[drive(skip)] String),
+    UDF(#[drive(skip)] String),
+    Stage(#[drive(skip)] String),
+}
+
+impl Display for GrantObjectName {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            GrantObjectName::Database(database_name) => {
+                write!(f, "DATABASE {database_name}")
+            }
+            GrantObjectName::Table(database_name, table_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, "TABLE {database_name}.{table_name}")
+                } else {
+                    write!(f, "TABLE {table_name}")
+                }
+            }
+            GrantObjectName::UDF(udf) => write!(f, " UDF {udf}"),
+            GrantObjectName::Stage(stage) => write!(f, " STAGE {stage}"),
+        }
+    }
+}
+
+impl Display for ShowObjectPrivilegesStmt {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "SHOW GRANTS ON {}", self.object)?;
+
+        if let Some(show_option) = &self.show_option {
+            write!(f, " {show_option}")?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub enum AccountMgrSource {
     Role {
@@ -158,48 +203,12 @@ impl Display for AccountMgrSource {
                 write!(f, " ")?;
                 write_comma_separated_list(f, privileges.iter().map(|p| p.to_string()))?;
                 write!(f, " ON")?;
-                match level {
-                    AccountMgrLevel::Global => write!(f, " *.*")?,
-                    AccountMgrLevel::Database(database_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.*")?;
-                        } else {
-                            write!(f, " *")?;
-                        }
-                    }
-                    AccountMgrLevel::Table(database_name, table_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.{table_name}")?;
-                        } else {
-                            write!(f, " {table_name}")?;
-                        }
-                    }
-                    AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}")?,
-                    AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}")?,
-                }
+                write!(f, " {}", level)?;
             }
             AccountMgrSource::ALL { level, .. } => {
                 write!(f, " ALL PRIVILEGES")?;
                 write!(f, " ON")?;
-                match level {
-                    AccountMgrLevel::Global => write!(f, " *.*")?,
-                    AccountMgrLevel::Database(database_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.*")?;
-                        } else {
-                            write!(f, " *")?;
-                        }
-                    }
-                    AccountMgrLevel::Table(database_name, table_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.{table_name}")?;
-                        } else {
-                            write!(f, " {table_name}")?;
-                        }
-                    }
-                    AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}")?,
-                    AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}")?,
-                }
+                write!(f, " {}", level)?;
             }
         }
         Ok(())
@@ -215,6 +224,30 @@ pub enum AccountMgrLevel {
     Stage(#[drive(skip)] String),
 }
 
+impl Display for AccountMgrLevel {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            AccountMgrLevel::Global => write!(f, " *.*"),
+            AccountMgrLevel::Database(database_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, " {database_name}.*")
+                } else {
+                    write!(f, " *")
+                }
+            }
+            AccountMgrLevel::Table(database_name, table_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, " {database_name}.{table_name}")
+                } else {
+                    write!(f, " {table_name}")
+                }
+            }
+            AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}"),
+            AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub enum SecondaryRolesOption {
     None,

diff --git a/src/query/ast/src/ast/visitors/visitor.rs b/src/query/ast/src/ast/visitors/visitor.rs
@@ -619,6 +619,8 @@ pub trait Visitor<'ast>: Sized {
     ) {
     }
 
+    fn visit_show_object_priv(&mut self, _show: &'ast ShowObjectPrivilegesStmt) {}
+
     fn visit_revoke(&mut self, _revoke: &'ast RevokeStmt) {}
 
     fn visit_create_udf(&mut self, _stmt: &'ast CreateUDFStmt) {}

diff --git a/src/query/ast/src/ast/visitors/visitor_mut.rs b/src/query/ast/src/ast/visitors/visitor_mut.rs
@@ -631,6 +631,8 @@ pub trait VisitorMut: Sized {
     ) {
     }
 
+    fn visit_show_object_priv(&mut self, _show: &mut ShowObjectPrivilegesStmt) {}
+
     fn visit_revoke(&mut self, _revoke: &mut RevokeStmt) {}
 
     fn visit_create_udf(&mut self, _stmt: &mut CreateUDFStmt) {}

diff --git a/src/query/ast/src/ast/visitors/walk.rs b/src/query/ast/src/ast/visitors/walk.rs
@@ -507,6 +507,7 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
             principal,
             show_options,
         } => visitor.visit_show_grant(principal, show_options),
+        Statement::ShowObjectPrivileges(stmt) => visitor.visit_show_object_priv(stmt),
         Statement::Revoke(stmt) => visitor.visit_revoke(stmt),
         Statement::CreateUDF(stmt) => visitor.visit_create_udf(stmt),
         Statement::DropUDF {

diff --git a/src/query/ast/src/ast/visitors/walk_mut.rs b/src/query/ast/src/ast/visitors/walk_mut.rs
@@ -502,6 +502,7 @@ pub fn walk_statement_mut<V: VisitorMut>(visitor: &mut V, statement: &mut Statem
             principal,
             show_options,
         } => visitor.visit_show_grant(principal, show_options),
+        Statement::ShowObjectPrivileges(stmt) => visitor.visit_show_object_priv(stmt),
         Statement::Revoke(stmt) => visitor.visit_revoke(stmt),
         Statement::CreateUDF(stmt) => visitor.visit_create_udf(stmt),
         Statement::DropUDF {

diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -43,6 +43,7 @@ use crate::rule;
 
 pub enum ShowGrantOption {
     PrincipalIdentity(PrincipalIdentity),
+    GrantObjectName(GrantObjectName),
     ShareGrantObjectName(ShareGrantObjectName),
     ShareName(String),
 }
@@ -1295,6 +1296,12 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
                 principal: None,
                 show_options: opt_limit,
             },
+            Some(ShowGrantOption::GrantObjectName(object)) => {
+                Statement::ShowObjectPrivileges(ShowObjectPrivilegesStmt {
+                    object,
+                    show_option: opt_limit,
+                })
+            }
         },
     );
     let revoke = map(
@@ -2909,6 +2916,40 @@ pub fn grant_share_object_name(i: Input) -> IResult<ShareGrantObjectName> {
     )(i)
 }
 
+pub fn on_object_name(i: Input) -> IResult<GrantObjectName> {
+    let database = map(
+        rule! {
+            DATABASE ~ #ident
+        },
+        |(_, database)| GrantObjectName::Database(database.to_string()),
+    );
+
+    // `db01`.'tb1' or `db01`.`tb1` or `db01`.tb1
+    let table = map(
+        rule! {
+            TABLE ~  #dot_separated_idents_1_to_2
+        },
+        |(_, (database, table))| {
+            GrantObjectName::Table(database.map(|db| db.to_string()), table.to_string())
+        },
+    );
+
+    let stage = map(rule! { STAGE ~ #ident}, |(_, stage_name)| {
+        GrantObjectName::Stage(stage_name.to_string())
+    });
+
+    let udf = map(rule! { UDF ~ #ident}, |(_, udf_name)| {
+        GrantObjectName::UDF(udf_name.to_string())
+    });
+
+    rule!(
+        #database : "DATABASE <database>"
+        | #table : "TABLE <database>.<table>"
+        | #stage : "STAGE <stage_name>"
+        | #udf : "UDF <udf_name>"
+    )(i)
+}
+
 pub fn grant_level(i: Input) -> IResult<AccountMgrLevel> {
     // *.*
     let global = map(rule! { "*" ~ "." ~ "*" }, |_| AccountMgrLevel::Global);
@@ -3027,9 +3068,9 @@ pub fn show_grant_option(i: Input) -> IResult<ShowGrantOption> {
 
     let share_object_name = map(
         rule! {
-            ON ~ #grant_share_object_name
+            ON ~ #on_object_name
         },
-        |(_, object_name)| ShowGrantOption::ShareGrantObjectName(object_name),
+        |(_, object_name)| ShowGrantOption::GrantObjectName(object_name),
     );
 
     let share_name = map(
@@ -3041,7 +3082,7 @@ pub fn show_grant_option(i: Input) -> IResult<ShowGrantOption> {
 
     rule!(
         #grant_role: "FOR  { ROLE <role_name> | [USER] <user> }"
-        | #share_object_name: "ON {DATABASE <db_name> | TABLE <db_name>.<table_name>}"
+        | #share_object_name: "ON {DATABASE <db_name> | TABLE <db_name>.<table_name> | UDF <udf_name> | STAGE <stage_name> }"
         | #share_name: "OF SHARE <share_name>"
     )(i)
 }

diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs
@@ -310,6 +310,7 @@ fn test_statement() {
         r#"SHOW GRANTS FOR USER 'test-grant';"#,
         r#"SHOW GRANTS FOR ROLE role1;"#,
         r#"SHOW GRANTS FOR ROLE 'role1';"#,
+        r#"SHOW GRANTS ON TABLE t;"#,
         r#"REVOKE SELECT, CREATE ON * FROM 'test-grant';"#,
         r#"REVOKE SELECT ON tb1 FROM ROLE role1;"#,
         r#"REVOKE SELECT ON tb1 FROM ROLE 'role1';"#,
@@ -911,6 +912,7 @@ fn test_statement_error() {
         r#"drop table :a"#,
         r#"drop table IDENTIFIER(a)"#,
         r#"drop table IDENTIFIER(:a)"#,
+        r#"SHOW GRANTS ON task t1;"#,
     ];
 
     for case in cases {

diff --git a/src/query/ast/tests/it/testdata/statement-error.txt b/src/query/ast/tests/it/testdata/statement-error.txt
@@ -916,3 +916,13 @@ error:
   | while parsing `DROP TABLE [IF EXISTS] [<database>.]<table>`
 
 
+---------- Input ----------
+SHOW GRANTS ON task t1;
+---------- Output ---------
+error: 
+  --> SQL:1:16
+  |
+1 | SHOW GRANTS ON task t1;
+  |                ^^^^ unexpected `task`, expecting `STAGE`, `TABLE`, `DATABASE`, or `UDF`
+
+