caisin/sea-orm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

First metric and tracing implementation

Browse Source
This commit is contained in:
Marco Napetti 2021-12-09 15:12:41 +01:00
parent 18ed683712
commit 6b94673dd9
14 changed files with 317 additions and 56 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
Cargo.toml
View File

@ -27,7 +27,7 @@ async-trait = { version = "^0.1" }
chrono = { version = "^0", optional = true }
futures = { version = "^0.3" }
futures-util = { version = "^0.3" }
log = { version = "^0.4", optional = true }
tracing = "0.1"
rust_decimal = { version = "^1", optional = true }
sea-orm-macros = { version = "^0.4.1", path = "sea-orm-macros", optional = true }
sea-query = { version = "^0.19.1", features = ["thread-safe"] }
@ -36,8 +36,9 @@ serde = { version = "^1.0", features = ["derive"] }
serde_json = { version = "^1", optional = true }
sqlx = { version = "^0.5", optional = true }
uuid = { version = "0.8", features = ["serde", "v4"], optional = true }
ouroboros = "0.11"
ouroboros = "0.14"
url = "^2.2"
once_cell = "1.8"
[dev-dependencies]
smol = { version = "^1.2" }
@ -47,12 +48,12 @@ tokio = { version = "^1.6", features = ["full"] }
actix-rt = { version = "2.2.0" }
maplit = { version = "^1" }
rust_decimal_macros = { version = "^1" }
env_logger = { version = "^0.9" }
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
sea-orm = { path = ".", features = ["debug-print"] }
pretty_assertions = { version = "^0.7" }
[features]
debug-print = ["log"]
debug-print = []
default = [
"macros",
"mock",
@ -60,6 +61,8 @@ default = [
"with-chrono",
"with-rust_decimal",
"with-uuid",
"runtime-tokio-rustls",#debug
"sqlx-all",#debug
]
macros = ["sea-orm-macros"]
mock = []

8
src/database/db_connection.rs
View File

@ -3,6 +3,7 @@ use crate::{
StatementBuilder, TransactionError,
};
use sea_query::{MysqlQueryBuilder, PostgresQueryBuilder, QueryBuilder, SqliteQueryBuilder};
use tracing::instrument;
use std::{future::Future, pin::Pin};
use url::Url;
@ -49,6 +50,7 @@ pub enum DatabaseBackend {
/// The same as [DatabaseBackend] just shorter :)
pub type DbBackend = DatabaseBackend;
#[derive(Debug)]
pub(crate) enum InnerConnection {
#[cfg(feature = "sqlx-mysql")]
MySql(PoolConnection<sqlx::MySql>),
@ -104,6 +106,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
}
}
#[instrument(level = "trace")]
async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
match self {
#[cfg(feature = "sqlx-mysql")]
@ -118,6 +121,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
}
}
#[instrument(level = "trace")]
async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
match self {
#[cfg(feature = "sqlx-mysql")]
@ -132,6 +136,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
}
}
#[instrument(level = "trace")]
async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
match self {
#[cfg(feature = "sqlx-mysql")]
@ -146,6 +151,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
}
}
#[instrument(level = "trace")]
fn stream(
&'a self,
stmt: Statement,
@ -167,6 +173,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
})
}
#[instrument(level = "trace")]
async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
match self {
#[cfg(feature = "sqlx-mysql")]
@ -185,6 +192,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseConnection {
/// Execute the function inside a transaction.
/// If the function returns an error, the transaction will be rolled back. If it does not return an error, the transaction will be committed.
#[instrument(level = "trace", skip(_callback))]
async fn transaction<F, T, E>(&self, _callback: F) -> Result<T, TransactionError<E>>
where
F: for<'c> FnOnce(

6
src/database/mock.rs
View File

@ -4,6 +4,7 @@ use crate::{
Statement,
};
use sea_query::{Value, ValueType, Values};
use tracing::instrument;
use std::{collections::BTreeMap, sync::Arc};
/// Defines a Mock database suitable for testing
@ -89,6 +90,7 @@ impl MockDatabase {
}
impl MockDatabaseTrait for MockDatabase {
#[instrument(level = "trace")]
fn execute(&mut self, counter: usize, statement: Statement) -> Result<ExecResult, DbErr> {
if let Some(transaction) = &mut self.transaction {
transaction.push(statement);
@ -104,6 +106,7 @@ impl MockDatabaseTrait for MockDatabase {
}
}
#[instrument(level = "trace")]
fn query(&mut self, counter: usize, statement: Statement) -> Result<Vec<QueryResult>, DbErr> {
if let Some(transaction) = &mut self.transaction {
transaction.push(statement);
@ -122,6 +125,7 @@ impl MockDatabaseTrait for MockDatabase {
}
}
#[instrument(level = "trace")]
fn begin(&mut self) {
if self.transaction.is_some() {
self.transaction
@ -133,6 +137,7 @@ impl MockDatabaseTrait for MockDatabase {
}
}
#[instrument(level = "trace")]
fn commit(&mut self) {
if self.transaction.is_some() {
if self.transaction.as_mut().unwrap().commit(self.db_backend) {
@ -144,6 +149,7 @@ impl MockDatabaseTrait for MockDatabase {
}
}
#[instrument(level = "trace")]
fn rollback(&mut self) {
if self.transaction.is_some() {
if self.transaction.as_mut().unwrap().rollback(self.db_backend) {

2
src/database/mod.rs
View File

@ -14,6 +14,7 @@ pub use db_connection::*;
pub use mock::*;
pub use statement::*;
pub use stream::*;
use tracing::instrument;
pub use transaction::*;
use crate::DbErr;
@ -42,6 +43,7 @@ pub struct ConnectOptions {
impl Database {
/// Method to create a [DatabaseConnection] on a database
#[instrument(level = "trace", skip(opt))]
pub async fn connect<C>(opt: C) -> Result<DatabaseConnection, DbErr>
where
C: Into<ConnectOptions>,

92
src/database/stream/query.rs
View File

@ -12,6 +12,8 @@ use futures::TryStreamExt;
#[cfg(feature = "sqlx-dep")]
use sqlx::{pool::PoolConnection, Executor};
use tracing::instrument;
use crate::{DbErr, InnerConnection, QueryResult, Statement};
/// Creates a stream from a [QueryResult]
@ -59,41 +61,71 @@ impl std::fmt::Debug for QueryStream {
}
impl QueryStream {
#[instrument(level = "trace")]
fn build(stmt: Statement, conn: InnerConnection) -> QueryStream {
QueryStreamBuilder {
stmt,
conn,
stream_builder: |conn, stmt| match conn {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(c) => {
let query = crate::driver::sqlx_mysql::sqlx_query(stmt);
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
stream_builder: |conn, stmt| {
match conn {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(c) => {
let query = crate::driver::sqlx_mysql::sqlx_query(stmt);
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(c) => {
let query = crate::driver::sqlx_postgres::sqlx_query(stmt);
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(c) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(stmt);
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "mock")]
InnerConnection::Mock(c) => c.fetch(stmt),
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(c) => {
let query = crate::driver::sqlx_postgres::sqlx_query(stmt);
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(c) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(stmt);
Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
}
#[cfg(feature = "mock")]
InnerConnection::Mock(c) => c.fetch(stmt),
},
}
}
.build()
}

45
src/database/stream/transaction.rs
View File

@ -11,6 +11,8 @@ use sqlx::Executor;
use futures::lock::MutexGuard;
use tracing::instrument;
use crate::{DbErr, InnerConnection, QueryResult, Statement};
/// `TransactionStream` cannot be used in a `transaction` closure as it does not impl `Send`.
@ -31,6 +33,7 @@ impl<'a> std::fmt::Debug for TransactionStream<'a> {
}
impl<'a> TransactionStream<'a> {
#[instrument(level = "trace")]
pub(crate) async fn build(
conn: MutexGuard<'a, InnerConnection>,
stmt: Statement,
@ -44,32 +47,56 @@ impl<'a> TransactionStream<'a> {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(c) => {
let query = crate::driver::sqlx_mysql::sqlx_query(stmt);
Box::pin(
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
) as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>;
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(c) => {
let query = crate::driver::sqlx_postgres::sqlx_query(stmt);
Box::pin(
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
) as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>;
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(c) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(stmt);
Box::pin(
let _start = std::time::SystemTime::now();
let res = Box::pin(
c.fetch(query)
.map_ok(Into::into)
.map_err(crate::sqlx_error_to_query_err),
)
as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>
) as Pin<Box<dyn Stream<Item = Result<QueryResult, DbErr>>>>;
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "mock")]
InnerConnection::Mock(c) => c.fetch(stmt),

45
src/database/transaction.rs
View File

@ -7,6 +7,7 @@ use crate::{sqlx_error_to_exec_err, sqlx_error_to_query_err};
use futures::lock::Mutex;
#[cfg(feature = "sqlx-dep")]
use sqlx::{pool::PoolConnection, TransactionManager};
use tracing::instrument;
use std::{future::Future, pin::Pin, sync::Arc};
// a Transaction is just a sugar for a connection where START TRANSACTION has been executed
@ -66,6 +67,7 @@ impl DatabaseTransaction {
Self::begin(Arc::new(Mutex::new(InnerConnection::Mock(inner))), backend).await
}
#[instrument(level = "trace")]
async fn begin(
conn: Arc<Mutex<InnerConnection>>,
backend: DbBackend,
@ -104,6 +106,7 @@ impl DatabaseTransaction {
/// Runs a transaction to completion returning an rolling back the transaction on
/// encountering an error if it fails
#[instrument(level = "trace", skip(callback))]
pub(crate) async fn run<F, T, E>(self, callback: F) -> Result<T, TransactionError<E>>
where
F: for<'b> FnOnce(
@ -125,6 +128,7 @@ impl DatabaseTransaction {
}
/// Commit a transaction atomically
#[instrument(level = "trace")]
pub async fn commit(mut self) -> Result<(), DbErr> {
self.open = false;
match *self.conn.lock().await {
@ -155,6 +159,7 @@ impl DatabaseTransaction {
}
/// rolls back a transaction in case error are encountered during the operation
#[instrument(level = "trace")]
pub async fn rollback(mut self) -> Result<(), DbErr> {
self.open = false;
match *self.conn.lock().await {
@ -185,6 +190,7 @@ impl DatabaseTransaction {
}
// the rollback is queued and will be performed on next async operation, like returning the connection to the pool
#[instrument(level = "trace")]
fn start_rollback(&mut self) {
if self.open {
if let Some(mut conn) = self.conn.try_lock() {
@ -229,6 +235,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
self.backend
}
#[instrument(level = "trace")]
async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
debug_print!("{}", stmt);
@ -236,17 +243,44 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
#[cfg(feature = "sqlx-mysql")]
InnerConnection::MySql(conn) => {
let query = crate::driver::sqlx_mysql::sqlx_query(&stmt);
query.execute(conn).await.map(Into::into)
let _start = std::time::SystemTime::now();
let res = query.execute(conn).await.map(Into::into);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-postgres")]
InnerConnection::Postgres(conn) => {
let query = crate::driver::sqlx_postgres::sqlx_query(&stmt);
query.execute(conn).await.map(Into::into)
let _start = std::time::SystemTime::now();
let res = query.execute(conn).await.map(Into::into);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "sqlx-sqlite")]
InnerConnection::Sqlite(conn) => {
let query = crate::driver::sqlx_sqlite::sqlx_query(&stmt);
query.execute(conn).await.map(Into::into)
let _start = std::time::SystemTime::now();
let res = query.execute(conn).await.map(Into::into);
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
}
#[cfg(feature = "mock")]
InnerConnection::Mock(conn) => return conn.execute(stmt),
@ -255,6 +289,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
_res.map_err(sqlx_error_to_exec_err)
}
#[instrument(level = "trace")]
async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
debug_print!("{}", stmt);
@ -285,6 +320,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
}
}
#[instrument(level = "trace")]
async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
debug_print!("{}", stmt);
@ -320,6 +356,7 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
_res.map_err(sqlx_error_to_query_err)
}
#[instrument(level = "trace")]
fn stream(
&'a self,
stmt: Statement,
@ -329,12 +366,14 @@ impl<'a> ConnectionTrait<'a> for DatabaseTransaction {
)
}
#[instrument(level = "trace")]
async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
DatabaseTransaction::begin(Arc::clone(&self.conn), self.backend).await
}
/// Execute the function inside a transaction.
/// If the function returns an error, the transaction will be rolled back. If it does not return an error, the transaction will be committed.
#[instrument(level = "trace", skip(_callback))]
async fn transaction<F, T, E>(&self, _callback: F) -> Result<T, TransactionError<E>>
where
F: for<'c> FnOnce(

9
src/driver/mock.rs
View File

@ -3,6 +3,7 @@ use crate::{
Statement, Transaction,
};
use futures::Stream;
use tracing::instrument;
use std::{
fmt::Debug,
pin::Pin,
@ -69,6 +70,7 @@ impl MockDatabaseConnector {
/// Cpnnect to the [MockDatabase]
#[allow(unused_variables)]
#[instrument(level = "trace")]
pub async fn connect(string: &str) -> Result<DatabaseConnection, DbErr> {
macro_rules! connect_mock_db {
( $syntax: expr ) => {
@ -117,6 +119,7 @@ impl MockDatabaseConnection {
}
/// Execute the SQL statement in the [MockDatabase]
#[instrument(level = "trace")]
pub fn execute(&self, statement: Statement) -> Result<ExecResult, DbErr> {
debug_print!("{}", statement);
let counter = self.execute_counter.fetch_add(1, Ordering::SeqCst);
@ -124,6 +127,7 @@ impl MockDatabaseConnection {
}
/// Return one [QueryResult] if the query was successful
#[instrument(level = "trace")]
pub fn query_one(&self, statement: Statement) -> Result<Option<QueryResult>, DbErr> {
debug_print!("{}", statement);
let counter = self.query_counter.fetch_add(1, Ordering::SeqCst);
@ -132,6 +136,7 @@ impl MockDatabaseConnection {
}
/// Return all [QueryResult]s if the query was successful
#[instrument(level = "trace")]
pub fn query_all(&self, statement: Statement) -> Result<Vec<QueryResult>, DbErr> {
debug_print!("{}", statement);
let counter = self.query_counter.fetch_add(1, Ordering::SeqCst);
@ -139,6 +144,7 @@ impl MockDatabaseConnection {
}
/// Return [QueryResult]s from a multi-query operation
#[instrument(level = "trace")]
pub fn fetch(
&self,
statement: &Statement,
@ -150,16 +156,19 @@ impl MockDatabaseConnection {
}
/// Create a statement block of SQL statements that execute together.
#[instrument(level = "trace")]
pub fn begin(&self) {
self.mocker.lock().unwrap().begin()
}
/// Commit a transaction atomically to the database
#[instrument(level = "trace")]
pub fn commit(&self) {
self.mocker.lock().unwrap().commit()
}
/// Roll back a faulty transaction
#[instrument(level = "trace")]
pub fn rollback(&self) {
self.mocker.lock().unwrap().rollback()
}

41
src/driver/sqlx_mysql.rs
View File

@ -7,6 +7,7 @@ use sqlx::{
sea_query::sea_query_driver_mysql!();
use sea_query_driver_mysql::bind_query;
use tracing::instrument;
use crate::{
debug_print, error::*, executor::*, ConnectOptions, DatabaseConnection, DatabaseTransaction,
@ -32,6 +33,7 @@ impl SqlxMySqlConnector {
}
/// Add configuration options for the MySQL database
#[instrument(level = "trace")]
pub async fn connect(options: ConnectOptions) -> Result<DatabaseConnection, DbErr> {
let mut opt = options
.url
@ -59,15 +61,25 @@ impl SqlxMySqlConnector {
impl SqlxMySqlPoolConnection {
/// Execute a [Statement] on a MySQL backend
#[instrument(level = "trace")]
pub async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.execute(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.execute(conn).await {
Ok(res) => Ok(res.into()),
Err(err) => Err(sqlx_error_to_exec_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Exec(
"Failed to acquire connection from pool.".to_owned(),
@ -76,18 +88,28 @@ impl SqlxMySqlPoolConnection {
}
/// Get one result from a SQL query. Returns [Option::None] if no match was found
#[instrument(level = "trace")]
pub async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_one(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_one(conn).await {
Ok(row) => Ok(Some(row.into())),
Err(err) => match err {
sqlx::Error::RowNotFound => Ok(None),
_ => Err(DbErr::Query(err.to_string())),
},
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -96,15 +118,25 @@ impl SqlxMySqlPoolConnection {
}
/// Get the results of a query returning them as a Vec<[QueryResult]>
#[instrument(level = "trace")]
pub async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_all(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_all(conn).await {
Ok(rows) => Ok(rows.into_iter().map(|r| r.into()).collect()),
Err(err) => Err(sqlx_error_to_query_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -113,6 +145,7 @@ impl SqlxMySqlPoolConnection {
}
/// Stream the results of executing a SQL query
#[instrument(level = "trace")]
pub async fn stream(&self, stmt: Statement) -> Result<QueryStream, DbErr> {
debug_print!("{}", stmt);
@ -126,6 +159,7 @@ impl SqlxMySqlPoolConnection {
}
/// Bundle a set of SQL statements that execute together.
#[instrument(level = "trace")]
pub async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
if let Ok(conn) = self.pool.acquire().await {
DatabaseTransaction::new_mysql(conn).await
@ -137,6 +171,7 @@ impl SqlxMySqlPoolConnection {
}
/// Create a MySQL transaction
#[instrument(level = "trace", skip(callback))]
pub async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
where
F: for<'b> FnOnce(

41
src/driver/sqlx_postgres.rs
View File

@ -7,6 +7,7 @@ use sqlx::{
sea_query::sea_query_driver_postgres!();
use sea_query_driver_postgres::bind_query;
use tracing::instrument;
use crate::{
debug_print, error::*, executor::*, ConnectOptions, DatabaseConnection, DatabaseTransaction,
@ -32,6 +33,7 @@ impl SqlxPostgresConnector {
}
/// Add configuration options for the MySQL database
#[instrument(level = "trace")]
pub async fn connect(options: ConnectOptions) -> Result<DatabaseConnection, DbErr> {
let mut opt = options
.url
@ -59,15 +61,25 @@ impl SqlxPostgresConnector {
impl SqlxPostgresPoolConnection {
/// Execute a [Statement] on a PostgreSQL backend
#[instrument(level = "trace")]
pub async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.execute(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.execute(conn).await {
Ok(res) => Ok(res.into()),
Err(err) => Err(sqlx_error_to_exec_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Exec(
"Failed to acquire connection from pool.".to_owned(),
@ -76,18 +88,28 @@ impl SqlxPostgresPoolConnection {
}
/// Get one result from a SQL query. Returns [Option::None] if no match was found
#[instrument(level = "trace")]
pub async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_one(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_one(conn).await {
Ok(row) => Ok(Some(row.into())),
Err(err) => match err {
sqlx::Error::RowNotFound => Ok(None),
_ => Err(DbErr::Query(err.to_string())),
},
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -96,15 +118,25 @@ impl SqlxPostgresPoolConnection {
}
/// Get the results of a query returning them as a Vec<[QueryResult]>
#[instrument(level = "trace")]
pub async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_all(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_all(conn).await {
Ok(rows) => Ok(rows.into_iter().map(|r| r.into()).collect()),
Err(err) => Err(sqlx_error_to_query_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -113,6 +145,7 @@ impl SqlxPostgresPoolConnection {
}
/// Stream the results of executing a SQL query
#[instrument(level = "trace")]
pub async fn stream(&self, stmt: Statement) -> Result<QueryStream, DbErr> {
debug_print!("{}", stmt);
@ -126,6 +159,7 @@ impl SqlxPostgresPoolConnection {
}
/// Bundle a set of SQL statements that execute together.
#[instrument(level = "trace")]
pub async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
if let Ok(conn) = self.pool.acquire().await {
DatabaseTransaction::new_postgres(conn).await
@ -137,6 +171,7 @@ impl SqlxPostgresPoolConnection {
}
/// Create a PostgreSQL transaction
#[instrument(level = "trace", skip(callback))]
pub async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
where
F: for<'b> FnOnce(

41
src/driver/sqlx_sqlite.rs
View File

@ -7,6 +7,7 @@ use sqlx::{
sea_query::sea_query_driver_sqlite!();
use sea_query_driver_sqlite::bind_query;
use tracing::instrument;
use crate::{
debug_print, error::*, executor::*, ConnectOptions, DatabaseConnection, DatabaseTransaction,
@ -32,6 +33,7 @@ impl SqlxSqliteConnector {
}
/// Add configuration options for the SQLite database
#[instrument(level = "trace")]
pub async fn connect(options: ConnectOptions) -> Result<DatabaseConnection, DbErr> {
let mut options = options;
let mut opt = options
@ -63,15 +65,25 @@ impl SqlxSqliteConnector {
impl SqlxSqlitePoolConnection {
/// Execute a [Statement] on a SQLite backend
#[instrument(level = "trace")]
pub async fn execute(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.execute(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.execute(conn).await {
Ok(res) => Ok(res.into()),
Err(err) => Err(sqlx_error_to_exec_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Exec(
"Failed to acquire connection from pool.".to_owned(),
@ -80,18 +92,28 @@ impl SqlxSqlitePoolConnection {
}
/// Get one result from a SQL query. Returns [Option::None] if no match was found
#[instrument(level = "trace")]
pub async fn query_one(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_one(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_one(conn).await {
Ok(row) => Ok(Some(row.into())),
Err(err) => match err {
sqlx::Error::RowNotFound => Ok(None),
_ => Err(DbErr::Query(err.to_string())),
},
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -100,15 +122,25 @@ impl SqlxSqlitePoolConnection {
}
/// Get the results of a query returning them as a Vec<[QueryResult]>
#[instrument(level = "trace")]
pub async fn query_all(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
debug_print!("{}", stmt);
let query = sqlx_query(&stmt);
if let Ok(conn) = &mut self.pool.acquire().await {
match query.fetch_all(conn).await {
let _start = std::time::SystemTime::now();
let res = match query.fetch_all(conn).await {
Ok(rows) => Ok(rows.into_iter().map(|r| r.into()).collect()),
Err(err) => Err(sqlx_error_to_query_err(err)),
};
if let Some(callback) = crate::metric::get_callback() {
let info = crate::metric::Info {
elapsed: _start.elapsed().unwrap_or_default(),
statement: &stmt,
};
callback(&info);
}
res
} else {
Err(DbErr::Query(
"Failed to acquire connection from pool.".to_owned(),
@ -117,6 +149,7 @@ impl SqlxSqlitePoolConnection {
}
/// Stream the results of executing a SQL query
#[instrument(level = "trace")]
pub async fn stream(&self, stmt: Statement) -> Result<QueryStream, DbErr> {
debug_print!("{}", stmt);
@ -130,6 +163,7 @@ impl SqlxSqlitePoolConnection {
}
/// Bundle a set of SQL statements that execute together.
#[instrument(level = "trace")]
pub async fn begin(&self) -> Result<DatabaseTransaction, DbErr> {
if let Ok(conn) = self.pool.acquire().await {
DatabaseTransaction::new_sqlite(conn).await
@ -141,6 +175,7 @@ impl SqlxSqlitePoolConnection {
}
/// Create a MySQL transaction
#[instrument(level = "trace", skip(callback))]
pub async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
where
F: for<'b> FnOnce(

2
src/lib.rs
View File

@ -284,6 +284,8 @@ pub mod schema;
#[cfg(feature = "macros")]
pub mod tests_cfg;
mod util;
/// Holds types and methods to perform metric collection
pub mod metric;
pub use database::*;
pub use driver::*;

28
src/metric.rs Normal file
View File

@ -0,0 +1,28 @@
use std::time::Duration;
use once_cell::sync::OnceCell;
type Callback = Box<dyn Fn(&Info<'_>) + Send + Sync>;
static METRIC: OnceCell<Callback> = OnceCell::new();
#[derive(Debug)]
/// Query execution infos
pub struct Info<'a> {
/// Query executiuon duration
pub elapsed: Duration,
/// Query data
pub statement: &'a crate::Statement,
}
/// Sets a new metric callback, returning it if already set
pub fn set_callback<F>(callback: F) -> Result<(), Callback>
where
F: Fn(&Info<'_>) + Send + Sync + 'static,
{
METRIC.set(Box::new(callback))
}
pub(crate) fn get_callback() -> Option<&'static Callback> {
METRIC.get()
}

2
src/util.rs
View File

@ -15,7 +15,7 @@
#[macro_export]
#[cfg(feature = "debug-print")]
macro_rules! debug_print {
($( $args:expr ),*) => { log::debug!( $( $args ),* ); }
($( $args:expr ),*) => { tracing::debug!( $( $args ),* ); }
}
#[macro_export]