Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions airflow-core/src/airflow/dag_processing/bundles/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations

import os
import warnings
from typing import TYPE_CHECKING

Expand Down Expand Up @@ -106,6 +107,24 @@ def _add_example_dag_bundle(bundle_config_list: list[_ExternalBundleConfig]):
)


def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundleConfig]):
from airflow import providers

for provider_path in providers.__path__:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

THat is a good approach but I assume we need to make a lookup to all providers via plugins_manager. There could also be (in theory) providers not coming from Apache Airflow and they might reside in other hierarchy outside of "airflow" package

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also some providers are nested one level down, like "common/sql" - so most probably via providers_manager we can discover the path/module name where stored... and if not we might need to extend provider.yaml to explicitly mark if examples are available.

for name in os.listdir(provider_path):
example_dag_folder = os.path.join(provider_path, name, "example_dags")
if os.path.isdir(example_dag_folder):
bundle_config_list.append(
_ExternalBundleConfig(
name=f"airflow-provider-{name}-example-dags",
classpath="airflow.dag_processing.bundles.local.LocalDagBundle",
kwargs={
"path": example_dag_folder,
},
)
)


def _is_safe_bundle_url(url: str) -> bool:
"""
Check if a bundle URL is safe to use.
Expand Down Expand Up @@ -191,6 +210,7 @@ def parse_config(self) -> None:
bundle_config_list = _parse_bundle_config(config_list)
if conf.getboolean("core", "LOAD_EXAMPLES"):
_add_example_dag_bundle(bundle_config_list)
_add_provider_example_dags_to_bundle(bundle_config_list)

for bundle_config in bundle_config_list:
if bundle_config.team_name and not conf.getboolean("core", "multi_team"):
Expand Down
8 changes: 4 additions & 4 deletions airflow-core/src/airflow/dag_processing/dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -618,12 +618,12 @@ def collect_dags(

files_to_parse = list_py_file_paths(dag_folder, safe_mode=safe_mode)

if include_examples:
from airflow import example_dags
# if include_examples:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change felt bold. Which is why I commented it out. If we were to let DagBag load examples, then we are duplicating the provider examples loading in both Dagbundle and in dagbag.

Will this affect backward compatibility?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even if it sfeels hard - I assume this is the core thing to remove.

I assume this is a victim of the surgery but later all Dags will be loaded from standard provider directly.

If there are legacy side effects we need to clean them up. So you are doing the "right" thing.
(And please remove commented-out code :-D before merge)

# from airflow import example_dags

example_dag_folder = next(iter(example_dags.__path__))
# example_dag_folder = next(iter(example_dags.__path__))

files_to_parse.extend(list_py_file_paths(example_dag_folder, safe_mode=safe_mode))
# files_to_parse.extend(list_py_file_paths(example_dag_folder, safe_mode=safe_mode))

for filepath in files_to_parse:
try:
Expand Down
1 change: 0 additions & 1 deletion airflow-core/src/airflow/example_dags/standard

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def test_201_and_400_requests(self, url_safe_serializer, session, test_client):
assert response.status_code == 201
parsing_requests = session.scalars(select(DagPriorityParsingRequest)).all()
assert len(parsing_requests) == 1
assert parsing_requests[0].bundle_name == "dags-folder"
assert parsing_requests[0].bundle_name == "example_dags"
assert parsing_requests[0].relative_fileloc == test_dag.relative_fileloc
_check_last_log(session, dag_id=None, event="reparse_dag_file", logical_date=None)

Expand All @@ -65,7 +65,7 @@ def test_201_and_400_requests(self, url_safe_serializer, session, test_client):
assert response.status_code == 409
parsing_requests = session.scalars(select(DagPriorityParsingRequest)).all()
assert len(parsing_requests) == 1
assert parsing_requests[0].bundle_name == "dags-folder"
assert parsing_requests[0].bundle_name == "example_dags"
assert parsing_requests[0].relative_fileloc == test_dag.relative_fileloc
_check_last_log(session, dag_id=None, event="reparse_dag_file", logical_date=None)

Expand Down
23 changes: 9 additions & 14 deletions airflow-core/tests/unit/cli/commands/test_dag_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,8 @@ class TestCliDags:

@classmethod
def setup_class(cls):
parse_and_sync_to_db(os.devnull, include_examples=True)
with conf_vars({("core", "load_examples"): "True"}):
parse_and_sync_to_db(os.devnull, include_examples=True)
cls.parser = cli_parser.get_parser()

@classmethod
Expand Down Expand Up @@ -216,9 +217,8 @@ def test_next_execution(self, tmp_path, stdout_capture):

# Rebuild Test DB for other tests
clear_db_dags()
parse_and_sync_to_db(os.devnull, include_examples=True)
self.setup_class()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR changes the example_dags loaded in Airflow, especially testcases. That leads us to add include_examples parameter, like we do for DagBag

We have two options

  1. Introducing the include_examples parameter to DagBundleManager just to satisfy testing needs
  2. setting conf_vars({("core", "load_examples"): "True"}): at testcase level wherever needed

Choosing (2)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd also propose to choose 2.

In my mind making code in a structure for testability is great but adding balconies for testing in productive code has the danger of future "mis-use"/"unintended use".

Side effect is that we need this very much often, so in regards of convenience maybe a dedicated decorator might be nice (but not critical, just nice) and in regards of performance might be needed to cache examples to prevent heavy and often re-parse as this will maybe impact performance of testing if re-parsed often.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reading more context of the change I assume this is a bit of a mis-conception how the previous code was constructured. Calling the "setup_class()" or a re-sync of dags is conceptually wrong. Not for your change but in general. If the test "messes up" the dags and side effects should be prevented then a proper fixture is needed and it should not be made in a way that the initialization is called again... hard to get the context and error prone as missing might make other test cases broken...

Not sure, might not be fixed here but would be a good issue in another PR to improve.


@conf_vars({("core", "load_examples"): "true"})
def test_cli_report(self, stdout_capture):
args = self.parser.parse_args(["dags", "report", "--output", "json"])
with stdout_capture as temp_stdout:
Expand All @@ -228,7 +228,6 @@ def test_cli_report(self, stdout_capture):
assert "airflow/example_dags/example_complex.py" in out
assert "example_complex" in out

@conf_vars({("core", "load_examples"): "true"})
def test_cli_get_dag_details(self, stdout_capture):
args = self.parser.parse_args(["dags", "details", "example_complex", "--output", "yaml"])
with stdout_capture as temp_stdout:
Expand All @@ -245,7 +244,6 @@ def test_cli_get_dag_details(self, stdout_capture):
for value in dag_details_values:
assert value in out

@conf_vars({("core", "load_examples"): "true"})
def test_cli_list_dags(self, stdout_capture):
args = self.parser.parse_args(["dags", "list", "--output", "json"])
with stdout_capture as temp_stdout:
Expand All @@ -256,11 +254,12 @@ def test_cli_list_dags(self, stdout_capture):
assert key in dag_list[0]
assert any("airflow/example_dags/example_complex.py" in d["fileloc"] for d in dag_list)

@conf_vars({("core", "load_examples"): "true"})
def test_cli_list_local_dags(self, stdout_capture):
# Clear the database
clear_db_dags()
args = self.parser.parse_args(["dags", "list", "--output", "json", "--local"])
args = self.parser.parse_args(
["dags", "list", "--output", "json", "--local", "--bundle-name", "example_dags"]
)
with stdout_capture as temp_stdout:
dag_command.dag_list_dags(args)
out = temp_stdout.getvalue()
Expand All @@ -269,7 +268,7 @@ def test_cli_list_local_dags(self, stdout_capture):
assert key in dag_list[0]
assert any("airflow/example_dags/example_complex.py" in d["fileloc"] for d in dag_list)
# Rebuild Test DB for other tests
parse_and_sync_to_db(os.devnull, include_examples=True)
self.setup_class()

@conf_vars({("core", "load_examples"): "false"})
def test_cli_list_local_dags_with_bundle_name(self, configure_testing_dag_bundle, stdout_capture):
Expand All @@ -290,9 +289,8 @@ def test_cli_list_local_dags_with_bundle_name(self, configure_testing_dag_bundle
str(TEST_DAGS_FOLDER / "test_example_bash_operator.py") in d["fileloc"] for d in dag_list
)
# Rebuild Test DB for other tests
parse_and_sync_to_db(os.devnull, include_examples=True)
self.setup_class()

@conf_vars({("core", "load_examples"): "true"})
def test_cli_list_dags_custom_cols(self, stdout_capture):
args = self.parser.parse_args(
["dags", "list", "--output", "json", "--columns", "dag_id,last_parsed_time"]
Expand All @@ -306,7 +304,6 @@ def test_cli_list_dags_custom_cols(self, stdout_capture):
for key in ["fileloc", "owners", "is_paused"]:
assert key not in dag_list[0]

@conf_vars({("core", "load_examples"): "true"})
def test_cli_list_dags_invalid_cols(self, stderr_capture):
args = self.parser.parse_args(["dags", "list", "--output", "json", "--columns", "dag_id,invalid_col"])
with stderr_capture as temp_stderr:
Expand Down Expand Up @@ -350,9 +347,8 @@ def test_cli_list_dags_prints_local_import_errors(

assert "Failed to load all files." in out
# Rebuild Test DB for other tests
parse_and_sync_to_db(os.devnull, include_examples=True)
self.setup_class()

@conf_vars({("core", "load_examples"): "true"})
@mock.patch("airflow.models.DagModel.get_dagmodel")
def test_list_dags_none_get_dagmodel(self, mock_get_dagmodel, stdout_capture):
mock_get_dagmodel.return_value = None
Expand All @@ -365,7 +361,6 @@ def test_list_dags_none_get_dagmodel(self, mock_get_dagmodel, stdout_capture):
assert key in dag_list[0]
assert any("airflow/example_dags/example_complex.py" in d["fileloc"] for d in dag_list)

@conf_vars({("core", "load_examples"): "true"})
def test_dagbag_dag_col(self, session):
dagbag = DBDagBag()
dag_details = dag_command._get_dagbag_dag_details(
Expand Down
3 changes: 2 additions & 1 deletion airflow-core/tests/unit/cli/commands/test_task_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,8 @@ def test_task_state(self):
)

def test_task_states_for_dag_run(self):
dag2 = DagBag().dags["example_python_operator"]
from airflow.providers.standard.example_dags.example_python_operator import dag as dag2

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That looks OK but is a hard wiring from core to standard provider. Have no better idea but smells a bit "not good".

Might we need a fixture on a general level allowing dags to load w/o tainting individual tests with hard-wired dependencies?


lazy_deserialized_dag2 = LazyDeserializedDAG.from_dag(dag2)

SerializedDagModel.write_dag(lazy_deserialized_dag2, bundle_name="testing")
Expand Down
18 changes: 10 additions & 8 deletions airflow-core/tests/unit/dag_processing/test_dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import zipfile
from copy import deepcopy
from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest import mock
from unittest.mock import patch

Expand All @@ -41,6 +42,7 @@
from airflow.models.dag import DagModel
from airflow.models.dagwarning import DagWarning, DagWarningType
from airflow.models.serialized_dag import SerializedDagModel
from airflow.providers.standard import example_dags as standard_example_dags

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar as above but a bit pmore problematic - if we make this a global import the python module loader needs to have the dependency available at moment of collecting tests already. Can you make this a local import of a small utility function in the test or better to be a fixture that loads it at test runtime to prevent a file level import? (Alongside with code in line 56 and then use it in line 216 and 482 etc)

from airflow.sdk import DAG, BaseOperator

from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH
Expand All @@ -51,7 +53,7 @@

pytestmark = pytest.mark.db_test

example_dags_folder = AIRFLOW_ROOT_PATH / "airflow-core" / "src" / "airflow" / "example_dags" / "standard"
standard_example_dags_folder = Path(standard_example_dags.__file__).parent

PY311 = sys.version_info >= (3, 11)
PY313 = sys.version_info >= (3, 13)
Expand Down Expand Up @@ -211,9 +213,9 @@ def test_get_existing_dag(self, tmp_path):
"""
Test that we're able to parse some example DAGs and retrieve them
"""
dagbag = DagBag(dag_folder=os.fspath(tmp_path), include_examples=True)
dagbag = DagBag(dag_folder=standard_example_dags_folder, include_examples=False)

some_expected_dag_ids = ["example_bash_operator", "example_branch_operator"]
some_expected_dag_ids = ["example_bash_operator", "example_python_operator"]

for dag_id in some_expected_dag_ids:
dag = dagbag.get_dag(dag_id)
Expand Down Expand Up @@ -477,7 +479,7 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True):
_TestDagBag.process_file_calls += 1
super().process_file(filepath, only_if_updated, safe_mode)

dagbag = _TestDagBag(include_examples=True)
dagbag = _TestDagBag(dag_folder=standard_example_dags_folder)
dagbag.process_file_calls

# Should not call process_file again, since it's already loaded during init.
Expand All @@ -489,9 +491,9 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True):
("file_to_load", "expected"),
(
pytest.param(
pathlib.Path(example_dags_folder) / "example_bash_operator.py",
pathlib.Path(standard_example_dags_folder) / "example_bash_operator.py",
{
"example_bash_operator": f"{example_dags_folder.relative_to(AIRFLOW_ROOT_PATH) / 'example_bash_operator.py'}"
"example_bash_operator": f"{standard_example_dags_folder.relative_to(AIRFLOW_ROOT_PATH) / 'example_bash_operator.py'}"
},
id="example_bash_operator",
),
Expand Down Expand Up @@ -553,7 +555,7 @@ def test_refresh_py_dag(self, mock_dagmodel, tmp_path):
Test that we can refresh an ordinary .py DAG
"""
dag_id = "example_bash_operator"
fileloc = str(example_dags_folder / "example_bash_operator.py")
fileloc = str(standard_example_dags_folder / "example_bash_operator.py")

mock_dagmodel.return_value = DagModel()
mock_dagmodel.return_value.last_expired = datetime.max.replace(tzinfo=timezone.utc)
Expand All @@ -567,7 +569,7 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True):
_TestDagBag.process_file_calls += 1
return super().process_file(filepath, only_if_updated, safe_mode)

dagbag = _TestDagBag(dag_folder=os.fspath(tmp_path), include_examples=True)
dagbag = _TestDagBag(dag_folder=standard_example_dags_folder, include_examples=False)

assert dagbag.process_file_calls == 1
dag = dagbag.get_dag(dag_id)
Expand Down
18 changes: 11 additions & 7 deletions airflow-core/tests/unit/models/test_dagcode.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,10 @@ def teardown_method(self):

def _write_two_example_dags(self, session):
example_dags = make_example_dags(example_dags_module)
bash_dag = example_dags["example_bash_operator"]
sync_dag_to_db(bash_dag, session=session)
dag_version = DagVersion.get_latest_version("example_bash_operator")
x = DagCode(dag_version, bash_dag.fileloc)
xcomargs_dag = example_dags["example_xcom_args"]
sync_dag_to_db(xcomargs_dag, session=session)
dag_version = DagVersion.get_latest_version("example_xcom_args")
x = DagCode(dag_version, xcomargs_dag.fileloc)
session.add(x)
session.commit()
xcom_dag = example_dags["example_xcom"]
Expand All @@ -83,7 +83,7 @@ def _write_two_example_dags(self, session):
x = DagCode(dag_version, xcom_dag.fileloc)
session.add(x)
session.commit()
return [bash_dag, xcom_dag]
return [xcomargs_dag, xcom_dag]

def _write_example_dags(self):
example_dags = make_example_dags(example_dags_module)
Expand Down Expand Up @@ -128,7 +128,9 @@ def test_code_can_be_read_when_no_access_to_file(self, testing_dag_bundle):
Test that code can be retrieved from DB when you do not have access to Code file.
Source Code should at least exist in one of DB or File.
"""
example_dag = make_example_dags(example_dags_module).get("example_bash_operator")
from airflow.providers.standard import example_dags

example_dag = make_example_dags(example_dags).get("example_bash_operator")
sync_dag_to_db(example_dag)

# Mock that there is no access to the Dag File
Expand All @@ -141,7 +143,9 @@ def test_code_can_be_read_when_no_access_to_file(self, testing_dag_bundle):

def test_db_code_created_on_serdag_change(self, session, testing_dag_bundle):
"""Test new DagCode is created in DB when ser dag is changed"""
example_dag = make_example_dags(example_dags_module).get("example_bash_operator")
from airflow.providers.standard import example_dags

example_dag = make_example_dags(example_dags).get("example_bash_operator")
sync_dag_to_db(example_dag, session=session).create_dagrun(
run_id="test1",
run_after=pendulum.datetime(2025, 1, 1, tz="UTC"),
Expand Down
22 changes: 11 additions & 11 deletions airflow-core/tests/unit/models/test_serialized_dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,14 +138,14 @@ def my_callable2():
def test_serialized_dag_is_updated_if_dag_is_changed(self, testing_dag_bundle):
"""Test Serialized DAG is updated if DAG is changed"""
example_dags = make_example_dags(example_dags_module)
example_bash_op_dag = example_dags.get("example_bash_operator")
example_params_trigger_ui = example_dags.get("example_params_trigger_ui")
dag_updated = SDM.write_dag(
dag=LazyDeserializedDAG.from_dag(example_bash_op_dag),
dag=LazyDeserializedDAG.from_dag(example_params_trigger_ui),
bundle_name="testing",
)
assert dag_updated is True

s_dag = SDM.get(example_bash_op_dag.dag_id)
s_dag = SDM.get(example_params_trigger_ui.dag_id)
s_dag.dag.create_dagrun(
run_id="test1",
run_after=pendulum.datetime(2025, 1, 1, tz="UTC"),
Expand All @@ -157,28 +157,28 @@ def test_serialized_dag_is_updated_if_dag_is_changed(self, testing_dag_bundle):
# Test that if DAG is not changed, Serialized DAG is not re-written and last_updated
# column is not updated
dag_updated = SDM.write_dag(
dag=LazyDeserializedDAG.from_dag(example_bash_op_dag),
dag=LazyDeserializedDAG.from_dag(example_params_trigger_ui),
bundle_name="testing",
)
s_dag_1 = SDM.get(example_bash_op_dag.dag_id)
s_dag_1 = SDM.get(example_params_trigger_ui.dag_id)

assert s_dag_1.dag_hash == s_dag.dag_hash
assert s_dag.created_at == s_dag_1.created_at
assert dag_updated is False

# Update DAG
example_bash_op_dag.tags.add("new_tag")
assert example_bash_op_dag.tags == {"example", "example2", "new_tag"}
example_params_trigger_ui.tags.add("new_tag")
assert example_params_trigger_ui.tags == {"example", "new_tag", "params"}

dag_updated = SDM.write_dag(
dag=LazyDeserializedDAG.from_dag(example_bash_op_dag),
dag=LazyDeserializedDAG.from_dag(example_params_trigger_ui),
bundle_name="testing",
)
s_dag_2 = SDM.get(example_bash_op_dag.dag_id)
s_dag_2 = SDM.get(example_params_trigger_ui.dag_id)

assert s_dag.created_at != s_dag_2.created_at
assert s_dag.dag_hash != s_dag_2.dag_hash
assert s_dag_2.data["dag"]["tags"] == ["example", "example2", "new_tag"]
assert s_dag_2.data["dag"]["tags"] == ["example", "new_tag", "params"]
assert dag_updated is True

def test_read_dags(self):
Expand All @@ -197,7 +197,7 @@ def test_read_all_dags_only_picks_the_latest_serdags(self, session):
serialized_dags = SDM.read_all_dags()
assert len(example_dags) == len(serialized_dags)

dag = example_dags.get("example_bash_operator")
dag = example_dags.get("example_params_trigger_ui")
SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(dag=dag)).create_dagrun(
run_id="test1",
run_after=pendulum.datetime(2025, 1, 1, tz="UTC"),
Expand Down
5 changes: 4 additions & 1 deletion devel-common/src/tests_common/test_utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,11 @@ def parse_and_sync_to_db(folder: Path | str, include_examples: bool = False):
from airflow.dag_processing.dagbag import sync_bag_to_db
except ImportError:
from airflow.models.dagbag import sync_bag_to_db # type: ignore[no-redef, attribute-defined]

sync_bag_to_db(dagbag, "dags-folder", None, session=session)
for bundle in DagBundlesManager().get_all_dag_bundles():
dagbag = DagBag(dag_folder=bundle.path, include_examples=include_examples)
sync_bag_to_db(dagbag, bundle.name, None, session=session)

elif AIRFLOW_V_3_0_PLUS:
dagbag.sync_to_db("dags-folder", None, session) # type: ignore[attr-defined]
else:
Expand Down
Loading
Loading