use crate::metrics::{
create_connection_use_time_metric, create_write_txn_duration_metric, ConnectionUseTimeMetric,
WriteTxnDurationMetric,
};
use crate::DatabaseIdentifier;
use sqlx::pool::PoolConnection;
use sqlx::{Pool, Sqlite, SqliteConnection, Transaction};
use std::ops::{Deref, DerefMut};
use std::time::Instant;
#[derive(Clone, Debug)]
pub struct DbRead<I: DatabaseIdentifier> {
pub(crate) pool: Pool<Sqlite>,
pub(crate) identifier: I,
use_time_metric: ConnectionUseTimeMetric,
write_txn_metric: WriteTxnDurationMetric,
}
impl<I: DatabaseIdentifier> DbRead<I> {
pub(crate) fn new(pool: Pool<Sqlite>, identifier: I) -> Self {
let use_time_metric = create_connection_use_time_metric(&identifier);
let write_txn_metric = create_write_txn_duration_metric(&identifier);
Self {
pool,
identifier,
use_time_metric,
write_txn_metric,
}
}
pub(crate) async fn timed_conn(&self) -> sqlx::Result<TimedConn> {
let conn = self.pool.acquire().await?;
Ok(TimedConn::new(conn, self.use_time_metric.clone()))
}
pub fn identifier(&self) -> &I {
&self.identifier
}
pub fn pool(&self) -> &Pool<Sqlite> {
&self.pool
}
pub async fn begin(&self) -> sqlx::Result<TxRead<I>> {
let tx = self.pool.begin().await?;
Ok(TxRead::new(tx, self.identifier.clone()))
}
}
#[derive(Clone, Debug)]
pub struct DbWrite<I: DatabaseIdentifier>(DbRead<I>);
impl<I: DatabaseIdentifier> DbWrite<I> {
pub(crate) fn new(pool: Pool<Sqlite>, identifier: I) -> Self {
Self(DbRead::new(pool, identifier))
}
pub fn identifier(&self) -> &I {
self.0.identifier()
}
pub fn pool(&self) -> &Pool<Sqlite> {
self.0.pool()
}
pub async fn begin(&self) -> sqlx::Result<TxWrite<I>> {
let tx = self.pool().begin().await?;
Ok(TxWrite {
inner: TxRead::new(tx, self.0.identifier.clone()),
start: std::time::Instant::now(),
write_txn_metric: self.0.write_txn_metric.clone(),
})
}
}
pub struct TxRead<I: DatabaseIdentifier> {
tx: Option<Transaction<'static, Sqlite>>,
identifier: I,
}
impl<I: DatabaseIdentifier> TxRead<I> {
pub(crate) fn new(tx: Transaction<'static, Sqlite>, identifier: I) -> Self {
Self {
tx: Some(tx),
identifier,
}
}
pub fn identifier(&self) -> &I {
&self.identifier
}
pub async fn close(mut self) -> sqlx::Result<()> {
self.tx
.take()
.expect("transaction already consumed")
.rollback()
.await
}
pub(crate) fn conn_mut(&mut self) -> &mut SqliteConnection {
self.tx.as_mut().expect("transaction already consumed")
}
pub(crate) fn tx_mut(&mut self) -> &mut Transaction<'static, Sqlite> {
self.tx.as_mut().expect("transaction already consumed")
}
}
pub struct TxWrite<I: DatabaseIdentifier> {
inner: TxRead<I>,
start: std::time::Instant,
write_txn_metric: WriteTxnDurationMetric,
}
impl<I: DatabaseIdentifier> TxWrite<I> {
pub fn identifier(&self) -> &I {
self.inner.identifier()
}
pub async fn commit(mut self) -> sqlx::Result<()> {
self.inner
.tx
.take()
.expect("transaction already consumed")
.commit()
.await?;
self.write_txn_metric
.record(self.start.elapsed().as_secs_f64());
Ok(())
}
pub async fn rollback(mut self) -> sqlx::Result<()> {
self.inner
.tx
.take()
.expect("transaction already consumed")
.rollback()
.await
}
pub(crate) fn conn_mut(&mut self) -> &mut SqliteConnection {
self.inner.conn_mut()
}
pub(crate) fn tx_mut(&mut self) -> &mut Transaction<'static, Sqlite> {
self.inner.tx_mut()
}
}
impl<I: DatabaseIdentifier> From<DbWrite<I>> for DbRead<I> {
fn from(write: DbWrite<I>) -> Self {
write.0
}
}
impl<I: DatabaseIdentifier> AsRef<DbRead<I>> for DbWrite<I> {
fn as_ref(&self) -> &DbRead<I> {
&self.0
}
}
impl<I: DatabaseIdentifier> AsRef<DbRead<I>> for DbRead<I> {
fn as_ref(&self) -> &DbRead<I> {
self
}
}
impl<I: DatabaseIdentifier> From<TxWrite<I>> for TxRead<I> {
fn from(write: TxWrite<I>) -> Self {
write.inner
}
}
impl<I: DatabaseIdentifier> AsRef<TxRead<I>> for TxWrite<I> {
fn as_ref(&self) -> &TxRead<I> {
&self.inner
}
}
impl<I: DatabaseIdentifier> AsMut<TxRead<I>> for TxWrite<I> {
fn as_mut(&mut self) -> &mut TxRead<I> {
&mut self.inner
}
}
pub(crate) struct TimedConn {
conn: PoolConnection<Sqlite>,
use_time_metric: ConnectionUseTimeMetric,
acquired_at: Instant,
}
impl TimedConn {
fn new(conn: PoolConnection<Sqlite>, use_time_metric: ConnectionUseTimeMetric) -> Self {
Self {
conn,
use_time_metric,
acquired_at: Instant::now(),
}
}
}
impl Deref for TimedConn {
type Target = SqliteConnection;
fn deref(&self) -> &Self::Target {
&self.conn
}
}
impl DerefMut for TimedConn {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.conn
}
}
impl Drop for TimedConn {
fn drop(&mut self) {
self.use_time_metric
.record(self.acquired_at.elapsed().as_secs_f64());
}
}