feat: show info on repl startup

Author: esadek

Cargo.lock

@@ -25,4 +25,5 @@ syntect = "5.3.0"
 terminal-colorsaurus = "1.0.3"
 
 [dev-dependencies]
+arrow-buffer = "58.1.0"
 tempfile = "3.27.0"

Cargo.toml

@@ -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 {
@@ -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<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() {

src/database.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");
+    }
+}