use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use super::{TenantId, TenantConfig};
use crate::error::ApiError;
#[derive(Debug, Clone)]
pub struct TenantInfo {
pub id: TenantId,
pub name: String,
pub features: Vec<String>,
pub metadata: HashMap<String, String>,
}
impl From<TenantConfig> for TenantInfo {
fn from(config: TenantConfig) -> Self {
Self {
id: config.id,
name: config.name,
features: config.features,
metadata: config.metadata,
}
}
}
#[derive(Debug, Clone)]
pub struct TenantContext {
info: TenantInfo,
}
impl TenantContext {
pub fn new(info: TenantInfo) -> Self {
Self { info }
}
pub fn tenant_id(&self) -> &TenantId {
&self.info.id
}
pub fn tenant_name(&self) -> &str {
&self.info.name
}
pub fn has_feature(&self, feature: &str) -> bool {
self.info.features.iter().any(|f| f == feature)
}
pub fn get_metadata(&self, key: &str) -> Option<&String> {
self.info.metadata.get(key)
}
}
#[async_trait]
pub trait TenantResolver: Send + Sync {
async fn resolve_from_subdomain(&self, subdomain: &str) -> Result<TenantId, ApiError>;
async fn resolve_from_header(&self, header_value: &str) -> Result<TenantId, ApiError>;
async fn get_tenant_config(&self, tenant_id: &TenantId) -> Result<TenantConfig, ApiError>;
}
pub struct InMemoryTenantResolver {
tenants: Arc<RwLock<HashMap<TenantId, TenantConfig>>>,
subdomain_map: Arc<RwLock<HashMap<String, TenantId>>>,
}
impl InMemoryTenantResolver {
pub fn new() -> Self {
Self {
tenants: Arc::new(RwLock::new(HashMap::new())),
subdomain_map: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn add_tenant(&self, config: TenantConfig) -> Result<(), ApiError> {
let mut tenants = self.tenants.write().await;
if let Some(ref subdomain) = config.subdomain {
let mut subdomain_map = self.subdomain_map.write().await;
subdomain_map.insert(subdomain.clone(), config.id.clone());
}
tenants.insert(config.id.clone(), config);
Ok(())
}
pub async fn remove_tenant(&self, tenant_id: &TenantId) -> Result<(), ApiError> {
let mut tenants = self.tenants.write().await;
if let Some(config) = tenants.remove(tenant_id) {
if let Some(subdomain) = config.subdomain {
let mut subdomain_map = self.subdomain_map.write().await;
subdomain_map.remove(&subdomain);
}
}
Ok(())
}
pub async fn list_tenants(&self) -> Vec<TenantConfig> {
let tenants = self.tenants.read().await;
tenants.values().cloned().collect()
}
}
impl Default for InMemoryTenantResolver {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl TenantResolver for InMemoryTenantResolver {
async fn resolve_from_subdomain(&self, subdomain: &str) -> Result<TenantId, ApiError> {
let subdomain_map = self.subdomain_map.read().await;
subdomain_map
.get(subdomain)
.cloned()
.ok_or_else(|| ApiError::NotFound(format!("Tenant not found for subdomain: {}", subdomain)))
}
async fn resolve_from_header(&self, header_value: &str) -> Result<TenantId, ApiError> {
Ok(TenantId::new(header_value))
}
async fn get_tenant_config(&self, tenant_id: &TenantId) -> Result<TenantConfig, ApiError> {
let tenants = self.tenants.read().await;
tenants
.get(tenant_id)
.cloned()
.ok_or_else(|| ApiError::NotFound(format!("Tenant not found: {}", tenant_id)))
}
}
#[cfg(feature = "database")]
pub struct PostgresTenantResolver {
pool: sqlx::PgPool,
}
#[cfg(feature = "database")]
impl PostgresTenantResolver {
pub fn new(pool: sqlx::PgPool) -> Self {
Self { pool }
}
pub async fn init(&self) -> Result<(), ApiError> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS tenants (
id VARCHAR(255) PRIMARY KEY,
name VARCHAR(255) NOT NULL,
subdomain VARCHAR(255) UNIQUE,
database_url TEXT,
features JSONB NOT NULL DEFAULT '[]',
metadata JSONB NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
is_active BOOLEAN NOT NULL DEFAULT TRUE
);
CREATE INDEX IF NOT EXISTS idx_tenants_subdomain ON tenants(subdomain);
CREATE INDEX IF NOT EXISTS idx_tenants_active ON tenants(is_active);
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
}
#[cfg(feature = "database")]
#[async_trait]
impl TenantResolver for PostgresTenantResolver {
async fn resolve_from_subdomain(&self, subdomain: &str) -> Result<TenantId, ApiError> {
let row = sqlx::query_as::<_, (String,)>(
"SELECT id FROM tenants WHERE subdomain = $1 AND is_active = TRUE"
)
.bind(subdomain)
.fetch_optional(&self.pool)
.await?
.ok_or_else(|| ApiError::NotFound(format!("Tenant not found for subdomain: {}", subdomain)))?;
Ok(TenantId::new(row.0))
}
async fn resolve_from_header(&self, header_value: &str) -> Result<TenantId, ApiError> {
Ok(TenantId::new(header_value))
}
async fn get_tenant_config(&self, tenant_id: &TenantId) -> Result<TenantConfig, ApiError> {
let row = sqlx::query_as::<_, (String, String, Option<String>, Option<String>, serde_json::Value, serde_json::Value, chrono::DateTime<chrono::Utc>, bool)>(
"SELECT id, name, subdomain, database_url, features, metadata, created_at, is_active FROM tenants WHERE id = $1"
)
.bind(tenant_id.as_str())
.fetch_optional(&self.pool)
.await?
.ok_or_else(|| ApiError::NotFound(format!("Tenant not found: {}", tenant_id)))?;
let features: Vec<String> = serde_json::from_value(row.4)
.map_err(|e| ApiError::InternalServerError(format!("Failed to parse features: {}", e)))?;
let metadata: HashMap<String, String> = serde_json::from_value(row.5)
.map_err(|e| ApiError::InternalServerError(format!("Failed to parse metadata: {}", e)))?;
Ok(TenantConfig {
id: TenantId::new(row.0),
name: row.1,
subdomain: row.2,
database_url: row.3,
features,
metadata,
created_at: row.6,
is_active: row.7,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_in_memory_resolver() {
let resolver = InMemoryTenantResolver::new();
let config = TenantConfig::new(
TenantId::new("tenant-1"),
"Test Tenant".to_string(),
)
.with_subdomain("test".to_string());
resolver.add_tenant(config).await.unwrap();
let tenant_id = resolver
.resolve_from_subdomain("test")
.await
.unwrap();
assert_eq!(tenant_id.as_str(), "tenant-1");
}
}