hamelin_lib 0.9.2

//! Regression tests for bugs discovered during development
//!
//! Each test documents a specific bug that was found and fixed.

use crate::err::TranslationErrors;
use crate::tree::ast::pipeline::Pipeline;
use crate::tree::ast::query::Query;
use crate::tree::ast::ParseWithErrors;
use crate::tree::tests::shared::catalog_with_tables;
use crate::type_check_with_provider;
use crate::types::array::Array;
use crate::types::map::Map;
use crate::types::struct_type::Struct;
use crate::types::{Type, DOUBLE, INT, STRING, TIMESTAMP};
use pretty_assertions::assert_eq;
use rstest::rstest;

/// Helper to create the minimal simba.sysmon_events schema needed for regression tests
fn sysmon_events_schema() -> Vec<(&'static str, Vec<(&'static str, Type)>)> {
    let sysmon_event_struct = Struct::default()
        .with_str("code", STRING)
        .with_str("original", STRING)
        .into();
    let winlog_struct = Struct::default()
        .with_str("event_data", Map::new(STRING, STRING).into())
        .into();
    let host_struct = Struct::default().with_str("name", STRING).into();

    vec![(
        "simba.sysmon_events",
        vec![
            ("timestamp", TIMESTAMP),
            ("event", sysmon_event_struct),
            ("winlog", winlog_struct),
            ("host", host_struct),
        ],
    )]
}

// ======================== Regression Test Cases ========================

/// Bug: DROP command lost fields added by SET commands that appeared before it.
///
/// Example:
/// ```hamelin
/// FROM table
/// | SET reg_key = ...
/// | DROP some_field
/// | WHERE regexp_like(reg_key, ...)  // Error: reg_key not found!
/// ```
///
/// Root cause: DROP created output schema with `base: None`, discarding the
/// environment chain that contained fields from previous SET commands.
///
/// Fix: Changed DROP to preserve the base chain: `base: input_schema.base.clone()`
///
/// ---
///
/// Bug: FROM command with aliased CTEs didn't flatten the CTE's environment chain,
/// causing fields to be lost.
///
/// Example:
/// ```hamelin
/// DEF cte = FROM table
/// | SET host = host.name
/// | DROP some_field;
///
/// FROM x = cte
/// | SELECT timestamp  // Error: timestamp not found!
/// ```
///
/// Root cause: FROM command used `env.fields` instead of flattening the entire
/// environment chain. When a CTE had a base chain (from SET commands), only the
/// top-level fields were exposed, not the fields from the base.
///
/// Fix: Added `flatten_environment_fields()` helper that recursively collects all
/// fields from the environment chain.
///
/// ---
///
/// Bug: JOIN/LOOKUP commands didn't flatten the right-side environment, losing fields
/// from the base chain when joining with CTEs that had SET fields.
///
/// Root cause: JOIN used `right_env.fields.clone()` instead of flattening the environment.
///
/// Fix: Changed JOIN to use `nest_into()` which flattens and nests in one operation.
///
/// ---
///
/// Bug: UNNEST command's flatten helper didn't use the new `Environment::flatten()` method.
///
/// Fix: Simplified `flatten_environment()` helper to use `Environment::flatten()`.
#[rstest]
// DROP preserves SET fields
#[case::drop_preserves_set_fields(
    sysmon_events_schema(),
    r#"
FROM simba.sysmon_events
| SET reg_key = coalesce(winlog.event_data["TargetObject"], '')
| DROP winlog.event_data
| WHERE regexp_like(lower(reg_key), '(?i).*test.*')
"#,
    Struct::default()
        .with_str("reg_key", STRING)
        .with_str("timestamp", TIMESTAMP)
        .with_str("event", Struct::default().with_str("code", STRING).with_str("original", STRING).into())
        .with_str("host", Struct::default().with_str("name", STRING).into())
)]
// FROM with CTEs flattens environment
#[case::from_cte_flattens_environment(
    sysmon_events_schema(),
    r#"
DEF registry_events = FROM simba.sysmon_events
| WHERE event.code == "13"
| SET reg_key = coalesce(winlog.event_data["TargetObject"], '')
| SET host = host.name
| DROP winlog.event_data, event.original;

FROM x = registry_events
| SELECT timestamp, host, reg_key
"#,
    Struct::default()
        .with_str("timestamp", TIMESTAMP)
        .with_str("host", STRING)
        .with_str("reg_key", STRING)
)]
#[case::from_aliased_ctes_preserves_all_fields(
    sysmon_events_schema(),
    r#"
DEF registry_events = FROM simba.sysmon_events
| WHERE event.code == "13"
| SET reg_key = coalesce(winlog.event_data["TargetObject"], '')
| SET host = host.name
| DROP winlog.event_data, event.original;

FROM x = registry_events, y = registry_events
| SELECT x.timestamp, y.timestamp, x.host, y.host
"#,
    Struct::default()
        .with_str("x", Struct::default().with_str("timestamp", Type::Timestamp).with_str("host", STRING).into())
        .with_str("y", Struct::default().with_str("timestamp", Type::Timestamp).with_str("host", STRING).into())
)]
#[case::window_finds_timestamp_after_cte_from(
    sysmon_events_schema(),
    r#"
DEF registry_events = FROM simba.sysmon_events
| WHERE event.code == "13"
| SET host = host.name
| DROP winlog.event_data;

DEF process_events = FROM simba.sysmon_events
| WHERE event.code == "1"
| SET host = host.name
| DROP winlog.event_data;

FROM registry_event = registry_events, process_event = process_events
| WINDOW process_event = last_value(process_event),
         registry_event = last_value(registry_event)
  BY host WITHIN 1s
"#,
    Struct::default()
        .with_str("process_event", Struct::default()
            .with_str("host", STRING)
            .with_str("timestamp", TIMESTAMP)
            .with_str("event", Struct::default()
                .with_str("code", STRING)
                .with_str("original", STRING)
                .into())
            .into())
        .with_str("registry_event", Struct::default()
            .with_str("host", STRING)
            .with_str("timestamp", TIMESTAMP)
            .with_str("event", Struct::default()
                .with_str("code", STRING)
                .with_str("original", STRING)
                .into())
            .into())
        .with_str("host", STRING)
        .with_str("timestamp", TIMESTAMP)
        .with_str("event", Struct::default().with_str("code", STRING).with_str("original", STRING).into())
)]
// JOIN preserves base chain
#[case::join_preserves_base_chain(
    sysmon_events_schema(),
    r#"
DEF enriched = FROM simba.sysmon_events
| WHERE event.code == "1"
| SET hostname = host.name
| SET process_name = winlog.event_data["Image"]
| DROP winlog.event_data;

FROM simba.sysmon_events
| WHERE event.code == "3"
| JOIN x = enriched ON host.name == x.hostname
| SELECT timestamp, x.process_name, x.hostname
"#,
    Struct::default()
        .with_str("timestamp", TIMESTAMP)
        .with_str("x", Struct::default().with_str("process_name", STRING).with_str("hostname", STRING).into())
)]
// UNNEST with SET fields
#[case::unnest_with_set_fields(
    sysmon_events_schema(),
    r#"
FROM simba.sysmon_events
| WHERE event.code == "1"
| SET hostname = host.name
| SET process_info = {pid: winlog.event_data["ProcessId"], name: winlog.event_data["Image"]}
| UNNEST process_info
| SELECT hostname, pid, name
"#,
    Struct::default()
        .with_str("hostname", STRING)
        .with_str("pid", STRING)
        .with_str("name", STRING)
)]
// JOIN with CTE, reference joined fields in SET, then DROP the joined struct
#[case::join_then_drop_joined_struct(
    vec![("events", vec![("timestamp", TIMESTAMP), ("count", INT)])],
    r#"
DEF baseline = FROM events
| AGG avg = avg(count), stddev = stddev(count);

FROM events
| AGG count = count() BY timestamp = timestamp@h
| JOIN baseline
| SET threshold_upper = baseline.avg + baseline.stddev
| SET threshold_lower = baseline.avg - baseline.stddev
| DROP baseline
"#,
    Struct::default()
        .with_str("threshold_lower", DOUBLE)
        .with_str("threshold_upper", DOUBLE)
        .with_str("timestamp", TIMESTAMP)
        .with_str("count", INT)
)]
// Rebinding a top-level struct field to a scalar should completely shadow the struct
#[case::rebind_struct_to_scalar_shadows_children(
    vec![(
        "logs",
        vec![
            ("timestamp", TIMESTAMP),
            ("message", STRING),
            ("host", Type::Struct(
                Struct::default()
                    .with_str("name", STRING)
                    .with_str("ip", STRING)
                    .with_str("os", Type::Struct(
                        Struct::default()
                            .with_str("name", STRING)
                            .with_str("version", STRING)
                    ))
            ))
        ]
    )],
    r#"
FROM logs
| SET host = host.name
"#,
    Struct::default()
        .with_str("host", STRING)
        .with_str("timestamp", TIMESTAMP)
        .with_str("message", STRING)
)]
// Rebinding a nested struct field to a scalar should completely shadow the struct
#[case::rebind_nested_struct_to_scalar_shadows_children(
    vec![(
        "events",
        vec![
            ("timestamp", TIMESTAMP),
            ("event", Type::Struct(
                Struct::default()
                    .with_str("id", STRING)
                    .with_str("host", Type::Struct(
                        Struct::default()
                            .with_str("name", STRING)
                            .with_str("ip", STRING)
                            .with_str("os", Type::Struct(
                                Struct::default()
                                    .with_str("name", STRING)
                                    .with_str("version", STRING)
                            ))
                    ))
            ))
        ]
    )],
    r#"
FROM events
| SET event.host = event.host.name
"#,
    Struct::default()
        .with_str("event", Type::Struct(
            Struct::default()
                .with_str("id", STRING)
                .with_str("host", STRING)
        ))
        .with_str("timestamp", TIMESTAMP)
)]
// Test rebinding after AGG with array_agg - matches pattern from DET_33LF1U2.hmln
#[case::rebind_after_array_agg(
    vec![(
        "events",
        vec![
            ("id", INT),
            ("data", Type::Struct(
                Struct::default()
                    .with_str("host", Type::Struct(
                        Struct::default()
                            .with_str("name", STRING)
                            .with_str("ip", STRING)
                    ))
            ))
        ]
    )],
    r#"
FROM events
| AGG all_data = array_agg(data)
| SET first_data = all_data[0]
| SET first_data.host = first_data.host.name
"#,
    Struct::default()
        .with_str("first_data", Type::Struct(
            Struct::default()
                .with_str("host", STRING)
        ))
        .with_str("all_data", Array::new(Type::Struct(
            Struct::default()
                .with_str("host", Type::Struct(
                    Struct::default()
                        .with_str("name", STRING)
                        .with_str("ip", STRING)
                ))
        )).into())
)]
// Test rebinding a struct field created in previous SET to a value from within that struct
#[case::rebind_struct_field_from_same_level(
    vec![(
        "events",
        vec![
            ("id", INT),
            ("name", STRING),
            ("ip", STRING)
        ]
    )],
    r#"
FROM events
| SET event = {host: {name: name, ip: ip}}
| SET event.host = event.host.name
"#,
    Struct::default()
        .with_str("event", Type::Struct(
            Struct::default()
                .with_str("host", STRING)
        ))
        .with_str("id", INT)
        .with_str("name", STRING)
        .with_str("ip", STRING)
)]
fn test_regressions(
    #[case] tables: Vec<(&str, Vec<(&str, Type)>)>,
    #[case] query_str: &str,
    #[case] expected_schema: Struct,
) -> Result<(), TranslationErrors> {
    let query = Query::parse_result(query_str)?;
    let catalog = catalog_with_tables(tables);
    let typed = type_check_with_provider(query, catalog).into_result()?;

    let actual_schema = typed.schema();

    assert_eq!(actual_schema, expected_schema);
    Ok(())
}

/// Verify that parse → Display → parse round-trips produce equivalent ASTs.
///
/// Any Hamelin syntax that fails to round-trip through Display can be added as a case here.
#[rstest]
#[case::from_aliased("FROM x = events")]
#[case::join_aliased("FROM users | JOIN o = orders ON o.user_id == user_id")]
#[case::join_qualified_aliased("FROM users | JOIN o = schema.orders ON o.id == id")]
#[case::lookup_aliased("FROM users | LOOKUP o = orders ON o.user_id == user_id")]
fn test_display_roundtrip(#[case] input: &str) {
    let original = Pipeline::parse_result(input).expect("input should parse");
    let displayed = original.to_string();
    let reparsed = Pipeline::parse_result(&displayed);

    assert!(
        reparsed.is_ok(),
        "Failed to re-parse displayed pipeline.\n  Input:     {input}\n  Displayed: {displayed}\n  Error: {}",
        reparsed.unwrap_err()
    );

    assert_eq!(
        reparsed.unwrap(),
        original,
        "Round-tripped AST differs from original.\n  Input:     {input}\n  Displayed: {displayed}"
    );
}