Use the new RepositoryFactory everywhere

Author: sandhose

crates/cli/src/commands/debug.rs

@@ -11,6 +11,7 @@ use figment::Figment;
 use mas_config::{
     ConfigurationSection, ConfigurationSectionExt, DatabaseConfig, MatrixConfig, PolicyConfig,
 };
+use mas_storage_pg::PgRepositoryFactory;
 use tracing::{info, info_span};
 
 use crate::util::{
@@ -48,7 +49,8 @@ impl Options {
                 if with_dynamic_data {
                     let database_config = DatabaseConfig::extract(figment)?;
                     let pool = database_pool_from_config(&database_config).await?;
-                    load_policy_factory_dynamic_data(&policy_factory, &pool).await?;
+                    let repository_factory = PgRepositoryFactory::new(pool.clone());
+                    load_policy_factory_dynamic_data(&policy_factory, &repository_factory).await?;
                 }
 
                 let _instance = policy_factory.instantiate().await?;

crates/cli/src/commands/server.rs

@@ -134,7 +134,7 @@ impl Options {
 
         load_policy_factory_dynamic_data_continuously(
             &policy_factory,
-            &pool,
+            PgRepositoryFactory::new(pool.clone()).boxed(),
             shutdown.soft_shutdown_token(),
             shutdown.task_tracker(),
         )
@@ -172,7 +172,7 @@ impl Options {
 
             info!("Starting task worker");
             mas_tasks::init(
-                &pool,
+                PgRepositoryFactory::new(pool.clone()),
                 &mailer,
                 homeserver_connection.clone(),
                 url_builder.clone(),
@@ -193,7 +193,7 @@ impl Options {
         // Initialize the activity tracker
         // Activity is flushed every minute
         let activity_tracker = ActivityTracker::new(
-            pool.clone(),
+            PgRepositoryFactory::new(pool.clone()).boxed(),
             Duration::from_secs(60),
             shutdown.task_tracker(),
             shutdown.soft_shutdown_token(),
@@ -215,7 +215,7 @@ impl Options {
         limiter.start();
 
         let graphql_schema = mas_handlers::graphql_schema(
-            &pool,
+            PgRepositoryFactory::new(pool.clone()).boxed(),
             &policy_factory,
             homeserver_connection.clone(),
             site_config.clone(),

crates/cli/src/commands/worker.rs

@@ -10,6 +10,7 @@ use clap::Parser;
 use figment::Figment;
 use mas_config::{AppConfig, ConfigurationSection};
 use mas_router::UrlBuilder;
+use mas_storage_pg::PgRepositoryFactory;
 use tracing::{info, info_span};
 
 use crate::{
@@ -63,7 +64,7 @@ impl Options {
 
         info!("Starting task scheduler");
         mas_tasks::init(
-            &pool,
+            PgRepositoryFactory::new(pool.clone()),
             &mailer,
             conn,
             url_builder,

crates/cli/src/util.rs

@@ -20,8 +20,7 @@ use mas_matrix::{HomeserverConnection, ReadOnlyHomeserverConnection};
 use mas_matrix_synapse::SynapseConnection;
 use mas_policy::PolicyFactory;
 use mas_router::UrlBuilder;
-use mas_storage::RepositoryAccess;
-use mas_storage_pg::PgRepository;
+use mas_storage::{BoxRepositoryFactory, RepositoryAccess, RepositoryFactory};
 use mas_templates::{SiteConfigExt, Templates};
 use sqlx::{
     ConnectOptions, Executor, PgConnection, PgPool,
@@ -400,14 +399,13 @@ pub async fn database_connection_from_config_with_options(
 // XXX: this could be put somewhere else?
 pub async fn load_policy_factory_dynamic_data_continuously(
     policy_factory: &Arc<PolicyFactory>,
-    pool: &PgPool,
+    repository_factory: BoxRepositoryFactory,
     cancellation_token: CancellationToken,
     task_tracker: &TaskTracker,
 ) -> Result<(), anyhow::Error> {
     let policy_factory = policy_factory.clone();
-    let pool = pool.clone();
 
-    load_policy_factory_dynamic_data(&policy_factory, &pool).await?;
+    load_policy_factory_dynamic_data(&policy_factory, &*repository_factory).await?;
 
     task_tracker.spawn(async move {
         let mut interval = tokio::time::interval(Duration::from_secs(60));
@@ -420,7 +418,9 @@ pub async fn load_policy_factory_dynamic_data_continuously(
                 _ = interval.tick() => {}
             }
 
-            if let Err(err) = load_policy_factory_dynamic_data(&policy_factory, &pool).await {
+            if let Err(err) =
+                load_policy_factory_dynamic_data(&policy_factory, &*repository_factory).await
+            {
                 tracing::error!(
                     error = ?err,
                     "Failed to load policy factory dynamic data"
@@ -438,9 +438,10 @@ pub async fn load_policy_factory_dynamic_data_continuously(
 #[tracing::instrument(name = "policy.load_dynamic_data", skip_all)]
 pub async fn load_policy_factory_dynamic_data(
     policy_factory: &PolicyFactory,
-    pool: &PgPool,
+    repository_factory: &(dyn RepositoryFactory + Send + Sync),
 ) -> Result<(), anyhow::Error> {
-    let mut repo = PgRepository::from_pool(pool)
+    let mut repo = repository_factory
+        .create()
         .await
         .context("Failed to acquire database connection")?;
 

crates/handlers/src/activity_tracker/mod.rs

@@ -11,8 +11,7 @@ use std::net::IpAddr;
 
 use chrono::{DateTime, Utc};
 use mas_data_model::{BrowserSession, CompatSession, Session};
-use mas_storage::Clock;
-use sqlx::PgPool;
+use mas_storage::{BoxRepositoryFactory, Clock};
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
 use ulid::Ulid;
 
@@ -61,12 +60,12 @@ impl ActivityTracker {
     /// time, when the cancellation token is cancelled.
     #[must_use]
     pub fn new(
-        pool: PgPool,
+        repository_factory: BoxRepositoryFactory,
         flush_interval: std::time::Duration,
         task_tracker: &TaskTracker,
         cancellation_token: CancellationToken,
     ) -> Self {
-        let worker = Worker::new(pool);
+        let worker = Worker::new(repository_factory);
         let (sender, receiver) = tokio::sync::mpsc::channel(MESSAGE_QUEUE_SIZE);
         let tracker = ActivityTracker { channel: sender };
 

crates/handlers/src/activity_tracker/worker.rs

@@ -7,12 +7,13 @@
 use std::{collections::HashMap, net::IpAddr};
 
 use chrono::{DateTime, Utc};
-use mas_storage::{RepositoryAccess, RepositoryError, user::BrowserSessionRepository};
+use mas_storage::{
+    BoxRepositoryFactory, RepositoryAccess, RepositoryError, user::BrowserSessionRepository,
+};
 use opentelemetry::{
     Key, KeyValue,
     metrics::{Counter, Gauge, Histogram},
 };
-use sqlx::PgPool;
 use tokio_util::sync::CancellationToken;
 use ulid::Ulid;
 
@@ -43,15 +44,15 @@ struct ActivityRecord {
 
 /// Handles writing activity records to the database.
 pub struct Worker {
-    pool: PgPool,
+    repository_factory: BoxRepositoryFactory,
     pending_records: HashMap<(SessionKind, Ulid), ActivityRecord>,
     pending_records_gauge: Gauge<u64>,
     message_counter: Counter<u64>,
     flush_time_histogram: Histogram<u64>,
 }
 
 impl Worker {
-    pub(crate) fn new(pool: PgPool) -> Self {
+    pub(crate) fn new(repository_factory: BoxRepositoryFactory) -> Self {
         let message_counter = METER
             .u64_counter("mas.activity_tracker.messages")
             .with_description("The number of messages received by the activity tracker")
@@ -89,7 +90,7 @@ impl Worker {
         pending_records_gauge.record(0, &[]);
 
         Self {
-            pool,
+            repository_factory,
             pending_records: HashMap::with_capacity(MAX_PENDING_RECORDS),
             pending_records_gauge,
             message_counter,
@@ -218,11 +219,7 @@ impl Worker {
     #[tracing::instrument(name = "activity_tracker.flush", skip(self))]
     async fn try_flush(&mut self) -> Result<(), RepositoryError> {
         let pending_records = &self.pending_records;
-
-        let mut repo = mas_storage_pg::PgRepository::from_pool(&self.pool)
-            .await
-            .map_err(RepositoryError::from_error)?
-            .boxed();
+        let mut repo = self.repository_factory.create().await?;
 
         let mut browser_sessions = Vec::new();
         let mut oauth2_sessions = Vec::new();

crates/handlers/src/graphql/mod.rs

@@ -32,12 +32,12 @@ use mas_data_model::{BrowserSession, Session, SiteConfig, User};
 use mas_matrix::HomeserverConnection;
 use mas_policy::{InstantiateError, Policy, PolicyFactory};
 use mas_router::UrlBuilder;
-use mas_storage::{BoxClock, BoxRepository, BoxRng, Clock, RepositoryError, SystemClock};
-use mas_storage_pg::PgRepository;
+use mas_storage::{
+    BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng, Clock, RepositoryError, SystemClock,
+};
 use opentelemetry_semantic_conventions::trace::{GRAPHQL_DOCUMENT, GRAPHQL_OPERATION_NAME};
 use rand::{SeedableRng, thread_rng};
 use rand_chacha::ChaChaRng;
-use sqlx::PgPool;
 use state::has_session_ended;
 use tracing::{Instrument, info_span};
 use ulid::Ulid;
@@ -69,7 +69,7 @@ pub struct ExtraRouterParameters {
 }
 
 struct GraphQLState {
-    pool: PgPool,
+    repository_factory: BoxRepositoryFactory,
     homeserver_connection: Arc<dyn HomeserverConnection>,
     policy_factory: Arc<PolicyFactory>,
     site_config: SiteConfig,
@@ -81,11 +81,7 @@ struct GraphQLState {
 #[async_trait::async_trait]
 impl state::State for GraphQLState {
     async fn repository(&self) -> Result<BoxRepository, RepositoryError> {
-        let repo = PgRepository::from_pool(&self.pool)
-            .await
-            .map_err(RepositoryError::from_error)?;
-
-        Ok(repo.boxed())
+        self.repository_factory.create().await
     }
 
     async fn policy(&self) -> Result<Policy, InstantiateError> {
@@ -128,7 +124,7 @@ impl state::State for GraphQLState {
 
 #[must_use]
 pub fn schema(
-    pool: &PgPool,
+    repository_factory: BoxRepositoryFactory,
     policy_factory: &Arc<PolicyFactory>,
     homeserver_connection: impl HomeserverConnection + 'static,
     site_config: SiteConfig,
@@ -137,7 +133,7 @@ pub fn schema(
     limiter: Limiter,
 ) -> Schema {
     let state = GraphQLState {
-        pool: pool.clone(),
+        repository_factory,
         policy_factory: Arc::clone(policy_factory),
         homeserver_connection: Arc::new(homeserver_connection),
         site_config,

crates/handlers/src/test_utils.rs

@@ -34,8 +34,11 @@ use mas_keystore::{Encrypter, JsonWebKey, JsonWebKeySet, Keystore, PrivateKey};
 use mas_matrix::{HomeserverConnection, MockHomeserverConnection};
 use mas_policy::{InstantiateError, Policy, PolicyFactory};
 use mas_router::{SimpleRoute, UrlBuilder};
-use mas_storage::{BoxClock, BoxRepository, BoxRng, clock::MockClock};
-use mas_storage_pg::{DatabaseError, PgRepository};
+use mas_storage::{
+    BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng, RepositoryError, RepositoryFactory,
+    clock::MockClock,
+};
+use mas_storage_pg::PgRepositoryFactory;
 use mas_templates::{SiteConfigExt, Templates};
 use oauth2_types::{registration::ClientRegistrationResponse, requests::AccessTokenResponse};
 use rand::SeedableRng;
@@ -92,7 +95,7 @@ pub(crate) async fn policy_factory(
 
 #[derive(Clone)]
 pub(crate) struct TestState {
-    pub pool: PgPool,
+    pub repository_factory: PgRepositoryFactory,
     pub templates: Templates,
     pub key_store: Keystore,
     pub cookie_manager: CookieManager,
@@ -209,7 +212,7 @@ impl TestState {
         let limiter = Limiter::new(&RateLimitingConfig::default()).unwrap();
 
         let graphql_state = TestGraphQLState {
-            pool: pool.clone(),
+            repository_factory: PgRepositoryFactory::new(pool.clone()).boxed(),
             policy_factory: Arc::clone(&policy_factory),
             homeserver_connection: Arc::clone(&homeserver_connection),
             site_config: site_config.clone(),
@@ -224,14 +227,14 @@ impl TestState {
         let graphql_schema = graphql::schema_builder().data(state).finish();
 
         let activity_tracker = ActivityTracker::new(
-            pool.clone(),
+            PgRepositoryFactory::new(pool.clone()).boxed(),
             std::time::Duration::from_secs(60),
             &task_tracker,
             shutdown_token.child_token(),
         );
 
         Ok(Self {
-            pool,
+            repository_factory: PgRepositoryFactory::new(pool),
             templates,
             key_store,
             cookie_manager,
@@ -256,7 +259,7 @@ impl TestState {
     /// Reset the test utils to a fresh state, with the same configuration.
     pub async fn reset(self) -> Self {
         let site_config = self.site_config.clone();
-        let pool = self.pool.clone();
+        let pool = self.repository_factory.pool();
         let task_tracker = self.task_tracker.clone();
 
         // This should trigger the cancellation drop guard
@@ -351,9 +354,8 @@ impl TestState {
         access_token
     }
 
-    pub async fn repository(&self) -> Result<BoxRepository, DatabaseError> {
-        let repo = PgRepository::from_pool(&self.pool).await?;
-        Ok(repo.boxed())
+    pub async fn repository(&self) -> Result<BoxRepository, RepositoryError> {
+        self.repository_factory.create().await
     }
 
     /// Returns a new random number generator.
@@ -393,7 +395,7 @@ impl TestState {
 }
 
 struct TestGraphQLState {
-    pool: PgPool,
+    repository_factory: BoxRepositoryFactory,
     homeserver_connection: Arc<MockHomeserverConnection>,
     site_config: SiteConfig,
     policy_factory: Arc<PolicyFactory>,
@@ -407,11 +409,7 @@ struct TestGraphQLState {
 #[async_trait::async_trait]
 impl graphql::State for TestGraphQLState {
     async fn repository(&self) -> Result<BoxRepository, mas_storage::RepositoryError> {
-        let repo = PgRepository::from_pool(&self.pool)
-            .await
-            .map_err(mas_storage::RepositoryError::from_error)?;
-
-        Ok(repo.boxed())
+        self.repository_factory.create().await
     }
 
     async fn policy(&self) -> Result<Policy, InstantiateError> {
@@ -451,7 +449,7 @@ impl graphql::State for TestGraphQLState {
 
 impl FromRef<TestState> for PgPool {
     fn from_ref(input: &TestState) -> Self {
-        input.pool.clone()
+        input.repository_factory.pool()
     }
 }
 
@@ -598,14 +596,14 @@ impl FromRequestParts<TestState> for BoxRng {
 }
 
 impl FromRequestParts<TestState> for BoxRepository {
-    type Rejection = ErrorWrapper<mas_storage_pg::DatabaseError>;
+    type Rejection = ErrorWrapper<RepositoryError>;
 
     async fn from_request_parts(
         _parts: &mut axum::http::request::Parts,
         state: &TestState,
     ) -> Result<Self, Self::Rejection> {
-        let repo = PgRepository::from_pool(&state.pool).await?;
-        Ok(repo.boxed())
+        let repo = state.repository_factory.create().await?;
+        Ok(repo)
     }
 }