cloudlysis/aggregate/src/query/projection.rs

use super::AggregateProjection;
use crate::types::{AggregateId, AggregateType, Event, TenantId, Version};
use serde_json::Value as JsonValue;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Debug, Clone)]
pub struct ProjectionConfig {
    pub batch_size: usize,
    pub projection_timeout_ms: u64,
}

impl Default for ProjectionConfig {
    fn default() -> Self {
        Self {
            batch_size: 100,
            projection_timeout_ms: 5000,
        }
    }
}

pub struct StateProjection {
    config: ProjectionConfig,
    handlers: Arc<RwLock<HashMap<String, ProjectionHandler>>>,
}

type ProjectionHandler = Box<dyn Fn(&Event) -> Option<AggregateProjection> + Send + Sync>;

impl StateProjection {
    pub fn new(config: ProjectionConfig) -> Self {
        Self {
            config,
            handlers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub fn new_default() -> Self {
        Self::new(ProjectionConfig::default())
    }

    pub async fn register_handler<F>(&self, aggregate_type: &str, handler: F)
    where
        F: Fn(&Event) -> Option<AggregateProjection> + Send + Sync + 'static,
    {
        let mut handlers = self.handlers.write().await;
        handlers.insert(aggregate_type.to_string(), Box::new(handler));
    }

    pub async fn project_event(&self, event: &Event) -> Option<AggregateProjection> {
        let handlers = self.handlers.read().await;
        let aggregate_type = event.aggregate_type.as_str();

        handlers.get(aggregate_type).and_then(|h| h(event))
    }

    pub async fn project_events(&self, events: &[Event]) -> Vec<AggregateProjection> {
        let mut projections = Vec::with_capacity(events.len().min(self.config.batch_size));

        for event in events.iter().take(self.config.batch_size) {
            if let Some(proj) = self.project_event(event).await {
                projections.push(proj);
            }
        }

        projections
    }

    pub fn default_projection_from_event(event: &Event) -> AggregateProjection {
        AggregateProjection::new(
            event.tenant_id.as_str(),
            event.aggregate_id.to_string(),
            event.aggregate_type.as_str(),
            event.version.as_u64(),
            event.payload.clone(),
        )
    }

    pub fn default_projection_from_state(
        tenant_id: &TenantId,
        aggregate_id: &AggregateId,
        aggregate_type: &AggregateType,
        version: &Version,
        state: &JsonValue,
    ) -> AggregateProjection {
        AggregateProjection::new(
            tenant_id.as_str(),
            aggregate_id.to_string(),
            aggregate_type.as_str(),
            version.as_u64(),
            state.clone(),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use serde_json::json;

    fn create_test_event(tenant: &str, version: u64, event_type: &str) -> Event {
        Event {
            event_id: uuid::Uuid::now_v7(),
            tenant_id: TenantId::new(tenant),
            aggregate_id: AggregateId::new_v7(),
            aggregate_type: AggregateType::from("Account"),
            version: Version::from(version),
            event_type: event_type.to_string(),
            payload: json!({"amount": 100}),
            timestamp: Utc::now(),
            command_id: uuid::Uuid::nil(),
            correlation_id: None,
            traceparent: None,
        }
    }

    #[tokio::test]
    async fn state_projection_registers_handler() {
        let projection = StateProjection::new_default();

        projection
            .register_handler("Account", |event| {
                Some(AggregateProjection::new(
                    event.tenant_id.as_str(),
                    event.aggregate_id.to_string(),
                    "Account",
                    event.version.as_u64(),
                    event.payload.clone(),
                ))
            })
            .await;

        let event = create_test_event("tenant-a", 1, "deposited");
        let result = projection.project_event(&event).await;

        assert!(result.is_some());
        let proj = result.unwrap();
        assert_eq!(proj.aggregate_type, "Account");
    }

    #[tokio::test]
    async fn state_projection_project_events_batch() {
        let projection = StateProjection::new_default();

        projection
            .register_handler("Account", |event| {
                Some(AggregateProjection::new(
                    event.tenant_id.as_str(),
                    event.aggregate_id.to_string(),
                    "Account",
                    event.version.as_u64(),
                    event.payload.clone(),
                ))
            })
            .await;

        let events = vec![
            create_test_event("tenant-a", 1, "deposited"),
            create_test_event("tenant-a", 1, "deposited"),
            create_test_event("tenant-a", 1, "deposited"),
        ];

        let projections = projection.project_events(&events).await;

        assert_eq!(projections.len(), 3);
    }

    #[tokio::test]
    async fn state_projection_no_handler_returns_none() {
        let projection = StateProjection::new_default();

        let event = create_test_event("tenant-a", 1, "deposited");
        let result = projection.project_event(&event).await;

        assert!(result.is_none());
    }

    #[test]
    fn default_projection_from_event() {
        let event = create_test_event("tenant-a", 5, "deposited");
        let proj = StateProjection::default_projection_from_event(&event);

        assert_eq!(proj.tenant_id, "tenant-a");
        assert_eq!(proj.version, 5);
        assert_eq!(proj.state["amount"], 100);
    }

    #[test]
    fn default_projection_from_state() {
        let tenant_id = TenantId::new("tenant-a");
        let aggregate_id = AggregateId::new_v7();
        let aggregate_type = AggregateType::from("Account");
        let version = Version::from(10);
        let state = json!({"balance": 1000});

        let proj = StateProjection::default_projection_from_state(
            &tenant_id,
            &aggregate_id,
            &aggregate_type,
            &version,
            &state,
        );

        assert_eq!(proj.tenant_id, "tenant-a");
        assert_eq!(proj.aggregate_type, "Account");
        assert_eq!(proj.version, 10);
        assert_eq!(proj.state["balance"], 1000);
    }

    #[test]
    fn projection_config_defaults() {
        let config = ProjectionConfig::default();
        assert_eq!(config.batch_size, 100);
        assert_eq!(config.projection_timeout_ms, 5000);
    }
}