Allow any valid sqlite or pg connection string (#355)

Author: jayvdb

.github/workflows/ci.yml

@@ -47,6 +47,7 @@ jobs:
         sudo cat /etc/postgresql/16/main/pg_hba.conf
         sudo service postgresql restart && sleep 3
         echo BUTANE_PG_CONNSTR="host=localhost user=postgres sslmode=disable port=5432" >> $GITHUB_ENV
+        echo "/usr/lib/postgresql/16/bin" >> $GITHUB_PATH
     - name: Setup PostgreSQL on MacOS
       if: runner.os == 'macOS'
       run: |
@@ -109,6 +110,8 @@ jobs:
       run: |
         set -ex
         make regenerate-example-migrations
+        # TODO: This file is created by one of the tests on Windows
+        rm -f butane_core/sqlite
         git add -A
         git diff --cached --exit-code
     - name: Run tests in examples

Cargo.lock

@@ -43,6 +43,7 @@ log = "0.4"
 maybe-async-cfg = { version = "0.2.5", default-features = false }
 nonempty = "0.11"
 paste = "1.0.11"
+postgres-native-tls = "0.5"
 pretty_assertions = "1.4"
 proc-macro2 = { version = "1.0", default-features = false }
 quote = { version = "1.0", default-features = false }
@@ -54,6 +55,7 @@ serde_json = "1.0"
 sqlparser = "0.56"
 syn = { version = "2", features = ["extra-traits", "full"] }
 tempfile = "3.10"
+test-log = { version = "0.2", features = ["log", "trace"] }
 thiserror = "2.0"
 tokio = { version = "1"}
 tokio-postgres = "0.7"

Cargo.toml

@@ -43,7 +43,7 @@ nonempty.workspace = true
 pin-project = "1"
 tokio = {workspace = true, optional = true, features = ["rt", "sync", "rt-multi-thread"]}
 tokio-postgres = { optional = true, workspace = true }
-postgres-native-tls = { version = "0.5", optional = true }
+postgres-native-tls = { optional = true, workspace = true }
 proc-macro2 = { workspace = true }
 quote = { workspace = true }
 rand = { optional = true, workspace = true }
@@ -66,6 +66,8 @@ paste = { workspace = true }
 pretty_assertions.workspace = true
 tempfile.workspace = true
 tokio = { workspace = true, features = ["macros"] }
+uuid.workspace = true
+whoami = "1.6"
 
 [[test]]
 name = "uuid"

butane_core/Cargo.toml

@@ -526,6 +526,9 @@ pub trait Backend: Send + Sync + DynClone {
 
 dyn_clone::clone_trait_object!(Backend);
 
+/// Regular expression to match PostgreSQL key-value pairs in a connection string.
+const PG_KEY_PAIR_RE: &str = r"(hostaddr|host|dbname|user|port)\s*=";
+
 /// Connection specification. Contains the name of a database backend
 /// and the backend-specific connection string. See [`connect`]
 /// to make a [`Connection`] from a `ConnectionSpec`.
@@ -549,7 +552,7 @@ impl ConnectionSpec {
         contents.push('\n');
         f.write_all(contents.as_bytes()).map_err(|e| e.into())
     }
-    /// Load a previously saved connection spec
+    /// Load a previously saved connection spec.
     pub fn load(path: impl AsRef<Path>) -> Result<Self> {
         let path = conn_complete_if_dir(path.as_ref());
         serde_json::from_reader(fs::File::open(path)?).map_err(|e| e.into())
@@ -560,28 +563,115 @@ impl ConnectionSpec {
             None => Err(crate::Error::UnknownBackend(self.backend_name.clone())),
         }
     }
+    /// Get the backend name.
+    pub const fn backend_name(&self) -> &String {
+        &self.backend_name
+    }
+    /// Get the database connection string.
+    pub const fn connection_string(&self) -> &String {
+        &self.conn_str
+    }
+    /// Get the connection string URI if it is a URI.
+    /// Returns None if the connection string is not a URI.
+    pub fn connection_string_uri(&self) -> Option<url::Url> {
+        url::Url::parse(self.connection_string()).ok()
+    }
+
+    fn is_pg_key_value_pairs(connection_string: &str) -> bool {
+        // host and hostaddr are optional in Postgresql, however this prevents connecting
+        // without one of them: https://github.com/sfackler/rust-postgres/issues/1239
+        let re = regex::Regex::new(PG_KEY_PAIR_RE).unwrap();
+        re.is_match(connection_string)
+    }
+
+    /// Add a query parameter to the connection string.
+    pub fn add_parameter(&mut self, name: &str, value: &str) -> Result<()> {
+        if self.backend_name() == "pg" && Self::is_pg_key_value_pairs(&self.conn_str) {
+            // Append using PostgreSQL key-value pair syntax.
+            self.conn_str.push_str(&format!(" {}={}", name, value));
+            return Ok(());
+        }
+        if self.connection_string_uri().is_none() {
+            return Err(crate::Error::UnknownConnectString(
+                "Cannot add query parameter to connection string that is not a postgres or URI"
+                    .to_string(),
+            ));
+        }
+
+        if self.conn_str.contains('?') {
+            self.conn_str.push_str(&format!("&{}={}", name, value));
+        } else {
+            self.conn_str.push_str(&format!("?{}={}", name, value));
+        }
+        Ok(())
+    }
 }
 
+/// Convert a string to a connection spec.
+///
+/// The string may be any connection string supported by the database backend.
+/// If it is a URL with a scheme of "sqlite", the backend will be set to "sqlite",
+/// and the scheme will be set to "file" as required by the sqlite engine.
+/// If it is a URI with a scheme of "file", the backend will be set to "sqlite".
+/// If it does not have a scheme, it will be set to "pg" if it contains any of
+/// `host=`, `hostaddr=`, `dbname=`, or `user=`.
+/// Otherwise it will default to "sqlite".
 impl TryFrom<&str> for ConnectionSpec {
     type Error = crate::Error;
     fn try_from(value: &str) -> Result<Self> {
-        let parsed = url::Url::parse(value)?;
-        if parsed.scheme() == "sqlite" {
-            let path = value.trim_start_matches("sqlite://");
-            Ok(ConnectionSpec {
-                backend_name: "sqlite".to_string(),
-                conn_str: path.to_string(),
-            })
-        } else if ["postgres", "postgresql"].contains(&parsed.scheme()) {
-            Ok(ConnectionSpec {
-                backend_name: "pg".to_string(),
-                conn_str: value.to_string(),
-            })
-        } else {
-            Ok(ConnectionSpec {
-                backend_name: parsed.scheme().to_string(),
-                conn_str: value.to_string(),
-            })
+        match url::Url::parse(value) {
+            Ok(parsed) => {
+                match parsed.scheme() {
+                    "sqlite" => {
+                        // SQLite uses file: URLs, but we want to accept sqlite:
+                        // in order that the URIs are more expressive.
+                        let value = value.replacen("sqlite:", "file:", 1);
+                        Ok(ConnectionSpec {
+                            backend_name: "sqlite".to_string(),
+                            conn_str: value,
+                        })
+                    }
+                    "file" => Ok(ConnectionSpec {
+                        backend_name: "sqlite".to_string(),
+                        conn_str: value.to_string(),
+                    }),
+                    "postgres" | "postgresql" => Ok(ConnectionSpec {
+                        backend_name: "pg".to_string(),
+                        conn_str: value.to_string(),
+                    }),
+                    _ => Ok(ConnectionSpec {
+                        backend_name: parsed.scheme().to_string(),
+                        conn_str: value.to_string(),
+                    }),
+                }
+            }
+            Err(url::ParseError::InvalidPort) => {
+                // This occurs when using a PostgreSQL multi-host connection string.
+                if value.starts_with("postgres") {
+                    Ok(ConnectionSpec {
+                        backend_name: "pg".to_string(),
+                        conn_str: value.to_string(),
+                    })
+                } else {
+                    Ok(ConnectionSpec {
+                        backend_name: "sqlite".to_string(),
+                        conn_str: value.to_string(),
+                    })
+                }
+            }
+            Err(_) => {
+                // Spaces are allowed between the key and the equals sign in a PostgreSQL connection string.
+                if Self::is_pg_key_value_pairs(value) {
+                    return Ok(ConnectionSpec {
+                        backend_name: "pg".to_string(),
+                        conn_str: value.to_string(),
+                    });
+                }
+                Ok(ConnectionSpec {
+                    backend_name: "sqlite".to_string(),
+                    conn_str: value.to_string(),
+                })
+            }
         }
     }
 }

butane_core/src/db/mod.rs

@@ -260,6 +260,8 @@ pub enum Error {
     UriParse(#[from] url::ParseError),
     #[error("Unknown backend {0}")]
     UnknownBackend(String),
+    #[error("Unknown connect string {0}")]
+    UnknownConnectString(String),
     #[error("Range error")]
     OutOfRange,
     #[error("Internal logic error {0}")]