hamelin_analysis 0.9.4

//! Test helper utilities for incremental tests

use std::{collections::HashMap, sync::Arc};

use chrono::{DateTime, TimeZone, Utc};
use hamelin_lib::tree::ast::query::Query;

use crate::incremental::TimeRange;
use hamelin_lib::tree::builder::query as build_query;
use hamelin_lib::type_check_with_provider;
use hamelin_lib::{
    provider::EnvironmentProvider,
    tree::{
        ast::identifier::{Identifier, SimpleIdentifier},
        builder::pipeline::PipelineBuilder,
        typed_ast::{pipeline::TypedPipeline, query::TypedStatement},
    },
    types::{array::Array, struct_type::Struct, INT, STRING, TIMESTAMP},
};

/// Standard timestamp field name used in tests
pub const TIMESTAMP_FIELD: &str = "timestamp";

/// Create a timestamp from a simple string format: "YYYY-MM-DD HH:MM:SS"
pub fn ts(s: &str) -> DateTime<Utc> {
    // Parse format: "YYYY-MM-DD HH:MM:SS"
    let parts: Vec<&str> = s.split(' ').collect();
    let date_parts: Vec<&str> = parts[0].split('-').collect();
    let time_parts: Vec<&str> = parts[1].split(':').collect();

    let year = date_parts[0].parse::<i32>().unwrap();
    let month = date_parts[1].parse::<u32>().unwrap();
    let day = date_parts[2].parse::<u32>().unwrap();
    let hour = time_parts[0].parse::<u32>().unwrap();
    let minute = time_parts[1].parse::<u32>().unwrap();
    let second = time_parts[2].parse::<u32>().unwrap();

    Utc.with_ymd_and_hms(year, month, day, hour, minute, second)
        .unwrap()
}

/// Create an inclusive time range (end included) from two timestamp strings
pub fn time_range(start: &str, end: &str) -> TimeRange {
    (ts(start)..=ts(end)).into()
}

/// Create an exclusive-end time range (end excluded) from two timestamp strings
pub fn time_range_exclusive(start: &str, end: &str) -> TimeRange {
    (ts(start)..ts(end)).into()
}

/// Create a stale ranges map with a single table
pub fn stale_ranges(
    table_name: &str,
    start: &str,
    end: &str,
) -> HashMap<Identifier, Option<TimeRange>> {
    let mut map = HashMap::new();
    let simple_id: SimpleIdentifier = table_name.parse().unwrap();
    map.insert(simple_id.into(), Some(time_range(start, end)));
    map
}

/// Create a stale ranges map with multiple tables
pub fn multi_stale_ranges(ranges: &[(&str, &str, &str)]) -> HashMap<Identifier, Option<TimeRange>> {
    let mut map = HashMap::new();
    for (table_name, start, end) in ranges {
        let simple_id: SimpleIdentifier = table_name.parse().unwrap();
        map.insert(simple_id.into(), Some(time_range(start, end)));
    }
    map
}

/// Get the timestamp field identifier
pub fn timestamp_field() -> SimpleIdentifier {
    TIMESTAMP_FIELD.parse().unwrap()
}

/// Mock environment provider for testing
///
/// Provides schema for test tables:
/// - `events`: timestamp (TIMESTAMP), user (STRING), severity (STRING), amount (INT), tags (ARRAY(STRING))
/// - `users`: id (STRING), name (STRING)
/// - `other_events`: timestamp (TIMESTAMP), data (STRING)
#[derive(Debug)]
pub struct MockIncrementalProvider;

impl EnvironmentProvider for MockIncrementalProvider {
    fn reflect_columns(&self, name: &Identifier) -> anyhow::Result<Struct> {
        let mut env = Struct::default();

        let events: Identifier = SimpleIdentifier::new("events").into();
        let users: Identifier = SimpleIdentifier::new("users").into();
        let other_events: Identifier = SimpleIdentifier::new("other_events").into();

        if name == &events {
            env = env
                .with_str("timestamp", TIMESTAMP)
                .with_str("user", STRING)
                .with_str("severity", STRING)
                .with_str("amount", INT)
                .with_str("user_id", STRING)
                .with_str("tags", Array::new(STRING).into());
        } else if name == &users {
            env = env.with_str("id", STRING).with_str("name", STRING);
        } else if name == &other_events {
            env = env
                .with_str("timestamp", TIMESTAMP)
                .with_str("data", STRING);
        } else {
            anyhow::bail!("Unknown table: {:?}", name);
        }

        Ok(env)
    }

    fn reflect_datasets(&self) -> anyhow::Result<Vec<Identifier>> {
        Ok(vec![
            SimpleIdentifier::new("events").into(),
            SimpleIdentifier::new("users").into(),
            SimpleIdentifier::new("other_events").into(),
        ])
    }
}

/// Build a typed pipeline using the mock provider
pub fn build_pipeline(builder: PipelineBuilder) -> Arc<TypedPipeline> {
    build_typed_query(build_query().main(builder.build()).build()).pipeline
}

/// Build a typed query using the mock provider
pub fn build_typed_query(query: Query) -> TypedStatement {
    type_check_with_provider(query, Arc::new(MockIncrementalProvider)).output
}