Registry driver hooks

src/reader/adbc.rs

@@ -27,9 +27,11 @@ pub struct AdbcReader<D: Driver> {
     // takes &self.
     connection: RefCell<<D::DatabaseType as Database>::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<crate::reader::registry::AdbcInfo>,
+    /// 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<D: Driver> AdbcReader<D> {
             _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<D: Driver> AdbcReader<D> {
             _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<D: Driver> AdbcReader<D> {
         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<ManagedDriver> {
 
         let reader = Self::new_with_database_opts(driver, dialect, opts)?
             .with_statement_opts(conn.ggsql.stmt_options.clone());
-        // adbc://<driver> with a recognized short name also gets quirks.
+        // adbc://<driver> 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 {

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<super::DialectRef> {
-    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=<name>`) 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();

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]

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<fn(&mut String)>,
     /// 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<fn(&mut String)> =
+    Some(crate::reader::odbc::snowflake::apply_workbench_credentials);
+#[cfg(not(feature = "odbc"))]
+const SNOWFLAKE_CREDENTIAL_PROVIDER: Option<fn(&mut String)> = 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=<scheme>` parameter (e.g. dialect=postgres), or use \
+         dialect=ansi for generic ANSI SQL.",
+        dbms_name.unwrap_or("<none>"),
+        driver_hint.unwrap_or("<none>"),
+    ))
+}
+
+/// 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=<scheme>`) parameter; see
+/// naming what was seen. Unknown backends are an **error** — the escape
+/// hatch is a `dialect=ansi` (or `dialect=<scheme>`) parameter; see
 /// [`dialect_override`].
 pub fn detect_or_err(
     dbms_name: Option<&str>,
     driver_hint: Option<&str>,
 ) -> crate::Result<crate::reader::DialectRef> {
-    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=<scheme>` parameter (e.g. dialect=postgres), or use \
-                 dialect=ansi for generic ANSI SQL.",
-                dbms_name.unwrap_or("<none>"),
-                driver_hint.unwrap_or("<none>"),
-            ))
-        })
+    detect_entry_or_err(dbms_name, driver_hint).map(|e| e.dialect())
 }
 
 #[cfg(test)]