arrow_flight/sql/metadata/
sql_info.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Helpers for building responses to [`CommandGetSqlInfo`] metadata requests.
19//!
20//! - [`SqlInfoDataBuilder`] - a builder for collecting sql infos
21//!   and building a conformant `RecordBatch` with sql info server metadata.
22//! - [`SqlInfoData`] - a helper type wrapping a `RecordBatch`
23//!   used for storing sql info server metadata.
24//! - [`GetSqlInfoBuilder`] - a builder for constructing [`CommandGetSqlInfo`] responses.
25//!
26
27use std::collections::{BTreeMap, HashMap};
28use std::sync::Arc;
29
30use arrow_arith::boolean::or;
31use arrow_array::array::{Array, UInt32Array, UnionArray};
32use arrow_array::builder::{
33    ArrayBuilder, BooleanBuilder, Int8Builder, Int32Builder, Int64Builder, ListBuilder, MapBuilder,
34    StringBuilder, UInt32Builder,
35};
36use arrow_array::{RecordBatch, Scalar};
37use arrow_data::ArrayData;
38use arrow_ord::cmp::eq;
39use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, UnionFields, UnionMode};
40use arrow_select::filter::filter_record_batch;
41use once_cell::sync::Lazy;
42
43use crate::error::Result;
44use crate::sql::{CommandGetSqlInfo, SqlInfo};
45
46/// Represents a dynamic value
47#[derive(Debug, Clone, PartialEq)]
48pub enum SqlInfoValue {
49    String(String),
50    Bool(bool),
51    BigInt(i64),
52    Bitmask(i32),
53    StringList(Vec<String>),
54    ListMap(BTreeMap<i32, Vec<i32>>),
55}
56
57impl From<&str> for SqlInfoValue {
58    fn from(value: &str) -> Self {
59        Self::String(value.to_string())
60    }
61}
62
63impl From<bool> for SqlInfoValue {
64    fn from(value: bool) -> Self {
65        Self::Bool(value)
66    }
67}
68
69impl From<i32> for SqlInfoValue {
70    fn from(value: i32) -> Self {
71        Self::Bitmask(value)
72    }
73}
74
75impl From<i64> for SqlInfoValue {
76    fn from(value: i64) -> Self {
77        Self::BigInt(value)
78    }
79}
80
81impl From<&[&str]> for SqlInfoValue {
82    fn from(values: &[&str]) -> Self {
83        let values = values.iter().map(|s| s.to_string()).collect();
84        Self::StringList(values)
85    }
86}
87
88impl From<Vec<String>> for SqlInfoValue {
89    fn from(values: Vec<String>) -> Self {
90        Self::StringList(values)
91    }
92}
93
94impl From<BTreeMap<i32, Vec<i32>>> for SqlInfoValue {
95    fn from(value: BTreeMap<i32, Vec<i32>>) -> Self {
96        Self::ListMap(value)
97    }
98}
99
100impl From<HashMap<i32, Vec<i32>>> for SqlInfoValue {
101    fn from(value: HashMap<i32, Vec<i32>>) -> Self {
102        Self::ListMap(value.into_iter().collect())
103    }
104}
105
106impl From<&HashMap<i32, Vec<i32>>> for SqlInfoValue {
107    fn from(value: &HashMap<i32, Vec<i32>>) -> Self {
108        Self::ListMap(
109            value
110                .iter()
111                .map(|(k, v)| (k.to_owned(), v.to_owned()))
112                .collect(),
113        )
114    }
115}
116
117/// Something that can be converted into u32 (the representation of a [`SqlInfo`] name)
118pub trait SqlInfoName {
119    fn as_u32(&self) -> u32;
120}
121
122impl SqlInfoName for SqlInfo {
123    fn as_u32(&self) -> u32 {
124        // SqlInfos are u32 in the flight spec, but for some reason
125        // SqlInfo repr is an i32, so convert between them
126        u32::try_from(i32::from(*self)).expect("SqlInfo fit into u32")
127    }
128}
129
130// Allow passing u32 directly into to with_sql_info
131impl SqlInfoName for u32 {
132    fn as_u32(&self) -> u32 {
133        *self
134    }
135}
136
137/// Handles creating the dense [`UnionArray`] described by [flightsql]
138///
139/// incrementally build types/offset of the dense union. See [Union Spec] for details.
140///
141/// ```text
142/// *  value: dense_union<
143/// *              string_value: utf8,
144/// *              bool_value: bool,
145/// *              bigint_value: int64,
146/// *              int32_bitmask: int32,
147/// *              string_list: list<string_data: utf8>
148/// *              int32_to_int32_list_map: map<key: int32, value: list<$data$: int32>>
149/// * >
150/// ```
151///[flightsql]: https://github.com/apache/arrow/blob/f9324b79bf4fc1ec7e97b32e3cce16e75ef0f5e3/format/FlightSql.proto#L32-L43
152///[Union Spec]: https://arrow.apache.org/docs/format/Columnar.html#dense-union
153struct SqlInfoUnionBuilder {
154    // Values for each child type
155    string_values: StringBuilder,
156    bool_values: BooleanBuilder,
157    bigint_values: Int64Builder,
158    int32_bitmask_values: Int32Builder,
159    string_list_values: ListBuilder<StringBuilder>,
160    int32_to_int32_list_map_values: MapBuilder<Int32Builder, ListBuilder<Int32Builder>>,
161    type_ids: Int8Builder,
162    offsets: Int32Builder,
163}
164
165/// [`DataType`] for the output union array
166static UNION_TYPE: Lazy<DataType> = Lazy::new(|| {
167    let fields = vec![
168        Field::new("string_value", DataType::Utf8, false),
169        Field::new("bool_value", DataType::Boolean, false),
170        Field::new("bigint_value", DataType::Int64, false),
171        Field::new("int32_bitmask", DataType::Int32, false),
172        // treat list as nullable b/c that is what the builders make
173        Field::new(
174            "string_list",
175            DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))),
176            true,
177        ),
178        Field::new(
179            "int32_to_int32_list_map",
180            DataType::Map(
181                Arc::new(Field::new(
182                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
183                    DataType::Struct(Fields::from(vec![
184                        Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
185                        Field::new(
186                            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
187                            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
188                            true,
189                        ),
190                    ])),
191                    false,
192                )),
193                false,
194            ),
195            true,
196        ),
197    ];
198
199    DataType::Union(UnionFields::from_fields(fields), UnionMode::Dense)
200});
201
202impl SqlInfoUnionBuilder {
203    pub fn new() -> Self {
204        Self {
205            string_values: StringBuilder::new(),
206            bool_values: BooleanBuilder::new(),
207            bigint_values: Int64Builder::new(),
208            int32_bitmask_values: Int32Builder::new(),
209            string_list_values: ListBuilder::new(StringBuilder::new()),
210            int32_to_int32_list_map_values: MapBuilder::new(
211                None,
212                Int32Builder::new(),
213                ListBuilder::new(Int32Builder::new()),
214            ),
215            type_ids: Int8Builder::new(),
216            offsets: Int32Builder::new(),
217        }
218    }
219
220    /// Returns the DataType created by this builder
221    pub fn schema() -> &'static DataType {
222        &UNION_TYPE
223    }
224
225    /// Append the specified value to this builder
226    pub fn append_value(&mut self, v: &SqlInfoValue) -> Result<()> {
227        // typeid is which child and len is the child array's length
228        // *after* adding the value
229        let (type_id, len) = match v {
230            SqlInfoValue::String(v) => {
231                self.string_values.append_value(v);
232                (0, self.string_values.len())
233            }
234            SqlInfoValue::Bool(v) => {
235                self.bool_values.append_value(*v);
236                (1, self.bool_values.len())
237            }
238            SqlInfoValue::BigInt(v) => {
239                self.bigint_values.append_value(*v);
240                (2, self.bigint_values.len())
241            }
242            SqlInfoValue::Bitmask(v) => {
243                self.int32_bitmask_values.append_value(*v);
244                (3, self.int32_bitmask_values.len())
245            }
246            SqlInfoValue::StringList(values) => {
247                // build list
248                for v in values {
249                    self.string_list_values.values().append_value(v);
250                }
251                // complete the list
252                self.string_list_values.append(true);
253                (4, self.string_list_values.len())
254            }
255            SqlInfoValue::ListMap(values) => {
256                // build map
257                for (k, v) in values.clone() {
258                    self.int32_to_int32_list_map_values.keys().append_value(k);
259                    self.int32_to_int32_list_map_values
260                        .values()
261                        .append_value(v.into_iter().map(Some));
262                }
263                // complete the list
264                self.int32_to_int32_list_map_values.append(true)?;
265                (5, self.int32_to_int32_list_map_values.len())
266            }
267        };
268
269        self.type_ids.append_value(type_id);
270        let len = i32::try_from(len).expect("offset fit in i32");
271        self.offsets.append_value(len - 1);
272        Ok(())
273    }
274
275    /// Complete the construction and build the [`UnionArray`]
276    pub fn finish(self) -> UnionArray {
277        let Self {
278            mut string_values,
279            mut bool_values,
280            mut bigint_values,
281            mut int32_bitmask_values,
282            mut string_list_values,
283            mut int32_to_int32_list_map_values,
284            mut type_ids,
285            mut offsets,
286        } = self;
287        let type_ids = type_ids.finish();
288        let offsets = offsets.finish();
289
290        // form the correct ArrayData
291
292        let len = offsets.len();
293        let null_bit_buffer = None;
294        let offset = 0;
295
296        let buffers = vec![
297            type_ids.into_data().buffers()[0].clone(),
298            offsets.into_data().buffers()[0].clone(),
299        ];
300
301        let child_data = vec![
302            string_values.finish().into_data(),
303            bool_values.finish().into_data(),
304            bigint_values.finish().into_data(),
305            int32_bitmask_values.finish().into_data(),
306            string_list_values.finish().into_data(),
307            int32_to_int32_list_map_values.finish().into_data(),
308        ];
309
310        let data = ArrayData::try_new(
311            UNION_TYPE.clone(),
312            len,
313            null_bit_buffer,
314            offset,
315            buffers,
316            child_data,
317        )
318        .expect("Correctly created UnionArray");
319
320        UnionArray::from(data)
321    }
322}
323
324/// Helper to create [`CommandGetSqlInfo`] responses.
325///
326/// [`CommandGetSqlInfo`] are metadata requests used by a Flight SQL
327/// server to communicate supported capabilities to Flight SQL clients.
328///
329/// Servers construct - usually static - [`SqlInfoData`] via the [`SqlInfoDataBuilder`],
330/// and build responses using [`CommandGetSqlInfo::into_builder`]
331#[derive(Debug, Clone, PartialEq, Default)]
332pub struct SqlInfoDataBuilder {
333    /// Use BTreeMap to ensure the values are sorted by value as
334    /// to make output consistent
335    ///
336    /// Use u32 to support "custom" sql info values that are not
337    /// part of the SqlInfo enum
338    infos: BTreeMap<u32, SqlInfoValue>,
339}
340
341impl SqlInfoDataBuilder {
342    /// Create a new SQL info builder
343    pub fn new() -> Self {
344        Self::default()
345    }
346
347    /// register the specific sql metadata item
348    pub fn append(&mut self, name: impl SqlInfoName, value: impl Into<SqlInfoValue>) {
349        self.infos.insert(name.as_u32(), value.into());
350    }
351
352    /// Encode the contents of this list according to the [FlightSQL spec]
353    ///
354    /// [FlightSQL spec]: https://github.com/apache/arrow/blob/f9324b79bf4fc1ec7e97b32e3cce16e75ef0f5e3/format/FlightSql.proto#L32-L43
355    pub fn build(self) -> Result<SqlInfoData> {
356        let mut name_builder = UInt32Builder::new();
357        let mut value_builder = SqlInfoUnionBuilder::new();
358
359        // `infos` is a `BTreeMap`, so it iterates in sorted order already
360        for (name, value) in &self.infos {
361            name_builder.append_value(*name);
362            value_builder.append_value(value)?
363        }
364
365        let batch = RecordBatch::try_from_iter(vec![
366            ("info_name", Arc::new(name_builder.finish()) as _),
367            ("value", Arc::new(value_builder.finish()) as _),
368        ])?;
369
370        Ok(SqlInfoData { batch })
371    }
372
373    /// Return the [`Schema`] for a GetSchema RPC call with [`crate::sql::CommandGetSqlInfo`]
374    pub fn schema() -> &'static Schema {
375        // It is always the same
376        &SQL_INFO_SCHEMA
377    }
378}
379
380/// A builder for [`SqlInfoData`] which is used to create [`CommandGetSqlInfo`] responses.
381///
382/// # Example
383/// ```
384/// # use arrow_flight::sql::{metadata::SqlInfoDataBuilder, SqlInfo, SqlSupportedTransaction};
385/// // Create the list of metadata describing the server
386/// let mut builder = SqlInfoDataBuilder::new();
387/// builder.append(SqlInfo::FlightSqlServerName, "server name");
388///     // ... add other SqlInfo here ..
389/// builder.append(
390///     SqlInfo::FlightSqlServerTransaction,
391///     SqlSupportedTransaction::Transaction as i32,
392/// );
393///
394/// // Create the batch to send back to the client
395/// let info_data = builder.build().unwrap();
396/// ```
397///
398/// [protos]: https://github.com/apache/arrow/blob/6d3d2fca2c9693231fa1e52c142ceef563fc23f9/format/FlightSql.proto#L71-L820
399pub struct SqlInfoData {
400    batch: RecordBatch,
401}
402
403impl SqlInfoData {
404    /// Return a  [`RecordBatch`] containing only the requested `u32`, if any
405    /// from [`CommandGetSqlInfo`]
406    pub fn record_batch(&self, info: impl IntoIterator<Item = u32>) -> Result<RecordBatch> {
407        let arr = self.batch.column(0);
408        let type_filter = info
409            .into_iter()
410            .map(|tt| {
411                let s = UInt32Array::from(vec![tt]);
412                eq(arr, &Scalar::new(&s))
413            })
414            .collect::<std::result::Result<Vec<_>, _>>()?
415            .into_iter()
416            // We know the arrays are of same length as they are produced from the same root array
417            .reduce(|filter, arr| or(&filter, &arr).unwrap());
418        if let Some(filter) = type_filter {
419            Ok(filter_record_batch(&self.batch, &filter)?)
420        } else {
421            Ok(self.batch.clone())
422        }
423    }
424
425    /// Return the schema of the RecordBatch that will be returned
426    /// from [`CommandGetSqlInfo`]
427    pub fn schema(&self) -> SchemaRef {
428        self.batch.schema()
429    }
430}
431
432/// A builder for a [`CommandGetSqlInfo`] response.
433pub struct GetSqlInfoBuilder<'a> {
434    /// requested `SqlInfo`s. If empty means return all infos.
435    info: Vec<u32>,
436    infos: &'a SqlInfoData,
437}
438
439impl CommandGetSqlInfo {
440    /// Create a builder suitable for constructing a response
441    pub fn into_builder(self, infos: &SqlInfoData) -> GetSqlInfoBuilder<'_> {
442        GetSqlInfoBuilder {
443            info: self.info,
444            infos,
445        }
446    }
447}
448
449impl GetSqlInfoBuilder<'_> {
450    /// Builds a `RecordBatch` with the correct schema for a [`CommandGetSqlInfo`] response
451    pub fn build(self) -> Result<RecordBatch> {
452        self.infos.record_batch(self.info)
453    }
454
455    /// Return the schema of the RecordBatch that will be returned
456    /// from [`CommandGetSqlInfo`]
457    pub fn schema(&self) -> SchemaRef {
458        self.infos.schema()
459    }
460}
461
462// The schema produced by [`SqlInfoData`]
463static SQL_INFO_SCHEMA: Lazy<Schema> = Lazy::new(|| {
464    Schema::new(vec![
465        Field::new("info_name", DataType::UInt32, false),
466        Field::new("value", SqlInfoUnionBuilder::schema().clone(), false),
467    ])
468});
469
470#[cfg(test)]
471mod tests {
472    use std::collections::HashMap;
473
474    use super::SqlInfoDataBuilder;
475    use crate::sql::metadata::tests::assert_batches_eq;
476    use crate::sql::{SqlInfo, SqlNullOrdering, SqlSupportedTransaction, SqlSupportsConvert};
477
478    #[test]
479    fn test_sql_infos() {
480        let mut convert: HashMap<i32, Vec<i32>> = HashMap::new();
481        convert.insert(
482            SqlSupportsConvert::SqlConvertInteger as i32,
483            vec![
484                SqlSupportsConvert::SqlConvertFloat as i32,
485                SqlSupportsConvert::SqlConvertReal as i32,
486            ],
487        );
488
489        let mut builder = SqlInfoDataBuilder::new();
490        // str
491        builder.append(SqlInfo::SqlIdentifierQuoteChar, r#"""#);
492        // bool
493        builder.append(SqlInfo::SqlDdlCatalog, false);
494        // i32
495        builder.append(
496            SqlInfo::SqlNullOrdering,
497            SqlNullOrdering::SqlNullsSortedHigh as i32,
498        );
499        // i64
500        builder.append(SqlInfo::SqlMaxBinaryLiteralLength, i32::MAX as i64);
501        // [str]
502        builder.append(SqlInfo::SqlKeywords, &["SELECT", "DELETE"] as &[&str]);
503        builder.append(SqlInfo::SqlSupportsConvert, &convert);
504
505        let batch = builder.build().unwrap().record_batch(None).unwrap();
506
507        let expected = vec![
508            "+-----------+----------------------------------------+",
509            "| info_name | value                                  |",
510            "+-----------+----------------------------------------+",
511            "| 500       | {bool_value=false}                     |",
512            "| 504       | {string_value=\"}                       |",
513            "| 507       | {int32_bitmask=0}                      |",
514            "| 508       | {string_list=[SELECT, DELETE]}         |",
515            "| 517       | {int32_to_int32_list_map={7: [6, 13]}} |",
516            "| 541       | {bigint_value=2147483647}              |",
517            "+-----------+----------------------------------------+",
518        ];
519
520        assert_batches_eq(&[batch], &expected);
521    }
522
523    #[test]
524    fn test_filter_sql_infos() {
525        let mut builder = SqlInfoDataBuilder::new();
526        builder.append(SqlInfo::FlightSqlServerName, "server name");
527        builder.append(
528            SqlInfo::FlightSqlServerTransaction,
529            SqlSupportedTransaction::Transaction as i32,
530        );
531        let data = builder.build().unwrap();
532
533        let batch = data.record_batch(None).unwrap();
534        assert_eq!(batch.num_rows(), 2);
535
536        let batch = data
537            .record_batch([SqlInfo::FlightSqlServerTransaction as u32])
538            .unwrap();
539        let mut ref_builder = SqlInfoDataBuilder::new();
540        ref_builder.append(
541            SqlInfo::FlightSqlServerTransaction,
542            SqlSupportedTransaction::Transaction as i32,
543        );
544        let ref_batch = ref_builder.build().unwrap().record_batch(None).unwrap();
545
546        assert_eq!(batch, ref_batch);
547    }
548}
arrow_flight/sql/metadata/sql_info.rs

arrow_flight/sql/metadata/
sql_info.rs
