axonml-server 0.6.2

//! Database Schema Initialization — Collection Bootstrap and Seed Users
//!
//! Defines the canonical collection name constants for all Aegis-DB document
//! collections used by the AxonML server and provides `Schema::init()` to
//! create them on startup. Collections are created idempotently (409 Conflict
//! is silently ignored by `Database::create_collection`).
//!
//! Collections managed:
//! - `axonml_users` — user accounts and authentication data
//! - `axonml_runs` — training run records
//! - `axonml_models` — model registry entries
//! - `axonml_model_versions` — versioned model artifacts
//! - `axonml_endpoints` — inference serving endpoints
//! - `axonml_datasets` — dataset metadata
//! - `axonml_notebooks` — training notebook documents
//! - `axonml_checkpoints` — training checkpoints
//!
//! Also provides `create_default_admin()` and `create_devops_admin()` for
//! bootstrapping seed admin users with pre-hashed passwords.
//!
//! # File
//! `crates/axonml-server/src/db/schema.rs`
//!
//! # Author
//! Andrew Jewell Sr. — AutomataNexus LLC
//! ORCID: 0009-0005-2158-7060
//!
//! # Updated
//! April 16, 2026 11:15 PM EST
//!
//! # Disclaimer
//! Use at own risk. This software is provided "as is", without warranty of any
//! kind, express or implied. The author and AutomataNexus shall not be held
//! liable for any damages arising from the use of this software.

// =============================================================================
// Imports
// =============================================================================

use super::{Database, DbError};
use tracing::info;

// =============================================================================
// Collection Name Constants
// =============================================================================

/// Collection names used by AxonML
pub const USERS_COLLECTION: &str = "axonml_users";
pub const RUNS_COLLECTION: &str = "axonml_runs";
pub const MODELS_COLLECTION: &str = "axonml_models";
pub const VERSIONS_COLLECTION: &str = "axonml_model_versions";
pub const ENDPOINTS_COLLECTION: &str = "axonml_endpoints";
pub const DATASETS_COLLECTION: &str = "axonml_datasets";
pub const NOTEBOOKS_COLLECTION: &str = "axonml_notebooks";
pub const CHECKPOINTS_COLLECTION: &str = "axonml_checkpoints";

// =============================================================================
// Schema Initialization
// =============================================================================

/// Schema definitions for all AxonML collections
pub struct Schema;

impl Schema {
    /// Initialize all database collections
    pub async fn init(db: &Database) -> Result<(), DbError> {
        info!("Initializing AxonML database schema...");

        // Create document store collections
        Self::create_users_collection(db).await?;
        Self::create_runs_collection(db).await?;
        Self::create_models_collection(db).await?;
        Self::create_model_versions_collection(db).await?;
        Self::create_endpoints_collection(db).await?;
        Self::create_datasets_collection(db).await?;
        Self::create_notebooks_collection(db).await?;
        Self::create_checkpoints_collection(db).await?;

        info!("Database schema initialized successfully");
        Ok(())
    }

    // -------------------------------------------------------------------------
    // Individual Collection Creation
    // -------------------------------------------------------------------------

    /// Create users collection
    async fn create_users_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(USERS_COLLECTION).await?;
        info!("Created {} collection", USERS_COLLECTION);
        Ok(())
    }

    /// Create training runs collection
    async fn create_runs_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(RUNS_COLLECTION).await?;
        info!("Created {} collection", RUNS_COLLECTION);
        Ok(())
    }

    /// Create models collection
    async fn create_models_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(MODELS_COLLECTION).await?;
        info!("Created {} collection", MODELS_COLLECTION);
        Ok(())
    }

    /// Create model versions collection
    async fn create_model_versions_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(VERSIONS_COLLECTION).await?;
        info!("Created {} collection", VERSIONS_COLLECTION);
        Ok(())
    }

    /// Create inference endpoints collection
    async fn create_endpoints_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(ENDPOINTS_COLLECTION).await?;
        info!("Created {} collection", ENDPOINTS_COLLECTION);
        Ok(())
    }

    /// Create datasets collection
    async fn create_datasets_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(DATASETS_COLLECTION).await?;
        info!("Created {} collection", DATASETS_COLLECTION);
        Ok(())
    }

    /// Create training notebooks collection
    async fn create_notebooks_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(NOTEBOOKS_COLLECTION).await?;
        info!("Created {} collection", NOTEBOOKS_COLLECTION);
        Ok(())
    }

    /// Create checkpoints collection
    async fn create_checkpoints_collection(db: &Database) -> Result<(), DbError> {
        db.create_collection(CHECKPOINTS_COLLECTION).await?;
        info!("Created {} collection", CHECKPOINTS_COLLECTION);
        Ok(())
    }

    // =========================================================================
    // Seed Users
    // =========================================================================

    /// Create default admin user if not exists
    pub async fn create_default_admin(db: &Database, password_hash: &str) -> Result<(), DbError> {
        // Check if admin exists
        let admin = db.doc_get(USERS_COLLECTION, "admin").await?;

        if admin.is_none() {
            let admin_data = serde_json::json!({
                "id": "admin",
                "email": "admin@axonml.local",
                "name": "Administrator",
                "password_hash": password_hash,
                "role": "admin",
                "mfa_enabled": false,
                "totp_secret": null,
                "webauthn_credentials": [],
                "recovery_codes": [],
                "email_pending": false,
                "email_verified": true,
                "verification_token": null,
                "created_at": chrono::Utc::now().to_rfc3339(),
                "updated_at": chrono::Utc::now().to_rfc3339()
            });

            db.doc_insert(USERS_COLLECTION, Some("admin"), admin_data)
                .await?;

            info!("Created default admin user");
        }

        Ok(())
    }

    /// Create DevOps admin user if not exists
    pub async fn create_devops_admin(db: &Database, password_hash: &str) -> Result<(), DbError> {
        // Check if DevOps user exists
        let devops = db.doc_get(USERS_COLLECTION, "devops").await?;

        if devops.is_none() {
            let devops_data = serde_json::json!({
                "id": "devops",
                "email": "DevOps@AutomataNexus.com",
                "name": "Andrew Jewell",
                "password_hash": password_hash,
                "role": "admin",
                "mfa_enabled": false,
                "totp_secret": null,
                "webauthn_credentials": [],
                "recovery_codes": [],
                "email_pending": false,
                "email_verified": true,
                "verification_token": null,
                "created_at": chrono::Utc::now().to_rfc3339(),
                "updated_at": chrono::Utc::now().to_rfc3339()
            });

            db.doc_insert(USERS_COLLECTION, Some("devops"), devops_data)
                .await?;

            info!("Created DevOps admin user");
        }

        Ok(())
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_collection_names() {
        assert_eq!(USERS_COLLECTION, "axonml_users");
        assert_eq!(RUNS_COLLECTION, "axonml_runs");
        assert_eq!(MODELS_COLLECTION, "axonml_models");
        assert_eq!(VERSIONS_COLLECTION, "axonml_model_versions");
        assert_eq!(ENDPOINTS_COLLECTION, "axonml_endpoints");
    }
}