register UDF in shared session context

Author: nikhilsinhaparseable

src/lib.rs

@@ -68,8 +68,8 @@ use once_cell::sync::Lazy;
 pub use openid;
 use parseable::PARSEABLE;
 use reqwest::{Client, ClientBuilder};
+pub use {clap, tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
 pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
-pub use {tracing_actix_web, tracing_opentelemetry, tracing_subscriber, clap};
 
 // It is very unlikely that panic will occur when dealing with locks.
 pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";

src/metastore/metastore_traits.rs

@@ -229,7 +229,7 @@ pub trait Metastore: std::fmt::Debug + Send + Sync {
         stream_name: &str,
         get_base: bool,
         tenant_id: &Option<String>,
-        is_migration: bool
+        is_migration: bool,
     ) -> Result<Bytes, MetastoreError>;
     async fn put_stream_json(
         &self,
@@ -247,7 +247,7 @@ pub trait Metastore: std::fmt::Debug + Send + Sync {
         stream_name: &str,
         mode: Option<Mode>,
         tenant_id: &Option<String>,
-        is_migration: bool
+        is_migration: bool,
     ) -> Result<Vec<Bytes>, MetastoreError>;
 
     /// Fetch manifests only for the explicitly requested date keys

src/metastore/metastores/object_store_metastore.rs

@@ -892,7 +892,7 @@ impl Metastore for ObjectStoreMetastore {
         stream_name: &str,
         mode: Option<Mode>,
         tenant_id: &Option<String>,
-        _is_migration: bool
+        _is_migration: bool,
     ) -> Result<Vec<Bytes>, MetastoreError> {
         let root = tenant_id.as_deref().unwrap_or("");
         let path = RelativePathBuf::from_iter([root, stream_name, STREAM_ROOT_DIRECTORY]);

src/query/mod.rs

@@ -32,7 +32,7 @@ use datafusion::execution::{
 };
 use datafusion::logical_expr::expr::Alias;
 use datafusion::logical_expr::{
-    Aggregate, Explain, Filter, LogicalPlan, PlanType, Projection, ToStringifiedPlan,
+    Aggregate, Explain, Filter, LogicalPlan, PlanType, Projection, ScalarUDF, ToStringifiedPlan,
 };
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
@@ -112,6 +112,13 @@ impl InMemorySessionContext {
         ctx.clone()
     }
 
+    pub fn register_udf(&self, udf: ScalarUDF) {
+        self.session_context
+            .write()
+            .expect("SessionContext should be writeable")
+            .register_udf(udf);
+    }
+
     pub fn add_schema(&self, tenant_id: &str) {
         self.session_context
             .write()