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

arrow_flight/sql/metadata/
xdbc_info.rs
