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

arrow_flight/sql/metadata/
sql_info.rs
