hamelin_lib/

lib.rs

extern crate core;

use std::cell::RefCell;
use std::fmt::{Display, Formatter};
use std::rc::Rc;
use std::sync::Arc;

use antlr::hamelinparser::{
    CommandContextAll, CommandEOFContextAttrs, ExpressionContextAll, ExpressionEOFContextAttrs,
    IdentifierContextAll, IdentifierEOFContextAll, IdentifierEOFContextAttrs,
    SimpleIdentifierContextAll, SimpleIdentifierEOFContextAttrs,
};
use antlr_rust::errors::ANTLRError;
use anyhow::anyhow;
use ast::command::within;
use ast::err::ContextualTranslationErrors;
use ast::ExpressionTranslationContext;
use ast::QueryTranslationContext;
use chrono::Utc;
use func::def::FunctionTranslationContext;
use func::utils::within_range;
use serde::Serialize;
use sql::expression::identifier::SimpleIdentifier;
use sql::expression::literal::{ColumnReference, TimestampLiteral};
use sql::expression::SQLExpression;
use sql::query::SQLQuery;
use sql::types::SQLTimestampTzType;
use translation::range_builder::RangeBuilder;
use translation::{ContextualResult, PendingQuery, Translation};
use tsify_next::Tsify;
use types::TIMESTAMP;

use crate::antlr::hamelinparser::{
    ExpressionEOFContextAll, HamelintypeContextAll, PipelineEOFContextAll, QueryContextAll,
    QueryEOFContextAll, QueryEOFContextAttrs,
};
use crate::ast::err::{Context, TranslationError, TranslationErrors};
use crate::ast::expression::HamelinExpression;
use crate::ast::query::HamelinQuery;
use crate::catalog::CatalogProvider;
use crate::func::registry::FunctionRegistry;
use crate::parser::{make_hamelin_parser_from_input, HamelinStringParser};
use crate::provider::EnvironmentProvider;
use crate::sql::statement::Statement;
use crate::translation::env::Environment;
use crate::translation::{
    DMLTranslation, ExpressionTranslation, QueryTranslation, StatementTranslation,
};

pub mod antlr;
pub mod ast;
pub mod catalog;
pub mod eval;
pub mod func;
pub mod incremental;
pub mod interner;
pub mod parser;
pub mod provider;
pub mod reverse_eval;
pub mod sql;
pub mod translation;
pub mod tree;
pub mod types;
mod write_utils;

pub const VERSION: &str = env!("CARGO_PKG_VERSION");

pub fn execute_resillient_parse<T>(
    input: String,
    parse_fn: fn(HamelinStringParser) -> Result<T, ANTLRError>,
) -> anyhow::Result<(T, TranslationErrors)> {
    let (node, errors) = {
        let (parser, errors) = make_hamelin_parser_from_input(input);
        parse_fn(parser)
            .map_err(|e| anyhow!(e.to_string()))
            .map(|node| (node, errors))?
    };
    let unwrapped = Rc::try_unwrap(errors)
        .expect("could not unwrap Rc<TranslationErrors> after parsing complete.")
        .into_inner();
    Ok((node, unwrapped))
}

pub fn execute_parse<T>(
    input: String,
    parse_fn: fn(HamelinStringParser) -> Result<T, ANTLRError>,
) -> Result<T, TranslationErrors> {
    let len = input.len();
    let (node, errors) = {
        let (parser, errors) = make_hamelin_parser_from_input(input);
        parse_fn(parser)
            .map_err(|e| {
                TranslationError::new(Context::new(0..=len - 1, "failed to load grammars"))
                    .with_source_boxed(e.to_string().into())
                    .single()
            })
            .map(|node| (node, errors))?
    };
    let unwrapped = Rc::try_unwrap(errors)
        .expect("could not unwrap Rc<TranslationErrors> after parsing complete.")
        .into_inner();
    unwrapped.or_ok(node)
}

pub fn parse_command(command: String) -> Result<Rc<CommandContextAll<'static>>, TranslationErrors> {
    execute_parse(command, |mut parser| parser.commandEOF())
        .map(|ctx| ctx.command().expect("required"))
}

pub fn parse_query(statement: String) -> Result<Rc<QueryContextAll<'static>>, TranslationErrors> {
    execute_parse(statement, |mut parser| parser.queryEOF())
        .map(|ctx| ctx.query().expect("required"))
}

pub fn resilient_parse_query(
    statement: String,
) -> anyhow::Result<(Rc<QueryEOFContextAll<'static>>, TranslationErrors)> {
    execute_resillient_parse(statement, |mut parser| parser.queryEOF())
}

pub fn parse_expression(
    expression: String,
) -> Result<Rc<ExpressionContextAll<'static>>, TranslationErrors> {
    execute_parse(expression, |mut parser| parser.expressionEOF())
        .map(|ctx| ctx.expression().expect("required"))
}

pub fn resilient_parse_expression(
    expression: String,
) -> anyhow::Result<(Rc<ExpressionEOFContextAll<'static>>, TranslationErrors)> {
    execute_resillient_parse(expression, |mut parser| parser.expressionEOF())
}

pub fn resilient_parse_pipeline(
    pipeline: String,
) -> anyhow::Result<(Rc<PipelineEOFContextAll<'static>>, TranslationErrors)> {
    execute_resillient_parse(pipeline, |mut parser| parser.pipelineEOF())
}

pub fn parse_identifier(
    identifier: String,
) -> Result<Rc<IdentifierContextAll<'static>>, TranslationErrors> {
    execute_parse(identifier, |mut parser| parser.identifierEOF())
        .map(|ctx| ctx.identifier().expect("required"))
}

pub fn resilient_parse_identifier(
    identifier: String,
) -> anyhow::Result<(Rc<IdentifierEOFContextAll<'static>>, TranslationErrors)> {
    execute_resillient_parse(identifier, |mut parser| parser.identifierEOF())
}

pub fn parse_simple_identifier(
    identifier: String,
) -> Result<Rc<SimpleIdentifierContextAll<'static>>, TranslationErrors> {
    execute_parse(identifier, |mut parser| parser.simpleIdentifierEOF())
        .map(|ctx| ctx.simpleIdentifier().expect("required"))
}

pub fn parse_type(
    type_string: String,
) -> Result<Rc<HamelintypeContextAll<'static>>, TranslationErrors> {
    execute_parse(type_string, |mut parser| parser.hamelintype())
}

// TODO: Refactor to usu a Rust RangeInclusive, because it's more idiomatic and easier to use
#[derive(Debug, Clone, Default)]
pub struct TimeRange {
    pub start: Option<chrono::DateTime<Utc>>,
    pub end: Option<chrono::DateTime<Utc>>,
}

#[derive(Clone)]
pub struct Compiler {
    pub expression_environment: Arc<Environment>,
    pub query_environment_provider: Arc<dyn EnvironmentProvider>,
    pub time_range_filter: Option<SQLExpression>,
    pub registry: Arc<FunctionRegistry>,
}

#[derive(Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub struct FunctionDescription {
    pub name: String,
    pub parameters: String,
}

impl Compiler {
    pub fn new() -> Self {
        Self {
            expression_environment: Arc::new(Environment::default()),
            query_environment_provider: Arc::new(CatalogProvider::default()),
            time_range_filter: None,
            registry: Arc::new(FunctionRegistry::default()),
        }
    }

    pub fn get_function_descriptions(&self) -> Vec<FunctionDescription> {
        self.registry
            .function_defs
            .iter()
            .flat_map(|(name, defs)| {
                defs.iter().map(|def| FunctionDescription {
                    name: name.clone(),
                    parameters: def.parameters.to_string(),
                })
            })
            .collect()
    }

    /// Use the given environment for compiling expressions
    pub fn set_environment(&mut self, environment: Arc<Environment>) {
        self.expression_environment = environment.clone();
    }

    /// Use the given environment provider for compiling queries.
    pub fn set_environment_provider(&mut self, provider: Arc<dyn EnvironmentProvider>) {
        self.query_environment_provider = provider;
    }

    pub fn set_time_range(&mut self, time_range: TimeRange) {
        let ident = ColumnReference::new(SimpleIdentifier::new("timestamp").into());
        let mut range = RangeBuilder::default();

        if let Some(from) = time_range.start {
            range = range.with_begin(
                TimestampLiteral::new(from).into(),
                SQLTimestampTzType::new(3).into(),
            )
        }

        if let Some(to) = time_range.end {
            range = range.with_end(
                TimestampLiteral::new(to).into(),
                SQLTimestampTzType::new(3).into(),
            );
        }

        self.time_range_filter = Some(within_range(ident.into(), range))
    }

    pub fn set_time_range_expression(
        &mut self,
        hamelin_range_expression: String,
    ) -> Result<(), ContextualTranslationErrors> {
        let templated = format!("WITHIN {}", hamelin_range_expression);
        let tree = parse_command(templated.clone())
            .map_err(|e| ContextualTranslationErrors::new(templated.clone(), e))?;
        let previous = PendingQuery::new(
            SQLQuery::default(),
            Environment::default().with_binding("timestamp".parse().unwrap(), TIMESTAMP),
        );

        let filter = if let CommandContextAll::WithinCommandContext(ctx) = tree.as_ref() {
            within::translate(
                ctx,
                &previous,
                QueryTranslationContext::new(
                    None,
                    self.query_environment_provider.clone(),
                    self.registry.clone(),
                    None,
                ),
            )
            .map_err(|e| ContextualTranslationErrors::new(templated.clone(), e))?
            .query
            .where_
            .expect("withins always have a where clause")
        } else {
            unreachable!()
        };
        self.time_range_filter = Some(filter);
        Ok(())
    }

    pub fn compile_expression(
        &self,
        expression: String,
    ) -> Result<ExpressionTranslation, ContextualTranslationErrors> {
        let ctx = parse_expression(expression.clone())
            .map_err(|e| ContextualTranslationErrors::new(expression.clone(), e))?;

        HamelinExpression::new(
            ctx,
            ExpressionTranslationContext::new(
                self.expression_environment.clone(),
                self.registry.clone(),
                FunctionTranslationContext::default(),
                None,
                Rc::new(RefCell::new(None)),
            ),
        )
        .translate()
        .map_err(|e| ContextualTranslationErrors::new(expression, e))
    }

    pub fn compile(
        &self,
        hmln: String,
    ) -> Result<StatementTranslation, ContextualTranslationErrors> {
        let ctx = parse_query(hmln.clone())
            .map_err(|e| ContextualTranslationErrors::new(hmln.clone(), e))?;

        let translation_context = QueryTranslationContext::new(
            self.time_range_filter.clone(),
            self.query_environment_provider.clone(),
            self.registry.clone(),
            None,
        );
        let pending = HamelinQuery::new(ctx.clone(), translation_context.clone()).translate();
        if !pending.errors.is_empty() {
            Err(ContextualTranslationErrors::new(hmln, pending.errors))
        } else {
            let cols = pending.translation.env.into_external_columns();
            let res = match pending.translation.statement {
                Statement::SQLQuery(sqlquery) => QueryTranslation {
                    translation: Translation {
                        sql: sqlquery.to_string(),
                        columns: cols,
                    },
                }
                .into(),
                Statement::DML(dml) => DMLTranslation {
                    translation: Translation {
                        sql: dml.to_string(),
                        columns: cols,
                    },
                }
                .into(),
            };

            Ok(res)
        }
    }

    pub fn compile_query(
        &self,
        hmln: String,
    ) -> Result<QueryTranslation, ContextualTranslationErrors> {
        match self.compile(hmln.clone())? {
            StatementTranslation::Query(query_translation) => Ok(query_translation),
            StatementTranslation::DML(_) => {
                let len = hmln.len();
                Err(ContextualTranslationErrors::new(
                    hmln,
                    TranslationError::new(Context::new(0..=len - 1, "statement has side effects"))
                        .single(),
                ))
            }
        }
    }

    pub fn compile_dml(&self, hmln: String) -> Result<DMLTranslation, ContextualTranslationErrors> {
        match self.compile(hmln.clone())? {
            StatementTranslation::DML(dmltranslation) => Ok(dmltranslation),
            StatementTranslation::Query(_) => {
                let len = hmln.len();
                Err(ContextualTranslationErrors::new(
                    hmln,
                    TranslationError::new(Context::new(0..=len - 1, "statement is a query"))
                        .single(),
                ))
            }
        }
    }

    pub fn compile_query_at(&self, query: String, at: Option<usize>) -> ContextualResult {
        // Travis wants me to panic if I am given a magic string value. That way he can test his panic handling?
        if query
            == "gQ3!mV@x2#Z9^LN7eKd$8wuT0pFzY*b&XHf5+v1RAoJ6MqPCrslijEkDWgBtUO4nchSmyV9Z$L&N^eXpQa"
        {
            panic!("Panic requested");
        }

        let mut res = ContextualResult::new(query.clone());

        let (ctx, errors) = match resilient_parse_query(query.clone()) {
            Ok(s) => s,
            Err(e) => {
                res.add_error(TranslationError::fatal(query.as_str(), e.into()));
                return res;
            }
        };
        res.add_errors(errors);

        let translation_context = QueryTranslationContext::new(
            self.time_range_filter.clone(),
            self.query_environment_provider.clone(),
            self.registry.clone(),
            at,
        );
        let pending = ctx
            .query()
            .map(|sctx| HamelinQuery::new(sctx, translation_context.clone()).translate())
            .unwrap_or_default();

        res.with_pending_result(pending)
            .with_completions(translation_context.completions.clone())
    }

    pub fn get_statement_datasets(
        &self,
        query: String,
    ) -> Result<Vec<String>, ContextualTranslationErrors> {
        let ctx = parse_query(query.clone())
            .map_err(|e| ContextualTranslationErrors::new(query.clone(), e))?;

        let pending = HamelinQuery::new(
            ctx.clone(),
            QueryTranslationContext::new(
                self.time_range_filter.clone(),
                self.query_environment_provider.clone(),
                self.registry.clone(),
                None,
            ),
        )
        .translate()
        .into_result()
        .map_err(|e| ContextualTranslationErrors::new(query.clone(), e))?;

        Ok(pending
            .statement
            .get_table_references()
            .into_iter()
            .map(|t| t.name.to_hamelin())
            .collect())
    }
}

impl Display for Compiler {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "Hamelin Compiler.\n")?;
        write!(
            f,
            "Environment (for expressions): {}\n",
            self.expression_environment
        )?;
        write!(
            f,
            "EnvironmentProvider (for queries): {:#?}\n",
            self.query_environment_provider
        )?;

        Ok(())
    }
}

#[derive(Clone)]
pub struct CompilerContext {}

#[cfg(test)]
mod tests {
    #[test]
    fn verify_send_sync() {
        fn verify_send_sync<T: Send + Sync>() {}
        verify_send_sync::<super::Compiler>();
    }
}