feat: make scripts table work again

src/frontend/Cargo.toml

@@ -11,6 +11,7 @@ testing = []
 
 [dependencies]
 api = { workspace = true }
+arc-swap = "1.0"
 arrow-flight.workspace = true
 async-compat = "0.2"
 async-stream.workspace = true

src/frontend/src/instance.rs

@@ -393,7 +393,7 @@ impl FrontendInstance for Instance {
             heartbeat_task.start().await?;
         }
 
-        self.script_executor.start(self).await?;
+        self.script_executor.start(self)?;
 
         futures::future::try_join_all(self.servers.values().map(start_server))
             .await

src/frontend/src/instance/script.rs

@@ -18,6 +18,7 @@ use async_trait::async_trait;
 use common_query::Output;
 use common_telemetry::timer;
 use servers::query_handler::ScriptHandler;
+use session::context::QueryContextRef;
 
 use crate::instance::Instance;
 use crate::metrics;
@@ -26,25 +27,25 @@ use crate::metrics;
 impl ScriptHandler for Instance {
     async fn insert_script(
         &self,
-        schema: &str,
+        query_ctx: QueryContextRef,
         name: &str,
         script: &str,
     ) -> servers::error::Result<()> {
         let _timer = timer!(metrics::METRIC_HANDLE_SCRIPTS_ELAPSED);
         self.script_executor
-            .insert_script(schema, name, script)
+            .insert_script(query_ctx, name, script)
             .await
     }
 
     async fn execute_script(
         &self,
-        schema: &str,
+        query_ctx: QueryContextRef,
         name: &str,
         params: HashMap<String, String>,
     ) -> servers::error::Result<Output> {
         let _timer = timer!(metrics::METRIC_RUN_SCRIPT_ELAPSED);
         self.script_executor
-            .execute_script(schema, name, params)
+            .execute_script(query_ctx, name, params)
             .await
     }
 }

src/frontend/src/script.rs

@@ -13,12 +13,17 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use catalog::CatalogManagerRef;
 use common_query::Output;
 use query::QueryEngineRef;
+use servers::query_handler::grpc::GrpcQueryHandler;
+use session::context::QueryContextRef;
 
-use crate::error::Result;
+use crate::error::{Error, Result};
+
+type FrontendGrpcQueryHandlerRef = Arc<dyn GrpcQueryHandler<Error = Error> + Send + Sync>;
 
 #[cfg(not(feature = "python"))]
 mod dummy {
@@ -34,13 +39,13 @@ mod dummy {
             Ok(Self {})
         }
 
-        pub async fn start(&self) -> Result<()> {
+        pub fn start(&self, instance: &Instance) -> Result<()> {
             Ok(())
         }
 
         pub async fn insert_script(
             &self,
-            _schema: &str,
+            _query_ctx: QueryContextRef,
             _name: &str,
             _script: &str,
         ) -> servers::error::Result<()> {
@@ -49,7 +54,7 @@ mod dummy {
 
         pub async fn execute_script(
             &self,
-            _schema: &str,
+            _query_ctx: QueryContextRef,
             _name: &str,
             _params: HashMap<String, String>,
         ) -> servers::error::Result<Output> {
@@ -63,10 +68,11 @@ mod python {
     use api::v1::ddl_request::Expr;
     use api::v1::greptime_request::Request;
     use api::v1::{CreateTableExpr, DdlRequest};
+    use arc_swap::ArcSwap;
     use catalog::RegisterSystemTableRequest;
     use common_error::ext::BoxedError;
     use common_meta::table_name::TableName;
-    use common_telemetry::logging::error;
+    use common_telemetry::{error, info};
     use operator::expr_factory;
     use script::manager::ScriptManager;
     use servers::query_handler::grpc::GrpcQueryHandler;
@@ -78,30 +84,76 @@ mod python {
     use crate::error::{CatalogSnafu, InvalidSystemTableDefSnafu, TableNotFoundSnafu};
     use crate::instance::Instance;
 
+    /// A placeholder for the real gRPC handler.
+    /// It is temporary and will be replaced soon.
+    struct DummyHandler;
+
+    impl DummyHandler {
+        fn arc() -> Arc<Self> {
+            Arc::new(Self {})
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl GrpcQueryHandler for DummyHandler {
+        type Error = Error;
+
+        async fn do_query(
+            &self,
+            _query: Request,
+            _ctx: QueryContextRef,
+        ) -> std::result::Result<Output, Self::Error> {
+            unreachable!();
+        }
+    }
+
     pub struct ScriptExecutor {
-        script_manager: ScriptManager,
+        script_manager: ScriptManager<Error>,
+        grpc_handler: ArcSwap<FrontendGrpcQueryHandlerRef>,
+        catalog_manager: CatalogManagerRef,
     }
 
     impl ScriptExecutor {
         pub async fn new(
             catalog_manager: CatalogManagerRef,
             query_engine: QueryEngineRef,
         ) -> Result<Self> {
+            let grpc_handler = DummyHandler::arc();
             Ok(Self {
-                script_manager: ScriptManager::new(catalog_manager, query_engine)
+                grpc_handler: ArcSwap::new(Arc::new(grpc_handler.clone() as _)),
+                script_manager: ScriptManager::new(grpc_handler as _, query_engine)
                     .await
                     .context(crate::error::StartScriptManagerSnafu)?,
+                catalog_manager,
             })
         }
 
-        pub async fn start(&self, instance: &Instance) -> Result<()> {
+        pub fn start(&self, instance: &Instance) -> Result<()> {
+            let handler = Arc::new(instance.clone());
+            self.grpc_handler.store(Arc::new(handler.clone() as _));
+            self.script_manager
+                .start(handler)
+                .context(crate::error::StartScriptManagerSnafu)?;
+
+            Ok(())
+        }
+
+        /// Create scripts table for the specific catalog if it's not exists.
+        /// The function is idempotent and safe to be called more than once for the same catalog
+        async fn create_scripts_table_if_need(&self, catalog: &str) -> Result<()> {
+            let scripts_table = self.script_manager.get_scripts_table(catalog);
+
+            if scripts_table.is_some() {
+                return Ok(());
+            }
+
             let RegisterSystemTableRequest {
                 create_table_request: request,
                 open_hook,
-            } = self.script_manager.create_table_request();
+            } = self.script_manager.create_table_request(catalog);
 
-            if let Some(table) = instance
-                .catalog_manager()
+            if let Some(table) = self
+                .catalog_manager
                 .table(
                     &request.catalog_name,
                     &request.schema_name,
@@ -111,9 +163,11 @@ mod python {
                 .context(CatalogSnafu)?
             {
                 if let Some(open_hook) = open_hook {
-                    (open_hook)(table).await.context(CatalogSnafu)?;
+                    (open_hook)(table.clone()).await.context(CatalogSnafu)?;
                 }
 
+                self.script_manager.insert_scripts_table(catalog, table);
+
                 return Ok(());
             }
 
@@ -125,7 +179,9 @@ mod python {
 
             let expr = Self::create_table_expr(request)?;
 
-            let _ = instance
+            let _ = self
+                .grpc_handler
+                .load()
                 .do_query(
                     Request::Ddl(DdlRequest {
                         expr: Some(Expr::CreateTable(expr)),
@@ -134,8 +190,8 @@ mod python {
                 )
                 .await?;
 
-            let table = instance
-                .catalog_manager()
+            let table = self
+                .catalog_manager
                 .table(
                     &table_name.catalog_name,
                     &table_name.schema_name,
@@ -148,9 +204,16 @@ mod python {
                 })?;
 
             if let Some(open_hook) = open_hook {
-                (open_hook)(table).await.context(CatalogSnafu)?;
+                (open_hook)(table.clone()).await.context(CatalogSnafu)?;
             }
 
+            info!(
+                "Created scripts table {}.",
+                table.table_info().full_table_name()
+            );
+
+            self.script_manager.insert_scripts_table(catalog, table);
+
             Ok(())
         }
 
@@ -196,16 +259,31 @@ mod python {
 
         pub async fn insert_script(
             &self,
-            schema: &str,
+            query_ctx: QueryContextRef,
             name: &str,
             script: &str,
         ) -> servers::error::Result<()> {
+            self.create_scripts_table_if_need(query_ctx.current_catalog())
+                .await
+                .map_err(|e| {
+                    error!(e; "Failed to create scripts table");
+                    servers::error::InternalSnafu {
+                        err_msg: e.to_string(),
+                    }
+                    .build()
+                })?;
+
             let _s = self
                 .script_manager
-                .insert_and_compile(schema, name, script)
+                .insert_and_compile(
+                    query_ctx.current_catalog(),
+                    query_ctx.current_schema(),
+                    name,
+                    script,
+                )
                 .await
                 .map_err(|e| {
-                    error!(e; "Instance failed to insert script");
+                    error!(e; "Failed to insert script");
                     BoxedError::new(e)
                 })
                 .context(servers::error::InsertScriptSnafu { name })?;
@@ -215,15 +293,30 @@ mod python {
 
         pub async fn execute_script(
             &self,
-            schema: &str,
+            query_ctx: QueryContextRef,
             name: &str,
             params: HashMap<String, String>,
         ) -> servers::error::Result<Output> {
+            self.create_scripts_table_if_need(query_ctx.current_catalog())
+                .await
+                .map_err(|e| {
+                    error!(e; "Failed to create scripts table");
+                    servers::error::InternalSnafu {
+                        err_msg: e.to_string(),
+                    }
+                    .build()
+                })?;
+
             self.script_manager
-                .execute(schema, name, params)
+                .execute(
+                    query_ctx.current_catalog(),
+                    query_ctx.current_schema(),
+                    name,
+                    params,
+                )
                 .await
                 .map_err(|e| {
-                    error!(e; "Instance failed to execute script");
+                    error!(e; "Failed to execute script");
                     BoxedError::new(e)
                 })
                 .context(servers::error::ExecuteScriptSnafu { name })

src/script/Cargo.toml

@@ -24,6 +24,8 @@ python = [
 ]
 
 [dependencies]
+api.workspace = true
+arc-swap = "1.0"
 arrow.workspace = true
 async-trait.workspace = true
 catalog = { workspace = true }
@@ -62,6 +64,7 @@ rustpython-vm = { git = "https://github.com/discord9/RustPython", optional = tru
     "default",
     "codegen",
 ] }
+servers.workspace = true
 session = { workspace = true }
 snafu = { version = "0.7", features = ["backtraces"] }
 sql = { workspace = true }
@@ -75,6 +78,7 @@ common-test-util = { workspace = true }
 criterion = { version = "0.4", features = ["html_reports", "async_tokio"] }
 log-store = { workspace = true }
 mito = { workspace = true }
+operator.workspace = true
 rayon = "1.0"
 ron = "0.7"
 serde = { version = "1.0", features = ["derive"] }