Add datafusion cli for iceberg

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

Author: liurenjie1024

Cargo.toml

@@ -52,24 +52,27 @@ arrow-select = { version = "54.2.0" }
 arrow-string = { version = "54.2.0" }
 async-std = "1.12"
 async-trait = "0.1.86"
-aws-config = "1"
+aws-config = "1.5"
 aws-sdk-glue = "1.39"
 bimap = "0.6"
 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"
 expect-test = "1"
+dirs = "6.0.0"
 fnv = "1.0.7"
 futures = "0.3"
 hive_metastore = "0.1"
 http = "1.1"
-iceberg = { version = "0.4.0", path = "./crates/iceberg" }
-iceberg-catalog-memory = { version = "0.4.0", path = "./crates/catalog/memory" }
-iceberg-catalog-rest = { version = "0.4.0", path = "./crates/catalog/rest" }
-iceberg-datafusion = { version = "0.4.0", path = "./crates/integrations/datafusion" }
+iceberg = { path = "./crates/iceberg" }
+iceberg-catalog-memory = { path = "./crates/catalog/memory" }
+iceberg-catalog-rest = { path = "./crates/catalog/rest" }
+iceberg-datafusion = { path = "./crates/integrations/datafusion" }
 itertools = "0.13"
 mockito = "1"
 murmur3 = "0.5.2"
@@ -99,6 +102,7 @@ tokio = { version = "1.44", default-features = false }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.8"
 typed-builder = "0.20"
+toml = "0.8.20"
 url = "2.5.4"
 uuid = { version = "1.14", features = ["v7"] }
 volo-thrift = "0.10"

crates/integrations/cli/Cargo.toml

@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[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}
+anyhow = {workspace = true}
+iceberg-datafusion = {workspace = true}
+toml = {workspace = true}
+iceberg-catalog-rest = {workspace = true}
+tracing = {workspace = true}
+tracing-subscriber = {workspace = true}
+dirs = {workspace = true}

crates/integrations/cli/README.md

@@ -0,0 +1,25 @@
+<!--
+  ~ Licensed to the Apache Software Foundation (ASF) under one
+  ~ or more contributor license agreements.  See the NOTICE file
+  ~ distributed with this work for additional information
+  ~ regarding copyright ownership.  The ASF licenses this file
+  ~ to you under the Apache License, Version 2.0 (the
+  ~ "License"); you may not use this file except in compliance
+  ~ with the License.  You may obtain a copy of the License at
+  ~
+  ~   http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing,
+  ~ software distributed under the License is distributed on an
+  ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  ~ KIND, either express or implied.  See the License for the
+  ~ specific language governing permissions and limitations
+  ~ under the License.
+-->
+
+
+# Introduction
+
+Iceberg CLI (`iceberg-cli`) is a small command line utility that runs SQL queries against tables,
+which is backed by the DataFusion engine.
+

crates/integrations/cli/src/catalog.rs

@@ -0,0 +1,155 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+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_catalog_rest::{RestCatalog, RestCatalogConfig};
+use iceberg_datafusion::IcebergCatalogProvider;
+use toml::{Table as TomlTable, Value};
+
+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 r#type = 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 r#type != "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>> {
+        tracing::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/lib.rs

@@ -0,0 +1,22 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#![doc = include_str!("../README.md")]
+pub const ICEBERG_CLI_VERSION: &str = env!("CARGO_PKG_VERSION");
+
+mod catalog;
+pub use catalog::*;

crates/integrations/cli/src/main.rs

@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+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::exec;
+use datafusion_cli::print_format::PrintFormat;
+use datafusion_cli::print_options::{MaxRows, PrintOptions};
+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 {
+    tracing_subscriber::fmt::init();
+
+    if let Err(e) = main_inner().await {
+        println!("Error: {e}");
+        return ExitCode::FAILURE;
+    }
+
+    ExitCode::SUCCESS
+}
+
+async fn main_inner() -> anyhow::Result<()> {
+    let args = Args::parse();
+
+    if !args.quiet {
+        println!("ICEBERG CLI v{}", ICEBERG_CLI_VERSION);
+    }
+
+    let session_config = SessionConfig::from_env()?.with_information_schema(true);
+
+    let 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?)
+}