yorishiro-server 0.50.0

pub mod commands;

use anyhow::{Context, Result};
use clap::{Subcommand, ValueEnum};
use sqlx::PgPool;
use uuid::Uuid;
use yorishiro_core::models::maintenance::{self, MaintenanceMode};
use yorishiro_core::models::tenancy::{self, MembershipRole};
use yorishiro_core::services::auth::ApiKeyScope;

use commands::{create_api_key, list_api_keys, resync_embeddings, revoke_api_key};

/// Subcommands under `yorishiro-server admin`.
/// API keys are stored only as SHA-256 hashes and user passwords only as argon2 hashes, so neither can be provisioned by hand in SQL: this CLI is the only bootstrap mechanism.
#[derive(Subcommand)]
pub enum AdminCommand {
    /// Create a new tenant.
    /// With --template, also creates a schema from the template and a default workspace linked to it (required for login to work).
    CreateTenant {
        name: String,
        /// Cap on the number of workspaces this tenant may create.
        /// Omit for unlimited.
        #[arg(long)]
        max_workspaces: Option<i32>,
        /// Built-in template ID to bootstrap a schema and default workspace.
        /// Without this, only a tenant is created (no workspace, no login possible).
        #[arg(long)]
        template: Option<String>,
    },
    /// List all tenants.
    ListTenants,
    /// Create an additional workspace under a tenant (see `admin list-tenants` for the tenant ID).
    CreateWorkspace {
        tenant_id: Uuid,
        name: String,
        /// Cap on the number of entities this workspace may hold.
        /// Omit for unlimited.
        #[arg(long)]
        max_entities: Option<i32>,
        /// Schema to associate with this workspace.
        /// Omit to leave it unset.
        #[arg(long)]
        schema_id: Option<Uuid>,
    },
    /// List workspaces under a tenant.
    ListWorkspaces { tenant_id: Uuid },
    /// Create a human user account.
    CreateUser {
        email: String,
        password: String,
        #[arg(long)]
        display_name: Option<String>,
    },
    /// Add (or change the role of) a user's membership in a tenant.
    AddMember {
        tenant_id: Uuid,
        user_id: Uuid,
        role: RoleArg,
    },
    /// List a tenant's members.
    ListMembers { tenant_id: Uuid },
    /// Create an invite token for an email to join a tenant with a given role.
    /// Signup is invite-only; there is no self-service, unauthenticated account creation.
    CreateInvite {
        tenant_id: Uuid,
        email: String,
        role: RoleArg,
        /// How long the invite stays redeemable.
        /// Defaults to 7 days.
        #[arg(long, default_value_t = 168)]
        ttl_hours: i64,
    },
    /// Issue a new API key for a workspace (see `admin list-workspaces <tenant-id>` for the workspace ID).
    CreateApiKey {
        workspace_id: Uuid,
        scope: ScopeArg,
        /// Attribute the key to a specific user (see `admin list-members <tenant-id>`).
        /// The requested scope is capped by that user's tenant role (owner/admin: schema, member: write, viewer: read); omit for an unattributed service key.
        #[arg(long)]
        user: Option<Uuid>,
    },
    /// List API keys for a workspace.
    ListApiKeys { workspace_id: Uuid },
    /// Revoke (delete) an API key (see `admin list-api-keys <workspace-id>` for the key ID).
    RevokeApiKey { key_id: Uuid },
    /// Re-sync embeddings for entities whose embedding is still missing.
    ResyncEmbeddings { workspace_id: Uuid },
    /// Put the deployment into maintenance, or take it out.
    ///
    /// `read-only` refuses writes with 423; `full-lock` refuses everything with 503; `off` serves normally.
    /// The state is shared by every node, and `/up`/`/health` keep answering so an orchestrator does not restart a server that is deliberately paused.
    Maintenance {
        #[arg(value_enum)]
        mode: MaintenanceArg,
        /// Seconds for the Retry-After header.
        #[arg(long, default_value_t = 300)]
        retry_after: u32,
        /// Shown to callers instead of the generic message.
        #[arg(long)]
        reason: Option<String>,
    },
    /// Show the current maintenance state.
    MaintenanceStatus,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum MaintenanceArg {
    Off,
    ReadOnly,
    FullLock,
}

impl From<MaintenanceArg> for MaintenanceMode {
    fn from(arg: MaintenanceArg) -> Self {
        match arg {
            MaintenanceArg::Off => Self::Off,
            MaintenanceArg::ReadOnly => Self::ReadOnly,
            MaintenanceArg::FullLock => Self::FullLock,
        }
    }
}

#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum ScopeArg {
    Read,
    Write,
    Schema,
    Migration,
}

impl From<ScopeArg> for ApiKeyScope {
    fn from(value: ScopeArg) -> Self {
        match value {
            ScopeArg::Read => ApiKeyScope::Read,
            ScopeArg::Write => ApiKeyScope::Write,
            ScopeArg::Schema => ApiKeyScope::Schema,
            ScopeArg::Migration => ApiKeyScope::Migration,
        }
    }
}

#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum RoleArg {
    Owner,
    Admin,
    Member,
    Viewer,
}

impl From<RoleArg> for MembershipRole {
    fn from(value: RoleArg) -> Self {
        match value {
            RoleArg::Owner => MembershipRole::Owner,
            RoleArg::Admin => MembershipRole::Admin,
            RoleArg::Member => MembershipRole::Member,
            RoleArg::Viewer => MembershipRole::Viewer,
        }
    }
}

/// Entry point for the admin subcommands.
/// Unlike a plain server start (no args), this operates on the database directly using the DATABASE_URL connection role (the admin role that can run migrations and is the only role with write access to `identity.tenants`/ `identity.users`/`identity.tenant_memberships`).
pub async fn run(command: AdminCommand) -> Result<()> {
    let database_url =
        std::env::var("DATABASE_URL").context("DATABASE_URL must be set for admin commands")?;
    let pool = PgPool::connect(&database_url)
        .await
        .context("failed to connect to database")?;
    sqlx::migrate!("./migrations")
        .set_ignore_missing(true)
        .run(&pool)
        .await?;

    run_with_pool(&pool, command).await
}

/// Execute an admin command against a pool whose migrations have already been applied.
pub async fn run_with_pool(pool: &PgPool, command: AdminCommand) -> Result<()> {
    // What a workspace created here gets stamped with.
    // Built from the environment rather than from a running provider: the admin commands do not start one, and loading a local ONNX model to read one number off it would make `create-tenant` wait on a model it never uses.
    let embedding_stamp: Option<(String, i32)> = std::env::var("YORISHIRO_EMBEDDING_DIMENSIONS")
        .ok()
        .and_then(|d| d.parse::<i32>().ok())
        .or(Some(1024))
        .map(|dimensions| (crate::embedding_model_name(), dimensions));

    match command {
        AdminCommand::CreateTenant {
            name,
            max_workspaces,
            template,
        } => {
            let tenant = tenancy::create_tenant(pool, &name, max_workspaces).await?;
            println!("tenant created");
            println!("  id:            {}", tenant.id);
            println!("  name:          {}", tenant.name);
            println!("  max_workspaces: {}", format_limit(tenant.max_workspaces));

            if let Some(template_id) = template {
                let definition = yorishiro_core::templates::get_template(&template_id)?;

                // Workspace first, then its schema.
                // A schema belongs to a workspace, so the workspace has to exist to own it; the workspace's own `schema_id` is linked afterwards, once there is a schema to point at.
                let workspace = tenancy::create_workspace(
                    pool,
                    tenant.id,
                    "default",
                    None,
                    None,
                    embedding_stamp.as_ref().map(|(m, d)| (m.as_str(), *d)),
                )
                .await?;

                let mut conn = pool.acquire().await.context("acquire connection")?;
                let (schema, _diff) = yorishiro_core::models::schemas::create_schema(
                    &mut conn,
                    tenant.id,
                    workspace.id,
                    definition,
                )
                .await?;
                drop(conn);
                tenancy::set_workspace_schema(pool, workspace.id, schema.id).await?;

                println!("schema created (from template '{template_id}')");
                println!("  id:      {}", schema.id);
                println!("  name:    {}", schema.name);
                println!("  version: {}", schema.version);

                println!("default workspace created");
                println!("  id:        {}", workspace.id);
                println!("  name:      {}", workspace.name);
                println!("  schema_id: {}", schema.id);
            } else {
                println!();
                println!("next steps:");
                println!("  1. create a schema (via REST API or --template)");
                println!(
                    "  2. admin create-workspace {} <name> --schema-id <id>",
                    tenant.id
                );
            }
        }
        AdminCommand::ListTenants => {
            let tenants = tenancy::list_tenants(pool).await?;
            if tenants.is_empty() {
                println!("no tenants (create one with `admin create-tenant <name>`)");
            }
            for tenant in tenants {
                println!(
                    "{}  {:<24} max_workspaces={}",
                    tenant.id,
                    tenant.name,
                    format_limit(tenant.max_workspaces)
                );
            }
        }
        AdminCommand::CreateWorkspace {
            tenant_id,
            name,
            max_entities,
            schema_id,
        } => {
            let workspace = tenancy::create_workspace(
                pool,
                tenant_id,
                &name,
                max_entities,
                schema_id,
                embedding_stamp.as_ref().map(|(m, d)| (m.as_str(), *d)),
            )
            .await
            .map_err(anyhow::Error::from)?;
            println!("workspace created");
            println!("  id:           {}", workspace.id);
            println!("  tenant id:    {}", workspace.tenant_id);
            println!("  name:         {}", workspace.name);
            println!("  max_entities: {}", format_limit(workspace.max_entities));
            if let Some(schema_id) = workspace.schema_id {
                println!("  schema id:    {schema_id}");
            }
        }
        AdminCommand::ListWorkspaces { tenant_id } => {
            let workspaces = tenancy::list_workspaces(pool, tenant_id)
                .await
                .map_err(anyhow::Error::from)?;
            if workspaces.is_empty() {
                println!("no workspaces for tenant {tenant_id}");
            }
            for workspace in workspaces {
                println!(
                    "{}  {:<24} max_entities={}",
                    workspace.id,
                    workspace.name,
                    format_limit(workspace.max_entities)
                );
            }
        }
        AdminCommand::CreateUser {
            email,
            password,
            display_name,
        } => {
            let mut conn = pool
                .acquire()
                .await
                .context("failed to acquire a connection")?;
            let user = tenancy::create_user(&mut *conn, &email, &password, display_name.as_deref())
                .await
                .map_err(anyhow::Error::from)?;
            println!("user created");
            println!("  id:    {}", user.id);
            println!("  email: {}", user.email);
        }
        AdminCommand::AddMember {
            tenant_id,
            user_id,
            role,
        } => {
            let mut conn = pool
                .acquire()
                .await
                .context("failed to acquire a connection")?;
            tenancy::add_member(&mut *conn, tenant_id, user_id, role.into())
                .await
                .map_err(anyhow::Error::from)?;
            println!("membership added: user {user_id} is now {role:?} of tenant {tenant_id}");
        }
        AdminCommand::ListMembers { tenant_id } => {
            let members = tenancy::list_members(pool, tenant_id)
                .await
                .map_err(anyhow::Error::from)?;
            if members.is_empty() {
                println!("no members for tenant {tenant_id}");
            }
            for member in members {
                println!("{}  {:<8?} {}", member.user_id, member.role, member.email);
            }
        }
        AdminCommand::CreateInvite {
            tenant_id,
            email,
            role,
            ttl_hours,
        } => {
            let (invite, token) = tenancy::create_invite(
                pool,
                tenant_id,
                &email,
                role.into(),
                chrono::Duration::hours(ttl_hours),
            )
            .await
            .map_err(anyhow::Error::from)?;
            println!("invite created (the plaintext token is shown ONLY once, send it now)");
            println!("  token:      {token}");
            println!("  invite id:  {}", invite.id);
            println!("  tenant id:  {}", invite.tenant_id);
            println!("  email:      {}", invite.email);
            println!("  role:       {:?}", invite.role);
            println!(
                "  expires at: {}",
                invite.expires_at.format("%Y-%m-%d %H:%M UTC")
            );
        }
        AdminCommand::CreateApiKey {
            workspace_id,
            scope,
            user,
        } => {
            let scope = ApiKeyScope::from(scope);
            let created = create_api_key(pool, workspace_id, scope, user).await?;
            println!("api key created (the plaintext key is shown ONLY once, store it now)");
            println!("  key:          {}", created.plaintext);
            println!("  key id:       {}", created.id);
            println!("  workspace id: {}", created.workspace_id);
            println!("  scope:        {scope:?}");
            if let Some(user_id) = created.user_id {
                println!("  user id:      {user_id}");
            }
        }
        AdminCommand::ListApiKeys { workspace_id } => {
            let keys = list_api_keys(pool, workspace_id).await?;
            if keys.is_empty() {
                println!("no api keys for workspace {workspace_id}");
            }
            for key in keys {
                println!(
                    "{}  {:<8} prefix={}  user={}  created={}  last_used={}",
                    key.id,
                    key.scope,
                    key.key_prefix,
                    key.user_id
                        .map(|id| id.to_string())
                        .unwrap_or_else(|| "-".into()),
                    key.created_at.format("%Y-%m-%d %H:%M"),
                    key.last_used_at
                        .map(|t| t.format("%Y-%m-%d %H:%M").to_string())
                        .unwrap_or_else(|| "never".into()),
                );
            }
        }
        AdminCommand::RevokeApiKey { key_id } => {
            revoke_api_key(pool, key_id).await?;
            println!("api key {key_id} revoked (takes effect on the next request)");
        }
        AdminCommand::ResyncEmbeddings { workspace_id } => {
            let provider = crate::build_embedding_provider()
                .context("embedding provider must be configured (see .env.example)")?;
            let report = resync_embeddings(pool, workspace_id, provider.as_ref()).await?;
            println!(
                "resync finished: {} entities had no embedding, {} synced, {} failed \
                 (entities whose entity_type has no x-embed field stay without embedding)",
                report.candidates, report.synced, report.failed,
            );
        }
        AdminCommand::Maintenance {
            mode,
            retry_after,
            reason,
        } => {
            let mode: MaintenanceMode = mode.into();
            let state = maintenance::set(pool, mode, retry_after, reason).await?;
            match state.mode {
                MaintenanceMode::Off => println!("maintenance off; serving normally"),
                MaintenanceMode::ReadOnly => println!(
                    "maintenance read-only: writes refused with 423, Retry-After {}s",
                    state.retry_after
                ),
                MaintenanceMode::FullLock => println!(
                    "maintenance full lock: all requests refused with 503, Retry-After {}s \
                     (/up and /health keep answering)",
                    state.retry_after
                ),
            }
            if let Some(reason) = state.reason {
                println!("reason shown to callers: {reason}");
            }
        }
        AdminCommand::MaintenanceStatus => {
            let mut conn = pool.acquire().await?;
            let state = maintenance::get(&mut *conn).await?;
            println!(
                "mode={} retry_after={}s reason={}",
                state.mode.as_db_str(),
                state.retry_after,
                state.reason.as_deref().unwrap_or("(none)")
            );
        }
    }

    Ok(())
}

fn format_limit(limit: Option<i32>) -> String {
    match limit {
        Some(n) => n.to_string(),
        None => "unlimited".to_string(),
    }
}

#[cfg(test)]
#[path = "../../tests/admin/mod.rs"]
mod tests;