uqa_sql/semantics/
runtime_scalars.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Runtime scalar argument and row-context rules.
8
9use crate::{expr::RowLookup, SQLError, ScalarExpr};
10use uqa_core::Value;
11
12pub fn deep_learn_arguments(
13    args: &[ScalarExpr],
14    evaluate: &mut dyn FnMut(&ScalarExpr) -> Result<Value, SQLError>,
15) -> Result<(String, String), SQLError> {
16    if args.len() != 2 {
17        return Err(SQLError::BadArity {
18            name: "deep_learn".into(),
19            expected: "2".into(),
20            actual: args.len(),
21        });
22    }
23    let model_name = match evaluate(&args[0])? {
24        Value::Str(s) => s,
25        other => {
26            return Err(SQLError::TypeMismatch(format!(
27                "deep_learn.model must be a string, got {other:?}"
28            )));
29        }
30    };
31    let training_source = match evaluate(&args[1])? {
32        Value::Str(s) => s,
33        other => {
34            return Err(SQLError::TypeMismatch(format!(
35                "deep_learn.training_set must be a table name or JSON string, got {other:?}"
36            )));
37        }
38    };
39    Ok((model_name, training_source))
40}
41
42pub enum DeepLearnSource<'a> {
43    Json(&'a str),
44    Table(&'a str),
45}
46pub fn deep_learn_source(source: &str) -> DeepLearnSource<'_> {
47    let trimmed = source.trim();
48    if trimmed.starts_with('{') {
49        DeepLearnSource::Json(trimmed)
50    } else {
51        DeepLearnSource::Table(source)
52    }
53}
54pub fn no_scalar_arguments(name: &str, arguments: &[Value]) -> Result<(), SQLError> {
55    if arguments.is_empty() {
56        Ok(())
57    } else {
58        Err(SQLError::BadArity {
59            name: name.into(),
60            expected: "0".into(),
61            actual: arguments.len(),
62        })
63    }
64}
65pub fn notification_arguments(arguments: &[Value]) -> Result<(&str, &str), SQLError> {
66    match arguments {
67        [channel, payload] => Ok((
68            notification_text_argument(channel, "channel")?,
69            notification_text_argument(payload, "payload")?,
70        )),
71        _ => Err(SQLError::BadArity {
72            name: "pg_notify".into(),
73            expected: "2".into(),
74            actual: arguments.len(),
75        }),
76    }
77}
78fn notification_text_argument<'a>(value: &'a Value, label: &str) -> Result<&'a str, SQLError> {
79    match value {
80        Value::Null => Ok(""),
81        Value::Str(value) => Ok(value),
82        Value::FixedChar(value) => Ok(value.trim_end_matches(' ')),
83        other => Err(SQLError::TypeMismatch(format!(
84            "pg_notify {label} must be text, got {other:?}"
85        ))),
86    }
87}
88
89pub fn merge_action_value(args: &[ScalarExpr], row: &dyn RowLookup) -> Result<Value, SQLError> {
90    Ok(merge_action_value_with_control(
91        args.len(),
92        row,
93        &uqa_core::memory::ProductionControl::uncontrolled(),
94    )?
95    .into_uncontrolled()
96    .expect("ordinary MERGE action output has no reservation"))
97}
98
99pub(crate) fn merge_action_value_with_control(
100    argument_count: usize,
101    row: &dyn RowLookup,
102    control: &uqa_core::memory::ProductionControl<'_>,
103) -> Result<uqa_core::memory::Produced<Value>, SQLError> {
104    control.check()?;
105    if argument_count != 0 {
106        return Err(SQLError::BadArity {
107            name: "merge_action".into(),
108            expected: "0".into(),
109            actual: argument_count,
110        });
111    }
112    let action = row
113        .internal_column(super::merge_action_attribute())
114        .ok_or_else(|| {
115            SQLError::Unsupported("merge_action() is only valid in MERGE RETURNING".into())
116        })?;
117    Ok(control.copy_value(action)?)
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use uqa_core::memory::{MemoryBudget, ProductionControl};
124    use uqa_core::CancellationToken;
125
126    struct ActionRow(Value);
127    impl RowLookup for ActionRow {
128        fn column(&self, _: &str) -> Option<&Value> {
129            None
130        }
131        fn qualified_column(&self, _: &str, _: &str) -> Option<&Value> {
132            None
133        }
134        fn internal_column(&self, column: crate::ast::InternalColumnRef) -> Option<&Value> {
135            (column == crate::semantics::merge_action_attribute()).then_some(&self.0)
136        }
137    }
138
139    #[test]
140    fn merge_action_output_retains_its_allowance_and_honors_cancellation() {
141        let row = ActionRow(Value::Str("UPDATE".into()));
142        let budget = MemoryBudget::new(256);
143        let original = CancellationToken::new();
144        let invoking = CancellationToken::new();
145        let control = ProductionControl::new(&budget, &original, &invoking);
146        let output = merge_action_value_with_control(0, &row, &control).unwrap();
147        assert_eq!(*output, row.0);
148        assert!(budget.used() > 0);
149        assert_eq!(budget.used(), output.reserved_bytes());
150        drop(output);
151        assert_eq!(budget.used(), 0);
152        let tiny = MemoryBudget::new(0);
153        let limited = ProductionControl::new(&tiny, &original, &invoking);
154        assert_eq!(
155            merge_action_value_with_control(0, &row, &limited)
156                .unwrap_err()
157                .sqlstate(),
158            Some("53200")
159        );
160        assert_eq!(tiny.used(), 0);
161        invoking.cancel();
162        assert_eq!(
163            merge_action_value_with_control(0, &row, &control)
164                .unwrap_err()
165                .sqlstate(),
166            Some("57014")
167        );
168        assert_eq!(budget.used(), 0);
169    }
170}
uqa_sql/semantics/runtime_scalars.rs

uqa_sql/semantics/
runtime_scalars.rs
