Python API¶

def get_python_source(self) -> str:
    """Return the source shipped to the environment when ``python_callable`` is a function.

    The runner executes it and calls the name of ``python_callable`` in the resulting namespace. Task
    decorators are removed, including the one named by ``custom_operator_name`` (``@task.<name>`` of a
    decorator built on ``@task.pixi``) and this provider's decorator functions used under their own names
    (``@pixi_task(...)``), and line numbers kept, so tracebacks point at the DAG file. Override it to wrap
    the function, for example by appending code that rebinds the name to a wrapper, as for
    ``@task.virtualenv``.
    """
    stripped = {*STRIPPED_DECORATORS, *_DECORATOR_FUNCTIONS}
    operator_name = getattr(self, "custom_operator_name", "")
    if operator_name.startswith("@"):
        stripped.add(operator_name[1:])
    return function_source(self.python_callable, stripped)[0]  # type: ignore[arg-type]

def callable_input(self, context: Mapping[str, Any] | None = None) -> bytes:
    """Return the runner's input: the callable, its arguments and the task context, in the ``serializer``'s format.

    :param context: the task context, of which the runner passes the JSON-safe part
        ([`serializable_context`][airflow.providers.pixi.utils.context.serializable_context]) to parameters named
        after its keys.
    """
    if isinstance(self.python_callable, str):
        module, _, name = self.python_callable.partition(":")
        spec: dict[str, Any] = {"module": module, "name": name}
    else:
        fn = inspect.unwrap(self.python_callable)
        spec = {
            "source": self.get_python_source(),
            "filename": inspect.getsourcefile(fn) or "<pixi task>",
            "name": fn.__name__,
        }
    spec.update(args=list(self.op_args), kwargs=dict(self.op_kwargs))
    if context_values := serializable_context(context):
        spec["context"] = context_values
    if self.serializer == "pickle":
        return pickle.dumps(spec, protocol=PICKLE_PROTOCOL)
    try:
        return json.dumps(spec).encode()
    except TypeError as e:
        raise TypeError(f"op_args and op_kwargs must be JSON-serializable, or pass serializer='pickle': {e}") from e

def add_pypi_dependencies(self, *requirements: str) -> None:
    """Add pip requirement strings, such as ``"mytracker>=1"``, to the PyPI dependencies of the inline manifest.

    For subclasses and mixins that need a package in the environment, whichever form the DAG gave
    ``pypi_dependencies`` in: a dict gets the ``[pypi-dependencies]`` entries of ``requirements``, pip
    strings become a list with ``requirements`` at the end. ``pypi_dependencies`` is assigned a new value,
    so a value saved before stays as it was and can be restored afterwards.

    :raises AirflowException: if the environment comes from ``pixi_project_path`` or ``pixi_toml_path``,
        whose manifest has to list the packages itself.
    :raises ValueError: for a requirement ``pypi_dependencies`` does not accept, or a package listed twice.
    """
    if not self.inline_manifest:
        raise AirflowException(
            self._project_pypi_error(list(requirements), self.pixi_project_path or self.pixi_toml_path)
        )
    added = pypi_dependencies_table(list(requirements))
    current = self.pypi_dependencies
    if isinstance(current, dict):
        listed = {canonicalize_name(name) for name in current}
        for name in added:
            if canonicalize_name(name) in listed:
                raise ValueError(f"{name} is listed twice in pypi_dependencies")
        self.pypi_dependencies = {**current, **added}
        return
    extended = [*([current] if isinstance(current, str) else current or []), *requirements]
    if not _needs_rendering(extended):
        pypi_dependencies_table(extended)
    self.pypi_dependencies = extended

def inline_manifest_toml(self) -> str:
    """Return the ``pixi.toml`` of an inline manifest, from the rendered arguments."""
    pypi_dependencies = pypi_dependencies_table(self.pypi_dependencies)
    dependencies = conda_dependencies(self.dependencies)
    if pypi_dependencies and "python" not in dependencies:
        dependencies["python"] = f"{sys.version_info.major}.{sys.version_info.minor}.*"
    return build_pixi_toml(
        channels=self.channels or DEFAULT_CHANNELS,
        platforms=self.platforms or self.default_platforms(),
        dependencies=dependencies,
        pypi_dependencies=pypi_dependencies,
    )

def default_platforms(self) -> list[str]:
    """Return the platforms of an inline manifest without ``platforms``: the one of the machine running pixi.

    Called when the task runs. Pixi solves the environment for every platform of the manifest, and a package
    missing on one of them fails the solve, so only the platform that runs it is listed.
    """
    return [local_platform()]

def pixi_run_options(self) -> list[str]:
    """Return the options of ``pixi run`` that follow ``--manifest-path``: the environment and the lock mode.

    Called when the task runs, as templated fields are then rendered.
    """
    environment = self._checked_environment(AirflowException)
    options = ["--environment", environment] if environment else []
    lock_mode = self._checked_lock_mode(AirflowException)
    if lock_mode:
        options.append(f"--{lock_mode}")
    return options

def pixi_run_command(self, pixi: str, manifest: str) -> list[str]:
    """Return ``pixi run`` with the manifest and [`pixi_run_options`][pixi_run_options].

    The command to run in the environment goes after it.
    """
    return [pixi, "run", "--manifest-path", manifest, *self.pixi_run_options()]

@contextlib.contextmanager
def local_manifest(self) -> Iterator[tuple[str, str]]:
    """Yield the manifest to pass to pixi on the worker and the directory to run in.

    An inline manifest is written to ``env_cache_path`` or to a temporary directory, removed afterwards
    unless ``cleanup_temp_manifest=False``. Runs at execution time, so templated paths are rendered. Relative
    paths are resolved against the directory of the DAG file.
    """
    if self._held_manifest is not None:
        yield self._held_manifest
        return
    if self.pixi_project_path is not None:
        project = self._local_path(self.pixi_project_path)
        self._check_project_manifest(project)
        yield project, project
        return
    if self.pixi_toml_path is not None:
        manifest = self._local_path(self.pixi_toml_path)
        self._check_project_manifest(manifest)
        yield manifest, os.path.dirname(manifest)
        return
    toml = self.inline_manifest_toml()
    if self.env_cache_path:
        digest = hashlib.sha256(toml.encode()).hexdigest()[:16]
        env_dir = os.path.join(self._local_path(self.env_cache_path), f"pixi-{digest}")
        os.makedirs(env_dir, exist_ok=True)
        manifest = os.path.join(env_dir, "pixi.toml")
        if not os.path.exists(manifest):
            _write_atomic(manifest, toml)
        yield manifest, env_dir
        return
    env_dir = tempfile.mkdtemp(prefix="airflow_pixi_")
    try:
        manifest = os.path.join(env_dir, "pixi.toml")
        Path(manifest).write_text(toml)
        yield manifest, env_dir
    finally:
        if self.cleanup_temp_manifest:
            shutil.rmtree(env_dir, ignore_errors=True)

@contextlib.contextmanager
def holding_manifest(self) -> Iterator[None]:
    """Keep one [`local_manifest`][local_manifest] for every use inside, such as the pokes of one sensor run."""
    with self.local_manifest() as held:
        self._held_manifest = held
        try:
            yield
        finally:
            self._held_manifest = None

def run_callable(self: Any, context: Mapping[str, Any] | None = None) -> Any:
    """Call ``python_callable`` in the environment and return its result.

    :param context: the task context, for the callable's parameters named after its keys.
    """
    pixi = resolve_pixi(self.pixi_binary)
    with (
        self.local_manifest() as (manifest, cwd),
        tempfile.TemporaryDirectory(prefix="airflow_pixi_io_") as io,
        self.pixi_run_env(os.environ, self.env_vars) as env,
    ):
        env.setdefault("PYTHONUNBUFFERED", "1")  # the callable's output reaches the log as it prints
        input_path, output_path = os.path.join(io, "input"), os.path.join(io, "output")
        Path(input_path).write_bytes(self.callable_input(context))
        command = self.pixi_run_command(pixi, manifest)
        self.log.info("Running %s in %s", " ".join(command), cwd)
        returncode, tail = self._run(
            [*command, "python", "-c", RUNNER_SCRIPT, self.serializer, input_path, output_path], cwd, env
        )
        if returncode in self.skip_on_exit_code:
            raise AirflowSkipException(f"pixi run exited with code {returncode}, which skips the task")
        if returncode != 0:
            error_path = Path(output_path + ERROR_SUFFIX)
            if error_path.exists():
                raise self._callable_error(json.loads(error_path.read_text()))
            raise AirflowException(f"pixi run exited with code {returncode}:\n" + "\n".join(tail))
        if not os.path.exists(output_path):
            raise self._missing_result()
        if self.serializer == "pickle":
            return pickle.loads(Path(output_path).read_bytes())
        return json.loads(Path(output_path).read_text())

@contextlib.contextmanager
def pixi_run_env(
    self, base: Mapping[str, str], overrides: Mapping[str, str] | None = None
) -> Iterator[dict[str, str]]:
    """Yield the environment of a pixi run, valid until the block exits.

    In increasing precedence: ``base`` (usually the worker's environment) without the worker's Python settings
    ([`WORKER_PYTHON_VARIABLES`][WORKER_PYTHON_VARIABLES]) and with ``PYTHONNOUSERSITE=1``, so the environment's
    Python imports only what the Pixi environment has; the cache directory Variables; ``env_from_variables`` and
    ``env_from_connections``; then ``overrides`` (such as ``env_vars``), which can set those Python variables again.
    With ``pixi_conn_id``, ``RATTLER_AUTH_FILE`` and ``NETRC`` point at temporary files with the credentials, merged
    with the files these variables named before (see
    [`pixi_auth_env`][airflow.providers.pixi.hooks.pixi.pixi_auth_env]), and removed when the block exits.
    """
    env = {name: value for name, value in base.items() if name not in WORKER_PYTHON_VARIABLES}
    env["PYTHONNOUSERSITE"] = "1"
    env.update(
        resolve_env(
            optional_variables={
                "PIXI_CACHE_DIR": self.pixi_cache_dir_variable,
                "UV_CACHE_DIR": self.uv_cache_dir_variable,
                "PIP_CACHE_DIR": self.pip_cache_dir_variable,
            },
            env_from_variables=self.env_from_variables,
            env_from_connections=self.env_from_connections,
        )
    )
    env.update(overrides or {})
    with pixi_auth_env(self.pixi_conn_id, env) as auth:
        env.update(auth)
        yield env

def pixi_command(self, pixi: str, manifest: str) -> list[str]:
    """Return the command that runs in place of ``bash_command``: ``pixi run ... bash -c <bash_command>``.

    Called when the task runs, as templated fields are then rendered. Override it to run something else
    through pixi, as [`PixiTaskOperator`][airflow.providers.pixi.operators.task.PixiTaskOperator] does.
    """
    return [*self.pixi_run_command(pixi, manifest), "bash", "-c", str(self.bash_command)]

def get_env(self, context: Any) -> dict[str, str]:
    """Return the environment of the run: ``BashOperator``'s, with the variables of this operator merged in.

    Called by ``execute``; see the class docstring for the precedence. Files written for ``pixi_conn_id``
    exist until ``execute`` returns.
    """
    if self._env_stack is None:
        raise AirflowException(f"{type(self).__name__}.get_env can only be called from execute")
    env = super().get_env(context)
    overrides = {**(self.env or {}), **(self.env_vars or {})}
    return self._env_stack.enter_context(self.pixi_run_env(env, overrides))

def pixi_command(self, pixi: str, manifest: str) -> list[str]:
    """Return ``pixi run ... <task> <task_args...>``, from the rendered ``task`` and ``task_args``."""
    task = str(self.task).strip()
    if not task:
        raise AirflowException("task rendered to an empty name")
    if not isinstance(self.task_args, (list, tuple)):
        raise AirflowException(f"task_args must render to a list, not {type(self.task_args).__name__}")
    return [*self.pixi_run_command(pixi, manifest), task, *map(str, self.task_args)]

def default_platforms(self) -> list[str]:
    """Return [`DEFAULT_POD_PLATFORMS`][DEFAULT_POD_PLATFORMS]: pods run on Linux nodes, not on the worker."""
    return list(DEFAULT_POD_PLATFORMS)

def pod_script(self) -> str:
    """Return the shell script the pod runs: check pixi, write the inputs, run the callable."""
    min_version = shlex.quote(str(MIN_PIXI_VERSION))
    pixi = shlex.quote(self.pixi_binary)
    lines = [
        "set -eu",
        f'version="$({pixi} --version)"',
        'version="${version##* }"',
        f'if [ "$(printf "%s\\n%s\\n" {min_version} "$version" | sort -V | head -n 1)" != {min_version} ]; then',
        f'  echo "pixi $version in the image is older than {MIN_PIXI_VERSION}, which this provider needs" >&2',
        "  exit 1",
        "fi",
        'work="$(mktemp -d)"',
        f'printf %s "${_INPUT_ENV}" | base64 -d > "$work/input"',
    ]
    if self.pixi_project_path is not None:
        self._check_project_manifest(self.pixi_project_path)
        lines += [f"manifest={_pod_path(self.pixi_project_path)}", 'cd "$manifest"']
    elif self.pixi_toml_path is not None:
        self._check_project_manifest(self.pixi_toml_path)
        directory = posixpath.dirname(self.pixi_toml_path) or "."
        lines += [f"manifest={_pod_path(self.pixi_toml_path)}", f"cd {shlex.quote(directory)}"]
    else:
        if self.env_cache_path:
            digest = hashlib.sha256(self.inline_manifest_toml().encode()).hexdigest()[:16]
            lines.append(f"dir={_pod_path(posixpath.join(self.env_cache_path, f'pixi-{digest}'))}")
        else:
            lines.append('dir="$work/env"')
        lines += [
            'mkdir -p "$dir"',
            'if [ ! -f "$dir/pixi.toml" ]; then',
            f'  printf %s "${_MANIFEST_ENV}" | base64 -d > "$dir/.pixi.toml.$$"',
            '  mv "$dir/.pixi.toml.$$" "$dir/pixi.toml"',
            "fi",
            'manifest="$dir/pixi.toml"',
            'cd "$dir"',
        ]
    options = "".join(f" {shlex.quote(option)}" for option in self.pixi_run_options())
    run = f'{pixi} run --manifest-path "$manifest"{options} python -c "${_RUNNER_ENV}"'
    lines += [
        f'{run} {shlex.quote(self.serializer)} "$work/input" "$work/output" json',
        f'[ -f "$work/output" ] || {{ echo {shlex.quote(str(self._missing_result()))} >&2; exit 1; }}',
    ]
    if self.do_xcom_push:
        xcom_dir = shlex.quote(self.xcom_dir)
        lines.append(f'mkdir -p {xcom_dir} && cp "$work/output" {xcom_dir}/return.json')
    return "\n".join(lines) + "\n"

def pod_env_vars(self, context: Mapping[str, Any] | None = None) -> list[k8s.V1EnvVar]:
    """Return the environment variables that carry the callable, its arguments and an inline manifest.

    :param context: the task context, of which the JSON-safe part reaches the callable.
    :raises AirflowException: if the input or the manifest is longer than
        [`MAX_ENV_VALUE_BYTES`][MAX_ENV_VALUE_BYTES], so the task fails before the pod is created.
    """
    env = {
        _INPUT_ENV: base64.b64encode(self.callable_input(context)).decode(),
        _RUNNER_ENV: RUNNER_SCRIPT,
        TERMINATION_LOG_ENV: self.termination_message_path,
    }
    self._check_env_size(
        _INPUT_ENV,
        env[_INPUT_ENV],
        "The callable with its op_args, op_kwargs and the task context",
        "Pass smaller arguments, such as a path or URL instead of the data, or put the code in the image and use "
        "pixi_project_path with a 'module.path:callable_name' callable.",
    )
    if self.inline_manifest:
        env[_MANIFEST_ENV] = base64.b64encode(self.inline_manifest_toml().encode()).decode()
        self._check_env_size(
            _MANIFEST_ENV,
            env[_MANIFEST_ENV],
            "The inline manifest",
            "Put the environment in a project in the image and use pixi_project_path.",
        )
    return [k8s.V1EnvVar(name=name, value=value) for name, value in env.items()]