aboutsummaryrefslogtreecommitdiffstats
path: root/crates/turtle/src/atuin_server/database
diff options
context:
space:
mode:
Diffstat (limited to 'crates/turtle/src/atuin_server/database')
-rw-r--r--crates/turtle/src/atuin_server/database/db/mod.rs274
-rw-r--r--crates/turtle/src/atuin_server/database/db/wrappers.rs32
-rw-r--r--crates/turtle/src/atuin_server/database/mod.rs99
-rw-r--r--crates/turtle/src/atuin_server/database/models.rs5
4 files changed, 0 insertions, 410 deletions
diff --git a/crates/turtle/src/atuin_server/database/db/mod.rs b/crates/turtle/src/atuin_server/database/db/mod.rs
deleted file mode 100644
index 77bd0c61..00000000
--- a/crates/turtle/src/atuin_server/database/db/mod.rs
+++ /dev/null
@@ -1,274 +0,0 @@
-use std::collections::HashMap;
-
-use rand::Rng;
-
-use crate::{
- atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus},
- atuin_server::database::{DbError, DbResult, DbSettings, models::User},
-};
-use sqlx::postgres::PgPoolOptions;
-
-use tracing::instrument;
-use uuid::Uuid;
-use wrappers::DbRecord;
-
-mod wrappers;
-
-const MIN_PG_VERSION: u32 = 14;
-
-#[derive(Clone)]
-pub(crate) struct ServerPostgres {
- pool: sqlx::Pool<sqlx::postgres::Postgres>,
- /// Optional read replica pool for read-only queries
- read_pool: Option<sqlx::Pool<sqlx::postgres::Postgres>>,
-}
-
-impl ServerPostgres {
- /// Returns the appropriate pool for read operations.
- /// Uses `read_pool` if available, otherwise falls back to the primary pool.
- fn read_pool(&self) -> &sqlx::Pool<sqlx::postgres::Postgres> {
- self.read_pool.as_ref().unwrap_or(&self.pool)
- }
-}
-
-impl ServerPostgres {
- pub(crate) async fn new(settings: &DbSettings) -> DbResult<Self> {
- let pool = PgPoolOptions::new()
- .max_connections(100)
- .connect(settings.db_uri.as_str())
- .await?;
-
- // Call server_version_num to get the DB server's major version number
- // The call returns None for servers older than 8.x.
- let pg_major_version: u32 =
- pool.acquire()
- .await?
- .server_version_num()
- .ok_or(DbError::Other(eyre::Report::msg(
- "could not get PostgreSQL version",
- )))?
- / 10000;
-
- if pg_major_version < MIN_PG_VERSION {
- return Err(DbError::Other(eyre::Report::msg(format!(
- "unsupported PostgreSQL version {pg_major_version}, minimum required is {MIN_PG_VERSION}"
- ))));
- }
-
- sqlx::migrate!("./db/server-pg-migrations")
- .run(&pool)
- .await
- .map_err(|error| DbError::Other(error.into()))?;
-
- // Create read replica pool if configured
- let read_pool = if let Some(read_db_uri) = &settings.read_db_uri {
- tracing::info!("Connecting to read replica database");
- let read_pool = PgPoolOptions::new()
- .max_connections(100)
- .connect(read_db_uri.as_str())
- .await?;
-
- // Verify the read replica is also a supported PostgreSQL version
- let read_pg_major_version: u32 = read_pool
- .acquire()
- .await?
- .server_version_num()
- .ok_or(DbError::Other(eyre::Report::msg(
- "could not get PostgreSQL version from read replica",
- )))?
- / 10000;
-
- if read_pg_major_version < MIN_PG_VERSION {
- return Err(DbError::Other(eyre::Report::msg(format!(
- "unsupported PostgreSQL version {read_pg_major_version} on read replica, minimum required is {MIN_PG_VERSION}"
- ))));
- }
-
- Some(read_pool)
- } else {
- None
- };
-
- Ok(Self { pool, read_pool })
- }
-
- #[instrument(skip_all)]
- pub(crate) async fn add_records(
- &self,
- user: &User,
- records: &[Record<EncryptedData>],
- ) -> DbResult<()> {
- let mut tx = self.pool.begin().await?;
-
- // We won't have uploaded this data if it wasn't the max. Therefore, we can deduce the max
- // idx without having to make further database queries. Doing the query on this small
- // amount of data should be much, much faster.
- //
- // Worst case, say we get this wrong. We end up caching data that isn't actually the max
- // idx, so clients upload again. The cache logic can be verified with a sql query anyway :)
-
- let mut heads = HashMap::<(HostId, &str), u64>::new();
-
- for i in records {
- let id = crate::atuin_common::utils::uuid_v7();
-
- let result = sqlx::query(
- "
- INSERT INTO store (id, client_id, host, idx, timestamp, version, tag, data, cek, user_id)
- VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
- ON conflict DO nothing
- ",
- )
- .bind(id)
- .bind(i.id)
- .bind(i.host.id)
- .bind(i.idx as i64)
- .bind(i.timestamp as i64) // throwing away some data, but i64 is still big in terms of time
- .bind(&i.version)
- .bind(&i.tag)
- .bind(&i.data.data)
- .bind(&i.data.content_encryption_key)
- .bind(user.id)
- .execute(&mut *tx)
- .await?;
-
- // Only update heads if we actually inserted the record
- if result.rows_affected() > 0 {
- heads
- .entry((i.host.id, &i.tag))
- .and_modify(|e| {
- if i.idx > *e {
- *e = i.idx;
- }
- })
- .or_insert(i.idx);
- }
- }
-
- // we've built the map of heads for this push, so commit it to the database
- for ((host, tag), idx) in heads {
- sqlx::query(
- "
- INSERT INTO store_idx_cache (user_id, host, tag, idx)
- VALUES ($1, $2, $3, $4)
- ON conflict(user_id, host, tag) DO update
- SET idx = greatest(store_idx_cache.idx, $4)
- ",
- )
- .bind(user.id)
- .bind(host)
- .bind(tag)
- .bind(idx as i64)
- .execute(&mut *tx)
- .await?;
- }
-
- tx.commit().await?;
-
- Ok(())
- }
-
- #[instrument(skip_all)]
- pub(crate) async fn next_records(
- &self,
- user: &User,
- host: HostId,
- tag: String,
- start: Option<RecordIdx>,
- count: u64,
- ) -> DbResult<Vec<Record<EncryptedData>>> {
- tracing::debug!("{:?} - {:?} - {:?}", host, tag, start);
- let start = start.unwrap_or(0);
-
- let records: Result<Vec<DbRecord>, DbError> = sqlx::query_as(
- "
- SELECT client_id, host, idx, timestamp, version, tag, data, cek FROM store
- WHERE user_id = $1
- AND tag = $2
- AND host = $3
- AND idx >= $4
- ORDER BY idx asc
- LIMIT $5
- ",
- )
- .bind(user.id)
- .bind(tag.clone())
- .bind(host)
- .bind(start as i64)
- .bind(count as i64)
- .fetch_all(self.read_pool())
- .await
- .map_err(Into::into);
-
- let ret = match records {
- Ok(records) => {
- let records: Vec<Record<EncryptedData>> = records
- .into_iter()
- .map(|f| {
- let record: Record<EncryptedData> = f.into();
- record
- })
- .collect();
-
- records
- }
- Err(DbError::NotFound) => {
- tracing::debug!("no records found in store: {:?}/{}", host, tag);
- return Ok(vec![]);
- }
- Err(e) => return Err(e),
- };
-
- Ok(ret)
- }
-
- pub(crate) async fn status(&self, user: &User) -> DbResult<RecordStatus> {
- // If IDX_CACHE_ROLLOUT is set, then we
- // 1. Read the value of the var, use it as a % chance of using the cache
- // 2. If we use the cache, just read from the cache table
- // 3. If we don't use the cache, read from the store table
- // IDX_CACHE_ROLLOUT should be between 0 and 100.
-
- let idx_cache_rollout =
- std::env::var("IDX_CACHE_ROLLOUT").unwrap_or_else(|_| "0".to_string());
- let idx_cache_rollout = idx_cache_rollout.parse::<f64>().unwrap_or(0.0);
- let use_idx_cache = rand::thread_rng().gen_bool(idx_cache_rollout / 100.0);
-
- let mut res: Vec<(Uuid, String, i64)> = if use_idx_cache {
- tracing::debug!("using idx cache for user {}", user.id);
- sqlx::query_as(
- "
- SELECT host, tag, idx
- FROM store_idx_cache
- WHERE user_id = $1
- ",
- )
- .bind(user.id)
- .fetch_all(self.read_pool())
- .await?
- } else {
- tracing::debug!("using aggregate query for user {}", user.id);
- sqlx::query_as(
- "
- SELECT host, tag, max(idx)
- FROM store
- WHERE user_id = $1
- GROUP BY host, tag
- ",
- )
- .bind(user.id)
- .fetch_all(self.read_pool())
- .await?
- };
-
- res.sort();
-
- let mut status = RecordStatus::new();
-
- for i in &res {
- status.set_raw(HostId(i.0), i.1.clone(), i.2 as u64);
- }
-
- Ok(status)
- }
-}
diff --git a/crates/turtle/src/atuin_server/database/db/wrappers.rs b/crates/turtle/src/atuin_server/database/db/wrappers.rs
deleted file mode 100644
index 0315e331..00000000
--- a/crates/turtle/src/atuin_server/database/db/wrappers.rs
+++ /dev/null
@@ -1,32 +0,0 @@
-use crate::atuin_common::record::{EncryptedData, Host, Record};
-use sqlx::{Row, postgres::PgRow};
-
-pub(crate) struct DbRecord(pub Record<EncryptedData>);
-
-impl<'a> ::sqlx::FromRow<'a, PgRow> for DbRecord {
- fn from_row(row: &'a PgRow) -> ::sqlx::Result<Self> {
- let timestamp: i64 = row.try_get("timestamp")?;
- let idx: i64 = row.try_get("idx")?;
-
- let data = EncryptedData {
- data: row.try_get("data")?,
- content_encryption_key: row.try_get("cek")?,
- };
-
- Ok(Self(Record {
- id: row.try_get("client_id")?,
- host: Host::new(row.try_get("host")?),
- idx: idx as u64,
- timestamp: timestamp as u64,
- version: row.try_get("version")?,
- tag: row.try_get("tag")?,
- data,
- }))
- }
-}
-
-impl From<DbRecord> for Record<EncryptedData> {
- fn from(other: DbRecord) -> Self {
- other.0
- }
-}
diff --git a/crates/turtle/src/atuin_server/database/mod.rs b/crates/turtle/src/atuin_server/database/mod.rs
deleted file mode 100644
index 43fe5c3b..00000000
--- a/crates/turtle/src/atuin_server/database/mod.rs
+++ /dev/null
@@ -1,99 +0,0 @@
-pub(crate) mod db;
-pub(crate) mod models;
-
-use std::fmt::{Debug, Display};
-
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug)]
-pub(crate) enum DbError {
- NotFound,
- Other(eyre::Report),
-}
-
-impl Display for DbError {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- match self {
- Self::NotFound => write!(f, "Not found"),
- Self::Other(report) => write!(f, "Other: {report}"),
- }
- }
-}
-
-impl From<time::error::ComponentRange> for DbError {
- fn from(error: time::error::ComponentRange) -> Self {
- Self::Other(error.into())
- }
-}
-
-impl From<time::error::Error> for DbError {
- fn from(error: time::error::Error) -> Self {
- Self::Other(error.into())
- }
-}
-
-impl From<sqlx::Error> for DbError {
- fn from(error: sqlx::Error) -> Self {
- match error {
- sqlx::Error::RowNotFound => Self::NotFound,
- error => Self::Other(error.into()),
- }
- }
-}
-
-impl std::error::Error for DbError {}
-
-pub(crate) type DbResult<T> = Result<T, DbError>;
-
-#[derive(Debug, PartialEq)]
-pub(crate) enum DbType {
- Postgres,
- Unknown,
-}
-
-#[derive(Clone, Deserialize, Serialize)]
-pub(crate) struct DbSettings {
- pub(crate) db_uri: String,
-
- /// Optional URI for read replicas. If set, read-only queries will use this connection.
- pub(crate) read_db_uri: Option<String>,
-}
-
-impl DbSettings {
- pub(crate) fn db_type(&self) -> DbType {
- if self.db_uri.starts_with("postgres://") || self.db_uri.starts_with("postgresql://") {
- DbType::Postgres
- } else {
- DbType::Unknown
- }
- }
-}
-
-fn redact_db_uri(uri: &str) -> String {
- url::Url::parse(uri).map_or_else(
- |_| uri.to_string(),
- |mut url| {
- url.set_password(Some("****")).expect("should be possible");
- url.to_string()
- },
- )
-}
-
-// Do our best to redact passwords so they're not logged in the event of an error.
-impl Debug for DbSettings {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- if self.db_type() == DbType::Postgres {
- let redacted_uri = redact_db_uri(&self.db_uri);
- let redacted_read_uri = self.read_db_uri.as_ref().map(|uri| redact_db_uri(uri));
- f.debug_struct("DbSettings")
- .field("db_uri", &redacted_uri)
- .field("read_db_uri", &redacted_read_uri)
- .finish()
- } else {
- f.debug_struct("DbSettings")
- .field("db_uri", &self.db_uri)
- .field("read_db_uri", &self.read_db_uri)
- .finish()
- }
- }
-}
diff --git a/crates/turtle/src/atuin_server/database/models.rs b/crates/turtle/src/atuin_server/database/models.rs
deleted file mode 100644
index 3fa6f471..00000000
--- a/crates/turtle/src/atuin_server/database/models.rs
+++ /dev/null
@@ -1,5 +0,0 @@
-use uuid::Uuid;
-
-pub(crate) struct User {
- pub(crate) id: Uuid,
-}
generated by cgit v1.3.1 (git 2.54.0) at 2026-08-28 09:32:35 +0000