substrait_explain/
cli.rs

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

substrait_explain/
cli.rs
