amql_engine/

pipeline.rs

//! QueryPipeline — composable builder for running AQL queries with pluggable middleware.
//!
//! External consumers (e.g. datum) construct a pipeline once, register resolvers,
//! extractors, and middleware, then call `query` or `select` repeatedly.

use crate::code_cache::CodeCache;
use crate::error::AqlError;
use crate::extractor::{BuiltinExtractor, ExtractorRegistry};
use crate::query::{unified_query, QueryResult};
use crate::resolver::{CodeResolver, ResolverRegistry};
use crate::store::{Annotation, AnnotationStore};
use crate::types::{ProjectRoot, Scope, SelectorStr};

/// Context passed to each middleware.
///
/// Contains the selector and scope that produced the results.
/// Read-only — middleware transforms results, not the query.
pub struct QueryContext<'a> {
    pub selector: &'a SelectorStr,
    pub scope: &'a Scope,
}

/// A middleware that can transform query results after AQL runs them.
///
/// Implement this trait to filter, sort, augment, or replace results.
/// Middleware are applied in registration order.
///
/// # Example
///
/// ```rust
/// use amql_engine::{QueryContext, QueryMiddleware, QueryResult};
///
/// struct OwnerFilter { owner: String }
///
/// impl QueryMiddleware for OwnerFilter {
///     fn name(&self) -> &str { "owner-filter" }
///     fn transform(&self, results: Vec<QueryResult>, _ctx: &QueryContext<'_>) -> Vec<QueryResult> {
///         results.into_iter()
///             .filter(|r| r.annotations.iter()
///                 .any(|a| a.attrs.get("owner")
///                     .and_then(|v| v.as_str())
///                     .map(|v| v == self.owner)
///                     .unwrap_or(false)))
///             .collect()
///     }
/// }
/// ```
pub trait QueryMiddleware: Send + Sync {
    /// Unique name for this middleware.
    fn name(&self) -> &str;
    /// Transform query results. Called after query execution, in registration order.
    fn transform(&self, results: Vec<QueryResult>, ctx: &QueryContext<'_>) -> Vec<QueryResult>;
}

/// Builder and runner for AQL queries.
///
/// Wraps all registries and state needed to run queries. Create once,
/// call `query` or `select` many times.
///
/// # Example
///
/// ```rust,no_run
/// use amql_engine::{QueryPipeline, ProjectRoot};
///
/// let root = ProjectRoot::from(std::path::Path::new("."));
/// let mut pipeline = QueryPipeline::builder(&root).build();
/// let results = pipeline.query("route[method=\"GET\"]", "src/").unwrap();
/// ```
pub struct QueryPipeline {
    root: ProjectRoot,
    store: AnnotationStore,
    cache: CodeCache,
    resolvers: ResolverRegistry,
    extractors: ExtractorRegistry,
    middleware: Vec<Box<dyn QueryMiddleware>>,
}

/// Builder for `QueryPipeline`.
pub struct QueryPipelineBuilder {
    root: ProjectRoot,
    resolvers: ResolverRegistry,
    extractors: ExtractorRegistry,
    middleware: Vec<Box<dyn QueryMiddleware>>,
}

impl QueryPipelineBuilder {
    /// Register an additional resolver. Called after `QueryPipeline::builder`.
    pub fn with_resolver(mut self, resolver: impl CodeResolver + 'static) -> Self {
        self.resolvers.register(Box::new(resolver));
        self
    }

    /// Register an additional extractor.
    pub fn with_extractor(mut self, extractor: impl BuiltinExtractor + 'static) -> Self {
        self.extractors.register(Box::new(extractor));
        self
    }

    /// Register a middleware. Applied in order after query execution.
    pub fn with_middleware(mut self, middleware: impl QueryMiddleware + 'static) -> Self {
        self.middleware.push(Box::new(middleware));
        self
    }

    /// Build the pipeline, loading all annotations from the project.
    pub fn build(self) -> QueryPipeline {
        let mut store = AnnotationStore::new(self.root.as_ref());
        store.load_all_from_locator();
        let cache = CodeCache::new(self.root.as_ref());
        QueryPipeline {
            root: self.root,
            store,
            cache,
            resolvers: self.resolvers,
            extractors: self.extractors,
            middleware: self.middleware,
        }
    }
}

impl QueryPipeline {
    /// Create a builder pre-loaded with all default resolvers and extractors.
    pub fn builder(root: &ProjectRoot) -> QueryPipelineBuilder {
        QueryPipelineBuilder {
            root: root.clone(),
            resolvers: ResolverRegistry::with_defaults(),
            extractors: ExtractorRegistry::with_defaults(),
            middleware: Vec::new(),
        }
    }

    /// Run a unified query (code + annotations) and apply middleware.
    pub fn query(&mut self, selector: &str, scope: &str) -> Result<Vec<QueryResult>, AqlError> {
        let selector_str = SelectorStr::from(selector);
        let scope_val = Scope::from(scope);
        let ctx = QueryContext {
            selector: &selector_str,
            scope: &scope_val,
        };

        let mut results = unified_query(
            selector,
            &scope_val,
            &mut self.cache,
            &mut self.store,
            &self.resolvers,
            None,
        )?;

        for mw in &self.middleware {
            results = mw.transform(results, &ctx);
        }

        Ok(results)
    }

    /// Select annotations by selector (no code parsing).
    pub fn select(&self, selector: &str, scope: Option<&str>) -> Result<Vec<Annotation>, AqlError> {
        self.store.select(selector, None, scope, None)
    }

    /// Reload all annotations from disk (call after external writes).
    pub fn reload(&mut self) {
        self.store.load_all_from_locator();
    }

    /// Return the extractor registry for this pipeline.
    pub fn extractors(&self) -> &ExtractorRegistry {
        &self.extractors
    }

    /// Return the project root for this pipeline.
    pub fn root(&self) -> &ProjectRoot {
        &self.root
    }
}