Start adding prepared statements

.github/actions/setup-rust/action.yml

@@ -33,6 +33,7 @@ runs:
         rustup install stable
         rustup default stable
         rustup component add clippy
+        rustup component add rustfmt
     - name: Install Cargo tools
       shell: bash
       run: |

crates/datafusion-app/Cargo.toml

@@ -31,11 +31,13 @@ opendal = { version = "0.51", features = [
   "services-huggingface",
 ], optional = true }
 parking_lot = "0.12.3"
+prost = "0.13.5"
 serde = { version = "1.0.197", features = ["derive"] }
+serde_json = { version = "1.0.140", features = ["indexmap"] }
 tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread"] }
 tokio-stream = { version = "0.1.15", features = ["net"] }
 tonic = { version = "0.12.3", optional = true }
-url = { version = "2.5.2", optional = true }
+url = { version = "2.5.2", features = ["serde"], optional = true }
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["async_tokio"] }

crates/datafusion-app/src/catalog/mod.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
+use std::{collections::HashMap, fs::File, path::Path, sync::Arc};
 
 use datafusion::{
     arrow::{
@@ -26,13 +26,18 @@ use datafusion::{
     catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider},
     common::Result,
     datasource::MemTable,
+    error::DataFusionError,
+    scalar::ScalarValue,
     DATAFUSION_VERSION,
 };
+use indexmap::IndexMap;
 
 use crate::config::ExecutionConfig;
 
+type PreparedStatementsMap = IndexMap<String, HashMap<String, ScalarValue>>;
+
 pub fn create_app_catalog(
-    _config: &ExecutionConfig,
+    config: &ExecutionConfig,
     app_name: &str,
     app_version: &str,
 ) -> Result<Arc<dyn CatalogProvider>> {
@@ -41,9 +46,47 @@ pub fn create_app_catalog(
     catalog.register_schema("meta", Arc::<MemorySchemaProvider>::clone(&meta_schema))?;
     let versions_table = try_create_meta_versions_table(app_name, app_version)?;
     meta_schema.register_table("versions".to_string(), versions_table)?;
+    #[cfg(feature = "flightsql")]
+    {
+        let flightsql_schema = Arc::new(MemorySchemaProvider::new());
+        catalog.register_schema("flightsql", flightsql_schema)?;
+        let db_path = config.db.path.to_file_path().map_err(|_| {
+            DataFusionError::External("error converting DB path to file path".to_string().into())
+        })?;
+        let prepared_statements_file = db_path
+            .join(app_name)
+            .join("flightsql")
+            .join("prepared_statements");
+        let prepared_statements = if let Ok(true) = prepared_statements_file.try_exists() {
+            let reader = File::open(prepared_statements_file)
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?;
+            let vals: PreparedStatementsMap = serde_json::from_reader(reader)
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?;
+        } else {
+        };
+    }
     Ok(Arc::new(catalog))
 }
 
+// fn create_flightsql_prepared_statements_table() -> Result<Arc<MemTable>> {
+//     let fields = vec![
+//         Field::new("datafusion", DataType::Utf8, false),
+//         Field::new("datafusion-app", DataType::Utf8, false),
+//     ];
+//     let schema = Arc::new(Schema::new(fields));
+//
+//     let batches = RecordBatch::try_new(
+//         Arc::<Schema>::clone(&schema),
+//         vec![
+//             Arc::new(app_version_arr),
+//             Arc::new(datafusion_version_arr),
+//             Arc::new(datafusion_app_version_arr),
+//         ],
+//     )?;
+//
+//     Ok(Arc::new(MemTable::try_new(schema, vec![vec![batches]])?))
+// }
+
 fn try_create_meta_versions_table(app_name: &str, app_version: &str) -> Result<Arc<MemTable>> {
     let fields = vec![
         Field::new(app_name, DataType::Utf8, false),

crates/datafusion-app/src/config.rs

@@ -24,6 +24,8 @@ use datafusion_udfs_wasm::WasmInputDataType;
 use serde::Deserialize;
 use std::collections::HashMap;
 
+use crate::db::{default_db_config, DbConfig};
+
 #[cfg(feature = "s3")]
 use {
     color_eyre::Result,
@@ -94,6 +96,7 @@ pub struct ExecutionConfig {
     #[cfg(feature = "observability")]
     #[serde(default)]
     pub observability: ObservabilityConfig,
+    pub db: DbConfig,
 }
 
 impl Default for ExecutionConfig {
@@ -111,6 +114,7 @@ impl Default for ExecutionConfig {
             catalog: default_catalog(),
             #[cfg(feature = "observability")]
             observability: default_observability(),
+            db: default_db_config(),
         }
     }
 }

src/db.rs → crates/datafusion-app/src/db.rs

@@ -27,8 +27,39 @@ use datafusion::{
     prelude::SessionContext,
 };
 use log::info;
+use serde::Deserialize;
+use url::Url;
 
-use crate::config::DbConfig;
+#[derive(Debug, Clone, Deserialize)]
+pub struct DbConfig {
+    #[serde(default = "default_db_path")]
+    pub path: Url,
+}
+
+impl Default for DbConfig {
+    fn default() -> Self {
+        default_db_config()
+    }
+}
+
+pub fn default_db_config() -> DbConfig {
+    DbConfig {
+        path: default_db_path(),
+    }
+}
+
+fn default_db_path() -> Url {
+    let base = directories::BaseDirs::new().expect("Base directories should be available");
+    let path = base
+        .data_dir()
+        .to_path_buf()
+        .join("dft/")
+        .to_str()
+        .unwrap()
+        .to_string();
+    let with_schema = format!("file://{path}");
+    Url::parse(&with_schema).unwrap()
+}
 
 pub async fn register_db(ctx: &SessionContext, db_config: &DbConfig) -> Result<()> {
     info!("registering tables to database");

crates/datafusion-app/src/flightsql.rs

@@ -19,30 +19,39 @@ use std::sync::Arc;
 
 use arrow_flight::{
     decode::FlightRecordBatchStream,
-    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetTables},
-    FlightInfo,
+    encode::FlightDataEncoderBuilder,
+    sql::{
+        client::FlightSqlServiceClient, ActionCreatePreparedStatementRequest,
+        ActionCreatePreparedStatementResult, Any, CommandGetDbSchemas, CommandGetTables,
+        CommandPreparedStatementQuery, ProstMessageExt,
+    },
+    Action, FlightDescriptor, FlightInfo, PutResult,
 };
-#[cfg(feature = "flightsql")]
 use base64::engine::{general_purpose::STANDARD, Engine as _};
+use color_eyre::eyre::{self, Result};
 use datafusion::{
+    arrow::array::RecordBatch,
+    common::ParamValues,
     error::{DataFusionError, Result as DFResult},
     physical_plan::stream::RecordBatchStreamAdapter,
     sql::parser::DFParser,
 };
+use futures::{stream, TryStreamExt};
 use log::{debug, error, info, warn};
-
-use color_eyre::eyre::{self, Result};
+use prost::Message;
 use tokio::sync::Mutex;
 use tokio_stream::StreamExt;
 use tonic::{transport::Channel, IntoRequest};
 
-#[cfg(feature = "flightsql")]
 use crate::config::BasicAuth;
 
 use crate::{
     config::FlightSQLConfig, flightsql_benchmarks::FlightSQLBenchmarkStats, ExecOptions, ExecResult,
 };
 
+pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement";
+pub(crate) static CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement";
+
 pub type FlightSQLClient = Arc<Mutex<Option<FlightSqlServiceClient<Channel>>>>;
 
 #[derive(Clone, Debug, Default)]
@@ -275,6 +284,78 @@ impl FlightSQLContext {
         }
     }
 
+    pub async fn create_prepared_statement(
+        &self,
+        query: String,
+    ) -> DFResult<ActionCreatePreparedStatementResult> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            let cmd = ActionCreatePreparedStatementRequest {
+                query,
+                transaction_id: None,
+            };
+            let action = Action {
+                r#type: CREATE_PREPARED_STATEMENT.to_string(),
+                body: cmd.as_any().encode_to_vec().into(),
+            };
+            let mut result = client
+                .inner_mut()
+                .do_action(action)
+                .await
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?
+                .into_inner();
+            let result = result
+                .message()
+                .await
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?
+                .unwrap();
+            let any = Any::decode(&*result.body)
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?;
+            let prepared_result: ActionCreatePreparedStatementResult = any.unpack()?.unwrap();
+            Ok(prepared_result)
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured.  Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
+    pub async fn bind_prepared_statement_params(
+        &self,
+        prepared_id: String,
+        params: RecordBatch,
+    ) -> DFResult<Option<PutResult>> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            let cmd = CommandPreparedStatementQuery {
+                prepared_statement_handle: self.handle.clone(),
+            };
+
+            let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
+            let flight_stream_builder = FlightDataEncoderBuilder::new()
+                .with_flight_descriptor(Some(descriptor))
+                .with_schema(params.schema());
+            let flight_data = flight_stream_builder
+                .build(futures::stream::iter([Ok(params)]))
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| DataFusionError::External(e.to_string().into()))?;
+
+            client
+                .do_put(stream::iter(flight_data))
+                .await?
+                .message()
+                .await
+                .map_err(|e| DataFusionError::External(e.to_string().into()))
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured. Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
     pub async fn do_get(&self, flight_info: FlightInfo) -> DFResult<Vec<FlightRecordBatchStream>> {
         let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;

crates/datafusion-app/src/lib.rs

@@ -17,6 +17,7 @@
 
 pub mod catalog;
 pub mod config;
+pub mod db;
 pub mod executor;
 pub mod extensions;
 #[cfg(feature = "flightsql")]
@@ -27,6 +28,7 @@ pub mod local;
 pub mod local_benchmarks;
 #[cfg(feature = "observability")]
 pub mod observability;
+pub mod prepared_statement;
 pub mod sql_utils;
 pub mod stats;
 pub mod tables;