multi tenancy part 1

backend/src/tenants.rs

@@ -0,0 +1,123 @@
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+
+use axum::{async_trait, extract::FromRequestParts, http::request::Parts};
+use diesel::{pg::PgConnection, prelude::*, sql_types::Text};
+use uuid::Uuid;
+
+use crate::{
+    db::PgPool,
+    error::{AppError, AppResult},
+    models::Tenant,
+    schema::tenants::dsl,
+    state::AppState,
+};
+
+pub struct TenantRepository;
+
+impl TenantRepository {
+    pub fn get_by_id(conn: &mut PgConnection, tenant_id: Uuid) -> AppResult<Tenant> {
+        dsl::tenants.find(tenant_id).first(conn).map_err(Into::into)
+    }
+
+    pub fn get_by_slug(conn: &mut PgConnection, slug: &str) -> AppResult<Tenant> {
+        dsl::tenants
+            .filter(dsl::slug.eq(slug))
+            .first(conn)
+            .map_err(Into::into)
+    }
+}
+
+#[derive(Clone)]
+pub struct TenantService {
+    pool: PgPool,
+    cache_by_id: Arc<RwLock<HashMap<Uuid, Tenant>>>,
+    cache_by_slug: Arc<RwLock<HashMap<String, Tenant>>>,
+}
+
+impl TenantService {
+    pub fn new(pool: PgPool) -> Self {
+        Self {
+            pool,
+            cache_by_id: Arc::new(RwLock::new(HashMap::new())),
+            cache_by_slug: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub fn get_by_id(&self, tenant_id: Uuid) -> AppResult<Tenant> {
+        if let Some(tenant) = self.cache_by_id.read().unwrap().get(&tenant_id) {
+            return Ok(tenant.clone());
+        }
+
+        let tenant = self.load(|conn| TenantRepository::get_by_id(conn, tenant_id))?;
+        self.store(&tenant);
+        Ok(tenant)
+    }
+
+    pub fn get_by_slug(&self, slug: &str) -> AppResult<Tenant> {
+        if let Some(tenant) = self.cache_by_slug.read().unwrap().get(slug) {
+            return Ok(tenant.clone());
+        }
+
+        let slug_owned = slug.to_owned();
+        let tenant = self.load(|conn| TenantRepository::get_by_slug(conn, &slug_owned))?;
+        self.store(&tenant);
+        Ok(tenant)
+    }
+
+    pub fn tenant_id_for_slug(&self, slug: &str) -> AppResult<Uuid> {
+        Ok(self.get_by_slug(slug)?.tenant_id)
+    }
+
+    fn load<F>(&self, loader: F) -> AppResult<Tenant>
+    where
+        F: FnOnce(&mut PgConnection) -> AppResult<Tenant>,
+    {
+        let mut conn = self
+            .pool
+            .get()
+            .map_err(|err| AppError::internal(format!("database pool error: {err}")))?;
+        let tenant = loader(&mut conn)?;
+        Ok(tenant)
+    }
+
+    fn store(&self, tenant: &Tenant) {
+        {
+            let mut by_id = self.cache_by_id.write().unwrap();
+            by_id.insert(tenant.tenant_id, tenant.clone());
+        }
+        {
+            let mut by_slug = self.cache_by_slug.write().unwrap();
+            by_slug.insert(tenant.slug.clone(), tenant.clone());
+        }
+    }
+}
+
+pub fn apply_tenant_guc(conn: &mut PgConnection, tenant_id: Uuid) -> AppResult<()> {
+    diesel::sql_query("SELECT set_config('papercrate.tenant_id', $1, true)")
+        .bind::<Text, _>(tenant_id.to_string())
+        .execute(conn)
+        .map(|_| ())
+        .map_err(AppError::from)
+}
+
+pub struct TenantContext {
+    pub tenant: Tenant,
+}
+
+#[async_trait]
+impl FromRequestParts<AppState> for TenantContext {
+    type Rejection = AppError;
+
+    async fn from_request_parts(
+        _parts: &mut Parts,
+        state: &AppState,
+    ) -> Result<Self, Self::Rejection> {
+        let tenant = state
+            .tenants
+            .get_by_slug(&state.config.default_tenant_slug)?;
+        Ok(Self { tenant })
+    }
+}