feat(mapper): rewrite std SQL functions for EQL cols

freshtonic · freshtonic · commit 3c20b44afb18 · 2025-05-08T23:12:04.000+10:00
e.g. `jsonb_query_path(eql_col, selector)` on an EQL column should be transformed to `eql_v1.jsonb_query_path(..)`
diff --git a/packages/eql-mapper/src/inference/infer_type_impls/function.rs b/packages/eql-mapper/src/inference/infer_type_impls/function.rs
@@ -2,7 +2,7 @@ use eql_mapper_macros::trace_infer;
 use sqltk::parser::ast::{Function, FunctionArguments};
 
 use crate::{
-    get_type_signature_for_special_cased_sql_function, inference::infer_type::InferType,
+    get_sql_function_def, inference::infer_type::InferType,
     CompoundIdent, FunctionSig, TypeError, TypeInferencer,
 };
 
@@ -23,9 +23,9 @@ impl<'ast> InferType<'ast, Function> for TypeInferencer<'ast> {
         let Function { name, args, .. } = function;
         let fn_name = CompoundIdent::from(&name.0);
 
-        match get_type_signature_for_special_cased_sql_function(&fn_name, args) {
-            Some(sig) => {
-                sig.instantiate(&*self).apply_constraints(self, function)?;
+        match get_sql_function_def(&fn_name, args) {
+            Some(sql_fn) => {
+                sql_fn.sig.instantiate(&*self).apply_constraints(self, function)?;
             }
             None => {
                 FunctionSig::instantiate_native(function).apply_constraints(self, function)?;
diff --git a/packages/eql-mapper/src/inference/sql_fn_macros.rs b/packages/eql-mapper/src/inference/sql_fn_macros.rs
@@ -21,10 +21,35 @@ macro_rules! sql_fn_args {
 
 #[macro_export]
 macro_rules! sql_fn {
+    ($name:ident $args:tt -> $return_kind:ident, rewrite) => {
+        $crate::SqlFunction::new(
+            stringify!($name),
+            FunctionSig::new($crate::sql_fn_args!($args), $crate::to_kind!($return_kind)),
+            crate::RewriteRule::AsEqlFunction,
+        )
+    };
+
     ($name:ident $args:tt -> $return_kind:ident) => {
         $crate::SqlFunction::new(
             stringify!($name),
             FunctionSig::new($crate::sql_fn_args!($args), $crate::to_kind!($return_kind)),
+            crate::RewriteRule::Ignore,
         )
     };
-}
+
+    ($schema:ident . $name:ident $args:tt -> $return_kind:ident, rewrite) => {
+        $crate::SqlFunction::new(
+            stringify!($schema . $name),
+            FunctionSig::new($crate::sql_fn_args!($args), $crate::to_kind!($return_kind)),
+            crate::RewriteRule::AsEqlFunction,
+        )
+    };
+
+    ($schema:ident . $name:ident $args:tt -> $return_kind:ident) => {
+        $crate::SqlFunction::new(
+            stringify!($schema . $name),
+            FunctionSig::new($crate::sql_fn_args!($args), $crate::to_kind!($return_kind)),
+            crate::RewriteRule::Ignore,
+        )
+    };
+}
diff --git a/packages/eql-mapper/src/inference/sql_functions.rs b/packages/eql-mapper/src/inference/sql_functions.rs
@@ -17,12 +17,22 @@ use super::TypeError;
 ///
 /// See [`SQL_FUNCTION_SIGNATURES`].
 #[derive(Debug)]
-pub(crate) struct SqlFunction(CompoundIdent, FunctionSig);
+pub(crate) struct SqlFunction {
+    pub(crate) name: CompoundIdent,
+    pub(crate) sig: FunctionSig,
+    pub(crate) rewrite_rule: RewriteRule,
+}
+
+#[derive(Debug)]
+pub(crate) enum RewriteRule {
+    Ignore,
+    AsEqlFunction,
+}
 
 /// A representation of the type of an argument or return type in a SQL function.
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 pub(crate) enum Kind {
-    /// A type that mjust be a native type
+    /// A type that must be a native type
     Native,
 
     /// A type that can be a native or EQL type. The `str` is the generic variable name.
@@ -178,8 +188,12 @@ fn get_function_arg_expr(fn_arg: &FunctionArg) -> &FunctionArgExpr {
 }
 
 impl SqlFunction {
-    fn new(ident: &str, sig: FunctionSig) -> Self {
-        Self(CompoundIdent::from(ident), sig)
+    fn new(ident: &str, sig: FunctionSig, rewrite_rule: RewriteRule) -> Self {
+        Self {
+            name: CompoundIdent::from(ident),
+            sig,
+            rewrite_rule,
+        }
     }
 }
 
@@ -202,38 +216,51 @@ impl From<&Vec<Ident>> for CompoundIdent {
 }
 
 /// SQL functions that are handled with special case type checking rules.
-static SQL_FUNCTION_SIGNATURES: LazyLock<HashMap<CompoundIdent, Vec<FunctionSig>>> =
-    LazyLock::new(|| {
-        // Notation: a single uppercase letter denotes an unknown type. Matching letters in a signature will be assigned
-        // *the same type variable* and thus must resolve to the same type. (🙏 Haskell)
-        //
-        // Eventually we should type check EQL types against their configured indexes instead of leaving that to the EQL
-        // extension in the database. I can imagine supporting type bounds in signatures here, such as: `T: Eq`
-        let sql_fns = vec![
-            sql_fn!(count(T) -> NATIVE),
-            sql_fn!(min(T) -> T),
-            sql_fn!(max(T) -> T),
-            sql_fn!(jsonb_path_query(T, T) -> T),
-            sql_fn!(jsonb_path_query_first(T, T) -> T),
-            sql_fn!(jsonb_path_exists(T, T) -> T),
-        ];
-
-        let mut sql_fns_by_name: HashMap<CompoundIdent, Vec<FunctionSig>> = HashMap::new();
-
-        for (key, chunk) in &sql_fns.into_iter().chunk_by(|sql_fn| sql_fn.0.clone()) {
-            sql_fns_by_name.insert(
-                key.clone(),
-                chunk.into_iter().map(|sql_fn| sql_fn.1).collect(),
-            );
-        }
+static SQL_FUNCTIONS: LazyLock<HashMap<CompoundIdent, Vec<SqlFunction>>> = LazyLock::new(|| {
+    // Notation: a single uppercase letter denotes an unknown type. Matching letters in a signature will be assigned
+    // *the same type variable* and thus must resolve to the same type. (🙏 Haskell)
+    //
+    // Eventually we should type check EQL types against their configured indexes instead of leaving that to the EQL
+    // extension in the database. I can imagine supporting type bounds in signatures here, such as: `T: Eq`
+    let sql_fns = vec![
+        // TODO: when search_path support is added to the resolver we should change these
+        // to their fully-qualified names.
+        sql_fn!(count(T) -> NATIVE),
+        sql_fn!(min(T) -> T, rewrite),
+        sql_fn!(max(T) -> T, rewrite),
+        sql_fn!(jsonb_path_query(T, T) -> T, rewrite),
+        sql_fn!(jsonb_path_query_first(T, T) -> T, rewrite),
+        sql_fn!(jsonb_path_exists(T, T) -> T, rewrite),
+        sql_fn!(jsonb_array_length(T) -> T, rewrite),
+        sql_fn!(jsonb_array_elements(T) -> T, rewrite),
+        sql_fn!(jsonb_array_elements_text(T) -> T, rewrite),
+        // These are typings for when customer SQL already contains references to EQL functions.
+        // They must be type checked but not rewritten.
+        sql_fn!(eql_v1.min(T) -> T),
+        sql_fn!(eql_v1.max(T) -> T),
+        sql_fn!(eql_v1.jsonb_path_query(T, T) -> T),
+        sql_fn!(eql_v1.jsonb_path_query_first(T, T) -> T),
+        sql_fn!(eql_v1.jsonb_path_exists(T, T) -> T),
+        sql_fn!(eql_v1.jsonb_array_length(T) -> T),
+        sql_fn!(eql_v1.jsonb_array_elements(T) -> T),
+        sql_fn!(eql_v1.jsonb_array_elements_text(T) -> T),
+    ];
+
+    let mut sql_fns_by_name: HashMap<CompoundIdent, Vec<SqlFunction>> = HashMap::new();
+
+    for (key, chunk) in &sql_fns.into_iter().chunk_by(|sql_fn| sql_fn.name.clone()) {
+        sql_fns_by_name.insert(key.clone(), chunk.into_iter().collect());
+    }
 
-        sql_fns_by_name
-    });
+    sql_fns_by_name
+});
 
-pub(crate) fn get_type_signature_for_special_cased_sql_function(
+pub(crate) fn get_sql_function_def(
     fn_name: &CompoundIdent,
     args: &FunctionArguments,
-) -> Option<&'static FunctionSig> {
-    let sigs = SQL_FUNCTION_SIGNATURES.get(fn_name)?;
-    sigs.iter().find(|sig| sig.is_applicable_to_args(args))
+) -> Option<&'static SqlFunction> {
+    let sql_fns = SQL_FUNCTIONS.get(fn_name)?;
+    sql_fns
+        .iter()
+        .find(|sql_fn| sql_fn.sig.is_applicable_to_args(args))
 }
diff --git a/packages/eql-mapper/src/lib.rs b/packages/eql-mapper/src/lib.rs
@@ -1391,6 +1391,42 @@ mod test {
         }
     }
 
+    #[test]
+    fn rewrite_standard_sql_fns_on_eql_types() {
+        // init_tracing();
+        let schema = resolver(schema! {
+            tables: {
+                employees: {
+                    id (PK),
+                    eql_col (EQL),
+                    native_col,
+                }
+            }
+        });
+
+        let statement = parse("
+            SELECT jsonb_path_query(eql_col, '$.secret'), jsonb_path_query(native_col, '$.not-secret') FROM employees
+        ");
+
+        match type_check(schema, &statement) {
+            Ok(typed) => {
+                match typed.transform(test_helpers::dummy_encrypted_json_selector(
+                    &statement,
+                    ast::Value::SingleQuotedString("$.secret".into()),
+                )) {
+                    Ok(statement) => {
+                        assert_eq!(
+                            statement.to_string(),
+                            "SELECT eql_v1.jsonb_path_query(eql_col, ROW('<encrypted-selector($.secret)>'::JSONB)), jsonb_path_query(native_col, '$.not-secret') FROM employees"
+                        );
+                    }
+                    Err(err) => panic!("transformation failed: {err}"),
+                }
+            }
+            Err(err) => panic!("type check failed: {err}"),
+        }
+    }
+
     #[test]
     fn jsonb_operator_arrow() {
         test_jsonb_operator("->");
diff --git a/packages/eql-mapper/src/test_helpers.rs b/packages/eql-mapper/src/test_helpers.rs
@@ -93,6 +93,7 @@ pub(crate) fn find_nodekey_for_value_node(
 
     visitor.found
 }
+
 #[macro_export]
 macro_rules! col {
     ((NATIVE)) => {
diff --git a/packages/eql-mapper/src/transformation_rules/mod.rs b/packages/eql-mapper/src/transformation_rules/mod.rs
@@ -15,6 +15,7 @@ mod fail_on_placeholder_change;
 mod group_by_eql_col;
 mod preserve_effective_aliases;
 mod replace_plaintext_eql_literals;
+mod rewrite_standard_sql_fns_on_eql_types;
 mod use_equivalent_eql_fns_on_eql_types;
 mod wrap_eql_cols_in_order_by_with_ore_fn;
 mod wrap_grouped_eql_col_in_aggregate_fn;
@@ -24,6 +25,7 @@ use std::marker::PhantomData;
 pub(crate) use fail_on_placeholder_change::*;
 pub(crate) use group_by_eql_col::*;
 pub(crate) use preserve_effective_aliases::*;
+pub(crate) use rewrite_standard_sql_fns_on_eql_types::*;
 pub(crate) use replace_plaintext_eql_literals::*;
 pub(crate) use use_equivalent_eql_fns_on_eql_types::*;
 pub(crate) use wrap_eql_cols_in_order_by_with_ore_fn::*;
diff --git a/packages/eql-mapper/src/transformation_rules/rewrite_standard_sql_fns_on_eql_types.rs b/packages/eql-mapper/src/transformation_rules/rewrite_standard_sql_fns_on_eql_types.rs
@@ -0,0 +1,75 @@
+use std::mem;
+use std::{collections::HashMap, sync::Arc};
+
+use sqltk::parser::ast::{Expr, Function, Ident, ObjectName};
+use sqltk::{AsNodeKey, NodeKey, NodePath, Visitable};
+
+use crate::{
+    get_sql_function_def, CompoundIdent, EqlMapperError, RewriteRule, SqlFunction, Type, Value,
+};
+
+use super::TransformationRule;
+
+#[derive(Debug)]
+pub struct RewriteStandardSqlFnsOnEqlTypes<'ast> {
+    node_types: Arc<HashMap<NodeKey<'ast>, Type>>,
+}
+
+impl<'ast> RewriteStandardSqlFnsOnEqlTypes<'ast> {
+    pub fn new(node_types: Arc<HashMap<NodeKey<'ast>, Type>>) -> Self {
+        Self { node_types }
+    }
+}
+
+impl<'ast> TransformationRule<'ast> for RewriteStandardSqlFnsOnEqlTypes<'ast> {
+    fn apply<N: Visitable>(
+        &mut self,
+        node_path: &NodePath<'ast>,
+        target_node: &mut N,
+    ) -> Result<bool, EqlMapperError> {
+        if self.would_edit(node_path, target_node) {
+            if let Some((_expr, function)) = node_path.last_2_as::<Expr, Function>() {
+                if matches!(
+                    self.node_types.get(&function.as_node_key()),
+                    Some(Type::Value(Value::Eql(_)))
+                ) {
+                    let function_name = CompoundIdent::from(&function.name.0);
+
+                    if let Some(SqlFunction {
+                        rewrite_rule: RewriteRule::AsEqlFunction,
+                        ..
+                    }) = get_sql_function_def(&function_name, &function.args)
+                    {
+                        let function = target_node.downcast_mut::<Function>().unwrap();
+                        let mut existing_name = mem::take(&mut function.name.0);
+                        existing_name.insert(0, Ident::new("eql_v1"));
+                        function.name = ObjectName(existing_name);
+                    }
+                }
+            }
+        }
+
+        Ok(false)
+    }
+
+    fn would_edit<N: Visitable>(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool {
+        if let Some((_expr, function)) = node_path.last_2_as::<Expr, Function>() {
+            if matches!(
+                self.node_types.get(&function.as_node_key()),
+                Some(Type::Value(Value::Eql(_)))
+            ) {
+                let function_name = CompoundIdent::from(&function.name.0);
+
+                if let Some(SqlFunction {
+                    rewrite_rule: RewriteRule::AsEqlFunction,
+                    ..
+                }) = get_sql_function_def(&function_name, &function.args)
+                {
+                    return true;
+                }
+            }
+        }
+
+        false
+    }
+}
diff --git a/packages/eql-mapper/src/type_checked_statement.rs b/packages/eql-mapper/src/type_checked_statement.rs
@@ -4,10 +4,7 @@ use sqltk::parser::ast::{self, Statement};
 use sqltk::{AsNodeKey, NodeKey, Transformable};
 
 use crate::{
-    DryRunnable, EqlMapperError, EqlValue, FailOnPlaceholderChange, GroupByEqlCol, Param,
-    PreserveEffectiveAliases, Projection, ReplacePlaintextEqlLiterals, TransformationRule, Type,
-    UseEquivalentSqlFuncForEqlTypes, Value, WrapEqlColsInOrderByWithOreFn,
-    WrapGroupedEqlColInAggregateFn,
+    DryRunnable, EqlMapperError, EqlValue, FailOnPlaceholderChange, GroupByEqlCol, Param, PreserveEffectiveAliases, Projection, ReplacePlaintextEqlLiterals, RewriteStandardSqlFnsOnEqlTypes, TransformationRule, Type, UseEquivalentSqlFuncForEqlTypes, Value, WrapEqlColsInOrderByWithOreFn, WrapGroupedEqlColInAggregateFn
 };
 
 /// A `TypeCheckedStatement` is returned from a successful call to [`crate::type_check`].
@@ -140,6 +137,7 @@ impl<'ast> TypeCheckedStatement<'ast> {
         encrypted_literals: HashMap<NodeKey<'ast>, sqltk::parser::ast::Value>,
     ) -> DryRunnable<impl TransformationRule<'_>> {
         DryRunnable::new((
+            RewriteStandardSqlFnsOnEqlTypes::new(Arc::clone(&self.node_types)),
             WrapGroupedEqlColInAggregateFn::new(Arc::clone(&self.node_types)),
             GroupByEqlCol::new(Arc::clone(&self.node_types)),
             WrapEqlColsInOrderByWithOreFn::new(Arc::clone(&self.node_types)),

Original file line number	Diff line number	Diff line change
`@@ -93,6 +93,7 @@ pub(crate) fn find_nodekey_for_value_node(`
`93`	`93`
`94`	`94`	`visitor.found`
`95`	`95`	`}`
	`96`	`+`
`96`	`97`	`#[macro_export]`
`97`	`98`	`macro_rules! col {`
`98`	`99`	`((NATIVE)) => {`