Backfill Definitions¶

def get_backfill_key(*, validate: bool = True) -> str:
    """Return the raw backfill key for the current job run.

    Reads the ``backfill_key`` job parameter and optionally validates
    it against the job's `BackfillDef` boundaries.

    For time-based backfills the key is an ISO-8601 date/time string;
    for `StaticBackfill` it is one of the declared keys (e.g.
    ``"us"``, ``"eu"``).

    Parameters
    ----------
    validate:
        When ``True`` (the default), verify that the key is valid for
        the job's `BackfillDef`.  A `ValueError` is raised if the
        key is out of range.  Ignored when the job has no backfill
        definition.

    Raises
    ------
    RuntimeError
        If ``backfill_key`` is missing or empty.  This indicates the
        job has no backfill definition and was not started via the
        backfill CLI.
    ValueError
        If *validate* is ``True`` and the backfill key is outside the
        backfill definition's boundaries.

    Returns
    -------
    str
        The raw backfill key string.
    """
    from databricks_bundle_decorators.context import params

    raw = params.get(BACKFILL_KEY_PARAM, "")
    if not raw:
        raise RuntimeError(
            "backfill_key is not set. "
            "This usually means the job was not invoked with a "
            "backfill_key parameter. Use @job(backfill=...) and "
            "the backfill CLI, or pass backfill_key explicitly."
        )
    if validate:
        job_name: str | None = params.get("__job_name__")
        _validate_backfill_key(raw, job_name)

    return raw

def get_run_logical_date(*, validate: bool = True) -> datetime:
    """Return the backfill key parsed as a timezone-aware ``datetime``.

    Convenience wrapper around `get_backfill_key` for time-based
    backfills (`DailyBackfill`, `WeeklyBackfill`, etc.).  Not
    suitable for `StaticBackfill` with non-date keys — use
    `get_backfill_key` instead.

    Parameters
    ----------
    validate:
        When ``True`` (the default), verify that the key is valid for
        the job's `BackfillDef`.  A `ValueError` is raised if the
        key is out of range.  Ignored when the job has no backfill
        definition.

    Raises
    ------
    RuntimeError
        If ``backfill_key`` is missing or empty.
    ValueError
        If the key cannot be parsed as an ISO-8601 date/time, or if
        *validate* is ``True`` and it falls outside the backfill
        definition's boundaries.

    Returns
    -------
    datetime
        Timezone-aware datetime representing the backfill key.
    """
    raw = get_backfill_key(validate=validate)
    return _parse_logical_date_str(raw)

def all_partitions(proxy: TaskProxy) -> _AllPartitionsProxy:
    """Wrap a `TaskProxy` so the downstream task receives all partitions.

    Use inside a ``@job`` body to indicate that the downstream task
    should read the **entire** dataset from the upstream task, across
    all partitions, rather than filtering to the current ``backfill_key``.

    Parameters
    ----------
    proxy:
        A `TaskProxy` returned by calling a ``@task``-decorated
        function inside a ``@job`` body.

    Returns
    -------
    `_AllPartitionsProxy`
        A wrapped proxy that records the all-partitions flag on the
        dependency edge.

    Example
    -------
    ::

        @job(backfill=DailyBackfill(start_date="2024-01-01"))
        def my_pipeline():
            @task(io_manager=io)
            def extract():
                ...

            @task
            def aggregate(data):
                ...

            data = extract()
            aggregate(all_partitions(data))
    """
    if not isinstance(proxy, TaskProxy):
        raise TypeError(
            f"all_partitions() expects a TaskProxy returned by calling "
            f"a @task-decorated function inside a @job body, "
            f"got {type(proxy).__name__!r}."
        )
    return _AllPartitionsProxy(proxy.task_key)