Generic FlightTableFactory with a default FlightSqlDriver

apache · Aug 14, 2024 · 25b9903 · 25b9903
1 parent 69c99a7
commit 25b9903
Show file tree

Hide file tree

Showing 16 changed files with 1,907 additions and 10 deletions.
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git 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
@@ -42,6 +42,7 @@ use datafusion::physical_plan::{collect, execute_stream, ExecutionPlanProperties
 use datafusion::sql::parser::{DFParser, Statement};
 use datafusion::sql::sqlparser::dialect::dialect_from_str;
 
+use datafusion::datasource::flight::config::FlightOptions;
 use datafusion::sql::sqlparser;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
@@ -386,6 +387,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/Cargo.toml b/datafusion/core/Cargo.toml
@@ -59,10 +59,12 @@ default = [
     "unicode_expressions",
     "compression",
     "parquet",
+    "flight",
 ]
 encoding_expressions = ["datafusion-functions/encoding_expressions"]
 # Used for testing ONLY: causes all values to hash to the same value (test for collisions)
 force_hash_collisions = ["datafusion-physical-plan/force_hash_collisions", "datafusion-common/force_hash_collisions"]
+flight = ["dep:arrow-flight", "dep:base64", "dep:tonic"]
 math_expressions = ["datafusion-functions/math_expressions"]
 parquet = ["datafusion-common/parquet", "dep:parquet"]
 pyarrow = ["datafusion-common/pyarrow", "parquet"]
@@ -83,6 +85,7 @@ ahash = { workspace = true }
 apache-avro = { version = "0.16", optional = true }
 arrow = { workspace = true }
 arrow-array = { workspace = true }
+arrow-flight = { workspace = true, optional = true }
 arrow-ipc = { workspace = true }
 arrow-schema = { workspace = true }
 async-compression = { version = "0.4.0", features = [
@@ -94,6 +97,7 @@ async-compression = { version = "0.4.0", features = [
     "tokio",
 ], optional = true }
 async-trait = { workspace = true }
+base64 = { version = "0.22", optional = true }
 bytes = { workspace = true }
 bzip2 = { version = "0.4.3", optional = true }
 chrono = { workspace = true }
@@ -128,11 +132,13 @@ parking_lot = { workspace = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = "1.0.15"
 pin-project-lite = "^0.2.7"
+prost = "0.12"
 rand = { workspace = true }
 sqlparser = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true }
 tokio-util = { version = "0.7.4", features = ["io"], optional = true }
+tonic = { version = "0.11", optional = true }
 url = { workspace = true }
 uuid = { version = "1.7", features = ["v4"] }
 xz2 = { version = "0.1", optional = true, features = ["static"] }
@@ -161,6 +167,7 @@ test-utils = { path = "../../test-utils" }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot", "fs"] }
 tokio-postgres = "0.7.7"
+tokio-stream = { version = "0.1.15", features = ["net"] }
 
 [target.'cfg(not(target_os = "windows"))'.dev-dependencies]
 nix = { version = "0.29.0", features = ["fs"] }