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.

    When the parameter is missing or empty **and** the job has a
    time-based `BackfillDef`, the key is auto-derived from the
    current time (e.g. today's date for `DailyBackfill`) and a
    warning is logged.  This allows cron-triggered and file-arrival
    runs to work without explicitly supplying the key.

    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 and no automatic
        key can be derived (e.g. the job has no backfill definition,
        or uses `StaticBackfill`).
    ValueError
        If *validate* is ``True`` and the backfill key is outside the
        backfill definition's boundaries.

    Returns
    -------
    str
        The raw backfill key string.
    """
    raw = params.get(BACKFILL_KEY_PARAM, "")
    if not raw:
        raw = _auto_derive_backfill_key()
    if validate:
        job_name: str | None = params.get("__job_name__")
        _validate_backfill_key(raw, job_name)

    return raw

def get_backfill_keys(*, validate: bool = True) -> list[str]:
    """Return all backfill keys for the current run.

    When neither ``lookback`` nor ``collect_schedule_gaps`` is
    configured, this returns a single-element list equivalent to
    ``[get_backfill_key()]``.

    With ``lookback=N``, the result includes *N* prior keys plus the
    current key.  This applies in **all** run modes (scheduled and
    explicit backfill).

    With ``collect_schedule_gaps=True``, keys between the previous
    cron fire date and the current key are included.  This only
    applies to auto-derived keys (scheduled runs); when
    ``backfill_key`` is explicitly provided (via ``dbxdec backfill
    --keys``), schedule gap logic is bypassed.

    When both are configured the result is the sorted union.

    Parameters
    ----------
    validate:
        When ``True`` (the default), verify that the primary key is
        valid for the job's `BackfillDef`.

    Returns
    -------
    list[str]
        Sorted list of backfill key strings (ascending).
    """
    primary_key = get_backfill_key(validate=validate)

    job_name: str | None = params.get("__job_name__")
    if job_name is None:
        return [primary_key]

    job_meta = _JOB_REGISTRY.get(job_name)
    if job_meta is None or job_meta.backfill is None:
        return [primary_key]

    backfill = job_meta.backfill

    # StaticBackfill has no lookback/schedule_gaps support
    if isinstance(backfill, StaticBackfill):
        return [primary_key]

    # --exact flag: bypass all multi-key expansion
    if params.get(EXACT_BACKFILL_PARAM, "") == "1":
        return [primary_key]

    lookback: int = getattr(backfill, "lookback", 0)
    collect_gaps: bool = getattr(backfill, "collect_schedule_gaps", False)

    if lookback == 0 and not collect_gaps:
        return [primary_key]

    all_keys: set[str] = {primary_key}

    # Lookback: always applies
    if lookback > 0:
        all_keys.update(_compute_lookback_keys(backfill, primary_key, lookback))

    # Schedule gaps: only when key was auto-derived (not explicitly provided)
    if collect_gaps:
        explicitly_provided = bool(params.get(BACKFILL_KEY_PARAM, ""))
        if not explicitly_provided:
            schedule_cron = _get_job_schedule_cron(job_name)
            if schedule_cron is not None:
                all_keys.update(
                    _compute_schedule_gap_keys(backfill, primary_key, schedule_cron)
                )

    result = sorted(all_keys)
    if len(result) > 1:
        _logger.info(
            "Expanded backfill_key %r to %d keys (lookback=%d, collect_schedule_gaps=%s): %s",
            primary_key,
            len(result),
            lookback,
            collect_gaps,
            result,
        )
    return result

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)