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

arrow_flight/sql/metadata/
xdbc_info.rs
