hamelin_analysis 0.16.4

//! Tests for time range calculation (query_range and replace_range)

use crate::incremental::{compute_incremental_ranges_for_pipeline, TimeRange};
use hamelin_lib::tree::{
    ast::dataset_identifier::QualifiedDatasetIdentifier,
    builder::{
        add, at_day, at_hour, call, eq, field_ref, hours, minutes, pattern_quantified, pipeline,
        string, table_ref, PipelineBuilder,
    },
};
use rstest::rstest;
use std::collections::HashMap;

use super::helpers::{
    build_pipeline, stale_ranges, test_default_space, time_range, timestamp_field,
    MockIncrementalProvider,
};

#[rstest]
#[case::basic_passthrough(
    pipeline()
        .from(|f| f.table_reference("events"))
        .where_cmd(eq(field_ref("severity"), string("high"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"), // query_range
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")  // replace_range
)]
#[case::set_passthrough(
    pipeline()
        .from(|f| f.table_reference("events"))
        .set_cmd(|l| l.named_field("risk_score", field_ref("severity"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")
)]
#[case::match_passthrough(
    pipeline()
        .match_cmd(|m| m.pattern(pattern_quantified(table_ref("events"), None))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")
)]
#[case::window_lookback_1h(
    pipeline()
        .from(|f| f.table_reference("events"))
        .window(|w| w
            .named_field("sum_amount", call("sum").arg(field_ref("amount")))
            .within(hours(-1))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 13:30:00", "2024-01-01 16:45:00"), // query_range: 1h lookback
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")  // replace_range: unchanged
)]
#[case::window_lookback_lookforward(
    pipeline()
        .from(|f| f.table_reference("events"))
        .window(|w| w
            .named_field("sum_amount", call("sum").arg(field_ref("amount")))
            .within(hours(-1)..minutes(30))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 13:30:00", "2024-01-01 17:15:00"), // query_range: 1h back, 30m forward
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")  // replace_range: unchanged
)]
#[case::agg_hourly(
    pipeline()
        .from(|f| f.table_reference("events"))
        .agg(|a| a
            .named_aggregate("max_amount", call("max").arg(field_ref("amount")))
            .named_group("timestamp", at_hour(field_ref("timestamp")))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"), // query_range: full hours
    time_range("2024-01-01 14:00:00", "2024-01-01 16:00:00")  // replace_range: truncated hours
)]
#[case::agg_daily(
    pipeline()
        .from(|f| f.table_reference("events"))
        .agg(|a| a
            .named_aggregate("count", call("count"))
            .named_group("timestamp", at_day(field_ref("timestamp")))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-03 16:45:00"),
    time_range("2024-01-01 00:00:00", "2024-01-04 00:00:00"), // query_range: full days
    time_range("2024-01-01 00:00:00", "2024-01-03 00:00:00")  // replace_range: truncated days
)]
#[case::agg_with_multiple_groupby(
    pipeline()
        .from(|f| f.table_reference("events"))
        .agg(|a| a
            .named_aggregate("count", call("count"))
            .named_group("timestamp", at_hour(field_ref("timestamp")))
            .named_group("user_id", field_ref("user.id"))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"), // query_range: full hours
    time_range("2024-01-01 14:00:00", "2024-01-01 16:00:00")  // replace_range: truncated hours
)]
#[case::window_then_agg(
    pipeline()
        .from(|f| f.table_reference("events"))
        .window(|w| w
            .named_field("sum_amount", call("sum").arg(field_ref("amount")))
            .within(hours(-1))
        )
        .agg(|a| a
            .named_aggregate("max_sum", call("max").arg(field_ref("sum_amount")))
            .named_group("timestamp", at_hour(field_ref("timestamp")))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 13:00:00", "2024-01-01 17:00:00"), // query_range: lookback + full hours
    time_range("2024-01-01 14:00:00", "2024-01-01 16:00:00")  // replace_range: truncated hours
)]
#[case::suppress_widens_query_to_bucket_boundaries(
    pipeline()
        .from(|f| f.table_reference("events"))
        .suppress(hours(1), |s| s),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"), // read full hours touched by stale outputs
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")  // replace_range unchanged on timestamp
)]
#[case::set_transforms_timestamp(
    pipeline()
        .from(|f| f.table_reference("events"))
        .set_cmd(|l| l.named_field("timestamp", add(field_ref("timestamp"), hours(1)))),
    stale_ranges("events", "2024-01-01 14:00:00", "2024-01-01 16:00:00"),
    time_range("2024-01-01 14:00:00", "2024-01-01 16:00:00"), // query_range: unchanged
    time_range("2024-01-01 15:00:00", "2024-01-01 17:00:00")  // replace_range: shifted by 1h
)]
#[case::select_truncates_timestamp(
    pipeline()
        .from(|f| f.table_reference("events"))
        .select(|s| s
            .named_field("timestamp", at_hour(field_ref("timestamp")))
            .field("user")
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"), // query_range: expands to cover truncation boundaries
    time_range("2024-01-01 14:00:00", "2024-01-01 16:00:00")  // replace_range: truncated
)]
fn test_time_range_calculation(
    #[case] pipeline_builder: PipelineBuilder,
    #[case] stale_ranges_map: HashMap<QualifiedDatasetIdentifier, Option<TimeRange>>,
    #[case] expected_query_range: TimeRange,
    #[case] expected_replace_range: TimeRange,
) {
    let pipeline = build_pipeline(pipeline_builder);
    let provider = MockIncrementalProvider;
    let result = compute_incremental_ranges_for_pipeline(
        &pipeline,
        stale_ranges_map,
        &timestamp_field(),
        None,
        false,
        &[],
        Some(&test_default_space()),
        &provider,
    );

    assert!(result.is_ok(), "Expected success, got error: {:?}", result);
    let ranges = result.unwrap();

    assert_eq!(
        ranges.query_range, expected_query_range,
        "Query range mismatch"
    );
    assert_eq!(
        ranges.replace_range, expected_replace_range,
        "Replace range mismatch"
    );
}

/// LOOKUP / JOIN are gated by `allow_lookups` in both range computation and strategy checks.
#[rstest]
#[case::join_passthrough(
    pipeline()
        .from(|f| f.table_reference("events"))
        .join("users", eq(field_ref("events.user_id"), field_ref("users.id"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")
)]
#[case::lookup_passthrough(
    pipeline()
        .from(|f| f.table_reference("events"))
        .lookup(
            "users",
            eq(field_ref("events.user_id"), field_ref("users.id")),
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")
)]
fn test_time_ranges_allow_lookups(
    #[case] pipeline_builder: PipelineBuilder,
    #[case] stale_ranges_map: HashMap<QualifiedDatasetIdentifier, Option<TimeRange>>,
    #[case] expected_query_range: TimeRange,
    #[case] expected_replace_range: TimeRange,
) {
    let pipeline = build_pipeline(pipeline_builder);
    let provider = MockIncrementalProvider;
    let result = compute_incremental_ranges_for_pipeline(
        &pipeline,
        stale_ranges_map,
        &timestamp_field(),
        None,
        true,
        &[],
        Some(&test_default_space()),
        &provider,
    );

    assert!(result.is_ok(), "Expected success, got error: {:?}", result);
    let ranges = result.unwrap();

    assert_eq!(
        ranges.query_range, expected_query_range,
        "Query range mismatch"
    );
    assert_eq!(
        ranges.replace_range, expected_replace_range,
        "Replace range mismatch"
    );
}

#[test]
fn qualified_stale_range_keys_match_unqualified_pipeline_refs() {
    use hamelin_lib::tree::ast::identifier::SimpleIdentifier;

    let pipeline = build_pipeline(pipeline().from(|f| f.table_reference("events")));
    let stale_ranges_map = stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00");

    let result = compute_incremental_ranges_for_pipeline(
        &pipeline,
        stale_ranges_map,
        &timestamp_field(),
        None,
        false,
        &[],
        Some(&SimpleIdentifier::new("test")),
        &MockIncrementalProvider,
    );

    assert!(
        result.is_ok(),
        "Expected success with qualified stale keys, got error: {:?}",
        result
    );
    let ranges = result.unwrap();
    assert_eq!(
        ranges.replace_range,
        time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")
    );
}

#[test]
fn qualified_pipeline_ref_does_not_match_wrong_stale_key() {
    use crate::incremental::IncrementalAnalysisError;
    use hamelin_lib::tree::builder::table_ref;

    let pipeline =
        build_pipeline(pipeline().from(|f| f.table_reference(table_ref("other:events"))));
    let stale_ranges_map = stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00");

    let result = compute_incremental_ranges_for_pipeline(
        &pipeline,
        stale_ranges_map,
        &timestamp_field(),
        None,
        false,
        &[],
        None,
        &MockIncrementalProvider,
    );

    assert!(
        matches!(
            result,
            Err(IncrementalAnalysisError::DependencyNotProvided(ref id))
            if id.to_string() == "other:events"
        ),
        "Expected DependencyNotProvided for other:events, got {:?}",
        result
    );
}