use chrono::{DateTime, Utc};
use serde::Serialize;
use sqlx::{Row, SqlitePool};
use uuid::Uuid;
use allowthem_core::error::AuthError;
use crate::cache::TenantMeta;
use crate::error::SaasError;
use crate::tenants::{TenantId, TenantStatus};
#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
pub struct TenantUsage {
pub period: String,
pub mau_count: i64,
pub limit_reached_at: Option<DateTime<Utc>>,
pub notified_at: Option<DateTime<Utc>>,
}
pub struct ControlDb {
pool: SqlitePool,
}
impl ControlDb {
pub async fn new(pool: SqlitePool) -> Result<Self, AuthError> {
sqlx::migrate!("./migrations")
.run(&pool)
.await
.map_err(sqlx::Error::from)?;
Ok(Self { pool })
}
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
pub async fn tenant_meta_by_slug(&self, slug: &str) -> Result<Option<TenantMeta>, SaasError> {
let row = sqlx::query("SELECT id, status, plan_id FROM tenants WHERE slug = ?1")
.bind(slug)
.fetch_optional(&self.pool)
.await?;
let Some(row) = row else { return Ok(None) };
let id_bytes: Vec<u8> = row.try_get("id")?;
let status: TenantStatus = row.try_get("status")?;
let plan_id: Vec<u8> = row.try_get("plan_id")?;
let id = Uuid::from_slice(&id_bytes).map_err(|_| SaasError::TenantNotFound)?;
Ok(Some(TenantMeta {
id: TenantId::from(id),
status,
plan_id,
}))
}
pub async fn most_recently_seen_tenants(&self, count: i64) -> Result<Vec<TenantId>, SaasError> {
let rows = sqlx::query(
"SELECT id FROM tenants \
WHERE status = 'active' AND last_seen_at IS NOT NULL \
ORDER BY last_seen_at DESC LIMIT ?1",
)
.bind(count)
.fetch_all(&self.pool)
.await?;
let mut result = Vec::with_capacity(rows.len());
for row in rows {
let bytes: Vec<u8> = row.try_get("id")?;
match Uuid::from_slice(&bytes) {
Ok(uuid) => result.push(TenantId::from(uuid)),
Err(_) => {
tracing::warn!("skipping tenant with undecodable UUID in most_recently_seen");
}
}
}
Ok(result)
}
pub async fn touch_last_seen(&self, tenant_id: &TenantId) -> Result<(), SaasError> {
sqlx::query("UPDATE tenants SET last_seen_at = datetime('now') WHERE id = ?1")
.bind(tenant_id.as_bytes())
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn usage_for_tenant(
&self,
tenant_id: &TenantId,
) -> Result<Vec<TenantUsage>, SaasError> {
let rows = sqlx::query_as::<_, TenantUsage>(
"SELECT period, mau_count, limit_reached_at, notified_at \
FROM tenant_usage \
WHERE tenant_id = ?1 \
ORDER BY period DESC",
)
.bind(tenant_id.as_bytes())
.fetch_all(&self.pool)
.await?;
Ok(rows)
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use sqlx::Row;
use std::str::FromStr;
pub async fn test_pool() -> SqlitePool {
let opts = sqlx::sqlite::SqliteConnectOptions::from_str("sqlite::memory:")
.unwrap()
.pragma("foreign_keys", "ON");
SqlitePool::connect_with(opts).await.unwrap()
}
#[tokio::test]
async fn control_db_runs_migrations() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await;
assert!(db.is_ok());
}
#[tokio::test]
async fn tenant_slug_unique() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await.unwrap();
let row = sqlx::query("SELECT id FROM tenant_plans LIMIT 1")
.fetch_one(db.pool())
.await
.unwrap();
let plan_id: Vec<u8> = row.get("id");
let id_a = uuid::Uuid::new_v4();
let id_b = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'Acme', 'acme', 'a@a.com', ?, 'active', 'acme.db')",
)
.bind(id_a.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await
.unwrap();
let res = sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'Acme 2', 'acme', 'b@b.com', ?, 'active', 'acme2.db')",
)
.bind(id_b.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await;
assert!(res.is_err(), "duplicate slug should be rejected");
}
#[tokio::test]
async fn tenant_status_check_rejects_invalid() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await.unwrap();
let row = sqlx::query("SELECT id FROM tenant_plans LIMIT 1")
.fetch_one(db.pool())
.await
.unwrap();
let plan_id: Vec<u8> = row.get("id");
let id = uuid::Uuid::new_v4();
let res = sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'Bad', 'bad-status', 'c@c.com', ?, 'banned', 'bad.db')",
)
.bind(id.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await;
assert!(res.is_err(), "invalid status should be rejected by CHECK");
}
#[tokio::test]
async fn member_role_check_rejects_invalid() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await.unwrap();
let row = sqlx::query("SELECT id FROM tenant_plans LIMIT 1")
.fetch_one(db.pool())
.await
.unwrap();
let plan_id: Vec<u8> = row.get("id");
let tenant_id = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'Role Test', 'role-test', 'd@d.com', ?, 'active', 'rtest.db')",
)
.bind(tenant_id.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await
.unwrap();
let member_id = uuid::Uuid::new_v4();
let res = sqlx::query(
"INSERT INTO tenant_members (id, tenant_id, email, role) \
VALUES (?, ?, 'e@e.com', 'superuser')",
)
.bind(member_id.as_bytes().as_ref())
.bind(tenant_id.as_bytes().as_ref())
.execute(db.pool())
.await;
assert!(res.is_err(), "invalid role should be rejected by CHECK");
}
#[tokio::test]
async fn usage_for_tenant_returns_records() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await.unwrap();
let row = sqlx::query("SELECT id FROM tenant_plans LIMIT 1")
.fetch_one(db.pool())
.await
.unwrap();
let plan_id: Vec<u8> = row.get("id");
let tid = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'U', 'usagetest', 'u@u.com', ?, 'active', 'u.db')",
)
.bind(tid.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await
.unwrap();
let uid = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenant_usage (id, tenant_id, period, mau_count) \
VALUES (?, ?, '2026-04', 42)",
)
.bind(uid.as_bytes().as_ref())
.bind(tid.as_bytes().as_ref())
.execute(db.pool())
.await
.unwrap();
let tenant_id = TenantId::from(tid);
let usage = db.usage_for_tenant(&tenant_id).await.unwrap();
assert_eq!(usage.len(), 1);
assert_eq!(usage[0].period, "2026-04");
assert_eq!(usage[0].mau_count, 42);
assert!(usage[0].limit_reached_at.is_none());
}
#[tokio::test]
async fn api_key_hash_unique() {
let pool = test_pool().await;
let db = ControlDb::new(pool).await.unwrap();
let row = sqlx::query("SELECT id FROM tenant_plans LIMIT 1")
.fetch_one(db.pool())
.await
.unwrap();
let plan_id: Vec<u8> = row.get("id");
let tenant_id = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenants (id, name, slug, owner_email, plan_id, status, db_path) \
VALUES (?, 'Key Test', 'key-test', 'f@f.com', ?, 'active', 'ktest.db')",
)
.bind(tenant_id.as_bytes().as_ref())
.bind(&plan_id)
.execute(db.pool())
.await
.unwrap();
let hash = vec![0u8; 32];
let key_id_a = uuid::Uuid::new_v4();
let key_id_b = uuid::Uuid::new_v4();
sqlx::query(
"INSERT INTO tenant_api_keys (id, tenant_id, name, key_hash, scope) \
VALUES (?, ?, 'key-a', ?, '[]')",
)
.bind(key_id_a.as_bytes().as_ref())
.bind(tenant_id.as_bytes().as_ref())
.bind(&hash)
.execute(db.pool())
.await
.unwrap();
let res = sqlx::query(
"INSERT INTO tenant_api_keys (id, tenant_id, name, key_hash, scope) \
VALUES (?, ?, 'key-b', ?, '[]')",
)
.bind(key_id_b.as_bytes().as_ref())
.bind(tenant_id.as_bytes().as_ref())
.bind(&hash)
.execute(db.pool())
.await;
assert!(res.is_err(), "duplicate key_hash should be rejected");
}
}