arrow_flight/sql/metadata/
sql_info.rs

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

arrow_flight/sql/metadata/
sql_info.rs
