From 83cc11ffb1ed3f28d197bce670ef4c3e5916b890 Mon Sep 17 00:00:00 2001
From: Dennis Zhuang <killme2008@gmail.com>
Date: Wed, 21 Feb 2024 15:56:45 +0800
Subject: [PATCH] feat: ensure procedure_state and migrate_region can be only
 called under greptime catalog

---
 Cargo.lock                                        |  1 +
 src/common/function/Cargo.toml                    |  1 +
 src/common/function/src/lib.rs                    |  1 +
 src/common/function/src/macros.rs                 | 13 +++++++++++++
 src/common/function/src/system/procedure_state.rs |  2 ++
 src/common/function/src/table/migrate_region.rs   |  2 ++
 src/common/query/src/error.rs                     |  5 +++++
 7 files changed, 25 insertions(+)
 create mode 100644 src/common/function/src/macros.rs

diff --git a/Cargo.lock b/Cargo.lock
index 4e1ecc4c649e..c9cbd2056adc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1785,6 +1785,7 @@ dependencies = [
  "arc-swap",
  "async-trait",
  "chrono-tz 0.6.3",
+ "common-catalog",
  "common-error",
  "common-macro",
  "common-meta",
diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml
index 7b716f07ab6f..a57eacefad4b 100644
--- a/src/common/function/Cargo.toml
+++ b/src/common/function/Cargo.toml
@@ -9,6 +9,7 @@ api.workspace = true
 arc-swap = "1.0"
 async-trait.workspace = true
 chrono-tz = "0.6"
+common-catalog.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
 common-meta.workspace = true
diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs
index 10fbf13a7a05..1d37d7068c98 100644
--- a/src/common/function/src/lib.rs
+++ b/src/common/function/src/lib.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod macros;
 pub mod scalars;
 mod system;
 mod table;
diff --git a/src/common/function/src/macros.rs b/src/common/function/src/macros.rs
new file mode 100644
index 000000000000..e94168c4043f
--- /dev/null
+++ b/src/common/function/src/macros.rs
@@ -0,0 +1,13 @@
+/// Ensure current function is invokded under `greptime` catalog.
+#[macro_export]
+macro_rules! ensure_greptime {
+    ($func_ctx: expr) => {{
+        use common_catalog::consts::DEFAULT_CATALOG_NAME;
+        snafu::ensure!(
+            $func_ctx.query_ctx.current_catalog() == DEFAULT_CATALOG_NAME,
+            common_query::error::PermissionDeniedSnafu {
+                err_msg: format!("current catalog is not {DEFAULT_CATALOG_NAME}")
+            }
+        );
+    }};
+}
diff --git a/src/common/function/src/system/procedure_state.rs b/src/common/function/src/system/procedure_state.rs
index 5be539258859..4f6305078465 100644
--- a/src/common/function/src/system/procedure_state.rs
+++ b/src/common/function/src/system/procedure_state.rs
@@ -69,6 +69,8 @@ impl Function for ProcedureStateFunction {
     }
 
     fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        crate::ensure_greptime!(func_ctx);
+
         ensure!(
             columns.len() == 1,
             InvalidFuncArgsSnafu {
diff --git a/src/common/function/src/table/migrate_region.rs b/src/common/function/src/table/migrate_region.rs
index 0c657a4120f6..6447c6de6b3d 100644
--- a/src/common/function/src/table/migrate_region.rs
+++ b/src/common/function/src/table/migrate_region.rs
@@ -78,6 +78,8 @@ impl Function for MigrateRegionFunction {
     }
 
     fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
+        crate::ensure_greptime!(func_ctx);
+
         let (region_ids, from_peers, to_peers, replay_timeouts) = match columns.len() {
             3 => {
                 let region_ids = cast_u64_vector(&columns[0])?;
diff --git a/src/common/query/src/error.rs b/src/common/query/src/error.rs
index 7b16ef6590af..758ec214b6a1 100644
--- a/src/common/query/src/error.rs
+++ b/src/common/query/src/error.rs
@@ -192,6 +192,9 @@ pub enum Error {
 
     #[snafu(display("Invalid function args: {}", err_msg))]
     InvalidFuncArgs { err_msg: String, location: Location },
+
+    #[snafu(display("Permission denied: {}", err_msg))]
+    PermissionDenied { err_msg: String, location: Location },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -234,6 +237,8 @@ impl ErrorExt for Error {
             Error::ProcedureService { source, .. } | Error::TableMutation { source, .. } => {
                 source.status_code()
             }
+
+            Error::PermissionDenied { .. } => StatusCode::PermissionDenied,
         }
     }