Skip to content

Commit

Permalink
Allow for specificaion of "specific" sample windows
Browse files Browse the repository at this point in the history
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'}`.
  • Loading branch information
QMalcolm committed Jan 29, 2025
1 parent 825de61 commit b3b057d
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 2 deletions.
12 changes: 11 additions & 1 deletion core/dbt/cli/option_types.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand Down
56 changes: 55 additions & 1 deletion tests/unit/cli/test_option_types.py
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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

0 comments on commit b3b057d

Please sign in to comment.