Skip to content

Commit f8ab910

Browse files
tobyhedefreshtonic
tobyhede
authored and
freshtonic
committed
Add test for jsonb_path_query inference
1 parent ae42841 commit f8ab910

File tree

11 files changed

+67
-33
lines changed

11 files changed

+67
-33
lines changed

‎Cargo.lock

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/cipherstash-proxy-integration/src/select/jsonb_path_query.rs

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
#[cfg(test)]
22
mod tests {
3-
use crate::common::{
4-
clear, connect_with_tls, insert, query_by, random_id, trace,
5-
};
3+
use crate::common::{clear, connect_with_tls, insert, query_by, random_id, trace};
64
use serde_json::Value;
75
use tracing::info;
86

@@ -45,8 +43,7 @@ mod tests {
4543
// info!("Results: {:?}", results);
4644

4745
let selector = "$.nested.string";
48-
let sql =
49-
format!("SELECT eql_v2.jsonb_path_query(encrypted_jsonb, $1)::jsonb FROM encrypted");
46+
let sql = format!("SELECT jsonb_path_query(encrypted_jsonb, $1)::jsonb FROM encrypted");
5047

5148
let actual = query_by::<Value>(&sql, &selector).await;
5249

‎packages/eql-mapper-macros/src/parse_type_spec.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -408,8 +408,8 @@ impl Parse for FunctionDecl {
408408
Ok(Self(quote! {
409409
crate::inference::unifier::FunctionDecl {
410410
name: sqltk::parser::ast::ObjectName(vec![
411-
sqltk::parser::ast::Ident::new(#schema),
412-
sqltk::parser::ast::Ident::new(#function_name),
411+
sqltk::parser::ast::ObjectNamePart::Identifier(sqltk::parser::ast::Ident::new(#schema)),
412+
sqltk::parser::ast::ObjectNamePart::Identifier(sqltk::parser::ast::Ident::new(#function_name)),
413413
]),
414414
inner: crate::inference::unifier::FunctionSignatureDecl::new(
415415
vec![#(#generic_args),*],

‎packages/eql-mapper/src/inference/infer_type_impls/expr.rs

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
use crate::{
2-
get_sql_binop_rule, inference::{unifier::Type, InferType, TypeError}, SqlIdent, TypeInferencer
2+
get_sql_binop_rule,
3+
inference::{unifier::Type, InferType, TypeError},
4+
SqlIdent, TypeInferencer,
35
};
46
use eql_mapper_macros::trace_infer;
57
use sqltk::parser::ast::{AccessExpr, Array, Expr, Ident, Subscript};
@@ -29,7 +31,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
2931

3032
Expr::QualifiedWildcard(object_name, _) => {
3133
self.unify_node_with_type(
32-
this_expr,
34+
return_val,
3335
self.resolve_qualified_wildcard(object_name)?,
3436
)?;
3537
}
@@ -98,7 +100,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
98100
}
99101

100102
Expr::BinaryOp { left, op, right } => {
101-
get_sql_binop_rule(op).apply_constraints(self, left, right, this_expr)?;
103+
get_sql_binop_rule(op).apply_constraints(self, left, right, return_val)?;
102104
}
103105

104106
//customer_name LIKE 'A%';
@@ -280,15 +282,15 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
280282
Some(operand) => {
281283
for cond_when in conditions {
282284
self.unify_nodes_with_type(
283-
this_expr,
285+
return_val,
284286
&**operand,
285287
self.unify_node_with_type(&cond_when.condition, self.fresh_tvar())?,
286288
)?;
287289
}
288290
}
289291
None => {
290292
for cond_when in conditions {
291-
self.unify_node_with_type(&cond_when.condition, Type::any_native())?;
293+
self.unify_node_with_type(&cond_when.condition, Type::native())?;
292294
}
293295
}
294296
}
@@ -301,7 +303,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
301303
self.unify_node_with_type(else_result, result_ty.clone())?;
302304
};
303305

304-
self.unify_node_with_type(this_expr, result_ty)?;
306+
self.unify_node_with_type(return_val, result_ty)?;
305307
}
306308

307309
Expr::Exists {
@@ -355,16 +357,16 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
355357
AccessExpr::Subscript(Subscript::Index { index }) => {
356358
access_ty = self.fresh_tvar();
357359
root_ty = Type::array(access_ty.clone());
358-
self.unify_node_with_type(index, Type::any_native())?;
360+
self.unify_node_with_type(index, Type::native())?;
359361
}
360362
AccessExpr::Subscript(Subscript::Slice {
361363
lower_bound,
362364
upper_bound,
363365
stride,
364366
}) => {
365-
self.unify_node_with_type(lower_bound, Type::any_native())?;
366-
self.unify_node_with_type(upper_bound, Type::any_native())?;
367-
self.unify_node_with_type(stride, Type::any_native())?;
367+
self.unify_node_with_type(lower_bound, Type::native())?;
368+
self.unify_node_with_type(upper_bound, Type::native())?;
369+
self.unify_node_with_type(stride, Type::native())?;
368370
access_ty = self.fresh_tvar();
369371
root_ty = Type::array(access_ty.clone());
370372
}
@@ -376,7 +378,7 @@ impl<'ast> InferType<'ast, Expr> for TypeInferencer<'ast> {
376378
}
377379
}
378380

379-
self.unify_node_with_type(this_expr, access_ty)?;
381+
self.unify_node_with_type(return_val, access_ty)?;
380382
self.unify_node_with_type(&**root, root_ty)?;
381383
}
382384

‎packages/eql-mapper/src/inference/sql_types/sql_decls.rs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
use std::{collections::HashMap, sync::LazyLock};
22

33
use eql_mapper_macros::{binary_operators, functions};
4-
use sqltk::parser::ast::{BinaryOperator, Ident, ObjectName};
4+
use sqltk::parser::ast::{BinaryOperator, Ident, ObjectName, ObjectNamePart};
55

66
use crate::unifier::{BinaryOpDecl, FunctionDecl};
77

88
use super::{SqlBinaryOp, SqlFunction};
99

10-
1110
/// SQL operators that can accept EQL types.
1211
static SQL_BINARY_OPERATORS: LazyLock<HashMap<BinaryOperator, BinaryOpDecl>> =
1312
LazyLock::new(|| {
@@ -82,7 +81,10 @@ static SQL_FUNCTION_TYPES: LazyLock<HashMap<ObjectName, FunctionDecl>> = LazyLoc
8281
pub(crate) fn get_sql_function(fn_name: &ObjectName) -> SqlFunction {
8382
// FIXME: this is a hack and we need proper schema resolution logic
8483
let fully_qualified_fn_name = if fn_name.0.len() == 1 {
85-
&ObjectName(vec![Ident::new("pg_catalog"), fn_name.0[0].clone()])
84+
&ObjectName(vec![
85+
ObjectNamePart::Identifier(Ident::new("pg_catalog")),
86+
fn_name.0[0].clone(),
87+
])
8688
} else {
8789
fn_name
8890
};

‎packages/eql-mapper/src/inference/sql_types/sql_function_types.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1-
use std::sync::Arc;
1+
use std::sync::{Arc, LazyLock};
22

3-
use sqltk::parser::ast::{Function, FunctionArg, FunctionArgExpr, FunctionArguments};
3+
use sqltk::parser::ast::{
4+
Function, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, ObjectNamePart,
5+
};
46

57
use crate::{
68
unifier::{FunctionDecl, Type, Unifier},
@@ -15,10 +17,13 @@ pub(crate) enum SqlFunction {
1517
Fallback,
1618
}
1719

20+
static PG_CATALOG: LazyLock<ObjectNamePart> =
21+
LazyLock::new(|| ObjectNamePart::Identifier(Ident::new("pg_catalog")));
22+
1823
impl SqlFunction {
1924
pub(crate) fn should_rewrite(&self) -> bool {
2025
match self {
21-
SqlFunction::Explicit(function_spec) => &function_spec.name.0[0].value == "pg_catalog",
26+
SqlFunction::Explicit(function_spec) => function_spec.name.0[0] == *PG_CATALOG,
2227
SqlFunction::Fallback => false,
2328
}
2429
}

‎packages/eql-mapper/src/inference/unifier/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,7 @@ mod test {
388388
fn constructor_with_var() {
389389
let mut unifier = Unifier::new(DepMut::new(TypeRegistry::new()));
390390

391-
shallow_init_types!{ &mut unifier, {
391+
shallow_init_types! { &mut unifier, {
392392
let lhs = Native;
393393
let rhs = T;
394394
}};

‎packages/eql-mapper/src/inference/unifier/type_env.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -199,8 +199,8 @@ mod test {
199199
use crate::{
200200
test_helpers,
201201
unifier::{
202-
Array, AssociatedType, Constructor, EqlTerm, EqlTrait, EqlTraits, EqlValue, InstantiateType,
203-
Type, Unifier, Value,
202+
Array, AssociatedType, Constructor, EqlTerm, EqlTrait, EqlTraits, EqlValue,
203+
InstantiateType, Type, Unifier, Value,
204204
},
205205
NativeValue, TableColumn, TypeError, TypeRegistry,
206206
};
@@ -267,8 +267,7 @@ mod test {
267267

268268
if let Type::Associated(associated) = &*instance.get_type(&tvar!(A))? {
269269
assert_eq!(
270-
associated
271-
.resolve_selector_target(&mut unifier)?.as_deref(),
270+
associated.resolve_selector_target(&mut unifier)?.as_deref(),
272271
Some(&Type::Constructor(Constructor::Value(Value::Eql(
273272
EqlTerm::JsonAccessor(EqlValue(
274273
TableColumn {

‎packages/eql-mapper/src/lib.rs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1775,4 +1775,28 @@ mod test {
17751775
Err(err) => panic!("type check failed: {err}"),
17761776
}
17771777
}
1778+
1779+
#[test]
1780+
fn jsonb_path_query_param_to_eql() {
1781+
// init_tracing();
1782+
let schema = resolver(schema! {
1783+
tables: {
1784+
patients: {
1785+
id,
1786+
notes (EQL: JsonLike),
1787+
}
1788+
}
1789+
});
1790+
1791+
let statement = parse("SELECT eql_v2.jsonb_path_query(notes, $1) as notes FROM patients");
1792+
1793+
let typed = type_check(schema, &statement)
1794+
.map_err(|err| err.to_string())
1795+
.unwrap();
1796+
1797+
assert_eq!(
1798+
typed.projection,
1799+
projection![(EQL(patients.notes: JsonLike) as notes)]
1800+
);
1801+
}
17781802
}

‎packages/eql-mapper/src/transformation_rules/cast_params_as_encrypted.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
use super::helpers::cast_as_encrypted;
22
use super::TransformationRule;
33
use crate::{Constructor, EqlMapperError, Type};
4-
use crate::{EqlMapperError, Type};
54
use sqltk::parser::ast::{Expr, Value, ValueWithSpan};
65
use sqltk::parser::tokenizer::Span;
76
use sqltk::{NodeKey, NodePath, Visitable};
@@ -57,7 +56,13 @@ impl<'ast> TransformationRule<'ast> for CastParamsAsEncrypted<'ast> {
5756
}
5857

5958
fn would_edit<N: Visitable>(&mut self, node_path: &NodePath<'ast>, _target_node: &N) -> bool {
60-
if let Some((node @ Expr::Value(Value::Placeholder(_)),)) = node_path.last_1_as() {
59+
if let Some((
60+
node @ Expr::Value(ValueWithSpan {
61+
value: Value::Placeholder(_),
62+
..
63+
}),
64+
)) = node_path.last_1_as()
65+
{
6166
if let Some(Type::Constructor(Constructor::Value(crate::Value::Eql(_)))) =
6267
self.node_types.get(&NodeKey::new(node))
6368
{

‎packages/eql-mapper/src/type_checked_statement.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ impl<'ast> TypeCheckedStatement<'ast> {
4848
statement: &'ast Statement,
4949
projection: Projection,
5050
params: Vec<(Param, Value)>,
51-
literals: Vec<(EqlValue, &'ast ast::Value)>,
51+
literals: Vec<(EqlTerm, &'ast ast::Value)>,
5252
node_types: Arc<HashMap<NodeKey<'ast>, Type>>,
5353
) -> Self {
5454
Self {

0 commit comments

Comments
 (0)