Skip to main content

datafusion_ducklake_cli/
runner.rs

1use std::fs;
2use std::io::{self, IsTerminal, Read, Write};
3use std::path::Path;
4use std::time::Instant;
5
6use datafusion_ducklake_provider::{DuckLakeSessionContext, register_ducklake_object_store};
7
8use crate::args::Cli;
9use crate::error::CliError;
10use crate::object_store::object_store_configs;
11use crate::output::format_dataframe;
12use crate::repl::run_repl;
13use crate::split::split_script;
14
15#[derive(Clone, Copy, Debug)]
16pub(crate) struct RunOptions {
17    pub quiet: bool,
18    pub timing: bool,
19    pub continue_on_error: bool,
20}
21
22pub async fn run(cli: Cli) -> Result<(), CliError> {
23    let ctx = DuckLakeSessionContext::new();
24    for config in object_store_configs(&cli)? {
25        register_ducklake_object_store(ctx.inner(), &config)?;
26    }
27
28    let options = RunOptions {
29        quiet: cli.quiet,
30        timing: cli.timing,
31        continue_on_error: cli.continue_on_error,
32    };
33    let init_options = RunOptions {
34        continue_on_error: false,
35        ..options
36    };
37
38    for init in &cli.init {
39        run_file(&ctx, init, init_options).await?;
40    }
41
42    if !cli.execute.is_empty() {
43        for sql in &cli.execute {
44            run_script_source(&ctx, "command line", sql, options).await?;
45        }
46        return Ok(());
47    }
48
49    if let Some(file) = &cli.file {
50        run_file(&ctx, file, options).await?;
51        return Ok(());
52    }
53
54    if !io::stdin().is_terminal() {
55        let mut sql = String::new();
56        io::stdin().read_to_string(&mut sql)?;
57        run_script_source(&ctx, "stdin", &sql, options).await?;
58        return Ok(());
59    }
60
61    run_repl(&ctx, &cli, options).await
62}
63
64async fn run_file(
65    ctx: &DuckLakeSessionContext,
66    path: &Path,
67    options: RunOptions,
68) -> Result<(), CliError> {
69    let sql = fs::read_to_string(path)?;
70    run_script_source(ctx, &path.display().to_string(), &sql, options).await
71}
72
73pub(crate) async fn run_script_source(
74    ctx: &DuckLakeSessionContext,
75    source: &str,
76    sql: &str,
77    options: RunOptions,
78) -> Result<(), CliError> {
79    for statement in split_script(sql) {
80        if let Err(err) = run_statement(ctx, &statement, options).await {
81            let message = format!("{source}: {err}");
82            if !options.continue_on_error {
83                return Err(CliError::Message(message));
84            }
85            eprintln!("{message}");
86        }
87    }
88    Ok(())
89}
90
91pub(crate) async fn run_statement(
92    ctx: &DuckLakeSessionContext,
93    statement: &str,
94    options: RunOptions,
95) -> Result<(), CliError> {
96    let started = Instant::now();
97    let dataframe = ctx.sql(statement).await?;
98    if let Some(rendered) = format_dataframe(dataframe, options.quiet).await? {
99        io::stdout().write_all(rendered.as_bytes())?;
100    }
101    if options.timing && !options.quiet {
102        eprintln!("Time: {:.3}s", started.elapsed().as_secs_f64());
103    }
104    Ok(())
105}