Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18"""Dag file utilities for finding and loading Dag files."""
20from __future__ import annotations
22import hashlib
23import os
24import re
25import zipfile
26from pathlib import Path
27from typing import TYPE_CHECKING, Any
29UNUSUAL_MODULE_PREFIX = "unusual_prefix_"
30MODIFIED_DAG_MODULE_NAME = f"{UNUSUAL_MODULE_PREFIX}{{path_hash}}_{{module_name}}"
32if TYPE_CHECKING:
33 from contextlib import AbstractContextManager
34 from typing import Protocol
36 class _MightContainDagCallable(Protocol):
37 def __call__(
38 self,
39 file_path: str | _DagDefinitionLike,
40 zip_file: zipfile.ZipFile | None = None,
41 ) -> bool: ...
43 class _ConfLike(Protocol):
44 def getimport(self, section: str, key: str, **kwargs: Any) -> Any: ...
46 class _DagDefinitionLike(Protocol):
47 """Structural view of a DagDefinition: read its bytes, or materialize it as a file."""
49 def read_bytes(self) -> bytes: ...
51 def as_file(self) -> AbstractContextManager[Path]: ...
54def get_unique_dag_module_name(file_path: str) -> str:
55 """Return a unique module name in the format unusual_prefix_{sha1 of module's file path}_{original module name}."""
56 if isinstance(file_path, str):
57 path_hash = hashlib.sha1(file_path.encode("utf-8"), usedforsecurity=False).hexdigest()
58 org_mod_name = re.sub(r"[.-]", "_", Path(file_path).stem)
59 return MODIFIED_DAG_MODULE_NAME.format(path_hash=path_hash, module_name=org_mod_name)
60 raise ValueError("file_path should be a string to generate unique module name")
63def accepts_dag_definition(func: _MightContainDagCallable) -> _MightContainDagCallable:
64 """
65 Mark a ``might_contain_dag_callable`` as accepting a Dag definition, not only a path.
67 A marked callable is handed the definition itself, so an archive member or any other
68 non-filesystem source is checked without being written to a temporary file first.
69 """
70 func.accepts_dag_definition = True # type: ignore[attr-defined]
71 return func
74@accepts_dag_definition
75def might_contain_dag_via_default_heuristic(
76 file_path: str | _DagDefinitionLike,
77 zip_file: zipfile.ZipFile | None = None,
78) -> bool:
79 """
80 Heuristic that guesses whether a Python file contains an Airflow DAG definition.
82 :param file_path: path to the file to check, or a DagDefinition-like object whose bytes
83 are read directly (nothing is read from disk).
84 :param zip_file: if passed, checks the named member inside the archive. Otherwise, check
85 the local filesystem.
86 :return: True, if file might contain DAGs.
87 """
88 if not isinstance(file_path, (str, os.PathLike)):
89 data = file_path.read_bytes()
90 elif zip_file:
91 with zip_file.open(file_path) as current_file:
92 data = current_file.read()
93 elif zipfile.is_zipfile(file_path):
94 return True
95 else:
96 with open(file_path, "rb") as dag_file:
97 data = dag_file.read()
98 data = data.lower()
99 if b"airflow" not in data:
100 return False
101 return any(s in data for s in (b"dag", b"asset"))
104def might_contain_dag(
105 file_path: str | _DagDefinitionLike,
106 safe_mode: bool = True,
107 zip_file: zipfile.ZipFile | None = None,
108 *,
109 conf: _ConfLike,
110) -> bool:
111 """
112 Check whether a source might contain Airflow DAGs.
114 ``file_path`` may be a filesystem path (optionally with a ``zip_file`` archive whose
115 member it names) or a DagDefinition-like object exposing ``read_bytes()`` and
116 ``as_file()``. Passing a definition lets the check run against an in-memory or
117 archive-backed source without materializing a file. When safe_mode is off (with False
118 value), this function always returns True.
120 A callable marked with :func:`accepts_dag_definition`, including the default heuristic,
121 is handed the definition and reads its bytes directly. Any other callable only
122 understands the legacy ``(file_path, zip_file)`` signature, so a definition is
123 materialized through its own ``as_file()`` and passed by path for compatibility.
124 """
125 if not safe_mode:
126 return True
128 might_contain_dag_callable: _MightContainDagCallable | None = None
129 try:
130 might_contain_dag_callable = conf.getimport(
131 "core",
132 "might_contain_dag_callable",
133 fallback=None,
134 )
135 except Exception as e:
136 import logging
138 logging.getLogger(__name__).warning(
139 "Failed to load might_contain_dag_callable from config, falling back to default heuristic: %s",
140 e,
141 )
143 if might_contain_dag_callable is None:
144 return might_contain_dag_via_default_heuristic(file_path, zip_file=zip_file)
146 if isinstance(file_path, (str, os.PathLike)) or getattr(
147 might_contain_dag_callable, "accepts_dag_definition", False
148 ):
149 return might_contain_dag_callable(file_path=file_path, zip_file=zip_file)
150 # Legacy callables only accept (file_path, zip_file); let the definition materialize itself.
151 with file_path.as_file() as materialized:
152 return might_contain_dag_callable(file_path=str(materialized), zip_file=None)