diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 460b6e520b..69ebd881da 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -66,6 +66,10 @@ Options: -e, --end TEXT The end datetime of the interval for which this command will be applied. --execution-time TEXT The execution time (defaults to now). + --time-zone TEXT IANA timezone for day-sized relative dates such as + '1 week ago', 'today', and 'yesterday' (for example, + America/Los_Angeles). Those values are converted to + UTC timestamps. Defaults to UTC. --help Show this message and exit. ``` @@ -83,6 +87,10 @@ Options: command will be applied. -e, --end TEXT The end datetime of the interval for which this command will be applied. + --time-zone TEXT IANA timezone for day-sized relative dates such as '1 + week ago', 'today', and 'yesterday' (for example, + America/Los_Angeles). Those values are converted to + UTC timestamps. Defaults to UTC. --help Show this message and exit. ``` @@ -205,6 +213,10 @@ Options: -e, --end TEXT The end datetime of the interval for which this command will be applied. --execution-time TEXT The execution time (defaults to now). + --time-zone TEXT IANA timezone for day-sized relative dates such as + '1 week ago', 'today', and 'yesterday' (for example, + America/Los_Angeles). Those values are converted to + UTC timestamps. Defaults to UTC. --limit INTEGER The number of rows which the query should be limited to. --help Show this message and exit. @@ -352,6 +364,12 @@ Options: -e, --end TEXT The end datetime of the interval for which this command will be applied. --execution-time TEXT The execution time (defaults to now). + --time-zone TEXT IANA timezone for day-sized relative dates + such as '1 week ago', 'today', and + 'yesterday' (for example, + America/Los_Angeles). Those values are + converted to UTC timestamps. Defaults to + UTC. --create-from TEXT The environment to create the target environment from if it doesn't exist. Default: prod. @@ -442,6 +460,10 @@ Options: -e, --end TEXT The end datetime of the interval for which this command will be applied. --execution-time TEXT The execution time (defaults to now). + --time-zone TEXT IANA timezone for day-sized relative dates such + as '1 week ago', 'today', and 'yesterday' (for + example, America/Los_Angeles). Those values are + converted to UTC timestamps. Defaults to UTC. --expand TEXT Whether or not to expand materialized models (defaults to False). If True, all referenced models are expanded as raw queries. Multiple @@ -508,6 +530,11 @@ Options: this command will be applied. -e, --end TEXT The end datetime of the interval for which this command will be applied. + --time-zone TEXT IANA timezone for day-sized relative dates + such as '1 week ago', 'today', and + 'yesterday' (for example, + America/Los_Angeles). Those values are + converted to UTC timestamps. Defaults to UTC. --skip-janitor Skip the janitor task. --ignore-cron Run for all missing intervals, ignoring individual cron schedules. diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index fc6cad8011..e05ed08897 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -4,6 +4,7 @@ import os import sys import typing as t +from datetime import datetime, tzinfo from pathlib import Path import click @@ -22,7 +23,12 @@ from sqlmesh.core.console import configure_console, get_console from sqlmesh.core.context import Context from sqlmesh.utils import Verbosity -from sqlmesh.utils.date import TimeLike +from sqlmesh.utils.date import ( + TimeLike, + is_categorical_relative_expression, + parse_time_zone, + to_datetime, +) from sqlmesh.utils.errors import MissingDependencyError, SQLMeshError logger = logging.getLogger(__name__) @@ -57,6 +63,57 @@ def parse_args(self, ctx: click.Context, args: t.List[str]) -> t.List[str]: return rest +def _resolve_day_sized( + value: t.Optional[TimeLike], + relative_tz: tzinfo, + relative_base: t.Optional[datetime] = None, +) -> t.Optional[TimeLike]: + if isinstance(value, str) and is_categorical_relative_expression(value): + return to_datetime(value, relative_base=relative_base, relative_tz=relative_tz) + return value + + +def _resolve_time_zone_args( + start: t.Optional[TimeLike], + end: t.Optional[TimeLike], + execution_time: t.Optional[TimeLike], + time_zone: t.Optional[str], +) -> t.Tuple[t.Optional[TimeLike], t.Optional[TimeLike], t.Optional[TimeLike]]: + """Turn day-sized relative bounds into UTC timestamps. + + Phrases such as "1 week ago", "today", and "yesterday" are anchored to midnight in + `time_zone`. Sub-day phrases and absolute timestamps are left unchanged so downstream + plan, render, and macro handling keep using one UTC timestamp. + """ + if not time_zone: + return start, end, execution_time + + relative_tz = parse_time_zone(time_zone) + if relative_tz is None: + return start, end, execution_time + + execution_time = _resolve_day_sized(execution_time, relative_tz) + relative_base = to_datetime(execution_time) if execution_time is not None else None + start = _resolve_day_sized(start, relative_tz, relative_base) + end = _resolve_day_sized(end, relative_tz, relative_base) + return start, end, execution_time + + +def _apply_time_zone_kwargs(kwargs: t.Dict[str, t.Any]) -> None: + start, end, execution_time = _resolve_time_zone_args( + kwargs.get("start"), + kwargs.get("end"), + kwargs.get("execution_time"), + kwargs.pop("time_zone", None), + ) + if "start" in kwargs: + kwargs["start"] = start + if "end" in kwargs: + kwargs["end"] = end + if "execution_time" in kwargs: + kwargs["execution_time"] = execution_time + + def _sqlmesh_version() -> str: try: from sqlmesh import __version__ @@ -283,6 +340,7 @@ def init( @opt.start_time @opt.end_time @opt.execution_time +@opt.time_zone @opt.expand @click.option( "--dialect", @@ -300,12 +358,16 @@ def render( start: TimeLike, end: TimeLike, execution_time: t.Optional[TimeLike] = None, + time_zone: t.Optional[str] = None, expand: t.Optional[t.Union[bool, t.Iterable[str]]] = None, dialect: t.Optional[str] = None, no_format: bool = False, **format_kwargs: t.Any, ) -> None: """Render a model's query, optionally expanding referenced models.""" + start, end, execution_time = _resolve_time_zone_args( # type: ignore[assignment] + start, end, execution_time, time_zone + ) model = ctx.obj.get_model(model, raise_if_missing=True) rendered = ctx.obj.render( @@ -338,6 +400,7 @@ def render( @opt.start_time @opt.end_time @opt.execution_time +@opt.time_zone @click.option( "--limit", type=int, @@ -352,9 +415,13 @@ def evaluate( start: TimeLike, end: TimeLike, execution_time: t.Optional[TimeLike] = None, + time_zone: t.Optional[str] = None, limit: t.Optional[int] = None, ) -> None: """Evaluate a model and return a dataframe with a default limit of 1000.""" + start, end, execution_time = _resolve_time_zone_args( # type: ignore[assignment] + start, end, execution_time, time_zone + ) df = ctx.obj.evaluate( model, start=start, @@ -422,6 +489,7 @@ def diff(ctx: click.Context, environment: t.Optional[str] = None) -> None: @opt.start_time @opt.end_time @opt.execution_time +@opt.time_zone @click.option( "--create-from", type=str, @@ -585,6 +653,7 @@ def plan( ) -> None: """Apply local changes to the target environment.""" context = ctx.obj + _apply_time_zone_kwargs(kwargs) restate_models = kwargs.pop("restate_model") or None select_models = kwargs.pop("select_model") or None allow_destructive_models = kwargs.pop("allow_destructive_model") or None @@ -609,6 +678,7 @@ def plan( @click.argument("environment", required=False) @opt.start_time @opt.end_time +@opt.time_zone @click.option("--skip-janitor", is_flag=True, help="Skip the janitor task.") @click.option( "--ignore-cron", @@ -637,6 +707,7 @@ def plan( def run(ctx: click.Context, environment: t.Optional[str] = None, **kwargs: t.Any) -> None: """Evaluate missing intervals for the target environment.""" context = ctx.obj + _apply_time_zone_kwargs(kwargs) select_models = kwargs.pop("select_model") or None completion_status = context.run(environment, select_models=select_models, **kwargs) if completion_status.is_failure: @@ -866,6 +937,7 @@ def test( @opt.start_time @opt.end_time @opt.execution_time +@opt.time_zone @click.pass_obj @error_handler @cli_analytics @@ -875,8 +947,12 @@ def audit( start: TimeLike, end: TimeLike, execution_time: t.Optional[TimeLike] = None, + time_zone: t.Optional[str] = None, ) -> None: """Run audits for the target model(s).""" + start, end, execution_time = _resolve_time_zone_args( # type: ignore[assignment] + start, end, execution_time, time_zone + ) if not obj.audit(models=models, start=start, end=end, execution_time=execution_time): exit(1) @@ -897,6 +973,7 @@ def audit( ) @opt.start_time @opt.end_time +@opt.time_zone @click.pass_context @error_handler @cli_analytics @@ -907,8 +984,12 @@ def check_intervals( select_model: t.List[str], start: TimeLike, end: TimeLike, + time_zone: t.Optional[str] = None, ) -> None: """Show missing intervals in an environment, respecting signals.""" + start, end, _ = _resolve_time_zone_args( # type: ignore[assignment] + start, end, None, time_zone + ) context = ctx.obj context.console.show_intervals( context.check_intervals( diff --git a/sqlmesh/cli/options.py b/sqlmesh/cli/options.py index 2e4642eb0e..1acf7c072b 100644 --- a/sqlmesh/cli/options.py +++ b/sqlmesh/cli/options.py @@ -37,6 +37,15 @@ help="The execution time (defaults to now).", ) +time_zone = click.option( + "--time-zone", + help=( + "IANA timezone for day-sized relative dates such as '1 week ago', 'today', and " + "'yesterday' (for example, America/Los_Angeles). Those values are converted to UTC " + "timestamps. Defaults to UTC." + ), +) + expand = click.option( "--expand", multiple=True, diff --git a/sqlmesh/utils/date.py b/sqlmesh/utils/date.py index f8df65c352..8c1d34cbf5 100644 --- a/sqlmesh/utils/date.py +++ b/sqlmesh/utils/date.py @@ -4,6 +4,7 @@ import time import typing as t import warnings +import zoneinfo from datetime import date, datetime, timedelta, timezone, tzinfo @@ -13,6 +14,7 @@ from sqlglot import exp from sqlmesh.utils import ttl_cache +from sqlmesh.utils.errors import ConfigError if t.TYPE_CHECKING: import pandas as pd @@ -46,6 +48,36 @@ } +def parse_time_zone(tz: t.Optional[str]) -> t.Optional[zoneinfo.ZoneInfo]: + """Parse an IANA timezone name. None means UTC/default behavior.""" + if not tz or tz == "UTC": + return None + + try: + return zoneinfo.ZoneInfo(tz) + except Exception as e: + available_timezones = zoneinfo.available_timezones() + + if available_timezones: + raise ConfigError(f"'{tz}' is not a valid IANA timezone name.") + raise ConfigError( + f"{e}. IANA time zone data is not available on your system. " + "`pip install tzdata` to leverage time zones." + ) + + +def _floor_to_local_midnight(dt: datetime, local_tz: tzinfo) -> datetime: + local_dt = dt.astimezone(local_tz) + return local_dt.replace(hour=0, minute=0, second=0, microsecond=0) + + +def _timezone_name(tzinfo_: tzinfo) -> str: + key = getattr(tzinfo_, "key", None) + if not isinstance(key, str) or not key: + raise ConfigError(f"Timezone {tzinfo_!r} has no IANA name.") + return key + + def now(minute_floor: bool = True) -> datetime: """ Current utc datetime with optional minute level accuracy / granularity. @@ -149,6 +181,7 @@ def to_datetime( relative_base: t.Optional[datetime] = None, check_categorical_relative_expression: bool = True, tz: t.Optional[tzinfo] = None, + relative_tz: t.Optional[tzinfo] = None, ) -> datetime: """Converts a value into a UTC datetime object. @@ -157,6 +190,10 @@ def to_datetime( relative_base: The datetime to reference for time expressions that are using relative terms. check_categorical_relative_expression: If True, takes into account the relative expressions that are categorical. tz: Timezone to convert datetime to, defaults to utc + relative_tz: Timezone used when `value` is a day-sized relative phrase such as + "1 week ago", "today", or "yesterday". The phrase is anchored to midnight in + this timezone, then converted to UTC. Sub-day phrases and absolute timestamps + ignore this argument. Raises: ValueError if value cannot be converted to a datetime. @@ -169,7 +206,13 @@ def to_datetime( elif isinstance(value, date): dt = datetime(value.year, value.month, value.day) elif isinstance(value, exp.Expr): - return to_datetime(value.name) + return to_datetime( + value.name, + relative_base=relative_base, + check_categorical_relative_expression=check_categorical_relative_expression, + tz=tz, + relative_tz=relative_tz, + ) else: try: epoch = float(value) @@ -179,17 +222,38 @@ def to_datetime( if epoch is None: relative_base = relative_base or now() expression = str(value) - if check_categorical_relative_expression and is_categorical_relative_expression( + parse_timezone = "UTC" + parse_relative_base: datetime = relative_base + use_relative_tz = bool( + relative_tz + and check_categorical_relative_expression + and is_categorical_relative_expression(expression) + ) + + if use_relative_tz: + assert relative_tz is not None + parse_timezone = _timezone_name(relative_tz) + parse_relative_base = _floor_to_local_midnight(relative_base, relative_tz) + elif check_categorical_relative_expression and is_categorical_relative_expression( expression ): - relative_base = relative_base.replace(hour=0, minute=0, second=0, microsecond=0) + parse_relative_base = relative_base.replace( + hour=0, minute=0, second=0, microsecond=0 + ) # note: we hardcode TIMEZONE: UTC to work around this bug: https://github.com/scrapinghub/dateparser/issues/896 - # where dateparser just silently fails if it cant interpret the contents of /etc/localtime - # this works because SQLMesh only deals with UTC, there is no concept of user local time + # where dateparser just silently fails if it cant interpret the contents of /etc/localtime. + # Day-sized phrases may instead use relative_tz; the result is still converted to UTC. dt = dateparser.parse( - expression, settings={"RELATIVE_BASE": relative_base, "TIMEZONE": "UTC"} + expression, + settings={"RELATIVE_BASE": parse_relative_base, "TIMEZONE": parse_timezone}, ) + + if use_relative_tz and dt is not None: + assert relative_tz is not None + if not dt.tzinfo: + dt = dt.replace(tzinfo=relative_tz) + dt = dt.astimezone(UTC) else: try: dt = datetime.strptime(str(value), DATE_INT_FMT) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index f1540727b1..5c67552513 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -518,6 +518,127 @@ def test_plan_dev_start_date(runner, tmp_path): assert "sqlmesh_example__dev.incremental_model: [2023-01-01" in result.output +@time_machine.travel("2023-01-20 03:00:00 UTC", tick=False) +def test_apply_time_zone_kwargs_converts_day_sized_bounds_to_utc(): + from sqlmesh.cli.main import _apply_time_zone_kwargs + + kwargs = { + "start": "today", + "end": "2 hours ago", + "execution_time": "2023-01-20 03:00:00", + "time_zone": "America/Los_Angeles", + } + _apply_time_zone_kwargs(kwargs) + + assert "time_zone" not in kwargs + assert kwargs["start"] == to_datetime("2023-01-19 08:00:00+00:00") + assert kwargs["execution_time"] == "2023-01-20 03:00:00" + assert kwargs["end"] == "2 hours ago" + + utc_kwargs = {"start": "today", "time_zone": "UTC"} + _apply_time_zone_kwargs(utc_kwargs) + assert utc_kwargs == {"start": "today"} + + +@time_machine.travel("2023-01-20 03:00:00 UTC", tick=False) +def test_render_today_with_time_zone_uses_local_midnight(runner, tmp_path): + create_example_project(tmp_path) + (tmp_path / "models" / "incremental_model.sql").write_text( + """MODEL ( + name sqlmesh_example.incremental_model, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_date + ), + start '2020-01-01', + cron '@daily', + grain (id, event_date) +); + +SELECT + id, + item_id, + event_date, + @start_tstz AS start_tstz, + @start_ds AS start_ds, +FROM sqlmesh_example.seed_model +WHERE event_date BETWEEN @start_date AND @end_date +""" + ) + + with_tz = runner.invoke( + cli, + [ + "--paths", + str(tmp_path), + "render", + "sqlmesh_example.incremental_model", + "--start", + "today", + "--end", + "today", + "--time-zone", + "America/Los_Angeles", + ], + ) + assert with_tz.exit_code == 0, with_tz.output + # 2023-01-20 03:00 UTC is still 2023-01-19 in Los Angeles, so local midnight is 08:00 UTC. + # @start_ds is the UTC date of that timestamp, not a separately localized calendar date. + assert "'2023-01-19 08:00:00+00:00'" in with_tz.output + assert "'2023-01-19'" in with_tz.output + + without_tz = runner.invoke( + cli, + [ + "--paths", + str(tmp_path), + "render", + "sqlmesh_example.incremental_model", + "--start", + "today", + "--end", + "today", + ], + ) + assert without_tz.exit_code == 0, without_tz.output + assert "'2023-01-20 00:00:00+00:00'" in without_tz.output + assert "'2023-01-20'" in without_tz.output + + +def test_plan_accepts_time_zone_option(runner, tmp_path): + create_example_project(tmp_path) + + result = runner.invoke( + cli, + ["--log-file-dir", tmp_path, "--paths", tmp_path, "plan", "--help"], + ) + assert result.exit_code == 0 + assert "--time-zone" in result.output + + +@time_machine.travel("2023-01-20 12:30:30 UTC", tick=False) +def test_plan_rejects_invalid_time_zone(runner, tmp_path): + create_example_project(tmp_path) + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "plan", + "dev", + "--start", + "1 week ago", + "--time-zone", + "Not/A_Timezone", + "--skip-tests", + "--no-prompts", + ], + ) + assert result.exit_code != 0 + + def test_plan_dev_end_date(runner, tmp_path): create_example_project(tmp_path) diff --git a/tests/utils/test_date.py b/tests/utils/test_date.py index c926507ccd..104432940f 100644 --- a/tests/utils/test_date.py +++ b/tests/utils/test_date.py @@ -6,6 +6,7 @@ from sqlglot import exp import pandas as pd # noqa: TID253 +from sqlmesh.utils.errors import ConfigError from sqlmesh.utils.date import ( UTC, TimeLike, @@ -15,6 +16,7 @@ is_relative, make_inclusive, make_ts_exclusive, + parse_time_zone, to_datetime, to_time_column, to_timestamp, @@ -77,10 +79,78 @@ def test_to_datetime_with_expressions(expression, result) -> None: assert to_datetime(expression) == result -def test_to_timestamp() -> None: +@time_machine.travel("2023-01-20 12:30:30 UTC", tick=False) +def test_to_datetime_with_relative_tz() -> None: + la_tz = parse_time_zone("America/Los_Angeles") + assert la_tz is not None + assert to_datetime("1 week ago") == datetime(2023, 1, 13, 0, 0, 0, tzinfo=UTC) + assert to_datetime("1 week ago", relative_tz=la_tz) == datetime( + 2023, 1, 13, 8, 0, 0, tzinfo=UTC + ) + assert to_datetime("today", relative_tz=la_tz) == datetime(2023, 1, 20, 8, 0, 0, tzinfo=UTC) + assert to_datetime("yesterday", relative_tz=la_tz) == datetime(2023, 1, 19, 8, 0, 0, tzinfo=UTC) + assert to_datetime("2020-01-01", relative_tz=la_tz) == datetime(2020, 1, 1, tzinfo=UTC) + + +@time_machine.travel("2023-01-20 15:00:00 UTC", tick=False) +def test_to_datetime_with_relative_tz_tokyo() -> None: + tokyo_tz = parse_time_zone("Asia/Tokyo") + assert to_datetime("1 day ago", relative_tz=tokyo_tz) == datetime( + 2023, 1, 19, 15, 0, 0, tzinfo=UTC + ) + + +@time_machine.travel("2023-01-20 15:00:00 UTC", tick=False) +def test_date_dict_uses_utc_timestamp_from_day_sized_relative() -> None: + tokyo_tz = parse_time_zone("Asia/Tokyo") + start = to_datetime("1 day ago", relative_tz=tokyo_tz) + assert start == datetime(2023, 1, 19, 15, 0, 0, tzinfo=UTC) + + resp = date_dict("2023-01-20 15:00:00", start, "2023-01-20") + assert resp["start_dt"] == start + assert resp["start_tstz"] == "2023-01-19 15:00:00+00:00" + assert resp["start_ds"] == "2023-01-19" + assert resp["start_date"] == date(2023, 1, 19) + assert resp["end_ds"] == "2023-01-20" + + +def test_parse_time_zone() -> None: + assert parse_time_zone(None) is None + assert parse_time_zone("UTC") is None + la_tz = parse_time_zone("America/Los_Angeles") + assert la_tz is not None + assert la_tz.key == "America/Los_Angeles" + + with pytest.raises(ConfigError, match="not a valid IANA timezone name"): + parse_time_zone("Not/A_Timezone") + + +def test_to_timestamp_from_date_string() -> None: assert to_timestamp("2020-01-01") == 1577836800000 +@time_machine.travel("2023-03-13 12:00:00 UTC", tick=False) +def test_to_datetime_relative_tz_spring_forward_dst() -> None: + la_tz = parse_time_zone("America/Los_Angeles") + assert to_datetime("1 day ago", relative_tz=la_tz) == datetime(2023, 3, 12, 8, 0, 0, tzinfo=UTC) + + +@time_machine.travel("2023-11-06 12:00:00 UTC", tick=False) +def test_to_datetime_relative_tz_fall_back_dst() -> None: + la_tz = parse_time_zone("America/Los_Angeles") + assert to_datetime("1 day ago", relative_tz=la_tz) == datetime(2023, 11, 5, 7, 0, 0, tzinfo=UTC) + + +@time_machine.travel("2023-01-20 12:30:00 UTC", tick=False) +def test_to_datetime_hour_relative_ignores_time_zone() -> None: + la_tz = parse_time_zone("America/Los_Angeles") + relative_base = to_datetime("2023-01-20 12:30:00") + without_tz = to_datetime("2 hours ago", relative_base=relative_base) + with_tz = to_datetime("2 hours ago", relative_base=relative_base, relative_tz=la_tz) + assert with_tz == without_tz + assert with_tz == datetime(2023, 1, 20, 10, 30, 0, tzinfo=UTC) + + @pytest.mark.parametrize( "start_in, end_in, start_out, end_out", [