refactor: extensible function registry

crates/catalog/src/mutator.rs

@@ -19,6 +19,9 @@ impl CatalogMutator {
         CatalogMutator { client: None }
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.client.is_none()
+    }
     pub fn new(client: Option<MetastoreClientHandle>) -> Self {
         CatalogMutator { client }
     }

crates/glaredb/src/highlighter.rs

@@ -5,7 +5,7 @@ use datafusion::sql::sqlparser::keywords::Keyword;
 use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
 use nu_ansi_term::{Color, Style};
 use reedline::{Highlighter, Hinter, SearchQuery, StyledText, ValidationResult, Validator};
-use sqlbuiltins::functions::FUNCTION_REGISTRY;
+use sqlbuiltins::functions::DEFAULT_BUILTIN_FUNCTIONS;
 
 use crate::local::is_client_cmd;
 
@@ -192,7 +192,8 @@ fn colorize_sql(query: &str, st: &mut StyledText, is_hint: bool) {
                         st.push((new_style().fg(Color::LightGreen), format!("{w}")))
                     }
                     // Functions
-                    other if FUNCTION_REGISTRY.contains(other) => {
+                    // todo: replace this with the session functions
+                    other if DEFAULT_BUILTIN_FUNCTIONS.contains(other) => {
                         st.push((colorize_function(), format!("{w}")));
                     }
                     _ => {
@@ -201,7 +202,7 @@ fn colorize_sql(query: &str, st: &mut StyledText, is_hint: bool) {
                 },
                 // TODO: add more keywords
                 _ => {
-                    if FUNCTION_REGISTRY.contains(&w.value) {
+                    if DEFAULT_BUILTIN_FUNCTIONS.contains(&w.value) {
                         st.push((colorize_function(), format!("{w}")));
                     } else {
                         st.push((new_style(), format!("{w}")))

crates/metastore/src/database.rs

@@ -38,7 +38,7 @@ use sqlbuiltins::builtins::{
     FIRST_NON_STATIC_OID,
     SCHEMA_DEFAULT,
 };
-use sqlbuiltins::functions::{BuiltinFunction, FUNCTION_REGISTRY};
+use sqlbuiltins::functions::{BuiltinFunction, DEFAULT_BUILTIN_FUNCTIONS};
 use sqlbuiltins::validation::{
     validate_database_tunnel_support,
     validate_object_name,
@@ -1033,6 +1033,42 @@ impl State {
                 // Update the new storage size
                 self.deployment.storage_size = update_deployment_storage.new_storage_size;
             }
+            Mutation::CreateFunction(f) => {
+                let schema_id = self.get_schema_id(DEFAULT_SCHEMA)?;
+                let oid = self
+                    .schema_objects
+                    .get(&schema_id)
+                    .and_then(|objs| objs.functions.get(&f.name))
+                    .copied();
+
+                let oid = match oid {
+                    Some(_) => return Err(MetastoreError::DuplicateName(f.name.clone())),
+                    None => self.next_oid(),
+                };
+
+
+                let ent = FunctionEntry {
+                    meta: EntryMeta {
+                        entry_type: EntryType::Function,
+                        id: oid,
+                        parent: schema_id,
+                        name: f.name.clone(),
+                        builtin: false,
+                        external: true,
+                        is_temp: false,
+                    },
+                    func_type: f.function_type,
+                    signature: Some(f.signature),
+                    user_defined: true,
+                };
+
+                self.entries.insert(oid, CatalogEntry::Function(ent))?;
+                self.schema_objects
+                    .entry(schema_id)
+                    .or_default()
+                    .functions
+                    .insert(f.name, oid);
+            }
         };
 
         Ok(())
@@ -1287,17 +1323,17 @@ impl BuiltinCatalog {
         let table_func_ents = Self::builtin_function_to_entries(
             &mut oid_gen,
             schema_id,
-            FUNCTION_REGISTRY.table_funcs_iter(),
+            DEFAULT_BUILTIN_FUNCTIONS.table_funcs_iter(),
         );
         let scalar_func_ents = Self::builtin_function_to_entries(
             &mut oid_gen,
             schema_id,
-            FUNCTION_REGISTRY.scalar_funcs_iter(),
+            DEFAULT_BUILTIN_FUNCTIONS.scalar_funcs_iter(),
         );
         let scalar_udf_ents = Self::builtin_function_to_entries(
             &mut oid_gen,
             schema_id,
-            FUNCTION_REGISTRY.scalar_udfs_iter(),
+            DEFAULT_BUILTIN_FUNCTIONS.scalar_udfs_iter(),
         );
 
         for func_ent in table_func_ents
@@ -1364,6 +1400,7 @@ impl BuiltinCatalog {
                     meta,
                     func_type: func.function_type(),
                     signature: func.signature(),
+                    user_defined: false,
                 })
             }
         }

crates/protogen/proto/metastore/catalog.proto

@@ -188,24 +188,26 @@ message TunnelEntry {
 // available are builtins. We can change this however we need and not have to
 // worry about backwards compatability.
 message FunctionEntry {
-  enum FunctionType {
-    // Unknown catalog entry. We should error if this is encountered.
-    UNKNOWN = 0;
-    AGGREGATE = 1;
-    SCALAR = 2;
-    TABLE_RETURNING = 3;
-  }
-
-  enum RuntimePreference {
-    UNSPECIFIED = 0;
-    LOCAL = 1;
-    REMOTE = 2;
-  }
-
   EntryMeta meta = 1;
   FunctionType func_type = 2;
   reserved 3;  // Function runtime preference (static)
   Signature signature = 4;
+  bool user_defined = 5;
+  // next: 6
+}
+
+enum RuntimePreference {
+  UNSPECIFIED = 0;
+  LOCAL = 1;
+  REMOTE = 2;
+}
+
+enum FunctionType {
+  // Unknown catalog entry. We should error if this is encountered.
+  UNKNOWN = 0;
+  AGGREGATE = 1;
+  SCALAR = 2;
+  TABLE_RETURNING = 3;
 }
 
 message CredentialsEntry {

crates/protogen/proto/metastore/service.proto

@@ -38,8 +38,9 @@ message Mutation {
     CreateCredentials create_credentials = 15;
     DropCredentials drop_credentials = 16;
     UpdateDeploymentStorage update_deployment_storage = 17;
+    CreateFunction create_function = 18;
   }
-  // next: 17
+  // next: 19
 }
 
 message DropDatabase {
@@ -80,6 +81,13 @@ message CreateTable {
   bool or_replace = 5;
 }
 
+message CreateFunction {
+  string name = 1;
+  repeated string aliases = 2;
+  catalog.Signature signature = 3;
+  catalog.FunctionType type = 4;
+}
+
 message CreateExternalTable {
   string schema = 1;
   string name = 2;

crates/protogen/src/metastore/types/catalog.rs

@@ -559,32 +559,32 @@ impl FunctionType {
 impl TryFrom<i32> for FunctionType {
     type Error = ProtoConvError;
     fn try_from(value: i32) -> Result<Self, Self::Error> {
-        catalog::function_entry::FunctionType::try_from(value)
+        catalog::FunctionType::try_from(value)
             .map_err(|_| ProtoConvError::UnknownEnumVariant("FunctionType", value))
             .and_then(|t| t.try_into())
     }
 }
 
-impl TryFrom<catalog::function_entry::FunctionType> for FunctionType {
+impl TryFrom<catalog::FunctionType> for FunctionType {
     type Error = ProtoConvError;
-    fn try_from(value: catalog::function_entry::FunctionType) -> Result<Self, Self::Error> {
+    fn try_from(value: catalog::FunctionType) -> Result<Self, Self::Error> {
         Ok(match value {
-            catalog::function_entry::FunctionType::Unknown => {
+            catalog::FunctionType::Unknown => {
                 return Err(ProtoConvError::ZeroValueEnumVariant("FunctionType"))
             }
-            catalog::function_entry::FunctionType::Aggregate => FunctionType::Aggregate,
-            catalog::function_entry::FunctionType::Scalar => FunctionType::Scalar,
-            catalog::function_entry::FunctionType::TableReturning => FunctionType::TableReturning,
+            catalog::FunctionType::Aggregate => FunctionType::Aggregate,
+            catalog::FunctionType::Scalar => FunctionType::Scalar,
+            catalog::FunctionType::TableReturning => FunctionType::TableReturning,
         })
     }
 }
 
-impl From<FunctionType> for catalog::function_entry::FunctionType {
+impl From<FunctionType> for catalog::FunctionType {
     fn from(value: FunctionType) -> Self {
         match value {
-            FunctionType::Aggregate => catalog::function_entry::FunctionType::Aggregate,
-            FunctionType::Scalar => catalog::function_entry::FunctionType::Scalar,
-            FunctionType::TableReturning => catalog::function_entry::FunctionType::TableReturning,
+            FunctionType::Aggregate => catalog::FunctionType::Aggregate,
+            FunctionType::Scalar => catalog::FunctionType::Scalar,
+            FunctionType::TableReturning => catalog::FunctionType::TableReturning,
         }
     }
 }
@@ -610,32 +610,28 @@ impl RuntimePreference {
 impl TryFrom<i32> for RuntimePreference {
     type Error = ProtoConvError;
     fn try_from(value: i32) -> Result<Self, Self::Error> {
-        let pref = catalog::function_entry::RuntimePreference::try_from(value)
+        let pref = catalog::RuntimePreference::try_from(value)
             .map_err(|_| ProtoConvError::UnknownEnumVariant("RuntimePreference", value))?;
         Ok(pref.into())
     }
 }
 
-impl From<catalog::function_entry::RuntimePreference> for RuntimePreference {
-    fn from(value: catalog::function_entry::RuntimePreference) -> Self {
+impl From<catalog::RuntimePreference> for RuntimePreference {
+    fn from(value: catalog::RuntimePreference) -> Self {
         match value {
-            catalog::function_entry::RuntimePreference::Unspecified => {
-                RuntimePreference::Unspecified
-            }
-            catalog::function_entry::RuntimePreference::Local => RuntimePreference::Local,
-            catalog::function_entry::RuntimePreference::Remote => RuntimePreference::Remote,
+            catalog::RuntimePreference::Unspecified => RuntimePreference::Unspecified,
+            catalog::RuntimePreference::Local => RuntimePreference::Local,
+            catalog::RuntimePreference::Remote => RuntimePreference::Remote,
         }
     }
 }
 
-impl From<RuntimePreference> for catalog::function_entry::RuntimePreference {
+impl From<RuntimePreference> for catalog::RuntimePreference {
     fn from(value: RuntimePreference) -> Self {
         match value {
-            RuntimePreference::Unspecified => {
-                catalog::function_entry::RuntimePreference::Unspecified
-            }
-            RuntimePreference::Local => catalog::function_entry::RuntimePreference::Local,
-            RuntimePreference::Remote => catalog::function_entry::RuntimePreference::Remote,
+            RuntimePreference::Unspecified => catalog::RuntimePreference::Unspecified,
+            RuntimePreference::Local => catalog::RuntimePreference::Local,
+            RuntimePreference::Remote => catalog::RuntimePreference::Remote,
         }
     }
 }
@@ -645,6 +641,7 @@ pub struct FunctionEntry {
     pub meta: EntryMeta,
     pub func_type: FunctionType,
     pub signature: Option<Signature>,
+    pub user_defined: bool,
 }
 
 impl TryFrom<catalog::FunctionEntry> for FunctionEntry {
@@ -655,6 +652,7 @@ impl TryFrom<catalog::FunctionEntry> for FunctionEntry {
             meta,
             func_type: value.func_type.try_into()?,
             signature: value.signature.map(|s| s.try_into()).transpose()?,
+            user_defined: value.user_defined,
         })
     }
 }
@@ -801,11 +799,12 @@ impl TryFrom<catalog::Signature> for Signature {
 
 impl From<FunctionEntry> for catalog::FunctionEntry {
     fn from(value: FunctionEntry) -> Self {
-        let func_type: catalog::function_entry::FunctionType = value.func_type.into();
+        let func_type: catalog::FunctionType = value.func_type.into();
         catalog::FunctionEntry {
             meta: Some(value.meta.into()),
             func_type: func_type as i32,
             signature: value.signature.map(|s| s.into()),
+            user_defined: value.user_defined,
         }
     }
 }