From b3b057d1b59e6f8e13d69e189d5954b177833cde Mon Sep 17 00:00:00 2001 From: Quigley Malcolm Date: Wed, 29 Jan 2025 09:56:27 -0600 Subject: [PATCH] Allow for specificaion of "specific" sample windows In most cases people will want to set "relative" sample windows, i.e. "3 days" to sample the last three days. However, there are some cases where people will want to "specific" sample windows for some chunk of historic time, i.e. `{'start': '2024-01-01', 'end': '2024-01-31'}`. --- core/dbt/cli/option_types.py | 12 ++++++- tests/unit/cli/test_option_types.py | 56 ++++++++++++++++++++++++++++- 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/core/dbt/cli/option_types.py b/core/dbt/cli/option_types.py index b90980840c0..006b28c7c19 100644 --- a/core/dbt/cli/option_types.py +++ b/core/dbt/cli/option_types.py @@ -1,5 +1,6 @@ from typing import Optional +import pytz from click import Choice, Context, Parameter, ParamType from dbt.config.utils import normalize_warn_error_options, parse_cli_yaml_string @@ -104,7 +105,16 @@ def convert( if isinstance(value, str): try: - return SampleWindow.from_relative_string(value) + # Try and identify if it's a "dict" or a "str" + if value.lstrip()[0] == "{": + param_option_name: str = param.opts[0] if param.opts else param.name # type: ignore + parsed_dict = parse_cli_yaml_string(value, param_option_name.strip("-")) + sample_window = SampleWindow.from_dict(parsed_dict) + sample_window.start = sample_window.start.replace(tzinfo=pytz.UTC) + sample_window.end = sample_window.end.replace(tzinfo=pytz.UTC) + return sample_window + else: + return SampleWindow.from_relative_string(value) except Exception as e: self.fail(e.__str__(), param, ctx) else: diff --git a/tests/unit/cli/test_option_types.py b/tests/unit/cli/test_option_types.py index 1067f64a3c3..bf43df8d13a 100644 --- a/tests/unit/cli/test_option_types.py +++ b/tests/unit/cli/test_option_types.py @@ -1,7 +1,13 @@ +from datetime import datetime +from typing import Union + +import freezegun import pytest +import pytz from click import BadParameter, Option -from dbt.cli.option_types import YAML +from dbt.cli.option_types import YAML, SampleWindowType +from dbt.event_time.sample_window import SampleWindow class TestYAML: @@ -24,3 +30,51 @@ def test_yaml_init_invalid_yaml_str(self, invalid_yaml_str): with pytest.raises(BadParameter) as e: YAML().convert(invalid_yaml_str, Option(["--vars"]), None) assert "--vars" in e.value.format_message() + + +class TestSampleWindowType: + @pytest.mark.parametrize( + "input,expected_result", + [ + ( + "{'start': '2025-01-24', 'end': '2025-01-27'}", + SampleWindow( + start=datetime(2025, 1, 24, 0, 0, 0, 0, pytz.UTC), + end=datetime(2025, 1, 27, 0, 0, 0, 0, pytz.UTC), + ), + ), + ( + "{'tart': '2025-01-24', 'bend': '2025-01-27'}", + BadParameter('Field "start" of type datetime is missing in SampleWindow instance'), + ), + ( + "{}", + BadParameter('Field "start" of type datetime is missing in SampleWindow instance'), + ), + ( + "cats", + BadParameter( + "Runtime Error\n Cannot load SAMPLE_WINDOW from 'cats'. Must be of form 'DAYS_INT GRAIN_SIZE'." + ), + ), + ], + ) + def test_convert(self, input: str, expected_result: Union[SampleWindow, Exception]): + try: + result = SampleWindowType().convert(input, Option(["--sample-window"]), None) + assert result == expected_result + except Exception as e: + assert str(e) == str(expected_result) + + # this had to be a seprate test case because the @freezegun.freeze_time + # was screwing up the instantiation of SampleWindow.from_dict calls for the + # other test cases + @freezegun.freeze_time("2025-01-28T02:03:0Z") + def test_convert_relative(self): + input = "3 days" + expected_result = SampleWindow( + start=datetime(2025, 1, 25, 2, 3, 0, 0, pytz.UTC), + end=datetime(2025, 1, 28, 2, 3, 0, 0, pytz.UTC), + ) + result = SampleWindowType().convert(input, Option(["--sample-window"]), None) + assert result == expected_result