erpc_analysis/

db_trait.rs

use async_trait::async_trait;
use std::collections::HashMap;
use thiserror::Error;

use crate::models::metrics::{GraphMetrics, NodeMetrics};
use crate::models::partitions::ComponentAnalysisResult;

/// Represents errors that can occur during analysis operations, including
/// database interactions.
#[derive(Error, Debug)]
pub enum AnalysisError {
    /// An error originating from the underlying database driver (neo4rs).
    #[error("Database Driver Error: {0}")]
    DriverError(#[from] neo4rs::Error),

    /// Failed to establish an initial connection to the database.
    #[error("Failed to connect to database: {0}")]
    ConnectionFailed(String),

    /// An error occurred while attempting to drop an existing
    /// GDS graph projection.
    #[error(
        "Failed to drop graph projection '{projection_name}': {source_error}"
    )]
    ProjectionDropFailed {
        projection_name: String,
        source_error: String,
    },

    /// An error occurred during the creation of a GDS graph projection.
    #[error(
        "Failed to create graph projection '{projection_name}': {source_error}"
    )]
    ProjectionCreationFailed {
        projection_name: String,
        source_error: String,
    },

    /// A generic error for database query execution failures not covered by
    /// other specific variants.
    #[error("Database query execution failed: {0}")]
    QueryFailed(String),

    /// Indicates that a requested GDS graph projection was not found.
    #[error("GDS graph projection '{0}' not found.")]
    ProjectionNotFound(String),

    /// Indicates that a GDS graph projection already exists when it
    /// was not expected.
    #[error("GDS graph projection '{0}' already exists.")]
    ProjectionAlreadyExists(String),

    /// An error occurred during the execution of a graph algorithm.
    #[error("Graph algorithm execution failed: {0}")]
    AlgorithmError(String),

    /// For configuration-related errors during analysis setup.
    #[error("Analysis configuration error: {0}")]
    ConfigurationError(String),

    /// For I/O errors that might occur (e.g., reading a script file
    /// for a query).
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),

    /// For any other general error encountered during analysis.
    #[error("An unexpected analysis error occurred: {0}")]
    Generic(String),
}

/// Parameters for creating a GDS graph projection.
#[derive(Debug, Clone)]
pub struct GraphProjectionParams {
    pub projection_name: String,
    pub node_label: String,
    pub relationship_types: HashMap<String, String>,
    pub relationship_properties_to_project: Option<Vec<String>>,
}

/// Trait defining the interface for database operations required by
/// the eRPC analysis engine.
#[async_trait]
pub trait AnalysisDatabase: Send + Sync {
    /// Creates or recreates a GDS graph projection.
    /// If a projection with the same name exists, it should ideally be dropped
    /// and recreated.
    async fn create_graph_projection(
        &self,
        params: &GraphProjectionParams,
    ) -> Result<(), AnalysisError>;

    /// Deletes an existing GDS graph projection if it exists.
    /// Should succeed even if the projection does not exist (making
    /// it idempotent).
    async fn delete_graph_projection(
        &self,
        projection_name: &str,
    ) -> Result<(), AnalysisError>;

    /// Checks if a GDS graph projection with the given name exists.
    async fn check_graph_projection_exists(
        &self,
        projection_name: &str,
    ) -> Result<bool, AnalysisError>;

    /// Calculates node-level degree metrics for all nodes in a
    /// given GDS projection.
    /// Returns a vector of NodeMetrics containing in-degree,
    /// out-degree, and total degree for each node in the graph.
    async fn calculate_node_degrees(
        &self,
        projection_name: &str,
    ) -> Result<Vec<NodeMetrics>, AnalysisError>;

    /// Calculates comprehensive graph metrics for a given GDS projection
    /// including basic counts, degree distribution, and degree statistics.
    async fn calculate_graph_metrics(
        &self,
        projection_name: &str,
    ) -> Result<GraphMetrics, AnalysisError>;

    /// Calculates weakly connected components using Neo4j GDS WCC algorithm.
    /// Returns analysis results containing components, sizes, and statistics.
    async fn calculate_weakly_connected_components(
        &self,
        projection_name: &str,
    ) -> Result<ComponentAnalysisResult, AnalysisError>;

    /// Calculates strongly connected components using Neo4j GDS SCC algorithm.
    /// Returns analysis results containing components, sizes, and statistics.
    async fn calculate_strongly_connected_components(
        &self,
        projection_name: &str,
    ) -> Result<ComponentAnalysisResult, AnalysisError>;
}

/// Mock database implementation for testing purposes
pub mod mock {
    use super::*;
    use std::collections::HashMap;
    use std::sync::RwLock;

    use crate::models::partitions::ConnectedComponent;

    /// Simple in-memory mock database for testing
    #[derive(Debug)]
    pub struct MockDatabase {
        pub projections: RwLock<HashMap<String, MockProjection>>,
        pub should_fail_on: Option<String>, // Operation that should fail
        pub call_count: RwLock<HashMap<String, usize>>, // Track method calls
    }

    #[derive(Debug, Clone)]
    pub struct MockProjection {
        pub name: String,
        pub node_count: i64,
        pub relationship_count: i64,
        pub nodes: Vec<NodeMetrics>,
    }

    impl Default for MockDatabase {
        fn default() -> Self {
            Self {
                projections: RwLock::new(HashMap::new()),
                should_fail_on: None,
                call_count: RwLock::new(HashMap::new()),
            }
        }
    }

    impl MockDatabase {
        pub fn new() -> Self {
            Self::default()
        }

        /// Add a test projection with predefined data
        pub fn with_projection(
            self,
            name: &str,
            nodes: Vec<NodeMetrics>,
        ) -> Self {
            let node_count = nodes.len() as i64;
            let relationship_count =
                nodes.iter().map(|n| n.total_degree).sum::<i64>();

            self.projections.write().unwrap().insert(
                name.to_string(),
                MockProjection {
                    name: name.to_string(),
                    node_count,
                    relationship_count,
                    nodes,
                },
            );
            self
        }

        /// Configure the mock to fail on a specific operation
        pub fn fail_on(mut self, operation: &str) -> Self {
            self.should_fail_on = Some(operation.to_string());
            self
        }

        /// Get the number of times a method was called
        pub fn get_call_count(&self, method: &str) -> usize {
            self.call_count
                .read()
                .unwrap()
                .get(method)
                .copied()
                .unwrap_or(0)
        }

        fn increment_call_count(&self, method: &str) {
            *self
                .call_count
                .write()
                .unwrap()
                .entry(method.to_string())
                .or_insert(0) += 1;
        }

        fn should_fail(&self, operation: &str) -> bool {
            self.should_fail_on
                .as_ref()
                .is_some_and(|fail_op| fail_op == operation)
        }
    }

    #[async_trait]
    impl AnalysisDatabase for MockDatabase {
        async fn create_graph_projection(
            &self,
            params: &GraphProjectionParams,
        ) -> Result<(), AnalysisError> {
            self.increment_call_count("create_graph_projection");

            if self.should_fail("create_graph_projection") {
                return Err(AnalysisError::ProjectionCreationFailed {
                    projection_name: params.projection_name.clone(),
                    source_error: "Mock failure".to_string(),
                });
            }

            // Mock successful creation
            self.projections.write().unwrap().insert(
                params.projection_name.clone(),
                MockProjection {
                    name: params.projection_name.clone(),
                    node_count: 0,
                    relationship_count: 0,
                    nodes: vec![],
                },
            );

            Ok(())
        }

        async fn delete_graph_projection(
            &self,
            projection_name: &str,
        ) -> Result<(), AnalysisError> {
            self.increment_call_count("delete_graph_projection");

            if self.should_fail("delete_graph_projection") {
                return Err(AnalysisError::ProjectionDropFailed {
                    projection_name: projection_name.to_string(),
                    source_error: "Mock failure".to_string(),
                });
            }

            // Always succeed (idempotent)
            self.projections.write().unwrap().remove(projection_name);
            Ok(())
        }

        async fn check_graph_projection_exists(
            &self,
            projection_name: &str,
        ) -> Result<bool, AnalysisError> {
            self.increment_call_count("check_graph_projection_exists");

            if self.should_fail("check_graph_projection_exists") {
                return Err(AnalysisError::QueryFailed(
                    "Mock failure".to_string(),
                ));
            }

            Ok(self
                .projections
                .read()
                .unwrap()
                .contains_key(projection_name))
        }

        async fn calculate_node_degrees(
            &self,
            projection_name: &str,
        ) -> Result<Vec<NodeMetrics>, AnalysisError> {
            self.increment_call_count("calculate_node_degrees");

            if self.should_fail("calculate_node_degrees") {
                return Err(AnalysisError::QueryFailed(
                    "Mock failure".to_string(),
                ));
            }

            match self.projections.read().unwrap().get(projection_name) {
                Some(projection) => Ok(projection.nodes.clone()),
                None => Err(AnalysisError::ProjectionNotFound(
                    projection_name.to_string(),
                )),
            }
        }

        async fn calculate_graph_metrics(
            &self,
            projection_name: &str,
        ) -> Result<GraphMetrics, AnalysisError> {
            self.increment_call_count("calculate_graph_metrics");

            if self.should_fail("calculate_graph_metrics") {
                return Err(AnalysisError::QueryFailed(
                    "Mock failure".to_string(),
                ));
            }

            match self.projections.read().unwrap().get(projection_name) {
                Some(projection) => {
                    let mut degree_distribution = HashMap::new();
                    let mut total_degree_sum = 0i64;
                    let mut max_degree = 0i64;
                    let mut min_degree = i64::MAX;

                    for node in &projection.nodes {
                        let degree = node.total_degree;
                        *degree_distribution.entry(degree).or_insert(0) += 1;
                        total_degree_sum += degree;
                        max_degree = max_degree.max(degree);
                        min_degree = min_degree.min(degree);
                    }

                    if min_degree == i64::MAX {
                        min_degree = 0;
                    }

                    let average_degree = if !projection.nodes.is_empty() {
                        total_degree_sum as f64 / projection.nodes.len() as f64
                    } else {
                        0.0
                    };

                    Ok(GraphMetrics {
                        node_count: Some(projection.node_count),
                        relationship_count: Some(
                            projection.relationship_count,
                        ),
                        degree_distribution: Some(degree_distribution),
                        average_degree: Some(average_degree),
                        max_degree: Some(max_degree),
                        min_degree: Some(min_degree),
                    })
                }
                None => Err(AnalysisError::ProjectionNotFound(
                    projection_name.to_string(),
                )),
            }
        }

        async fn calculate_weakly_connected_components(
            &self,
            projection_name: &str,
        ) -> Result<ComponentAnalysisResult, AnalysisError> {
            self.increment_call_count("calculate_weakly_connected_components");

            if self.should_fail("calculate_weakly_connected_components") {
                return Err(AnalysisError::AlgorithmError(
                    "Mock failure".to_string(),
                ));
            }

            match self.projections.read().unwrap().get(projection_name) {
                Some(projection) => {
                    // Mock WCC analysis: create realistic mock components
                    let mut components = Vec::new();

                    if !projection.nodes.is_empty() {
                        // For mock purposes, create varied component structures
                        if projection.nodes.len() > 2 {
                            // First component with first half of nodes
                            let mid = projection.nodes.len() / 2;
                            let first_component_fingerprints: Vec<String> =
                                projection.nodes[0..mid]
                                    .iter()
                                    .map(|n| n.fingerprint.clone())
                                    .collect();

                            components.push(ConnectedComponent {
                                component_id: 0,
                                relay_fingerprints:
                                    first_component_fingerprints.clone(),
                                size: first_component_fingerprints.len(),
                            });

                            // Second component with remaining nodes
                            let second_component_fingerprints: Vec<String> =
                                projection.nodes[mid..]
                                    .iter()
                                    .map(|n| n.fingerprint.clone())
                                    .collect();

                            components.push(ConnectedComponent {
                                component_id: 1,
                                relay_fingerprints:
                                    second_component_fingerprints.clone(),
                                size: second_component_fingerprints.len(),
                            });
                        } else {
                            // Single component with all nodes
                            let all_fingerprints: Vec<String> = projection
                                .nodes
                                .iter()
                                .map(|n| n.fingerprint.clone())
                                .collect();

                            components.push(ConnectedComponent {
                                component_id: 0,
                                relay_fingerprints: all_fingerprints.clone(),
                                size: all_fingerprints.len(),
                            });
                        }
                    }

                    // Sort components by size (largest first)
                    components.sort_by(|a, b| b.size.cmp(&a.size));

                    // Calculate statistics
                    let total_components = components.len();
                    let largest_component_size =
                        components.first().map(|c| c.size).unwrap_or(0);
                    let smallest_component_size =
                        components.last().map(|c| c.size).unwrap_or(0);

                    // Calculate size distribution
                    let mut component_size_distribution = HashMap::new();
                    for component in &components {
                        *component_size_distribution
                            .entry(component.size)
                            .or_insert(0) += 1;
                    }

                    // Calculate isolation ratio
                    let total_nodes = projection.nodes.len();
                    let isolation_ratio = if total_nodes > 0 {
                        (largest_component_size as f64 / total_nodes as f64)
                            * 100.0
                    } else {
                        0.0
                    };

                    Ok(ComponentAnalysisResult {
                        components,
                        total_components: Some(total_components),
                        largest_component_size: Some(largest_component_size),
                        smallest_component_size: Some(smallest_component_size),
                        component_size_distribution: Some(
                            component_size_distribution,
                        ),
                        isolation_ratio: Some(isolation_ratio),
                    })
                }
                None => Err(AnalysisError::ProjectionNotFound(
                    projection_name.to_string(),
                )),
            }
        }

        async fn calculate_strongly_connected_components(
            &self,
            projection_name: &str,
        ) -> Result<ComponentAnalysisResult, AnalysisError> {
            self.increment_call_count(
                "calculate_strongly_connected_components",
            );

            if self.should_fail("calculate_strongly_connected_components") {
                return Err(AnalysisError::AlgorithmError(
                    "Mock failure".to_string(),
                ));
            }

            match self.projections.read().unwrap().get(projection_name) {
                Some(projection) => {
                    // Mock SCC analysis: create realistic mock components
                    let mut components = Vec::new();

                    if !projection.nodes.is_empty() {
                        // For mock purposes, create varied component structures
                        if projection.nodes.len() > 2 {
                            // First component with first half of nodes
                            let mid = projection.nodes.len() / 2;
                            let first_component_fingerprints: Vec<String> =
                                projection.nodes[0..mid]
                                    .iter()
                                    .map(|n| n.fingerprint.clone())
                                    .collect();

                            components.push(ConnectedComponent {
                                component_id: 0,
                                relay_fingerprints:
                                    first_component_fingerprints.clone(),
                                size: first_component_fingerprints.len(),
                            });

                            // Second component with remaining nodes
                            let second_component_fingerprints: Vec<String> =
                                projection.nodes[mid..]
                                    .iter()
                                    .map(|n| n.fingerprint.clone())
                                    .collect();

                            components.push(ConnectedComponent {
                                component_id: 1,
                                relay_fingerprints:
                                    second_component_fingerprints.clone(),
                                size: second_component_fingerprints.len(),
                            });
                        } else {
                            // Single component with all nodes
                            let all_fingerprints: Vec<String> = projection
                                .nodes
                                .iter()
                                .map(|n| n.fingerprint.clone())
                                .collect();

                            components.push(ConnectedComponent {
                                component_id: 0,
                                relay_fingerprints: all_fingerprints.clone(),
                                size: all_fingerprints.len(),
                            });
                        }
                    }

                    // Sort components by size (largest first)
                    components.sort_by(|a, b| b.size.cmp(&a.size));

                    // Calculate statistics
                    let total_components = components.len();
                    let largest_component_size =
                        components.first().map(|c| c.size).unwrap_or(0);
                    let smallest_component_size =
                        components.last().map(|c| c.size).unwrap_or(0);

                    // Calculate size distribution correctly
                    let mut component_size_distribution = HashMap::new();
                    for component in &components {
                        *component_size_distribution
                            .entry(component.size)
                            .or_insert(0) += 1;
                    }

                    // Calculate isolation ratio
                    let total_nodes = projection.nodes.len();
                    let isolation_ratio = if total_nodes > 0 {
                        (largest_component_size as f64 / total_nodes as f64)
                            * 100.0
                    } else {
                        0.0
                    };

                    Ok(ComponentAnalysisResult {
                        components,
                        total_components: Some(total_components),
                        largest_component_size: Some(largest_component_size),
                        smallest_component_size: Some(smallest_component_size),
                        component_size_distribution: Some(
                            component_size_distribution,
                        ),
                        isolation_ratio: Some(isolation_ratio),
                    })
                }
                None => Err(AnalysisError::ProjectionNotFound(
                    projection_name.to_string(),
                )),
            }
        }
    }
}

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

    #[tokio::test]
    async fn test_mock_database_basic_operations() {
        let db = MockDatabase::new();

        // Test projection doesn't exist initially
        let exists = db.check_graph_projection_exists("test").await.unwrap();
        assert!(!exists);

        // Create projection
        let params = GraphProjectionParams {
            projection_name: "test".to_string(),
            node_label: "Relay".to_string(),
            relationship_types: HashMap::new(),
            relationship_properties_to_project: None,
        };

        db.create_graph_projection(&params).await.unwrap();

        // Check it exists now
        let exists = db.check_graph_projection_exists("test").await.unwrap();
        assert!(exists);

        // Delete projection
        db.delete_graph_projection("test").await.unwrap();

        // Check it doesn't exist anymore
        let exists = db.check_graph_projection_exists("test").await.unwrap();
        assert!(!exists);
    }

    #[tokio::test]
    async fn test_mock_with_test_data() {
        let nodes = vec![
            NodeMetrics {
                fingerprint: "RELAY001".to_string(),
                in_degree: 3,
                out_degree: 2,
                total_degree: 5,
            },
            NodeMetrics {
                fingerprint: "RELAY002".to_string(),
                in_degree: 1,
                out_degree: 4,
                total_degree: 5,
            },
        ];

        let db = MockDatabase::new().with_projection("test_proj", nodes);

        let result = db.calculate_node_degrees("test_proj").await.unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].fingerprint, "RELAY001");
        assert_eq!(result[0].total_degree, 5);
    }

    #[tokio::test]
    async fn test_mock_wcc_edge_cases() {
        let db = MockDatabase::new();

        // Test projection not found
        let result = db
            .calculate_weakly_connected_components("nonexistent")
            .await;
        assert!(result.is_err());

        // Test with only 2 nodes (should create single component)
        let nodes = vec![
            NodeMetrics {
                fingerprint: "RELAY001".to_string(),
                in_degree: 1,
                out_degree: 1,
                total_degree: 2,
            },
            NodeMetrics {
                fingerprint: "RELAY002".to_string(),
                in_degree: 1,
                out_degree: 1,
                total_degree: 2,
            },
        ];

        let db_with_data =
            MockDatabase::new().with_projection("single_component", nodes);

        let result = db_with_data
            .calculate_weakly_connected_components("single_component")
            .await
            .unwrap();

        // Should have 1 component with both nodes
        assert_eq!(result.total_components.unwrap(), 1);
        assert_eq!(result.components.len(), 1);
        assert_eq!(result.components[0].size, 2);
        assert_eq!(result.isolation_ratio.unwrap(), 100.0);
    }

    #[tokio::test]
    async fn test_mock_scc_edge_cases_and_failures() {
        let db = MockDatabase::new();

        // Test projection not found
        let result = db
            .calculate_strongly_connected_components("nonexistent")
            .await;
        assert!(result.is_err());

        // Test with only 2 nodes (should create single component)
        let nodes = vec![
            NodeMetrics {
                fingerprint: "SCC_SINGLE_001".to_string(),
                in_degree: 1,
                out_degree: 1,
                total_degree: 2,
            },
            NodeMetrics {
                fingerprint: "SCC_SINGLE_002".to_string(),
                in_degree: 1,
                out_degree: 1,
                total_degree: 2,
            },
        ];

        let db_single =
            MockDatabase::new().with_projection("single_scc", nodes);

        let result = db_single
            .calculate_strongly_connected_components("single_scc")
            .await
            .unwrap();

        // Should have only 1 component
        assert_eq!(result.total_components.unwrap(), 1);
        assert_eq!(result.components.len(), 1);
        assert_eq!(result.components[0].size, 2);
        assert_eq!(result.isolation_ratio.unwrap(), 100.0);

        // Test failure scenario
        let fail_nodes = vec![NodeMetrics {
            fingerprint: "FAIL_RELAY001".to_string(),
            in_degree: 1,
            out_degree: 1,
            total_degree: 2,
        }];

        let db_fail = MockDatabase::new()
            .with_projection("fail_test", fail_nodes)
            .fail_on("calculate_strongly_connected_components");

        let fail_result = db_fail
            .calculate_strongly_connected_components("fail_test")
            .await;
        assert!(fail_result.is_err());
        assert_eq!(
            db_fail.get_call_count("calculate_strongly_connected_components"),
            1
        );
    }
}