From 75f6c0aa8580529a5bd620b45dc03a765f928cbc Mon Sep 17 00:00:00 2001 From: Kevin Costner <120246174+kevincostner17@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:49:13 +0530 Subject: [PATCH] fix(stream): reject unknown --timestamp columns freshdata stream accepted a --timestamp that named no input column: it exited 0 with no warning and time-series cleaning was quietly skipped for every batch. The stream CLI now checks its column-name flags (--timestamp, and with it --watermark, --entity-id and --ordered-dedupe-keys, plus --target-column and --id-columns) against the first batch before anything is cleaned or written. An unknown name exits 1 with a one-line "freshdata: error: ..." message, a did-you-mean hint when one is close, and no output or .partial file. stream-kafka checks --target-column and --id-columns the same way. --- src/freshdata/streaming/_cli.py | 69 ++++++++++- tests/test_streaming_cli_columns.py | 185 ++++++++++++++++++++++++++++ 2 files changed, 250 insertions(+), 4 deletions(-) create mode 100644 tests/test_streaming_cli_columns.py diff --git a/src/freshdata/streaming/_cli.py b/src/freshdata/streaming/_cli.py index 804a8d68..1070ec15 100644 --- a/src/freshdata/streaming/_cli.py +++ b/src/freshdata/streaming/_cli.py @@ -11,6 +11,7 @@ import argparse import contextlib +import difflib import json import os from collections.abc import Iterator @@ -180,6 +181,60 @@ def _timeseries_config(args: argparse.Namespace) -> Any: return TimeSeriesCleanConfig(**kwargs) +def _column_options(args: argparse.Namespace) -> list[tuple[str, list[str]]]: + """The column-name flags given on the command line, as ``(flag, [names])``. + + The time-series flags only count when ``--timestamp`` is given, because + without it they are not used at all. + """ + opts: list[tuple[str, list[str]]] = [] + if getattr(args, "target_column", None): + opts.append(("--target-column", [args.target_column])) + if getattr(args, "id_columns", None): + opts.append(("--id-columns", list(args.id_columns))) + if getattr(args, "timestamp", None): + opts.append(("--timestamp", [args.timestamp])) + if getattr(args, "watermark", None): + opts.append(("--watermark", [args.watermark])) + if getattr(args, "entity_id", None): + opts.append(("--entity-id", list(args.entity_id))) + if getattr(args, "ordered_dedupe_keys", None): + opts.append(("--ordered-dedupe-keys", list(args.ordered_dedupe_keys))) + return opts + + +def _check_column_options(columns: Any, args: argparse.Namespace) -> None: + """Raise :class:`ValueError` if a column-name flag names no column in *columns*.""" + present = list(columns) + names = [str(c) for c in present] + problems: list[str] = [] + for flag, wanted in _column_options(args): + for name in wanted: + if name in present: + continue + match = difflib.get_close_matches(name, names, n=1) + hint = f" (did you mean {match[0]!r}?)" if match else "" + problems.append(f"{flag} column {name!r} not found in input{hint}") + if problems: + shown = ", ".join(names[:20]) + (", ..." if len(names) > 20 else "") + raise ValueError(f"{'; '.join(problems)}; input columns: {shown}") + + +def _checked_batches(batches: Iterator[pd.DataFrame], + args: argparse.Namespace) -> Iterator[pd.DataFrame]: + """Pass *batches* through, checking the column-name flags against the first one. + + The check runs before the first batch is cleaned or written, so a bad column + name stops the run before any output (or ``.partial`` file) exists. + """ + first = True + for batch in batches: + if first: + first = False + _check_column_options(batch.columns, args) + yield batch + + def _write_exceptions(cleaner: StreamingCleaner, path: str | None, sanitize_formulas: bool = True) -> None: """Persist any quarantined (late/anomalous) rows to *path* (CSV or Parquet).""" @@ -230,8 +285,10 @@ def _run_stream(cleaner: StreamingCleaner, batches: Iterator[pd.DataFrame], def cmd_stream(args: argparse.Namespace) -> int: cleaner = StreamingCleaner(**_stream_options(args)) sanitize = getattr(args, "sanitize_formulas", True) - batches = _read_chunks(args.input, args.batch_size, - preserve_leading_zeros=cleaner.config.preserve_leading_zeros) + batches = _checked_batches( + _read_chunks(args.input, args.batch_size, + preserve_leading_zeros=cleaner.config.preserve_leading_zeros), + args) return _run_stream(cleaner, batches, _BatchWriter(args.output, sanitize_formulas=sanitize), args.report, args.quiet, @@ -240,10 +297,14 @@ def cmd_stream(args: argparse.Namespace) -> int: def cmd_stream_kafka(args: argparse.Namespace) -> int: + from ._connectors import kafka_batches + cleaner = StreamingCleaner(**_stream_options(args)) - batches = cleaner.clean_kafka( # validates the kafka dependency + # Same source as ``cleaner.clean_kafka`` (which validates the kafka dependency), + # with the column-name flags checked against the first raw batch. + batches = cleaner.clean_batches(_checked_batches(kafka_batches( topic=args.topic, bootstrap_servers=args.bootstrap_servers, - batch_size=args.batch_size, max_batches=args.max_batches) + batch_size=args.batch_size, max_batches=args.max_batches), args)) writer = _BatchWriter(args.output, sanitize_formulas=getattr(args, "sanitize_formulas", True)) if args.report: diff --git a/tests/test_streaming_cli_columns.py b/tests/test_streaming_cli_columns.py new file mode 100644 index 00000000..f0d69639 --- /dev/null +++ b/tests/test_streaming_cli_columns.py @@ -0,0 +1,185 @@ +"""``freshdata stream`` / ``stream-kafka`` reject column-name flags that name no column.""" + +from __future__ import annotations + +import json +import sys +import types +from pathlib import Path + +import pandas as pd +import pytest + +from freshdata.enterprise.cli import main + + +def _events_csv(path: Path, rows: int = 30) -> Path: + pd.DataFrame({ + "event_time": pd.date_range("2024-01-01", periods=rows, freq="min").astype(str), + "sensor": ["a", "b"] * (rows // 2), + "value": [float(i) for i in range(rows)], + }).to_csv(path, index=False) + return path + + +def _assert_no_output(out: Path) -> None: + assert not out.exists() + assert not Path(f"{out}.partial").exists() + + +def test_unknown_timestamp_exits_1_with_one_line_error(tmp_path, capsys): + src = _events_csv(tmp_path / "events.csv") + out = tmp_path / "out.csv" + reports = tmp_path / "rep" + + rc = main(["stream", str(src), "-o", str(out), "--report", str(reports), + "--timestamp", "event_tim", "--anomaly", "mad", "--quiet"]) + + assert rc == 1 + err = capsys.readouterr().err + lines = err.strip().splitlines() + assert len(lines) == 1 + assert lines[0].startswith("freshdata: error: ") + assert "--timestamp column 'event_tim' not found" in lines[0] + _assert_no_output(out) + assert not (reports / "summary.json").exists() + assert not list(reports.glob("batch_*.json")) + + +def test_unknown_timestamp_suggests_close_match(tmp_path, capsys): + src = _events_csv(tmp_path / "events.csv") + + rc = main(["stream", str(src), "-o", str(tmp_path / "out.csv"), + "--timestamp", "event_tim", "--quiet"]) + + assert rc == 1 + assert "(did you mean 'event_time'?)" in capsys.readouterr().err + + +def test_unknown_timestamp_without_close_match_has_no_suggestion(tmp_path, capsys): + src = _events_csv(tmp_path / "events.csv") + + rc = main(["stream", str(src), "-o", str(tmp_path / "out.csv"), + "--timestamp", "zzzz", "--quiet"]) + + assert rc == 1 + err = capsys.readouterr().err + assert "did you mean" not in err + assert "input columns: event_time, sensor, value" in err + + +def test_unknown_timestamp_leaves_no_parquet_output(tmp_path): + pytest.importorskip("pyarrow") + src = _events_csv(tmp_path / "events.csv") + out = tmp_path / "out.parquet" + + rc = main(["stream", str(src), "-o", str(out), "--batch-size", "10", + "--timestamp", "ts", "--quiet"]) + + assert rc == 1 + _assert_no_output(out) + + +def test_valid_timestamp_still_streams(tmp_path, capsys): + src = _events_csv(tmp_path / "events.csv") + out = tmp_path / "out.csv" + reports = tmp_path / "rep" + + rc = main(["stream", str(src), "-o", str(out), "--report", str(reports), + "--batch-size", "10", "--timestamp", "event_time", + "--entity-id", "sensor", "--watermark", "event_time", "--quiet"]) + + assert rc == 0 + assert capsys.readouterr().err == "" + assert len(pd.read_csv(out)) == 30 + assert not Path(f"{out}.partial").exists() + summary = json.loads((reports / "summary.json").read_text()) + assert "time_series" in summary["streaming"] + + +@pytest.mark.parametrize("flag_args, flag", [ + (["--timestamp", "event_time", "--watermark", "evt"], "--watermark"), + (["--timestamp", "event_time", "--entity-id", "sensor", "sensr"], "--entity-id"), + (["--timestamp", "event_time", "--ordered-dedupe-keys", "id"], "--ordered-dedupe-keys"), + (["--target-column", "valu"], "--target-column"), + (["--id-columns", "sensor", "id"], "--id-columns"), +]) +def test_other_unknown_column_flags_exit_1(tmp_path, capsys, flag_args, flag): + src = _events_csv(tmp_path / "events.csv") + out = tmp_path / "out.csv" + + rc = main(["stream", str(src), "-o", str(out), "--quiet", *flag_args]) + + assert rc == 1 + err = capsys.readouterr().err + assert err.startswith("freshdata: error: ") + assert f"{flag} column " in err + _assert_no_output(out) + + +def test_all_unknown_columns_are_reported_together(tmp_path, capsys): + src = _events_csv(tmp_path / "events.csv") + + rc = main(["stream", str(src), "-o", str(tmp_path / "out.csv"), "--quiet", + "--timestamp", "event_tim", "--target-column", "valu"]) + + assert rc == 1 + err = capsys.readouterr().err + assert len(err.strip().splitlines()) == 1 + assert "--timestamp column 'event_tim'" in err + assert "--target-column column 'valu'" in err + + +def test_time_series_flags_are_not_checked_without_timestamp(tmp_path): + src = _events_csv(tmp_path / "events.csv") + out = tmp_path / "out.csv" + + rc = main(["stream", str(src), "-o", str(out), "--quiet", "--watermark", "nope"]) + + assert rc == 0 + assert out.exists() + + +def _fake_kafka(monkeypatch): + records = [json.dumps({"customer_id": i, "amount": float(i % 5)}).encode() + for i in range(200)] + fake = types.ModuleType("kafka") + + class KafkaConsumer: + def __init__(self, topic, **kwargs): + self._messages = [type("M", (), {"value": r})() for r in records] + + def __iter__(self): + return iter(self._messages) + + fake.KafkaConsumer = KafkaConsumer # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "kafka", fake) + + +def test_stream_kafka_unknown_id_column_exits_1(tmp_path, monkeypatch, capsys): + _fake_kafka(monkeypatch) + out = tmp_path / "out.csv" + + rc = main(["stream-kafka", "--topic", "events", "--bootstrap-servers", "localhost:9092", + "--batch-size", "100", "--max-batches", "2", + "--id-columns", "customer", "-o", str(out), "--quiet"]) + + assert rc == 1 + err = capsys.readouterr().err + assert err.startswith("freshdata: error: ") + assert "--id-columns column 'customer' not found" in err + assert "(did you mean 'customer_id'?)" in err + _assert_no_output(out) + + +def test_stream_kafka_valid_columns_still_work(tmp_path, monkeypatch): + _fake_kafka(monkeypatch) + out = tmp_path / "out.csv" + + rc = main(["stream-kafka", "--topic", "events", "--bootstrap-servers", "localhost:9092", + "--batch-size", "100", "--max-batches", "2", + "--id-columns", "customer_id", "--target-column", "amount", + "-o", str(out), "--quiet"]) + + assert rc == 0 + assert len(pd.read_csv(out)) == 200