Skip to main content

pubky_homeserver/persistence/sql/
connection_string.rs

1use std::{fmt::Display, str::FromStr};
2
3use serde::{Deserialize, Serialize};
4
5/// A connection string for a  postgres database.
6/// See <https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING-URIS>
7#[derive(Debug, Clone, PartialEq)]
8pub struct ConnectionString(url::Url);
9
10impl ConnectionString {
11    /// Create a new connection string from a string.
12    /// This function validates that the connection string is a postgres connection string.
13    pub fn new(con_string: &str) -> anyhow::Result<Self> {
14        Self::validated(url::Url::parse(con_string)?)
15    }
16
17    /// Shared validation: ensures the URL uses a postgres scheme.
18    fn validated(url: url::Url) -> anyhow::Result<Self> {
19        let cs = Self(url);
20        if !cs.is_postgres() {
21            anyhow::bail!("Only postgres database urls are supported");
22        }
23        Ok(cs)
24    }
25
26    /// Get the connection string as a str.
27    pub fn as_str(&self) -> &str {
28        self.0.as_str()
29    }
30
31    fn is_postgres(&self) -> bool {
32        self.0.scheme() == "postgres" || self.0.scheme() == "postgresql"
33    }
34
35    /// Get the database name
36    /// For postgres, this is the database name directly
37    pub fn database_name(&self) -> &str {
38        self.0.path().trim_start_matches("/")
39    }
40
41    /// Set the database name, clearing any `dbname` query parameter that would
42    /// otherwise override the path. See
43    /// <https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING-URIS>
44    pub fn set_database_name(&mut self, db_name: &str) {
45        self.0.set_path(db_name);
46        self.remove_query_param("dbname");
47    }
48
49    /// Remove all occurrences of a query parameter by key.
50    fn remove_query_param(&mut self, key: &str) {
51        if self.0.query().is_none() {
52            return;
53        }
54        let pairs: Vec<_> = self
55            .0
56            .query_pairs()
57            .filter(|(k, _)| k != key)
58            .map(|(k, v)| (k.into_owned(), v.into_owned()))
59            .collect();
60        if pairs.is_empty() {
61            self.0.set_query(None);
62        } else {
63            self.0.query_pairs_mut().clear().extend_pairs(&pairs);
64        }
65    }
66}
67
68impl TryFrom<url::Url> for ConnectionString {
69    type Error = anyhow::Error;
70
71    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
72        Self::validated(url)
73    }
74}
75
76impl FromStr for ConnectionString {
77    type Err = anyhow::Error;
78
79    fn from_str(s: &str) -> Result<Self, Self::Err> {
80        Self::new(s)
81    }
82}
83
84impl Display for ConnectionString {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        write!(f, "{}", self.0)
87    }
88}
89
90impl Serialize for ConnectionString {
91    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
92    where
93        S: serde::Serializer,
94    {
95        serializer.serialize_str(self.as_str())
96    }
97}
98
99impl<'de> Deserialize<'de> for ConnectionString {
100    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
101    where
102        D: serde::Deserializer<'de>,
103    {
104        let s = String::deserialize(deserializer)?;
105        Self::new(&s).map_err(serde::de::Error::custom)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_valid_postgres_url() {
115        let _: ConnectionString = "postgres://localhost:5432/pubky_homeserver"
116            .parse()
117            .unwrap();
118    }
119
120    #[test]
121    fn test_non_postgres_url_rejected() {
122        let result: Result<ConnectionString, _> = "sqlite:///path/to/sqlite.db".parse();
123        assert!(result.is_err(), "sqlite URLs should be rejected");
124    }
125
126    #[test]
127    fn set_database_name_changes_path() {
128        let mut cs = ConnectionString::new("postgres://user:pass@localhost:5432/original").unwrap();
129        cs.set_database_name("new_db");
130        assert_eq!(cs.database_name(), "new_db");
131    }
132
133    #[test]
134    fn set_database_name_strips_dbname_query_param() {
135        let mut cs =
136            ConnectionString::new("postgres://user:pass@localhost:5432/postgres?dbname=postgres")
137                .unwrap();
138        cs.set_database_name("pubky_test_abc123");
139        assert_eq!(cs.database_name(), "pubky_test_abc123");
140        assert!(
141            !cs.as_str().contains("dbname="),
142            "dbname query param should be removed, got: {}",
143            cs.as_str()
144        );
145    }
146
147    #[test]
148    fn set_database_name_preserves_other_query_params() {
149        let mut cs = ConnectionString::new(
150            "postgres://user:pass@localhost:5432/postgres?dbname=postgres&sslmode=require",
151        )
152        .unwrap();
153        cs.set_database_name("pubky_test_abc123");
154        assert_eq!(cs.database_name(), "pubky_test_abc123");
155        assert!(
156            !cs.as_str().contains("dbname="),
157            "dbname should be removed, got: {}",
158            cs.as_str()
159        );
160        assert!(
161            cs.as_str().contains("sslmode=require"),
162            "other params should be preserved, got: {}",
163            cs.as_str()
164        );
165    }
166}