hamelin_lib 0.21.4

//! Query builder for constructing [`Query`] AST nodes.
//!
//! - **Tabular `DEF`s** — [`QueryBuilder::def_pipeline`] / [`QueryBuilderWithMain::def_pipeline`].
//! - **Scalar `DEF`s** — [`QueryBuilder::def_expression`].
//! - **Merging transformed sub-queries** — [`QueryBuilder::merge_as_cte`] and [`QueryBuilder::merge_as_main`].
//!
//! Compiler passes use the same tabular `DEF` shape as user code (`DefBody::Pipeline`).

use std::sync::Arc;

use crate::tree::ast::expression::Expression;
use crate::tree::ast::identifier::{Identifier, ParsedIdentifier};
use crate::tree::ast::node::Span;
use crate::tree::ast::pipeline::Pipeline;
use crate::tree::ast::query::{DefBody, DefStatement, Query, QueryKind, ValidQuery};
use crate::tree::typed_ast::pipeline::TypedPipeline;

use super::pipeline::PipelineBuilder;

/// Trait for types that can be converted to `Arc<Pipeline>` for use in query building.
pub trait IntoPipelineRc {
    fn into_pipeline_rc(self) -> Arc<Pipeline>;
}

impl IntoPipelineRc for Pipeline {
    fn into_pipeline_rc(self) -> Arc<Pipeline> {
        Arc::new(self)
    }
}

impl IntoPipelineRc for Arc<Pipeline> {
    fn into_pipeline_rc(self) -> Arc<Pipeline> {
        self
    }
}

impl IntoPipelineRc for PipelineBuilder {
    fn into_pipeline_rc(self) -> Arc<Pipeline> {
        Arc::new(self.build())
    }
}

impl IntoPipelineRc for Arc<TypedPipeline> {
    fn into_pipeline_rc(self) -> Arc<Pipeline> {
        self.ast.clone()
    }
}

impl IntoPipelineRc for &Arc<TypedPipeline> {
    fn into_pipeline_rc(self) -> Arc<Pipeline> {
        self.ast.clone()
    }
}

/// Builder for a query before the main pipeline is set.
#[derive(Clone)]
pub struct QueryBuilder {
    defs: Vec<DefStatement>,
}

impl QueryBuilder {
    pub fn new() -> Self {
        Self { defs: Vec::new() }
    }

    /// Set the main pipeline
    pub fn main(self, pipeline: impl IntoPipelineRc) -> QueryBuilderWithMain {
        QueryBuilderWithMain {
            defs: self.defs,
            main_pipeline: pipeline.into_pipeline_rc(),
        }
    }

    /// `DEF name = <pipeline> ;`
    pub fn def_pipeline(
        mut self,
        name: impl Into<Identifier>,
        pipeline: impl IntoPipelineRc,
    ) -> Self {
        self.defs.push(DefStatement {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            body: DefBody::Pipeline(pipeline.into_pipeline_rc()),
        });
        self
    }

    /// `DEF name = <expression> ;`
    pub fn def_expression(
        mut self,
        name: impl Into<Identifier>,
        expression: impl Into<Arc<Expression>>,
    ) -> Self {
        self.defs.push(DefStatement {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            body: DefBody::Expression(expression.into()),
        });
        self
    }

    /// Merge a [`Query`]'s `DEF`s into this builder, then add its main pipeline as a new tabular
    /// `DEF`.
    pub fn merge_as_cte(mut self, query: Query, name: impl Into<Identifier>) -> Self {
        if let QueryKind::Valid(valid) = query.kind {
            self.defs.extend(valid.defs);
            self.defs.push(DefStatement {
                span: Span::NONE,
                name: ParsedIdentifier::Valid(name.into()),
                body: DefBody::Pipeline(valid.main_pipeline),
            });
        }
        self
    }

    /// Merge a [`Query`]'s `DEF`s into this builder, using its main as this builder's main.
    pub fn merge_as_main(self, query: Query) -> Query {
        if let QueryKind::Valid(valid) = query.kind {
            let mut defs = self.defs;
            defs.extend(valid.defs);
            Query {
                span: Span::NONE,
                kind: ValidQuery {
                    span: Span::NONE,
                    defs,
                    main_pipeline: valid.main_pipeline,
                }
                .into(),
            }
        } else {
            // Error query - just return it as-is
            query
        }
    }
}

impl Default for QueryBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for Query (main pipeline has been set)
#[derive(Clone)]
pub struct QueryBuilderWithMain {
    defs: Vec<DefStatement>,
    main_pipeline: Arc<Pipeline>,
}

impl QueryBuilderWithMain {
    /// `DEF name = <pipeline> ;` (appended before the main body in the built [`ValidQuery`].
    pub fn def_pipeline(
        mut self,
        name: impl Into<Identifier>,
        pipeline: impl IntoPipelineRc,
    ) -> Self {
        self.defs.push(DefStatement {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            body: DefBody::Pipeline(pipeline.into_pipeline_rc()),
        });
        self
    }

    pub fn build(self) -> Query {
        Query {
            span: Span::NONE,
            kind: ValidQuery {
                span: Span::NONE,
                defs: self.defs,
                main_pipeline: self.main_pipeline,
            }
            .into(),
        }
    }
}

/// Create a query builder.
///
/// Use `.def_pipeline()` and `.def_expression()` to add DEFs, then `.main()` to set the main pipeline,
/// then `.build()` to create the Query.
pub fn query() -> QueryBuilder {
    QueryBuilder::new()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::ast::query::DefBody;
    use crate::tree::ast::query::QueryKind;
    use crate::tree::builder::pipeline;

    fn count_pipeline_defs(valid: &crate::tree::ast::query::ValidQuery) -> usize {
        valid
            .defs
            .iter()
            .filter(|d| matches!(d.body, DefBody::Pipeline(_)))
            .count()
    }

    #[test]
    fn test_query_builder_main_then_build() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("events")))
            .build();

        match &q.kind {
            QueryKind::Valid(valid) => {
                assert!(valid.defs.is_empty());
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_def_pipeline_then_main() {
        let q = query()
            .def_pipeline("cte1", pipeline().from(|f| f.table_reference("events")))
            .def_pipeline("cte2", pipeline().from(|f| f.table_reference("logs")))
            .main(pipeline().from(|f| f.table_reference("cte1")))
            .build();

        match &q.kind {
            QueryKind::Valid(valid) => {
                assert_eq!(count_pipeline_defs(valid), 2);
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_builder_main_then_def_pipeline() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("cte1")))
            .def_pipeline("cte1", pipeline().from(|f| f.table_reference("events")))
            .build();

        match &q.kind {
            QueryKind::Valid(valid) => {
                assert_eq!(count_pipeline_defs(valid), 1);
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_main_then_def_pipeline_shorthand() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("events")))
            .def_pipeline("cte1", pipeline().from(|f| f.table_reference("logs")))
            .build();

        match &q.kind {
            QueryKind::Valid(valid) => {
                assert_eq!(count_pipeline_defs(valid), 1);
            }
            _ => panic!("Expected valid query"),
        }
    }
}