Adds support for connecting to schemas other than public

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"]
 

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
 

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

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<uri>[^?]+)(?P<schema>\?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 {