hamelin_lib 0.7.13

//! Query builder for constructing Query AST nodes
//!
//! This module provides a fluent API for building queries with WITH clauses.
//!
//! The builder uses separate types to ensure the main pipeline is set exactly once:
//! - `QueryBuilder`: Can add WITH clauses, must call `.main()` before `.build()`
//! - `QueryBuilderWithMain`: Can add WITH clauses, can call `.build()`

use std::sync::Arc;

use crate::tree::ast::identifier::{Identifier, ParsedIdentifier};
use crate::tree::ast::node::Span;
use crate::tree::ast::pipeline::Pipeline;
use crate::tree::ast::query::{Query, QueryKind, ValidQuery, WithClause};
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 Query (main pipeline not yet set)
#[derive(Clone)]
pub struct QueryBuilder {
    with_clauses: Vec<WithClause>,
}

impl QueryBuilder {
    /// Create a new query builder without a main pipeline
    pub fn new() -> Self {
        Self {
            with_clauses: Vec::new(),
        }
    }

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

    /// Add a WITH clause
    pub fn with(mut self, name: impl Into<Identifier>, pipeline: impl IntoPipelineRc) -> Self {
        self.with_clauses.push(WithClause {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            pipeline: pipeline.into_pipeline_rc(),
        });
        self
    }

    /// Merge a Query's CTEs into this builder, then add its main as a new WITH clause.
    ///
    /// Use this when transforming existing WITH clauses that may have generated
    /// additional CTEs during transformation.
    pub fn merge_as_cte(mut self, query: Query, name: impl Into<Identifier>) -> Self {
        if let QueryKind::Valid(valid) = query.kind {
            // Add all CTEs from the query
            self.with_clauses.extend(valid.with_clauses);
            // Add the query's main pipeline as a new CTE with the given name
            self.with_clauses.push(WithClause {
                span: Span::NONE,
                name: ParsedIdentifier::Valid(name.into()),
                pipeline: valid.main_pipeline,
            });
        }
        self
    }

    /// Merge a Query's CTEs into this builder, using its main as this builder's main.
    ///
    /// Use this for the final main pipeline transformation.
    pub fn merge_as_main(mut self, query: Query) -> Query {
        if let QueryKind::Valid(valid) = query.kind {
            // Add all CTEs from the query
            self.with_clauses.extend(valid.with_clauses);
            // Use the query's main as our main
            Query {
                span: Span::NONE,
                kind: ValidQuery {
                    span: Span::NONE,
                    with_clauses: self.with_clauses,
                    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 {
    with_clauses: Vec<WithClause>,
    main_pipeline: Arc<Pipeline>,
}

impl QueryBuilderWithMain {
    /// Add a WITH clause
    pub fn with(mut self, name: impl Into<Identifier>, pipeline: impl IntoPipelineRc) -> Self {
        self.with_clauses.push(WithClause {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            pipeline: pipeline.into_pipeline_rc(),
        });
        self
    }

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

/// Create a query builder.
///
/// Use `.with()` to add CTEs, 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::builder::pipeline;

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

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

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

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

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

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

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

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