hamelin_lib 0.21.8

//! Pipeline builder for constructing Pipeline AST nodes
//!
//! This module provides a fluent API for building pipelines of commands.

use std::sync::Arc;

use crate::tree::ast::clause::TableReference;
use crate::tree::ast::command::Command;
use crate::tree::ast::dataset_identifier::DatasetIdentifier;
use crate::tree::ast::identifier::{Identifier, SimpleIdentifier};
use crate::tree::ast::node::Span;
use crate::tree::ast::pipeline::Pipeline;

use super::command::*;
use super::IntoExpressionBuilder;

// Implement for all command builders
impl From<SetCommandBuilder> for Command {
    fn from(value: SetCommandBuilder) -> Self {
        value.build()
    }
}

impl From<WhereCommandBuilder> for Command {
    fn from(value: WhereCommandBuilder) -> Self {
        value.build()
    }
}

impl From<SelectCommandBuilder> for Command {
    fn from(value: SelectCommandBuilder) -> Self {
        value.build()
    }
}

impl From<DropCommandBuilder> for Command {
    fn from(value: DropCommandBuilder) -> Self {
        value.build()
    }
}

impl From<FromCommandBuilder> for Command {
    fn from(value: FromCommandBuilder) -> Self {
        value.build()
    }
}

impl From<UnionCommandBuilder> for Command {
    fn from(value: UnionCommandBuilder) -> Self {
        value.build()
    }
}

impl From<LimitCommandBuilder> for Command {
    fn from(value: LimitCommandBuilder) -> Self {
        value.build()
    }
}

impl From<WithinCommandBuilder> for Command {
    fn from(value: WithinCommandBuilder) -> Self {
        value.build()
    }
}

impl From<SortCommandBuilder> for Command {
    fn from(value: SortCommandBuilder) -> Self {
        value.build()
    }
}

impl From<ValidParseCommandBuilder> for Command {
    fn from(value: ValidParseCommandBuilder) -> Self {
        value.build()
    }
}

impl From<AggCommandBuilder> for Command {
    fn from(value: AggCommandBuilder) -> Self {
        value.build()
    }
}

impl From<WindowCommandBuilder> for Command {
    fn from(value: WindowCommandBuilder) -> Self {
        value.build()
    }
}

impl From<SuppressCommandBuilder> for Command {
    fn from(value: SuppressCommandBuilder) -> Self {
        value.build()
    }
}

impl From<AppendCommandBuilder> for Command {
    fn from(value: AppendCommandBuilder) -> Self {
        value.build()
    }
}

impl From<JoinCommandBuilderWithOn> for Command {
    fn from(value: JoinCommandBuilderWithOn) -> Self {
        value.build()
    }
}

impl From<LookupCommandBuilderWithOn> for Command {
    fn from(value: LookupCommandBuilderWithOn) -> Self {
        value.build()
    }
}

impl From<UnnestCommandBuilder> for Command {
    fn from(value: UnnestCommandBuilder) -> Self {
        value.build()
    }
}

impl From<RowsCommandBuilder> for Command {
    fn from(value: RowsCommandBuilder) -> Self {
        value.build()
    }
}

impl From<NestCommandBuilder> for Command {
    fn from(value: NestCommandBuilder) -> Self {
        value.build()
    }
}

impl From<MatchCommandBuilder> for Command {
    fn from(value: MatchCommandBuilder) -> Self {
        value.build()
    }
}

/// Builder for Pipeline
#[derive(Debug)]
pub struct PipelineBuilder {
    commands: Vec<Command>,
    span: Span,
}

impl PipelineBuilder {
    pub fn new() -> Self {
        Self {
            commands: Vec::new(),
            span: Span::NONE,
        }
    }

    /// Set the span for this pipeline
    pub fn at(mut self, span: Span) -> Self {
        self.span = span;
        self
    }

    /// Add a command to the pipeline
    pub fn command(mut self, cmd: impl Into<Command>) -> Self {
        self.commands.push(cmd.into());
        self
    }

    /// Add a SET command using a closure
    pub fn set_cmd<F>(self, f: F) -> Self
    where
        F: FnOnce(SetCommandBuilder) -> SetCommandBuilder,
    {
        self.command(f(set_command()))
    }

    /// Add a WHERE command
    pub fn where_cmd(self, condition: impl IntoExpressionBuilder) -> Self {
        self.command(where_command(condition))
    }

    /// Add a SELECT command using a closure
    pub fn select<F>(self, f: F) -> Self
    where
        F: FnOnce(SelectCommandBuilder) -> SelectCommandBuilder,
    {
        self.command(f(select_command()))
    }

    /// Add a DROP command using a closure
    pub fn drop<F>(self, f: F) -> Self
    where
        F: FnOnce(DropCommandBuilder) -> DropCommandBuilder,
    {
        self.command(f(drop_command()))
    }

    /// Add a FROM command using a closure
    pub fn from<F>(self, f: F) -> Self
    where
        F: FnOnce(FromCommandBuilder) -> FromCommandBuilder,
    {
        self.command(f(from_command()))
    }

    /// Add a UNION command using a closure
    pub fn union<F>(self, f: F) -> Self
    where
        F: FnOnce(UnionCommandBuilder) -> UnionCommandBuilder,
    {
        self.command(f(union_command()))
    }

    /// Add a LIMIT command
    pub fn limit(self, count: impl IntoExpressionBuilder) -> Self {
        self.command(limit_command(count))
    }

    /// Add a WITHIN command
    pub fn within(self, duration: impl IntoExpressionBuilder) -> Self {
        self.command(within_command(duration))
    }

    /// Add a SORT command using a closure
    pub fn sort<F>(self, f: F) -> Self
    where
        F: FnOnce(SortCommandBuilder) -> SortCommandBuilder,
    {
        self.command(f(sort_command()))
    }

    /// Add a PARSE command using a closure
    ///
    /// Examples:
    /// - `pipeline.parse(|p| p.pattern("prefix-*-suffix").identifier("field1"))`
    /// - `pipeline.parse(|p| p.pattern("*:*").identifier("key").identifier("value").nodrop(true))`
    pub fn parse<F>(self, f: F) -> Self
    where
        F: FnOnce(ParseCommandBuilder) -> ValidParseCommandBuilder,
    {
        self.command(f(parse_command()))
    }

    /// Add an AGG command using a closure
    pub fn agg<F>(self, f: F) -> Self
    where
        F: FnOnce(AggCommandBuilder) -> AggCommandBuilder,
    {
        self.command(f(agg_command()))
    }

    /// Add a WINDOW command using a closure
    pub fn window<F>(self, f: F) -> Self
    where
        F: FnOnce(WindowCommandBuilder) -> WindowCommandBuilder,
    {
        self.command(f(window_command()))
    }

    /// Add a SUPPRESS command: `interval` is required (e.g. `hours(1)`); use the closure for `BY` keys.
    pub fn suppress<F>(self, interval: impl IntoExpressionBuilder, f: F) -> Self
    where
        F: FnOnce(SuppressCommandBuilder) -> SuppressCommandBuilder,
    {
        self.command(f(suppress_command(interval)))
    }

    /// Add an APPEND command - takes table as parameter
    pub fn append(
        self,
        table: impl Into<TableReference>,
        f: impl FnOnce(AppendCommandBuilder) -> AppendCommandBuilder,
    ) -> Self {
        self.command(f(append_command(table)))
    }

    /// Add a JOIN command with a space-qualified dataset reference
    pub fn join_dataset(
        self,
        dataset: DatasetIdentifier,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(join_dataset_reference(dataset).on(condition))
    }

    /// Add a JOIN command with aliased space-qualified dataset
    pub fn join_dataset_aliased(
        self,
        alias: impl Into<SimpleIdentifier>,
        dataset: DatasetIdentifier,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(JoinCommandBuilder::table_alias_dataset(alias, dataset).on(condition))
    }

    /// Add a LOOKUP command with a space-qualified dataset reference
    pub fn lookup_dataset(
        self,
        dataset: DatasetIdentifier,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(lookup_dataset_reference(dataset).on(condition))
    }

    /// Add a JOIN command with table reference
    pub fn join(
        self,
        table: impl Into<TableReference>,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(join_table_reference(table).on(condition))
    }

    /// Add a JOIN command with aliased table
    pub fn join_aliased(
        self,
        alias: impl Into<SimpleIdentifier>,
        table: impl Into<TableReference>,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(join_table_alias(alias, table).on(condition))
    }

    /// Add a LOOKUP command with table reference
    pub fn lookup(
        self,
        table: impl Into<TableReference>,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(lookup_table_reference(table).on(condition))
    }

    /// Add a LOOKUP command with aliased table
    pub fn lookup_aliased(
        self,
        alias: impl Into<SimpleIdentifier>,
        table: impl Into<TableReference>,
        condition: impl IntoExpressionBuilder,
    ) -> Self {
        self.command(lookup_table_alias(alias, table).on(condition))
    }

    /// Add an EXPLODE command using a closure
    ///
    /// Examples:
    /// - `pipeline.explode(|e| e.expression(field_ref("items")))` - explode expression
    /// - `pipeline.explode(|e| e.assignment("item", field_ref("items")))` - explode with assignment
    pub fn explode<F>(self, f: F) -> Self
    where
        F: FnOnce(ExplodeCommandBuilder<NoClause>) -> ExplodeCommandBuilder<WithClause>,
    {
        self.command(f(explode_command()))
    }

    /// Add an UNNEST command
    pub fn unnest(self, expression: impl IntoExpressionBuilder) -> Self {
        self.command(unnest_command(expression))
    }

    /// Add a ROWS command
    pub fn rows(self, expression: impl IntoExpressionBuilder) -> Self {
        self.command(rows_command(expression))
    }

    /// Add a NEST command
    pub fn nest(self, identifier: impl Into<Identifier>) -> Self {
        self.command(nest_command(identifier))
    }

    /// Add a MATCH command using a closure
    pub fn match_cmd<F>(self, f: F) -> Self
    where
        F: FnOnce(MatchCommandBuilder) -> MatchCommandBuilder,
    {
        self.command(f(match_command()))
    }

    pub fn build(&self) -> Pipeline {
        Pipeline {
            span: self.span,
            commands: self.commands.iter().cloned().map(Arc::new).collect(),
        }
    }
}

impl From<Arc<Command>> for Command {
    fn from(cmd: Arc<Command>) -> Self {
        (*cmd).clone()
    }
}

impl From<PipelineBuilder> for Pipeline {
    fn from(value: PipelineBuilder) -> Self {
        value.build()
    }
}

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

/// Create a pipeline builder
pub fn pipeline() -> PipelineBuilder {
    PipelineBuilder::new()
}

/// Create a pipeline from an expression (wraps it in a SELECT command)
pub fn pipeline_from_expression(expr: impl IntoExpressionBuilder) -> Pipeline {
    Pipeline::from(expr.into_expression_builder().build())
}