feat: ensure procedure_state and migrate_region can be only called un…

…der greptime catalog

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

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;

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}")
+            }
+        );
+    }};
+}

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 {

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])?;

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,
         }
     }