Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,4 +25,5 @@ syntect = "5.3.0"
terminal-colorsaurus = "1.0.3"

[dev-dependencies]
arrow-buffer = "58.1.0"
tempfile = "3.27.0"
178 changes: 175 additions & 3 deletions src/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
pub version: Option<String>,
}

/// 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<RecordBatch> = match reader.collect::<Result<_, _>>() {
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::<arrow_array::types::UInt32Type>())
else {
continue;
};
let Some(info_values) = batch
.column_by_name("info_value")
.and_then(|col| col.as_any().downcast_ref::<UnionArray>())
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::<i32>() 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<ManagedConnection, String> {
match source {
Expand Down Expand Up @@ -193,6 +279,92 @@ 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<u32> = rows.iter().map(|(name, _)| *name).collect();
let string_values: Vec<Option<&str>> = 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::<Vec<_>>();

let type_ids = ScalarBuffer::from(vec![0i8; rows.len()]);
let offsets = ScalarBuffer::from((0..rows.len() as i32).collect::<Vec<_>>());
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() {
Expand Down
66 changes: 66 additions & 0 deletions src/repl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,27 @@ 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));
Expand Down Expand Up @@ -93,3 +113,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");
}
}
Loading