Skip to main content

agentic_core/storage/
backend.rs

1//! Database URL classification and sanitization.
2
3use std::time::Duration;
4
5/// Database backend selected by a connection URL.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum DatabaseBackend {
8    /// `PostgreSQL`, using either the `postgres` or `postgresql` URI scheme.
9    Postgres,
10    /// `SQLite`.
11    Sqlite,
12    /// Another configured backend.
13    Other,
14}
15
16impl DatabaseBackend {
17    /// Parses a database URL and classifies its normalized URI scheme.
18    ///
19    /// # Errors
20    ///
21    /// Returns [`url::ParseError`] when `database_url` is not a valid absolute URL.
22    pub fn from_url(database_url: &str) -> Result<Self, url::ParseError> {
23        let url = url::Url::parse(database_url)?;
24        Ok(match url.scheme() {
25            "postgres" | "postgresql" => Self::Postgres,
26            "sqlite" => Self::Sqlite,
27            _ => Self::Other,
28        })
29    }
30
31    /// Returns the backend name used in diagnostics.
32    #[must_use]
33    pub const fn display_name(self) -> &'static str {
34        match self {
35            Self::Postgres => "PostgreSQL",
36            Self::Sqlite => "SQLite",
37            Self::Other => "configured",
38        }
39    }
40
41    pub(crate) fn from_connection(connection: &sqlx::AnyConnection) -> Self {
42        match connection.backend_name() {
43            "PostgreSQL" => Self::Postgres,
44            "SQLite" => Self::Sqlite,
45            _ => Self::Other,
46        }
47    }
48}
49
50pub(crate) async fn configure_postgres_timeouts(
51    connection: &mut sqlx::AnyConnection,
52    lock_timeout: Duration,
53    statement_timeout: Duration,
54) -> Result<(), sqlx::Error> {
55    let lock_timeout_ms = format!("{}ms", lock_timeout.as_millis());
56    sqlx::query("SELECT set_config('lock_timeout', $1, false)")
57        .bind(lock_timeout_ms)
58        .execute(&mut *connection)
59        .await?;
60    let statement_timeout_ms = format!("{}ms", statement_timeout.as_millis());
61    sqlx::query("SELECT set_config('statement_timeout', $1, false)")
62        .bind(statement_timeout_ms)
63        .execute(connection)
64        .await?;
65    Ok(())
66}
67
68pub(crate) fn redact_database_urls(message: &str) -> String {
69    const DATABASE_SCHEMES: [&str; 3] = ["postgresql://", "postgres://", "mysql://"];
70
71    let mut redacted = message.to_owned();
72    let mut lowercase = message.to_ascii_lowercase();
73    for scheme in DATABASE_SCHEMES {
74        let mut search_from = 0;
75        let replacement = format!("{scheme}[redacted]");
76        while let Some(offset) = lowercase[search_from..].find(scheme) {
77            let start = search_from + offset;
78            let end = redacted[start..]
79                .find(char::is_whitespace)
80                .map_or(redacted.len(), |length| start + length);
81            redacted.replace_range(start..end, &replacement);
82            lowercase.replace_range(start..end, &replacement);
83            search_from = start + replacement.len();
84        }
85    }
86    redacted
87}