diff --git a/src/reader/adbc.rs b/src/reader/adbc.rs index dc918ece..65505a10 100644 --- a/src/reader/adbc.rs +++ b/src/reader/adbc.rs @@ -27,9 +27,11 @@ pub struct AdbcReader { // takes &self. connection: RefCell<::ConnectionType>, dialect: crate::reader::DialectRef, - /// Registry-declared driver quirks (error-text conventions), when the - /// reader was built for a registered backend. `None` for ad-hoc drivers. - adbc_info: Option, + /// The registry entry for this backend, when the reader was built from + /// a connection string for a registered backend — source of the + /// driver quirks (error-text conventions) honored in `execute_sql` and + /// `register`. `None` for ad-hoc drivers (`new`, `with_dialect`). + entry: Option<&'static crate::reader::registry::DatabaseEntry>, registered_tables: crate::reader::RegisteredTables, // Driver-specific statement options (from `stmt.`-prefixed URI params) // applied to every statement created in execute_sql. @@ -62,7 +64,7 @@ impl AdbcReader { _database: database, connection: RefCell::new(connection), dialect, - adbc_info: None, + entry: None, registered_tables: crate::reader::RegisteredTables::new(), statement_opts: Vec::new(), }) @@ -93,7 +95,7 @@ impl AdbcReader { _database: database, connection: RefCell::new(connection), dialect, - adbc_info: None, + entry: None, registered_tables: crate::reader::RegisteredTables::new(), statement_opts: Vec::new(), }) @@ -112,11 +114,11 @@ impl AdbcReader { self } - /// Attach the registry's driver quirks for this backend, so error-text - /// conventions (DDL-without-result-set signaling, ingest schema - /// mismatches) are honored where the driver requires them. - fn with_adbc_info(mut self, info: crate::reader::registry::AdbcInfo) -> Self { - self.adbc_info = Some(info); + /// Attach the registry entry for this backend, so its driver quirks + /// (DDL-without-result-set signaling, ingest schema mismatches) are + /// honored where the driver requires them. + fn with_entry(mut self, entry: &'static crate::reader::registry::DatabaseEntry) -> Self { + self.entry = Some(entry); self } @@ -371,10 +373,10 @@ impl AdbcReader { let reader = Self::new_with_database_opts(driver, dialect, opts)? .with_statement_opts(conn.ggsql.stmt_options.clone()); - // adbc:// with a recognized short name also gets quirks. + // adbc:// with a recognized short name also resolves an entry. Ok( - match adbc_info_for_scheme(if scheme == "adbc" { body } else { scheme }) { - Some(info) => reader.with_adbc_info(info), + match crate::reader::registry::by_scheme(if scheme == "adbc" { body } else { scheme }) { + Some(entry) => reader.with_entry(entry), None => reader, }, ) @@ -546,7 +548,7 @@ where // - `ddl_retry_error` (Databricks): the query path cannot // run resultless statements at all; retry via // execute_update, the path meant for them. - let info = self.adbc_info; + let info = self.entry.and_then(|e| e.adbc); if info .and_then(|i| i.ddl_empty_result_error) .is_some_and(|needle| msg.contains(needle)) @@ -695,7 +697,8 @@ where Ok(_) => break, Err(e) => { let schema_mismatch = self - .adbc_info + .entry + .and_then(|e| e.adbc) .and_then(|i| i.ingest_schema_align_error) .is_some_and(|needle| e.to_string().contains(needle)); if attempts == 0 && schema_mismatch { diff --git a/src/reader/odbc/mod.rs b/src/reader/odbc/mod.rs index 8cb3cbdf..44af717c 100644 --- a/src/reader/odbc/mod.rs +++ b/src/reader/odbc/mod.rs @@ -6,7 +6,7 @@ #[allow(dead_code)] pub(crate) mod ffi; -mod snowflake; +pub(crate) mod snowflake; mod wrapper; use crate::reader::Reader; @@ -18,15 +18,6 @@ use ffi::*; use std::sync::Arc; use wrapper::{Connection, Statement}; -/// Detect the backend SQL dialect from the DBMS name and connection string. -/// -/// Delegates to the shared matcher in [`crate::reader::dialects`]; the -/// `Driver=` value from the ODBC connection string serves as the driver hint. -fn detect_dialect(dbms_name: Option<&str>, conn_str: &str) -> Result { - let driver = super::connection::extract_odbc_value(conn_str, "driver"); - super::registry::detect_or_err(dbms_name, driver.as_deref()) -} - /// Pull ggsql-owned keys (`Dialect=`) out of an ODBC connection /// string before connecting. `Dialect=` is ggsql's escape hatch for unknown /// backends, not a driver option — drivers reject unknown keys. @@ -82,16 +73,15 @@ impl OdbcReader { (None, None) => None, }; - if snowflake::is_snowflake(&conn_str) { - if let Some(resolved) = snowflake::resolve_connection_name(&conn_str) { - conn_str = resolved; - } - } - - if snowflake::is_snowflake(&conn_str) && !snowflake::has_token(&conn_str) { - if let Some(token) = snowflake::detect_workbench_token() { - conn_str = snowflake::inject_snowflake_token(&conn_str, &token); - } + // Backend-specific credential resolution (Snowflake Workbench OAuth + // tokens, `ConnectionName=` lookup), dispatched through the registry + // on the `Driver=` value — the DBMS name isn't known until after + // connecting, which is too late to supply credentials. + let driver = super::connection::extract_odbc_value(&conn_str, "driver"); + if let Some(apply) = super::registry::detect(None, driver.as_deref()) + .and_then(|e| e.odbc_credential_provider) + { + apply(&mut conn_str); } let connection = match Connection::connect(wrapper::odbc_env()?, &conn_str) { @@ -110,20 +100,28 @@ impl OdbcReader { }; let dbms_name = connection.dbms_name(); - - let dialect = match dialect { - Some(d) => d, - None => detect_dialect(dbms_name.as_deref(), &conn_str)?, + let driver = super::connection::extract_odbc_value(&conn_str, "driver"); + + // One registry resolution per connection: dialect and fetch behavior + // both derive from the same entry. With a pinned dialect the entry + // is best-effort (fetch quirks only); without one the entry must + // resolve so the dialect can come from it. + let (entry, dialect) = match dialect { + Some(d) => ( + super::registry::detect(dbms_name.as_deref(), driver.as_deref()), + d, + ), + None => { + let entry = + super::registry::detect_entry_or_err(dbms_name.as_deref(), driver.as_deref())?; + (Some(entry), entry.dialect()) + } }; // Fetch behavior is a per-backend capability: Oracle ODBC rejects // block cursors (`odbc_row_array_size`; HY090 at SQLFetch time, and // the failed fetch leaves the cursor unusable) and cannot convert // DECIMAL to integer C types (`odbc_numeric_as_double`). - let entry = super::registry::detect( - dbms_name.as_deref(), - super::connection::extract_odbc_value(&conn_str, "driver").as_deref(), - ); let batch_size = entry .and_then(|e| e.odbc_row_array_size) .unwrap_or(BATCH_SIZE); @@ -841,13 +839,14 @@ mod tests { #[test] fn test_detect_dialect_from_dbms_name() { - let d = detect_dialect(Some("Snowflake"), "anything").unwrap(); + let d = crate::reader::registry::detect_or_err(Some("Snowflake"), None).unwrap(); assert!(!d.sql_greatest(&["a", "b"]).is_empty()); - let d = detect_dialect(None, "Driver=Snowflake;Server=foo").unwrap(); + let d = + crate::reader::registry::detect_or_err(None, Some("Snowflake ODBC Driver")).unwrap(); assert!(!d.sql_greatest(&["a", "b"]).is_empty()); - let err = detect_dialect(None, "Driver=SomeOther;Server=localhost") + let err = crate::reader::registry::detect_or_err(None, Some("SomeOther")) .err() .unwrap() .to_string(); diff --git a/src/reader/odbc/snowflake.rs b/src/reader/odbc/snowflake.rs index 865f7038..773a86df 100644 --- a/src/reader/odbc/snowflake.rs +++ b/src/reader/odbc/snowflake.rs @@ -1,9 +1,19 @@ //! Snowflake Workbench credential detection and connection resolution. -pub(super) fn is_snowflake(conn_str: &str) -> bool { - crate::reader::connection::extract_odbc_value(conn_str, "driver") - .map(|d| d.to_lowercase().contains("snowflake")) - .unwrap_or(false) +/// Registry `odbc_credential_provider` hook for Snowflake: rewrite the +/// connection string to supply credentials before connecting. Resolves a +/// `ConnectionName=` reference from `~/.snowflake/connections.toml`, then +/// injects a Posit Workbench OAuth token when one is available and the +/// string doesn't carry a token already. +pub(crate) fn apply_workbench_credentials(conn_str: &mut String) { + if let Some(resolved) = resolve_connection_name(conn_str) { + *conn_str = resolved; + } + if !has_token(conn_str) { + if let Some(token) = detect_workbench_token() { + *conn_str = inject_snowflake_token(conn_str, &token); + } + } } pub(super) fn has_token(conn_str: &str) -> bool { @@ -169,11 +179,22 @@ mod tests { use super::*; #[test] - fn test_is_snowflake() { - assert!(is_snowflake( - "Driver=Snowflake;Server=foo.snowflakecomputing.com" - )); - assert!(!is_snowflake("Driver={PostgreSQL};Server=localhost")); + fn test_credential_hook_dispatches_on_driver() { + // The registry dispatches the credential hook on the Driver= value. + let driver = crate::reader::connection::extract_odbc_value( + "Driver=Snowflake;Server=foo.snowflakecomputing.com", + "driver", + ); + let entry = crate::reader::registry::detect(None, driver.as_deref()).unwrap(); + assert_eq!(entry.scheme, "snowflake"); + assert!(entry.odbc_credential_provider.is_some()); + + let driver = crate::reader::connection::extract_odbc_value( + "Driver={PostgreSQL};Server=localhost", + "driver", + ); + let entry = crate::reader::registry::detect(None, driver.as_deref()).unwrap(); + assert!(entry.odbc_credential_provider.is_none()); } #[test] diff --git a/src/reader/registry.rs b/src/reader/registry.rs index 0b079acf..6397ba7c 100644 --- a/src/reader/registry.rs +++ b/src/reader/registry.rs @@ -173,6 +173,12 @@ pub struct DatabaseEntry { /// to SQL_C_SLONG/SBIGINT. Elsewhere integer-scale numerics keep their /// Int64 type (and precision past 2^53). pub odbc_numeric_as_double: bool, + /// Pre-connect hook rewriting the ODBC connection string to supply + /// credentials the driver cannot obtain itself (Snowflake: resolve + /// `ConnectionName=` from `~/.snowflake/connections.toml`, inject a + /// Posit Workbench OAuth token). Dispatched on the `Driver=` value, so + /// it fires before the DBMS name is known. + pub odbc_credential_provider: Option, /// In-process reader to prefer when its cargo feature is compiled in /// (duckdb, sqlite). `None` for backends reached only through external /// ADBC/ODBC drivers. @@ -203,6 +209,7 @@ macro_rules! entry { odbc_dbq_style: false, odbc_row_array_size: None, odbc_numeric_as_double: false, + odbc_credential_provider: None, native_reader: None, } }; @@ -233,6 +240,14 @@ macro_rules! adbc { use DetectPattern::Contains as Has; use DriverUri::*; +/// Snowflake's ODBC credential hook, present only when the `odbc` feature +/// (and with it the `odbc::snowflake` module) is compiled in. +#[cfg(feature = "odbc")] +const SNOWFLAKE_CREDENTIAL_PROVIDER: Option = + Some(crate::reader::odbc::snowflake::apply_workbench_credentials); +#[cfg(not(feature = "odbc"))] +const SNOWFLAKE_CREDENTIAL_PROVIDER: Option = None; + /// All supported backends. **Order matters for detection**: entries are /// scanned top to bottom, so more specific patterns must precede generic /// ones (`microsoft sql server` before anything matching `sql`, Redshift @@ -281,14 +296,17 @@ pub static REGISTRY: &[DatabaseEntry] = &[ &MySqlDialect, adbc!("adbc_driver_mysql", "mysql", MySqlGoDsn, true) ), - entry!( - "snowflake", - &[], - "Snowflake", - &[Has("snowflake")], - &SnowflakeDialect, - adbc!("adbc_driver_snowflake", "snowflake", Passthrough, true) - ), + DatabaseEntry { + odbc_credential_provider: SNOWFLAKE_CREDENTIAL_PROVIDER, + ..entry!( + "snowflake", + &[], + "Snowflake", + &[Has("snowflake")], + &SnowflakeDialect, + adbc!("adbc_driver_snowflake", "snowflake", Passthrough, true) + ) + }, entry!( "bigquery", &[], @@ -535,28 +553,40 @@ pub fn resolve_dialect( detect_or_err(None, Some(body)) } +/// Error for an unrecognized backend, naming what was seen. Shared by +/// [`detect_or_err`] and [`detect_entry_or_err`]. +fn unknown_backend_error(dbms_name: Option<&str>, driver_hint: Option<&str>) -> crate::GgsqlError { + crate::GgsqlError::ReaderError(format!( + "Unrecognized database backend (DBMS name: {}, driver: {}). \ + ggsql does not know which SQL dialect to use. If the backend \ + is close to a supported one, pin the dialect explicitly with \ + a `dialect=` parameter (e.g. dialect=postgres), or use \ + dialect=ansi for generic ANSI SQL.", + dbms_name.unwrap_or(""), + driver_hint.unwrap_or(""), + )) +} + +/// Like [`detect`], but errors on unknown backends with a message naming +/// what was seen — silently falling back to ANSI produced broken SQL too +/// often. Returns the full entry so callers can derive dialect, ODBC +/// quirks, and more from a single resolution. +pub fn detect_entry_or_err( + dbms_name: Option<&str>, + driver_hint: Option<&str>, +) -> crate::Result<&'static DatabaseEntry> { + detect(dbms_name, driver_hint).ok_or_else(|| unknown_backend_error(dbms_name, driver_hint)) +} + /// Detect a dialect from an ODBC DBMS name and/or driver string, or error -/// naming what was seen. Unknown backends are an **error** — silently -/// falling back to ANSI produced broken SQL too often. The escape hatch is -/// a `dialect=ansi` (or `dialect=`) parameter; see +/// naming what was seen. Unknown backends are an **error** — the escape +/// hatch is a `dialect=ansi` (or `dialect=`) parameter; see /// [`dialect_override`]. pub fn detect_or_err( dbms_name: Option<&str>, driver_hint: Option<&str>, ) -> crate::Result { - detect(dbms_name, driver_hint) - .map(|e| e.dialect()) - .ok_or_else(|| { - crate::GgsqlError::ReaderError(format!( - "Unrecognized database backend (DBMS name: {}, driver: {}). \ - ggsql does not know which SQL dialect to use. If the backend \ - is close to a supported one, pin the dialect explicitly with \ - a `dialect=` parameter (e.g. dialect=postgres), or use \ - dialect=ansi for generic ANSI SQL.", - dbms_name.unwrap_or(""), - driver_hint.unwrap_or(""), - )) - }) + detect_entry_or_err(dbms_name, driver_hint).map(|e| e.dialect()) } #[cfg(test)]