hamelin_lib 0.21.6

use std::collections::HashMap;
use std::sync::Arc;

use crate::catalog::{Catalog, CatalogProvider, Column, HamelinType};
use crate::tree::ast::dataset_identifier::QualifiedDatasetIdentifier;
use crate::tree::ast::identifier::SimpleIdentifier;
use crate::tree::ast::query::Query;
use crate::tree::options::{TemplateParameterKind, TypeCheckOptions};
use crate::tree::typed_ast::query::TypedStatement;
use crate::types::INTERVAL;
use crate::{parse, type_check_with_options};

#[test]
fn templated_from_intersects_row_schemas() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:t_alice").unwrap(),
        vec![
            Column::new("id", HamelinType::Int),
            Column::new("x", HamelinType::String),
        ],
    );
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:t_bob").unwrap(),
        vec![Column::new("id", HamelinType::Int)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["t_alice".to_string(), "t_bob".to_string()]),
    );
    let q = parse(r#"FROM db:${which}"#).into_result().expect("parse");
    let typed: TypedStatement = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect("typecheck");
    let schema = typed.schema();
    assert!(schema.lookup(&SimpleIdentifier::new("id")).is_some());
    assert!(schema.lookup(&SimpleIdentifier::new("x")).is_none());
}

#[test]
fn templated_from_resolves_unqualified_ctes() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:base").unwrap(),
        vec![Column::new("id", HamelinType::Int)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["selected".to_string()]),
    );
    let q = parse("DEF selected = FROM db:base; FROM ${which}")
        .into_result()
        .expect("parse");
    let typed: TypedStatement = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect("typecheck");

    assert!(typed
        .schema()
        .lookup(&SimpleIdentifier::new("id"))
        .is_some());
}

#[test]
fn templated_from_preserves_cte_schema_uncertainty() {
    let provider = Arc::new(CatalogProvider::try_from(Catalog::default()).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["selected".to_string()]),
    );
    let q = parse("DEF selected = FROM db:missing; FROM ${which} | SELECT foo")
        .into_result()
        .expect("parse");
    let errors = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .errors;

    assert_eq!(errors.len(), 1, "{errors}");
    assert_eq!(
        errors.iter().next().unwrap().primary.message,
        "Dataset not found: db:missing"
    );
}

#[test]
fn templated_from_preserves_common_cte_schema_uncertainty() {
    let provider = Arc::new(CatalogProvider::try_from(Catalog::default()).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["a".to_string(), "b".to_string()]),
    );
    let q =
        parse("DEF a = FROM db:missing_a; DEF b = FROM db:missing_b; FROM ${which} | SELECT foo")
            .into_result()
            .expect("parse");
    let errors = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .errors;

    assert_eq!(errors.len(), 2, "{errors}");
    assert!(errors
        .iter()
        .all(|error| error.primary.message.starts_with("Dataset not found:")));
}

#[test]
fn templated_from_keeps_errors_independent_in_any_cte() {
    let provider = Arc::new(CatalogProvider::try_from(Catalog::default()).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["a".to_string(), "b".to_string()]),
    );
    let q = parse("DEF a = FROM db:missing_a; DEF b = SET x = 1; FROM ${which} | SELECT foo")
        .into_result()
        .expect("parse");
    let errors = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .errors;

    assert_eq!(errors.len(), 2, "{errors}");
    assert!(errors
        .iter()
        .any(|error| error.primary.message == "Field not found"));
}

#[test]
fn templated_join_requires_explicit_alias() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:base").unwrap(),
        vec![Column::new("id", HamelinType::Int)],
    );
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:t_alice").unwrap(),
        vec![Column::new("id", HamelinType::Int)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["t_alice".to_string()]),
    );
    let q = parse(r#"FROM db:base | JOIN db:${which} ON true"#)
        .into_result()
        .expect("parse");
    let err = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect_err("templated JOIN without alias should fail");
    assert!(err.to_string().contains("require an explicit alias"));
}

#[test]
fn templated_match_requires_explicit_alias() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:t_alice").unwrap(),
        vec![Column::new("id", HamelinType::Int)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "which".to_string(),
        TemplateParameterKind::IdentifierFragment(vec!["t_alice".to_string()]),
    );
    let q = parse(r#"MATCH db:${which}"#).into_result().expect("parse");
    let err = type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect_err("templated MATCH without alias should fail");
    assert!(err.to_string().contains("require an explicit alias"));
}

#[test]
fn templated_ts_trunc_typechecks_with_interval_primitive() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:events").unwrap(),
        vec![Column::new("timestamp", HamelinType::Timestamp)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "timeslice".to_string(),
        TemplateParameterKind::Primitive(INTERVAL),
    );
    let q = parse("FROM db:events | SELECT bucket = timestamp @${timeslice}")
        .into_result()
        .expect("parse");
    type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect("typecheck");
}

#[test]
fn interval_primitive_works_as_leaf_expression() {
    let mut catalog = Catalog::default();
    catalog.0.insert(
        QualifiedDatasetIdentifier::from_canonical_str("db:events").unwrap(),
        vec![Column::new("timestamp", HamelinType::Timestamp)],
    );
    let provider = Arc::new(CatalogProvider::try_from(catalog).unwrap());
    let mut tmpl = HashMap::new();
    tmpl.insert(
        "timeslice".to_string(),
        TemplateParameterKind::Primitive(INTERVAL),
    );
    let q = parse("FROM db:events | WHERE timestamp > now() - ${timeslice}")
        .into_result()
        .expect("parse");
    type_check_with_options::<Query>(
        Arc::new(q),
        TypeCheckOptions::builder()
            .provider(provider)
            .maybe_template_parameters(Some(Arc::new(tmpl)))
            .build(),
    )
    .into_result()
    .expect("typecheck");
}