compliance-scanner-agent/compliance-mcp/src/database.rs

use mongodb::{Client, Collection};

use compliance_core::models::*;

#[derive(Clone, Debug)]
pub struct Database {
    inner: mongodb::Database,
}

impl Database {
    pub async fn connect(uri: &str, db_name: &str) -> Result<Self, mongodb::error::Error> {
        let client = Client::with_uri_str(uri).await?;
        let db = client.database(db_name);
        db.run_command(mongodb::bson::doc! { "ping": 1 }).await?;
        tracing::info!("MCP server connected to MongoDB '{db_name}'");
        Ok(Self { inner: db })
    }

    pub fn findings(&self) -> Collection<Finding> {
        self.inner.collection("findings")
    }

    pub fn sbom_entries(&self) -> Collection<SbomEntry> {
        self.inner.collection("sbom_entries")
    }

    pub fn dast_findings(&self) -> Collection<DastFinding> {
        self.inner.collection("dast_findings")
    }

    pub fn dast_scan_runs(&self) -> Collection<DastScanRun> {
        self.inner.collection("dast_scan_runs")
    }

    pub fn pentest_sessions(&self) -> Collection<PentestSession> {
        self.inner.collection("pentest_sessions")
    }

    pub fn attack_chain_nodes(&self) -> Collection<AttackChainNode> {
        self.inner.collection("attack_chain_nodes")
    }

    pub fn pentest_messages(&self) -> Collection<PentestMessage> {
        self.inner.collection("pentest_messages")
    }
}