diff --git a/columnq/src/table/database.rs b/columnq/src/table/database.rs index 9c2932d..3fc6a40 100644 --- a/columnq/src/table/database.rs +++ b/columnq/src/table/database.rs @@ -10,6 +10,7 @@ pub enum DatabaseLoader { feature = "database-postgres" ))] mod imp { + use crate::table::TableLoadOption; use crate::table::{self, TableSource}; use connectorx::prelude::*; use log::debug; @@ -39,7 +40,16 @@ mod imp { t: &TableSource, ) -> Result { debug!("loading database table data..."); - let queries = CXQuery::naked(format!("SELECT * FROM {}", t.name)); + + let table_name = match &t.option { + Some(TableLoadOption::mysql { table }) => table.clone(), + Some(TableLoadOption::postgres { table }) => table.clone(), + Some(TableLoadOption::sqlite { table }) => table.clone(), + _ => None, + } + .unwrap_or(t.name.clone()); + + let queries = CXQuery::naked(format!("SELECT * FROM {}", table_name)); let source = SourceConn::try_from(t.get_uri_str()) .context(SourceSnafu) .context(table::LoadDatabaseSnafu)?; diff --git a/columnq/src/table/mod.rs b/columnq/src/table/mod.rs index 3926bee..49798c9 100644 --- a/columnq/src/table/mod.rs +++ b/columnq/src/table/mod.rs @@ -323,9 +323,15 @@ pub enum TableLoadOption { delta(TableOptionDelta), arrow {}, arrows {}, - mysql {}, - sqlite {}, - postgres {}, + mysql { + table: Option, + }, + sqlite { + table: Option, + }, + postgres { + table: Option, + }, } impl TableLoadOption { @@ -420,6 +426,10 @@ impl std::fmt::Display for TableIoSource { } } +fn table_name_from_path(path: &uriparse::Path) -> Option { + Some(path.segments()[0].to_string()) +} + #[derive(Deserialize, Clone, Debug, Eq, PartialEq)] #[serde(deny_unknown_fields)] pub struct TableSource { @@ -542,9 +552,16 @@ impl TableSource { let uri = URIReference::try_from(uri.as_str()).ok()?; let scheme = uri.scheme()?; match scheme.as_str() { - "mysql" => Some(TableLoadOption::mysql {}), - "sqlite" => Some(TableLoadOption::sqlite {}), - "postgresql" => Some(TableLoadOption::postgres {}), + "mysql" => Some(TableLoadOption::mysql { + table: table_name_from_path(uri.path()), + }), + "sqlite" => Some(TableLoadOption::sqlite { + // for sqlite, db uri only contains the path to the db file + table: None, + }), + "postgresql" => Some(TableLoadOption::postgres { + table: table_name_from_path(uri.path()), + }), _ => None, } } @@ -595,9 +612,9 @@ impl TableSource { None => { // database sources doesn't have suffix extension, parse scheme instead match TableSource::parse_option(&self.io_source) { - Some(TableLoadOption::mysql {}) => "mysql", - Some(TableLoadOption::sqlite {}) => "sqlite", - Some(TableLoadOption::postgres {}) => "postgres", + Some(TableLoadOption::mysql { .. }) => "mysql", + Some(TableLoadOption::sqlite { .. }) => "sqlite", + Some(TableLoadOption::postgres { .. }) => "postgres", _ => { return Err(Error::Extension { msg: format!("unsupported extension in uri: {uri}"), @@ -961,6 +978,18 @@ schema: ); } + #[test] + fn test_table_name_from_path() { + assert_eq!( + table_name_from_path( + &URIReference::try_from("mysql://root:123456@1.1.1.1:3306/test") + .unwrap() + .path() + ), + Some("test".to_string()), + ); + } + #[cfg(feature = "database-sqlite")] #[tokio::test] async fn test_load_sqlite_table() {