From 69b5e199ce374b2861d532dae8c2978a2fa80721 Mon Sep 17 00:00:00 2001 From: Emil Sadek Date: Mon, 13 Jul 2026 19:50:35 -0700 Subject: [PATCH 1/2] feat: show info on repl startup --- Cargo.lock | 1 + Cargo.toml | 1 + src/database.rs | 179 +++++++++++++++++++++++++++++++++++++++++++++++- src/repl.rs | 64 +++++++++++++++++ 4 files changed, 242 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9fe8082..242d84e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -601,6 +601,7 @@ dependencies = [ "adbc_driver_manager", "arrow", "arrow-array", + "arrow-buffer", "arrow-cast", "arrow-schema", "clap", diff --git a/Cargo.toml b/Cargo.toml index 96fa9ab..edde745 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,4 +25,5 @@ syntect = "5.3.0" terminal-colorsaurus = "1.0.3" [dev-dependencies] +arrow-buffer = "58.1.0" tempfile = "3.27.0" diff --git a/src/database.rs b/src/database.rs index 27bc236..d558e37 100644 --- a/src/database.rs +++ b/src/database.rs @@ -2,13 +2,99 @@ // SPDX-License-Identifier: Apache-2.0 use crate::cli::ConnectionSource; -use adbc_core::options::{AdbcVersion, OptionDatabase, OptionValue}; -use adbc_core::{Database, Driver, LOAD_FLAG_DEFAULT, Statement}; +use adbc_core::options::{AdbcVersion, InfoCode, OptionDatabase, OptionValue}; +use adbc_core::{Connection, Database, Driver, LOAD_FLAG_DEFAULT, Statement}; use adbc_driver_manager::profile::{ ConnectionProfile, ConnectionProfileProvider, FilesystemProfileProvider, process_profile_value, }; use adbc_driver_manager::{ManagedConnection, ManagedDriver}; -use arrow_array::RecordBatch; +use arrow_array::cast::AsArray; +use arrow_array::{Array, RecordBatch, UnionArray}; +use std::collections::HashSet; + +/// Vendor (database product) metadata reported by an ADBC driver. +#[derive(Debug, Default)] +pub struct VendorInfo { + pub name: Option, + pub version: Option, +} + +/// Query the connected driver for the database vendor name and version. +/// +/// This is best-effort: any failure (the driver does not implement `get_info`, +/// returns unexpected data, etc.) yields a `VendorInfo` with `None` fields so +/// that callers can continue without the metadata. +pub fn get_vendor_info(connection: &impl Connection) -> VendorInfo { + let codes = HashSet::from([InfoCode::VendorName, InfoCode::VendorVersion]); + let reader = match connection.get_info(Some(codes)) { + Ok(reader) => reader, + Err(_) => return VendorInfo::default(), + }; + let batches: Vec = match reader.collect::>() { + Ok(batches) => batches, + Err(_) => return VendorInfo::default(), + }; + parse_vendor_info(&batches) +} + +/// Extract the vendor name and version from the record batches returned by +/// [`Connection::get_info`]. +/// +/// The batches follow the ADBC `get_info` schema: an `info_name` (`u32`) column +/// and an `info_value` dense-union column. Vendor name/version are utf8 values +/// stored in the union's `string_value` child (type id 0). +fn parse_vendor_info(batches: &[RecordBatch]) -> VendorInfo { + // Derive the info codes from the same enum used in the `get_info` request + // so the parse cannot drift from the query. + let vendor_name_code = u32::from(&InfoCode::VendorName); + let vendor_version_code = u32::from(&InfoCode::VendorVersion); + + let mut info = VendorInfo::default(); + + for batch in batches { + let Some(info_names) = batch + .column_by_name("info_name") + .and_then(|col| col.as_primitive_opt::()) + else { + continue; + }; + let Some(info_values) = batch + .column_by_name("info_value") + .and_then(|col| col.as_any().downcast_ref::()) + else { + continue; + }; + + for row in 0..batch.num_rows() { + if info_names.is_null(row) { + continue; + } + let code = info_names.value(row); + if code != vendor_name_code && code != vendor_version_code { + continue; + } + + let value = info_values.value(row); + let Some(strings) = value.as_string_opt::() else { + continue; + }; + if strings.is_empty() || strings.is_null(0) { + continue; + } + let text = strings.value(0).to_string(); + + // The guard above ensures `code` is one of the two vendor codes, + // so `else` unambiguously means the version. + if code == vendor_name_code { + info.name = Some(text); + } else { + info.version = Some(text); + } + } + } + + info +} pub fn initialize_connection(source: ConnectionSource) -> Result { match source { @@ -193,6 +279,93 @@ fn merge_options( #[cfg(test)] mod tests { use super::*; + use adbc_core::schemas::GET_INFO_SCHEMA; + use arrow_array::{StringArray, UInt32Array, UnionArray}; + use arrow_buffer::ScalarBuffer; + use arrow_schema::DataType; + use std::sync::Arc; + + /// Build a RecordBatch matching the ADBC `get_info` schema from a list of + /// `(info_name, string_value)` pairs. `None` values are stored as null + /// entries in the union's `string_value` child. + fn make_info_batch(rows: &[(u32, Option<&str>)]) -> RecordBatch { + let info_names: Vec = rows.iter().map(|(name, _)| *name).collect(); + let string_values: Vec> = rows.iter().map(|(_, value)| *value).collect(); + + // The `info_value` union child fields, per GET_INFO_SCHEMA. All rows use + // the `string_value` child (type_id 0) in this helper. + let DataType::Union(union_fields, _) = GET_INFO_SCHEMA.field(1).data_type().clone() else { + panic!("info_value must be a union"); + }; + + let string_child = Arc::new(StringArray::from(string_values.clone())); + let children = union_fields + .iter() + .map(|(type_id, field)| -> arrow_array::ArrayRef { + if type_id == 0 { + string_child.clone() + } else { + // Empty child of the correct type for the unused union branches. + arrow_array::new_empty_array(field.data_type()) + } + }) + .collect::>(); + + let type_ids = ScalarBuffer::from(vec![0i8; rows.len()]); + let offsets = ScalarBuffer::from((0..rows.len() as i32).collect::>()); + let union = + UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap(); + + RecordBatch::try_new( + GET_INFO_SCHEMA.clone(), + vec![Arc::new(UInt32Array::from(info_names)), Arc::new(union)], + ) + .unwrap() + } + + #[test] + fn test_parse_vendor_info_name_and_version() { + let batch = make_info_batch(&[(0, Some("DuckDB")), (1, Some("v1.1.0"))]); + let info = parse_vendor_info(&[batch]); + assert_eq!(info.name.as_deref(), Some("DuckDB")); + assert_eq!(info.version.as_deref(), Some("v1.1.0")); + } + + #[test] + fn test_parse_vendor_info_name_only() { + let batch = make_info_batch(&[(0, Some("PostgreSQL"))]); + let info = parse_vendor_info(&[batch]); + assert_eq!(info.name.as_deref(), Some("PostgreSQL")); + assert_eq!(info.version, None); + } + + #[test] + fn test_parse_vendor_info_ignores_other_codes() { + // 100 = DriverName, 101 = DriverVersion; not vendor fields. + let batch = make_info_batch(&[ + (100, Some("ADBC DuckDB Driver")), + (1, Some("v1.1.0")), + (0, Some("DuckDB")), + ]); + let info = parse_vendor_info(&[batch]); + assert_eq!(info.name.as_deref(), Some("DuckDB")); + assert_eq!(info.version.as_deref(), Some("v1.1.0")); + } + + #[test] + fn test_parse_vendor_info_null_value() { + let batch = make_info_batch(&[(0, Some("DuckDB")), (1, None)]); + let info = parse_vendor_info(&[batch]); + assert_eq!(info.name.as_deref(), Some("DuckDB")); + assert_eq!(info.version, None); + } + + #[test] + fn test_parse_vendor_info_empty() { + let info = parse_vendor_info(&[]); + assert_eq!(info.name, None); + assert_eq!(info.version, None); + } #[test] fn test_build_database_options_empty() { diff --git a/src/repl.rs b/src/repl.rs index 81d9c67..665fae7 100644 --- a/src/repl.rs +++ b/src/repl.rs @@ -50,7 +50,25 @@ impl Prompt for SqlPrompt { } } +/// Build the REPL startup banner. The first line always shows `databow` and its +/// version. When the driver reports a vendor name, a second line shows the +/// connected database, including its version when available. +fn format_banner(version: &str, vendor: &database::VendorInfo) -> String { + let mut banner = format!("databow {version}"); + if let Some(name) = &vendor.name { + banner.push('\n'); + match &vendor.version { + Some(vendor_version) => banner.push_str(&format!("Connected to {name} {vendor_version}")), + None => banner.push_str(&format!("Connected to {name}")), + } + } + banner +} + pub fn run_repl(mut connection: impl Connection, table_mode: TableMode) { + let vendor = database::get_vendor_info(&connection); + println!("{}", format_banner(env!("CARGO_PKG_VERSION"), &vendor)); + let mut line_editor = Reedline::create() .with_highlighter(Box::new(SyntectHighlighter::new())) .with_validator(Box::new(SqlValidator)); @@ -93,3 +111,49 @@ pub fn run_repl(mut connection: impl Connection, table_mode: TableMode) { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::database::VendorInfo; + + #[test] + fn test_format_banner_version_only() { + let vendor = VendorInfo::default(); + assert_eq!(format_banner("0.1.2", &vendor), "databow 0.1.2"); + } + + #[test] + fn test_format_banner_with_vendor_name_and_version() { + let vendor = VendorInfo { + name: Some("DuckDB".to_string()), + version: Some("v1.1.0".to_string()), + }; + assert_eq!( + format_banner("0.1.2", &vendor), + "databow 0.1.2\nConnected to DuckDB v1.1.0" + ); + } + + #[test] + fn test_format_banner_with_vendor_name_only() { + let vendor = VendorInfo { + name: Some("PostgreSQL".to_string()), + version: None, + }; + assert_eq!( + format_banner("0.1.2", &vendor), + "databow 0.1.2\nConnected to PostgreSQL" + ); + } + + #[test] + fn test_format_banner_version_present_without_name_is_ignored() { + // A version with no name should not produce a second line. + let vendor = VendorInfo { + name: None, + version: Some("v1.1.0".to_string()), + }; + assert_eq!(format_banner("0.1.2", &vendor), "databow 0.1.2"); + } +} From 49c69e137bbf166460fa96f39070810266e7855e Mon Sep 17 00:00:00 2001 From: Emil Sadek Date: Mon, 13 Jul 2026 19:54:45 -0700 Subject: [PATCH 2/2] chore: apply rustfmt formatting --- src/database.rs | 3 +-- src/repl.rs | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/database.rs b/src/database.rs index d558e37..64a377e 100644 --- a/src/database.rs +++ b/src/database.rs @@ -313,8 +313,7 @@ mod tests { let type_ids = ScalarBuffer::from(vec![0i8; rows.len()]); let offsets = ScalarBuffer::from((0..rows.len() as i32).collect::>()); - let union = - UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap(); + let union = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap(); RecordBatch::try_new( GET_INFO_SCHEMA.clone(), diff --git a/src/repl.rs b/src/repl.rs index 665fae7..2e83b03 100644 --- a/src/repl.rs +++ b/src/repl.rs @@ -58,7 +58,9 @@ fn format_banner(version: &str, vendor: &database::VendorInfo) -> String { if let Some(name) = &vendor.name { banner.push('\n'); match &vendor.version { - Some(vendor_version) => banner.push_str(&format!("Connected to {name} {vendor_version}")), + Some(vendor_version) => { + banner.push_str(&format!("Connected to {name} {vendor_version}")) + } None => banner.push_str(&format!("Connected to {name}")), } }