diff --git a/Cargo.lock b/Cargo.lock index a9d1ddb..0c639b0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2395,6 +2395,7 @@ dependencies = [ "parquet 23.0.0", "polars", "prql-compiler", + "regex", "tokio", "url", ] diff --git a/Cargo.toml b/Cargo.toml index 3ec09f9..21d64be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,16 +20,17 @@ dotenvy = "0.15.3" duckdb = { version = "0.5.1", features = ["bundled", "modern-full"], optional = true } env_logger = "0.9.0" log = "0.4.17" -parquet = { version = "23" , optional = true} +parquet = { version = "23" , optional = true } polars = { version = "0.24.2", optional = true, features = ["docs-selection"] } prql-compiler = { version = "0.2.8" } +regex = { version = "1.6.0", optional = true } tokio = { version = "1.19", features = ["rt-multi-thread", "macros"] } url = "2" [features] #default = ["datafusion"] default = ["datafusion", "duckdb"] -datafusion = ["dep:datafusion"] +datafusion = ["dep:datafusion", "dep:regex"] duckdb = ["dep:duckdb", "dep:arrow", "dep:parquet"] polars = ["dep:polars"] diff --git a/README.md b/README.md index 79a67c3..0baa280 100644 --- a/README.md +++ b/README.md @@ -217,12 +217,13 @@ One noteworthy limitation of this approach is that you can only query tables in the postgres database and not views. By default you will be connected to the "public" schema and can reference tables -there within your query. If you want to query tables from another schema -then you currently have to reference these through aliased `--from` parameters -like so: +there within your query. You can specify a different schema to connect to using +the "?currentSchema=schema" paramter. If you want to query tables from another schema +outside of that then you currently have to reference these through aliased +`--from` parameters like so: - $ pq -d postgresql://username:password@host:port/database \ - --from alias=schema.table 'from alias | take 10' + $ pq -d postgresql://username:password@host:port/database?currentSchema=schema \ + --from alias=other_schema.table 'from alias | take 10' ### Environment Variables diff --git a/TODO.md b/TODO.md index 3a159c1..47cc813 100644 --- a/TODO.md +++ b/TODO.md @@ -23,7 +23,7 @@ * [x] Add support for DuckDB database files * [x] Add support for PostgreSQL through DuckDB * [x] Add support for SQLite through DuckDB -* [ ] Add support for ?currentSchema=schema option postgres URI +* [x] Add support for ?currentSchema=schema option postgres URI * [ ] Add tests * [ ] Publish to crates.io * [ ] Publish to homebrew diff --git a/src/backends/duckdb.rs b/src/backends/duckdb.rs index d32e20d..091b1d1 100644 --- a/src/backends/duckdb.rs +++ b/src/backends/duckdb.rs @@ -13,6 +13,7 @@ use duckdb::{ types::{FromSql, ValueRef}, Connection, }; +use regex::Regex; use crate::{get_dest_from_to, get_sql_from_query, OutputFormat, OutputWriter, SourcesType}; @@ -63,8 +64,12 @@ pub fn query( if parts.len() == 1 { parts.insert(0, "public"); } - let table = parts.pop().ok_or(anyhow!("Couldn't extract table name from {source}."))?; - let schema = parts.pop().ok_or(anyhow!("Couldn't extract schema name from {source}."))?; + let table = parts + .pop() + .ok_or(anyhow!("Couldn't extract table name from {source}."))?; + let schema = parts + .pop() + .ok_or(anyhow!("Couldn't extract schema name from {source}."))?; format!("postgres_scan('{database}', '{schema}', '{table}')") } else { format!(r#"'{source}'"#) @@ -90,10 +95,26 @@ pub fn query( con } else if database.starts_with("postgres") { let con = Connection::open_in_memory()?; + // Check if a schema was specified + let re = Regex::new(r"^(?P[^?]+)(?P\?currentSchema=.+)?$")?; + let caps = re + .captures(database) + .ok_or(anyhow!("Couldn't match regex!"))?; + let uri = caps + .name("uri") + .ok_or(anyhow!("Couldn't extract URI!"))? + .as_str(); + debug!("uri={:?}", uri); + let schema_param = caps + .name("schema") + .map_or("?currentSchema=public", |p| p.as_str()); + let schema = schema_param.split("=").last().map_or("public", |p| p); + debug!("schema={:?}", schema); // Install and load the postgres_scanner extension let load_extension = "INSTALL postgres_scanner; LOAD postgres_scanner;"; con.execute_batch(load_extension)?; - let attach_sql = format!("CALL postgres_attach('{database}')"); + let attach_sql = format!("CALL postgres_attach('{uri}', source_schema='{schema}')"); + debug!("attach_sql={:?}", attach_sql); con.execute_batch(&attach_sql)?; con } else {