Skip to main content

substrait_explain/extensions/
args.rs

1//! Text-format data structures used by registered advanced extension handlers.
2//!
3//! These types describe the arguments accepted by custom relation types,
4//! enhancements, and optimization hints. Relation extensions can additionally
5//! describe output columns.
6//!
7//! The interface presented to extension handlers is structured rather than
8//! textual: handlers read and write values such as [`ExtensionArgs`], [`Expr`],
9//! and [`proto::Type`]. `substrait-explain` handles the surrounding
10//! parsing/textification. Some values need plan context before they reach a
11//! handler; for example, an expression argument like `add($0, $1)` is parsed
12//! using [`SimpleExtensions`](crate::extensions::SimpleExtensions) to resolve
13//! the text function name to the protobuf function anchor, and formatted by
14//! resolving that anchor back to a text name.
15//!
16//! The extension-facing interface for Substrait objects (e.g. [`proto::Type`])
17//! should map directly to Substrait protobuf concepts. Sometimes that means
18//! storing the protobuf type directly, as named output columns do with
19//! [`proto::Type`]; sometimes it means using a small wrapper, as
20//! expression-compatible arguments do with [`Expr`] around
21//! [`proto::Expression`].
22//!
23//! Untyped scalar literals (e.g. `2`, `2.435`, `'string'`) are kept as
24//! extension scalar values so text rendering can preserve scalar syntax even in
25//! verbose output, while handlers that accept expressions can still widen them
26//! into default Substrait literal expressions.
27
28use std::collections::HashSet;
29use std::fmt;
30use std::slice::Iter as SliceIter;
31use std::vec::IntoIter as VecIntoIter;
32
33use indexmap::IndexMap;
34use substrait::proto;
35use substrait::proto::expression::field_reference::ReferenceType;
36use substrait::proto::expression::literal::LiteralType;
37use substrait::proto::expression::{RexType, reference_segment};
38
39use super::{Explainable, ExtensionError};
40use crate::textify::expressions::Reference;
41
42/// Kind of relation addendum in the text format.
43///
44/// Addenda are `+`-prefixed lines attached to relations. They are syntax-level
45/// constructs, distinct from [`crate::extensions::registry::ExtensionType`],
46/// which describes registry namespaces.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub(crate) enum AddendumKind {
49    Enhancement,
50    Optimization,
51    ExtensionTable,
52}
53
54impl AddendumKind {
55    pub(crate) fn prefix(self) -> &'static str {
56        match self {
57            AddendumKind::Enhancement => "Enh",
58            AddendumKind::Optimization => "Opt",
59            AddendumKind::ExtensionTable => "Ext",
60        }
61    }
62}
63
64/// A Substrait expression carried as an extension argument or output column.
65///
66/// Boxed because `proto::Expression` is large (multiple `Vec` fields in
67/// variants like `ScalarFunction`).
68#[derive(Debug, Clone)]
69pub struct Expr(Box<proto::Expression>);
70
71impl Expr {
72    /// Create a direct field-reference expression (`$N`).
73    pub fn field(index: i32) -> Self {
74        Reference(index).into()
75    }
76
77    /// Borrow the underlying Substrait expression protobuf.
78    pub fn as_proto(&self) -> &proto::Expression {
79        self.0.as_ref()
80    }
81
82    /// Clone the underlying Substrait expression protobuf.
83    pub fn to_proto(&self) -> proto::Expression {
84        self.as_proto().clone()
85    }
86
87    /// If this expression is a direct field reference (`$N`), return it.
88    pub fn as_direct_reference(&self) -> Option<i32> {
89        let Some(RexType::Selection(field_ref)) = self.as_proto().rex_type.as_ref() else {
90            return None;
91        };
92        let Some(ReferenceType::DirectReference(segment)) = field_ref.reference_type.as_ref()
93        else {
94            return None;
95        };
96        let Some(reference_segment::ReferenceType::StructField(field)) =
97            segment.reference_type.as_ref()
98        else {
99            return None;
100        };
101        if field.child.is_some() {
102            return None;
103        }
104        Some(field.field)
105    }
106}
107
108impl From<proto::Expression> for Expr {
109    fn from(expr: proto::Expression) -> Self {
110        Expr(Box::new(expr))
111    }
112}
113
114impl From<proto::expression::Literal> for Expr {
115    fn from(literal: proto::expression::Literal) -> Self {
116        proto::Expression {
117            rex_type: Some(RexType::Literal(literal)),
118        }
119        .into()
120    }
121}
122
123impl From<Reference> for Expr {
124    fn from(reference: Reference) -> Self {
125        proto::Expression::from(reference).into()
126    }
127}
128
129impl From<Expr> for proto::Expression {
130    fn from(expr: Expr) -> Self {
131        *expr.0
132    }
133}
134
135impl From<i64> for Expr {
136    fn from(value: i64) -> Self {
137        proto::expression::Literal {
138            literal_type: Some(LiteralType::I64(value)),
139            nullable: false,
140            type_variation_reference: 0,
141        }
142        .into()
143    }
144}
145
146impl From<f64> for Expr {
147    fn from(value: f64) -> Self {
148        proto::expression::Literal {
149            literal_type: Some(LiteralType::Fp64(value)),
150            nullable: false,
151            type_variation_reference: 0,
152        }
153        .into()
154    }
155}
156
157impl From<bool> for Expr {
158    fn from(value: bool) -> Self {
159        proto::expression::Literal {
160            literal_type: Some(LiteralType::Boolean(value)),
161            nullable: false,
162            type_variation_reference: 0,
163        }
164        .into()
165    }
166}
167
168impl From<String> for Expr {
169    fn from(value: String) -> Self {
170        proto::expression::Literal {
171            literal_type: Some(LiteralType::String(value)),
172            nullable: false,
173            type_variation_reference: 0,
174        }
175        .into()
176    }
177}
178
179impl From<&str> for Expr {
180    fn from(value: &str) -> Self {
181        value.to_string().into()
182    }
183}
184
185/// Represents extension arguments plus optional output columns.
186///
187/// Named arguments are stored in an [`IndexMap`] whose iteration order
188/// determines display order. Extension [`super::Explainable::to_args()`]
189/// implementations should insert named arguments in the order they should
190/// appear in the text format.
191#[derive(Debug, Clone, Default)]
192pub struct ExtensionArgs {
193    /// Positional arguments.
194    pub positional: Vec<ExtensionValue>,
195    /// Named arguments, displayed in the order they were inserted
196    pub named: IndexMap<String, ExtensionValue>,
197    /// Output columns for custom relation types.
198    pub output_columns: Vec<ExtensionColumn>,
199}
200
201/// [`ArgsAccess`] provides a view of the arguments in the text form of an
202/// extension relation or advanced extension.
203///
204/// Positional arguments are available via [`Self::positional`], named arguments
205/// via [`Self::get_named`] or [`Self::expect_named`] which can do type
206/// conversion as well, and output columns via [`Self::output_columns`].
207///
208/// [`ArgsAccess`] tracks which arguments are accessed; when
209/// [`Explainable::from_args`] returns, any unaccessed arguments will be raised
210/// as [`ExtensionError::InvalidArgument`] errors.
211pub struct ArgsAccess<'a> {
212    args: &'a ExtensionArgs,
213    handled: HashSet<&'a str>,
214    positional_handled: bool,
215}
216
217impl<'a> ArgsAccess<'a> {
218    pub(crate) fn new(args: &'a ExtensionArgs) -> Self {
219        Self {
220            args,
221            handled: HashSet::new(),
222            positional_handled: false,
223        }
224    }
225
226    /// Returns the positional arguments in source order and marks them as handled.
227    pub fn positional(&mut self) -> &'a [ExtensionValue] {
228        self.positional_handled = true;
229        &self.args.positional
230    }
231
232    /// Returns the output columns for a custom relation.
233    ///
234    /// Output columns are not included in the unhandled-argument check.
235    pub fn output_columns(&self) -> &'a [ExtensionColumn] {
236        &self.args.output_columns
237    }
238
239    /// Returns a named argument without converting it, or `None` if it is absent.
240    ///
241    /// A present argument is marked as handled.
242    pub fn get_named_arg(&mut self, name: &str) -> Option<&'a ExtensionValue> {
243        match self.args.named.get_key_value(name) {
244            Some((k, value)) => {
245                self.handled.insert(k);
246                Some(value)
247            }
248            None => None,
249        }
250    }
251
252    /// Returns a named argument converted to `T`, or `None` if it is absent.
253    ///
254    /// A present argument is marked as handled.
255    pub fn get_named<T>(&mut self, name: &str) -> Result<Option<T>, ExtensionError>
256    where
257        T: TryFrom<&'a ExtensionValue>,
258        T::Error: Into<ExtensionError>,
259    {
260        self.get_named_arg(name)
261            .map(|value| {
262                T::try_from(value).map_err(|error| ExtensionError::NamedArgumentConversion {
263                    name: name.to_string(),
264                    source: Box::new(error.into()),
265                })
266            })
267            .transpose()
268    }
269
270    /// Fetch the named argument `name`, with an expected type of
271    /// [`TupleValue`], and convert each element to type `T`.
272    ///
273    /// Returns `Ok(None)` if the argument is absent. Returns an error if the
274    /// argument is not a tuple, or if any element of the tuple fails
275    /// conversion.
276    pub fn get_named_tuple<T>(&mut self, name: &str) -> Result<Option<Vec<T>>, ExtensionError>
277    where
278        T: TryFrom<&'a ExtensionValue>,
279        T::Error: Into<ExtensionError>,
280    {
281        self.get_named_arg(name)
282            .map(|value| {
283                let tuple = <&TupleValue>::try_from(value)?;
284                tuple
285                    .into_iter()
286                    .map(|element| T::try_from(element).map_err(Into::into))
287                    .collect::<Result<Vec<_>, ExtensionError>>()
288            })
289            .transpose()
290            .map_err(|source| ExtensionError::NamedArgumentConversion {
291                name: name.to_string(),
292                source: Box::new(source),
293            })
294    }
295
296    /// Returns a required named argument converted to `T`.
297    ///
298    /// A present argument is marked as handled.
299    pub fn expect_named<T>(&mut self, name: &str) -> Result<T, ExtensionError>
300    where
301        T: TryFrom<&'a ExtensionValue>,
302        T::Error: Into<ExtensionError>,
303    {
304        self.get_named(name)?
305            .ok_or_else(|| ExtensionError::MissingArgument {
306                name: name.to_string(),
307            })
308    }
309
310    /// Rejects arguments that the decoder did not handle.
311    pub(crate) fn finish(self) -> Result<(), ExtensionError> {
312        if !self.positional_handled && !self.args.positional.is_empty() {
313            return Err(ExtensionError::InvalidArgument(format!(
314                "Unhandled positional arguments: {}",
315                self.args.positional.len()
316            )));
317        }
318
319        let mut unhandled_args = Vec::new();
320        for name in self.args.named.keys() {
321            if !self.handled.contains(name.as_str()) {
322                unhandled_args.push(name.as_str());
323            }
324        }
325
326        if unhandled_args.is_empty() {
327            Ok(())
328        } else {
329            // Sort for stable error messages.
330            unhandled_args.sort();
331            Err(ExtensionError::InvalidArgument(format!(
332                "Unknown named arguments: {}",
333                unhandled_args.join(", ")
334            )))
335        }
336    }
337}
338
339/// A tuple-valued extension argument.
340///
341/// Tuple values preserve positional order and can be iterated by value or by
342/// reference.
343#[derive(Debug, Clone)]
344pub struct TupleValue(Vec<ExtensionValue>);
345
346impl TupleValue {
347    pub fn len(&self) -> usize {
348        self.0.len()
349    }
350
351    pub fn is_empty(&self) -> bool {
352        self.0.is_empty()
353    }
354
355    pub fn iter(&self) -> SliceIter<'_, ExtensionValue> {
356        self.0.iter()
357    }
358}
359
360impl<'a> IntoIterator for &'a TupleValue {
361    type Item = &'a ExtensionValue;
362    type IntoIter = SliceIter<'a, ExtensionValue>;
363
364    fn into_iter(self) -> Self::IntoIter {
365        self.0.iter()
366    }
367}
368
369impl IntoIterator for TupleValue {
370    type Item = ExtensionValue;
371    type IntoIter = VecIntoIter<ExtensionValue>;
372
373    fn into_iter(self) -> Self::IntoIter {
374        self.0.into_iter()
375    }
376}
377
378impl FromIterator<ExtensionValue> for TupleValue {
379    fn from_iter<I: IntoIterator<Item = ExtensionValue>>(iter: I) -> Self {
380        TupleValue(iter.into_iter().collect())
381    }
382}
383
384impl From<Vec<ExtensionValue>> for TupleValue {
385    fn from(items: Vec<ExtensionValue>) -> Self {
386        TupleValue(items)
387    }
388}
389
390/// Represents a value in extension arguments.
391///
392/// These values are the structured form of text-format extension arguments,
393/// fully resolved - i.e. any additional context (such as function anchors etc)
394/// are part of this struct itself.
395#[derive(Debug, Clone)]
396pub enum ExtensionValue {
397    /// Untyped literals. These are not input or output with types (e.g. `2`,
398    /// not `2:i64`), and suitable for protobuf extension fields that are not
399    /// substrait types.
400    String(String),
401    Integer(i64),
402    Float(f64),
403    Boolean(bool),
404    /// An untyped null literal.
405    Null,
406    /// A conversion or formatting error preserved for best-effort textification.
407    Error(ExtensionError),
408
409    /// Substrait expression value, including typed literals and field references.
410    ///
411    /// Use `TryFrom<&ExtensionValue> for Expr` when a handler accepts either an
412    /// expression or a scalar value widened into an expression.
413    Expr(Expr),
414    /// Enum value (e.g. &CORE, &Inner) — the string holds the identifier
415    /// without the `&` prefix
416    Enum(String),
417    /// Tuple of values, e.g. (&HASH, &RANGE) or (42, 'hello')
418    Tuple(TupleValue),
419    // TODO: Consider adding support for types as arguments. May need dedicated
420    // syntax (`:typename`, perhaps?), as type names may not be distinguishable
421    // from identifiers
422}
423
424/// The variant kind of an [`ExtensionValue`], used in diagnostics.
425#[derive(Debug, Clone, Copy, PartialEq, Eq)]
426pub enum ExtensionValueKind {
427    String,
428    Integer,
429    Float,
430    Boolean,
431    Null,
432    Error,
433    Reference,
434    Enum,
435    Tuple,
436    Expression,
437}
438
439impl fmt::Display for ExtensionValueKind {
440    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441        match self {
442            ExtensionValueKind::String => write!(f, "string"),
443            ExtensionValueKind::Integer => write!(f, "integer"),
444            ExtensionValueKind::Float => write!(f, "float"),
445            ExtensionValueKind::Boolean => write!(f, "boolean"),
446            ExtensionValueKind::Null => write!(f, "null"),
447            ExtensionValueKind::Error => write!(f, "error"),
448            ExtensionValueKind::Reference => write!(f, "reference"),
449            ExtensionValueKind::Enum => write!(f, "enum"),
450            ExtensionValueKind::Tuple => write!(f, "tuple"),
451            ExtensionValueKind::Expression => write!(f, "expression"),
452        }
453    }
454}
455
456impl ExtensionValue {
457    /// Return the variant kind of this value for structured diagnostics.
458    pub fn kind(&self) -> ExtensionValueKind {
459        match self {
460            ExtensionValue::String(_) => ExtensionValueKind::String,
461            ExtensionValue::Integer(_) => ExtensionValueKind::Integer,
462            ExtensionValue::Float(_) => ExtensionValueKind::Float,
463            ExtensionValue::Boolean(_) => ExtensionValueKind::Boolean,
464            ExtensionValue::Null => ExtensionValueKind::Null,
465            ExtensionValue::Error(_) => ExtensionValueKind::Error,
466            ExtensionValue::Expr(_) => ExtensionValueKind::Expression,
467            ExtensionValue::Enum(_) => ExtensionValueKind::Enum,
468            ExtensionValue::Tuple(_) => ExtensionValueKind::Tuple,
469        }
470    }
471}
472
473impl From<ExtensionError> for ExtensionValue {
474    fn from(error: ExtensionError) -> Self {
475        ExtensionValue::Error(error)
476    }
477}
478
479impl From<Expr> for ExtensionValue {
480    fn from(expr: Expr) -> Self {
481        ExtensionValue::Expr(expr)
482    }
483}
484
485impl From<proto::Expression> for ExtensionValue {
486    fn from(expr: proto::Expression) -> Self {
487        Expr::from(expr).into()
488    }
489}
490
491impl From<proto::expression::Literal> for ExtensionValue {
492    fn from(literal: proto::expression::Literal) -> Self {
493        Expr::from(literal).into()
494    }
495}
496
497impl From<Reference> for ExtensionValue {
498    fn from(reference: Reference) -> Self {
499        Expr::from(reference).into()
500    }
501}
502
503impl From<i64> for ExtensionValue {
504    fn from(value: i64) -> Self {
505        ExtensionValue::Integer(value)
506    }
507}
508
509impl From<f64> for ExtensionValue {
510    fn from(value: f64) -> Self {
511        ExtensionValue::Float(value)
512    }
513}
514
515impl From<bool> for ExtensionValue {
516    fn from(value: bool) -> Self {
517        ExtensionValue::Boolean(value)
518    }
519}
520
521impl From<String> for ExtensionValue {
522    fn from(value: String) -> Self {
523        ExtensionValue::String(value)
524    }
525}
526
527impl From<&str> for ExtensionValue {
528    fn from(value: &str) -> Self {
529        ExtensionValue::String(value.to_string())
530    }
531}
532
533impl<T> From<Vec<T>> for ExtensionValue
534where
535    T: Into<ExtensionValue>,
536{
537    fn from(values: Vec<T>) -> Self {
538        ExtensionValue::Tuple(values.into_iter().map(Into::into).collect())
539    }
540}
541
542impl ExtensionError {
543    fn invalid_type(expected: ExtensionValueKind, actual: &ExtensionValue) -> Self {
544        match actual {
545            ExtensionValue::Error(source) => Self::ArgumentConversion {
546                expected,
547                source: Box::new(source.clone()),
548            },
549            _ => Self::InvalidArgumentType {
550                expected,
551                actual: actual.kind(),
552            },
553        }
554    }
555}
556
557impl<'a> TryFrom<&'a ExtensionValue> for &'a str {
558    type Error = ExtensionError;
559
560    fn try_from(value: &'a ExtensionValue) -> Result<&'a str, Self::Error> {
561        match value {
562            ExtensionValue::String(s) => Ok(s),
563            v => Err(ExtensionError::invalid_type(ExtensionValueKind::String, v)),
564        }
565    }
566}
567
568impl TryFrom<&ExtensionValue> for String {
569    type Error = ExtensionError;
570
571    fn try_from(value: &ExtensionValue) -> Result<String, Self::Error> {
572        <&str>::try_from(value).map(ToOwned::to_owned)
573    }
574}
575
576impl TryFrom<ExtensionValue> for String {
577    type Error = ExtensionError;
578
579    fn try_from(value: ExtensionValue) -> Result<String, Self::Error> {
580        String::try_from(&value)
581    }
582}
583
584/// Helper for extracting the identifier from an [`ExtensionValue::Enum`].
585pub struct EnumValue(pub String);
586
587impl<'a> TryFrom<&'a ExtensionValue> for EnumValue {
588    type Error = ExtensionError;
589
590    fn try_from(value: &'a ExtensionValue) -> Result<EnumValue, Self::Error> {
591        match value {
592            ExtensionValue::Enum(s) => Ok(EnumValue(s.clone())),
593            v => Err(ExtensionError::invalid_type(ExtensionValueKind::Enum, v)),
594        }
595    }
596}
597
598impl<'a> TryFrom<&'a ExtensionValue> for &'a TupleValue {
599    type Error = ExtensionError;
600
601    fn try_from(value: &'a ExtensionValue) -> Result<&'a TupleValue, Self::Error> {
602        match value {
603            ExtensionValue::Tuple(tv) => Ok(tv),
604            v => Err(ExtensionError::invalid_type(ExtensionValueKind::Tuple, v)),
605        }
606    }
607}
608
609impl TryFrom<&ExtensionValue> for i64 {
610    type Error = ExtensionError;
611
612    fn try_from(value: &ExtensionValue) -> Result<i64, Self::Error> {
613        match value {
614            ExtensionValue::Integer(i) => Ok(*i),
615            v => Err(ExtensionError::invalid_type(ExtensionValueKind::Integer, v)),
616        }
617    }
618}
619
620impl TryFrom<&ExtensionValue> for f64 {
621    type Error = ExtensionError;
622
623    fn try_from(value: &ExtensionValue) -> Result<f64, Self::Error> {
624        match value {
625            ExtensionValue::Float(f) => Ok(*f),
626            v => Err(ExtensionError::invalid_type(ExtensionValueKind::Float, v)),
627        }
628    }
629}
630
631impl TryFrom<&ExtensionValue> for bool {
632    type Error = ExtensionError;
633
634    fn try_from(value: &ExtensionValue) -> Result<bool, Self::Error> {
635        match value {
636            ExtensionValue::Boolean(b) => Ok(*b),
637            v => Err(ExtensionError::invalid_type(ExtensionValueKind::Boolean, v)),
638        }
639    }
640}
641
642impl TryFrom<&ExtensionValue> for Reference {
643    type Error = ExtensionError;
644
645    fn try_from(value: &ExtensionValue) -> Result<Reference, Self::Error> {
646        match value {
647            ExtensionValue::Expr(expr) => expr
648                .as_direct_reference()
649                .map(Reference)
650                .ok_or_else(|| ExtensionError::invalid_type(ExtensionValueKind::Reference, value)),
651            v => Err(ExtensionError::invalid_type(
652                ExtensionValueKind::Reference,
653                v,
654            )),
655        }
656    }
657}
658
659impl TryFrom<&ExtensionValue> for Expr {
660    type Error = ExtensionError;
661
662    fn try_from(value: &ExtensionValue) -> Result<Expr, Self::Error> {
663        match value {
664            ExtensionValue::Expr(e) => Ok(e.clone()),
665            // Untyped extension scalars are intentionally expression-compatible:
666            // `arg=2` carries no syntax that distinguishes "configuration
667            // integer" from "i64 literal expression". Scalar-specific
668            // extraction (`i64`, `&str`, `bool`, etc.) still requires the scalar
669            // variants, while expression extraction widens them to default
670            // non-nullable Substrait literal expressions.
671            ExtensionValue::Integer(i) => Ok(Expr::from(*i)),
672            ExtensionValue::Float(f) => Ok(Expr::from(*f)),
673            ExtensionValue::String(s) => Ok(Expr::from(s.as_str())),
674            ExtensionValue::Boolean(b) => Ok(Expr::from(*b)),
675            v => Err(ExtensionError::invalid_type(
676                ExtensionValueKind::Expression,
677                v,
678            )),
679        }
680    }
681}
682
683/// Represents an output column specification.
684///
685/// These values mirror the text-format output column forms. Named columns keep
686/// the parsed Substrait type protobuf so handlers can convert directly to
687/// relation schemas.
688#[derive(Debug, Clone)]
689pub enum ExtensionColumn {
690    /// Named column with a parsed Substrait type (e.g. `name:i64?`).
691    Named {
692        /// Column name as it appears in the extension relation output.
693        name: String,
694        /// Parsed Substrait type for the column.
695        ///
696        /// This uses the protobuf field name, hence the raw identifier.
697        r#type: proto::Type,
698    },
699    /// Expression-compatible output column, including field references.
700    Expr(Expr),
701}
702
703impl ExtensionColumn {
704    /// Create an expression output column that references an existing input field (`$N`).
705    pub fn field(index: i32) -> Self {
706        Self::Expr(Expr::field(index))
707    }
708}
709
710impl ExtensionArgs {
711    /// Decodes these arguments as an [`Explainable`] extension value.
712    ///
713    /// If decoding succeeds, this method rejects any named or positional
714    /// arguments the implementation did not access. If decoding fails, it
715    /// returns that error without checking for unhandled arguments.
716    pub fn parse<T>(&self) -> Result<T, ExtensionError>
717    where
718        T: Explainable,
719    {
720        let mut access = ArgsAccess::new(self);
721        let value = T::from_args(&mut access)?;
722        access.finish()?;
723        Ok(value)
724    }
725
726    /// Push a positional extension argument.
727    pub fn push<T>(&mut self, value: T)
728    where
729        T: Into<ExtensionValue>,
730    {
731        self.positional.push(value.into());
732    }
733
734    /// Insert a named extension argument, returning any previous value.
735    pub fn insert<K, V>(&mut self, name: K, value: V) -> Option<ExtensionValue>
736    where
737        K: Into<String>,
738        V: Into<ExtensionValue>,
739    {
740        self.named.insert(name.into(), value.into())
741    }
742}
743
744#[cfg(test)]
745mod tests {
746    use super::*;
747
748    fn assert_named_invalid_type(
749        error: &ExtensionError,
750        name: &str,
751        expected: ExtensionValueKind,
752        actual: ExtensionValueKind,
753    ) {
754        assert!(
755            matches!(
756                error,
757                ExtensionError::NamedArgumentConversion {
758                    name: actual_name,
759                    source,
760                } if actual_name == name
761                    && matches!(
762                        source.as_ref(),
763                        ExtensionError::InvalidArgumentType {
764                            expected: actual_expected,
765                            actual: actual_actual,
766                        } if *actual_expected == expected && *actual_actual == actual
767                    )
768            ),
769            "unexpected error: {error:?}"
770        );
771    }
772
773    #[test]
774    fn get_named_converts_present_values_and_returns_none_for_missing_values() {
775        let mut args = ExtensionArgs::default();
776        args.insert("count", 8_i64);
777        let mut access = ArgsAccess::new(&args);
778
779        assert_eq!(access.get_named::<i64>("count").unwrap(), Some(8));
780        assert_eq!(access.get_named::<i64>("missing").unwrap(), None);
781        assert!(access.finish().is_ok());
782    }
783
784    #[test]
785    fn get_named_contextualizes_conversion_errors() {
786        let mut args = ExtensionArgs::default();
787        args.insert("count", ExtensionValue::Null);
788        let mut access = ArgsAccess::new(&args);
789
790        let error = access
791            .get_named::<i64>("count")
792            .expect_err("null should not convert to i64");
793
794        assert_eq!(
795            error.to_string(),
796            "Invalid named argument 'count': Invalid argument: expected integer, got null"
797        );
798        assert!(access.finish().is_ok());
799    }
800
801    #[test]
802    fn get_named_tuple_converts_present_values_and_returns_none_for_missing_values() {
803        let mut args = ExtensionArgs::default();
804        args.insert(
805            "names",
806            ExtensionValue::Tuple(vec!["first".into(), "second".into()].into()),
807        );
808        let mut access = ArgsAccess::new(&args);
809
810        assert_eq!(access.get_named_tuple::<String>("missing").unwrap(), None);
811        assert_eq!(
812            access.get_named_tuple::<String>("names").unwrap(),
813            Some(vec!["first".to_string(), "second".to_string()])
814        );
815        assert!(access.finish().is_ok());
816    }
817
818    #[test]
819    fn get_named_tuple_contextualizes_invalid_outer_value() {
820        let mut args = ExtensionArgs::default();
821        args.insert("names", "not a tuple");
822        let mut access = ArgsAccess::new(&args);
823
824        let error = access
825            .get_named_tuple::<String>("names")
826            .expect_err("string should not convert to tuple");
827
828        assert_named_invalid_type(
829            &error,
830            "names",
831            ExtensionValueKind::Tuple,
832            ExtensionValueKind::String,
833        );
834        assert!(access.finish().is_ok());
835    }
836
837    #[test]
838    fn get_named_tuple_contextualizes_invalid_element() {
839        let mut args = ExtensionArgs::default();
840        args.insert(
841            "names",
842            ExtensionValue::Tuple(vec!["first".into(), 2_i64.into()].into()),
843        );
844        let mut access = ArgsAccess::new(&args);
845
846        let error = access
847            .get_named_tuple::<String>("names")
848            .expect_err("integer should not convert to string");
849
850        assert_named_invalid_type(
851            &error,
852            "names",
853            ExtensionValueKind::String,
854            ExtensionValueKind::Integer,
855        );
856        assert!(access.finish().is_ok());
857    }
858
859    #[test]
860    fn vector_encodes_as_tuple() {
861        let encoded: ExtensionValue = vec![1_i64, 2_i64].into();
862        let tuple = <&TupleValue>::try_from(&encoded).unwrap();
863        let values = tuple
864            .iter()
865            .map(i64::try_from)
866            .collect::<Result<Vec<_>, _>>()
867            .unwrap();
868
869        assert_eq!(values, vec![1, 2]);
870    }
871
872    #[test]
873    fn error_value_reports_expected_type_and_source_when_extracted() {
874        let value = ExtensionValue::Error(ExtensionError::Custom("bad value".to_string()));
875
876        let error = i64::try_from(&value).expect_err("error value should not convert");
877
878        assert_eq!(
879            error.to_string(),
880            "Cannot convert argument to integer: bad value"
881        );
882        assert!(matches!(
883            &error,
884            ExtensionError::ArgumentConversion {
885                expected: ExtensionValueKind::Integer,
886                source,
887            } if matches!(source.as_ref(), ExtensionError::Custom(message) if message == "bad value")
888        ));
889    }
890
891    #[test]
892    fn expect_named_reports_missing_argument_name() {
893        let args = ExtensionArgs::default();
894        let mut access = ArgsAccess::new(&args);
895
896        let error = access
897            .expect_named::<i64>("count")
898            .expect_err("missing argument should fail");
899
900        assert_eq!(error.to_string(), "Missing required argument: count");
901        assert!(access.finish().is_ok());
902    }
903}