You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2023/05/27 10:14:45 UTC
[arrow-rs] branch master updated: feat(flight): add sql-info helpers (#4266)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 77aa8f5b2 feat(flight): add sql-info helpers (#4266)
77aa8f5b2 is described below
commit 77aa8f5b2645a91724048f5c1d644c6b52880028
Author: Robert Pack <42...@users.noreply.github.com>
AuthorDate: Sat May 27 12:14:39 2023 +0200
feat(flight): add sql-info helpers (#4266)
* feat: baseline sql-info helpers
* chore: clippy
* chore: add license to files
* docs: add some basic docstrings
* Update arrow-flight/src/sql/sql_info.rs
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
* fix: move flight info
* test: add simple filter test
* fix: docs link
* fix: one more docs link
* fix: one more one more docs link
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
arrow-flight/Cargo.toml | 8 +-
arrow-flight/examples/flight_sql_server.rs | 104 +++++---
arrow-flight/src/sql/mod.rs | 4 +
arrow-flight/src/sql/server.rs | 19 +-
arrow-flight/src/sql/sql_info.rs | 376 +++++++++++++++++++++++++++++
5 files changed, 464 insertions(+), 47 deletions(-)
diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml
index e22642b2a..206cc6505 100644
--- a/arrow-flight/Cargo.toml
+++ b/arrow-flight/Cargo.toml
@@ -31,15 +31,17 @@ arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
# Cast is needed to work around https://github.com/apache/arrow-rs/issues/3389
arrow-cast = { workspace = true }
+arrow-data = { workspace = true }
arrow-ipc = { workspace = true }
arrow-schema = { workspace = true }
base64 = { version = "0.21", default-features = false, features = ["std"] }
-tonic = { version = "0.9", default-features = false, features = ["transport", "codegen", "prost"] }
bytes = { version = "1", default-features = false }
+futures = { version = "0.3", default-features = false, features = ["alloc"] }
+once_cell = { version = "1", optional = true }
paste = { version = "1.0" }
prost = { version = "0.11", default-features = false, features = ["prost-derive"] }
tokio = { version = "1.0", default-features = false, features = ["macros", "rt", "rt-multi-thread"] }
-futures = { version = "0.3", default-features = false, features = ["alloc"] }
+tonic = { version = "0.9", default-features = false, features = ["transport", "codegen", "prost"] }
# CLI-related dependencies
clap = { version = "4.1", default-features = false, features = ["std", "derive", "env", "help", "error-context", "usage"], optional = true }
@@ -51,7 +53,7 @@ all-features = true
[features]
default = []
-flight-sql-experimental = []
+flight-sql-experimental = ["once_cell"]
tls = ["tonic/tls"]
# Enable CLI tools
diff --git a/arrow-flight/examples/flight_sql_server.rs b/arrow-flight/examples/flight_sql_server.rs
index 01632285c..27ae5d854 100644
--- a/arrow-flight/examples/flight_sql_server.rs
+++ b/arrow-flight/examples/flight_sql_server.rs
@@ -15,22 +15,10 @@
// specific language governing permissions and limitations
// under the License.
-use arrow_array::builder::StringBuilder;
-use arrow_array::{ArrayRef, RecordBatch};
-use arrow_flight::sql::{
- ActionBeginSavepointRequest, ActionBeginSavepointResult,
- ActionBeginTransactionResult, ActionCancelQueryRequest, ActionCancelQueryResult,
- ActionCreatePreparedStatementResult, ActionEndSavepointRequest,
- ActionEndTransactionRequest, Any, CommandStatementSubstraitPlan, ProstMessageExt,
- SqlInfo,
-};
-use arrow_flight::{
- Action, FlightData, FlightEndpoint, HandshakeRequest, HandshakeResponse, IpcMessage,
- Location, SchemaAsIpc, Ticket,
-};
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
-use futures::{stream, Stream};
+use futures::{stream, Stream, TryStreamExt};
+use once_cell::sync::Lazy;
use prost::Message;
use std::pin::Pin;
use std::sync::Arc;
@@ -38,22 +26,30 @@ use tonic::transport::Server;
use tonic::transport::{Certificate, Identity, ServerTlsConfig};
use tonic::{Request, Response, Status, Streaming};
+use arrow_array::builder::StringBuilder;
+use arrow_array::{ArrayRef, RecordBatch};
+use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::flight_descriptor::DescriptorType;
+use arrow_flight::sql::sql_info::SqlInfoList;
+use arrow_flight::sql::{
+ server::FlightSqlService, ActionBeginSavepointRequest, ActionBeginSavepointResult,
+ ActionBeginTransactionRequest, ActionBeginTransactionResult,
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, ActionCreatePreparedSubstraitPlanRequest,
+ ActionEndSavepointRequest, ActionEndTransactionRequest, Any, CommandGetCatalogs,
+ CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
+ CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
+ CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
+ CommandStatementSubstraitPlan, CommandStatementUpdate, ProstMessageExt, SqlInfo,
+ TicketStatementQuery,
+};
use arrow_flight::utils::batches_to_flight_data;
use arrow_flight::{
- flight_service_server::FlightService,
- flight_service_server::FlightServiceServer,
- sql::{
- server::FlightSqlService, ActionBeginTransactionRequest,
- ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
- ActionCreatePreparedSubstraitPlanRequest, CommandGetCatalogs,
- CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
- CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
- CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
- CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
- CommandStatementQuery, CommandStatementUpdate, TicketStatementQuery,
- },
- FlightDescriptor, FlightInfo,
+ flight_service_server::FlightService, flight_service_server::FlightServiceServer,
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, Location, SchemaAsIpc, Ticket,
};
use arrow_ipc::writer::IpcWriteOptions;
use arrow_schema::{ArrowError, DataType, Field, Schema};
@@ -68,6 +64,15 @@ const FAKE_TOKEN: &str = "uuid_token";
const FAKE_HANDLE: &str = "uuid_handle";
const FAKE_UPDATE_RESULT: i64 = 1;
+static INSTANCE_SQL_INFO: Lazy<SqlInfoList> = Lazy::new(|| {
+ SqlInfoList::new()
+ // Server information
+ .with_sql_info(SqlInfo::FlightSqlServerName, "Example Flight SQL Server")
+ .with_sql_info(SqlInfo::FlightSqlServerVersion, "1")
+ // 1.3 comes from https://github.com/apache/arrow/blob/f9324b79bf4fc1ec7e97b32e3cce16e75ef0f5e3/format/Schema.fbs#L24
+ .with_sql_info(SqlInfo::FlightSqlServerArrowVersion, "1.3")
+});
+
#[derive(Clone)]
pub struct FlightSqlServiceImpl {}
@@ -283,12 +288,38 @@ impl FlightSqlService for FlightSqlServiceImpl {
async fn get_flight_info_sql_info(
&self,
- _query: CommandGetSqlInfo,
- _request: Request<FlightDescriptor>,
+ query: CommandGetSqlInfo,
+ request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
- Err(Status::unimplemented(
- "get_flight_info_sql_info not implemented",
- ))
+ let flight_descriptor = request.into_inner();
+ let ticket = Ticket {
+ ticket: query.encode_to_vec().into(),
+ };
+
+ let options = IpcWriteOptions::default();
+
+ // encode the schema into the correct form
+ let IpcMessage(schema) = SchemaAsIpc::new(SqlInfoList::schema(), &options)
+ .try_into()
+ .expect("valid sql_info schema");
+
+ let endpoint = vec![FlightEndpoint {
+ ticket: Some(ticket),
+ // we assume users wnating to use this helper would reasonably
+ // never need to be distributed across multile endpoints?
+ location: vec![],
+ }];
+
+ let flight_info = FlightInfo {
+ schema,
+ flight_descriptor: Some(flight_descriptor),
+ endpoint,
+ total_records: -1,
+ total_bytes: -1,
+ ordered: false,
+ };
+
+ Ok(tonic::Response::new(flight_info))
}
async fn get_flight_info_primary_keys(
@@ -394,10 +425,15 @@ impl FlightSqlService for FlightSqlServiceImpl {
async fn do_get_sql_info(
&self,
- _query: CommandGetSqlInfo,
+ query: CommandGetSqlInfo,
_request: Request<Ticket>,
) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
- Err(Status::unimplemented("do_get_sql_info not implemented"))
+ let batch = INSTANCE_SQL_INFO.filter(&query.info).encode();
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(Arc::new(SqlInfoList::schema().clone()))
+ .build(futures::stream::once(async { batch }))
+ .map_err(Status::from);
+ Ok(Response::new(Box::pin(stream)))
}
async fn do_get_primary_keys(
diff --git a/arrow-flight/src/sql/mod.rs b/arrow-flight/src/sql/mod.rs
index 797ddfc9e..2c193b78b 100644
--- a/arrow-flight/src/sql/mod.rs
+++ b/arrow-flight/src/sql/mod.rs
@@ -84,6 +84,7 @@ pub use gen::SqlSupportedPositionedCommands;
pub use gen::SqlSupportedResultSetConcurrency;
pub use gen::SqlSupportedResultSetType;
pub use gen::SqlSupportedSubqueries;
+pub use gen::SqlSupportedTransaction;
pub use gen::SqlSupportedTransactions;
pub use gen::SqlSupportedUnions;
pub use gen::SqlSupportsConvert;
@@ -92,8 +93,11 @@ pub use gen::SupportedSqlGrammar;
pub use gen::TicketStatementQuery;
pub use gen::UpdateDeleteRules;
+pub use sql_info::SqlInfoList;
+
pub mod client;
pub mod server;
+pub mod sql_info;
/// ProstMessageExt are useful utility methods for prost::Message types
pub trait ProstMessageExt: prost::Message + Default {
diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs
index 89eb70e23..a33b5b92d 100644
--- a/arrow-flight/src/sql/server.rs
+++ b/arrow-flight/src/sql/server.rs
@@ -19,30 +19,29 @@
use std::pin::Pin;
-use crate::sql::{Any, Command};
use futures::Stream;
use prost::Message;
use tonic::{Request, Response, Status, Streaming};
use super::{
- super::{
- flight_service_server::FlightService, Action, ActionType, Criteria, Empty,
- FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse,
- PutResult, SchemaResult, Ticket,
- },
ActionBeginSavepointRequest, ActionBeginSavepointResult,
ActionBeginTransactionRequest, ActionBeginTransactionResult,
ActionCancelQueryRequest, ActionCancelQueryResult,
ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
ActionCreatePreparedStatementResult, ActionCreatePreparedSubstraitPlanRequest,
- ActionEndSavepointRequest, ActionEndTransactionRequest, CommandGetCatalogs,
- CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
- CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
- CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ ActionEndSavepointRequest, ActionEndTransactionRequest, Any, Command,
+ CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas,
+ CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
CommandStatementSubstraitPlan, CommandStatementUpdate, DoPutUpdateResult,
ProstMessageExt, SqlInfo, TicketStatementQuery,
};
+use crate::{
+ flight_service_server::FlightService, Action, ActionType, Criteria, Empty,
+ FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse,
+ PutResult, SchemaResult, Ticket,
+};
pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement";
pub(crate) static CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement";
diff --git a/arrow-flight/src/sql/sql_info.rs b/arrow-flight/src/sql/sql_info.rs
new file mode 100644
index 000000000..e0b7df70c
--- /dev/null
+++ b/arrow-flight/src/sql/sql_info.rs
@@ -0,0 +1,376 @@
+// 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.
+
+//! Auxiliary module to handle [`crate::sql::CommandGetSqlInfo`] queries.
+//!
+//! [`crate::sql::CommandGetSqlInfo`] represents metadata requests againsts the Flight SQL server.
+//! Via this mechanism, the server can communicate supported capabilities to generic
+//! Flight SQL clients.
+//!
+//! Servers construct a [`SqlInfoList`] by adding infos via `with_sql_info`.
+//! The availabe configuration options are defined in the [Flight SQL protos][protos].
+//!
+//! [protos]: https://github.com/apache/arrow/blob/6d3d2fca2c9693231fa1e52c142ceef563fc23f9/format/FlightSql.proto#L71-L820
+
+use std::{borrow::Cow, collections::BTreeMap, sync::Arc};
+
+use arrow_array::array::{Array, UnionArray};
+use arrow_array::builder::{
+ ArrayBuilder, BooleanBuilder, Int32Builder, Int64Builder, Int8Builder, ListBuilder,
+ StringBuilder, UInt32Builder,
+};
+use arrow_array::RecordBatch;
+use arrow_data::ArrayData;
+use arrow_schema::{DataType, Field, Schema, UnionFields, UnionMode};
+use once_cell::sync::Lazy;
+
+use super::SqlInfo;
+use crate::error::Result;
+
+/// Represents a dynamic value
+#[derive(Debug, Clone, PartialEq)]
+pub enum SqlInfoValue {
+ String(String),
+ Bool(bool),
+ BigInt(i64),
+ Bitmask(i32),
+ StringList(Vec<String>),
+ // TODO support more exotic metadata that requires the map of lists
+ //ListMap(BTreeMap<i32, Vec<i32>>),
+}
+
+impl From<&str> for SqlInfoValue {
+ fn from(value: &str) -> Self {
+ Self::String(value.to_string())
+ }
+}
+
+impl From<bool> for SqlInfoValue {
+ fn from(value: bool) -> Self {
+ Self::Bool(value)
+ }
+}
+
+impl From<i32> for SqlInfoValue {
+ fn from(value: i32) -> Self {
+ Self::Bitmask(value)
+ }
+}
+
+impl From<i64> for SqlInfoValue {
+ fn from(value: i64) -> Self {
+ Self::BigInt(value)
+ }
+}
+
+impl From<&[&str]> for SqlInfoValue {
+ fn from(values: &[&str]) -> Self {
+ let values = values.iter().map(|s| s.to_string()).collect();
+ Self::StringList(values)
+ }
+}
+
+/// Something that can be converted into u32 (the represenation of a [`SqlInfo`] name)
+pub trait SqlInfoName {
+ fn as_u32(&self) -> u32;
+}
+
+impl SqlInfoName for SqlInfo {
+ fn as_u32(&self) -> u32 {
+ // SqlInfos are u32 in the flight spec, but for some reason
+ // SqlInfo repr is an i32, so convert between them
+ u32::try_from(i32::from(*self)).expect("SqlInfo fit into u32")
+ }
+}
+
+// Allow passing u32 directly into to with_sql_info
+impl SqlInfoName for u32 {
+ fn as_u32(&self) -> u32 {
+ *self
+ }
+}
+
+/// Handles creating the dense [`UnionArray`] described by [flightsql]
+///
+///
+/// NOT YET COMPLETE: The int32_to_int32_list_map
+///
+/// ```text
+/// * value: dense_union<
+/// * string_value: utf8,
+/// * bool_value: bool,
+/// * bigint_value: int64,
+/// * int32_bitmask: int32,
+/// * string_list: list<string_data: utf8>
+/// * int32_to_int32_list_map: map<key: int32, value: list<$data$: int32>>
+/// * >
+/// ```
+///[flightsql]: (https://github.com/apache/arrow/blob/f9324b79bf4fc1ec7e97b32e3cce16e75ef0f5e3/format/FlightSql.proto#L32-L43
+struct SqlInfoUnionBuilder {
+ // Values for each child type
+ string_values: StringBuilder,
+ bool_values: BooleanBuilder,
+ bigint_values: Int64Builder,
+ int32_bitmask_values: Int32Builder,
+ string_list_values: ListBuilder<StringBuilder>,
+
+ /// incrementally build types/offset of the dense union,
+ ///
+ /// See [Union Spec] for details.
+ ///
+ /// [Union Spec]: https://arrow.apache.org/docs/format/Columnar.html#dense-union
+ type_ids: Int8Builder,
+ offsets: Int32Builder,
+}
+
+/// [`DataType`] for the output union array
+static UNION_TYPE: Lazy<DataType> = Lazy::new(|| {
+ let fields = vec![
+ Field::new("string_value", DataType::Utf8, false),
+ Field::new("bool_value", DataType::Boolean, false),
+ Field::new("bigint_value", DataType::Int64, false),
+ Field::new("int32_bitmask", DataType::Int32, false),
+ // treat list as nullable b/c that is what the builders make
+ Field::new(
+ "string_list",
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
+ true,
+ ),
+ ];
+
+ // create "type ids", one for each type, assume they go from 0 .. num_fields
+ let type_ids: Vec<i8> = (0..fields.len()).map(|v| v as i8).collect();
+
+ DataType::Union(UnionFields::new(type_ids, fields), UnionMode::Dense)
+});
+
+impl SqlInfoUnionBuilder {
+ pub fn new() -> Self {
+ Self {
+ string_values: StringBuilder::new(),
+ bool_values: BooleanBuilder::new(),
+ bigint_values: Int64Builder::new(),
+ int32_bitmask_values: Int32Builder::new(),
+ string_list_values: ListBuilder::new(StringBuilder::new()),
+ type_ids: Int8Builder::new(),
+ offsets: Int32Builder::new(),
+ }
+ }
+
+ /// Returns the DataType created by this builder
+ pub fn schema() -> &'static DataType {
+ &UNION_TYPE
+ }
+
+ /// Append the specified value to this builder
+ pub fn append_value(&mut self, v: &SqlInfoValue) {
+ // typeid is which child and len is the child array's length
+ // *after* adding the value
+ let (type_id, len) = match v {
+ SqlInfoValue::String(v) => {
+ self.string_values.append_value(v);
+ (0, self.string_values.len())
+ }
+ SqlInfoValue::Bool(v) => {
+ self.bool_values.append_value(*v);
+ (1, self.bool_values.len())
+ }
+ SqlInfoValue::BigInt(v) => {
+ self.bigint_values.append_value(*v);
+ (2, self.bigint_values.len())
+ }
+ SqlInfoValue::Bitmask(v) => {
+ self.int32_bitmask_values.append_value(*v);
+ (3, self.int32_bitmask_values.len())
+ }
+ SqlInfoValue::StringList(values) => {
+ // build list
+ for v in values {
+ self.string_list_values.values().append_value(v);
+ }
+ // complete the list
+ self.string_list_values.append(true);
+ (4, self.string_list_values.len())
+ }
+ };
+
+ self.type_ids.append_value(type_id);
+ let len = i32::try_from(len).expect("offset fit in i32");
+ self.offsets.append_value(len - 1);
+ }
+
+ /// Complete the construction and build the [`UnionArray`]
+ pub fn finish(self) -> UnionArray {
+ let Self {
+ mut string_values,
+ mut bool_values,
+ mut bigint_values,
+ mut int32_bitmask_values,
+ mut string_list_values,
+ mut type_ids,
+ mut offsets,
+ } = self;
+ let type_ids = type_ids.finish();
+ let offsets = offsets.finish();
+
+ // form the correct ArrayData
+
+ let len = offsets.len();
+ let null_bit_buffer = None;
+ let offset = 0;
+
+ let buffers = vec![
+ type_ids.into_data().buffers()[0].clone(),
+ offsets.into_data().buffers()[0].clone(),
+ ];
+
+ let child_data = vec![
+ string_values.finish().into_data(),
+ bool_values.finish().into_data(),
+ bigint_values.finish().into_data(),
+ int32_bitmask_values.finish().into_data(),
+ string_list_values.finish().into_data(),
+ ];
+
+ let data = ArrayData::try_new(
+ UNION_TYPE.clone(),
+ len,
+ null_bit_buffer,
+ offset,
+ buffers,
+ child_data,
+ )
+ .expect("Correctly created UnionArray");
+
+ UnionArray::from(data)
+ }
+}
+
+/// A list of SQL info names and valies
+#[derive(Debug, Clone, PartialEq)]
+pub struct SqlInfoList {
+ /// Use BTreeMap to ensure the values are sorted by value as
+ /// to make output consistent
+ ///
+ /// Use u32 to support "custom" sql info values that are not
+ /// part of the SqlInfo enum
+ infos: BTreeMap<u32, SqlInfoValue>,
+}
+
+impl Default for SqlInfoList {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SqlInfoList {
+ pub fn new() -> Self {
+ Self {
+ infos: BTreeMap::new(),
+ }
+ }
+
+ /// register the specific sql metadata item
+ pub fn with_sql_info(
+ mut self,
+ name: impl SqlInfoName,
+ value: impl Into<SqlInfoValue>,
+ ) -> Self {
+ self.infos.insert(name.as_u32(), value.into());
+ self
+ }
+
+ /// Filter this info list keeping only the info values specified
+ /// in `infos`.
+ ///
+ /// Returns self if infos is empty (no filtering)
+ pub fn filter(&self, info: &[u32]) -> Cow<'_, Self> {
+ if info.is_empty() {
+ Cow::Borrowed(self)
+ } else {
+ let infos: BTreeMap<_, _> = info
+ .iter()
+ .filter_map(|name| self.infos.get(name).map(|v| (*name, v.clone())))
+ .collect();
+ Cow::Owned(Self { infos })
+ }
+ }
+
+ /// Encode the contents of this info list according to the FlightSQL spec
+ pub fn encode(&self) -> Result<RecordBatch> {
+ let mut name_builder = UInt32Builder::new();
+ let mut value_builder = SqlInfoUnionBuilder::new();
+
+ for (&name, value) in self.infos.iter() {
+ name_builder.append_value(name);
+ value_builder.append_value(value)
+ }
+
+ let batch = RecordBatch::try_from_iter(vec![
+ ("info_name", Arc::new(name_builder.finish()) as _),
+ ("value", Arc::new(value_builder.finish()) as _),
+ ])?;
+ Ok(batch)
+ }
+
+ /// Return the [`Schema`] for a GetSchema RPC call with [`crate::sql::CommandGetSqlInfo`]
+ pub fn schema() -> &'static Schema {
+ // It is always the same
+ &SQL_INFO_SCHEMA
+ }
+}
+
+// The schema produced by [`SqlInfoList`]
+static SQL_INFO_SCHEMA: Lazy<Schema> = Lazy::new(|| {
+ Schema::new(vec![
+ Field::new("info_name", DataType::UInt32, false),
+ Field::new("value", SqlInfoUnionBuilder::schema().clone(), false),
+ ])
+});
+
+#[cfg(test)]
+mod tests {
+ use super::SqlInfoList;
+ use crate::sql::{SqlInfo, SqlSupportedTransaction};
+
+ #[test]
+ fn test_filter_sql_infos() {
+ let info_list = SqlInfoList::new()
+ .with_sql_info(SqlInfo::FlightSqlServerName, "server name")
+ .with_sql_info(
+ SqlInfo::FlightSqlServerTransaction,
+ SqlSupportedTransaction::Transaction as i32,
+ );
+
+ let batch = info_list.encode().unwrap();
+ assert_eq!(batch.num_rows(), 2);
+
+ let batch = info_list
+ .filter(&[SqlInfo::FlightSqlServerTransaction as u32])
+ .encode()
+ .unwrap();
+ let ref_batch = SqlInfoList::new()
+ .with_sql_info(
+ SqlInfo::FlightSqlServerTransaction,
+ SqlSupportedTransaction::Transaction as i32,
+ )
+ .encode()
+ .unwrap();
+
+ assert_eq!(batch, ref_batch);
+ }
+}