substrait_explain/
cli.rs

1use std::fs;
2use std::io::{self, Read, Write};
3use std::process::ExitCode;
4
5use anyhow::{Context, Result};
6use clap::{Parser, Subcommand};
7use prost::Message;
8
9use crate::extensions::ExtensionRegistry;
10use crate::{FormatError, OutputOptions, Visibility, format_with_registry, parse_with_registry};
11
12/// The outcome of a CLI operation.
13///
14/// Distinguishes between complete success and "soft failures" like formatting
15/// issues where output was still written but there were problems.
16#[derive(Debug)]
17pub enum Outcome {
18    /// Operation completed successfully with no issues.
19    Success,
20    /// Output was written, but there were formatting issues.
21    HadFormattingIssues(Vec<FormatError>),
22}
23
24#[derive(Parser)]
25#[command(name = "substrait-explain")]
26#[command(about = "A CLI for parsing and formatting Substrait query plans")]
27#[command(version)]
28pub struct Cli {
29    #[command(subcommand)]
30    pub command: Commands,
31}
32
33impl Cli {
34    /// Run the CLI and return an exit code.
35    ///
36    /// Errors are printed to stderr.
37    pub fn run(self) -> ExitCode {
38        self.run_with_extensions(ExtensionRegistry::default())
39    }
40
41    /// Run the CLI with a custom extension registry and return an exit code.
42    ///
43    /// Use this when embedding the CLI in a binary that registers custom
44    /// extension relation types:
45    ///
46    /// ```rust,ignore
47    /// let mut registry = ExtensionRegistry::new();
48    /// registry.register_relation::<MyCustomScan>().unwrap();
49    /// Cli::parse().run_with_extensions(registry)
50    /// ```
51    pub fn run_with_extensions(self, registry: ExtensionRegistry) -> ExitCode {
52        match self.run_inner(&registry) {
53            Ok(Outcome::Success) => ExitCode::SUCCESS,
54            Ok(Outcome::HadFormattingIssues(errors)) => {
55                eprintln!("Formatting issues:");
56                for error in errors {
57                    eprintln!("  {error}");
58                }
59                ExitCode::FAILURE
60            }
61            Err(e) => {
62                eprintln!("Error: {e:?}");
63                ExitCode::FAILURE
64            }
65        }
66    }
67
68    fn run_inner(self, registry: &ExtensionRegistry) -> Result<Outcome> {
69        match &self.command {
70            Commands::Convert {
71                input,
72                output,
73                from,
74                to,
75                show_literal_types,
76                verbose,
77            } => {
78                let reader = get_reader(input)
79                    .with_context(|| format!("Failed to open input file: {input}"))?;
80                let writer = get_writer(output)
81                    .with_context(|| format!("Failed to create output file: {output}"))?;
82                let options = self.create_output_options(*show_literal_types);
83                let from_format = self.resolve_input_format(from, input)?;
84                let to_format = self.resolve_output_format(to, output)?;
85                self.run_convert_with_io(
86                    reader,
87                    writer,
88                    &from_format,
89                    &to_format,
90                    &options,
91                    *verbose,
92                    registry,
93                )
94            }
95
96            Commands::Validate {
97                input,
98                output,
99                verbose,
100            } => {
101                let reader = get_reader(input)
102                    .with_context(|| format!("Failed to open input file: {input}"))?;
103                let writer = get_writer(output)
104                    .with_context(|| format!("Failed to create output file: {output}"))?;
105                self.run_validate_with_io(reader, writer, *verbose, registry)
106            }
107        }
108    }
109
110    /// Run CLI with provided readers and writers for testing
111    pub fn run_with_io<R: Read, W: Write>(
112        &self,
113        reader: R,
114        writer: W,
115        registry: &ExtensionRegistry,
116    ) -> Result<Outcome> {
117        match &self.command {
118            Commands::Convert {
119                input,
120                output,
121                from,
122                to,
123                show_literal_types,
124                verbose,
125                ..
126            } => {
127                let options = self.create_output_options(*show_literal_types);
128                let from_format = self.resolve_input_format(from, input)?;
129                let to_format = self.resolve_output_format(to, output)?;
130                self.run_convert_with_io(
131                    reader,
132                    writer,
133                    &from_format,
134                    &to_format,
135                    &options,
136                    *verbose,
137                    registry,
138                )
139            }
140
141            Commands::Validate { verbose, .. } => {
142                self.run_validate_with_io(reader, writer, *verbose, registry)
143            }
144        }
145    }
146
147    fn create_output_options(&self, show_literal_types: bool) -> OutputOptions {
148        let mut options = OutputOptions::default();
149
150        if show_literal_types {
151            options.literal_types = Visibility::Always;
152        }
153
154        options
155    }
156
157    fn resolve_input_format(&self, format: &Option<Format>, input_path: &str) -> Result<Format> {
158        match format {
159            Some(fmt) => Ok(fmt.clone()),
160            None => Format::from_extension(input_path).ok_or_else(|| {
161                anyhow::anyhow!(
162                    "Could not auto-detect input format from file extension. \
163                     Please specify format explicitly with -f/--from. \
164                     Supported formats: text, json, yaml, protobuf/proto/pb"
165                )
166            }),
167        }
168    }
169
170    fn resolve_output_format(&self, format: &Option<Format>, output_path: &str) -> Result<Format> {
171        match format {
172            Some(fmt) => Ok(fmt.clone()),
173            None => Format::from_extension(output_path).ok_or_else(|| {
174                anyhow::anyhow!(
175                    "Could not auto-detect output format from file extension. \
176                     Please specify format explicitly with -t/--to. \
177                     Supported formats: text, json, yaml, protobuf/proto/pb"
178                )
179            }),
180        }
181    }
182
183    // TODO: this could use a refactor; the too_many_arguments tells us
184    // something useful here. We could perhaps add a type containing (registry,
185    // formats, options) or something
186    #[allow(clippy::too_many_arguments)]
187    fn run_convert_with_io<R: Read, W: Write>(
188        &self,
189        reader: R,
190        writer: W,
191        from: &Format,
192        to: &Format,
193        options: &OutputOptions,
194        verbose: bool,
195        registry: &ExtensionRegistry,
196    ) -> Result<Outcome> {
197        // Read input based on format
198        let plan = from.read_plan(reader, registry).with_context(|| {
199            format!(
200                "Failed to parse input as {} format",
201                format!("{from:?}").to_lowercase()
202            )
203        })?;
204
205        // Write output based on format
206        let outcome = to
207            .write_plan(writer, &plan, options, registry)
208            .with_context(|| {
209                format!(
210                    "Failed to write output as {} format",
211                    format!("{to:?}").to_lowercase()
212                )
213            })?;
214
215        if verbose && matches!(outcome, Outcome::Success) {
216            eprintln!("Successfully converted from {from:?} to {to:?}");
217        }
218
219        Ok(outcome)
220    }
221
222    fn run_validate_with_io<R: Read, W: Write>(
223        &self,
224        reader: R,
225        writer: W,
226        verbose: bool,
227        registry: &ExtensionRegistry,
228    ) -> Result<Outcome> {
229        let plan = Format::Text
230            .read_plan(reader, registry)
231            .with_context(|| "Failed to parse input as Substrait text format")?;
232
233        let outcome = Format::Text
234            .write_plan(writer, &plan, &OutputOptions::default(), registry)
235            .with_context(|| "Failed to format plan as Substrait text format")?;
236
237        if verbose && matches!(outcome, Outcome::Success) {
238            eprintln!("Successfully validated plan");
239        }
240
241        Ok(outcome)
242    }
243}
244
245#[derive(Subcommand)]
246pub enum Commands {
247    /// Convert between different Substrait plan formats
248    ///
249    /// Format auto-detection:
250    ///   If -f/--from or -t/--to are not specified, formats will be auto-detected
251    ///   from file extensions:
252    ///     .substrait, .txt    -> text format
253    ///     .json               -> json format
254    ///     .yaml, .yml         -> yaml format
255    ///     .pb, .proto, .protobuf -> protobuf format
256    ///
257    /// Plan formats:
258    ///   text     - Human-readable Substrait text format
259    ///   json     - JSON serialized protobuf
260    ///   yaml     - YAML serialized protobuf
261    ///   protobuf - Binary protobuf format
262    Convert {
263        /// Input file (use - for stdin)
264        #[arg(short, long, default_value = "-")]
265        input: String,
266        /// Output file (use - for stdout)
267        #[arg(short, long, default_value = "-")]
268        output: String,
269        /// Input format: text, json, yaml, protobuf/proto/pb (auto-detected from file extension if not specified)
270        #[arg(short = 'f', long)]
271        from: Option<Format>,
272        /// Output format: text, json, yaml, protobuf/proto/pb (auto-detected from file extension if not specified)
273        #[arg(short = 't', long)]
274        to: Option<Format>,
275        /// Show literal types (text output only)
276        #[arg(long)]
277        show_literal_types: bool,
278        /// Verbose output
279        #[arg(short, long)]
280        verbose: bool,
281    },
282    /// Validate text format by parsing and formatting (roundtrip test)
283    Validate {
284        /// Input file (use - for stdin)
285        #[arg(short, long, default_value = "-")]
286        input: String,
287        /// Output file (use - for stdout)
288        #[arg(short, long, default_value = "-")]
289        output: String,
290        /// Verbose output
291        #[arg(short, long)]
292        verbose: bool,
293    },
294}
295
296#[derive(Clone, Debug, PartialEq)]
297pub enum Format {
298    Text,
299    Json,
300    Yaml,
301    Protobuf,
302}
303
304impl std::str::FromStr for Format {
305    type Err = String;
306
307    fn from_str(s: &str) -> Result<Self, Self::Err> {
308        match s.to_lowercase().as_str() {
309            "text" => Ok(Format::Text),
310            "json" => Ok(Format::Json),
311            "yaml" => Ok(Format::Yaml),
312            "protobuf" | "proto" | "pb" => Ok(Format::Protobuf),
313            _ => Err(format!(
314                "Invalid format: '{s}'. Supported formats: text, json, yaml, protobuf/proto/pb"
315            )),
316        }
317    }
318}
319
320impl Format {
321    /// Detect format from file extension
322    pub fn from_extension(path: &str) -> Option<Format> {
323        if path == "-" {
324            return None; // stdin/stdout - no extension
325        }
326
327        let extension = std::path::Path::new(path)
328            .extension()
329            .and_then(|ext| ext.to_str())
330            .map(|ext| ext.to_lowercase());
331
332        match extension.as_deref() {
333            Some("substrait") | Some("txt") => Some(Format::Text),
334            Some("json") => Some(Format::Json),
335            Some("yaml") | Some("yml") => Some(Format::Yaml),
336            Some("pb") | Some("proto") | Some("protobuf") => Some(Format::Protobuf),
337            _ => None,
338        }
339    }
340
341    pub fn read_plan<R: Read>(
342        &self,
343        reader: R,
344        registry: &ExtensionRegistry,
345    ) -> Result<substrait::proto::Plan> {
346        match self {
347            Format::Text => {
348                let input_text = read_text_input(reader)?;
349                Ok(parse_with_registry(&input_text, registry)?)
350            }
351            Format::Json => {
352                let input_text = read_text_input(reader)?;
353                let pool = crate::json::build_descriptor_pool(&registry.descriptors())?;
354                crate::json::parse_json(&input_text, &pool)
355            }
356            Format::Yaml => {
357                #[cfg(feature = "serde")]
358                {
359                    let input_text = read_text_input(reader)?;
360                    Ok(serde_yaml::from_str(&input_text)?)
361                }
362                #[cfg(not(feature = "serde"))]
363                {
364                    Err("YAML support requires the 'serde' feature. Install with: cargo install substrait-explain --features cli,serde".into())
365                }
366            }
367            Format::Protobuf => {
368                let input_bytes = read_binary_input(reader)?;
369                Ok(substrait::proto::Plan::decode(&input_bytes[..])?)
370            }
371        }
372    }
373
374    pub fn write_plan<W: Write>(
375        &self,
376        writer: W,
377        plan: &substrait::proto::Plan,
378        options: &OutputOptions,
379        registry: &ExtensionRegistry,
380    ) -> Result<Outcome> {
381        match self {
382            Format::Text => {
383                let (text, errors) = format_with_registry(plan, options, registry);
384
385                // Write output first (best-effort)
386                write_text_output(writer, &text)?;
387
388                // Return outcome based on whether there were formatting issues
389                if errors.is_empty() {
390                    Ok(Outcome::Success)
391                } else {
392                    Ok(Outcome::HadFormattingIssues(errors))
393                }
394            }
395            Format::Json => {
396                #[cfg(feature = "serde")]
397                {
398                    let json = serde_json::to_string_pretty(plan)?;
399                    write_text_output(writer, &json)?;
400                    Ok(Outcome::Success)
401                }
402                #[cfg(not(feature = "serde"))]
403                {
404                    Err("JSON support requires the 'serde' feature. Install with: cargo install substrait-explain --features cli,serde".into())
405                }
406            }
407            Format::Yaml => {
408                #[cfg(feature = "serde")]
409                {
410                    let yaml = serde_yaml::to_string(plan)?;
411                    write_text_output(writer, &yaml)?;
412                    Ok(Outcome::Success)
413                }
414                #[cfg(not(feature = "serde"))]
415                {
416                    Err("YAML support requires the 'serde' feature. Install with: cargo install substrait-explain --features cli,serde".into())
417                }
418            }
419            Format::Protobuf => {
420                let bytes = plan.encode_to_vec();
421                write_binary_output(writer, &bytes)?;
422                Ok(Outcome::Success)
423            }
424        }
425    }
426}
427
428/// Read text input from reader
429fn read_text_input<R: Read>(mut reader: R) -> Result<String> {
430    let mut buffer = String::new();
431    reader.read_to_string(&mut buffer)?;
432    Ok(buffer)
433}
434
435/// Read binary input from reader
436fn read_binary_input<R: Read>(mut reader: R) -> Result<Vec<u8>> {
437    let mut buffer = Vec::new();
438    reader.read_to_end(&mut buffer)?;
439    Ok(buffer)
440}
441
442/// Write text output to writer
443fn write_text_output<W: Write>(mut writer: W, content: &str) -> Result<()> {
444    writer.write_all(content.as_bytes())?;
445    Ok(())
446}
447
448/// Write binary output to writer
449fn write_binary_output<W: Write>(mut writer: W, content: &[u8]) -> Result<()> {
450    writer.write_all(content)?;
451    Ok(())
452}
453
454/// Helper function to get reader from file path (or stdin if "-")
455fn get_reader(path: &str) -> Result<Box<dyn Read>> {
456    if path == "-" {
457        Ok(Box::new(io::stdin()))
458    } else {
459        Ok(Box::new(fs::File::open(path)?))
460    }
461}
462
463/// Helper function to get writer from file path (or stdout if "-")
464fn get_writer(path: &str) -> Result<Box<dyn Write>> {
465    if path == "-" {
466        Ok(Box::new(io::stdout()))
467    } else {
468        Ok(Box::new(fs::File::create(path)?))
469    }
470}
471
472#[cfg(test)]
473mod tests {
474    use std::io::Cursor;
475
476    use substrait::proto::expression::RexType;
477    use substrait::proto::plan_rel;
478    use substrait::proto::rel::RelType;
479
480    use super::*;
481    use crate::extensions::{Explainable, ExtensionArgs, ExtensionColumn, ExtensionError};
482    use crate::fixtures::parse_type;
483    use crate::parse;
484
485    const BASIC_PLAN: &str = r#"=== Plan
486Root[result]
487  Project[$0, $1]
488    Read[data => a:i64, b:string]
489"#;
490
491    const PLAN_WITH_EXTENSIONS: &str = r#"=== Extensions
492URNs:
493  @  1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml
494Functions:
495  # 10 @  1: gt
496
497=== Plan
498Root[result]
499  Filter[gt($2, 100):boolean => $0, $1, $2]
500    Project[$0, $1, $2]
501      Read[data => a:i64, b:string, c:i32]
502"#;
503
504    #[test]
505    fn test_convert_text_to_text() {
506        let input = Cursor::new(BASIC_PLAN);
507        let mut output = Vec::new();
508
509        let cli = Cli {
510            command: Commands::Convert {
511                input: "input.substrait".to_string(),
512                output: "output.substrait".to_string(),
513                from: Some(Format::Text),
514                to: Some(Format::Text),
515                show_literal_types: false,
516                verbose: false,
517            },
518        };
519
520        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
521            .unwrap();
522
523        let output_content = String::from_utf8(output).unwrap();
524        assert!(output_content.contains("=== Plan"));
525        assert!(output_content.contains("Root[result]"));
526        assert!(output_content.contains("Project[$0, $1]"));
527        assert!(output_content.contains("Read[data => a:i64, b:string]"));
528    }
529
530    #[test]
531    fn test_convert_text_to_json() {
532        let input = Cursor::new(BASIC_PLAN);
533        let mut output = Vec::new();
534
535        let cli = Cli {
536            command: Commands::Convert {
537                input: "input.substrait".to_string(),
538                output: "output.json".to_string(),
539                from: Some(Format::Text),
540                to: Some(Format::Json),
541                show_literal_types: false,
542                verbose: false,
543            },
544        };
545
546        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
547            .unwrap();
548
549        let output_content = String::from_utf8(output).unwrap();
550        assert!(output_content.contains("\"relations\""));
551        assert!(output_content.contains("\"root\""));
552        assert!(output_content.contains("\"project\""));
553        assert!(output_content.contains("\"read\""));
554    }
555
556    #[test]
557    fn test_convert_json_to_text() {
558        // First convert text to JSON
559        let input = Cursor::new(BASIC_PLAN);
560        let mut json_output = Vec::new();
561
562        let cli_to_json = Cli {
563            command: Commands::Convert {
564                input: "input.substrait".to_string(),
565                output: "output.json".to_string(),
566                from: Some(Format::Text),
567                to: Some(Format::Json),
568                show_literal_types: false,
569                verbose: false,
570            },
571        };
572
573        cli_to_json
574            .run_with_io(input, &mut json_output, &ExtensionRegistry::default())
575            .unwrap();
576
577        // Now convert JSON back to text
578        let json_input = Cursor::new(json_output);
579        let mut text_output = Vec::new();
580
581        let cli_to_text = Cli {
582            command: Commands::Convert {
583                input: "input.json".to_string(),
584                output: "output.substrait".to_string(),
585                from: Some(Format::Json),
586                to: Some(Format::Text),
587                show_literal_types: false,
588                verbose: false,
589            },
590        };
591
592        cli_to_text
593            .run_with_io(json_input, &mut text_output, &ExtensionRegistry::default())
594            .unwrap();
595
596        let output_content = String::from_utf8(text_output).unwrap();
597        assert!(output_content.contains("=== Plan"));
598        assert!(output_content.contains("Root[result]"));
599    }
600
601    #[test]
602    fn test_convert_with_protobuf_output() {
603        let input = Cursor::new(BASIC_PLAN);
604        let mut output = Vec::new();
605
606        let cli = Cli {
607            command: Commands::Convert {
608                input: "input.substrait".to_string(),
609                output: "output.pb".to_string(),
610                from: Some(Format::Text),
611                to: Some(Format::Protobuf),
612                show_literal_types: false,
613                verbose: false,
614            },
615        };
616
617        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
618            .unwrap();
619
620        // Protobuf output should be binary, so we just check that it's not empty
621        assert!(!output.is_empty());
622
623        // Should not contain readable text
624        let output_string = String::from_utf8_lossy(&output);
625        assert!(!output_string.contains("=== Plan"));
626    }
627
628    #[test]
629    fn test_validate_command() {
630        let input = Cursor::new(BASIC_PLAN);
631        let mut output = Vec::new();
632
633        let cli = Cli {
634            command: Commands::Validate {
635                input: String::new(),
636                output: String::new(),
637                verbose: false,
638            },
639        };
640
641        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
642            .unwrap();
643
644        let output_content = String::from_utf8(output).unwrap();
645        assert!(output_content.contains("=== Plan"));
646        assert!(output_content.contains("Root[result]"));
647        assert!(output_content.contains("Project[$0, $1]"));
648        assert!(output_content.contains("Read[data => a:i64, b:string]"));
649    }
650
651    #[test]
652    fn test_validate_with_extensions() {
653        let input = Cursor::new(PLAN_WITH_EXTENSIONS);
654        let mut output = Vec::new();
655
656        let cli = Cli {
657            command: Commands::Validate {
658                input: String::new(),
659                output: String::new(),
660                verbose: false,
661            },
662        };
663
664        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
665            .unwrap();
666
667        let output_content = String::from_utf8(output).unwrap();
668        assert!(output_content.contains("=== Extensions"));
669        assert!(output_content.contains("=== Plan"));
670        assert!(output_content.contains("Root[result]"));
671        assert!(output_content.contains("Filter[gt($2, 100):boolean"));
672    }
673
674    #[test]
675    fn test_convert_with_formatting_options() {
676        let input = Cursor::new(BASIC_PLAN);
677        let mut output = Vec::new();
678
679        let cli = Cli {
680            command: Commands::Convert {
681                input: "input.substrait".to_string(),
682                output: "output.substrait".to_string(),
683                from: Some(Format::Text),
684                to: Some(Format::Text),
685                show_literal_types: true,
686                verbose: false,
687            },
688        };
689
690        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
691            .unwrap();
692
693        let output_content = String::from_utf8(output).unwrap();
694        assert!(output_content.contains("=== Plan"));
695        assert!(output_content.contains("Root[result]"));
696    }
697
698    #[test]
699    fn test_auto_detect_from_extension() {
700        // Test auto-detection of text format
701        assert_eq!(Format::from_extension("plan.substrait"), Some(Format::Text));
702        assert_eq!(Format::from_extension("plan.txt"), Some(Format::Text));
703
704        // Test auto-detection of JSON format
705        assert_eq!(Format::from_extension("plan.json"), Some(Format::Json));
706
707        // Test auto-detection of YAML format
708        assert_eq!(Format::from_extension("plan.yaml"), Some(Format::Yaml));
709        assert_eq!(Format::from_extension("plan.yml"), Some(Format::Yaml));
710
711        // Test auto-detection of protobuf format
712        assert_eq!(Format::from_extension("plan.pb"), Some(Format::Protobuf));
713        assert_eq!(Format::from_extension("plan.proto"), Some(Format::Protobuf));
714        assert_eq!(
715            Format::from_extension("plan.protobuf"),
716            Some(Format::Protobuf)
717        );
718
719        // Test unknown extensions
720        assert_eq!(Format::from_extension("plan.unknown"), None);
721        assert_eq!(Format::from_extension("plan"), None);
722
723        // Test stdin/stdout
724        assert_eq!(Format::from_extension("-"), None);
725    }
726
727    #[test]
728    fn test_convert_with_auto_detection() {
729        let input = Cursor::new(BASIC_PLAN);
730        let mut output = Vec::new();
731
732        let cli = Cli {
733            command: Commands::Convert {
734                input: "input.substrait".to_string(),
735                output: "output.json".to_string(),
736                from: None, // Auto-detect from extension
737                to: None,   // Auto-detect from extension
738                show_literal_types: false,
739                verbose: false,
740            },
741        };
742
743        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
744            .unwrap();
745
746        let output_content = String::from_utf8(output).unwrap();
747        assert!(output_content.contains("\"relations\""));
748        assert!(output_content.contains("\"root\""));
749        assert!(output_content.contains("\"project\""));
750        assert!(output_content.contains("\"read\""));
751    }
752
753    #[test]
754    fn test_auto_detection_error_unknown_input_extension() {
755        let input = Cursor::new(BASIC_PLAN);
756        let mut output = Vec::new();
757
758        let cli = Cli {
759            command: Commands::Convert {
760                input: "input.unknown".to_string(),
761                output: "output.json".to_string(),
762                from: None, // Should fail auto-detection
763                to: None,
764                show_literal_types: false,
765                verbose: false,
766            },
767        };
768
769        let result = cli.run_with_io(input, &mut output, &ExtensionRegistry::default());
770        assert!(result.is_err());
771        assert!(
772            result
773                .unwrap_err()
774                .to_string()
775                .contains("Could not auto-detect input format")
776        );
777    }
778
779    #[test]
780    fn test_auto_detection_error_unknown_output_extension() {
781        let input = Cursor::new(BASIC_PLAN);
782        let mut output = Vec::new();
783
784        let cli = Cli {
785            command: Commands::Convert {
786                input: "input.substrait".to_string(),
787                output: "output.unknown".to_string(),
788                from: None,
789                to: None, // Should fail auto-detection
790                show_literal_types: false,
791                verbose: false,
792            },
793        };
794
795        let result = cli.run_with_io(input, &mut output, &ExtensionRegistry::default());
796        assert!(result.is_err());
797        assert!(
798            result
799                .unwrap_err()
800                .to_string()
801                .contains("Could not auto-detect output format")
802        );
803    }
804
805    #[test]
806    fn test_explicit_format_overrides_auto_detection() {
807        let input = Cursor::new(BASIC_PLAN);
808        let mut output = Vec::new();
809
810        let cli = Cli {
811            command: Commands::Convert {
812                input: "input.json".to_string(), // Would auto-detect as JSON
813                output: "output.pb".to_string(), // Would auto-detect as Protobuf
814                from: Some(Format::Text),        // Explicit override
815                to: Some(Format::Text),          // Explicit override
816                show_literal_types: false,
817                verbose: false,
818            },
819        };
820
821        cli.run_with_io(input, &mut output, &ExtensionRegistry::default())
822            .unwrap();
823
824        let output_content = String::from_utf8(output).unwrap();
825        assert!(output_content.contains("=== Plan"));
826        assert!(output_content.contains("Root[result]"));
827    }
828
829    #[test]
830    fn test_protobuf_roundtrip() {
831        // Convert text to protobuf
832        let input = Cursor::new(BASIC_PLAN);
833        let mut protobuf_output = Vec::new();
834
835        let cli_to_protobuf = Cli {
836            command: Commands::Convert {
837                input: "input.substrait".to_string(),
838                output: "output.pb".to_string(),
839                from: Some(Format::Text),
840                to: Some(Format::Protobuf),
841                show_literal_types: false,
842                verbose: false,
843            },
844        };
845
846        cli_to_protobuf
847            .run_with_io(input, &mut protobuf_output, &ExtensionRegistry::default())
848            .unwrap();
849
850        // Convert protobuf back to text
851        let protobuf_input = Cursor::new(protobuf_output);
852        let mut text_output = Vec::new();
853
854        let cli_to_text = Cli {
855            command: Commands::Convert {
856                input: "input.pb".to_string(),
857                output: "output.substrait".to_string(),
858                from: Some(Format::Protobuf),
859                to: Some(Format::Text),
860                show_literal_types: false,
861                verbose: false,
862            },
863        };
864
865        cli_to_text
866            .run_with_io(
867                protobuf_input,
868                &mut text_output,
869                &ExtensionRegistry::default(),
870            )
871            .unwrap();
872
873        let output_content = String::from_utf8(text_output).unwrap();
874        assert!(output_content.contains("=== Plan"));
875        assert!(output_content.contains("Root[result]"));
876        assert!(output_content.contains("Read[data => a:i64, b:string]"));
877    }
878
879    // -----------------------------------------------------------------
880    // Minimal test extension for verifying registry-aware CLI parsing
881    // -----------------------------------------------------------------
882
883    /// A minimal ExtensionLeaf with one named argument, used to verify
884    /// that CLI commands pass the registry through to the text parser.
885    #[derive(Clone, PartialEq, prost::Message)]
886    struct TestSource {
887        #[prost(string, tag = "1")]
888        tag: String,
889    }
890
891    impl prost::Name for TestSource {
892        const NAME: &'static str = "TestSource";
893        const PACKAGE: &'static str = "test";
894        fn full_name() -> String {
895            "test.TestSource".to_string()
896        }
897        fn type_url() -> String {
898            "type.googleapis.com/test.TestSource".to_string()
899        }
900    }
901
902    impl Explainable for TestSource {
903        fn name() -> &'static str {
904            "TestSource"
905        }
906
907        fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
908            let mut extractor = args.extractor();
909            let tag: &str = extractor.expect_named_arg("tag")?;
910            extractor.check_exhausted()?;
911            Ok(TestSource {
912                tag: tag.to_string(),
913            })
914        }
915
916        fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
917            let mut args = ExtensionArgs::default();
918            args.insert("tag", self.tag.clone());
919            args.output_columns.push(ExtensionColumn::Named {
920                name: "val".to_string(),
921                r#type: parse_type("i64"),
922            });
923            Ok(args)
924        }
925    }
926
927    fn make_extension_registry() -> ExtensionRegistry {
928        let mut registry = ExtensionRegistry::new();
929        registry.register_relation::<TestSource>().unwrap();
930        registry
931    }
932
933    const PLAN_WITH_CUSTOM_EXTENSION: &str = r#"=== Plan
934Root[val]
935  ExtensionLeaf:TestSource[tag='hello' => val:i64]
936"#;
937
938    #[test]
939    fn test_convert_text_to_text_with_extension_registry() {
940        let registry = make_extension_registry();
941        let input = Cursor::new(PLAN_WITH_CUSTOM_EXTENSION);
942        let mut output = Vec::new();
943
944        let cli = Cli {
945            command: Commands::Convert {
946                input: "input.substrait".to_string(),
947                output: "output.substrait".to_string(),
948                from: Some(Format::Text),
949                to: Some(Format::Text),
950                show_literal_types: false,
951                verbose: false,
952            },
953        };
954
955        cli.run_with_io(input, &mut output, &registry).unwrap();
956
957        let output_content = String::from_utf8(output).unwrap();
958        assert_eq!(output_content, PLAN_WITH_CUSTOM_EXTENSION);
959    }
960
961    #[test]
962    fn test_convert_text_to_json_with_extension_registry() {
963        let registry = make_extension_registry();
964        let input = Cursor::new(PLAN_WITH_CUSTOM_EXTENSION);
965        let mut output = Vec::new();
966
967        let cli = Cli {
968            command: Commands::Convert {
969                input: "input.substrait".to_string(),
970                output: "output.json".to_string(),
971                from: Some(Format::Text),
972                to: Some(Format::Json),
973                show_literal_types: false,
974                verbose: false,
975            },
976        };
977
978        cli.run_with_io(input, &mut output, &registry).unwrap();
979
980        let output_content = String::from_utf8(output).unwrap();
981        assert!(output_content.contains("\"extensionLeaf\""));
982    }
983
984    #[test]
985    fn test_validate_with_extension_registry() {
986        let registry = make_extension_registry();
987        let input = Cursor::new(PLAN_WITH_CUSTOM_EXTENSION);
988        let mut output = Vec::new();
989
990        let cli = Cli {
991            command: Commands::Validate {
992                input: String::new(),
993                output: String::new(),
994                verbose: false,
995            },
996        };
997
998        cli.run_with_io(input, &mut output, &registry).unwrap();
999
1000        let output_content = String::from_utf8(output).unwrap();
1001        assert_eq!(output_content, PLAN_WITH_CUSTOM_EXTENSION);
1002    }
1003
1004    #[test]
1005    fn test_convert_text_fails_without_extension_registry() {
1006        // Without the registry, parsing a plan with custom extensions should fail
1007        let input = Cursor::new(PLAN_WITH_CUSTOM_EXTENSION);
1008        let mut output = Vec::new();
1009
1010        let cli = Cli {
1011            command: Commands::Convert {
1012                input: "input.substrait".to_string(),
1013                output: "output.substrait".to_string(),
1014                from: Some(Format::Text),
1015                to: Some(Format::Text),
1016                show_literal_types: false,
1017                verbose: false,
1018            },
1019        };
1020
1021        let result = cli.run_with_io(input, &mut output, &ExtensionRegistry::default());
1022        assert!(result.is_err());
1023    }
1024
1025    /// Creates a plan with an invalid function reference that will cause formatting errors.
1026    fn make_plan_with_invalid_function_ref() -> substrait::proto::Plan {
1027        const VALID_PLAN: &str = r#"=== Extensions
1028URNs:
1029  @  1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml
1030Functions:
1031  # 10 @  1: equal
1032
1033=== Plan
1034Root[result]
1035  Filter[equal($0, 42:i32):boolean => $0]
1036    Read[data => a:i32]
1037"#;
1038
1039        let mut plan = parse(VALID_PLAN).expect("Failed to parse valid plan");
1040
1041        // Navigate to the function and corrupt its reference
1042        let rel_root = plan.relations.first_mut().unwrap();
1043        let plan_rel::RelType::Root(root) = rel_root.rel_type.as_mut().unwrap() else {
1044            panic!("Expected Root relation");
1045        };
1046        let rel = root.input.as_mut().unwrap();
1047        let RelType::Filter(filter) = rel.rel_type.as_mut().unwrap() else {
1048            panic!("Expected Filter relation");
1049        };
1050        let condition = filter.condition.as_mut().unwrap();
1051        let RexType::ScalarFunction(func) = condition.rex_type.as_mut().unwrap() else {
1052            panic!("Expected ScalarFunction");
1053        };
1054        func.function_reference = 999; // Invalid - doesn't exist in extensions
1055
1056        plan
1057    }
1058
1059    #[test]
1060    fn test_write_plan_reports_formatting_issues() {
1061        let plan = make_plan_with_invalid_function_ref();
1062        let mut output = Vec::new();
1063
1064        let result = Format::Text.write_plan(
1065            &mut output,
1066            &plan,
1067            &OutputOptions::default(),
1068            &ExtensionRegistry::default(),
1069        );
1070
1071        // Should succeed but report formatting issues
1072        let outcome = result.expect("write_plan should not return hard error");
1073        assert!(
1074            matches!(outcome, Outcome::HadFormattingIssues(ref errors) if !errors.is_empty()),
1075            "Expected HadFormattingIssues with errors, got {outcome:?}"
1076        );
1077        // Output should still be written (best-effort formatting)
1078        assert!(
1079            !output.is_empty(),
1080            "Output should be written even with issues"
1081        );
1082    }
1083}
substrait_explain/cli.rs

substrait_explain/
cli.rs
