jason810496 commented on code in PR #72369: URL: https://github.com/apache/airflow/pull/72369#discussion_r3949567547
########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG Review Comment: Should the `sdk.DAG` here be `LazyDeserializedDAG` instead? ########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class DagDefinition(ABC): + """Abstract base class for a DAG source definition.""" + + @property + @abstractmethod + def freshness_token(self) -> str: + """Opaque, generalized token representing the current state of the source.""" + + @abstractmethod + def get_relative_loc(self, root: Path | None = None) -> str: + """Get relative location of the definition to a root directory.""" + + @abstractmethod + def read_bytes(self) -> bytes: + """Read and return the content of the resource as bytes.""" + + def read_text(self, encoding: str = "utf-8") -> str: + """Read and return the content of the resource as a string.""" + return self.read_bytes().decode(encoding) + + @abstractmethod + def as_file(self) -> contextlib.AbstractContextManager[Path]: + """ + Return a context manager yielding a Path pointing to a local file. + + For file-backed resources, this is the actual file path. + For others, a temp file is created and cleaned up. + """ + + @abstractmethod + def __repr__(self) -> str: + """Return string representation used by import error and warning objects.""" + + +@dataclass +class FileDagDefinition(DagDefinition): + """A DAG definition backed by a file on the local filesystem.""" + + path: Path + + @property + def freshness_token(self) -> str: + try: + stat = self.path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return str(self.path) + try: + return str(self.path.relative_to(root)) + except ValueError: + return str(self.path) + + def read_bytes(self) -> bytes: + return self.path.read_bytes() + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + yield self.path + + def __repr__(self) -> str: + return str(self.path) + + +@dataclass +class DagImportError: + """Structured error information for DAG import failures.""" + + source_reference: str + message: str + error_type: str = "import" + line_number: int | None = None + column_number: int | None = None + context: str | None = None + suggestion: str | None = None + stacktrace: str | None = None + + def format_message(self) -> str: + """Format the error as a human-readable string.""" + parts = [f"Error in {self.source_reference}"] + if self.line_number is not None: + loc = f"line {self.line_number}" + if self.column_number is not None: + loc += f", column {self.column_number}" + parts.append(f"Location: {loc}") + parts.append(f"Error ({self.error_type}): {self.message}") + if self.context: + parts.append(f"Context:\n{self.context}") + if self.suggestion: + parts.append(f"Suggestion: {self.suggestion}") + return "\n".join(parts) + + +@dataclass +class DagImportWarning: + """Warning information for non-fatal issues during DAG import.""" + + source_reference: str + message: str + warning_type: str = "import" + line_number: int | None = None + context: dict[str, Any] | None = None + + +@dataclass +class DagImportResult: + """Result of importing DAGs from a definition.""" + + definition: DagDefinition | None = None + dags: list[DAG] = field(default_factory=list) + errors: list[DagImportError] = field(default_factory=list) + skipped_definitions: list[DagDefinition] = field(default_factory=list) + warnings: list[DagImportWarning] = field(default_factory=list) + dependencies: list[DagDefinition] = field(default_factory=list) + + +@dataclass +class DagSourceCode: + """Raw source code and its language identifier for a DAG definition.""" + + source_code: str + language: str + + +class AbstractDagImporter(ABC): + """Abstract base class for DAG importers.""" + + @property + @abstractmethod + def supported_extensions(self) -> list[str]: + """Return file extensions this importer handles (e.g., ['.py', '.zip']).""" + + @abstractmethod + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a DAG definition.""" + + def can_handle(self, definition: DagDefinition | str | Path) -> bool: + """Check if this importer can handle the given definition.""" + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) Review Comment: IIUC, there isn't `DagDefinition.path` attribute. ########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class DagDefinition(ABC): + """Abstract base class for a DAG source definition.""" + + @property + @abstractmethod + def freshness_token(self) -> str: + """Opaque, generalized token representing the current state of the source.""" + + @abstractmethod + def get_relative_loc(self, root: Path | None = None) -> str: + """Get relative location of the definition to a root directory.""" + + @abstractmethod + def read_bytes(self) -> bytes: + """Read and return the content of the resource as bytes.""" + + def read_text(self, encoding: str = "utf-8") -> str: + """Read and return the content of the resource as a string.""" + return self.read_bytes().decode(encoding) + + @abstractmethod + def as_file(self) -> contextlib.AbstractContextManager[Path]: + """ + Return a context manager yielding a Path pointing to a local file. + + For file-backed resources, this is the actual file path. + For others, a temp file is created and cleaned up. + """ + + @abstractmethod + def __repr__(self) -> str: + """Return string representation used by import error and warning objects.""" + + +@dataclass +class FileDagDefinition(DagDefinition): + """A DAG definition backed by a file on the local filesystem.""" + + path: Path + + @property + def freshness_token(self) -> str: + try: + stat = self.path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return str(self.path) + try: + return str(self.path.relative_to(root)) + except ValueError: + return str(self.path) + + def read_bytes(self) -> bytes: + return self.path.read_bytes() + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + yield self.path + + def __repr__(self) -> str: + return str(self.path) + + +@dataclass +class DagImportError: + """Structured error information for DAG import failures.""" + + source_reference: str + message: str + error_type: str = "import" + line_number: int | None = None + column_number: int | None = None + context: str | None = None + suggestion: str | None = None + stacktrace: str | None = None + + def format_message(self) -> str: + """Format the error as a human-readable string.""" + parts = [f"Error in {self.source_reference}"] + if self.line_number is not None: + loc = f"line {self.line_number}" + if self.column_number is not None: + loc += f", column {self.column_number}" + parts.append(f"Location: {loc}") + parts.append(f"Error ({self.error_type}): {self.message}") + if self.context: + parts.append(f"Context:\n{self.context}") + if self.suggestion: + parts.append(f"Suggestion: {self.suggestion}") + return "\n".join(parts) + + +@dataclass +class DagImportWarning: + """Warning information for non-fatal issues during DAG import.""" + + source_reference: str + message: str + warning_type: str = "import" + line_number: int | None = None + context: dict[str, Any] | None = None + + +@dataclass +class DagImportResult: + """Result of importing DAGs from a definition.""" + + definition: DagDefinition | None = None + dags: list[DAG] = field(default_factory=list) + errors: list[DagImportError] = field(default_factory=list) + skipped_definitions: list[DagDefinition] = field(default_factory=list) + warnings: list[DagImportWarning] = field(default_factory=list) + dependencies: list[DagDefinition] = field(default_factory=list) + + +@dataclass +class DagSourceCode: + """Raw source code and its language identifier for a DAG definition.""" + + source_code: str + language: str + + +class AbstractDagImporter(ABC): + """Abstract base class for DAG importers.""" + + @property + @abstractmethod + def supported_extensions(self) -> list[str]: + """Return file extensions this importer handles (e.g., ['.py', '.zip']).""" + + @abstractmethod + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a DAG definition.""" + + def can_handle(self, definition: DagDefinition | str | Path) -> bool: + """Check if this importer can handle the given definition.""" + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) + ) + return Path(path).suffix.lower() in self.supported_extensions if path else False + + def list_dag_definitions( + self, + bundle_name: str, + bundle_path: Path, + *, + safe_mode: bool = True, + ) -> Iterator[DagDefinition]: + """ + List DAG definitions in a bundle that this importer can handle. + + Override this method to customize definition discovery for your importer. + The default implementation finds files matching supported_extensions + and respects .airflowignore files. + """ + try: + from airflow.configuration import conf Review Comment: ```suggestion from airflow.sdk.configuration import conf ``` Additionally, the `ImportError`exception should be unreachable. ########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class DagDefinition(ABC): + """Abstract base class for a DAG source definition.""" + + @property + @abstractmethod + def freshness_token(self) -> str: + """Opaque, generalized token representing the current state of the source.""" + + @abstractmethod + def get_relative_loc(self, root: Path | None = None) -> str: + """Get relative location of the definition to a root directory.""" + + @abstractmethod + def read_bytes(self) -> bytes: + """Read and return the content of the resource as bytes.""" + + def read_text(self, encoding: str = "utf-8") -> str: + """Read and return the content of the resource as a string.""" + return self.read_bytes().decode(encoding) + + @abstractmethod + def as_file(self) -> contextlib.AbstractContextManager[Path]: + """ + Return a context manager yielding a Path pointing to a local file. + + For file-backed resources, this is the actual file path. + For others, a temp file is created and cleaned up. + """ + + @abstractmethod + def __repr__(self) -> str: + """Return string representation used by import error and warning objects.""" + + +@dataclass +class FileDagDefinition(DagDefinition): + """A DAG definition backed by a file on the local filesystem.""" + + path: Path + + @property + def freshness_token(self) -> str: + try: + stat = self.path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return str(self.path) + try: + return str(self.path.relative_to(root)) + except ValueError: + return str(self.path) + + def read_bytes(self) -> bytes: + return self.path.read_bytes() + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + yield self.path + + def __repr__(self) -> str: + return str(self.path) + + +@dataclass +class DagImportError: + """Structured error information for DAG import failures.""" + + source_reference: str + message: str + error_type: str = "import" + line_number: int | None = None + column_number: int | None = None + context: str | None = None + suggestion: str | None = None + stacktrace: str | None = None + + def format_message(self) -> str: + """Format the error as a human-readable string.""" + parts = [f"Error in {self.source_reference}"] + if self.line_number is not None: + loc = f"line {self.line_number}" + if self.column_number is not None: + loc += f", column {self.column_number}" + parts.append(f"Location: {loc}") + parts.append(f"Error ({self.error_type}): {self.message}") + if self.context: + parts.append(f"Context:\n{self.context}") + if self.suggestion: + parts.append(f"Suggestion: {self.suggestion}") + return "\n".join(parts) + + +@dataclass +class DagImportWarning: + """Warning information for non-fatal issues during DAG import.""" + + source_reference: str + message: str + warning_type: str = "import" + line_number: int | None = None + context: dict[str, Any] | None = None + + +@dataclass +class DagImportResult: + """Result of importing DAGs from a definition.""" + + definition: DagDefinition | None = None + dags: list[DAG] = field(default_factory=list) + errors: list[DagImportError] = field(default_factory=list) + skipped_definitions: list[DagDefinition] = field(default_factory=list) + warnings: list[DagImportWarning] = field(default_factory=list) + dependencies: list[DagDefinition] = field(default_factory=list) + + +@dataclass +class DagSourceCode: + """Raw source code and its language identifier for a DAG definition.""" + + source_code: str + language: str + + +class AbstractDagImporter(ABC): + """Abstract base class for DAG importers.""" + + @property + @abstractmethod + def supported_extensions(self) -> list[str]: + """Return file extensions this importer handles (e.g., ['.py', '.zip']).""" + + @abstractmethod + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a DAG definition.""" + + def can_handle(self, definition: DagDefinition | str | Path) -> bool: + """Check if this importer can handle the given definition.""" + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) + ) + return Path(path).suffix.lower() in self.supported_extensions if path else False + + def list_dag_definitions( + self, + bundle_name: str, + bundle_path: Path, + *, + safe_mode: bool = True, + ) -> Iterator[DagDefinition]: + """ + List DAG definitions in a bundle that this importer can handle. + + Override this method to customize definition discovery for your importer. + The default implementation finds files matching supported_extensions + and respects .airflowignore files. + """ + try: + from airflow.configuration import conf + + ignore_file_syntax = conf.get_mandatory_value("core", "DAG_IGNORE_FILE_SYNTAX", fallback="glob") + except ImportError: + ignore_file_syntax = "glob" + + supported_exts = [ext.lower() for ext in self.supported_extensions] + + for file_path in find_path_from_directory(bundle_path, ".airflowignore", ignore_file_syntax): + path = Path(file_path) + + if not path.is_file(): + continue + + # Check if this importer handles this file extension + if path.suffix.lower() not in supported_exts: + continue + + yield FileDagDefinition(path=path) + + @abstractmethod + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Retrieve the raw source code and its language identifier for the specified DAG definition.""" + + +class DagImporterRegistry: + """ + Registry for DAG importers. Singleton that manages importers by file extension. + + Each file extension can only be handled by one importer at a time. If multiple + importers claim the same extension, the last registered one wins and a warning + is logged. The built-in PythonDagImporter handles .py and ZipImporter handles .zip files. + """ + + _importers: dict[str, AbstractDagImporter] + + def __init__(self, register_defaults: bool = True): + self._importers = {} + if register_defaults: + self._register_default_importers() + + def _register_default_importers(self) -> None: + from airflow.sdk.importers.python_importer import PythonDagImporter + from airflow.sdk.importers.zip_importer import ZipImporter + + self.register(PythonDagImporter()) + self.register(ZipImporter()) + + def register(self, importer: AbstractDagImporter) -> None: + """ + Register an importer for its supported extensions. + + Each extension can only have one importer. If an extension is already registered, + the new importer will override it and a warning will be logged. + """ + for ext in importer.supported_extensions: + ext_lower = ext.lower() + if ext_lower in self._importers: + existing = self._importers[ext_lower] + log.warning( + "Extension '%s' already registered by %s, overriding with %s", + ext, + type(existing).__name__, + type(importer).__name__, + ) + self._importers[ext_lower] = importer + + def _get_suffix(self, definition: DagDefinition | str | Path) -> str | None: + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) + ) + return Path(path).suffix.lower() if path else None Review Comment: How about adding a common utility? Since this is almost identical with `can_handle`. ########## task-sdk/src/airflow/sdk/importers/python_importer.py: ########## @@ -0,0 +1,254 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Python DAG importer - imports DAGs from Python files.""" + +from __future__ import annotations + +import importlib.machinery +import importlib.util +import logging +import os +import sys +import traceback +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk.execution_time.timeout import timeout +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagImportWarning, + DagSourceCode, +) + +if TYPE_CHECKING: + from types import ModuleType + + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class PythonDagImporter(AbstractDagImporter): + """ + Importer for Python DAG files. + + This is the default importer registered with the DagImporterRegistry. It handles + .py files containing Python DAGs. + """ + + supported_extensions = [".py"] + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """ + Import DAGs from a Python DAG definition. + + :param definition: The definition to import from. + :param bundle_path: Root path of the DAG bundle. + :param bundle_name: Name of the DAG bundle. + :param safe_mode: If True, skip files that don't appear to contain DAGs. + :return: DagImportResult with imported DAGs and any errors. + """ + from airflow.sdk.definitions._internal.contextmanager import DagContext + + result = DagImportResult(definition=definition) + + # Clear any autoregistered dags from previous imports + DagContext.autoregistered_dags.clear() + + # Capture warnings during import + captured_warnings: list[warnings.WarningMessage] = [] + + try: + with warnings.catch_warnings(record=True) as captured_warnings: + with definition.as_file() as local_path: + filepath = os.fspath(local_path) + modules = self._load_modules_from_file( + filepath, + safe_mode, + result, + bundle_path=bundle_path, + bundle_name=bundle_name, + ) + except TypeError: + # Configuration errors (e.g., invalid timeout type) should propagate + raise + except Exception as e: + result.errors.append( + DagImportError( + source_reference=repr(definition), + message=str(e), + error_type="import", + stacktrace=traceback.format_exc(), + ) + ) + return result + + # Convert captured warnings to DagImportWarning + for warn_msg in captured_warnings: + category = warn_msg.category.__name__ + if (module := warn_msg.category.__module__) != "builtins": + category = f"{module}.{category}" + result.warnings.append( + DagImportWarning( + source_reference=repr(definition), + message=str(warn_msg.message), + warning_type=category, + line_number=warn_msg.lineno, + ) + ) + + # Process imported modules to extract DAGs + self._process_modules( + modules, + result, + bundle_path=bundle_path, + bundle_name=bundle_name, + ) + + return result + + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Retrieve the raw source code for the Python definition.""" + from airflow.sdk.importers.base import DagSourceCode + + return DagSourceCode( + source_code=definition.read_text(encoding="utf-8"), + language="python", + ) + + def _load_modules_from_file( + self, + filepath: str, + safe_mode: bool, + result: DagImportResult, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + ) -> list[ModuleType]: + from airflow import settings + from airflow.sdk._shared.module_loading.dag_file import get_unique_dag_module_name, might_contain_dag + from airflow.sdk.definitions._internal.contextmanager import DagContext + + definition = result.definition + + if not might_contain_dag(filepath, safe_mode): + log.debug("File %s assumed to contain no DAGs. Skipping.", filepath) + if definition is not None: + result.skipped_definitions.append(definition) + return [] + + log.debug("Importing %s (bundle: %s)", filepath, bundle_name) + mod_name = get_unique_dag_module_name(filepath) + + if mod_name in sys.modules: + del sys.modules[mod_name] + + DagContext.current_autoregister_module_name = mod_name + + def parse(mod_name: str, filepath: str) -> list[ModuleType]: + from airflow.configuration import conf + + try: + loader = importlib.machinery.SourceFileLoader(mod_name, filepath) + spec = importlib.util.spec_from_loader(mod_name, loader) + new_module = importlib.util.module_from_spec(spec) # type: ignore[arg-type] + sys.modules[spec.name] = new_module # type: ignore[union-attr] + loader.exec_module(new_module) + return [new_module] + except KeyboardInterrupt: + sys.modules.pop(mod_name, None) + raise + except BaseException as e: + sys.modules.pop(mod_name, None) + DagContext.autoregistered_dags.clear() + log.exception("Failed to import: %s", filepath) + if conf and conf.getboolean("core", "dagbag_import_error_tracebacks"): + stacktrace = traceback.format_exc( + limit=-conf.getint("core", "dagbag_import_error_traceback_depth") + ) Review Comment: Let's cache the value from the `conf` as private class property, since it will lookup the again for each `parse` call (if we're using custom backend, this means each external call per `parse`). ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: Review Comment: May I ask why do we need to create a temp file for each `as_file` ctx manager call? ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): + super().__init__() + self._internal_importers: dict[str, AbstractDagImporter] = {} + if internal_importers is None: + from airflow.sdk.importers.python_importer import PythonDagImporter + + self._internal_importers[".py"] = PythonDagImporter() + else: + from airflow.sdk._shared.module_loading import import_string + + for ext, cfg in internal_importers.items(): + ext_lower = ext if ext.startswith(".") else f".{ext}" + if isinstance(cfg, AbstractDagImporter): + self._internal_importers[ext_lower.lower()] = cfg + elif isinstance(cfg, dict) and "classpath" in cfg: + importer_class = import_string(cfg["classpath"]) + self._internal_importers[ext_lower.lower()] = importer_class(**cfg.get("kwargs", {})) + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a ZIP archive DAG definition by routing internal files.""" + result = DagImportResult(definition=definition) + + with definition.as_file() as local_zip_path: + try: + valid_members: dict[str, bytes] = {} + with zipfile.ZipFile(local_zip_path) as z: + for member_name in z.namelist(): + # ZipSlip prevention: check for directory traversal attempts + if ".." in Path(member_name).parts or Path(member_name).is_absolute(): + log.warning( + "Skipping zip member %r in %s: directory traversal patterns detected", + member_name, + local_zip_path, + ) + continue + suffix_lower = Path(member_name).suffix.lower() + if suffix_lower in self._internal_importers: + valid_members[member_name] = z.read(member_name) + except Exception as e: + result.errors.append( + DagImportError( + source_reference=definition.get_relative_loc(bundle_path), + message=f"Failed to read ZIP archive: {e}", + error_type="zip_read_error", + ) + ) + return result + + # Extract matching members into a single temp directory to avoid per-file churn + with tempfile.TemporaryDirectory(prefix="airflow_zip_") as temp_dir: + temp_dir_path = Path(temp_dir) + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + extracted_file.parent.mkdir(parents=True, exist_ok=True) + extracted_file.write_bytes(content) + + with _temporary_sys_path(str(temp_dir_path)): + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + suffix_lower = extracted_file.suffix.lower() + + nested_def = ZipFileDagDefinition( + zip_path=local_zip_path, + file_path=member_name, + _content=content, + _temp_path=extracted_file, + ) + + importer = self._internal_importers[suffix_lower] + if not importer.can_handle(nested_def): + continue + + member_result = importer.import_definition( + definition=nested_def, + bundle_path=bundle_path, + bundle_name=bundle_name, + safe_mode=safe_mode, + ) + + result.dags.extend(member_result.dags) + result.errors.extend(member_result.errors) + result.warnings.extend(member_result.warnings) + result.skipped_definitions.extend(member_result.skipped_definitions) + result.dependencies.extend(member_result.dependencies) + + return result + + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + if isinstance(definition, ZipFileDagDefinition): + suffix = Path(definition.file_path).suffix.lower() + if suffix in self._internal_importers: + return self._internal_importers[suffix].get_source_code(definition) + raise ValueError(f"No internal importer registered for zip member {definition.file_path}") + + # If definition is the zip archive itself, route to code member(s) + with definition.as_file() as local_zip_path: + with zipfile.ZipFile(local_zip_path) as z: + candidates = [ + name + for name in z.namelist() + if Path(name).suffix.lower() in self._internal_importers + and not name.startswith("__MACOSX") + and ".." not in Path(name).parts + and not Path(name).is_absolute() + ] + if not candidates: + raise ValueError(f"No code files found inside ZIP archive {definition}") + if len(candidates) == 1: + nested_def = ZipFileDagDefinition(zip_path=local_zip_path, file_path=candidates[0]) + importer = self._internal_importers[Path(candidates[0]).suffix.lower()] + return importer.get_source_code(nested_def) + + parts = [] + primary_language = "python" Review Comment: ```suggestion ``` primary_language will always be set in the loop. ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): + super().__init__() + self._internal_importers: dict[str, AbstractDagImporter] = {} + if internal_importers is None: + from airflow.sdk.importers.python_importer import PythonDagImporter + + self._internal_importers[".py"] = PythonDagImporter() + else: + from airflow.sdk._shared.module_loading import import_string + + for ext, cfg in internal_importers.items(): + ext_lower = ext if ext.startswith(".") else f".{ext}" + if isinstance(cfg, AbstractDagImporter): + self._internal_importers[ext_lower.lower()] = cfg + elif isinstance(cfg, dict) and "classpath" in cfg: + importer_class = import_string(cfg["classpath"]) + self._internal_importers[ext_lower.lower()] = importer_class(**cfg.get("kwargs", {})) + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a ZIP archive DAG definition by routing internal files.""" + result = DagImportResult(definition=definition) + + with definition.as_file() as local_zip_path: + try: + valid_members: dict[str, bytes] = {} + with zipfile.ZipFile(local_zip_path) as z: + for member_name in z.namelist(): + # ZipSlip prevention: check for directory traversal attempts + if ".." in Path(member_name).parts or Path(member_name).is_absolute(): + log.warning( + "Skipping zip member %r in %s: directory traversal patterns detected", + member_name, + local_zip_path, + ) + continue + suffix_lower = Path(member_name).suffix.lower() + if suffix_lower in self._internal_importers: + valid_members[member_name] = z.read(member_name) + except Exception as e: + result.errors.append( + DagImportError( + source_reference=definition.get_relative_loc(bundle_path), + message=f"Failed to read ZIP archive: {e}", + error_type="zip_read_error", + ) + ) + return result + + # Extract matching members into a single temp directory to avoid per-file churn + with tempfile.TemporaryDirectory(prefix="airflow_zip_") as temp_dir: + temp_dir_path = Path(temp_dir) + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + extracted_file.parent.mkdir(parents=True, exist_ok=True) + extracted_file.write_bytes(content) + + with _temporary_sys_path(str(temp_dir_path)): + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + suffix_lower = extracted_file.suffix.lower() + + nested_def = ZipFileDagDefinition( + zip_path=local_zip_path, + file_path=member_name, + _content=content, + _temp_path=extracted_file, + ) + + importer = self._internal_importers[suffix_lower] + if not importer.can_handle(nested_def): + continue Review Comment: Additionally, it seems we could move the validation forward to fail fast. ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): Review Comment: ```suggestion def __init__(self, internal_importers: dict[str, AbstractDagImporter | dict[str,str]] | None = None): ``` Btw, does `ZipImporter` will be construct with either way in prod code path? ########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class DagDefinition(ABC): + """Abstract base class for a DAG source definition.""" + + @property + @abstractmethod + def freshness_token(self) -> str: + """Opaque, generalized token representing the current state of the source.""" + + @abstractmethod + def get_relative_loc(self, root: Path | None = None) -> str: + """Get relative location of the definition to a root directory.""" + + @abstractmethod + def read_bytes(self) -> bytes: + """Read and return the content of the resource as bytes.""" + + def read_text(self, encoding: str = "utf-8") -> str: + """Read and return the content of the resource as a string.""" + return self.read_bytes().decode(encoding) + + @abstractmethod + def as_file(self) -> contextlib.AbstractContextManager[Path]: + """ + Return a context manager yielding a Path pointing to a local file. + + For file-backed resources, this is the actual file path. + For others, a temp file is created and cleaned up. + """ + + @abstractmethod + def __repr__(self) -> str: + """Return string representation used by import error and warning objects.""" + + +@dataclass +class FileDagDefinition(DagDefinition): + """A DAG definition backed by a file on the local filesystem.""" + + path: Path + + @property + def freshness_token(self) -> str: + try: + stat = self.path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return str(self.path) + try: + return str(self.path.relative_to(root)) + except ValueError: + return str(self.path) + + def read_bytes(self) -> bytes: + return self.path.read_bytes() + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + yield self.path + + def __repr__(self) -> str: + return str(self.path) + + +@dataclass +class DagImportError: + """Structured error information for DAG import failures.""" + + source_reference: str + message: str + error_type: str = "import" + line_number: int | None = None + column_number: int | None = None + context: str | None = None + suggestion: str | None = None + stacktrace: str | None = None + + def format_message(self) -> str: + """Format the error as a human-readable string.""" + parts = [f"Error in {self.source_reference}"] + if self.line_number is not None: + loc = f"line {self.line_number}" + if self.column_number is not None: + loc += f", column {self.column_number}" + parts.append(f"Location: {loc}") + parts.append(f"Error ({self.error_type}): {self.message}") + if self.context: + parts.append(f"Context:\n{self.context}") + if self.suggestion: + parts.append(f"Suggestion: {self.suggestion}") + return "\n".join(parts) + + +@dataclass +class DagImportWarning: + """Warning information for non-fatal issues during DAG import.""" + + source_reference: str + message: str + warning_type: str = "import" + line_number: int | None = None + context: dict[str, Any] | None = None + + +@dataclass +class DagImportResult: + """Result of importing DAGs from a definition.""" + + definition: DagDefinition | None = None + dags: list[DAG] = field(default_factory=list) + errors: list[DagImportError] = field(default_factory=list) + skipped_definitions: list[DagDefinition] = field(default_factory=list) + warnings: list[DagImportWarning] = field(default_factory=list) + dependencies: list[DagDefinition] = field(default_factory=list) Review Comment: IIUC, the `DagImportResult` _should_ match the `DagFileParsingResult`. https://github.com/apache/airflow/blob/b3b62fa8cf3c32cb45bb2815385a2af498300754/airflow-core/src/airflow/dag_processing/processor.py#L133-L145 ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): + super().__init__() + self._internal_importers: dict[str, AbstractDagImporter] = {} + if internal_importers is None: + from airflow.sdk.importers.python_importer import PythonDagImporter + + self._internal_importers[".py"] = PythonDagImporter() + else: + from airflow.sdk._shared.module_loading import import_string + + for ext, cfg in internal_importers.items(): + ext_lower = ext if ext.startswith(".") else f".{ext}" + if isinstance(cfg, AbstractDagImporter): + self._internal_importers[ext_lower.lower()] = cfg + elif isinstance(cfg, dict) and "classpath" in cfg: + importer_class = import_string(cfg["classpath"]) + self._internal_importers[ext_lower.lower()] = importer_class(**cfg.get("kwargs", {})) + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a ZIP archive DAG definition by routing internal files.""" + result = DagImportResult(definition=definition) + + with definition.as_file() as local_zip_path: Review Comment: - `with definition.as_file() as local_zip_path:` -> write a temp zip file - `valid_members[member_name] = z.read(member_name)` -> store all valid file content in memory - `extracted_file.write_bytes(content)` -> write a temp file again for each file content that we just stored in memory Would it be possible to improve the overall read path? ########## task-sdk/src/airflow/sdk/importers/python_importer.py: ########## @@ -0,0 +1,254 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Python DAG importer - imports DAGs from Python files.""" + +from __future__ import annotations + +import importlib.machinery +import importlib.util +import logging +import os +import sys +import traceback +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk.execution_time.timeout import timeout +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagImportWarning, + DagSourceCode, +) + +if TYPE_CHECKING: + from types import ModuleType + + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class PythonDagImporter(AbstractDagImporter): + """ + Importer for Python DAG files. + + This is the default importer registered with the DagImporterRegistry. It handles + .py files containing Python DAGs. + """ + + supported_extensions = [".py"] + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """ + Import DAGs from a Python DAG definition. + + :param definition: The definition to import from. + :param bundle_path: Root path of the DAG bundle. + :param bundle_name: Name of the DAG bundle. + :param safe_mode: If True, skip files that don't appear to contain DAGs. + :return: DagImportResult with imported DAGs and any errors. + """ + from airflow.sdk.definitions._internal.contextmanager import DagContext + + result = DagImportResult(definition=definition) + + # Clear any autoregistered dags from previous imports + DagContext.autoregistered_dags.clear() + + # Capture warnings during import + captured_warnings: list[warnings.WarningMessage] = [] + + try: + with warnings.catch_warnings(record=True) as captured_warnings: + with definition.as_file() as local_path: + filepath = os.fspath(local_path) + modules = self._load_modules_from_file( + filepath, + safe_mode, + result, + bundle_path=bundle_path, + bundle_name=bundle_name, + ) + except TypeError: + # Configuration errors (e.g., invalid timeout type) should propagate + raise + except Exception as e: + result.errors.append( + DagImportError( + source_reference=repr(definition), + message=str(e), + error_type="import", + stacktrace=traceback.format_exc(), + ) + ) + return result + + # Convert captured warnings to DagImportWarning + for warn_msg in captured_warnings: + category = warn_msg.category.__name__ + if (module := warn_msg.category.__module__) != "builtins": + category = f"{module}.{category}" + result.warnings.append( + DagImportWarning( + source_reference=repr(definition), + message=str(warn_msg.message), + warning_type=category, + line_number=warn_msg.lineno, + ) + ) + + # Process imported modules to extract DAGs + self._process_modules( + modules, + result, + bundle_path=bundle_path, + bundle_name=bundle_name, + ) + + return result + + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Retrieve the raw source code for the Python definition.""" + from airflow.sdk.importers.base import DagSourceCode + + return DagSourceCode( + source_code=definition.read_text(encoding="utf-8"), + language="python", + ) + + def _load_modules_from_file( + self, + filepath: str, + safe_mode: bool, + result: DagImportResult, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + ) -> list[ModuleType]: + from airflow import settings + from airflow.sdk._shared.module_loading.dag_file import get_unique_dag_module_name, might_contain_dag + from airflow.sdk.definitions._internal.contextmanager import DagContext + + definition = result.definition + + if not might_contain_dag(filepath, safe_mode): + log.debug("File %s assumed to contain no DAGs. Skipping.", filepath) + if definition is not None: + result.skipped_definitions.append(definition) + return [] + + log.debug("Importing %s (bundle: %s)", filepath, bundle_name) + mod_name = get_unique_dag_module_name(filepath) + + if mod_name in sys.modules: + del sys.modules[mod_name] + + DagContext.current_autoregister_module_name = mod_name + + def parse(mod_name: str, filepath: str) -> list[ModuleType]: + from airflow.configuration import conf Review Comment: Shouldn't the `conf` be safe to import at top level? ########## task-sdk/src/airflow/sdk/importers/base.py: ########## @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Abstract base class for DAG importers.""" + +from __future__ import annotations + +import contextlib +import logging +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.sdk._shared.module_loading.file_discovery import find_path_from_directory + +if TYPE_CHECKING: + from airflow.sdk import DAG + +log = logging.getLogger(__name__) + + +class DagDefinition(ABC): + """Abstract base class for a DAG source definition.""" + + @property + @abstractmethod + def freshness_token(self) -> str: + """Opaque, generalized token representing the current state of the source.""" + + @abstractmethod + def get_relative_loc(self, root: Path | None = None) -> str: + """Get relative location of the definition to a root directory.""" + + @abstractmethod + def read_bytes(self) -> bytes: + """Read and return the content of the resource as bytes.""" + + def read_text(self, encoding: str = "utf-8") -> str: + """Read and return the content of the resource as a string.""" + return self.read_bytes().decode(encoding) + + @abstractmethod + def as_file(self) -> contextlib.AbstractContextManager[Path]: + """ + Return a context manager yielding a Path pointing to a local file. + + For file-backed resources, this is the actual file path. + For others, a temp file is created and cleaned up. + """ + + @abstractmethod + def __repr__(self) -> str: + """Return string representation used by import error and warning objects.""" + + +@dataclass +class FileDagDefinition(DagDefinition): + """A DAG definition backed by a file on the local filesystem.""" + + path: Path + + @property + def freshness_token(self) -> str: + try: + stat = self.path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return str(self.path) + try: + return str(self.path.relative_to(root)) + except ValueError: + return str(self.path) + + def read_bytes(self) -> bytes: + return self.path.read_bytes() + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + yield self.path + + def __repr__(self) -> str: + return str(self.path) + + +@dataclass +class DagImportError: + """Structured error information for DAG import failures.""" + + source_reference: str + message: str + error_type: str = "import" + line_number: int | None = None + column_number: int | None = None + context: str | None = None + suggestion: str | None = None + stacktrace: str | None = None + + def format_message(self) -> str: + """Format the error as a human-readable string.""" + parts = [f"Error in {self.source_reference}"] + if self.line_number is not None: + loc = f"line {self.line_number}" + if self.column_number is not None: + loc += f", column {self.column_number}" + parts.append(f"Location: {loc}") + parts.append(f"Error ({self.error_type}): {self.message}") + if self.context: + parts.append(f"Context:\n{self.context}") + if self.suggestion: + parts.append(f"Suggestion: {self.suggestion}") + return "\n".join(parts) + + +@dataclass +class DagImportWarning: + """Warning information for non-fatal issues during DAG import.""" + + source_reference: str + message: str + warning_type: str = "import" + line_number: int | None = None + context: dict[str, Any] | None = None + + +@dataclass +class DagImportResult: + """Result of importing DAGs from a definition.""" + + definition: DagDefinition | None = None + dags: list[DAG] = field(default_factory=list) + errors: list[DagImportError] = field(default_factory=list) + skipped_definitions: list[DagDefinition] = field(default_factory=list) + warnings: list[DagImportWarning] = field(default_factory=list) + dependencies: list[DagDefinition] = field(default_factory=list) + + +@dataclass +class DagSourceCode: + """Raw source code and its language identifier for a DAG definition.""" + + source_code: str + language: str + + +class AbstractDagImporter(ABC): + """Abstract base class for DAG importers.""" + + @property + @abstractmethod + def supported_extensions(self) -> list[str]: + """Return file extensions this importer handles (e.g., ['.py', '.zip']).""" + + @abstractmethod + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a DAG definition.""" + + def can_handle(self, definition: DagDefinition | str | Path) -> bool: + """Check if this importer can handle the given definition.""" + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) + ) + return Path(path).suffix.lower() in self.supported_extensions if path else False + + def list_dag_definitions( + self, + bundle_name: str, + bundle_path: Path, + *, + safe_mode: bool = True, + ) -> Iterator[DagDefinition]: + """ + List DAG definitions in a bundle that this importer can handle. + + Override this method to customize definition discovery for your importer. + The default implementation finds files matching supported_extensions + and respects .airflowignore files. + """ + try: + from airflow.configuration import conf + + ignore_file_syntax = conf.get_mandatory_value("core", "DAG_IGNORE_FILE_SYNTAX", fallback="glob") + except ImportError: + ignore_file_syntax = "glob" + + supported_exts = [ext.lower() for ext in self.supported_extensions] + + for file_path in find_path_from_directory(bundle_path, ".airflowignore", ignore_file_syntax): + path = Path(file_path) + + if not path.is_file(): + continue + + # Check if this importer handles this file extension + if path.suffix.lower() not in supported_exts: + continue + + yield FileDagDefinition(path=path) + + @abstractmethod + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Retrieve the raw source code and its language identifier for the specified DAG definition.""" + + +class DagImporterRegistry: + """ + Registry for DAG importers. Singleton that manages importers by file extension. + + Each file extension can only be handled by one importer at a time. If multiple + importers claim the same extension, the last registered one wins and a warning + is logged. The built-in PythonDagImporter handles .py and ZipImporter handles .zip files. + """ + + _importers: dict[str, AbstractDagImporter] + + def __init__(self, register_defaults: bool = True): + self._importers = {} + if register_defaults: + self._register_default_importers() + + def _register_default_importers(self) -> None: + from airflow.sdk.importers.python_importer import PythonDagImporter + from airflow.sdk.importers.zip_importer import ZipImporter + + self.register(PythonDagImporter()) + self.register(ZipImporter()) + + def register(self, importer: AbstractDagImporter) -> None: + """ + Register an importer for its supported extensions. + + Each extension can only have one importer. If an extension is already registered, + the new importer will override it and a warning will be logged. + """ + for ext in importer.supported_extensions: + ext_lower = ext.lower() + if ext_lower in self._importers: + existing = self._importers[ext_lower] + log.warning( + "Extension '%s' already registered by %s, overriding with %s", + ext, + type(existing).__name__, + type(importer).__name__, + ) + self._importers[ext_lower] = importer + + def _get_suffix(self, definition: DagDefinition | str | Path) -> str | None: + path = ( + definition + if isinstance(definition, (str, Path)) + else getattr(definition, "path", getattr(definition, "file_path", None)) + ) + return Path(path).suffix.lower() if path else None + + def get_importer(self, definition: DagDefinition | str | Path) -> AbstractDagImporter | None: + """Get the appropriate importer for a definition or file, or None if unsupported.""" + suffix = self._get_suffix(definition) + if suffix and suffix in self._importers: + return self._importers[suffix] + + for importer in set(self._importers.values()): Review Comment: Would it be better to construct a lookup table in the register time instead of computing them all during each call time. ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): + super().__init__() + self._internal_importers: dict[str, AbstractDagImporter] = {} + if internal_importers is None: + from airflow.sdk.importers.python_importer import PythonDagImporter + + self._internal_importers[".py"] = PythonDagImporter() + else: + from airflow.sdk._shared.module_loading import import_string + + for ext, cfg in internal_importers.items(): + ext_lower = ext if ext.startswith(".") else f".{ext}" + if isinstance(cfg, AbstractDagImporter): + self._internal_importers[ext_lower.lower()] = cfg + elif isinstance(cfg, dict) and "classpath" in cfg: + importer_class = import_string(cfg["classpath"]) + self._internal_importers[ext_lower.lower()] = importer_class(**cfg.get("kwargs", {})) + + def import_definition( + self, + definition: DagDefinition, + *, + bundle_path: Path | None = None, + bundle_name: str | None = None, + safe_mode: bool = True, + ) -> DagImportResult: + """Import DAGs from a ZIP archive DAG definition by routing internal files.""" + result = DagImportResult(definition=definition) + + with definition.as_file() as local_zip_path: + try: + valid_members: dict[str, bytes] = {} + with zipfile.ZipFile(local_zip_path) as z: + for member_name in z.namelist(): + # ZipSlip prevention: check for directory traversal attempts + if ".." in Path(member_name).parts or Path(member_name).is_absolute(): + log.warning( + "Skipping zip member %r in %s: directory traversal patterns detected", + member_name, + local_zip_path, + ) + continue + suffix_lower = Path(member_name).suffix.lower() + if suffix_lower in self._internal_importers: + valid_members[member_name] = z.read(member_name) + except Exception as e: + result.errors.append( + DagImportError( + source_reference=definition.get_relative_loc(bundle_path), + message=f"Failed to read ZIP archive: {e}", + error_type="zip_read_error", + ) + ) + return result + + # Extract matching members into a single temp directory to avoid per-file churn + with tempfile.TemporaryDirectory(prefix="airflow_zip_") as temp_dir: + temp_dir_path = Path(temp_dir) + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + extracted_file.parent.mkdir(parents=True, exist_ok=True) + extracted_file.write_bytes(content) + + with _temporary_sys_path(str(temp_dir_path)): + for member_name, content in valid_members.items(): + extracted_file = temp_dir_path / member_name + suffix_lower = extracted_file.suffix.lower() + + nested_def = ZipFileDagDefinition( + zip_path=local_zip_path, + file_path=member_name, + _content=content, + _temp_path=extracted_file, + ) + + importer = self._internal_importers[suffix_lower] Review Comment: ```suggestion if (importer := self._internal_importers.get(suffix_lower, None)) is None: continue ``` ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -0,0 +1,244 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Zip archive DAG importer.""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import tempfile +import threading +import zipfile +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagDefinition, + DagImportError, + DagImportResult, + DagSourceCode, +) + +log = logging.getLogger(__name__) + +_sys_path_lock = threading.Lock() + + [email protected] +def _temporary_sys_path(path: str) -> Iterator[None]: + """Safely prepend a path to sys.path with synchronization and restoration.""" + with _sys_path_lock: + already_present = path in sys.path + if not already_present: + sys.path.insert(0, path) + try: + yield + finally: + if not already_present: + with contextlib.suppress(ValueError): + sys.path.remove(path) + + +@dataclass +class ZipFileDagDefinition(DagDefinition): + """A DAG definition backed by a file inside a ZIP archive.""" + + zip_path: Path + file_path: str + _content: bytes | None = field(default=None, repr=False, compare=False) + _temp_path: Path | None = field(default=None, repr=False, compare=False) + + @property + def freshness_token(self) -> str: + try: + stat = self.zip_path.stat() + return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}" + except OSError: + return "" + + def get_relative_loc(self, root: Path | None = None) -> str: + if root is None: + return f"{self.zip_path}:{self.file_path}" + try: + rel_zip = self.zip_path.relative_to(root) + return f"{rel_zip}:{self.file_path}" + except ValueError: + return f"{self.zip_path}:{self.file_path}" + + def read_bytes(self) -> bytes: + if self._content is None: + with zipfile.ZipFile(self.zip_path) as z: + self._content = z.read(self.file_path) + return self._content + + @contextlib.contextmanager + def as_file(self) -> Iterator[Path]: + if self._temp_path is not None and self._temp_path.exists(): + yield self._temp_path + return + + suffix = Path(self.file_path).suffix + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + f.write(self.read_bytes()) + temp_path = Path(f.name) + try: + yield temp_path + finally: + with contextlib.suppress(OSError): + temp_path.unlink() + + def __repr__(self) -> str: + return f"{self.zip_path}:{self.file_path}" + + +class ZipImporter(AbstractDagImporter): + """Composite importer responsible for routing archive members to internal importers.""" + + supported_extensions = [".zip"] + + def __init__(self, internal_importers: dict[str, Any] | None = None): + super().__init__() + self._internal_importers: dict[str, AbstractDagImporter] = {} + if internal_importers is None: + from airflow.sdk.importers.python_importer import PythonDagImporter + + self._internal_importers[".py"] = PythonDagImporter() + else: + from airflow.sdk._shared.module_loading import import_string + + for ext, cfg in internal_importers.items(): + ext_lower = ext if ext.startswith(".") else f".{ext}" + if isinstance(cfg, AbstractDagImporter): + self._internal_importers[ext_lower.lower()] = cfg + elif isinstance(cfg, dict) and "classpath" in cfg: + importer_class = import_string(cfg["classpath"]) + self._internal_importers[ext_lower.lower()] = importer_class(**cfg.get("kwargs", {})) Review Comment: If yes, we will need to deal with the error handling here. Or perhaps adding pre-validation before the actual import. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
