Skip to main content

atuin_server_database/
lib.rs

1#![forbid(unsafe_code)]
2
3pub mod models;
4
5use std::fmt::Debug;
6
7use async_trait::async_trait;
8use atuin_domain::record::{EncryptedData, Record, RecordIdx, RecordSeriesKey, RecordStatus};
9use serde::{Deserialize, Serialize};
10
11use self::models::{NewSession, NewUser, Session, User};
12
13#[derive(Debug, derive_more::Display, derive_more::Error, derive_more::From)]
14#[display("{self:?}")]
15pub enum DbError {
16    #[from(skip)]
17    NotFound,
18    #[from(time::error::ComponentRange, time::error::Error)]
19    Other(eyre::Report),
20}
21
22impl From<sqlx::Error> for DbError {
23    fn from(error: sqlx::Error) -> Self {
24        match error {
25            sqlx::Error::RowNotFound => Self::NotFound,
26            error => Self::Other(error.into()),
27        }
28    }
29}
30
31pub type DbResult<T> = Result<T, DbError>;
32
33#[derive(Debug, PartialEq)]
34pub enum DbType {
35    Postgres,
36    Sqlite,
37    Unknown,
38}
39
40#[derive(Clone, Deserialize, Serialize)]
41pub struct DbSettings {
42    pub db_uri: String,
43    /// Optional URI for read replicas. If set, read-only queries will use this connection.
44    pub read_db_uri: Option<String>,
45}
46
47impl DbSettings {
48    pub fn db_type(&self) -> DbType {
49        if self.db_uri.starts_with("postgres://") || self.db_uri.starts_with("postgresql://") {
50            DbType::Postgres
51        } else if self.db_uri.starts_with("sqlite:") {
52            DbType::Sqlite
53        } else {
54            DbType::Unknown
55        }
56    }
57}
58
59fn redact_db_uri(uri: &str) -> String {
60    url::Url::parse(uri)
61        .map(|mut url| {
62            let _ = url.set_password(Some("****"));
63            url.to_string()
64        })
65        .unwrap_or_else(|_| uri.to_string())
66}
67
68// Do our best to redact passwords so they're not logged in the event of an error.
69impl Debug for DbSettings {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        if self.db_type() == DbType::Postgres {
72            let redacted_uri = redact_db_uri(&self.db_uri);
73            let redacted_read_uri = self.read_db_uri.as_ref().map(|uri| redact_db_uri(uri));
74            f.debug_struct("DbSettings")
75                .field("db_uri", &redacted_uri)
76                .field("read_db_uri", &redacted_read_uri)
77                .finish()
78        } else {
79            f.debug_struct("DbSettings")
80                .field("db_uri", &self.db_uri)
81                .field("read_db_uri", &self.read_db_uri)
82                .finish()
83        }
84    }
85}
86
87#[async_trait]
88pub trait Database: Sized + Clone + Send + Sync + 'static {
89    async fn new(settings: &DbSettings) -> DbResult<Self>;
90
91    async fn get_session(&self, token: &str) -> DbResult<Session>;
92    async fn get_session_user(&self, token: &str) -> DbResult<User>;
93    async fn add_session(&self, session: &NewSession) -> DbResult<()>;
94
95    async fn get_user(&self, username: &str) -> DbResult<User>;
96    async fn get_user_session(&self, u: &User) -> DbResult<Session>;
97    async fn add_user(&self, user: &NewUser) -> DbResult<i64>;
98
99    async fn update_user_password(&self, u: &User) -> DbResult<()>;
100
101    async fn delete_user(&self, u: &User) -> DbResult<()>;
102    async fn delete_store(&self, user: &User) -> DbResult<()>;
103
104    async fn add_records(&self, user: &User, record: &[Record<EncryptedData>]) -> DbResult<()>;
105    async fn next_records(
106        &self,
107        user: &User,
108        series: &RecordSeriesKey,
109        start: Option<RecordIdx>,
110        count: u64,
111    ) -> DbResult<Vec<Record<EncryptedData>>>;
112
113    // Return the tail record ID for each store, so (HostID, Tag, TailRecordID)
114    async fn status(&self, user: &User) -> DbResult<RecordStatus>;
115}