diff --git a/README.md b/README.md index 961fb4d..1865fae 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,30 @@ Execute SQL queries: └─────────┴───────────┴────────────────┴───────────────┴───────────────────┴─────────────┴────────┴──────┘ ``` +Inspect the database with colon-prefixed commands instead of SQL: + +``` +> :get-objects warehouse.main.% +┌───────────┬───────────┬──────────┬────────────┐ +│ catalog │ db_schema │ table │ table_type │ +├───────────┼───────────┼──────────┼────────────┤ +│ warehouse │ main │ adelie │ VIEW │ +│ warehouse │ main │ penguins │ BASE TABLE │ +└───────────┴───────────┴──────────┴────────────┘ +> :get-schema penguins +┌────────────────┬─────────┬──────────┐ +│ column │ type │ nullable │ +├────────────────┼─────────┼──────────┤ +│ species │ Utf8 │ true │ +│ island │ Utf8 │ true │ +│ bill_length_mm │ Float64 │ true │ +│ body_mass_g │ Int32 │ true │ +│ year │ Int32 │ true │ +└────────────────┴─────────┴──────────┘ +``` + +Type `:help` for the full list of commands. + ### Non-interactive Usage Execute a query directly and exit: diff --git a/docs/reference.md b/docs/reference.md index 25c203a..d7b7a1f 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -114,6 +114,50 @@ The output format is inferred from the file extension: | `.csv` | CSV | | `.arrow`, `.ipc`| Arrow IPC | +## Commands + +A line starting with `:` is a command instead of a SQL query. Commands work in the interactive shell and with [`--query`](/reference/#-query), [`--file`](/reference/#-file) and standard input. + +| Command | Alias | Description | +|-----------------------------------------|-------|-----------------------------------| +| `:get-objects []` | `:go` | List catalogs, schemas and tables | +| `:get-schema ` | `:gs` | Show the columns of a table | +| `:help` | `:h` | List the commands | +| `:quit` | `:q` | Exit databow | + +`:get-objects` calls the ADBC `GetObjects` method, so it works the same way on every driver. Without an identifier it lists every catalog, schema and table the connection exposes. The parts of the identifier are ADBC search patterns, where `%` and `_` are wildcards: + +```console +> :get-objects warehouse.main.% +┌───────────┬───────────┬──────────┬────────────┐ +│ catalog │ db_schema │ table │ table_type │ +├───────────┼───────────┼──────────┼────────────┤ +│ warehouse │ main │ adelie │ VIEW │ +│ warehouse │ main │ penguins │ BASE TABLE │ +└───────────┴───────────┴──────────┴────────────┘ +``` + +`:get-schema` calls `GetTableSchema` and shows the Arrow schema of one table. The table name must match exactly: + +```console +> :get-schema penguins +┌────────────────┬─────────┬──────────┐ +│ column │ type │ nullable │ +├────────────────┼─────────┼──────────┤ +│ species │ Utf8 │ true │ +│ island │ Utf8 │ true │ +│ bill_length_mm │ Float64 │ true │ +│ body_mass_g │ Int32 │ true │ +│ year │ Int32 │ true │ +└────────────────┴─────────┴──────────┘ +``` + +Both commands produce a table like any query, so [`--mode`](/reference/#-mode) and [`--output`](/reference/#-output) work: + +```sh +databow --profile warehouse --query ":get-objects" --output objects.json +``` + ## --help Print the help message diff --git a/docs/tutorial.md b/docs/tutorial.md index 59ec6ae..74068d8 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -108,6 +108,33 @@ $ databow --profile warehouse --mode ascii-markdown | Adelie | 152 | ``` +### Commands + +A line starting with `:` is a command instead of a query. `:get-objects` lists catalogs, schemas and tables. `:get-schema` shows the columns of one table. + +```console +$ databow --profile warehouse +> :get-objects warehouse.main.% +┌───────────┬───────────┬──────────┬────────────┐ +│ catalog │ db_schema │ table │ table_type │ +├───────────┼───────────┼──────────┼────────────┤ +│ warehouse │ main │ adelie │ VIEW │ +│ warehouse │ main │ penguins │ BASE TABLE │ +└───────────┴───────────┴──────────┴────────────┘ +> :get-schema penguins +┌────────────────┬─────────┬──────────┐ +│ column │ type │ nullable │ +├────────────────┼─────────┼──────────┤ +│ species │ Utf8 │ true │ +│ island │ Utf8 │ true │ +│ bill_length_mm │ Float64 │ true │ +│ body_mass_g │ Int32 │ true │ +│ year │ Int32 │ true │ +└────────────────┴─────────┴──────────┘ +``` + +Type `:help` for the full list, and `:quit` to exit. The [commands reference](/reference/#commands) documents each one. + ## Non-interactive Usage The [`--query` argument](/reference/#-query) can be used to execute a query and exit: diff --git a/src/command.rs b/src/command.rs new file mode 100644 index 0000000..6bd9b8e --- /dev/null +++ b/src/command.rs @@ -0,0 +1,297 @@ +use crate::database; +use adbc_core::Connection; +use arrow_array::RecordBatch; + +pub const HELP: &str = "\ +Commands: + Execute a SQL query + :get-objects, :go [] List catalogs, schemas and tables, optionally + filtered by an identifier. Filters are ADBC + search patterns, where % and _ are wildcards. + :get-schema, :gs Show the columns of a table. The table name + must match exactly. + :help, :h Show this message + :quit, :q Exit databow (Ctrl-D also works)"; + +#[derive(Debug, PartialEq)] +pub enum Command { + Query(String), + GetObjects { + catalog: Option, + db_schema: Option, + table: Option, + }, + GetSchema { + catalog: Option, + db_schema: Option, + table: String, + }, + Help, + Quit, +} + +pub fn parse(line: &str) -> Result { + let trimmed = line.trim_end().trim_end_matches(';'); + + let Some(rest) = trimmed.trim_start().strip_prefix(':') else { + return Ok(Command::Query(trimmed.to_string())); + }; + + let (name, argument) = match rest.split_once(char::is_whitespace) { + Some((name, argument)) => (name, argument.trim()), + None => (rest, ""), + }; + + if argument.split_whitespace().count() > 1 { + return Err(format!( + "Command ':{name}' takes a single identifier, got '{argument}'" + )); + } + + match name { + "help" | "h" => Ok(Command::Help), + "quit" | "q" => Ok(Command::Quit), + "get-objects" | "go" => { + let (catalog, db_schema, table) = parse_identifier(argument)?; + Ok(Command::GetObjects { + catalog, + db_schema, + table, + }) + } + "get-schema" | "gs" => { + let (catalog, db_schema, table) = parse_identifier(argument)?; + let Some(table) = table else { + return Err( + "Command ':get-schema' requires a table name, e.g. ':get-schema my_table'" + .to_string(), + ); + }; + Ok(Command::GetSchema { + catalog, + db_schema, + table, + }) + } + _ => Err(format!( + "Unknown command ':{name}'. Type :help for a list of commands." + )), + } +} + +type Identifier = (Option, Option, Option); + +fn parse_identifier(argument: &str) -> Result { + if argument.is_empty() { + return Ok((None, None, None)); + } + + let parts: Vec<&str> = argument.split('.').collect(); + let (catalog, db_schema, table) = match parts.as_slice() { + [table] => ("", "", *table), + [db_schema, table] => ("", *db_schema, *table), + [catalog, db_schema, table] => (*catalog, *db_schema, *table), + _ => { + return Err(format!( + "Invalid identifier '{argument}': expected [catalog.][schema.]table. Identifiers containing '.' are not supported." + )); + } + }; + + Ok((optional(catalog), optional(db_schema), optional(table))) +} + +fn optional(part: &str) -> Option { + if part.is_empty() { + None + } else { + Some(part.to_string()) + } +} + +pub fn run(connection: &mut impl Connection, command: Command) -> Result, String> { + match command { + Command::Query(sql) => database::execute_query(connection, &sql), + Command::GetObjects { + catalog, + db_schema, + table, + } => database::get_objects( + connection, + catalog.as_deref(), + db_schema.as_deref(), + table.as_deref(), + ), + Command::GetSchema { + catalog, + db_schema, + table, + } => { + database::get_table_schema(connection, catalog.as_deref(), db_schema.as_deref(), &table) + } + Command::Help | Command::Quit => Ok(Vec::new()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn objects(catalog: Option<&str>, db_schema: Option<&str>, table: Option<&str>) -> Command { + Command::GetObjects { + catalog: catalog.map(str::to_string), + db_schema: db_schema.map(str::to_string), + table: table.map(str::to_string), + } + } + + #[test] + fn test_parse_query() { + assert_eq!( + parse("SELECT 1;").unwrap(), + Command::Query("SELECT 1".to_string()) + ); + } + + #[test] + fn test_parse_query_with_cast_operator() { + assert_eq!( + parse("SELECT 1::int").unwrap(), + Command::Query("SELECT 1::int".to_string()) + ); + } + + #[test] + fn test_parse_query_with_leading_whitespace() { + assert_eq!( + parse(" SELECT 1;").unwrap(), + Command::Query(" SELECT 1".to_string()) + ); + } + + #[test] + fn test_parse_commented_out_command_is_a_query() { + assert_eq!( + parse("-- :help").unwrap(), + Command::Query("-- :help".to_string()) + ); + } + + #[test] + fn test_parse_help_aliases() { + assert_eq!(parse(":help").unwrap(), Command::Help); + assert_eq!(parse(":h").unwrap(), Command::Help); + assert_eq!(parse(" :help ").unwrap(), Command::Help); + assert_eq!(parse(":help;").unwrap(), Command::Help); + } + + #[test] + fn test_parse_quit_aliases() { + assert_eq!(parse(":quit").unwrap(), Command::Quit); + assert_eq!(parse(":q").unwrap(), Command::Quit); + } + + #[test] + fn test_parse_get_objects_without_identifier() { + assert_eq!(parse(":get-objects").unwrap(), objects(None, None, None)); + assert_eq!(parse(":go").unwrap(), objects(None, None, None)); + } + + #[test] + fn test_parse_get_objects_identifier_arities() { + assert_eq!(parse(":go t").unwrap(), objects(None, None, Some("t"))); + assert_eq!( + parse(":go s.t").unwrap(), + objects(None, Some("s"), Some("t")) + ); + assert_eq!( + parse(":go c.s.t").unwrap(), + objects(Some("c"), Some("s"), Some("t")) + ); + } + + #[test] + fn test_parse_get_objects_empty_segments_are_none() { + assert_eq!( + parse(":go .s.t").unwrap(), + objects(None, Some("s"), Some("t")) + ); + assert_eq!( + parse(":go c..t").unwrap(), + objects(Some("c"), None, Some("t")) + ); + assert_eq!(parse(":go s.").unwrap(), objects(None, Some("s"), None)); + } + + #[test] + fn test_parse_get_objects_too_many_parts() { + let err = parse(":go a.b.c.d").unwrap_err(); + assert!(err.contains("a.b.c.d"), "{err}"); + } + + #[test] + fn test_parse_get_schema_identifier_arities() { + assert_eq!( + parse(":get-schema t").unwrap(), + Command::GetSchema { + catalog: None, + db_schema: None, + table: "t".to_string(), + } + ); + assert_eq!( + parse(":gs c.s.t").unwrap(), + Command::GetSchema { + catalog: Some("c".to_string()), + db_schema: Some("s".to_string()), + table: "t".to_string(), + } + ); + } + + #[test] + fn test_parse_get_schema_requires_a_table() { + assert!( + parse(":get-schema") + .unwrap_err() + .contains("requires a table") + ); + assert!(parse(":gs").unwrap_err().contains("requires a table")); + assert!(parse(":gs s.").unwrap_err().contains("requires a table")); + } + + #[test] + fn test_parse_unknown_command() { + let err = parse(":nope").unwrap_err(); + assert!(err.contains(":nope"), "{err}"); + assert!(err.contains(":help"), "{err}"); + } + + #[test] + fn test_parse_does_not_prefix_match_aliases() { + assert!(parse(":g").is_err()); + assert!(parse(":get").is_err()); + } + + #[test] + fn test_parse_rejects_multiple_arguments() { + let err = parse(":go a b").unwrap_err(); + assert!(err.contains(":go"), "{err}"); + } + + #[test] + fn test_help_lists_every_command_and_alias() { + for name in [ + ":get-objects", + ":go", + ":get-schema", + ":gs", + ":help", + ":h", + ":quit", + ":q", + ] { + assert!(HELP.contains(name), "help is missing {name}"); + } + } +} diff --git a/src/database.rs b/src/database.rs index 9dd5c4f..e217236 100644 --- a/src/database.rs +++ b/src/database.rs @@ -2,15 +2,18 @@ // SPDX-License-Identifier: Apache-2.0 use crate::cli::ConnectionSource; -use adbc_core::options::{AdbcVersion, InfoCode, OptionDatabase, OptionValue}; +use adbc_core::options::{AdbcVersion, InfoCode, ObjectDepth, 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, ManagedDatabase, ManagedDriver}; use arrow_array::cast::AsArray; -use arrow_array::{Array, RecordBatch, UnionArray}; +use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, StringArray, UnionArray}; +use arrow_cast::display::array_value_to_string; +use arrow_schema::{DataType, Field, Schema}; use std::collections::HashSet; +use std::sync::Arc; /// Vendor (database product) metadata reported by an ADBC driver. #[derive(Debug, Default)] @@ -308,11 +311,128 @@ fn merge_options( 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 arrow_array::{ListArray, StringArray, StructArray, UInt32Array, UnionArray}; + use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, TimeUnit}; use std::sync::Arc; + type TableFixture<'a> = (&'a str, &'a str); + type SchemaFixture<'a> = (Option<&'a str>, Option>>); + type CatalogFixture<'a> = (Option<&'a str>, Option>>); + + fn make_objects_batch(catalogs: &[CatalogFixture]) -> RecordBatch { + let mut catalog_names: Vec> = Vec::new(); + let mut schema_names: Vec> = Vec::new(); + let mut schema_offsets: Vec = vec![0]; + let mut schema_validity: Vec = Vec::new(); + let mut table_names: Vec<&str> = Vec::new(); + let mut table_types: Vec<&str> = Vec::new(); + let mut table_offsets: Vec = vec![0]; + let mut table_validity: Vec = Vec::new(); + + for (catalog_name, schemas) in catalogs { + catalog_names.push(*catalog_name); + let Some(schemas) = schemas else { + schema_validity.push(false); + schema_offsets.push(schema_names.len() as i32); + continue; + }; + schema_validity.push(true); + for (schema_name, tables) in schemas { + schema_names.push(*schema_name); + let Some(tables) = tables else { + table_validity.push(false); + table_offsets.push(table_names.len() as i32); + continue; + }; + table_validity.push(true); + for (table_name, table_type) in tables { + table_names.push(table_name); + table_types.push(table_type); + } + table_offsets.push(table_names.len() as i32); + } + schema_offsets.push(schema_names.len() as i32); + } + + let tables = StructArray::from(vec![ + ( + Arc::new(Field::new("table_name", DataType::Utf8, true)), + Arc::new(StringArray::from(table_names)) as ArrayRef, + ), + ( + Arc::new(Field::new("table_type", DataType::Utf8, true)), + Arc::new(StringArray::from(table_types)) as ArrayRef, + ), + ]); + let table_lists = ListArray::new( + Arc::new(Field::new("l", tables.data_type().clone(), true)), + OffsetBuffer::new(table_offsets.into()), + Arc::new(tables), + Some(NullBuffer::from(table_validity)), + ); + + let schemas = StructArray::from(vec![ + ( + Arc::new(Field::new("db_schema_name", DataType::Utf8, true)), + Arc::new(StringArray::from(schema_names)) as ArrayRef, + ), + ( + Arc::new(Field::new( + "db_schema_tables", + table_lists.data_type().clone(), + true, + )), + Arc::new(table_lists) as ArrayRef, + ), + ]); + let schema_lists = ListArray::new( + Arc::new(Field::new("l", schemas.data_type().clone(), true)), + OffsetBuffer::new(schema_offsets.into()), + Arc::new(schemas), + Some(NullBuffer::from(schema_validity)), + ); + + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("catalog_name", DataType::Utf8, true), + Field::new("catalog_db_schemas", schema_lists.data_type().clone(), true), + ])), + vec![ + Arc::new(StringArray::from(catalog_names)), + Arc::new(schema_lists), + ], + ) + .unwrap() + } + + fn rendered_rows(batch: &RecordBatch) -> Vec { + (0..batch.num_rows()) + .map(|row| { + (0..batch.num_columns()) + .map(|column_index| { + let column = batch.column(column_index); + if column.is_null(row) { + "NULL".to_string() + } else { + array_value_to_string(column, row).unwrap() + } + }) + .collect::>() + .join(" | ") + }) + .collect() + } + + fn column_names(batch: &RecordBatch) -> Vec { + batch + .schema() + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect() + } + /// 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. @@ -486,6 +606,162 @@ mod tests { let merged = merge_options(base, overrides); assert_eq!(merged.len(), 2); } + + #[test] + fn test_flatten_objects_rows_are_sorted() { + let batch = make_objects_batch(&[ + ( + Some("zeta"), + Some(vec![(Some("main"), Some(vec![("t1", "BASE TABLE")]))]), + ), + ( + Some("alpha"), + Some(vec![ + (Some("s2"), Some(vec![("b", "VIEW"), ("a", "BASE TABLE")])), + (Some("s1"), Some(vec![])), + ]), + ), + ]); + + assert_eq!( + rendered_rows(&flatten_objects(&[batch]).unwrap()), + vec![ + "alpha | s1 | NULL | NULL", + "alpha | s2 | a | BASE TABLE", + "alpha | s2 | b | VIEW", + "zeta | main | t1 | BASE TABLE", + ] + ); + } + + #[test] + fn test_flatten_objects_null_children() { + let batch = make_objects_batch(&[ + (Some("c1"), None), + (Some("c2"), Some(vec![(Some("s"), None)])), + ( + None, + Some(vec![(None, Some(vec![("orphan", "BASE TABLE")]))]), + ), + ]); + + assert_eq!( + rendered_rows(&flatten_objects(&[batch]).unwrap()), + vec![ + "NULL | NULL | orphan | BASE TABLE", + "c1 | NULL | NULL | NULL", + "c2 | s | NULL | NULL", + ] + ); + } + + #[test] + fn test_flatten_objects_multiple_batches() { + let first = make_objects_batch(&[( + Some("b"), + Some(vec![(Some("main"), Some(vec![("t", "BASE TABLE")]))]), + )]); + let second = make_objects_batch(&[( + Some("a"), + Some(vec![(Some("main"), Some(vec![("t", "VIEW")]))]), + )]); + + assert_eq!( + rendered_rows(&flatten_objects(&[first, second]).unwrap()), + vec!["a | main | t | VIEW", "b | main | t | BASE TABLE"] + ); + } + + #[test] + fn test_flatten_objects_empty() { + let batch = flatten_objects(&[]).unwrap(); + assert_eq!(batch.num_rows(), 0); + assert_eq!( + column_names(&batch), + vec!["catalog", "db_schema", "table", "table_type"] + ); + } + + #[test] + fn test_flatten_objects_without_db_schemas_column() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "catalog_name", + DataType::Utf8, + true, + )])), + vec![Arc::new(StringArray::from(vec![Some("only")]))], + ) + .unwrap(); + + assert_eq!( + rendered_rows(&flatten_objects(&[batch]).unwrap()), + vec!["only | NULL | NULL | NULL"] + ); + } + + #[test] + fn test_flatten_objects_rejects_unexpected_type() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("catalog_name", DataType::Utf8, true), + Field::new("catalog_db_schemas", DataType::Utf8, true), + ])), + vec![ + Arc::new(StringArray::from(vec![Some("c")])), + Arc::new(StringArray::from(vec![Some("not a list")])), + ], + ) + .unwrap(); + + let err = flatten_objects(&[batch]).unwrap_err(); + assert!(err.contains("catalog_db_schemas"), "{err}"); + } + + #[test] + fn test_schema_to_batch() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + let batch = schema_to_batch(&schema).unwrap(); + + assert_eq!(column_names(&batch), vec!["column", "type", "nullable"]); + assert_eq!( + rendered_rows(&batch), + vec!["a | Int32 | false", "b | Utf8 | true"] + ); + } + + #[test] + fn test_schema_to_batch_type_rendering() { + let schema = Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + Field::new("amount", DataType::Decimal128(10, 2), true), + Field::new("tags", DataType::new_list(DataType::Utf8, true), true), + ]); + let batch = schema_to_batch(&schema).unwrap(); + + assert_eq!( + rendered_rows(&batch), + vec![ + "ts | Timestamp(µs, \"UTC\") | true", + "amount | Decimal128(10, 2) | true", + "tags | List(Utf8) | true", + ] + ); + } + + #[test] + fn test_schema_to_batch_empty_schema() { + let batch = schema_to_batch(&Schema::empty()).unwrap(); + assert_eq!(batch.num_rows(), 0); + assert_eq!(column_names(&batch), vec!["column", "type", "nullable"]); + } } pub fn execute_query( @@ -514,3 +790,209 @@ pub fn execute_query( Ok(batches) } + +pub fn get_objects( + connection: &impl Connection, + catalog: Option<&str>, + db_schema: Option<&str>, + table_name: Option<&str>, +) -> Result, String> { + let reader = connection + .get_objects( + ObjectDepth::Tables, + catalog, + db_schema, + table_name, + None, + None, + ) + .map_err(|e| format!("Failed to get objects: {e}"))?; + + let batches: Vec = reader + .collect::>() + .map_err(|e| format!("Failed to collect objects: {e}"))?; + + Ok(vec![flatten_objects(&batches)?]) +} + +pub fn get_table_schema( + connection: &impl Connection, + catalog: Option<&str>, + db_schema: Option<&str>, + table_name: &str, +) -> Result, String> { + let schema = connection + .get_table_schema(catalog, db_schema, table_name) + .map_err(|e| format!("Failed to get table schema: {e}"))?; + + Ok(vec![schema_to_batch(&schema)?]) +} + +type ObjectRow = ( + Option, + Option, + Option, + Option, +); + +fn flatten_objects(batches: &[RecordBatch]) -> Result { + let mut rows: Vec = Vec::new(); + + for batch in batches { + let catalog_names = batch.column_by_name("catalog_name"); + let schema_lists = batch.column_by_name("catalog_db_schemas"); + + for row in 0..batch.num_rows() { + let catalog = optional_string(catalog_names, row)?; + + let Some(schemas) = list_element(schema_lists, row, "catalog_db_schemas")? else { + rows.push((catalog, None, None, None)); + continue; + }; + let Some(schemas) = schemas.as_struct_opt() else { + return Err(unexpected_type("catalog_db_schemas", schemas.data_type())); + }; + + let schema_names = schemas.column_by_name("db_schema_name"); + let table_lists = schemas.column_by_name("db_schema_tables"); + + for schema_index in 0..schemas.len() { + if schemas.is_null(schema_index) { + rows.push((catalog.clone(), None, None, None)); + continue; + } + let db_schema = optional_string(schema_names, schema_index)?; + + let Some(tables) = list_element(table_lists, schema_index, "db_schema_tables")? + else { + rows.push((catalog.clone(), db_schema, None, None)); + continue; + }; + let Some(tables) = tables.as_struct_opt() else { + return Err(unexpected_type("db_schema_tables", tables.data_type())); + }; + + let table_names = tables.column_by_name("table_name"); + let table_types = tables.column_by_name("table_type"); + + for table_index in 0..tables.len() { + if tables.is_null(table_index) { + continue; + } + rows.push(( + catalog.clone(), + db_schema.clone(), + optional_string(table_names, table_index)?, + optional_string(table_types, table_index)?, + )); + } + } + } + } + + rows.sort_by(|left, right| (&left.0, &left.1, &left.2).cmp(&(&right.0, &right.1, &right.2))); + + let mut catalogs = Vec::with_capacity(rows.len()); + let mut db_schemas = Vec::with_capacity(rows.len()); + let mut tables = Vec::with_capacity(rows.len()); + let mut table_types = Vec::with_capacity(rows.len()); + for (catalog, db_schema, table, table_type) in rows { + catalogs.push(catalog); + db_schemas.push(db_schema); + tables.push(table); + table_types.push(table_type); + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("catalog", DataType::Utf8, true), + Field::new("db_schema", DataType::Utf8, true), + Field::new("table", DataType::Utf8, true), + Field::new("table_type", DataType::Utf8, true), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(catalogs)), + Arc::new(StringArray::from(db_schemas)), + Arc::new(StringArray::from(tables)), + Arc::new(StringArray::from(table_types)), + ], + ) + .map_err(|e| format!("Failed to build the object table: {e}")) +} + +fn schema_to_batch(table_schema: &Schema) -> Result { + let field_count = table_schema.fields().len(); + let mut names = Vec::with_capacity(field_count); + let mut types = Vec::with_capacity(field_count); + let mut nullable = Vec::with_capacity(field_count); + + for field in table_schema.fields() { + names.push(field.name().to_string()); + types.push(field.data_type().to_string()); + nullable.push(field.is_nullable()); + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("column", DataType::Utf8, false), + Field::new("type", DataType::Utf8, false), + Field::new("nullable", DataType::Boolean, false), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(types)), + Arc::new(BooleanArray::from(nullable)), + ], + ) + .map_err(|e| format!("Failed to build the schema table: {e}")) +} + +fn optional_string(column: Option<&ArrayRef>, index: usize) -> Result, String> { + let Some(column) = column else { + return Ok(None); + }; + if column.is_null(index) { + return Ok(None); + } + array_value_to_string(column, index) + .map(Some) + .map_err(|e| format!("Failed to read object metadata: {e}")) +} + +fn list_element( + column: Option<&ArrayRef>, + index: usize, + name: &str, +) -> Result, String> { + let Some(column) = column else { + return Ok(None); + }; + + let element = if let Some(list) = column.as_list_opt::() { + if list.is_null(index) { + return Ok(None); + } + list.value(index) + } else if let Some(list) = column.as_list_opt::() { + if list.is_null(index) { + return Ok(None); + } + list.value(index) + } else { + return Err(unexpected_type(name, column.data_type())); + }; + + if element.is_empty() { + Ok(None) + } else { + Ok(Some(element)) + } +} + +fn unexpected_type(name: &str, data_type: &DataType) -> String { + format!("Unexpected type for column '{name}': {data_type}") +} diff --git a/src/highlighter.rs b/src/highlighter.rs index 56c54ad..1f3ad7f 100644 --- a/src/highlighter.rs +++ b/src/highlighter.rs @@ -53,6 +53,11 @@ impl Highlighter for SyntectHighlighter { fn highlight(&self, line: &str, _cursor: usize) -> StyledText { let mut styled = StyledText::new(); + if line.trim_start().starts_with(':') { + styled.push((Style::new(), line.to_string())); + return styled; + } + let syntax = self .syntax_set .find_syntax_by_extension("sql") diff --git a/src/main.rs b/src/main.rs index 90a1aed..5d936b1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 mod cli; +mod command; mod database; mod highlighter; mod output; @@ -48,17 +49,39 @@ fn main() { QuerySource::Interactive => unreachable!(), }; - let mut connection = database::initialize_connection(connection).unwrap_or_else(|e| { - eprintln!("{e}"); + let trimmed = sql.trim(); + if trimmed.starts_with(':') && trimmed.contains('\n') { + eprintln!("Error: commands cannot be combined with SQL in --file or stdin input"); exit(1); - }); - let batches = database::execute_query(&mut connection, &sql).unwrap_or_else(|e| { + } + + let parsed = command::parse(&sql).unwrap_or_else(|e| { eprintln!("{e}"); exit(1); }); - if let Err(e) = output_results(&batches, table_mode, output_path.as_deref()) { - eprintln!("{e}"); - exit(1); + + match parsed { + command::Command::Help => { + println!("{}", command::HELP); + if output_path.is_some() { + eprintln!("No output file was written: ':help' produces no results"); + } + } + command::Command::Quit => {} + parsed => { + let mut connection = database::initialize_connection(connection).unwrap_or_else(|e| { + eprintln!("{e}"); + exit(1); + }); + let batches = command::run(&mut connection, parsed).unwrap_or_else(|e| { + eprintln!("{e}"); + exit(1); + }); + if let Err(e) = output_results(&batches, table_mode, output_path.as_deref()) { + eprintln!("{e}"); + exit(1); + } + } } } diff --git a/src/repl.rs b/src/repl.rs index 2e83b03..df83175 100644 --- a/src/repl.rs +++ b/src/repl.rs @@ -1,6 +1,7 @@ // Copyright 2026 Columnar Technologies Inc. // SPDX-License-Identifier: Apache-2.0 +use crate::command::{self, Command}; use crate::database; use crate::highlighter::SyntectHighlighter; use crate::table::{TableMode, print_batches}; @@ -15,7 +16,7 @@ struct SqlValidator; impl Validator for SqlValidator { fn validate(&self, line: &str) -> ValidationResult { let trimmed = line.trim(); - if trimmed.is_empty() || trimmed.ends_with(';') { + if trimmed.is_empty() || trimmed.ends_with(';') || trimmed.starts_with(':') { ValidationResult::Complete } else { ValidationResult::Incomplete @@ -87,16 +88,18 @@ pub fn run_repl(mut connection: impl Connection, table_mode: TableMode) { continue; } - let sql = buffer.trim_end().trim_end_matches(';'); - let batches = match database::execute_query(&mut connection, sql) { - Ok(batches) => batches, - Err(err) => { - eprintln!("{err}"); - continue; - } - }; - if let Err(err) = print_batches(&batches, table_mode) { - eprintln!("Failed to print batches: {err}"); + match command::parse(&buffer) { + Err(err) => eprintln!("{err}"), + Ok(Command::Help) => println!("{}", command::HELP), + Ok(Command::Quit) => break, + Ok(parsed) => match command::run(&mut connection, parsed) { + Ok(batches) => { + if let Err(err) = print_batches(&batches, table_mode) { + eprintln!("Failed to print batches: {err}"); + } + } + Err(err) => eprintln!("{err}"), + }, } } Ok(Signal::CtrlC) => { @@ -119,6 +122,22 @@ mod tests { use super::*; use crate::database::VendorInfo; + #[test] + fn test_validator_completes_commands_without_a_semicolon() { + assert!(matches!( + SqlValidator.validate(":help"), + ValidationResult::Complete + )); + assert!(matches!( + SqlValidator.validate(" :get-schema t"), + ValidationResult::Complete + )); + assert!(matches!( + SqlValidator.validate("SELECT 1"), + ValidationResult::Incomplete + )); + } + #[test] fn test_format_banner_version_only() { let vendor = VendorInfo::default(); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 302c35c..78f8368 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -479,3 +479,261 @@ fn test_timestamp_with_time_zone() { stdout ); } + +fn duckdb_command(uri: &str, query: &str) -> std::process::Output { + Command::new("cargo") + .args([ + "run", "--", "--driver", "duckdb", "--uri", uri, "--query", query, + ]) + .output() + .expect("Failed to execute command") +} + +#[test] +fn test_help_command() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--query", ":help"]) + .output() + .expect("Failed to execute command"); + + assert!( + output.status.success(), + ":help should succeed. stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!(stdout.contains(":get-objects"), "stdout: {stdout}"); + assert!(stdout.contains(":get-schema"), "stdout: {stdout}"); + assert!(stdout.contains(":quit"), "stdout: {stdout}"); +} + +#[test] +fn test_help_command_does_not_need_a_working_driver() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "nosuchdriver", "--query", ":help"]) + .output() + .expect("Failed to execute command"); + + assert!(output.status.success()); + assert!(String::from_utf8_lossy(&output.stdout).contains(":get-objects")); +} + +#[test] +fn test_quit_command() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--query", ":quit"]) + .output() + .expect("Failed to execute command"); + + assert!(output.status.success()); + assert!(String::from_utf8_lossy(&output.stdout).is_empty()); +} + +#[test] +fn test_get_objects_command() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--query", ":get-objects"]) + .output() + .expect("Failed to execute command"); + + assert!( + output.status.success(), + ":get-objects should succeed. stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + let stdout = String::from_utf8_lossy(&output.stdout); + for header in ["catalog", "db_schema", "table", "table_type"] { + assert!( + stdout.contains(header), + "missing {header}. stdout: {stdout}" + ); + } +} + +#[test] +fn test_get_objects_and_get_schema_for_a_created_table() { + let directory = tempfile::tempdir().expect("Failed to create temp dir"); + let database = directory.path().join("objects.duckdb"); + let uri = database.to_string_lossy().to_string(); + + let created = duckdb_command( + &uri, + "CREATE TABLE penguins (species TEXT, body_mass_g INT); CREATE VIEW adelie AS SELECT 1 AS x;", + ); + assert!( + created.status.success(), + "setup should succeed. stderr: {}", + String::from_utf8_lossy(&created.stderr) + ); + + let objects = duckdb_command(&uri, ":get-objects objects.main.%"); + assert!( + objects.status.success(), + ":get-objects should succeed. stderr: {}", + String::from_utf8_lossy(&objects.stderr) + ); + let stdout = String::from_utf8_lossy(&objects.stdout); + assert!(stdout.contains("penguins"), "stdout: {stdout}"); + assert!(stdout.contains("BASE TABLE"), "stdout: {stdout}"); + assert!(stdout.contains("adelie"), "stdout: {stdout}"); + assert!(stdout.contains("VIEW"), "stdout: {stdout}"); + + let schema = duckdb_command(&uri, ":get-schema penguins"); + assert!( + schema.status.success(), + ":get-schema should succeed. stderr: {}", + String::from_utf8_lossy(&schema.stderr) + ); + let stdout = String::from_utf8_lossy(&schema.stdout); + for expected in [ + "column", + "type", + "nullable", + "species", + "Utf8", + "body_mass_g", + "Int32", + ] { + assert!( + stdout.contains(expected), + "missing {expected}. stdout: {stdout}" + ); + } +} + +#[test] +fn test_get_schema_honors_table_mode() { + let directory = tempfile::tempdir().expect("Failed to create temp dir"); + let database = directory.path().join("mode.duckdb"); + let uri = database.to_string_lossy().to_string(); + + let created = duckdb_command(&uri, "CREATE TABLE t (a INT)"); + assert!(created.status.success()); + + let output = Command::new("cargo") + .args([ + "run", + "--", + "--driver", + "duckdb", + "--uri", + &uri, + "--mode", + "ascii-markdown", + "--query", + ":get-schema t", + ]) + .output() + .expect("Failed to execute command"); + + assert!(output.status.success()); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!(stdout.contains("| a "), "stdout: {stdout}"); + assert!(stdout.contains("|---"), "stdout: {stdout}"); +} + +#[test] +fn test_get_objects_writes_output_file() { + let directory = tempfile::tempdir().expect("Failed to create temp dir"); + let database = directory.path().join("output.duckdb"); + let uri = database.to_string_lossy().to_string(); + let result = directory.path().join("objects.json"); + + let created = duckdb_command(&uri, "CREATE TABLE t (a INT)"); + assert!(created.status.success()); + + let output = Command::new("cargo") + .args([ + "run", + "--", + "--driver", + "duckdb", + "--uri", + &uri, + "--query", + ":get-objects output.main.t", + "--output", + &result.to_string_lossy(), + ]) + .output() + .expect("Failed to execute command"); + + assert!( + output.status.success(), + "writing objects to a file should succeed. stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + let contents = std::fs::read_to_string(&result).expect("Failed to read output file"); + assert!(contents.contains("\"table\":\"t\""), "contents: {contents}"); + assert!( + contents.contains("\"table_type\":\"BASE TABLE\""), + "contents: {contents}" + ); +} + +#[test] +fn test_unknown_command_errors() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--query", ":nope"]) + .output() + .expect("Failed to execute command"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains(":nope"), "stderr: {stderr}"); + assert!(stderr.contains(":help"), "stderr: {stderr}"); +} + +#[test] +fn test_get_schema_without_a_table_errors() { + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--query", ":get-schema"]) + .output() + .expect("Failed to execute command"); + + assert!(!output.status.success()); + assert!( + String::from_utf8_lossy(&output.stderr).contains("requires a table name"), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); +} + +#[test] +fn test_command_mixed_with_sql_in_a_file_errors() { + use std::io::Write; + + let mut temp_file = NamedTempFile::new().expect("Failed to create temp file"); + let file_path = temp_file.path().to_string_lossy().to_string(); + temp_file + .write_all(b":get-objects\nSELECT 1;\n") + .expect("Failed to write to temp file"); + + let output = Command::new("cargo") + .args(["run", "--", "--driver", "duckdb", "--file", &file_path]) + .output() + .expect("Failed to execute command"); + + assert!(!output.status.success()); + assert!( + String::from_utf8_lossy(&output.stderr).contains("cannot be combined with SQL"), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); +} + +#[test] +fn test_command_from_stdin() { + let output = Command::new("bash") + .arg("-c") + .arg("echo ':get-objects' | cargo run -- --driver duckdb") + .output() + .expect("Failed to execute command"); + + assert!( + output.status.success(), + ":get-objects from stdin should succeed. stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("db_schema")); +}