feat: infer driver from URI scheme

Author: esadek

README.md

@@ -121,7 +121,7 @@ Usage: databow [OPTIONS]
 
 Options:
       --profile <profile>    Connection profile name or path
-      --driver <driver>      Driver name (required if --profile not specified)
+      --driver <driver>      Driver name (inferred from --uri scheme or profile if not specified)
       --uri <uri>            Database uniform resource identifier
       --username <username>  Database user username
       --password <password>  Database user password

src/cli.rs

@@ -17,7 +17,7 @@ pub enum QuerySource {
 #[derive(Debug, Clone)]
 pub enum ConnectionSource {
     Direct {
-        driver_name: String,
+        driver_name: Option<String>,
         uri: Option<String>,
         username: Option<String>,
         password: Option<String>,
@@ -52,7 +52,7 @@ pub fn parse_args() -> AppConfig {
             .conflicts_with("driver"),
         Arg::new("driver")
             .long("driver")
-            .help("Driver name (required if --profile not specified)"),
+            .help("Driver name (inferred from --uri scheme or profile if not specified)"),
         Arg::new("uri")
             .long("uri")
             .help("Database uniform resource identifier"),
@@ -123,14 +123,28 @@ pub fn parse_args() -> AppConfig {
         }
     } else if let Some(driver_name) = driver_name {
         ConnectionSource::Direct {
-            driver_name,
+            driver_name: Some(driver_name),
+            uri,
+            username,
+            password,
+            options,
+        }
+    } else if uri.as_deref().is_some_and(uri_has_driver_scheme) {
+        // The driver can be inferred from the URI scheme (e.g. `--uri duckdb://...`
+        // implies the `duckdb` driver), so neither --driver nor --profile is
+        // required. The raw URI is handed to `ManagedDatabase::from_uri`, which
+        // performs the scheme parsing and driver selection.
+        ConnectionSource::Direct {
+            driver_name: None,
             uri,
             username,
             password,
             options,
         }
     } else {
-        eprintln!("Error: Either --profile or --driver is required");
+        eprintln!(
+            "Error: Provide --driver, --profile, or a --uri with a scheme (e.g. duckdb://...)"
+        );
         exit(1);
     };
 
@@ -163,6 +177,13 @@ pub fn parse_args() -> AppConfig {
     }
 }
 
+fn uri_has_driver_scheme(uri: &str) -> bool {
+    let Some(idx) = uri.find(':') else {
+        return false;
+    };
+    idx != 0
+}
+
 fn parse_option(option: &str) -> Result<(String, String), String> {
     let parts: Vec<&str> = option.splitn(2, '=').collect();
     if parts.len() != 2 {
@@ -184,15 +205,15 @@ mod tests {
     #[test]
     fn test_connection_source_direct() {
         let source = ConnectionSource::Direct {
-            driver_name: "duckdb".to_string(),
+            driver_name: Some("duckdb".to_string()),
             uri: Some(":memory:".to_string()),
             username: None,
             password: None,
             options: vec![],
         };
         match source {
             ConnectionSource::Direct { driver_name, .. } => {
-                assert_eq!(driver_name, "duckdb");
+                assert_eq!(driver_name, Some("duckdb".to_string()));
             }
             _ => panic!("Expected Direct variant"),
         }
@@ -216,6 +237,36 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_uri_has_driver_scheme_with_scheme() {
+        assert!(uri_has_driver_scheme("duckdb://:memory:"));
+        assert!(uri_has_driver_scheme("postgresql://host:5432/db"));
+        assert!(uri_has_driver_scheme("sqlite:file::memory:"));
+        assert!(uri_has_driver_scheme("sqlite::memory:"));
+        // Bare `driver:` form: accepted, matching `parse_driver_uri`.
+        assert!(uri_has_driver_scheme("duckdb:"));
+    }
+
+    #[test]
+    fn test_uri_has_driver_scheme_arbitrary_scheme() {
+        assert!(uri_has_driver_scheme("DuckDB://x"));
+    }
+
+    #[test]
+    fn test_uri_has_driver_scheme_no_scheme() {
+        // No colon at all.
+        assert!(!uri_has_driver_scheme("plain_path"));
+        // Empty scheme (leading colon).
+        assert!(!uri_has_driver_scheme(":memory:"));
+    }
+
+    #[test]
+    fn test_uri_has_driver_scheme_accepts_profile_scheme() {
+        // `profile://` is a valid scheme; `from_uri` resolves it to a profile.
+        assert!(uri_has_driver_scheme("profile://my_database"));
+        assert!(uri_has_driver_scheme("PROFILE://my_database"));
+    }
+
     #[test]
     fn test_connection_source_profile_with_path() {
         let source = ConnectionSource::Profile {
@@ -272,7 +323,7 @@ mod tests {
     fn test_app_config_with_direct_connection() {
         let config = AppConfig {
             connection: ConnectionSource::Direct {
-                driver_name: "test_driver".to_string(),
+                driver_name: Some("test_driver".to_string()),
                 uri: Some("test_uri".to_string()),
                 username: Some("test_user".to_string()),
                 password: Some("test_pass".to_string()),
@@ -291,7 +342,7 @@ mod tests {
                 password,
                 options,
             } => {
-                assert_eq!(driver_name, "test_driver");
+                assert_eq!(driver_name, &Some("test_driver".to_string()));
                 assert_eq!(*uri, Some("test_uri".to_string()));
                 assert_eq!(*username, Some("test_user".to_string()));
                 assert_eq!(*password, Some("test_pass".to_string()));

src/database.rs

@@ -7,7 +7,7 @@ use adbc_core::{Database, Driver, LOAD_FLAG_DEFAULT, Statement};
 use adbc_driver_manager::profile::{
     ConnectionProfile, ConnectionProfileProvider, FilesystemProfileProvider, process_profile_value,
 };
-use adbc_driver_manager::{ManagedConnection, ManagedDriver};
+use adbc_driver_manager::{ManagedConnection, ManagedDatabase, ManagedDriver};
 use arrow_array::RecordBatch;
 
 pub fn initialize_connection(source: ConnectionSource) -> Result<ManagedConnection, String> {
@@ -30,30 +30,58 @@ pub fn initialize_connection(source: ConnectionSource) -> Result<ManagedConnecti
 }
 
 fn initialize_direct_connection(
-    driver_name: String,
+    driver_name: Option<String>,
     uri: Option<String>,
     username: Option<String>,
     password: Option<String>,
     options: Vec<(String, String)>,
 ) -> Result<ManagedConnection, String> {
-    let mut driver = ManagedDriver::load_from_name(
-        &driver_name,
-        None,
-        AdbcVersion::default(),
-        LOAD_FLAG_DEFAULT,
-        None,
-    )
-    .map_err(|e| format!("Failed to load driver '{}': {}", driver_name, e))?;
-
-    let db_options = build_database_options(uri, username, password, options);
-
-    let database = driver
-        .new_database_with_opts(db_options)
-        .map_err(|e| format!("Failed to create database handle: {e}"))?;
+    match driver_name {
+        // Explicit driver: load it by name and hand it the URI as an option.
+        Some(driver_name) => {
+            let mut driver = ManagedDriver::load_from_name(
+                &driver_name,
+                None,
+                AdbcVersion::default(),
+                LOAD_FLAG_DEFAULT,
+                None,
+            )
+            .map_err(|e| format!("Failed to load driver '{}': {}", driver_name, e))?;
+
+            let db_options = build_database_options(uri, username, password, options);
+
+            let database = driver
+                .new_database_with_opts(db_options)
+                .map_err(|e| format!("Failed to create database handle: {e}"))?;
+
+            database
+                .new_connection()
+                .map_err(|e| format!("Failed to create connection: {e}"))
+        }
+        // No driver given: infer it from the URI scheme via the driver manager.
+        None => {
+            let uri = uri
+                .ok_or_else(|| "Internal error: URI is required to infer the driver".to_string())?;
+
+            // The URI itself selects the driver; the remaining CLI options are
+            // layered on top (they override any options derived from the URI).
+            let opts = build_database_options(None, username, password, options);
+
+            let database = ManagedDatabase::from_uri_with_opts(
+                &uri,
+                None,
+                AdbcVersion::default(),
+                LOAD_FLAG_DEFAULT,
+                None,
+                opts,
+            )
+            .map_err(|e| format!("Failed to load driver from URI '{}': {}", uri, e))?;
 
-    database
-        .new_connection()
-        .map_err(|e| format!("Failed to create connection: {e}"))
+            database
+                .new_connection()
+                .map_err(|e| format!("Failed to create connection: {e}"))
+        }
+    }
 }
 
 fn initialize_profile_connection(

tests/integration_test.rs

@@ -60,6 +60,41 @@ fn test_query_argument() {
     assert!(stdout.contains("42"));
 }
 
+#[test]
+fn test_driver_inferred_from_uri_scheme() {
+    // No --driver or --profile: the driver should be inferred from the
+    // `sqlite:` URI scheme.
+    let output = Command::new("cargo")
+        .args([
+            "run",
+            "--",
+            "--uri",
+            "sqlite::memory:",
+            "--query",
+            "SELECT 42 AS answer",
+        ])
+        .output()
+        .expect("Failed to execute command");
+
+    assert!(output.status.success());
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert!(stdout.contains("answer"));
+    assert!(stdout.contains("42"));
+}
+
+#[test]
+fn test_no_driver_no_scheme_uri_errors() {
+    // A URI with no scheme and no --driver/--profile should still error.
+    let output = Command::new("cargo")
+        .args(["run", "--", "--uri", "plain_path", "--query", "SELECT 1"])
+        .output()
+        .expect("Failed to execute command");
+
+    assert!(!output.status.success());
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert!(stderr.contains("--driver"));
+}
+
 #[test]
 fn test_file_argument() {
     // Create a temporary SQL file
@@ -241,6 +276,41 @@ driver = "duckdb"
     assert!(stdout.contains("42"));
 }
 
+#[test]
+fn test_uri_profile_scheme_from_file_path() {
+    use std::io::Write;
+
+    // A `profile://<path>` URI (no --driver/--profile) should be resolved by
+    // the driver manager's `from_uri`, which loads the referenced profile.
+    let mut temp_file = NamedTempFile::with_suffix(".toml").expect("Failed to create temp file");
+    let profile_path = temp_file.path().to_string_lossy().to_string();
+
+    temp_file
+        .write_all(
+            br#"profile_version = 1
+driver = "duckdb"
+
+[Options]
+"#,
+        )
+        .expect("Failed to write to temp file");
+
+    let uri = format!("profile://{profile_path}");
+    let output = Command::new("cargo")
+        .args(["run", "--", "--uri", &uri, "--query", "SELECT 42 AS answer"])
+        .output()
+        .expect("Failed to execute command");
+
+    assert!(
+        output.status.success(),
+        "profile:// URI should resolve the profile. stderr: {}",
+        String::from_utf8_lossy(&output.stderr)
+    );
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert!(stdout.contains("answer"));
+    assert!(stdout.contains("42"));
+}
+
 #[test]
 fn test_profile_with_uri_override() {
     use std::io::Write;