arrow_flight/sql/metadata/
xdbc_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 [`CommandGetXdbcTypeInfo`] metadata requests.
19//!
20//! - [`XdbcTypeInfo`] - a typed struct that holds the xdbc info corresponding to expected schema.
21//! - [`XdbcTypeInfoDataBuilder`] - a builder for collecting type infos
22//!   and building a conformant `RecordBatch`.
23//! - [`XdbcTypeInfoData`] - a helper type wrapping a `RecordBatch`
24//!   used for storing xdbc server metadata.
25//! - [`GetXdbcTypeInfoBuilder`] - a builder for constructing [`CommandGetXdbcTypeInfo`] responses.
26//!
27use std::sync::{Arc, LazyLock};
28
29use arrow_array::builder::{BooleanBuilder, Int32Builder, ListBuilder, StringBuilder};
30use arrow_array::{ArrayRef, Int32Array, ListArray, RecordBatch, Scalar};
31use arrow_ord::cmp::eq;
32use arrow_schema::{DataType, Field, Schema, SchemaRef};
33use arrow_select::filter::filter_record_batch;
34use arrow_select::take::take;
35
36use super::lexsort_to_indices;
37use crate::error::*;
38use crate::sql::{CommandGetXdbcTypeInfo, Nullable, Searchable, XdbcDataType, XdbcDatetimeSubcode};
39
40/// Data structure representing type information for xdbc types.
41#[derive(Debug, Clone, Default)]
42pub struct XdbcTypeInfo {
43    /// The name of the type
44    pub type_name: String,
45    /// The data type of the type
46    pub data_type: XdbcDataType,
47    /// The column size of the type
48    pub column_size: Option<i32>,
49    /// The prefix of the type
50    pub literal_prefix: Option<String>,
51    /// The suffix of the type
52    pub literal_suffix: Option<String>,
53    /// The create parameters of the type
54    pub create_params: Option<Vec<String>>,
55    /// The nullability of the type
56    pub nullable: Nullable,
57    /// Whether the type is case sensitive
58    pub case_sensitive: bool,
59    /// Whether the type is searchable
60    pub searchable: Searchable,
61    /// Whether the type is unsigned
62    pub unsigned_attribute: Option<bool>,
63    /// Whether the type has fixed precision and scale
64    pub fixed_prec_scale: bool,
65    /// Whether the type is auto-incrementing
66    pub auto_increment: Option<bool>,
67    /// The local type name of the type
68    pub local_type_name: Option<String>,
69    /// The minimum scale of the type
70    pub minimum_scale: Option<i32>,
71    /// The maximum scale of the type
72    pub maximum_scale: Option<i32>,
73    /// The SQL data type of the type
74    pub sql_data_type: XdbcDataType,
75    /// The optional datetime subcode of the type
76    pub datetime_subcode: Option<XdbcDatetimeSubcode>,
77    /// The number precision radix of the type
78    pub num_prec_radix: Option<i32>,
79    /// The interval precision of the type
80    pub interval_precision: Option<i32>,
81}
82
83/// Helper to create [`CommandGetXdbcTypeInfo`] responses.
84///
85/// [`CommandGetXdbcTypeInfo`] are metadata requests used by a Flight SQL
86/// server to communicate supported capabilities to Flight SQL clients.
87///
88/// Servers construct - usually static - [`XdbcTypeInfoData`] via the [`XdbcTypeInfoDataBuilder`],
89/// and build responses using [`CommandGetXdbcTypeInfo::into_builder`].
90pub struct XdbcTypeInfoData {
91    batch: RecordBatch,
92}
93
94impl XdbcTypeInfoData {
95    /// Return the raw (not encoded) RecordBatch that will be returned
96    /// from [`CommandGetXdbcTypeInfo`]
97    pub fn record_batch(&self, data_type: impl Into<Option<i32>>) -> Result<RecordBatch> {
98        if let Some(dt) = data_type.into() {
99            let scalar = Int32Array::from(vec![dt]);
100            let filter = eq(self.batch.column(1), &Scalar::new(&scalar))?;
101            Ok(filter_record_batch(&self.batch, &filter)?)
102        } else {
103            Ok(self.batch.clone())
104        }
105    }
106
107    /// Return the schema of the RecordBatch that will be returned
108    /// from [`CommandGetXdbcTypeInfo`]
109    pub fn schema(&self) -> SchemaRef {
110        self.batch.schema()
111    }
112}
113
114/// A builder for [`XdbcTypeInfoData`] which is used to create [`CommandGetXdbcTypeInfo`] responses.
115///
116/// # Example
117/// ```
118/// use arrow_flight::sql::{Nullable, Searchable, XdbcDataType};
119/// use arrow_flight::sql::metadata::{XdbcTypeInfo, XdbcTypeInfoDataBuilder};
120/// // Create the list of metadata describing the server. Since this would not change at
121/// // runtime, using LazyLock or similar patterns to construct the list is a common approach.
122/// let mut builder = XdbcTypeInfoDataBuilder::new();
123/// builder.append(XdbcTypeInfo {
124///     type_name: "INTEGER".into(),
125///     data_type: XdbcDataType::XdbcInteger,
126///     column_size: Some(32),
127///     literal_prefix: None,
128///     literal_suffix: None,
129///     create_params: None,
130///     nullable: Nullable::NullabilityNullable,
131///     case_sensitive: false,
132///     searchable: Searchable::Full,
133///     unsigned_attribute: Some(false),
134///     fixed_prec_scale: false,
135///     auto_increment: Some(false),
136///     local_type_name: Some("INTEGER".into()),
137///     minimum_scale: None,
138///     maximum_scale: None,
139///     sql_data_type: XdbcDataType::XdbcInteger,
140///     datetime_subcode: None,
141///     num_prec_radix: Some(2),
142///     interval_precision: None,
143/// });
144/// let info_list = builder.build().unwrap();
145///
146/// // to access the underlying record batch
147/// let batch = info_list.record_batch(None);
148/// ```
149pub struct XdbcTypeInfoDataBuilder {
150    infos: Vec<XdbcTypeInfo>,
151}
152
153impl Default for XdbcTypeInfoDataBuilder {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159impl XdbcTypeInfoDataBuilder {
160    /// Create a new instance of [`XdbcTypeInfoDataBuilder`].
161    pub fn new() -> Self {
162        Self { infos: Vec::new() }
163    }
164
165    /// Append a new row
166    pub fn append(&mut self, info: XdbcTypeInfo) {
167        self.infos.push(info);
168    }
169
170    /// Create helper structure for handling xdbc metadata requests.
171    pub fn build(self) -> Result<XdbcTypeInfoData> {
172        let mut type_name_builder = StringBuilder::new();
173        let mut data_type_builder = Int32Builder::new();
174        let mut column_size_builder = Int32Builder::new();
175        let mut literal_prefix_builder = StringBuilder::new();
176        let mut literal_suffix_builder = StringBuilder::new();
177        let mut create_params_builder = ListBuilder::new(StringBuilder::new());
178        let mut nullable_builder = Int32Builder::new();
179        let mut case_sensitive_builder = BooleanBuilder::new();
180        let mut searchable_builder = Int32Builder::new();
181        let mut unsigned_attribute_builder = BooleanBuilder::new();
182        let mut fixed_prec_scale_builder = BooleanBuilder::new();
183        let mut auto_increment_builder = BooleanBuilder::new();
184        let mut local_type_name_builder = StringBuilder::new();
185        let mut minimum_scale_builder = Int32Builder::new();
186        let mut maximum_scale_builder = Int32Builder::new();
187        let mut sql_data_type_builder = Int32Builder::new();
188        let mut datetime_subcode_builder = Int32Builder::new();
189        let mut num_prec_radix_builder = Int32Builder::new();
190        let mut interval_precision_builder = Int32Builder::new();
191
192        self.infos.into_iter().for_each(|info| {
193            type_name_builder.append_value(info.type_name);
194            data_type_builder.append_value(info.data_type as i32);
195            column_size_builder.append_option(info.column_size);
196            literal_prefix_builder.append_option(info.literal_prefix);
197            literal_suffix_builder.append_option(info.literal_suffix);
198            if let Some(params) = info.create_params {
199                if !params.is_empty() {
200                    for param in params {
201                        create_params_builder.values().append_value(param);
202                    }
203                    create_params_builder.append(true);
204                } else {
205                    create_params_builder.append_null();
206                }
207            } else {
208                create_params_builder.append_null();
209            }
210            nullable_builder.append_value(info.nullable as i32);
211            case_sensitive_builder.append_value(info.case_sensitive);
212            searchable_builder.append_value(info.searchable as i32);
213            unsigned_attribute_builder.append_option(info.unsigned_attribute);
214            fixed_prec_scale_builder.append_value(info.fixed_prec_scale);
215            auto_increment_builder.append_option(info.auto_increment);
216            local_type_name_builder.append_option(info.local_type_name);
217            minimum_scale_builder.append_option(info.minimum_scale);
218            maximum_scale_builder.append_option(info.maximum_scale);
219            sql_data_type_builder.append_value(info.sql_data_type as i32);
220            datetime_subcode_builder.append_option(info.datetime_subcode.map(|code| code as i32));
221            num_prec_radix_builder.append_option(info.num_prec_radix);
222            interval_precision_builder.append_option(info.interval_precision);
223        });
224
225        let type_name = Arc::new(type_name_builder.finish());
226        let data_type = Arc::new(data_type_builder.finish());
227        let column_size = Arc::new(column_size_builder.finish());
228        let literal_prefix = Arc::new(literal_prefix_builder.finish());
229        let literal_suffix = Arc::new(literal_suffix_builder.finish());
230        let (field, offsets, values, nulls) = create_params_builder.finish().into_parts();
231        // Re-defined the field to be non-nullable
232        let new_field = Arc::new(field.as_ref().clone().with_nullable(false));
233        let create_params = Arc::new(ListArray::new(new_field, offsets, values, nulls)) as ArrayRef;
234        let nullable = Arc::new(nullable_builder.finish());
235        let case_sensitive = Arc::new(case_sensitive_builder.finish());
236        let searchable = Arc::new(searchable_builder.finish());
237        let unsigned_attribute = Arc::new(unsigned_attribute_builder.finish());
238        let fixed_prec_scale = Arc::new(fixed_prec_scale_builder.finish());
239        let auto_increment = Arc::new(auto_increment_builder.finish());
240        let local_type_name = Arc::new(local_type_name_builder.finish());
241        let minimum_scale = Arc::new(minimum_scale_builder.finish());
242        let maximum_scale = Arc::new(maximum_scale_builder.finish());
243        let sql_data_type = Arc::new(sql_data_type_builder.finish());
244        let datetime_subcode = Arc::new(datetime_subcode_builder.finish());
245        let num_prec_radix = Arc::new(num_prec_radix_builder.finish());
246        let interval_precision = Arc::new(interval_precision_builder.finish());
247
248        let batch = RecordBatch::try_new(
249            Arc::clone(&GET_XDBC_INFO_SCHEMA),
250            vec![
251                type_name,
252                data_type,
253                column_size,
254                literal_prefix,
255                literal_suffix,
256                create_params,
257                nullable,
258                case_sensitive,
259                searchable,
260                unsigned_attribute,
261                fixed_prec_scale,
262                auto_increment,
263                local_type_name,
264                minimum_scale,
265                maximum_scale,
266                sql_data_type,
267                datetime_subcode,
268                num_prec_radix,
269                interval_precision,
270            ],
271        )?;
272
273        // Order batch by data_type and then by type_name
274        let sort_cols = batch.project(&[1, 0])?;
275        let indices = lexsort_to_indices(sort_cols.columns());
276        let columns = batch
277            .columns()
278            .iter()
279            .map(|c| take(c, &indices, None))
280            .collect::<std::result::Result<Vec<_>, _>>()?;
281
282        Ok(XdbcTypeInfoData {
283            batch: RecordBatch::try_new(batch.schema(), columns)?,
284        })
285    }
286
287    /// Return the [`Schema`] for a GetSchema RPC call with [`CommandGetXdbcTypeInfo`]
288    pub fn schema(&self) -> SchemaRef {
289        Arc::clone(&GET_XDBC_INFO_SCHEMA)
290    }
291}
292
293/// A builder for a [`CommandGetXdbcTypeInfo`] response.
294pub struct GetXdbcTypeInfoBuilder<'a> {
295    data_type: Option<i32>,
296    infos: &'a XdbcTypeInfoData,
297}
298
299impl CommandGetXdbcTypeInfo {
300    /// Create a builder suitable for constructing a response
301    pub fn into_builder(self, infos: &XdbcTypeInfoData) -> GetXdbcTypeInfoBuilder<'_> {
302        GetXdbcTypeInfoBuilder {
303            data_type: self.data_type,
304            infos,
305        }
306    }
307}
308
309impl GetXdbcTypeInfoBuilder<'_> {
310    /// Builds a `RecordBatch` with the correct schema for a [`CommandGetXdbcTypeInfo`] response
311    pub fn build(self) -> Result<RecordBatch> {
312        self.infos.record_batch(self.data_type)
313    }
314
315    /// Return the schema of the RecordBatch that will be returned
316    /// from [`CommandGetXdbcTypeInfo`]
317    pub fn schema(&self) -> SchemaRef {
318        self.infos.schema()
319    }
320}
321
322/// The schema for GetXdbcTypeInfo
323static GET_XDBC_INFO_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
324    Arc::new(Schema::new(vec![
325        Field::new("type_name", DataType::Utf8, false),
326        Field::new("data_type", DataType::Int32, false),
327        Field::new("column_size", DataType::Int32, true),
328        Field::new("literal_prefix", DataType::Utf8, true),
329        Field::new("literal_suffix", DataType::Utf8, true),
330        Field::new(
331            "create_params",
332            DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, false))),
333            true,
334        ),
335        Field::new("nullable", DataType::Int32, false),
336        Field::new("case_sensitive", DataType::Boolean, false),
337        Field::new("searchable", DataType::Int32, false),
338        Field::new("unsigned_attribute", DataType::Boolean, true),
339        Field::new("fixed_prec_scale", DataType::Boolean, false),
340        Field::new("auto_increment", DataType::Boolean, true),
341        Field::new("local_type_name", DataType::Utf8, true),
342        Field::new("minimum_scale", DataType::Int32, true),
343        Field::new("maximum_scale", DataType::Int32, true),
344        Field::new("sql_data_type", DataType::Int32, false),
345        Field::new("datetime_subcode", DataType::Int32, true),
346        Field::new("num_prec_radix", DataType::Int32, true),
347        Field::new("interval_precision", DataType::Int32, true),
348    ]))
349});
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    use crate::sql::metadata::tests::assert_batches_eq;
355
356    #[test]
357    fn test_create_batch() {
358        let mut builder = XdbcTypeInfoDataBuilder::new();
359        builder.append(XdbcTypeInfo {
360            type_name: "VARCHAR".into(),
361            data_type: XdbcDataType::XdbcVarchar,
362            column_size: Some(i32::MAX),
363            literal_prefix: Some("'".into()),
364            literal_suffix: Some("'".into()),
365            create_params: Some(vec!["length".into()]),
366            nullable: Nullable::NullabilityNullable,
367            case_sensitive: true,
368            searchable: Searchable::Full,
369            unsigned_attribute: None,
370            fixed_prec_scale: false,
371            auto_increment: None,
372            local_type_name: Some("VARCHAR".into()),
373            minimum_scale: None,
374            maximum_scale: None,
375            sql_data_type: XdbcDataType::XdbcVarchar,
376            datetime_subcode: None,
377            num_prec_radix: None,
378            interval_precision: None,
379        });
380        builder.append(XdbcTypeInfo {
381            type_name: "INTEGER".into(),
382            data_type: XdbcDataType::XdbcInteger,
383            column_size: Some(32),
384            literal_prefix: None,
385            literal_suffix: None,
386            create_params: None,
387            nullable: Nullable::NullabilityNullable,
388            case_sensitive: false,
389            searchable: Searchable::Full,
390            unsigned_attribute: Some(false),
391            fixed_prec_scale: false,
392            auto_increment: Some(false),
393            local_type_name: Some("INTEGER".into()),
394            minimum_scale: None,
395            maximum_scale: None,
396            sql_data_type: XdbcDataType::XdbcInteger,
397            datetime_subcode: None,
398            num_prec_radix: Some(2),
399            interval_precision: None,
400        });
401        builder.append(XdbcTypeInfo {
402            type_name: "INTERVAL".into(),
403            data_type: XdbcDataType::XdbcInterval,
404            column_size: Some(i32::MAX),
405            literal_prefix: Some("'".into()),
406            literal_suffix: Some("'".into()),
407            create_params: None,
408            nullable: Nullable::NullabilityNullable,
409            case_sensitive: false,
410            searchable: Searchable::Full,
411            unsigned_attribute: None,
412            fixed_prec_scale: false,
413            auto_increment: None,
414            local_type_name: Some("INTERVAL".into()),
415            minimum_scale: None,
416            maximum_scale: None,
417            sql_data_type: XdbcDataType::XdbcInterval,
418            datetime_subcode: Some(XdbcDatetimeSubcode::XdbcSubcodeUnknown),
419            num_prec_radix: None,
420            interval_precision: None,
421        });
422        let infos = builder.build().unwrap();
423
424        let batch = infos.record_batch(None).unwrap();
425        let expected = vec![
426            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
427            "| type_name | data_type | column_size | literal_prefix | literal_suffix | create_params | nullable | case_sensitive | searchable | unsigned_attribute | fixed_prec_scale | auto_increment | local_type_name | minimum_scale | maximum_scale | sql_data_type | datetime_subcode | num_prec_radix | interval_precision |",
428            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
429            "| INTEGER   | 4         | 32          |                |                |               | 1        | false          | 3          | false              | false            | false          | INTEGER         |               |               | 4             |                  | 2              |                    |",
430            "| INTERVAL  | 10        | 2147483647  | '              | '              |               | 1        | false          | 3          |                    | false            |                | INTERVAL        |               |               | 10            | 0                |                |                    |",
431            "| VARCHAR   | 12        | 2147483647  | '              | '              | [length]      | 1        | true           | 3          |                    | false            |                | VARCHAR         |               |               | 12            |                  |                |                    |",
432            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
433        ];
434        assert_batches_eq(&[batch], &expected);
435
436        let batch = infos.record_batch(Some(10)).unwrap();
437        let expected = vec![
438            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
439            "| type_name | data_type | column_size | literal_prefix | literal_suffix | create_params | nullable | case_sensitive | searchable | unsigned_attribute | fixed_prec_scale | auto_increment | local_type_name | minimum_scale | maximum_scale | sql_data_type | datetime_subcode | num_prec_radix | interval_precision |",
440            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
441            "| INTERVAL  | 10        | 2147483647  | '              | '              |               | 1        | false          | 3          |                    | false            |                | INTERVAL        |               |               | 10            | 0                |                |                    |",
442            "+-----------+-----------+-------------+----------------+----------------+---------------+----------+----------------+------------+--------------------+------------------+----------------+-----------------+---------------+---------------+---------------+------------------+----------------+--------------------+",
443        ];
444        assert_batches_eq(&[batch], &expected);
445    }
446}
arrow_flight/sql/metadata/xdbc_info.rs

arrow_flight/sql/metadata/
xdbc_info.rs
