4371276901

Build Type

push

github

Committed by GitHub

Commit Message

fix: Add sources names used in the query (#1190)

Run Details

28 of 28 new or added lines in 3 files covered. (100.0%)

27837 of 39159 relevant lines covered (71.09%)

73963.24 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

62.52

/dozer-sql/src/pipeline/builder.rs

use crate::pipeline::aggregation::factory::AggregationProcessorFactory;
use crate::pipeline::builder::PipelineError::InvalidQuery;
use crate::pipeline::errors::PipelineError;
use crate::pipeline::expression::builder::{ExpressionBuilder, NameOrAlias};
use crate::pipeline::product::set_factory::SetProcessorFactory;
use crate::pipeline::selection::factory::SelectionProcessorFactory;
use crate::pipeline::window::builder::relation_is_a_window;
use dozer_core::app::AppPipeline;
use dozer_core::app::PipelineEntryPoint;
use dozer_core::appsource::AppSourceId;
use dozer_core::node::PortHandle;
use dozer_core::DEFAULT_PORT_HANDLE;
use sqlparser::ast::{Join, SetOperator, SetQuantifier, TableFactor, TableWithJoins};

use sqlparser::{
    ast::{Query, Select, SetExpr, Statement},
    dialect::AnsiDialect,
    parser::Parser,
};
use std::collections::HashMap;
use std::sync::Arc;

use super::errors::UnsupportedSqlError;
use super::product::factory::FromProcessorFactory;
use super::window::factory::WindowProcessorFactory;

#[derive(Debug, Clone, Default)]
pub struct SchemaSQLContext {}

#[derive(Debug, Clone)]
pub struct OutputNodeInfo {
    // Name to connect in dag
    pub node: String,
    // Port to connect in dag
    pub port: PortHandle,
    // If this table is originally from a source or created in transforms
    pub is_derived: bool,
    // TODO add:indexes to the tables
}

pub struct TableInfo {
    pub name: NameOrAlias,
    pub override_name: Option<String>,
    pub is_derived: bool,
}
/// The struct contains some contexts during query to pipeline.
#[derive(Debug, Clone, Default)]
pub struct QueryContext {
    // Internal tables map, used to store the tables that are created by the queries
    pub pipeline_map: HashMap<(usize, String), OutputNodeInfo>,

    // Output tables map that are marked with "INTO" used to store the tables, these can be exposed to sinks.
    pub output_tables_map: HashMap<String, OutputNodeInfo>,

    // Used Sources
    pub used_sources: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct IndexedTableWithJoins {
    pub relation: (NameOrAlias, TableFactor),
    pub joins: Vec<(NameOrAlias, Join)>,
}
pub fn statement_to_pipeline(
    sql: &str,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    override_name: Option<String>,
) -> Result<QueryContext, PipelineError> {
    let dialect = AnsiDialect {};
    let mut ctx = QueryContext::default();

    let ast = Parser::parse_sql(&dialect, sql).unwrap();
    let query_name = NameOrAlias(format!("query_{}", uuid::Uuid::new_v4()), None);

    for (idx, statement) in ast.iter().enumerate() {
        match statement {
            Statement::Query(query) => {
                query_to_pipeline(
                    &TableInfo {
                        name: query_name.clone(),
                        is_derived: false,
                        override_name: override_name.clone(),
                    },
                    query,
                    pipeline,
                    &mut ctx,
                    false,
                    idx,
                )?;
            }
            s => {
                return Err(PipelineError::UnsupportedSqlError(
                    UnsupportedSqlError::GenericError(s.to_string()),
                ))
            }
        }
    }

    Ok(ctx)
}

fn query_to_pipeline(
    table_info: &TableInfo,
    query: &Query,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    query_ctx: &mut QueryContext,
    stateful: bool,
    pipeline_idx: usize,
) -> Result<(), PipelineError> {
    // return error if there is unsupported syntax
    if !query.order_by.is_empty() {
        return Err(PipelineError::UnsupportedSqlError(
            UnsupportedSqlError::OrderByError,
        ));
    }

    if query.limit.is_some() || query.offset.is_some() {
        return Err(PipelineError::UnsupportedSqlError(
            UnsupportedSqlError::LimitOffsetError,
        ));
    }

    // Attach the first pipeline if there is with clause
    if let Some(with) = &query.with {
        if with.recursive {
            return Err(PipelineError::UnsupportedSqlError(
                UnsupportedSqlError::Recursive,
            ));
        }

        for table in &with.cte_tables {
            if table.from.is_some() {
                return Err(PipelineError::UnsupportedSqlError(
                    UnsupportedSqlError::CteFromError,
                ));
            }
            let table_name = table.alias.name.to_string();
            if query_ctx
                .pipeline_map
                .contains_key(&(pipeline_idx, table_name.clone()))
            {
                return Err(InvalidQuery(format!(
                    "WITH query name {table_name:?} specified more than once"
                )));
            }
            query_to_pipeline(
                &TableInfo {
                    name: NameOrAlias(table_name.clone(), Some(table_name)),
                    is_derived: true,
                    override_name: None,
                },
                &table.query,
                pipeline,
                query_ctx,
                true,
                pipeline_idx,
            )?;
        }
    };

    match *query.body.clone() {
        SetExpr::Select(select) => {
            select_to_pipeline(
                table_info,
                *select,
                pipeline,
                query_ctx,
                stateful,
                pipeline_idx,
            )?;
        }
        SetExpr::Query(query) => {
            let query_name = format!("subquery_{}", uuid::Uuid::new_v4());
            let mut ctx = QueryContext::default();
            query_to_pipeline(
                &TableInfo {
                    name: NameOrAlias(query_name, None),
                    is_derived: true,
                    override_name: None,
                },
                &query,
                pipeline,
                &mut ctx,
                stateful,
                pipeline_idx,
            )?
        }
        SetExpr::SetOperation {
            op,
            set_quantifier,
            left,
            right,
        } => match op {
            SetOperator::Union => {
                set_to_pipeline(
                    table_info,
                    left,
                    right,
                    set_quantifier,
                    pipeline,
                    query_ctx,
                    stateful,
                    pipeline_idx,
                )?;
            }
            _ => return Err(PipelineError::InvalidOperator(op.to_string())),
        },
        _ => {
            return Err(PipelineError::UnsupportedSqlError(
                UnsupportedSqlError::GenericError("Unsupported query body structure".to_string()),
            ))
        }
    };
    Ok(())
}

fn select_to_pipeline(
    table_info: &TableInfo,
    select: Select,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    query_ctx: &mut QueryContext,
    stateful: bool,
    pipeline_idx: usize,
) -> Result<String, PipelineError> {
    // FROM clause
    if select.from.len() != 1 {
        return Err(PipelineError::UnsupportedSqlError(
            UnsupportedSqlError::FromCommaSyntax,
        ));
    }

    let input_tables = get_input_tables(&select.from[0], pipeline, query_ctx, pipeline_idx)?;

    let (input_nodes, output_node, mut used_sources) = add_from_to_pipeline(
        pipeline,
        &input_tables,
        &mut query_ctx.pipeline_map,
        pipeline_idx,
    )?;

    query_ctx.used_sources.append(&mut used_sources);

    let gen_agg_name = format!("agg_{}", uuid::Uuid::new_v4());
    let gen_selection_name = format!("select_{}", uuid::Uuid::new_v4());
    let (gen_product_name, product_output_port) = output_node;

    for (source_name, processor_name, processor_port) in input_nodes.iter() {
        if let Some(table_info) = query_ctx
            .pipeline_map
            .get(&(pipeline_idx, source_name.clone()))
        {
            pipeline.connect_nodes(
                &table_info.node,
                Some(table_info.port),
                processor_name,
                Some(*processor_port as PortHandle),
                true,
            )?;
            // If not present in pipeline_map, insert into used_sources as this is coming from source
        } else {
            query_ctx.used_sources.push(source_name.clone());
        }
    }

    let aggregation = AggregationProcessorFactory::new(select.clone(), stateful);

    pipeline.add_processor(Arc::new(aggregation), &gen_agg_name, vec![]);

    // Where clause
    if let Some(selection) = select.selection {
        let selection = SelectionProcessorFactory::new(selection);

        pipeline.add_processor(Arc::new(selection), &gen_selection_name, vec![]);

        pipeline.connect_nodes(
            &gen_product_name,
            Some(product_output_port),
            &gen_selection_name,
            Some(DEFAULT_PORT_HANDLE),
            true,
        )?;

        pipeline.connect_nodes(
            &gen_selection_name,
            Some(DEFAULT_PORT_HANDLE),
            &gen_agg_name,
            Some(DEFAULT_PORT_HANDLE),
            true,
        )?;
    } else {
        pipeline.connect_nodes(
            &gen_product_name,
            Some(product_output_port),
            &gen_agg_name,
            Some(DEFAULT_PORT_HANDLE),
            true,
        )?;
    }

    query_ctx.pipeline_map.insert(
        (pipeline_idx, table_info.name.0.to_string()),
        OutputNodeInfo {
            node: gen_agg_name.clone(),
            port: DEFAULT_PORT_HANDLE,
            is_derived: table_info.is_derived,
        },
    );

    let output_table_name = if let Some(into) = select.into {
        Some(into.name.to_string())
    } else {
        table_info.override_name.clone()
    };
    if let Some(table_name) = output_table_name {
        query_ctx.output_tables_map.insert(
            table_name,
            OutputNodeInfo {
                node: gen_agg_name.clone(),
                port: DEFAULT_PORT_HANDLE,
                is_derived: false,
            },
        );
    }

    Ok(gen_agg_name)
}

#[allow(clippy::type_complexity)]
fn add_from_to_pipeline(
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    input_tables: &IndexedTableWithJoins,
    pipeline_map: &mut HashMap<(usize, String), OutputNodeInfo>,
    pipeline_idx: usize,
) -> Result<
    (
        Vec<(String, String, PortHandle)>,
        (String, PortHandle),
        Vec<String>,
    ),
    PipelineError,
> {
    // the sources names that are used in this pipeline
    let mut used_sources = vec![];

    let mut product_entry_points = vec![];
    let mut input_nodes = vec![];

    let (relation_name_or_alias, relation) = input_tables.relation.clone();

    let product_processor_name = format!("product_{}", uuid::Uuid::new_v4());
    let product_processor = FromProcessorFactory::new(input_tables.clone());

    if relation_is_a_window(&relation).map_err(PipelineError::WindowError)? {
        let window_processor = WindowProcessorFactory::new(relation);
        let window_processor_name = format!("window_{}", uuid::Uuid::new_v4());
        let window_source_name = window_processor.get_source_name()?;
        let mut window_entry_points = vec![];

        if is_an_entry_point(&window_source_name, pipeline_map, pipeline_idx) {
            let entry_point = PipelineEntryPoint::new(
                AppSourceId::new(window_source_name.clone(), None),
                DEFAULT_PORT_HANDLE as PortHandle,
            );

            window_entry_points.push(entry_point);
            used_sources.push(window_source_name);
        } else {
            input_nodes.push((
                window_source_name,
                window_processor_name.clone(),
                DEFAULT_PORT_HANDLE as PortHandle,
            ));
        }

        pipeline.add_processor(
            Arc::new(window_processor),
            &window_processor_name,
            window_entry_points,
        );

        pipeline.connect_nodes(
            &window_processor_name,
            Some(DEFAULT_PORT_HANDLE as PortHandle),
            &product_processor_name,
            Some(0 as PortHandle),
            true,
        )?;
    } else {
        let product_input_name = relation_name_or_alias.0;

        if is_an_entry_point(&product_input_name, pipeline_map, pipeline_idx) {
            let entry_point = PipelineEntryPoint::new(
                AppSourceId::new(product_input_name.clone(), None),
                0 as PortHandle,
            );

            product_entry_points.push(entry_point);
            used_sources.push(product_input_name);
        } else {
            input_nodes.push((
                product_input_name,
                product_processor_name.clone(),
                0 as PortHandle,
            ));
        }
    }

    for (index, (join_relation_alias, join)) in input_tables.joins.iter().enumerate() {
        let (relation_name_or_alias, relation) =
            (join_relation_alias.clone(), join.relation.clone());

        if relation_is_a_window(&relation).map_err(PipelineError::WindowError)? {
            let window_processor = WindowProcessorFactory::new(relation.clone());
            let window_processor_name = format!("window_{}", uuid::Uuid::new_v4());
            let window_input_name = window_processor.get_source_name()?;
            let mut window_entry_points = vec![];

            if is_an_entry_point(&window_input_name, pipeline_map, pipeline_idx) {
                let entry_point = PipelineEntryPoint::new(
                    AppSourceId::new(window_input_name.clone(), None),
                    DEFAULT_PORT_HANDLE as PortHandle,
                );

                window_entry_points.push(entry_point.clone());
                used_sources.push(window_input_name);
            } else {
                input_nodes.push((
                    window_input_name,
                    window_processor_name.clone(),
                    DEFAULT_PORT_HANDLE as PortHandle,
                ));
            }

            pipeline.add_processor(
                Arc::new(window_processor),
                &window_processor_name,
                window_entry_points,
            );

            pipeline.connect_nodes(
                &window_processor_name,
                Some(DEFAULT_PORT_HANDLE as PortHandle),
                &product_processor_name,
                Some((index + 1) as PortHandle),
                true,
            )?;
        } else {
            let product_input_name = relation_name_or_alias.0;

            if is_an_entry_point(&product_input_name, pipeline_map, pipeline_idx) {
                let entry_point = PipelineEntryPoint::new(
                    AppSourceId::new(product_input_name.clone(), None),
                    (index + 1) as PortHandle,
                );

                product_entry_points.push(entry_point.clone());
                used_sources.push(product_input_name);
            } else {
                input_nodes.push((
                    product_input_name,
                    product_processor_name.clone(),
                    (index + 1) as PortHandle,
                ));
            }
        }
    }

    pipeline.add_processor(
        Arc::new(product_processor),
        &product_processor_name,
        product_entry_points,
    );

    Ok((
        input_nodes,
        (product_processor_name, DEFAULT_PORT_HANDLE as PortHandle),
        used_sources,
    ))
}

#[allow(clippy::too_many_arguments)]
fn set_to_pipeline(
    table_info: &TableInfo,
    left_select: Box<SetExpr>,
    right_select: Box<SetExpr>,
    set_quantifier: SetQuantifier,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    query_ctx: &mut QueryContext,
    stateful: bool,
    pipeline_idx: usize,
) -> Result<String, PipelineError> {
    let gen_left_set_name = format!("set_left_{}", uuid::Uuid::new_v4());
    let left_table_info = TableInfo {
        name: NameOrAlias(gen_left_set_name.clone(), None),
        override_name: None,
        is_derived: false,
    };
    let gen_right_set_name = format!("set_right_{}", uuid::Uuid::new_v4());
    let right_table_info = TableInfo {
        name: NameOrAlias(gen_right_set_name.clone(), None),
        override_name: None,
        is_derived: false,
    };

    let _left_pipeline_name = match *left_select {
        SetExpr::Select(select) => select_to_pipeline(
            &left_table_info,
            *select,
            pipeline,
            query_ctx,
            stateful,
            pipeline_idx,
        )?,
        SetExpr::SetOperation {
            op: _,
            set_quantifier,
            left,
            right,
        } => set_to_pipeline(
            &left_table_info,
            left,
            right,
            set_quantifier,
            pipeline,
            query_ctx,
            stateful,
            pipeline_idx,
        )?,
        _ => {
            return Err(PipelineError::InvalidQuery(
                "Invalid UNION left Query".to_string(),
            ))
        }
    };

    let _right_pipeline_name = match *right_select {
        SetExpr::Select(select) => select_to_pipeline(
            &right_table_info,
            *select,
            pipeline,
            query_ctx,
            stateful,
            pipeline_idx,
        )?,
        SetExpr::SetOperation {
            op: _,
            set_quantifier,
            left,
            right,
        } => set_to_pipeline(
            &right_table_info,
            left,
            right,
            set_quantifier,
            pipeline,
            query_ctx,
            stateful,
            pipeline_idx,
        )?,
        _ => {
            return Err(PipelineError::InvalidQuery(
                "Invalid UNION right Query".to_string(),
            ))
        }
    };

    let left_pipeline_output_node = match query_ctx
        .pipeline_map
        .get(&(pipeline_idx, gen_left_set_name))
    {
        Some(pipeline) => pipeline,
        None => {
            return Err(PipelineError::InvalidQuery(
                "Invalid UNION left Query".to_string(),
            ))
        }
    };

    let right_pipeline_output_node = match query_ctx
        .pipeline_map
        .get(&(pipeline_idx, gen_right_set_name))
    {
        Some(pipeline) => pipeline,
        None => {
            return Err(PipelineError::InvalidQuery(
                "Invalid UNION Right Query".to_string(),
            ))
        }
    };

    let set_proc_fac = SetProcessorFactory::new(set_quantifier);

    let mut gen_set_name = format!("set_{}", uuid::Uuid::new_v4());

    if table_info.override_name.is_some() {
        gen_set_name = table_info.override_name.to_owned().unwrap();
    }

    pipeline.add_processor(Arc::new(set_proc_fac), &gen_set_name, vec![]);

    pipeline.connect_nodes(
        &left_pipeline_output_node.node,
        Some(left_pipeline_output_node.port),
        &gen_set_name,
        Some(0 as PortHandle),
        true,
    )?;

    pipeline.connect_nodes(
        &right_pipeline_output_node.node,
        Some(right_pipeline_output_node.port),
        &gen_set_name,
        Some(1 as PortHandle),
        true,
    )?;

    for (_, table_name) in query_ctx.pipeline_map.keys() {
        query_ctx.output_tables_map.remove_entry(table_name);
    }

    query_ctx.pipeline_map.insert(
        (pipeline_idx, table_info.name.0.to_string()),
        OutputNodeInfo {
            node: gen_set_name.clone(),
            port: DEFAULT_PORT_HANDLE,
            is_derived: table_info.is_derived,
        },
    );

    Ok(gen_set_name)
}

/// Returns a vector of input port handles and relative table name
///
/// # Errors
///
/// This function will return an error if it's not possible to get an input name.
pub fn get_input_tables(
    from: &TableWithJoins,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    query_ctx: &mut QueryContext,
    pipeline_idx: usize,
) -> Result<IndexedTableWithJoins, PipelineError> {
    let name = get_from_source(&from.relation, pipeline, query_ctx, pipeline_idx)?;
    let mut joins = vec![];

    for join in from.joins.iter() {
        let input_name = get_from_source(&join.relation, pipeline, query_ctx, pipeline_idx)?;
        joins.push((input_name.clone(), join.clone()));
    }

    Ok(IndexedTableWithJoins {
        relation: (name, from.relation.clone()),
        joins,
    })
}

pub fn get_input_names(input_tables: &IndexedTableWithJoins) -> Vec<NameOrAlias> {
    let mut input_names = vec![];
    input_names.push(input_tables.relation.0.clone());

    for join in &input_tables.joins {
        input_names.push(join.0.clone());
    }
    input_names
}

pub fn get_entry_points(
    input_tables: &IndexedTableWithJoins,
    pipeline_map: &mut HashMap<(usize, String), OutputNodeInfo>,
    pipeline_idx: usize,
) -> Result<Vec<PipelineEntryPoint>, PipelineError> {
    let mut endpoints = vec![];

    let input_names = get_input_names(input_tables);

    for (input_port, table) in input_names.iter().enumerate() {
        let name = table.0.clone();
        if !pipeline_map.contains_key(&(pipeline_idx, name.clone())) {
            endpoints.push(PipelineEntryPoint::new(
                AppSourceId::new(name, None),
                input_port as PortHandle,
            ));
        }
    }

    Ok(endpoints)
}

pub fn is_an_entry_point(
    name: &str,
    pipeline_map: &mut HashMap<(usize, String), OutputNodeInfo>,
    pipeline_idx: usize,
) -> bool {
    if !pipeline_map.contains_key(&(pipeline_idx, name.to_owned())) {
        return true;
    }
    false
}

pub fn get_from_source(
    relation: &TableFactor,
    pipeline: &mut AppPipeline<SchemaSQLContext>,
    query_ctx: &mut QueryContext,
    pipeline_idx: usize,
) -> Result<NameOrAlias, PipelineError> {
    match relation {
        TableFactor::Table { name, alias, .. } => {
            let input_name = name
                .0
                .iter()
                .map(ExpressionBuilder::normalize_ident)
                .collect::<Vec<String>>()
                .join(".");
            let alias_name = alias
                .as_ref()
                .map(|a| ExpressionBuilder::fullname_from_ident(&[a.name.clone()]));

            Ok(NameOrAlias(input_name, alias_name))
        }
        TableFactor::Derived {
            lateral: _,
            subquery,
            alias,
        } => {
            let name = format!("derived_{}", uuid::Uuid::new_v4());
            let alias_name = alias.as_ref().map(|alias_ident| {
                ExpressionBuilder::fullname_from_ident(&[alias_ident.name.clone()])
            });

            let name_or = NameOrAlias(name, alias_name);
            query_to_pipeline(
                &TableInfo {
                    name: name_or.clone(),
                    is_derived: true,
                    override_name: None,
                },
                subquery,
                pipeline,
                query_ctx,
                false,
                pipeline_idx,
            )?;

            Ok(name_or)
        }
        _ => Err(PipelineError::UnsupportedSqlError(
            UnsupportedSqlError::JoinTable,
        )),
    }
}

#[cfg(test)]
mod tests {
    use dozer_core::app::AppPipeline;

    use super::statement_to_pipeline;

    #[test]
    fn parse_sql_pipeline() {
        let sql = r#"
                SELECT
                    a.name as "Genre",
                    SUM(amount) as "Gross Revenue(in $)"
                INTO gross_revenue_stats
                FROM
                (
                    SELECT
                        c.name,
                        f.title,
                        p.amount
                    FROM film f
                    LEFT JOIN film_category fc
                        ON fc.film_id = f.film_id
                    LEFT JOIN category c
                        ON fc.category_id = c.category_id
                    LEFT JOIN inventory i
                        ON i.film_id = f.film_id
                    LEFT JOIN rental r
                        ON r.inventory_id = i.inventory_id
                    LEFT JOIN payment p
                        ON p.rental_id = r.rental_id
                    WHERE p.amount IS NOT NULL
                ) a
                GROUP BY name;

                SELECT
                f.name, f.title, p.amount
                INTO film_amounts
                FROM film f
                LEFT JOIN film_category fc;

                WITH tbl as (select id from a)
                select id
                into cte_table
                from tbl;

                WITH tbl as (select id from  a),
                tbl2 as (select id from tbl)
                select id
                into nested_cte_table
                from tbl2;

                WITH cte_table1 as (select id_dt1 from (select id_t1 from table_1) as derived_table_1),
                cte_table2 as (select id_ct1 from cte_table1)
                select id_ct2
                into nested_derived_table
                from cte_table2;

                with tbl as (select id, ticker from stocks)
                select tbl.id
                into nested_stocks_table
                from  stocks join tbl on tbl.id = stocks.id;
            "#;

        let context = statement_to_pipeline(sql, &mut AppPipeline::new(), None).unwrap();

        // Should create as many output tables as into statements
        let mut output_keys = context.output_tables_map.keys().collect::<Vec<_>>();
        output_keys.sort();
        let mut expected_keys = vec![
            "gross_revenue_stats",
            "film_amounts",
            "cte_table",
            "nested_cte_table",
            "nested_derived_table",
            "nested_stocks_table",
        ];
        expected_keys.sort();
        assert_eq!(output_keys, expected_keys);
    }
}

getdozer / dozer / 4371276901

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous