agentic_core/storage/
backend.rs1use std::time::Duration;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum DatabaseBackend {
8 Postgres,
10 Sqlite,
12 Other,
14}
15
16impl DatabaseBackend {
17 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 #[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}