Add datafusion cli for iceberg

Signed-off-by: Ray Liu <[email protected]>

Author: liurenjie1024

Cargo.lock

@@ -59,8 +59,11 @@ bitvec = "1.0.1"
 bytes = "1.6"
 chrono = "0.4.38"
 ctor = "0.2.8"
+clap = { version = "4.5.32", features = ["derive", "cargo"] }
 datafusion = "45"
+datafusion-cli = "45"
 derive_builder = "0.20"
+dirs = "6.0.0"
 env_logger = "0.11.0"
 fnv = "1.0.7"
 futures = "0.3"
@@ -95,6 +98,7 @@ tempfile = "3.18"
 thrift = "0.17.0"  
 tokio = { version = "1.36", default-features = false }
 typed-builder = "0.20"
+toml = "0.8.20"
 url = "2.5.4"
 uuid = { version = "1.13.1", features = ["v7"] }
 volo-thrift = "0.10"

Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "iceberg-cli"
+description = "Apache iceberg client"
+version.workspace = true
+edition.workspace = true
+homepage.workspace = true
+repository.workspace = true
+license.workspace = true
+rust-version.workspace = true
+readme = "README.md"
+
+[dependencies]
+clap = {workspace = true}
+datafusion-cli = {workspace = true}
+datafusion = {workspace = true}
+tokio = {workspace = true}
+env_logger = {workspace = true}
+anyhow = {workspace = true}
+iceberg-datafusion = {workspace = true}
+toml = {workspace = true}
+iceberg-catalog-rest = {workspace = true}
+log = {workspace = true}
+dirs = {workspace = true}

crates/integrations/cli/Cargo.toml

@@ -0,0 +1,7 @@
+# Introduction
+
+Iceberg CLI (`iceberg-cli`) is a small command line utility that runs SQL queries against tables,
+which is backed by the DataFusion engine.
+
+# Configuration
+

crates/integrations/cli/README.md

@@ -0,0 +1,130 @@
+use std::any::Any;
+use std::collections::HashMap;
+use std::fs::read_to_string;
+use std::path::Path;
+use std::sync::Arc;
+use anyhow::anyhow;
+use datafusion::catalog::{CatalogProvider, CatalogProviderList};
+use iceberg_datafusion::IcebergCatalogProvider;
+use toml::{Table as TomlTable, Value};
+use iceberg_catalog_rest::{RestCatalog, RestCatalogConfig};
+use log::error;
+
+const CONFIG_NAME_CATALOGS: &str = "catalogs";
+
+#[derive(Debug)]
+pub struct IcebergCatalogList {
+    catalogs: HashMap<String, Arc<IcebergCatalogProvider>>,
+}
+
+impl IcebergCatalogList {
+    pub async fn parse(path: &Path) -> anyhow::Result<Self> {
+        let toml_table: TomlTable =  toml::from_str(&read_to_string(path)?)?;
+        Self::parse_table(&toml_table).await
+    }
+
+    pub async fn parse_table(configs: &TomlTable) -> anyhow::Result<Self> {
+        if let Value::Array(catalogs_config) =
+            configs.get(CONFIG_NAME_CATALOGS).ok_or_else(|| {
+                anyhow::Error::msg(format!("{CONFIG_NAME_CATALOGS} entry not found in config"))
+            })?
+        {
+            let mut catalogs = HashMap::with_capacity(catalogs_config.len());
+            for config in catalogs_config {
+                if let Value::Table(table_config) = config {
+                    let (name, catalog_provider) = IcebergCatalogList::parse_one(table_config)
+                        .await?;
+                    catalogs.insert(name, catalog_provider);
+                } else {
+                    return Err(anyhow!("{CONFIG_NAME_CATALOGS} entry must be a table"));
+                }
+            }
+            Ok(Self { catalogs })
+        } else {
+            Err(anyhow!("{CONFIG_NAME_CATALOGS} must be an array of table!"))
+        }
+    }
+
+    async fn parse_one(config: &TomlTable)
+        -> anyhow::Result<(String, Arc<IcebergCatalogProvider>)> {
+        let name = config
+            .get("name")
+            .ok_or_else(|| anyhow::anyhow!("name not found for catalog"))?
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("name is not string"))?;
+
+        let typ = config
+            .get("type")
+            .ok_or_else(|| anyhow::anyhow!("type not found for catalog"))?
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("type is not string"))?;
+
+        if typ != "rest" {
+            return Err(anyhow::anyhow!("Only rest catalog is supported for now!"));
+        }
+
+        let catalog_config = config
+            .get("config")
+            .ok_or_else(|| anyhow::anyhow!("config not found for catalog {name}"))?
+            .as_table()
+            .ok_or_else(|| anyhow::anyhow!("config is not table for catalog {name}"))?;
+
+        let uri = catalog_config
+            .get("uri")
+            .ok_or_else(|| anyhow::anyhow!("uri not found for catalog {name}"))?
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("uri is not string"))?;
+
+        let warehouse = catalog_config
+            .get("warehouse")
+            .ok_or_else(|| anyhow::anyhow!("warehouse not found for catalog {name}"))?
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("warehouse is not string for catalog {name}"))?;
+
+        let props_table = catalog_config
+            .get("props")
+            .ok_or_else(|| anyhow::anyhow!("props not found for catalog {name}"))?
+            .as_table()
+            .ok_or_else(|| anyhow::anyhow!("props is not table for catalog {name}"))?;
+
+        let mut props = HashMap::with_capacity(props_table.len());
+        for (key, value) in props_table {
+            let value_str = value.as_str()
+                .ok_or_else(|| anyhow::anyhow!("props {key} is not string"))?;
+            props.insert(key.to_string(), value_str.to_string());
+        }
+
+        let rest_catalog_config = RestCatalogConfig::builder()
+            .uri(uri.to_string())
+            .warehouse(warehouse.to_string())
+            .props(props)
+            .build();
+
+        Ok((name.to_string(), Arc::new(IcebergCatalogProvider::try_new(Arc::new(
+            RestCatalog::new(rest_catalog_config))).await?)))
+    }
+}
+
+impl CatalogProviderList for IcebergCatalogList {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn register_catalog(
+        &self,
+        _name: String,
+        _catalog: Arc<dyn CatalogProvider>,
+    ) -> Option<Arc<dyn CatalogProvider>> {
+        error!("Registering catalog is not supported yet");
+        None
+    }
+
+    fn catalog_names(&self) -> Vec<String> {
+        self.catalogs.keys().cloned().collect()
+    }
+
+    fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>> {
+        self.catalogs.get(name)
+            .map(|c| c.clone() as Arc<dyn CatalogProvider>)
+    }
+}

crates/integrations/cli/src/catalog.rs

@@ -0,0 +1,6 @@
+
+#![doc = include_str!("../README.md")]
+pub const ICEBERG_CLI_VERSION: &str = env!("CARGO_PKG_VERSION");
+
+mod catalog;
+pub use catalog::*;

crates/integrations/cli/src/lib.rs

@@ -0,0 +1,93 @@
+use std::path::PathBuf;
+use std::process::ExitCode;
+use std::str::FromStr;
+use std::sync::Arc;
+
+use clap::Parser;
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_cli::print_format::PrintFormat;
+use datafusion_cli::print_options::{MaxRows, PrintOptions};
+use datafusion_cli::exec;
+use iceberg_cli::{IcebergCatalogList, ICEBERG_CLI_VERSION};
+
+#[derive(Debug, Parser, PartialEq)]
+#[clap(author, version, about, long_about= None)]
+struct Args {
+    #[clap(
+        short = 'r',
+        long,
+        help = "Parse catalog config instead of using ~/.icebergrc",
+    )]
+    rc: Option<String>,
+
+    #[clap(long, value_enum, default_value_t = PrintFormat::Automatic)]
+    format: PrintFormat,
+
+    #[clap(
+        short,
+        long,
+        help = "Reduce printing other than the results and work quietly"
+    )]
+    quiet: bool,
+
+    #[clap(
+        long,
+        help = "The max number of rows to display for 'Table' format\n[possible values: numbers(0/10/...), inf(no limit)]",
+        default_value = "40"
+    )]
+    maxrows: MaxRows,
+
+    #[clap(long, help = "Enables console syntax highlighting")]
+    color: bool,
+}
+
+#[tokio::main]
+/// Calls [`main_inner`], then handles printing errors and returning the correct exit code
+pub async fn main() -> ExitCode {
+    if let Err(e) = main_inner().await {
+        println!("Error: {e}");
+        return ExitCode::FAILURE;
+    }
+
+    ExitCode::SUCCESS
+}
+
+async fn main_inner() -> anyhow::Result<()> {
+    env_logger::init();
+    let args = Args::parse();
+
+    if !args.quiet {
+        println!("DataFusion CLI v{}", ICEBERG_CLI_VERSION);
+    }
+
+    let session_config = SessionConfig::from_env()?.with_information_schema(true);
+
+    let mut rt_builder = RuntimeEnvBuilder::new();
+
+    let runtime_env = rt_builder.build_arc()?;
+
+    // enable dynamic file query
+    let ctx = SessionContext::new_with_config_rt(session_config, runtime_env).enable_url_table();
+    ctx.refresh_catalogs().await?;
+
+    let mut print_options = PrintOptions {
+        format: args.format,
+        quiet: args.quiet,
+        maxrows: args.maxrows,
+        color: args.color,
+    };
+
+    let rc = match args.rc {
+        Some(file) => PathBuf::from_str(&file)?,
+        None => dirs::home_dir()
+            .map(|h| h.join(".icebergrc"))
+            .ok_or_else(|| anyhow::anyhow!("cannot find home directory"))?,
+    };
+
+    let catalogs = Arc::new(IcebergCatalogList::parse(&rc).await?);
+    ctx.register_catalog_list(catalogs);
+
+    Ok(exec::exec_from_repl(&ctx, &mut print_options)
+        .await?)
+}

crates/integrations/cli/src/main.rs

@@ -18,5 +18,5 @@
 # iceberg-rust use unstable rust to run linters, such as `clippy` and `rustfmt`. But this will not affect downstream users, 
 # and only MSRV is required.
 [toolchain]
-channel = "nightly-2024-06-10"
+channel = "nightly-2025-03-10"
 components = ["rustfmt", "clippy"]