arminius_lib_db/

lib.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0
//! # Arminius Database Layer
//! The common interface for application code to execute
//! DB queries.
//!
//! In essence, this crate contains three types:
//!   - [`DefaultDb`]: executes queries outside of a transaction, can
//!     [start](`DefaultDb::begin_transaction`) a transaction.
//!   - [`TransactedDb`]: executes queries inside of a transaction, can
//!     [commit](`TransactedDb::commit`) and [roll back](`TransactedDb::rollback`) a transaction.
//!   - [`Database`]: a trait implemented by both [`DefaultDb`] and [`TransactedDb`].
//!
//! This crate does not provide any queries by itself. These are meant to be
//! added by other crates via extension traits on [`Database`].
//!
//! ## Writing queries
//! ### Conventions
//! Queries are meant to be defined in their own crate. The name of such crate
//! should adhere to the pattern `arminius-lib-db-<table>`.
//!
//! As a rule of thumb, one such crate should only be responsible for one
//! specific table. This allows for fine-grained control over what components
//! are allowed to execute which queries. However, there might be situations
//! where we deviate from this and have more than one table managed by an
//! extension crate, because they relate very close to each other (`room` and
//! `event` being an example for such). However, we should not end up having
//! multiple crates for one table.
//!
//! Such crates should not depend on other crates, with the following
//! exceptions:
//! - `arminius-lib-db`
//! - `arminius-lib-model`
//! - `sqlx`
//!
//! SELECT queries should return their results using one of the types in
//! `arminius-lib-model`, so that data is represented in a use case-agnostic way.
//! This allows queries to be reused in different situations.
//!
//! ### Registering queries via extension trait
//! Queries should be added as a trait:
//!
//! ```ignore
//! /// the trait that contains the queries
//! // This annotation will suggest imports as _ (preferred for extension traits).
//!  #[rust_analyzer::prefer_underscore_import]
//! trait FooQueries : Database {
//! 	// note that queries functions are not abstract
//! 	async fn create_foo(&mut self, ...) -> Result<arminius_lib_model::FooId> {
//! 		sqlx::query!("select ...").fetch_all(self.connection())
//! 	}
//!
//! 	// more queries methods here....
//! }
//! ```
//!
//! Make queries available to Database trait
//! via a blanket implementation:
//!
//! ```ignore
//! impl<T: Database> ExampleQueries for T {}
//! ```

pub mod query;
mod tests;

// TODO: remove those when db layer is done
use sqlx::{
	Acquire as _, ConnectOptions as _, Connection as _, PgPool, Postgres, Transaction,
	migrate::Migrator,
	pool::PoolConnection,
	postgres::{PgConnectOptions, PgListener, PgPoolOptions},
};
/// We re-export some SQLX types, so these can be used
pub use sqlx::{PgConnection as RawPgConnection, PgTransaction as RawPgTransaction};
use url::Url;

/// Provides access to a connection.
/// Should be called by extension traits but not from application code.
///
/// Usually we don't need to use this directly, so prefer underscore for imports.
#[rust_analyzer::prefer_underscore_import]
pub trait Database {
	/// Acquires a new connection from the pool.
	///
	/// [`TransactedDb::commit`]
	async fn connection(&mut self) -> Result<&mut RawPgConnection, DatabaseError>;

	/// Starts a new transaction, acquiring a fresh connection from the pool
	/// Once a transaction has been started, it is the callers'
	/// responsibility to make changes permanent by committing it.
	/// If the transaction is dropped without being committed, all changes
	/// are rolled back.
	async fn begin_transaction<'d>(&'d mut self) -> Result<TransactedDb<'d>, DatabaseError>;
}

/// Default implementation that works without a transaction
pub struct DefaultDb {
	/// db pool
	pool: PgPool,
	/// when starting a transaction,
	acquired: Option<PoolConnection<Postgres>>,
}

impl DefaultDb {
	/// Creates a new db instance.
	#[must_use]
	pub fn new(pool: PgPool) -> DefaultDb {
		DefaultDb { pool, acquired: None }
	}

	/// Connects to a postgresql database, using environment variables.
	/// If the `DATABASE_URL` variable is set, connection options are derived
	/// from it. Otherwise, refer to [PgConnectOptions::new_without_pgpass()]
	/// for the specific environment variables used for configuration.
	/// If the `ARMINIUS_LAZY` environment variable is set to `1` or `true`,
	/// a lazy connection pool is created, meaning that an actual connection
	/// attempt is deferred until the first query is executed.
	pub async fn from_env() -> Result<Self, DatabaseError> {
		let connect_options = if let Ok(url) = std::env::var("DATABASE_URL") {
			PgConnectOptions::from_url(
				&Url::parse(&url).map_err(|err| sqlx::Error::Configuration(err.into()))?,
			)?
		} else {
			// Fall back to the "standard" postgres variables (PGHOST, PGUSER,
			// etc). If these are not set, some hardcoded defaults are used.
			// See [`PgConnectOptions::new_without_pgpass()`] for details.
			PgConnectOptions::new()
		};
		// FIXME(#1377): Normalize the content of the variable.
		let lazy = std::env::var("ARMINIUS_LAZY")
			.is_ok_and(|var| var == "1" || var.to_lowercase() == "true");

		Self::connect(connect_options, PgPoolOptions::new(), lazy).await
	}

	/// Connect to a PostgreSQL database endpoint. Creates a connection pool
	/// with the given [PgConnectOptions] and [PgPoolOptions].
	pub async fn connect(
		pg_connect_options: PgConnectOptions,
		pg_pool_options: PgPoolOptions,
		lazy: bool,
	) -> Result<Self, DatabaseError> {
		let pool = if lazy {
			pg_pool_options.connect_lazy_with(pg_connect_options)
		} else {
			pg_pool_options.connect_with(pg_connect_options).await?
		};
		Ok(Self::new(pool))
	}

	/// Run the migrations defined in the `migrations` directory.
	pub async fn run_migrations(&self) -> Result<(), sqlx::Error> {
		sqlx::migrate!().run(&self.pool).await?;
		Ok(())
	}

	/// Run the migrations from the given `Migrator`
	pub async fn run_migrations_with(&self, migrator: &Migrator) -> Result<(), sqlx::Error> {
		migrator.run(&self.pool).await?;
		Ok(())
	}

	/// Create a `PgListener` using this instance's connection pool
	pub async fn listener(&self) -> Result<PgListener, DatabaseError> {
		Ok(PgListener::connect_with(&self.pool).await?)
	}

	/// Run a database connectivity check
	pub async fn ping(&self) -> Result<(), DatabaseError> {
		self.clone().connection().await?.ping().await?;
		Ok(())
	}
}

impl Database for DefaultDb {
	async fn connection(&mut self) -> Result<&mut RawPgConnection, DatabaseError> {
		Ok(match &mut self.acquired {
			Some(acquired) => acquired,
			acquired @ None => acquired.insert(self.pool.acquire().await?),
		})
	}

	async fn begin_transaction<'c>(&'c mut self) -> Result<TransactedDb<'c>, DatabaseError> {
		let transaction = self.pool.begin().await?;
		Ok(TransactedDb { transaction })
	}
}

impl Clone for DefaultDb {
	fn clone(&self) -> Self {
		Self { pool: self.pool.clone(), acquired: None }
	}
}

/// for whatever reason...
impl std::fmt::Debug for DefaultDb {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.debug_struct("DefaultDb").finish_non_exhaustive()
	}
}

/// Exposes the inner connection pool.
/// Unfortunately, this is still required because
/// some tests are still using sqlx directly.
#[cfg(feature = "expose_pool")]
impl AsRef<PgPool> for DefaultDb {
	fn as_ref(&self) -> &PgPool {
		&self.pool
	}
}

/// DB Queries inside a transaction
#[derive(Debug)]
pub struct TransactedDb<'c> {
	/// what shall I say... the transaction.
	transaction: Transaction<'c, Postgres>,
}

/// Operations only valid inside a transaction
impl<'c> TransactedDb<'c> {
	/// Attempts to commit the current transaction.
	pub async fn commit(self) -> Result<(), DatabaseError> {
		self.transaction.commit().await.map_err(DatabaseError::from)
	}

	/// Attempts to rollback the current transaction.
	pub async fn rollback(self) -> sqlx::Result<()> {
		self.transaction.rollback().await
	}
}

impl<'c> Database for TransactedDb<'c> {
	async fn connection(&mut self) -> Result<&mut RawPgConnection, DatabaseError> {
		Ok(self.transaction.acquire().await?)
	}

	/// Starts a nested transaction
	async fn begin_transaction<'d>(&'d mut self) -> Result<TransactedDb<'d>, DatabaseError> {
		let transaction = self.transaction.begin().await?;
		Ok(TransactedDb { transaction })
	}
}

/// A generic Error type to be used for functions that perform DB operations.
///
/// Note regarding sqlx Errors: Most sqlx errors are of technical nature
/// (like ConnectionPool exhausted, connection unexpectedly closed,
/// timeouts). These are likely not handled specifically by calling
/// code, and mostly just cause the operation to fail.
///
/// Some sqlx errors (such as ConstraintViolation) may actually be useful
/// for the caller. For example, in a signup process, we may very well want to
/// handle the case where an email address is already registered. In this case,
/// the recommended pattern is to create a new error type containing the errors
/// you want the caller to handle explicitly, and include a variant that
/// holds the DatabaseError. This makes it very clear to the caller what
/// application-level conditions may occur.
///
/// Example:
/// ```rust
/// #[derive(thiserror::Error, Debug)]
/// enum SignupError {
/// 	// This typically comes from a sqlx constraint violation, but we want
/// 	// the caller to handle this explicitly.
/// 	#[error("username already taken")]
/// 	UsernameAlreadyTaken,
///
/// 	// All other "uninteresting" sqlx errors will be signaled via the DatabaseError.
/// 	#[error(transparent)]
/// 	DatabaseError(arminius_lib_db::DatabaseError),
/// }
/// ```
///
/// Trivial query functions (like simple SELECT queries) that do not have any
/// non-technical error cases may of course return this type directly.
#[derive(thiserror::Error, Debug)]
pub enum DatabaseError {
	/// An sqlx error.
	#[error(transparent)]
	SqlxError(sqlx::Error),

	/// Unexpected content read from the database. E.g. a JSON field
	/// failed to deserialize, a field did not conform to a specific format
	/// (like an email address) etc
	#[error("Invalid database content: {0}")]
	InvalidData(String),
}

/// Allows to convenient convert to a `Result<T, DatabaseError::InvalidData>`.
///
/// Extension traits don't need to be used directly so prefer importing `as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait ResultExt<T> {
	/// Maps a [`Result<T, E>`] to [`Result<T, DatabaseError>`]
	fn invalid_data(self, details: &str) -> Result<T, DatabaseError>;
}

/// Transforms the `Result<T, E>` into a [`Result<T, DatabaseError>`].
impl<T, E: std::error::Error> ResultExt<T> for Result<T, E> {
	fn invalid_data(self, details: &str) -> Result<T, DatabaseError> {
		self.map_err(|err| {
			let message = format!("{details}: {err}");
			// Emitting the tracing event here ensures that the error gets
			// emitted from inside the query span, giving the correct context
			// for the error, most importantly the arguments to the query
			tracing::error!("Error in SQL data: {message}");
			DatabaseError::InvalidData(message)
		})
	}
}

impl From<sqlx::Error> for DatabaseError {
	fn from(value: sqlx::Error) -> Self {
		// Emitting the tracing event here ensures that the error gets emitted
		// from inside the query span, giving the correct context for the
		// error, most importantly the arguments to the query
		tracing::error!("Error executing SQL query: {value:?}");
		Self::SqlxError(value)
	}
}