helios_persistence/composite/

router.rs

//! Query routing logic for composite storage.
//!
//! This module determines how queries are routed to backends based on
//! detected features and backend capabilities.
//!
//! # Routing Rules
//!
//! The router applies these rules (from tests):
//! - Chained parameters → Graph backend
//! - `_text`/`_content` → Search backend
//! - `:above`/`:below`/`:in`/`:not-in` → Terminology service
//! - Default → Primary backend
//! - Writes → Primary only
//! - `_include`/`_revinclude` → Primary backend (for reference resolution)

use std::collections::{HashMap, HashSet};

use crate::types::{SearchModifier, SearchParameter, SearchQuery};

use super::analyzer::{QueryAnalysis, QueryAnalyzer, QueryFeature};
use super::config::{BackendEntry, BackendRole, CompositeConfig};

/// Strategy for merging results from multiple backends.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum MergeStrategy {
    /// Results must match all backends (AND).
    /// Used when all backends must agree on matches.
    #[default]
    Intersection,

    /// Results from any backend (OR).
    /// Used for broad searches.
    Union,

    /// Primary results, enriched by secondaries.
    /// Used when primary is authoritative but secondaries add metadata.
    PrimaryEnriched,

    /// Filter secondary results through primary.
    /// Used when secondary finds candidates, primary validates.
    SecondaryFiltered,
}

/// Routing decision for a query.
#[derive(Debug, Clone)]
pub struct RoutingDecision {
    /// Primary backend ID to execute main query.
    pub primary_target: String,

    /// Additional backends for specific features.
    /// Maps feature to backend ID.
    pub auxiliary_targets: HashMap<QueryFeature, String>,

    /// Query parts for each backend.
    pub query_parts: HashMap<String, QueryPart>,

    /// Execution order for query parts.
    pub execution_order: Vec<ExecutionStep>,

    /// Strategy for merging results.
    pub merge_strategy: MergeStrategy,

    /// Analysis used for this decision.
    pub analysis: QueryAnalysis,
}

impl RoutingDecision {
    /// Returns all unique backend IDs involved in this decision.
    pub fn all_backends(&self) -> HashSet<&str> {
        let mut backends = HashSet::new();
        backends.insert(self.primary_target.as_str());
        for backend_id in self.auxiliary_targets.values() {
            backends.insert(backend_id.as_str());
        }
        backends
    }

    /// Returns true if this decision uses multiple backends.
    pub fn is_multi_backend(&self) -> bool {
        !self.auxiliary_targets.is_empty()
    }

    /// Returns true if a specific backend is used.
    pub fn uses_backend(&self, backend_id: &str) -> bool {
        self.primary_target == backend_id
            || self.auxiliary_targets.values().any(|b| b == backend_id)
    }
}

/// A part of a query to execute on a specific backend.
#[derive(Debug, Clone)]
pub struct QueryPart {
    /// Backend ID for this part.
    pub backend_id: String,

    /// Parameters for this part.
    pub parameters: Vec<SearchParameter>,

    /// Feature being handled.
    pub feature: QueryFeature,

    /// Whether this part returns only IDs (not full resources).
    pub returns_ids_only: bool,
}

impl QueryPart {
    /// Creates a new query part.
    pub fn new(backend_id: impl Into<String>, feature: QueryFeature) -> Self {
        Self {
            backend_id: backend_id.into(),
            parameters: Vec::new(),
            feature,
            returns_ids_only: false,
        }
    }

    /// Adds parameters to this part.
    pub fn with_parameters(mut self, params: Vec<SearchParameter>) -> Self {
        self.parameters = params;
        self
    }

    /// Sets whether this part returns IDs only.
    pub fn with_ids_only(mut self, ids_only: bool) -> Self {
        self.returns_ids_only = ids_only;
        self
    }
}

/// An execution step in the query plan.
#[derive(Debug, Clone)]
pub enum ExecutionStep {
    /// Execute query part on a backend.
    Execute {
        /// Backend ID.
        backend_id: String,
        /// The query part to execute.
        part_feature: QueryFeature,
    },

    /// Wait for previous steps to complete.
    Barrier(Vec<String>),

    /// Merge results from multiple backends.
    Merge {
        /// Backend IDs to merge from.
        inputs: Vec<String>,
        /// Merge strategy.
        strategy: MergeStrategy,
    },

    /// Filter results through another backend.
    Filter {
        /// Backend to filter with.
        backend_id: String,
        /// Source of IDs to filter.
        source: String,
    },

    /// Resolve includes from primary.
    ResolveIncludes {
        /// Backend for include resolution.
        backend_id: String,
    },
}

/// Backend type for routing (matches test expectations).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BackendType {
    /// Primary storage backend.
    Primary,
    /// Search optimization backend.
    Search,
    /// Graph query backend.
    Graph,
    /// Terminology service.
    Terminology,
    /// Archive storage.
    Archive,
}

impl From<BackendRole> for BackendType {
    fn from(role: BackendRole) -> Self {
        match role {
            BackendRole::Primary => BackendType::Primary,
            BackendRole::Search => BackendType::Search,
            BackendRole::Graph => BackendType::Graph,
            BackendRole::Terminology => BackendType::Terminology,
            BackendRole::Archive => BackendType::Archive,
        }
    }
}

/// Simple routing result for tests.
#[derive(Debug)]
pub struct QueryRouting {
    /// Primary backend type.
    pub primary_backend: BackendType,
    /// Auxiliary backend types.
    pub auxiliary_backends: HashSet<BackendType>,
}

/// Query router that determines execution plan.
pub struct QueryRouter {
    config: CompositeConfig,
    analyzer: QueryAnalyzer,
}

impl QueryRouter {
    /// Creates a new router with the given configuration.
    pub fn new(config: CompositeConfig) -> Self {
        Self {
            config,
            analyzer: QueryAnalyzer::new(),
        }
    }

    /// Routes a query to appropriate backends.
    pub fn route(&self, query: &SearchQuery) -> Result<RoutingDecision, RoutingError> {
        // 1. Analyze query features
        let analysis = self.analyzer.analyze(query);

        // 2. Get primary backend
        let primary = self
            .config
            .primary()
            .ok_or(RoutingError::NoPrimaryBackend)?;

        // 3. Find capable backends for each specialized feature
        let mut auxiliary_targets = HashMap::new();
        let mut query_parts = HashMap::new();

        // Route specialized features
        for feature in &analysis.specialized_features {
            if let Some(backend) = self.find_backend_for_feature(*feature, &analysis) {
                if backend.id != primary.id {
                    auxiliary_targets.insert(*feature, backend.id.clone());

                    // Create query part for this feature
                    let params = analysis
                        .feature_params
                        .get(feature)
                        .cloned()
                        .unwrap_or_default();

                    query_parts.insert(
                        backend.id.clone(),
                        QueryPart::new(&backend.id, *feature)
                            .with_parameters(params)
                            .with_ids_only(true),
                    );
                }
            }
        }

        // Create primary query part with remaining parameters
        let primary_params = analysis
            .feature_params
            .get(&QueryFeature::BasicSearch)
            .cloned()
            .unwrap_or_default();

        query_parts.insert(
            primary.id.clone(),
            QueryPart::new(&primary.id, QueryFeature::BasicSearch)
                .with_parameters(primary_params)
                .with_ids_only(false),
        );

        // 4. Build execution order
        let execution_order =
            self.build_execution_order(&analysis, &auxiliary_targets, &primary.id);

        // 5. Determine merge strategy
        let merge_strategy = self.determine_merge_strategy(&analysis, &auxiliary_targets);

        Ok(RoutingDecision {
            primary_target: primary.id.clone(),
            auxiliary_targets,
            query_parts,
            execution_order,
            merge_strategy,
            analysis,
        })
    }

    /// Simple route function matching test expectations.
    pub fn route_simple(&self, query: &SearchQuery) -> QueryRouting {
        let features = self.analyzer.analyze(query).features;

        let mut routing = QueryRouting {
            primary_backend: BackendType::Primary,
            auxiliary_backends: HashSet::new(),
        };

        // Route chained search to graph
        if features.contains(&QueryFeature::ChainedSearch)
            || features.contains(&QueryFeature::ReverseChaining)
        {
            routing.auxiliary_backends.insert(BackendType::Graph);
        }

        // Route full-text to search
        if features.contains(&QueryFeature::FullTextSearch) {
            routing.auxiliary_backends.insert(BackendType::Search);
        }

        // Route terminology to terminology service
        if features.contains(&QueryFeature::TerminologySearch) {
            routing.auxiliary_backends.insert(BackendType::Terminology);
        }

        routing
    }

    /// Finds the best backend for a feature.
    fn find_backend_for_feature(
        &self,
        feature: QueryFeature,
        _analysis: &QueryAnalysis,
    ) -> Option<&BackendEntry> {
        // First check custom routing rules
        for rule in &self.config.routing_rules {
            if rule.triggers.contains(&feature) {
                if let Some(backend) = self.config.backend(&rule.target_backend) {
                    if backend.enabled {
                        return Some(backend);
                    }
                }
            }
        }

        // Then check by role mapping
        let preferred_role = match feature {
            QueryFeature::ChainedSearch | QueryFeature::ReverseChaining => Some(BackendRole::Graph),
            QueryFeature::FullTextSearch => Some(BackendRole::Search),
            QueryFeature::TerminologySearch => Some(BackendRole::Terminology),
            _ => None,
        };

        if let Some(role) = preferred_role {
            // Find backend with this role
            let mut candidates: Vec<_> = self.config.backends_with_role(role).collect();
            candidates.sort_by_key(|b| b.priority);

            if let Some(backend) = candidates.first() {
                return Some(*backend);
            }
        }

        // Fall back to primary or any backend with the capability
        if let Some(cap) = feature.required_capability() {
            let mut capable: Vec<_> = self.config.backends_with_capability(cap).collect();
            capable.sort_by_key(|b| b.priority);
            return capable.first().copied();
        }

        self.config.primary()
    }

    /// Builds the execution order.
    fn build_execution_order(
        &self,
        analysis: &QueryAnalysis,
        auxiliary_targets: &HashMap<QueryFeature, String>,
        primary_id: &str,
    ) -> Vec<ExecutionStep> {
        let mut steps = Vec::new();

        // If there are auxiliary backends, execute them first in parallel
        if !auxiliary_targets.is_empty() {
            // Execute all auxiliary queries
            for (feature, backend_id) in auxiliary_targets {
                steps.push(ExecutionStep::Execute {
                    backend_id: backend_id.clone(),
                    part_feature: *feature,
                });
            }

            // Barrier to wait for auxiliary results
            let aux_backends: Vec<_> = auxiliary_targets.values().cloned().collect();
            steps.push(ExecutionStep::Barrier(aux_backends.clone()));

            // Merge or filter with primary
            if auxiliary_targets.len() > 1 {
                steps.push(ExecutionStep::Merge {
                    inputs: aux_backends,
                    strategy: MergeStrategy::Intersection,
                });
            }
        }

        // Execute primary query
        steps.push(ExecutionStep::Execute {
            backend_id: primary_id.to_string(),
            part_feature: QueryFeature::BasicSearch,
        });

        // If there were auxiliary results, filter through them
        if !auxiliary_targets.is_empty() {
            steps.push(ExecutionStep::Filter {
                backend_id: primary_id.to_string(),
                source: "auxiliary_results".to_string(),
            });
        }

        // Resolve includes if needed
        if analysis.has_includes() {
            steps.push(ExecutionStep::ResolveIncludes {
                backend_id: primary_id.to_string(),
            });
        }

        steps
    }

    /// Determines the merge strategy based on query analysis.
    fn determine_merge_strategy(
        &self,
        _analysis: &QueryAnalysis,
        auxiliary_targets: &HashMap<QueryFeature, String>,
    ) -> MergeStrategy {
        if auxiliary_targets.is_empty() {
            return MergeStrategy::Intersection;
        }

        // If using graph or terminology, filter secondary through primary
        if auxiliary_targets.contains_key(&QueryFeature::ChainedSearch)
            || auxiliary_targets.contains_key(&QueryFeature::ReverseChaining)
            || auxiliary_targets.contains_key(&QueryFeature::TerminologySearch)
        {
            return MergeStrategy::SecondaryFiltered;
        }

        // If using full-text search, intersect results
        if auxiliary_targets.contains_key(&QueryFeature::FullTextSearch) {
            return MergeStrategy::Intersection;
        }

        MergeStrategy::Intersection
    }

    /// Decomposes a query into backend-specific parts.
    pub fn decompose_query(&self, query: &SearchQuery) -> Vec<QueryPart> {
        let _analysis = self.analyzer.analyze(query);
        let mut parts = Vec::new();
        let mut primary_params = Vec::new();
        let mut search_params = Vec::new();
        let mut graph_params = Vec::new();
        let mut term_params = Vec::new();

        for param in &query.parameters {
            // Full-text goes to search backend
            if param.name == "_text" || param.name == "_content" {
                search_params.push(param.clone());
            }
            // Chained goes to graph
            else if !param.chain.is_empty() {
                graph_params.push(param.clone());
            }
            // Terminology modifiers
            else if matches!(
                param.modifier,
                Some(SearchModifier::Above)
                    | Some(SearchModifier::Below)
                    | Some(SearchModifier::In)
                    | Some(SearchModifier::NotIn)
            ) {
                term_params.push(param.clone());
            }
            // Default to primary
            else {
                primary_params.push(param.clone());
            }
        }

        if !primary_params.is_empty() {
            parts.push(
                QueryPart::new("primary", QueryFeature::BasicSearch)
                    .with_parameters(primary_params),
            );
        }

        if !search_params.is_empty() {
            parts.push(
                QueryPart::new("search", QueryFeature::FullTextSearch)
                    .with_parameters(search_params)
                    .with_ids_only(true),
            );
        }

        if !graph_params.is_empty() {
            parts.push(
                QueryPart::new("graph", QueryFeature::ChainedSearch)
                    .with_parameters(graph_params)
                    .with_ids_only(true),
            );
        }

        if !term_params.is_empty() {
            parts.push(
                QueryPart::new("terminology", QueryFeature::TerminologySearch)
                    .with_parameters(term_params)
                    .with_ids_only(true),
            );
        }

        parts
    }
}

/// Routing errors.
#[derive(Debug, Clone, thiserror::Error)]
pub enum RoutingError {
    /// No primary backend configured.
    #[error("no primary backend configured")]
    NoPrimaryBackend,

    /// No backend capable of handling required features.
    #[error("no backend capable of handling feature: {feature:?}")]
    NoCapableBackend {
        /// The feature that cannot be handled.
        feature: QueryFeature,
    },

    /// Backend unavailable.
    #[error("backend '{backend_id}' is unavailable")]
    BackendUnavailable {
        /// The unavailable backend ID.
        backend_id: String,
    },
}

/// Convenience function to route a query (for tests).
pub fn route_query(query: &SearchQuery) -> QueryRouting {
    let config = CompositeConfig::default();
    let router = QueryRouter::new(config);
    router.route_simple(query)
}

/// Convenience function to decompose a query (for tests).
pub fn decompose_query(query: &SearchQuery) -> Vec<QueryPart> {
    let config = CompositeConfig::default();
    let router = QueryRouter::new(config);
    router.decompose_query(query)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::BackendKind;
    use crate::types::{ChainedParameter, SearchModifier, SearchParamType, SearchValue};

    fn test_config() -> CompositeConfig {
        CompositeConfig::builder()
            .primary("sqlite", BackendKind::Sqlite)
            .search_backend("es", BackendKind::Elasticsearch)
            .graph_backend("neo4j", BackendKind::Neo4j)
            .build()
            .unwrap()
    }

    #[test]
    fn test_route_simple_query_to_primary() {
        let query = SearchQuery::new("Patient").with_parameter(SearchParameter {
            name: "_id".to_string(),
            param_type: SearchParamType::Token,
            modifier: None,
            values: vec![SearchValue::eq("patient-123")],
            chain: vec![],
            components: vec![],
        });

        let routing = route_query(&query);
        assert_eq!(routing.primary_backend, BackendType::Primary);
        assert!(routing.auxiliary_backends.is_empty());
    }

    #[test]
    fn test_route_chained_search_to_graph() {
        let query = SearchQuery::new("Observation").with_parameter(SearchParameter {
            name: "name".to_string(),
            param_type: SearchParamType::String,
            modifier: None,
            values: vec![SearchValue::string("Smith")],
            chain: vec![ChainedParameter {
                reference_param: "subject".to_string(),
                target_type: Some("Patient".to_string()),
                target_param: "name".to_string(),
            }],
            components: vec![],
        });

        let routing = route_query(&query);
        assert!(routing.auxiliary_backends.contains(&BackendType::Graph));
    }

    #[test]
    fn test_route_fulltext_to_search() {
        let query = SearchQuery::new("Patient").with_parameter(SearchParameter {
            name: "_text".to_string(),
            param_type: SearchParamType::String,
            modifier: None,
            values: vec![SearchValue::string("cardiac")],
            chain: vec![],
            components: vec![],
        });

        let routing = route_query(&query);
        assert!(routing.auxiliary_backends.contains(&BackendType::Search));
    }

    #[test]
    fn test_route_terminology_to_terminology_service() {
        let query = SearchQuery::new("Observation").with_parameter(SearchParameter {
            name: "code".to_string(),
            param_type: SearchParamType::Token,
            modifier: Some(SearchModifier::Below),
            values: vec![SearchValue::token(Some("http://loinc.org"), "8867-4")],
            chain: vec![],
            components: vec![],
        });

        let routing = route_query(&query);
        assert!(
            routing
                .auxiliary_backends
                .contains(&BackendType::Terminology)
        );
    }

    #[test]
    fn test_route_complex_query_to_multiple_backends() {
        let query = SearchQuery::new("Observation")
            .with_parameter(SearchParameter {
                name: "name".to_string(),
                param_type: SearchParamType::String,
                modifier: None,
                values: vec![SearchValue::string("Smith")],
                chain: vec![ChainedParameter {
                    reference_param: "subject".to_string(),
                    target_type: Some("Patient".to_string()),
                    target_param: "name".to_string(),
                }],
                components: vec![],
            })
            .with_parameter(SearchParameter {
                name: "_text".to_string(),
                param_type: SearchParamType::String,
                modifier: None,
                values: vec![SearchValue::string("cardiac")],
                chain: vec![],
                components: vec![],
            })
            .with_parameter(SearchParameter {
                name: "code".to_string(),
                param_type: SearchParamType::Token,
                modifier: Some(SearchModifier::Below),
                values: vec![SearchValue::token(Some("http://loinc.org"), "8867-4")],
                chain: vec![],
                components: vec![],
            });

        let routing = route_query(&query);
        assert!(!routing.auxiliary_backends.is_empty());
    }

    #[test]
    fn test_decompose_query() {
        let query = SearchQuery::new("Observation")
            .with_parameter(SearchParameter {
                name: "code".to_string(),
                param_type: SearchParamType::Token,
                modifier: None,
                values: vec![SearchValue::token(Some("http://loinc.org"), "8867-4")],
                chain: vec![],
                components: vec![],
            })
            .with_parameter(SearchParameter {
                name: "_text".to_string(),
                param_type: SearchParamType::String,
                modifier: None,
                values: vec![SearchValue::string("cardiac")],
                chain: vec![],
                components: vec![],
            });

        let parts = decompose_query(&query);

        assert!(!parts.is_empty());
        assert!(parts.iter().any(|p| p.backend_id == "primary"));
        assert!(parts.iter().any(|p| p.backend_id == "search"));
    }

    #[test]
    fn test_routing_decision_with_config() {
        let config = test_config();
        let router = QueryRouter::new(config);

        let query = SearchQuery::new("Patient").with_parameter(SearchParameter {
            name: "_text".to_string(),
            param_type: SearchParamType::String,
            modifier: None,
            values: vec![SearchValue::string("cardiac")],
            chain: vec![],
            components: vec![],
        });

        let decision = router.route(&query).unwrap();

        assert_eq!(decision.primary_target, "sqlite");
        assert!(
            decision
                .auxiliary_targets
                .contains_key(&QueryFeature::FullTextSearch)
        );
    }
}