From a6d4042f892371ec2ecab08c9500b3a9ed2050d1 Mon Sep 17 00:00:00 2001 From: Wendell Smith Date: Fri, 7 Aug 2026 17:31:22 -0400 Subject: [PATCH] refactor(extensions)!: centralize argument extraction --- API.md | 36 +++++--- examples/extensions.rs | 18 ++-- src/cli.rs | 8 +- src/extensions/args.rs | 144 +++++++++++++++++-------------- src/extensions/examples.rs | 51 +++++------ src/extensions/mod.rs | 3 +- src/extensions/registry.rs | 140 +++++++++++++++++++++--------- src/lib.rs | 2 +- tests/adv_extension_roundtrip.rs | 22 ++--- tests/extension_roundtrip.rs | 64 ++++++-------- tests/extension_table.rs | 10 +-- tests/json_parsing.rs | 10 +-- 12 files changed, 280 insertions(+), 228 deletions(-) diff --git a/API.md b/API.md index 7f625704..c18c1311 100644 --- a/API.md +++ b/API.md @@ -101,7 +101,7 @@ Substrait has two extension mechanisms: - **Simple extensions** declare extension functions, types, and type variations. `substrait-explain` reads and writes these declarations in the `=== - Extensions` section and uses them to resolve anchors while parsing and +Extensions` section and uses them to resolve anchors while parsing and formatting expressions and types. No YAML files are read; substrait-explain relies on the protobufs / text format itself for function names, and does not validate type signatures exist / match extensions. @@ -137,9 +137,10 @@ extension arguments. `substrait-explain` handles parsing and rendering the text syntax around those arguments: - `name()` - The extension name used in text (e.g., `"ParquetScan"`) -- `from_args(args)` - Parse text arguments into your type -- `to_args(&self, context)` - Convert your type to text arguments, optionally - using information about relation inputs. +- `from_args(args)` - Build your extension value from text arguments with + `ArgsAccess` +- `to_args(&self, context)` - Convert your extension value to text arguments, + optionally using information about relation inputs The extension API works across three representations: @@ -151,7 +152,7 @@ The extension API works across three representations: Untyped scalar extension literals such as `2` or `'path'` are represented as scalar `ExtensionValue` variants and render without expression type suffixes, even in verbose output. The same values can still be requested as `Expr` through -`ArgsExtractor`, which widens them to default non-nullable Substrait literal +`ArgsAccess`, which widens them to default non-nullable Substrait literal expressions. Typed literals, field references, function calls, and casts are represented as expression values. @@ -164,11 +165,21 @@ The `to_args` method receives an `ExtensionContext`. Relation extensions get one emitted column count, including any output mapping applied by that child. Other extension namespaces receive an empty input slice. -Use `ArgsExtractor` for convenient argument parsing: +### Reading Extension Arguments -- `extractor.expect_named::(name)` - Required argument -- `extractor.get_named::(name)?` - Optional argument, returning `Option` -- `extractor.check_exhausted()` - Verify no unexpected arguments +The registry passes [`ArgsAccess`] to your [`Explainable::from_args`] +implementation. It provides access to positional arguments, named arguments, +and the relation's output columns declared by custom relations. + +If `from_args` returns successfully, the registry rejects any positional or +named arguments the implementation did not access. If `from_args` returns an +error, the registry preserves that error without checking for unhandled +arguments. Output columns are relation metadata and are not included in this +check. + +Call [`ExtensionArgs::parse`](extensions::ExtensionArgs::parse) to perform the +same checked conversion directly, such as when testing an `Explainable` +implementation. ### Extension Namespaces @@ -189,7 +200,8 @@ Register extensions to the appropriate namespace: ```rust,no_run # use prost::{Message, Name}; # use substrait_explain::extensions::{ -# Explainable, ExtensionArgs, ExtensionContext, ExtensionError, ExtensionRegistry, +# ArgsAccess, Explainable, ExtensionArgs, ExtensionContext, ExtensionError, +# ExtensionRegistry, # }; #[derive(Clone, PartialEq, Message)] pub struct MySourceConfig { @@ -204,7 +216,9 @@ impl Name for MySourceConfig { } # impl Explainable for MySourceConfig { # fn name() -> &'static str { "MySource" } -# fn from_args(_: &ExtensionArgs) -> Result { Ok(Self::default()) } +# fn from_args(_: &mut ArgsAccess<'_>) -> Result { +# Ok(Self::default()) +# } # fn to_args( # &self, # _context: &ExtensionContext<'_>, diff --git a/examples/extensions.rs b/examples/extensions.rs index 349e9969..c417cf84 100644 --- a/examples/extensions.rs +++ b/examples/extensions.rs @@ -15,8 +15,8 @@ use prost::{Message, Name}; use substrait::proto::{self, Plan, PlanRel, Rel, plan_rel, rel}; use substrait_explain::extensions::any::AnyRef; use substrait_explain::extensions::{ - AnyConvertible, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, - ExtensionRegistry, + AnyConvertible, ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, + ExtensionError, ExtensionRegistry, }; use substrait_explain::{OutputOptions, Parser, format_with_registry}; @@ -69,19 +69,15 @@ impl Explainable for ParquetScanConfig { "TypedParquetScan" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); + fn from_args(args: &mut ArgsAccess<'_>) -> Result { // path is required - let path: &str = extractor.expect_named("path")?; + let path: &str = args.expect_named("path")?; // batch_size and use_dictionary are optional, with default values - let batch_size: i64 = extractor.get_named("batch_size")?.unwrap_or(1024); - let use_dictionary: bool = extractor.get_named("use_dictionary")?.unwrap_or(true); - - // Validate there are no other named arguments - extractor.check_exhausted()?; + let batch_size: i64 = args.get_named("batch_size")?.unwrap_or(1024); + let use_dictionary: bool = args.get_named("use_dictionary")?.unwrap_or(true); let selected_columns = args - .output_columns + .output_columns() .iter() .map(|column| match column { ExtensionColumn::Named { name, r#type } => Ok(ParquetColumn { diff --git a/src/cli.rs b/src/cli.rs index 44bf40e6..c8257006 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -480,7 +480,7 @@ mod tests { use super::*; use crate::extensions::{ - Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, + ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, }; use crate::fixtures::parse_type; use crate::parse; @@ -907,10 +907,8 @@ Root[result] "TestSource" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let tag: &str = extractor.expect_named("tag")?; - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let tag: &str = args.expect_named("tag")?; Ok(TestSource { tag: tag.to_string(), }) diff --git a/src/extensions/args.rs b/src/extensions/args.rs index f60b1746..9a19a9bb 100644 --- a/src/extensions/args.rs +++ b/src/extensions/args.rs @@ -26,9 +26,9 @@ //! into default Substrait literal expressions. use std::collections::HashSet; +use std::fmt; use std::slice::Iter as SliceIter; use std::vec::IntoIter as VecIntoIter; -use std::{fmt, thread}; use indexmap::IndexMap; use substrait::proto; @@ -36,7 +36,7 @@ use substrait::proto::expression::field_reference::ReferenceType; use substrait::proto::expression::literal::LiteralType; use substrait::proto::expression::{RexType, reference_segment}; -use super::ExtensionError; +use super::{Explainable, ExtensionError}; use crate::textify::expressions::Reference; /// Kind of relation addendum in the text format. @@ -198,43 +198,60 @@ pub struct ExtensionArgs { pub output_columns: Vec, } -/// Helper struct for extracting named arguments with validation. +/// [`ArgsAccess`] provides a view of the arguments in the text form of an +/// extension relation or advanced extension. /// -/// Tracks which arguments have been consumed. Callers **must** call -/// [`check_exhausted`](ArgsExtractor::check_exhausted) before dropping to -/// verify no unexpected arguments remain. In debug builds, dropping without -/// calling `check_exhausted` will panic. This catches [`Explainable`](super::Explainable) -/// implementations that forget to reject unexpected named arguments. -pub struct ArgsExtractor<'a> { +/// Positional arguments are available via [`Self::positional`], named arguments +/// via [`Self::get_named`] or [`Self::expect_named`] which can do type +/// conversion as well, and output columns via [`Self::output_columns`]. +/// +/// [`ArgsAccess`] tracks which arguments are accessed; when +/// [`Explainable::from_args`] returns, any unaccessed arguments will be raised +/// as [`ExtensionError::InvalidArgument`] errors. +pub struct ArgsAccess<'a> { args: &'a ExtensionArgs, - consumed: HashSet<&'a str>, - checked: bool, + handled: HashSet<&'a str>, + positional_handled: bool, } -impl<'a> ArgsExtractor<'a> { - /// Create a new extractor for the given arguments - pub fn new(args: &'a ExtensionArgs) -> Self { +impl<'a> ArgsAccess<'a> { + pub(crate) fn new(args: &'a ExtensionArgs) -> Self { Self { args, - consumed: HashSet::new(), - checked: false, + handled: HashSet::new(), + positional_handled: false, } } - /// Get a named argument value, marking it as consumed if found. + /// Returns the positional arguments in source order and marks them as handled. + pub fn positional(&mut self) -> &'a [ExtensionValue] { + self.positional_handled = true; + &self.args.positional + } + + /// Returns the output columns for a custom relation. + /// + /// Output columns are not included in the unhandled-argument check. + pub fn output_columns(&self) -> &'a [ExtensionColumn] { + &self.args.output_columns + } + + /// Returns a named argument without converting it, or `None` if it is absent. + /// + /// A present argument is marked as handled. pub fn get_named_arg(&mut self, name: &str) -> Option<&'a ExtensionValue> { match self.args.named.get_key_value(name) { Some((k, value)) => { - self.consumed.insert(k); + self.handled.insert(k); Some(value) } None => None, } } - /// Get a named argument converted to `T`, or `None` if it is absent. + /// Returns a named argument converted to `T`, or `None` if it is absent. /// - /// Marks the argument as consumed if it exists in the source args. + /// A present argument is marked as handled. pub fn get_named(&mut self, name: &str) -> Result, ExtensionError> where T: TryFrom<&'a ExtensionValue>, @@ -250,9 +267,9 @@ impl<'a> ArgsExtractor<'a> { .transpose() } - /// Get a required named argument converted to `T`. + /// Returns a required named argument converted to `T`. /// - /// Marks the argument as consumed if found. + /// A present argument is marked as handled. pub fn expect_named(&mut self, name: &str) -> Result where T: TryFrom<&'a ExtensionValue>, @@ -264,48 +281,35 @@ impl<'a> ArgsExtractor<'a> { }) } - /// Check that all named arguments in the source have been consumed, - /// returning an error if not. - /// - /// Must be called before the extractor is dropped, to validate that all - /// args are correctly handled. In debug builds, dropping without calling - /// this method will panic. - pub fn check_exhausted(&mut self) -> Result<(), ExtensionError> { - self.checked = true; + /// Rejects arguments that the decoder did not handle. + pub(crate) fn finish(self) -> Result<(), ExtensionError> { + if !self.positional_handled && !self.args.positional.is_empty() { + return Err(ExtensionError::InvalidArgument(format!( + "Unhandled positional arguments: {}", + self.args.positional.len() + ))); + } - let mut unknown_args = Vec::new(); + let mut unhandled_args = Vec::new(); for name in self.args.named.keys() { - if !self.consumed.contains(name.as_str()) { - unknown_args.push(name.as_str()); + if !self.handled.contains(name.as_str()) { + unhandled_args.push(name.as_str()); } } - if unknown_args.is_empty() { + if unhandled_args.is_empty() { Ok(()) } else { - // Sort for stable error messages - unknown_args.sort(); + // Sort for stable error messages. + unhandled_args.sort(); Err(ExtensionError::InvalidArgument(format!( "Unknown named arguments: {}", - unknown_args.join(", ") + unhandled_args.join(", ") ))) } } } -impl Drop for ArgsExtractor<'_> { - fn drop(&mut self) { - if self.checked || thread::panicking() { - return; - } - // If we get here, the caller forgot to call check_exhausted(). - debug_assert!( - false, - "ArgsExtractor dropped without calling check_exhausted()" - ); - } -} - /// A tuple-valued extension argument. /// /// Tuple values preserve positional order and can be iterated by value or by @@ -661,6 +665,21 @@ impl ExtensionColumn { } impl ExtensionArgs { + /// Decodes these arguments as an [`Explainable`] extension value. + /// + /// If decoding succeeds, this method rejects any named or positional + /// arguments the implementation did not access. If decoding fails, it + /// returns that error without checking for unhandled arguments. + pub fn parse(&self) -> Result + where + T: Explainable, + { + let mut access = ArgsAccess::new(self); + let value = T::from_args(&mut access)?; + access.finish()?; + Ok(value) + } + /// Push a positional extension argument. pub fn push(&mut self, value: T) where @@ -677,11 +696,6 @@ impl ExtensionArgs { { self.named.insert(name.into(), value.into()) } - - /// Create an extractor for validating named arguments - pub fn extractor(&self) -> ArgsExtractor<'_> { - ArgsExtractor::new(self) - } } #[cfg(test)] @@ -692,20 +706,20 @@ mod tests { fn get_named_converts_present_values_and_returns_none_for_missing_values() { let mut args = ExtensionArgs::default(); args.insert("count", 8_i64); - let mut extractor = args.extractor(); + let mut access = ArgsAccess::new(&args); - assert_eq!(extractor.get_named::("count").unwrap(), Some(8)); - assert_eq!(extractor.get_named::("missing").unwrap(), None); - assert!(extractor.check_exhausted().is_ok()); + assert_eq!(access.get_named::("count").unwrap(), Some(8)); + assert_eq!(access.get_named::("missing").unwrap(), None); + assert!(access.finish().is_ok()); } #[test] fn get_named_contextualizes_conversion_errors() { let mut args = ExtensionArgs::default(); args.insert("count", ExtensionValue::Null); - let mut extractor = args.extractor(); + let mut access = ArgsAccess::new(&args); - let error = extractor + let error = access .get_named::("count") .expect_err("null should not convert to i64"); @@ -713,7 +727,7 @@ mod tests { error.to_string(), "Invalid named argument 'count': Invalid argument: expected integer, got null" ); - assert!(extractor.check_exhausted().is_ok()); + assert!(access.finish().is_ok()); } #[test] @@ -738,13 +752,13 @@ mod tests { #[test] fn expect_named_reports_missing_argument_name() { let args = ExtensionArgs::default(); - let mut extractor = args.extractor(); + let mut access = ArgsAccess::new(&args); - let error = extractor + let error = access .expect_named::("count") .expect_err("missing argument should fail"); assert_eq!(error.to_string(), "Missing required argument: count"); - assert!(extractor.check_exhausted().is_ok()); + assert!(access.finish().is_ok()); } } diff --git a/src/extensions/examples.rs b/src/extensions/examples.rs index 756fc2e5..0689118e 100644 --- a/src/extensions/examples.rs +++ b/src/extensions/examples.rs @@ -15,7 +15,7 @@ //! This hidden module is crate-owned example support for doctests and //! integration tests, not a stable extension API. -use crate::extensions::args::{EnumValue, ExtensionArgs, ExtensionValue}; +use crate::extensions::args::{ArgsAccess, EnumValue, ExtensionArgs, ExtensionValue}; use crate::extensions::registry::{ Explainable, ExtensionContext, ExtensionError, ExtensionRegistry, }; @@ -121,10 +121,10 @@ impl Explainable for PartitionHint { "PartitionHint" } - fn from_args(args: &ExtensionArgs) -> Result { - // Positional arguments are PartitionStrategy enum values. + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + // Read each positional argument as a PartitionStrategy enum. let strategies: Result, ExtensionError> = args - .positional + .positional() .iter() .map(|val| { let EnumValue(ident) = EnumValue::try_from(val)?; @@ -139,9 +139,7 @@ impl Explainable for PartitionHint { }) .collect(); - let mut extractor = args.extractor(); - let count: i64 = extractor.get_named("count")?.unwrap_or_default(); - extractor.check_exhausted()?; + let count: i64 = args.get_named("count")?.unwrap_or_default(); Ok(PartitionHint { strategies: strategies?, @@ -219,16 +217,14 @@ impl Explainable for PlanHint { "PlanHint" } - fn from_args(args: &ExtensionArgs) -> Result { - if !args.positional.is_empty() { + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + if !args.positional().is_empty() { return Err(ExtensionError::InvalidArgument( "PlanHint does not accept positional arguments".to_owned(), )); } - let mut extractor = args.extractor(); - let hint: String = extractor.expect_named::<&str>("hint")?.to_owned(); - extractor.check_exhausted()?; + let hint: String = args.expect_named::<&str>("hint")?.to_owned(); Ok(PlanHint { hint }) } @@ -303,24 +299,22 @@ impl Explainable for BlobStoreRead { "BlobStoreRead" } - fn from_args(args: &ExtensionArgs) -> Result { - if args.positional.len() != 1 { + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + if args.positional().len() != 1 { return Err(ExtensionError::InvalidArgument(format!( "BlobStoreRead expects exactly 1 positional path argument, got {}", - args.positional.len() + args.positional().len() ))); } - if !args.output_columns.is_empty() { + if !args.output_columns().is_empty() { return Err(ExtensionError::InvalidArgument( "BlobStoreRead output columns belong in Read:Extension[...]".to_owned(), )); } - let path = <&str>::try_from(&args.positional[0])?.to_owned(); - let mut extractor = args.extractor(); - let limit: i64 = extractor.get_named("limit")?.unwrap_or_default(); - let include_archived: bool = extractor.get_named("include_archived")?.unwrap_or(false); - extractor.check_exhausted()?; + let path = <&str>::try_from(&args.positional()[0])?.to_owned(); + let limit: i64 = args.get_named("limit")?.unwrap_or_default(); + let include_archived: bool = args.get_named("include_archived")?.unwrap_or(false); Ok(Self { path, @@ -401,7 +395,7 @@ mod tests { fn from_args_round_trip() { let original = make_hint(vec![PartitionStrategy::Hash, PartitionStrategy::Range], 16); let args = original.to_args(&ExtensionContext::default()).unwrap(); - let decoded = PartitionHint::from_args(&args).unwrap(); + let decoded = args.parse::().unwrap(); assert_eq!(original, decoded); } @@ -410,7 +404,7 @@ mod tests { let mut args = ExtensionArgs::default(); args.positional .push(ExtensionValue::Enum("BOGUS".to_owned())); - assert!(PartitionHint::from_args(&args).is_err()); + assert!(args.parse::().is_err()); } #[test] @@ -418,7 +412,7 @@ mod tests { // An integer positional arg where an enum is expected should fail. let mut args = ExtensionArgs::default(); args.push(1_i64); - let result = PartitionHint::from_args(&args); + let result = args.parse::(); assert!( result.is_err(), "expected error for non-enum positional arg, got {result:?}" @@ -427,10 +421,9 @@ mod tests { #[test] fn from_args_rejects_extra_named_args() { - // check_exhausted should reject unknown named args. let mut args = ExtensionArgs::default(); args.insert("unknown_key", 99_i64); - let result = PartitionHint::from_args(&args); + let result = args.parse::(); assert!( result.is_err(), "expected error for unknown named arg, got {result:?}" @@ -441,7 +434,7 @@ mod tests { fn from_args_empty_strategies_roundtrip() { let original = make_hint(vec![], 0); let args = original.to_args(&ExtensionContext::default()).unwrap(); - let decoded = PartitionHint::from_args(&args).unwrap(); + let decoded = args.parse::().unwrap(); assert_eq!(original, decoded); assert!(decoded.strategies.is_empty()); assert_eq!(decoded.count, 0); @@ -473,7 +466,7 @@ mod tests { hint: "use_index".to_owned(), }; let args = original.to_args(&ExtensionContext::default()).unwrap(); - let decoded = PlanHint::from_args(&args).unwrap(); + let decoded = args.parse::().unwrap(); assert_eq!(original, decoded); } @@ -507,7 +500,7 @@ mod tests { }; let args = original.to_args(&ExtensionContext::default()).unwrap(); - let decoded = BlobStoreRead::from_args(&args).unwrap(); + let decoded = args.parse::().unwrap(); assert_eq!(original, decoded); } diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index c18bce33..33f87d79 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -19,7 +19,8 @@ pub mod examples; pub use any::{Any, AnyRef}; pub(crate) use args::AddendumKind; pub use args::{ - EnumValue, Expr, ExtensionArgs, ExtensionColumn, ExtensionValue, ExtensionValueKind, TupleValue, + ArgsAccess, EnumValue, Expr, ExtensionArgs, ExtensionColumn, ExtensionValue, + ExtensionValueKind, TupleValue, }; pub use registry::{ AnyConvertible, Explainable, Extension, ExtensionContext, ExtensionError, ExtensionInput, diff --git a/src/extensions/registry.rs b/src/extensions/registry.rs index 72232c1b..1f12e376 100644 --- a/src/extensions/registry.rs +++ b/src/extensions/registry.rs @@ -26,8 +26,8 @@ //! //! ```rust //! use substrait_explain::extensions::{ -//! Any, AnyConvertible, AnyRef, Explainable, ExtensionArgs, ExtensionContext, ExtensionError, -//! ExtensionRegistry, +//! Any, AnyConvertible, AnyRef, ArgsAccess, Explainable, ExtensionArgs, ExtensionContext, +//! ExtensionError, ExtensionRegistry, //! }; //! //! // Define a custom extension type @@ -60,10 +60,8 @@ //! "ParquetScan" //! } //! -//! fn from_args(args: &ExtensionArgs) -> Result { -//! let mut extractor = args.extractor(); -//! let path: &str = extractor.expect_named("path")?; -//! extractor.check_exhausted()?; +//! fn from_args(args: &mut ArgsAccess<'_>) -> Result { +//! let path: &str = args.expect_named("path")?; //! Ok(CustomScanConfig { //! path: path.to_string(), //! }) @@ -94,7 +92,7 @@ use substrait::proto::r#type::{Nullability, Struct}; use thiserror::Error; use crate::extensions::any::{Any, AnyRef}; -use crate::extensions::args::{ExtensionArgs, ExtensionColumn, ExtensionValueKind}; +use crate::extensions::args::{ArgsAccess, ExtensionArgs, ExtensionColumn, ExtensionValueKind}; /// Type of extension in the registry, used for namespace separation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -173,7 +171,7 @@ pub enum ExtensionError { #[error("Extension '{name}' not found in registry")] NotFound { name: String }, - /// Required argument not present (from ArgsExtractor) + /// A required named argument was not present. #[error("Missing required argument: {name}")] MissingArgument { name: String }, @@ -335,16 +333,26 @@ impl ExtensionProtoConvert> for NamedStruct { } } -/// Trait for types that participate in text explanations. +/// Maps a registered extension type to and from text-format arguments. +/// +/// Implement this trait for each protobuf extension type you register with an +/// [`ExtensionRegistry`]. pub trait Explainable: Sized { - /// Canonical textual name for this extension. This is what appears in - /// Substrait text plans and how the registry identifies the type. + /// Returns the name used in the text format and the extension registry. fn name() -> &'static str; - /// Parse extension arguments into this type - fn from_args(args: &ExtensionArgs) -> Result; - - /// Convert this type to extension arguments + /// Builds your extension value from its text-format arguments. + /// + /// Read every named and positional argument the extension accepts from + /// `args`. After this method returns `Ok`, [`ExtensionArgs::parse`] and the + /// registry reject any arguments you did not access. If it returns `Err`, + /// they return that error without checking for unhandled arguments. Output + /// columns are not included in this check. + fn from_args(args: &mut ArgsAccess<'_>) -> Result; + + /// Builds the text-format arguments for this value. + /// + /// The [`context`](ExtensionContext) provides the relation's inputs. fn to_args(&self, context: &ExtensionContext<'_>) -> Result; } @@ -391,7 +399,7 @@ struct ExtensionAdapter(PhantomData); impl ExtensionConverter for ExtensionAdapter { fn parse_detail(&self, args: &ExtensionArgs) -> Result { - T::from_args(args)?.to_any() + args.parse::()?.to_any() } fn textify_detail( @@ -779,11 +787,9 @@ mod tests { "TestExtension" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let path: String = extractor.expect_named::<&str>("path")?.to_string(); - let batch_size: i64 = extractor.expect_named("batch_size")?; - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let path: String = args.expect_named::<&str>("path")?.to_string(); + let batch_size: i64 = args.expect_named("batch_size")?; Ok(TestExtension { path: path.to_string(), @@ -862,6 +868,64 @@ mod tests { ); } + #[test] + fn argument_error_takes_precedence_over_unhandled_arguments() { + let mut registry = ExtensionRegistry::new(); + registry.register_relation::().unwrap(); + + let mut args = ExtensionArgs::default(); + args.insert("unexpected", true); + + let error = registry + .parse_extension("TestExtension", &args) + .expect_err("missing path should return an error"); + + assert!(matches!( + error, + ExtensionError::MissingArgument { name } if name == "path" + )); + } + + #[test] + fn successful_decode_rejects_unhandled_named_arguments() { + let mut registry = ExtensionRegistry::new(); + registry.register_relation::().unwrap(); + + let mut args = ExtensionArgs::default(); + args.insert("path", "data.parquet"); + args.insert("batch_size", 2048_i64); + args.insert("unexpected", true); + + let error = registry + .parse_extension("TestExtension", &args) + .expect_err("unhandled argument should fail"); + + assert_eq!( + error.to_string(), + "Invalid argument: Unknown named arguments: unexpected" + ); + } + + #[test] + fn successful_decode_rejects_unhandled_positional_arguments() { + let mut registry = ExtensionRegistry::new(); + registry.register_relation::().unwrap(); + + let mut args = ExtensionArgs::default(); + args.push(true); + args.insert("path", "data.parquet"); + args.insert("batch_size", 2048_i64); + + let error = registry + .parse_extension("TestExtension", &args) + .expect_err("unhandled positional argument should fail"); + + assert_eq!( + error.to_string(), + "Invalid argument: Unhandled positional arguments: 1" + ); + } + #[test] fn test_extension_table_registry_basic() { let mut registry = ExtensionRegistry::new(); @@ -911,17 +975,17 @@ mod tests { r#type: parse_type("i32"), }); - // Test retrieval - use extractor - let mut extractor = args.extractor(); + let mut access = ArgsAccess::new(&args); - let path = extractor.get_named_arg("path").unwrap(); + let path = access.get_named_arg("path").unwrap(); assert_eq!(<&str>::try_from(path).unwrap(), "data/*.parquet"); - let batch_size = extractor.get_named_arg("batch_size").unwrap(); + let batch_size = access.get_named_arg("batch_size").unwrap(); assert_eq!(i64::try_from(batch_size).unwrap(), 1024); + assert_eq!(access.positional().len(), 1); - // Verify they were consumed - assert!(extractor.check_exhausted().is_ok()); + // Output columns are not part of argument validation. + assert!(access.finish().is_ok()); assert_eq!(args.positional.len(), 1); assert_eq!(args.output_columns.len(), 1); @@ -936,20 +1000,18 @@ mod tests { let result = registry.parse_extension("NonExistent", &args); assert!(matches!(result, Err(ExtensionError::NotFound { .. }))); - // Missing argument let args = ExtensionArgs::default(); - let mut extractor = args.extractor(); - let result = extractor.get_named_arg("missing"); + let mut access = ArgsAccess::new(&args); + let result = access.get_named_arg("missing"); assert!(result.is_none()); - assert!(extractor.check_exhausted().is_ok()); + assert!(access.finish().is_ok()); - // Type check example let mut args = ExtensionArgs::default(); args.insert("test", 42_i64); - let mut extractor = args.extractor(); - let result = extractor.get_named_arg("test"); + let mut access = ArgsAccess::new(&args); + let result = access.get_named_arg("test"); assert_eq!(i64::try_from(result.unwrap()).unwrap(), 42); - assert!(extractor.check_exhausted().is_ok()); + assert!(access.finish().is_ok()); } // Mock enhancement type for testing namespace separation @@ -986,10 +1048,8 @@ mod tests { "TestEnhancement" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let hint: String = extractor.expect_named::<&str>("hint")?.to_string(); - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let hint: String = args.expect_named::<&str>("hint")?.to_string(); Ok(TestEnhancement { hint }) } @@ -1129,7 +1189,7 @@ mod tests { "ConflictingExtension" } - fn from_args(_args: &ExtensionArgs) -> Result { + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(ConflictingExtension) } diff --git a/src/lib.rs b/src/lib.rs index 9d1ddf72..dfd6ecfd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,7 +2,7 @@ // Used as links in API.md #[cfg(doc)] -pub use extensions::{AnyConvertible, Explainable, ExtensionRegistry}; +pub use extensions::{AnyConvertible, ArgsAccess, Explainable, ExtensionRegistry}; pub mod extensions; pub mod grammar; diff --git a/tests/adv_extension_roundtrip.rs b/tests/adv_extension_roundtrip.rs index be63db6a..4d0492ff 100644 --- a/tests/adv_extension_roundtrip.rs +++ b/tests/adv_extension_roundtrip.rs @@ -168,7 +168,7 @@ Root[result] mod opt_fixture { use prost::Name; use substrait_explain::extensions::{ - Explainable, ExtensionArgs, ExtensionContext, ExtensionError, + ArgsAccess, Explainable, ExtensionArgs, ExtensionContext, ExtensionError, }; #[derive(Clone, PartialEq, prost::Message)] @@ -195,10 +195,8 @@ mod opt_fixture { "PlanHint" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let hint: String = extractor.expect_named::<&str>("hint")?.to_owned(); - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let hint: String = args.expect_named::<&str>("hint")?.to_owned(); Ok(PlanHint { hint }) } @@ -527,7 +525,7 @@ Root[result] mod extension_child_fixture { use prost::Name; use substrait_explain::extensions::{ - Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, + ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, }; #[derive(Clone, PartialEq, prost::Message)] @@ -551,11 +549,7 @@ mod extension_child_fixture { "TwoColumnScan" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - extractor.check_exhausted()?; - // Output columns are validated by the parser; we just ignore them here. - let _ = &args.output_columns; + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(TwoColumnScan {}) } @@ -850,7 +844,7 @@ Root[result] mod adv_ext_with_columns_fixture { use prost::Name; use substrait_explain::extensions::{ - Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, + ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, }; #[derive(Clone, PartialEq, prost::Message)] @@ -874,9 +868,7 @@ mod adv_ext_with_columns_fixture { "EnhancementWithColumns" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - extractor.check_exhausted()?; + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(EnhancementWithColumns {}) } diff --git a/tests/extension_roundtrip.rs b/tests/extension_roundtrip.rs index 82891a5d..b1606e11 100644 --- a/tests/extension_roundtrip.rs +++ b/tests/extension_roundtrip.rs @@ -10,8 +10,8 @@ use substrait::proto::expression::literal::LiteralType; use substrait::proto::r#type::Nullability; use substrait_explain::extensions::examples::PartitionHint; use substrait_explain::extensions::{ - EnumValue, Explainable, Expr, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, - ExtensionProtoConvert, ExtensionRegistry, ExtensionValue, TupleValue, + ArgsAccess, EnumValue, Explainable, Expr, ExtensionArgs, ExtensionColumn, ExtensionContext, + ExtensionError, ExtensionProtoConvert, ExtensionRegistry, ExtensionValue, TupleValue, }; use substrait_explain::{Parser, format_with_registry}; @@ -49,17 +49,14 @@ impl Explainable for UserTableConfig { "UserTable" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let table_name: &str = extractor.expect_named("name")?; - let version: i64 = extractor.get_named("version")?.unwrap_or(1); - let is_temporary: bool = extractor.get_named("temp")?.unwrap_or(false); - - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let table_name: &str = args.expect_named("name")?; + let version: i64 = args.get_named("version")?.unwrap_or(1); + let is_temporary: bool = args.get_named("temp")?.unwrap_or(false); // Extract columns from output columns to populate tracked_columns let mut tracked_columns = Vec::new(); - for col in &args.output_columns { + for col in args.output_columns() { match col { ExtensionColumn::Named { name, .. } => { tracked_columns.push(name.clone()); @@ -161,10 +158,8 @@ impl Explainable for FilterConfig { "TestFilter" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let expression: String = extractor.expect_named::<&str>("expr")?.to_string(); - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let expression: String = args.expect_named::<&str>("expr")?.to_string(); Ok(FilterConfig { expression }) } @@ -230,13 +225,12 @@ impl Explainable for LiteralConfig { "LiteralTest" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let path: String = extractor.expect_named::<&str>("path")?.to_string(); - let big: i64 = extractor.expect_named("big")?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let path: String = args.expect_named::<&str>("path")?.to_string(); + let big: i64 = args.expect_named("big")?; // Manually handle ratio to support both Integer and Float types - let ratio = match extractor.get_named_arg("ratio") { + let ratio = match args.get_named_arg("ratio") { Some(ExtensionValue::Float(f)) => *f, Some(ExtensionValue::Integer(i)) => *i as f64, Some(v) => { @@ -252,9 +246,7 @@ impl Explainable for LiteralConfig { } }; - let enabled: bool = extractor.expect_named("enabled")?; - - extractor.check_exhausted()?; + let enabled: bool = args.expect_named("enabled")?; Ok(LiteralConfig { path, @@ -344,9 +336,7 @@ impl Explainable for EmptySource { "EmptySource" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - extractor.check_exhausted()?; + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(EmptySource { marker: "empty".to_string(), }) @@ -425,9 +415,7 @@ impl Explainable for PassThroughWrapper { "PassThrough" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - extractor.check_exhausted()?; + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(PassThroughWrapper {}) } @@ -464,9 +452,7 @@ impl Explainable for BinaryMerge { "BinaryMerge" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - extractor.check_exhausted()?; + fn from_args(_args: &mut ArgsAccess<'_>) -> Result { Ok(BinaryMerge {}) } @@ -634,14 +620,14 @@ impl Explainable for TupleSortHint { "TupleSortHint" } - fn from_args(args: &ExtensionArgs) -> Result { - if args.positional.len() != 1 { + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + if args.positional().len() != 1 { return Err(ExtensionError::InvalidArgument(format!( "expected 1 positional tuple arg, got {}", - args.positional.len() + args.positional().len() ))); } - let tv = <&TupleValue>::try_from(&args.positional[0])?; + let tv = <&TupleValue>::try_from(&args.positional()[0])?; let directions = tv .iter() .map(|v| { @@ -649,8 +635,6 @@ impl Explainable for TupleSortHint { Ok(s) }) .collect::, ExtensionError>>()?; - let mut extractor = args.extractor(); - extractor.check_exhausted()?; Ok(TupleSortHint { directions }) } @@ -708,10 +692,14 @@ Root[result] #[test] fn test_tuple_sort_hint_from_args_rejects_non_tuple() { + let mut registry = ExtensionRegistry::new(); + registry + .register_enhancement::() + .expect("register_enhancement"); let mut args = ExtensionArgs::default(); args.positional .push(ExtensionValue::Enum("ASC".to_string())); - let result = TupleSortHint::from_args(&args); + let result = registry.parse_enhancement("TupleSortHint", &args); assert!( result.is_err(), "expected error when positional arg is not a tuple" diff --git a/tests/extension_table.rs b/tests/extension_table.rs index ff135752..8ed4c54f 100644 --- a/tests/extension_table.rs +++ b/tests/extension_table.rs @@ -4,7 +4,8 @@ mod common; use prost::{Message, Name}; use substrait_explain::extensions::{ - Explainable, ExtensionArgs, ExtensionContext, ExtensionError, ExtensionRegistry, examples, + ArgsAccess, Explainable, ExtensionArgs, ExtensionContext, ExtensionError, ExtensionRegistry, + examples, }; use substrait_explain::{Parser, format_with_registry}; @@ -32,11 +33,8 @@ impl Explainable for UserTable { "UserTable" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut extractor = args.extractor(); - let name: &str = extractor.expect_named("name")?; - - extractor.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let name: &str = args.expect_named("name")?; Ok(Self { name: name.to_string(), diff --git a/tests/json_parsing.rs b/tests/json_parsing.rs index 72763e06..55da726d 100644 --- a/tests/json_parsing.rs +++ b/tests/json_parsing.rs @@ -10,7 +10,7 @@ use std::io::Cursor; use substrait::proto; use substrait_explain::cli::{Cli, Commands, Format}; use substrait_explain::extensions::{ - Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, + ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError, ExtensionRegistry, }; use substrait_explain::json::{build_descriptor_pool, parse_json}; @@ -38,11 +38,9 @@ impl Explainable for ParquetScanConfig { "ParquetScan" } - fn from_args(args: &ExtensionArgs) -> Result { - let mut x = args.extractor(); - let path: &str = x.expect_named("path")?; - let batch_size: i64 = x.get_named("batch_size")?.unwrap_or(1024); - x.check_exhausted()?; + fn from_args(args: &mut ArgsAccess<'_>) -> Result { + let path: &str = args.expect_named("path")?; + let batch_size: i64 = args.get_named("batch_size")?.unwrap_or(1024); Ok(Self { path: path.to_string(), batch_size,