Move datafusion-cli main_inner impl into the datafusion-cli lib.

datafusion-cli/Cargo.toml

@@ -66,6 +66,7 @@ parking_lot = { workspace = true }
 parquet = { workspace = true, default-features = false }
 regex = { workspace = true }
 rustyline = "18.0"
+thiserror = "2.0.18"
 tokio = { workspace = true, features = ["macros", "parking_lot", "rt", "rt-multi-thread", "signal", "sync"] }
 url = { workspace = true }
 

datafusion-cli/examples/cli-custom-udf.rs

@@ -0,0 +1,59 @@
+// 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 arrow::datatypes::DataType;
+use datafusion::arrow::array::{ArrayRef, StringArray};
+use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf};
+use datafusion::prelude::SessionContext;
+use datafusion_cli::entry_point::{CliError, CliSession};
+use datafusion_common::cast::as_string_array;
+use mimalloc::MiMalloc;
+use std::sync::Arc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+/// In this example we want to reuse the datafusion-cli binary argument, hen extend the `SessionContext` with custom udf.
+///
+/// 1. Declares a `hello`` udf function.
+/// 2. Construct a `CliSession`
+/// 3. Registers the udf function with the `SessionContext` so the user can input `select hello(1)` at the prompt.
+/// 4. Runs the cli using [`dataframe_cli::CliSession::run`], printing any errors then exits.
+#[tokio::main]
+pub async fn main() -> Result<(), CliError> {
+    let custom_udf = create_udf(
+        "hello",
+        vec![DataType::Utf8],
+        DataType::Utf8,
+        Volatility::Immutable,
+        Arc::new(|args: &[ColumnarValue]| {
+            assert_eq!(args.len(), 1);
+            let args = ColumnarValue::values_to_arrays(args).unwrap();
+            let vals = as_string_array(&args[0]).expect("cast failed");
+            let array = vals
+                .iter()
+                .map(|v| v.map(|v| format!("hello {v}")))
+                .collect::<StringArray>();
+            Ok(ColumnarValue::from(Arc::new(array) as ArrayRef))
+        }),
+    );
+    let cli_session = CliSession::try_from_args(std::env::args())?;
+    let ctx: &SessionContext = cli_session.session_context();
+    ctx.register_udf(custom_udf);
+    cli_session.run().await?;
+    Ok(())
+}

datafusion-cli/examples/cli-session-context.rs

@@ -18,6 +18,7 @@
 //! Shows an example of a custom session context that unions the input plan with itself.
 //! To run this example, use `cargo run --example cli-session-context` from within the `datafusion-cli` directory.
 
+use std::env;
 use std::sync::Arc;
 
 use datafusion::{
@@ -28,9 +29,9 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_cli::{
-    cli_context::CliSessionContext, exec::exec_from_repl,
-    object_storage::instrumented::InstrumentedObjectStoreRegistry,
-    print_options::PrintOptions,
+    cli_context::CliSessionContext,
+    entry_point::{CliError, CliSession},
+    exec::exec_from_repl,
 };
 use object_store::ObjectStore;
 
@@ -39,14 +40,6 @@ struct MyUnionerContext {
     ctx: SessionContext,
 }
 
-impl Default for MyUnionerContext {
-    fn default() -> Self {
-        Self {
-            ctx: SessionContext::new(),
-        }
-    }
-}
-
 #[async_trait::async_trait]
 impl CliSessionContext for MyUnionerContext {
     fn task_ctx(&self) -> Arc<TaskContext> {
@@ -83,16 +76,13 @@ impl CliSessionContext for MyUnionerContext {
 
 #[tokio::main]
 /// Runs the example.
-pub async fn main() {
-    let my_ctx = MyUnionerContext::default();
-
-    let mut print_options = PrintOptions {
-        format: datafusion_cli::print_format::PrintFormat::Automatic,
-        quiet: false,
-        maxrows: datafusion_cli::print_options::MaxRows::Unlimited,
-        color: true,
-        instrumented_registry: Arc::new(InstrumentedObjectStoreRegistry::new()),
-    };
-
-    exec_from_repl(&my_ctx, &mut print_options).await.unwrap();
+pub async fn main() -> Result<(), CliError> {
+    let CliSession {
+        ctx,
+        args: _,
+        mut print_options,
+    } = CliSession::try_from_args(env::args())?;
+    let my_ctx = MyUnionerContext { ctx };
+    exec_from_repl(&my_ctx, &mut print_options).await?;
+    Ok(())
 }