From 820021bee0bb94f76445fc8fe6a8c11fe046b3f0 Mon Sep 17 00:00:00 2001
From: Costi Ciudatu <ciudatu@adobe.com>
Date: Mon, 12 Aug 2024 19:17:04 +0300
Subject: [PATCH] [datafusion-cli] support for flight tables

---
 datafusion-cli/Cargo.toml                     |  1 +
 datafusion-cli/src/exec.rs                    |  3 +
 .../core/src/datasource/flight/config.rs      | 62 +++++++++++++++++++
 datafusion/core/src/datasource/flight/mod.rs  |  1 +
 4 files changed, 67 insertions(+)
 create mode 100644 datafusion/core/src/datasource/flight/config.rs

diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml
index cbd9ffd0febab..d078868506d1c 100644
--- a/datafusion-cli/Cargo.toml
+++ b/datafusion-cli/Cargo.toml
@@ -44,6 +44,7 @@ datafusion = { path = "../datafusion/core", version = "41.0.0", features = [
     "regex_expressions",
     "unicode_expressions",
     "compression",
+    "flight",
 ] }
 dirs = "4.0.0"
 env_logger = "0.9"
diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs
index 178bce6f2fe65..e0eea32c95565 100644
--- a/datafusion-cli/src/exec.rs
+++ b/datafusion-cli/src/exec.rs
@@ -46,6 +46,7 @@ use datafusion::sql::sqlparser;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
 use tokio::signal;
+use datafusion::datasource::flight::config::FlightOptions;
 
 /// run and execute SQL statements and commands, against a context with the given print options
 pub async fn exec_from_commands(
@@ -383,6 +384,8 @@ pub(crate) async fn register_object_store_and_config_extensions(
     let mut table_options = ctx.session_state().default_table_options().clone();
     if let Some(format) = format {
         table_options.set_config_format(format);
+    } else {
+        table_options.extensions.insert(FlightOptions::default())
     }
     table_options.alter_with_string_hash_map(options)?;
 
diff --git a/datafusion/core/src/datasource/flight/config.rs b/datafusion/core/src/datasource/flight/config.rs
new file mode 100644
index 0000000000000..d2f63a74e2e76
--- /dev/null
+++ b/datafusion/core/src/datasource/flight/config.rs
@@ -0,0 +1,62 @@
+// 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.
+
+//! Only meant for datafusion-cli to accept the `flight` namespace
+
+use datafusion_common::config::{ConfigEntry, ConfigExtension, ExtensionOptions};
+use std::any::Any;
+use std::collections::HashMap;
+
+/// Collects the config entries
+#[derive(Default, Debug, Clone)]
+pub struct FlightOptions {
+    inner: HashMap<String, String>,
+}
+
+impl ConfigExtension for FlightOptions {
+    const PREFIX: &'static str = "flight";
+}
+
+impl ExtensionOptions for FlightOptions {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+
+    fn cloned(&self) -> Box<dyn ExtensionOptions> {
+        Box::new(self.clone())
+    }
+
+    fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> {
+        self.inner.insert(key.into(), value.into());
+        Ok(())
+    }
+
+    fn entries(&self) -> Vec<ConfigEntry> {
+        self.inner
+            .iter()
+            .map(|(key, value)| ConfigEntry {
+                key: key.to_owned(),
+                value: Some(value.to_owned()).filter(|s| !s.is_empty()),
+                description: "",
+            })
+            .collect()
+    }
+}
diff --git a/datafusion/core/src/datasource/flight/mod.rs b/datafusion/core/src/datasource/flight/mod.rs
index 05d45c72d5b97..b7c9a5a986f72 100644
--- a/datafusion/core/src/datasource/flight/mod.rs
+++ b/datafusion/core/src/datasource/flight/mod.rs
@@ -38,6 +38,7 @@ use datafusion_physical_plan::{ExecutionMode, ExecutionPlan, PlanProperties};
 
 use crate::datasource::physical_plan::FlightExec;
 
+pub mod config;
 pub mod sql;
 
 /// Generic Arrow Flight data source. Requires a [FlightDriver] that allows implementors