diff --git a/extension/partiql-extension-ion-functions/src/lib.rs b/extension/partiql-extension-ion-functions/src/lib.rs index 914942b4..034e190d 100644 --- a/extension/partiql-extension-ion-functions/src/lib.rs +++ b/extension/partiql-extension-ion-functions/src/lib.rs @@ -11,6 +11,7 @@ use partiql_logical as logical; use partiql_value::Value; use std::borrow::Cow; +use partiql_catalog::context::SessionContext; use std::error::Error; use std::fmt::Debug; use std::fs::File; @@ -98,7 +99,11 @@ impl BaseTableFunctionInfo for ReadIonFunction { pub(crate) struct EvalFnReadIon {} impl BaseTableExpr for EvalFnReadIon { - fn evaluate(&self, args: &[Cow]) -> BaseTableExprResult { + fn evaluate<'c>( + &self, + args: &[Cow], + _ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c> { if let Some(arg1) = args.first() { match arg1.as_ref() { Value::String(path) => parse_ion_file(path), @@ -155,9 +160,10 @@ fn parse_ion_buff<'a, I: 'a + ToIonDataSource>(input: I) -> BaseTableExprResult< mod tests { use super::*; + use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, Extension, PartiqlCatalog}; use partiql_eval::env::basic::MapBindings; - use partiql_eval::eval::{BasicContext, SystemContext}; + use partiql_eval::eval::BasicContext; use partiql_eval::plan::EvaluationMode; use partiql_parser::{Parsed, ParserResult}; use partiql_value::{bag, tuple, DateTime, Value}; diff --git a/partiql-catalog/src/context.rs b/partiql-catalog/src/context.rs new file mode 100644 index 00000000..2094f51d --- /dev/null +++ b/partiql-catalog/src/context.rs @@ -0,0 +1,27 @@ +use partiql_value::{BindingsName, DateTime, Tuple, Value}; +use std::any::Any; +use std::fmt::Debug; + +pub trait Bindings: Debug { + fn get(&self, name: &BindingsName) -> Option<&T>; +} + +impl Bindings for Tuple { + fn get(&self, name: &BindingsName) -> Option<&Value> { + self.get(name) + } +} + +#[derive(Debug)] +pub struct SystemContext { + pub now: DateTime, +} + +/// Represents a session context that is used during evaluation of a plan. +pub trait SessionContext<'a>: Debug { + fn bindings(&self) -> &dyn Bindings; + + fn system_context(&self) -> &SystemContext; + + fn user_context(&self, name: &str) -> Option<&(dyn Any)>; +} diff --git a/partiql-catalog/src/lib.rs b/partiql-catalog/src/lib.rs index e0ffa04b..62f92a21 100644 --- a/partiql-catalog/src/lib.rs +++ b/partiql-catalog/src/lib.rs @@ -4,6 +4,7 @@ use partiql_types::PartiqlType; use partiql_value::Value; use std::borrow::Cow; +use crate::context::SessionContext; use std::collections::HashMap; use std::error::Error; use std::fmt::Debug; @@ -13,6 +14,8 @@ use unicase::UniCase; pub mod call_defs; +pub mod context; + pub trait Extension: Debug { fn name(&self) -> String; fn load(&self, catalog: &mut dyn Catalog) -> Result<(), Box>; @@ -49,7 +52,11 @@ pub type BaseTableExprResult<'a> = Result, BaseTableExprResultError>; pub trait BaseTableExpr: Debug { - fn evaluate(&self, args: &[Cow]) -> BaseTableExprResult; + fn evaluate<'c>( + &self, + args: &[Cow], + ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c>; } pub trait BaseTableFunctionInfo: Debug { diff --git a/partiql-conformance-tests/tests/mod.rs b/partiql-conformance-tests/tests/mod.rs index 548a43df..46510126 100644 --- a/partiql-conformance-tests/tests/mod.rs +++ b/partiql-conformance-tests/tests/mod.rs @@ -3,13 +3,12 @@ use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval as eval; use partiql_eval::error::{EvalErr, PlanErr}; -use partiql_eval::eval::{ - BasicContext, EvalContext, EvalPlan, EvalResult, Evaluated, SystemContext, -}; +use partiql_eval::eval::{BasicContext, EvalContext, EvalPlan, EvalResult, Evaluated}; use partiql_logical as logical; use partiql_parser::{Parsed, ParserError, ParserResult}; use partiql_value::DateTime; +use partiql_catalog::context::SystemContext; use thiserror::Error; mod test_value; @@ -60,7 +59,7 @@ pub(crate) fn compile( #[track_caller] #[inline] -pub(crate) fn evaluate<'a, 'c>(mut plan: EvalPlan, ctx: &'c dyn EvalContext<'c>) -> EvalResult { +pub(crate) fn evaluate<'c>(mut plan: EvalPlan, ctx: &'c dyn EvalContext<'c>) -> EvalResult { plan.execute_mut(ctx) } diff --git a/partiql-eval/benches/bench_eval.rs b/partiql-eval/benches/bench_eval.rs index 89d514ce..f1d7ecf5 100644 --- a/partiql-eval/benches/bench_eval.rs +++ b/partiql-eval/benches/bench_eval.rs @@ -2,10 +2,11 @@ use std::borrow::Cow; use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan; use partiql_eval::plan::EvaluationMode; use partiql_logical as logical; @@ -13,7 +14,7 @@ use partiql_logical::BindingsOp::{Project, ProjectAll}; use partiql_logical::{ BinaryOp, BindingsOp, JoinKind, LogicalPlan, PathComponent, ValueExpr, VarRefType, }; -use partiql_value::{bag, list, tuple, BindingsName, Value}; +use partiql_value::{bag, list, tuple, BindingsName, DateTime, Value}; fn data() -> MapBindings { let hr = tuple![( @@ -134,7 +135,11 @@ fn eval_plan(logical: &LogicalPlan) -> EvalPlan { } fn evaluate(mut plan: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = plan.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = plan.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql-eval/src/env.rs b/partiql-eval/src/env.rs index 102529ca..7d1f2cd6 100644 --- a/partiql-eval/src/env.rs +++ b/partiql-eval/src/env.rs @@ -1,17 +1,8 @@ +use partiql_catalog::context::Bindings; use partiql_value::{BindingsName, Tuple, Value}; use std::fmt::Debug; use unicase::UniCase; -pub trait Bindings: Debug { - fn get(&self, name: &BindingsName) -> Option<&T>; -} - -impl Bindings for Tuple { - fn get(&self, name: &BindingsName) -> Option<&Value> { - self.get(name) - } -} - pub mod basic { use super::*; use std::collections::HashMap; diff --git a/partiql-eval/src/eval/evaluable.rs b/partiql-eval/src/eval/evaluable.rs index 9ec06511..95c5db2b 100644 --- a/partiql-eval/src/eval/evaluable.rs +++ b/partiql-eval/src/eval/evaluable.rs @@ -40,7 +40,7 @@ pub enum EvalType { /// `Evaluable` represents each evaluation operator in the evaluation plan as an evaluable entity. pub trait Evaluable: Debug { - fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value; + fn evaluate<'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value; fn update_input(&mut self, input: Value, branch_num: u8, ctx: &dyn EvalContext); fn get_vars(&self) -> Option<&[String]> { None diff --git a/partiql-eval/src/eval/expr/base_table.rs b/partiql-eval/src/eval/expr/base_table.rs index 8268d2ca..7a9e955b 100644 --- a/partiql-eval/src/eval/expr/base_table.rs +++ b/partiql-eval/src/eval/expr/base_table.rs @@ -31,7 +31,7 @@ impl EvalExpr for EvalFnBaseTableExpr { .iter() .map(|arg| arg.evaluate(bindings, ctx)) .collect_vec(); - let results = self.expr.evaluate(&args); + let results = self.expr.evaluate(&args, ctx.as_session()); let result = match results { Ok(it) => { let bag: Result = it.collect(); diff --git a/partiql-eval/src/eval/expr/path.rs b/partiql-eval/src/eval/expr/path.rs index 08b7de4d..c0bf8801 100644 --- a/partiql-eval/src/eval/expr/path.rs +++ b/partiql-eval/src/eval/expr/path.rs @@ -1,5 +1,3 @@ -use crate::env::Bindings; - pub use core::borrow::Borrow; use crate::eval::expr::{BindError, BindEvalExpr, EvalExpr}; @@ -8,6 +6,7 @@ use crate::eval::EvalContext; use partiql_value::Value::Missing; use partiql_value::{BindingsName, Tuple, Value}; +use partiql_catalog::context::Bindings; use std::borrow::Cow; use std::fmt::{Debug, Formatter}; diff --git a/partiql-eval/src/eval/mod.rs b/partiql-eval/src/eval/mod.rs index 9f608bff..6463fa1a 100644 --- a/partiql-eval/src/eval/mod.rs +++ b/partiql-eval/src/eval/mod.rs @@ -11,14 +11,14 @@ use petgraph::dot::Dot; use petgraph::prelude::StableGraph; use petgraph::{Directed, Outgoing}; -use partiql_value::{DateTime, Value}; +use partiql_value::Value; use crate::env::basic::MapBindings; -use crate::env::Bindings; use petgraph::graph::NodeIndex; use crate::error::{EvalErr, EvaluationError}; +use partiql_catalog::context::{Bindings, SessionContext, SystemContext}; use petgraph::visit::EdgeRef; use crate::eval::evaluable::{EvalType, Evaluable}; @@ -143,17 +143,9 @@ pub struct Evaluated { pub result: Value, } -#[derive(Debug)] -pub struct SystemContext { - pub now: DateTime, -} - /// Represents an evaluation context that is used during evaluation of a plan. -pub trait EvalContext<'a>: Debug { - fn bindings(&self) -> &dyn Bindings; - - fn system_context(&self) -> &SystemContext; - fn user_context(&self, name: &str) -> Option<&'a (dyn Any + 'a)>; +pub trait EvalContext<'a>: SessionContext<'a> + Debug { + fn as_session(&'a self) -> &'a dyn SessionContext<'a>; fn add_error(&self, error: EvaluationError); fn has_errors(&self) -> bool; @@ -180,8 +172,7 @@ impl<'a> BasicContext<'a> { } } } - -impl<'a> EvalContext<'a> for BasicContext<'a> { +impl<'a> SessionContext<'a> for BasicContext<'a> { fn bindings(&self) -> &dyn Bindings { &self.bindings } @@ -190,9 +181,14 @@ impl<'a> EvalContext<'a> for BasicContext<'a> { &self.sys } - fn user_context(&self, name: &str) -> Option<&'a (dyn Any + 'a)> { + fn user_context(&self, name: &str) -> Option<&(dyn Any)> { self.user.get(name).copied() } +} +impl<'a> EvalContext<'a> for BasicContext<'a> { + fn as_session(&'a self) -> &'a dyn SessionContext<'a> { + self + } fn add_error(&self, error: EvaluationError) { self.errors.borrow_mut().push(error) @@ -219,7 +215,7 @@ impl<'a, 'c> NestedContext<'a, 'c> { } } -impl<'a, 'c> EvalContext<'a> for NestedContext<'a, 'c> { +impl<'a, 'c> SessionContext<'a> for NestedContext<'a, 'c> { fn bindings(&self) -> &dyn Bindings { &self.bindings } @@ -227,8 +223,18 @@ impl<'a, 'c> EvalContext<'a> for NestedContext<'a, 'c> { delegate! { to self.parent { fn system_context(&self) -> &SystemContext; - fn user_context(&self, name: &str) -> Option<&'a (dyn Any + 'a)>; + fn user_context(&self, name: &str) -> Option<& (dyn Any )>; + } + } +} +impl<'a, 'c> EvalContext<'a> for NestedContext<'a, 'c> { + fn as_session(&'a self) -> &'a dyn SessionContext<'a> { + self + } + + delegate! { + to self.parent { fn add_error(&self, error: EvaluationError); fn has_errors(&self) -> bool; fn errors(&self) -> Vec; diff --git a/partiql-eval/src/lib.rs b/partiql-eval/src/lib.rs index c973340a..824d4079 100644 --- a/partiql-eval/src/lib.rs +++ b/partiql-eval/src/lib.rs @@ -9,13 +9,14 @@ mod tests { use crate::env::basic::MapBindings; use crate::plan; + use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use rust_decimal_macros::dec; use partiql_logical as logical; use partiql_logical::BindingsOp::{Distinct, Project, ProjectAll, ProjectValue}; - use crate::eval::{BasicContext, SystemContext}; + use crate::eval::BasicContext; use crate::plan::EvaluationMode; use partiql_logical::{ BagExpr, BetweenExpr, BinaryOp, BindingsOp, CoalesceExpr, ExprQuery, IsTypeExpr, JoinKind, @@ -2215,7 +2216,8 @@ mod tests { mod clause_from { use crate::eval::evaluable::{EvalScan, Evaluable}; use crate::eval::expr::{EvalGlobalVarRef, EvalPath, EvalPathComponent}; - use crate::eval::{BasicContext, SystemContext}; + use crate::eval::BasicContext; + use partiql_catalog::context::SystemContext; use partiql_value::{bag, list, BindingsName, DateTime}; use super::*; @@ -2354,7 +2356,7 @@ mod tests { use crate::eval::evaluable::{EvalUnpivot, Evaluable}; use crate::eval::expr::EvalGlobalVarRef; - use crate::eval::{BasicContext, SystemContext}; + use crate::eval::BasicContext; use super::*; diff --git a/partiql-logical-planner/src/lib.rs b/partiql-logical-planner/src/lib.rs index 872ff96f..fc13857a 100644 --- a/partiql-logical-planner/src/lib.rs +++ b/partiql-logical-planner/src/lib.rs @@ -38,10 +38,11 @@ impl<'c> LogicalPlanner<'c> { mod tests { use assert_matches::assert_matches; use partiql_ast_passes::error::AstTransformationError; + use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use partiql_eval::env::basic::MapBindings; - use partiql_eval::eval::{BasicContext, SystemContext}; + use partiql_eval::eval::BasicContext; use partiql_eval::plan; use partiql_eval::plan::EvaluationMode; diff --git a/partiql/benches/bench_agg.rs b/partiql/benches/bench_agg.rs index b4bfd6cb..d28b4c62 100644 --- a/partiql/benches/bench_agg.rs +++ b/partiql/benches/bench_agg.rs @@ -2,16 +2,17 @@ use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::Itertools; +use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan::{EvaluationMode, EvaluatorPlanner}; use partiql_logical::{BindingsOp, LogicalPlan}; use partiql_logical_planner::LogicalPlanner; use partiql_parser::{Parser, ParserResult}; -use partiql_value::{tuple, Bag, Value}; +use partiql_value::{tuple, Bag, DateTime, Value}; fn numbers() -> impl Iterator { (0..1000i64).map(Value::from) @@ -49,7 +50,11 @@ fn plan(catalog: &dyn Catalog, logical: &LogicalPlan) -> EvalPlan { } #[inline] pub(crate) fn evaluate(mut eval: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = eval.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = eval.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql/benches/bench_eval_multi_like.rs b/partiql/benches/bench_eval_multi_like.rs index 6348808f..8de96226 100644 --- a/partiql/benches/bench_eval_multi_like.rs +++ b/partiql/benches/bench_eval_multi_like.rs @@ -2,17 +2,18 @@ use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::Itertools; +use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use rand::{Rng, SeedableRng}; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan::{EvaluationMode, EvaluatorPlanner}; use partiql_logical::{BindingsOp, LogicalPlan}; use partiql_logical_planner::LogicalPlanner; use partiql_parser::{Parser, ParserResult}; -use partiql_value::{tuple, Bag, Value}; +use partiql_value::{tuple, Bag, DateTime, Value}; // Benchmarks: // - parsing, @@ -347,7 +348,11 @@ fn plan(catalog: &dyn Catalog, logical: &LogicalPlan) -> EvalPlan { } #[inline] pub(crate) fn evaluate(mut eval: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = eval.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = eval.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql/src/lib.rs b/partiql/src/lib.rs index daf64fa5..602fd95c 100644 --- a/partiql/src/lib.rs +++ b/partiql/src/lib.rs @@ -1,11 +1,12 @@ #[cfg(test)] mod tests { use partiql_ast_passes::error::AstTransformationError; + use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval as eval; use partiql_eval::env::basic::MapBindings; use partiql_eval::error::{EvalErr, PlanErr}; - use partiql_eval::eval::{BasicContext, EvalPlan, EvalResult, Evaluated, SystemContext}; + use partiql_eval::eval::{BasicContext, EvalPlan, EvalResult, Evaluated}; use partiql_eval::plan::EvaluationMode; use partiql_logical as logical; use partiql_parser::{Parsed, ParserError, ParserResult}; @@ -101,14 +102,14 @@ mod tests { fn order_by_count() { let query = "select foo, count(1) as n from << - { 'foo': 'foo' }, - { 'foo': 'bar' }, - { 'foo': 'qux' }, - { 'foo': 'bar' }, - { 'foo': 'baz' }, - { 'foo': 'bar' }, - { 'foo': 'baz' } - >> group by foo order by n desc"; + { 'foo': 'foo' }, + { 'foo': 'bar' }, + { 'foo': 'qux' }, + { 'foo': 'bar' }, + { 'foo': 'baz' }, + { 'foo': 'bar' }, + { 'foo': 'baz' } + >> group by foo order by n desc"; let res = eval(query, EvaluationMode::Permissive); assert!(res.is_ok()); diff --git a/partiql/tests/user_context.rs b/partiql/tests/user_context.rs new file mode 100644 index 00000000..e9ee53ce --- /dev/null +++ b/partiql/tests/user_context.rs @@ -0,0 +1,215 @@ +use std::any::Any; +use std::borrow::Cow; +use std::cell::RefCell; + +use std::error::Error; + +use thiserror::Error; + +use partiql_catalog::call_defs::{CallDef, CallSpec, CallSpecArg}; +use partiql_catalog::context::{SessionContext, SystemContext}; +use partiql_catalog::{ + BaseTableExpr, BaseTableExprResult, BaseTableExprResultError, BaseTableFunctionInfo, Catalog, + Extension, PartiqlCatalog, TableFunction, +}; +use partiql_eval::env::basic::MapBindings; +use partiql_eval::eval::BasicContext; +use partiql_eval::plan::EvaluationMode; +use partiql_parser::{Parsed, ParserResult}; +use partiql_value::{bag, tuple, DateTime, Value}; + +use partiql_logical as logical; + +#[derive(Debug)] +pub struct UserCtxTestExtension {} + +impl partiql_catalog::Extension for UserCtxTestExtension { + fn name(&self) -> String { + "ion".into() + } + + fn load(&self, catalog: &mut dyn Catalog) -> Result<(), Box> { + match catalog + .add_table_function(TableFunction::new(Box::new(TestUserContextFunction::new()))) + { + Ok(_) => Ok(()), + Err(e) => Err(Box::new(e) as Box), + } + } +} + +#[derive(Debug)] +pub(crate) struct TestUserContextFunction { + call_def: CallDef, +} + +impl TestUserContextFunction { + pub fn new() -> Self { + TestUserContextFunction { + call_def: CallDef { + names: vec!["test_user_context"], + overloads: vec![CallSpec { + input: vec![CallSpecArg::Positional], + output: Box::new(|args| { + logical::ValueExpr::Call(logical::CallExpr { + name: logical::CallName::ByName("test_user_context".to_string()), + arguments: args, + }) + }), + }], + }, + } + } +} + +impl BaseTableFunctionInfo for TestUserContextFunction { + fn call_def(&self) -> &CallDef { + &self.call_def + } + + fn plan_eval(&self) -> Box { + Box::new(EvalTestCtxTable {}) + } +} + +#[derive(Error, Debug)] +#[non_exhaustive] +pub enum UserCtxError { + #[error("unknown error")] + Unknown, +} + +#[derive(Debug)] +pub(crate) struct EvalTestCtxTable {} + +impl BaseTableExpr for EvalTestCtxTable { + fn evaluate<'c>( + &self, + args: &[Cow], + ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c> { + if let Some(arg1) = args.first() { + match arg1.as_ref() { + Value::String(name) => generated_data(name.to_string(), ctx), + _ => { + let error = UserCtxError::Unknown; + Err(Box::new(error) as BaseTableExprResultError) + } + } + } else { + let error = UserCtxError::Unknown; + Err(Box::new(error) as BaseTableExprResultError) + } + } +} + +struct TestDataGen<'a> { + ctx: &'a dyn SessionContext<'a>, + name: String, +} + +impl<'a> Iterator for TestDataGen<'a> { + type Item = Result; + + fn next(&mut self) -> Option { + if let Some(cv) = self.ctx.user_context(&self.name) { + if let Some(counter) = cv.downcast_ref::() { + let mut n = counter.data.borrow_mut(); + + if *n > 0 { + *n -= 1; + + let idx: u8 = (5 - *n) as u8; + let id = format!("id_{idx}"); + let m = idx % 2; + + return Some(Ok(tuple![("foo", m), ("bar", id)].into())); + } + } + } + None + } +} + +fn generated_data<'a>(name: String, ctx: &'a dyn SessionContext<'a>) -> BaseTableExprResult<'a> { + Ok(Box::new(TestDataGen { ctx, name })) +} + +#[derive(Debug)] +pub struct Counter { + data: RefCell, +} + +#[track_caller] +#[inline] +pub(crate) fn parse(statement: &str) -> ParserResult { + partiql_parser::Parser::default().parse(statement) +} + +#[track_caller] +#[inline] +pub(crate) fn lower( + catalog: &dyn Catalog, + parsed: &Parsed, +) -> partiql_logical::LogicalPlan { + let planner = partiql_logical_planner::LogicalPlanner::new(catalog); + planner.lower(parsed).expect("lower") +} + +#[track_caller] +#[inline] +pub(crate) fn evaluate( + catalog: &dyn Catalog, + logical: partiql_logical::LogicalPlan, + bindings: MapBindings, + ctx_vals: &[(String, &(dyn Any))], +) -> Value { + let mut planner = + partiql_eval::plan::EvaluatorPlanner::new(EvaluationMode::Permissive, catalog); + + let mut plan = planner.compile(&logical).expect("Expect no plan error"); + + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let mut ctx = BasicContext::new(bindings, sys); + for (k, v) in ctx_vals { + ctx.user.insert(k.to_string(), *v); + } + if let Ok(out) = plan.execute_mut(&ctx) { + out.result + } else { + Value::Missing + } +} +#[test] +fn test_context() { + let expected: Value = bag![ + tuple![("foo", 1), ("bar", "id_1")], + tuple![("foo", 0), ("bar", "id_2")], + tuple![("foo", 1), ("bar", "id_3")], + tuple![("foo", 0), ("bar", "id_4")], + tuple![("foo", 1), ("bar", "id_5")], + ] + .into(); + + let query = "SELECT foo, bar from test_user_context('counter') as data"; + + let mut catalog = PartiqlCatalog::default(); + let ext = UserCtxTestExtension {}; + ext.load(&mut catalog).expect("extension load to succeed"); + + let parsed = parse(query); + let lowered = lower(&catalog, &parsed.expect("parse")); + let bindings = Default::default(); + + let counter = Counter { + data: RefCell::new(5), + }; + let ctx: [(String, &dyn Any); 1] = [("counter".to_string(), &counter)]; + let out = evaluate(&catalog, lowered, bindings, &ctx); + + assert!(out.is_bag()); + assert_eq!(&out, &expected); + assert_eq!(*counter.data.borrow(), 0); +}