AkurAI-Drive / backend/src/telemetry.rs

Latest change 19126b373cd637ce1c055d17f2198a09424fe140 - fix(drive): cargo fmt + collapsible-if clippy fix (AKURAI-DRIVE-6) by AkurAI Agent

//! Privacy-safe Drive activation telemetry.
//!
//! Event payloads carry only canonical IDs and bounded aggregate dimensions.
//! Filenames, object keys, request URLs, and user-supplied organization labels
//! never cross this boundary. Legacy storage rows without a canonical tenant ID
//! remain intentionally untelemetrized rather than being guessed from a label.

use std::sync::{LazyLock, Mutex};

use akurai_telemetry::{
    PrivacyClassification, ReceiverError, TelemetryContext, TelemetryDelivery,
    TelemetryDeliveryConfig, TelemetryEnvelope, TelemetryReceiver,
};

use crate::db::with_db;

const DAY_SECS: i64 = 24 * 60 * 60;
const EVENT_UPLOAD: &str = "drive.upload_observed.v1";
const EVENT_RETRIEVAL: &str = "drive.retrieval_requested.v1";
const EVENT_FACT: &str = "drive.organization_fact.v1";

#[derive(Clone)]
pub struct StorageTelemetryContext {
    pub organization_id: String,
    pub workspace_id: String,
    pub plan: String,
}

impl StorageTelemetryContext {
    fn contract(&self) -> Option<TelemetryContext> {
        (!self.organization_id.is_empty() && !self.workspace_id.is_empty()).then(|| {
            TelemetryContext::new(
                &self.organization_id,
                &self.workspace_id,
                &crate::config::get().drive_installation_id,
                "akurai-drive",
            )
        })
    }
}

struct SqliteReceiver;

impl TelemetryReceiver for SqliteReceiver {
    fn receive(&self, event: TelemetryEnvelope) -> Result<(), ReceiverError> {
        with_db(|conn| {
            conn.execute(
                "INSERT OR IGNORE INTO telemetry_events
                 (event_id, schema_version, organization_id, workspace_id, installation_id, app_id,
                  event_name, occurred_at, attributes)
                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
                rusqlite::params![
                    event.event_id,
                    event.schema_version,
                    event.organization_id,
                    event.workspace_id,
                    event.installation_id,
                    event.app_id,
                    event.event_name,
                    event.occurred_at,
                    serde_json::to_string(&event.attributes).unwrap_or_else(|_| "[]".to_string()),
                ],
            )
            .map_err(|error| ReceiverError::new(error.to_string()))?;

            let day = event.occurred_at.div_euclid(DAY_SECS) * DAY_SECS;
            let mut fact = DailyFact::from_attributes(&event.attributes);
            fact.active_day =
                event.event_name == EVENT_UPLOAD || event.event_name == EVENT_RETRIEVAL;
            if event.event_name == EVENT_UPLOAD {
                fact.first_upload_at = Some(event.occurred_at);
            }
            if event.event_name == EVENT_RETRIEVAL {
                fact.first_retrieval_at = Some(event.occurred_at);
            }
            if event.event_name == EVENT_FACT || fact.active_day {
                upsert_daily_fact(conn, &event.organization_id, &event.workspace_id, day, fact)?;
            }
            Ok(())
        })
    }
}

struct DailyFact {
    active_day: bool,
    first_upload_at: Option<i64>,
    first_retrieval_at: Option<i64>,
    storage_used_bytes: Option<i64>,
    quota_bytes: Option<i64>,
    quota_pressure: Option<String>,
}

impl DailyFact {
    fn from_attributes(attributes: &[(String, String)]) -> Self {
        let value = |name: &str| {
            attributes
                .iter()
                .find_map(|(key, value)| (key == name).then_some(value.as_str()))
        };
        Self {
            active_day: value("active_day") == Some("1"),
            first_upload_at: None,
            first_retrieval_at: None,
            storage_used_bytes: value("storage_used_bytes").and_then(|value| value.parse().ok()),
            quota_bytes: value("quota_bytes").and_then(|value| value.parse().ok()),
            quota_pressure: value("quota_pressure").map(str::to_string),
        }
    }
}

fn upsert_daily_fact(
    conn: &rusqlite::Connection,
    organization_id: &str,
    workspace_id: &str,
    day: i64,
    fact: DailyFact,
) -> Result<(), ReceiverError> {
    conn.execute(
        "INSERT INTO organization_daily_facts
         (organization_id, workspace_id, day, active_day, first_upload_at, first_retrieval_at,
          storage_used_bytes, quota_bytes, quota_pressure)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
         ON CONFLICT(organization_id, workspace_id, day) DO UPDATE SET
           active_day = MAX(active_day, excluded.active_day),
           first_upload_at = COALESCE(organization_daily_facts.first_upload_at, excluded.first_upload_at),
           first_retrieval_at = COALESCE(organization_daily_facts.first_retrieval_at, excluded.first_retrieval_at),
           storage_used_bytes = COALESCE(excluded.storage_used_bytes, organization_daily_facts.storage_used_bytes),
           quota_bytes = COALESCE(excluded.quota_bytes, organization_daily_facts.quota_bytes),
           quota_pressure = COALESCE(excluded.quota_pressure, organization_daily_facts.quota_pressure)",
        rusqlite::params![
            organization_id,
            workspace_id,
            day,
            i64::from(fact.active_day),
            fact.first_upload_at,
            fact.first_retrieval_at,
            fact.storage_used_bytes,
            fact.quota_bytes,
            fact.quota_pressure,
        ],
    )
    .map_err(|error| ReceiverError::new(error.to_string()))?;
    Ok(())
}

static DELIVERY: LazyLock<Mutex<TelemetryDelivery>> = LazyLock::new(|| {
    Mutex::new(
        TelemetryDelivery::new_with_config(
            SqliteReceiver,
            TelemetryDeliveryConfig {
                queue_capacity: 256,
                dedupe_capacity: 512,
                ..TelemetryDeliveryConfig::default()
            },
        )
        .expect("valid Drive telemetry delivery configuration"),
    )
});

fn emit(
    context: &StorageTelemetryContext,
    name: &str,
    now: i64,
    attributes: Vec<(String, String)>,
) {
    let Some(context) = context.contract() else {
        return;
    };
    let day = now.div_euclid(DAY_SECS);
    let event_id = format!(
        "drive.{name}.{}.{}.{}",
        context.organization_id, context.workspace_id, day
    );
    if let Ok(event) = TelemetryEnvelope::new(
        event_id,
        &context,
        name,
        now,
        PrivacyClassification::Operational,
        attributes,
    ) && let Ok(delivery) = DELIVERY.lock()
    {
        let _ = delivery.sender().emit(event);
    }
}

/// Record a successful object listing with at least one object as observed Drive value.
pub fn observe_upload(context: &StorageTelemetryContext, storage_used_bytes: u64) {
    let now = chrono::Utc::now().timestamp();
    emit(
        context,
        EVENT_UPLOAD,
        now,
        vec![("active_day".to_string(), "1".to_string())],
    );
    let (quota_bytes, quota_pressure) = quota(context.plan.as_str(), storage_used_bytes);
    let mut attributes = vec![
        ("active_day".to_string(), "1".to_string()),
        (
            "storage_used_bytes".to_string(),
            storage_used_bytes.to_string(),
        ),
        ("quota_pressure".to_string(), quota_pressure.to_string()),
    ];
    if let Some(quota_bytes) = quota_bytes {
        attributes.push(("quota_bytes".to_string(), quota_bytes.to_string()));
    }
    emit(context, EVENT_FACT, now, attributes);
}

/// Record a successful authorization for a retrieval; direct S3 GETs never reveal object keys.
pub fn observe_retrieval(context: &StorageTelemetryContext) {
    let now = chrono::Utc::now().timestamp();
    emit(
        context,
        EVENT_RETRIEVAL,
        now,
        vec![("active_day".to_string(), "1".to_string())],
    );
}

fn quota(plan: &str, storage_used_bytes: u64) -> (Option<u64>, &'static str) {
    let quota = match plan {
        "starter" => Some(100 * 1024_u64.pow(3)),
        "pro" => Some(1024_u64.pow(4)),
        _ => None,
    };
    let pressure = match quota {
        Some(quota) if storage_used_bytes.saturating_mul(100) >= quota.saturating_mul(90) => {
            "critical"
        }
        Some(quota) if storage_used_bytes.saturating_mul(100) >= quota.saturating_mul(75) => {
            "elevated"
        }
        Some(_) => "normal",
        None => "unknown",
    };
    (quota, pressure)
}