Skip to main content

atuin_server_database/
lib.rs

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