refactor(storage): add redis structs to storage impls (#1910)

This commit is contained in:
Sampras Lopes
2023-08-11 12:50:38 +05:30
committed by GitHub
parent 4311f2ce5e
commit 3e269663c3
9 changed files with 306 additions and 133 deletions

View File

@ -1,37 +1,13 @@
use async_bb8_diesel::{AsyncConnection, ConnectionError};
use bb8::{CustomizeConnection, PooledConnection};
use bb8::PooledConnection;
use diesel::PgConnection;
use error_stack::{IntoReport, ResultExt};
#[cfg(feature = "kms")]
use external_services::kms;
#[cfg(feature = "kms")]
use external_services::kms::decrypt::KmsDecrypt;
#[cfg(not(feature = "kms"))]
use masking::PeekInterface;
use crate::{configs::settings::Database, errors};
use crate::errors;
pub type PgPool = bb8::Pool<async_bb8_diesel::ConnectionManager<PgConnection>>;
pub type PgPooledConn = async_bb8_diesel::Connection<PgConnection>;
#[derive(Debug)]
struct TestTransaction;
#[async_trait::async_trait]
impl CustomizeConnection<PgPooledConn, ConnectionError> for TestTransaction {
#[allow(clippy::unwrap_used)]
async fn on_acquire(&self, conn: &mut PgPooledConn) -> Result<(), ConnectionError> {
use diesel::Connection;
conn.run(|conn| {
conn.begin_test_transaction().unwrap();
Ok(())
})
.await
}
}
#[allow(clippy::expect_used)]
pub async fn redis_connection(
conf: &crate::configs::settings::Settings,
@ -41,41 +17,6 @@ pub async fn redis_connection(
.expect("Failed to create Redis Connection Pool")
}
#[allow(clippy::expect_used)]
pub async fn diesel_make_pg_pool(
database: &Database,
test_transaction: bool,
#[cfg(feature = "kms")] kms_client: &kms::KmsClient,
) -> PgPool {
#[cfg(feature = "kms")]
let password = database
.password
.clone()
.decrypt_inner(kms_client)
.await
.expect("Failed to KMS decrypt database password");
#[cfg(not(feature = "kms"))]
let password = &database.password.peek();
let database_url = format!(
"postgres://{}:{}@{}:{}/{}",
database.username, password, database.host, database.port, database.dbname
);
let manager = async_bb8_diesel::ConnectionManager::<PgConnection>::new(database_url);
let mut pool = bb8::Pool::builder()
.max_size(database.pool_size)
.connection_timeout(std::time::Duration::from_secs(database.connection_timeout));
if test_transaction {
pool = pool.connection_customizer(Box::new(TestTransaction));
}
pool.build(manager)
.await
.expect("Failed to create PostgreSQL connection pool")
}
pub async fn pg_connection_read(
store: &crate::services::Store,
) -> errors::CustomResult<

View File

@ -11,7 +11,7 @@ use external_services::kms::{self, decrypt::KmsDecrypt};
#[cfg(not(feature = "kms"))]
use masking::PeekInterface;
use redis_interface::{errors as redis_errors, PubsubInterface, RedisValue};
use storage_impl::diesel as diesel_impl;
use storage_impl::{diesel as diesel_impl, DatabaseStore};
use tokio::sync::oneshot;
pub use self::{api::*, encryption::*};
@ -184,9 +184,9 @@ impl Store {
#[cfg(not(feature = "olap"))]
diesel_store: diesel_impl::store::Store::new(
#[cfg(not(feature = "kms"))]
&config.master_database.clone().into(),
config.master_database.clone().into(),
#[cfg(feature = "kms")]
&config
config
.master_database
.clone()
.decrypt_inner(kms_client)
@ -197,24 +197,26 @@ impl Store {
.await,
#[cfg(feature = "olap")]
diesel_store: diesel_impl::store::ReplicaStore::new(
#[cfg(not(feature = "kms"))]
&config.master_database.clone().into(),
#[cfg(feature = "kms")]
&config
.master_database
.clone()
.decrypt_inner(kms_client)
.await
.expect("Failed to decrypt master database"),
#[cfg(not(feature = "kms"))]
&config.replica_database.clone().into(),
#[cfg(feature = "kms")]
&config
.replica_database
.clone()
.decrypt_inner(kms_client)
.await
.expect("Failed to decrypt replica database"),
(
#[cfg(not(feature = "kms"))]
config.master_database.clone().into(),
#[cfg(feature = "kms")]
config
.master_database
.clone()
.decrypt_inner(kms_client)
.await
.expect("Failed to decrypt master database"),
#[cfg(not(feature = "kms"))]
config.replica_database.clone().into(),
#[cfg(feature = "kms")]
config
.replica_database
.clone()
.decrypt_inner(kms_client)
.await
.expect("Failed to decrypt replica database"),
),
test_transaction,
)
.await,