fix: dynamically get python udf from scripts table

Signed-off-by: xxxuuu <[email protected]>
GreptimeTeam · Apr 22, 2024 · 9a06553 · 9a06553
1 parent 8f2ce4a
commit 9a06553
Show file tree

Hide file tree

Showing 24 changed files with 610 additions and 324 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs
@@ -275,11 +275,12 @@ impl StartCommand {
             cached_meta_backend.clone(),
             catalog_manager,
             Arc::new(DatanodeClients::default()),
-            meta_client,
+            meta_client.clone(),
         )
         .with_plugin(plugins.clone())
         .with_cache_invalidator(multi_cache_invalidator)
         .with_heartbeat_task(heartbeat_task)
+        .with_meta_client(meta_client.clone())
         .try_build()
         .await
         .context(StartFrontendSnafu)?;

diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs
@@ -17,6 +17,7 @@ use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 
 use once_cell::sync::Lazy;
+use session::context::QueryContextRef;
 
 use crate::function::FunctionRef;
 use crate::scalars::aggregate::{AggregateFunctionMetaRef, AggregateFunctions};
@@ -28,10 +29,17 @@ use crate::scalars::timestamp::TimestampFunction;
 use crate::system::SystemFunction;
 use crate::table::TableFunction;
 
+pub trait FunctionProvider {
+    fn get_function(&self, name: &str, query_ctx: QueryContextRef) -> Option<FunctionRef>;
+}
+
+pub type FunctionProviderRef = Arc<dyn FunctionProvider + Sync + Send>;
+
 #[derive(Default)]
 pub struct FunctionRegistry {
     functions: RwLock<HashMap<String, FunctionRef>>,
     aggregate_functions: RwLock<HashMap<String, AggregateFunctionMetaRef>>,
+    dynamic_providers: RwLock<Vec<FunctionProviderRef>>,
 }
 
 impl FunctionRegistry {
@@ -51,6 +59,10 @@ impl FunctionRegistry {
             .insert(func.name(), func);
     }
 
+    pub fn register_provider(&self, provider: FunctionProviderRef) {
+        self.dynamic_providers.write().unwrap().push(provider);
+    }
+
     pub fn get_aggr_function(&self, name: &str) -> Option<AggregateFunctionMetaRef> {
         self.aggregate_functions.read().unwrap().get(name).cloned()
     }
@@ -59,6 +71,20 @@ impl FunctionRegistry {
         self.functions.read().unwrap().get(name).cloned()
     }
 
+    pub fn get_function_fallback(
+        &self,
+        name: &str,
+        query_ctx: QueryContextRef,
+    ) -> Option<FunctionRef> {
+        for provider in self.dynamic_providers.read().unwrap().iter() {
+            let fun = provider.get_function(name, query_ctx.clone());
+            if fun.is_some() {
+                return fun;
+            }
+        }
+        None
+    }
+
     pub fn functions(&self) -> Vec<FunctionRef> {
         self.functions.read().unwrap().values().cloned().collect()
     }

diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
@@ -151,6 +151,7 @@ impl Instance {
             .enable_heartbeat()
             .enable_procedure()
             .enable_access_cluster_info()
+            .enable_lock()
             .channel_manager(channel_manager)
             .ddl_channel_manager(ddl_channel_manager)
             .build();

diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs
@@ -21,6 +21,7 @@ use common_meta::datanode_manager::DatanodeManagerRef;
 use common_meta::ddl::ProcedureExecutorRef;
 use common_meta::key::TableMetadataManager;
 use common_meta::kv_backend::KvBackendRef;
+use meta_client::client::MetaClientRef;
 use operator::delete::Deleter;
 use operator::insert::Inserter;
 use operator::procedure::ProcedureServiceOperator;
@@ -46,6 +47,7 @@ pub struct FrontendBuilder {
     plugins: Option<Plugins>,
     procedure_executor: ProcedureExecutorRef,
     heartbeat_task: Option<HeartbeatTask>,
+    meta_client: Option<MetaClientRef>,
 }
 
 impl FrontendBuilder {
@@ -63,6 +65,7 @@ impl FrontendBuilder {
             plugins: None,
             procedure_executor,
             heartbeat_task: None,
+            meta_client: None,
         }
     }
 
@@ -87,6 +90,13 @@ impl FrontendBuilder {
         }
     }
 
+    pub fn with_meta_client(self, meta_client: MetaClientRef) -> Self {
+        Self {
+            meta_client: Some(meta_client),
+            ..self
+        }
+    }
+
     pub async fn try_build(self) -> Result<Instance> {
         let kv_backend = self.kv_backend;
         let datanode_manager = self.datanode_manager;
@@ -137,7 +147,12 @@ impl FrontendBuilder {
         .query_engine();
 
         let script_executor = Arc::new(
-            ScriptExecutor::new(self.catalog_manager.clone(), query_engine.clone()).await?,
+            ScriptExecutor::new(
+                self.catalog_manager.clone(),
+                query_engine.clone(),
+                self.meta_client,
+            )
+            .await?,
         );
 
         let statement_executor = Arc::new(StatementExecutor::new(