kacpermuda commented on code in PR #75678: URL: https://github.com/apache/airflow/pull/75678#discussion_r4410718694
########## providers/openlineage/tests/system/openlineage/example_openlineage_policy_exclude_datasets_versioned_dag.py: ########## @@ -0,0 +1,125 @@ +# 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. +""" +Dag exercising the OpenLineage emission policy where Airflow versions differ. + +The comment above each task says which rules and authoring calls apply to it, and what its events must show. + +The rules are in ``AIRFLOW__OPENLINEAGE__EMISSION_POLICY`` of ``airflow-e2e-tests/docker/openlineage.yml``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import datetime + +from airflow import DAG +from airflow.providers.common.compat.assets import Asset, AssetAlias +from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector +from airflow.providers.common.compat.sdk import AirflowSkipException +from airflow.providers.standard.operators.bash import BashOperator +from airflow.providers.standard.operators.python import PythonOperator +from airflow.utils.trigger_rule import TriggerRule + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import AIRFLOW_VERSION, get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + + +class QuietHook: + """Hook that no rule is scoped to.""" + + +def _inlets(stems: Sequence[str]) -> list[Asset]: + return [Asset(uri=f"file://host1/{stem}_in.txt") for stem in stems] + + +def _outlets(stems: Sequence[str]) -> list[Asset]: + return [Asset(uri=f"file://host1/{stem}_out.txt") for stem in stems] + + +ALIAS = AssetAlias("exclude-datasets-alias") + + +def _runtime_outlets(outlet_events): + for stem in ["dag_excluded_runtime", "keep"]: + outlet_events[ALIAS].add(Asset(uri=f"file://host1/{stem}_out.txt")) + + +def _collector_cap(): + # The first 5 assets are excluded and the cap keeps the first 67, so the last 3 of the 70 are lost, not recovered. + stems = [f"dag_excluded_{i:02d}" for i in range(1, 6)] + [f"cap_keep_{i:02d}" for i in range(6, 71)] + collector = get_hook_lineage_collector() + hook = QuietHook() + for stem in stems: + collector.add_input_asset(context=hook, uri=f"file://host1/{stem}_in.txt") + collector.add_output_asset(context=hook, uri=f"file://host1/{stem}_out.txt") + + +def _skip(): + raise AirflowSkipException("Intentional skip") + + +DAG_ID = "openlineage_policy_exclude_datasets_versioned_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + # The Dag list drops dag_excluded from the declared outlets and from the outlets added at runtime through + # the alias. COMPLETE keeps keep_declared and, on Airflow 3.1+, the runtime keep. + t_runtime_outlets = PythonOperator( + task_id="t_runtime_outlets", + python_callable=_runtime_outlets, + inlets=_inlets(["dag_excluded_declared", "keep_declared"]), + outlets=[*_outlets(["dag_excluded_declared", "keep_declared"]), ALIAS], + ) + # Registers 70 assets. On Airflow 3.2+ the collector keeps the first 67, then the Dag list drops the 5 + # dag_excluded ones, leaving 62. Before 3.2 there is no cap, leaving 65. Review Comment: done, added comment about it for now ########## providers/openlineage/tests/system/openlineage/example_openlineage_policy_exclude_datasets_authoring_dag.py: ########## @@ -0,0 +1,284 @@ +# 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. +""" +Dag exercising ``exclude_datasets`` authoring overrides via ``extend_global_openlineage_emission_policy``. + +Hook rules are regexes on the names of the stand-in classes below. Every task reports an input and an output per +stem, and the expected events check both. + +The comment above each task says which rules and authoring calls apply to it, and what its events must show. + +The rules are in ``AIRFLOW__OPENLINEAGE__EMISSION_POLICY`` of ``airflow-e2e-tests/docker/openlineage.yml``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import datetime + +from airflow import DAG +from airflow.providers.common.compat.assets import Asset +from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector +from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy +from airflow.providers.standard.operators.python import PythonOperator + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + + +class QuietHook: + """Hook that no rule is scoped to.""" + + +class AuthHook: + """Hook with a Dag + hook rule that authoring replaces.""" + + +class LockedHook: + """Hook with a locked Dag + task + hook rule.""" + + +class AuthDroppedHook: + """Hook whose datasets a Dag + hook rule with ``hook_lineage: false`` drops.""" + + +class AuthLockedDroppedHook: + """Hook whose datasets a locked Dag + task + hook rule with ``hook_lineage: false`` drops.""" + + +def _register_stems(hook: object, stems: Sequence[str]): + collector = get_hook_lineage_collector() + for stem in stems: + collector.add_input_asset(context=hook, uri=f"file://host1/{stem}_in.txt") + collector.add_output_asset(context=hook, uri=f"file://host1/{stem}_out.txt") + + +def _register_uris(hook: object, inputs: Sequence[str], outputs: Sequence[str]): + collector = get_hook_lineage_collector() + for uri in inputs: + collector.add_input_asset(context=hook, uri=uri) + for uri in outputs: + collector.add_output_asset(context=hook, uri=uri) + + +def _inlets(stems: Sequence[str]) -> list[Asset]: + return [Asset(uri=f"file://host1/{stem}_in.txt") for stem in stems] + + +def _outlets(stems: Sequence[str]) -> list[Asset]: + return [Asset(uri=f"file://host1/{stem}_out.txt") for stem in stems] + + +def _authored_dag(): + _register_stems(QuietHook(), ["dag_authored", "conf_excluded"]) + _register_stems(AuthHook(), ["conf_hook_excluded"]) + + +def _dag_then_task(): + _register_stems(QuietHook(), ["dag_authored", "task_authored", "conf_excluded"]) + + +def _after_dag_call(): + _register_stems(QuietHook(), ["conf_excluded", "dag_authored"]) + _register_stems(AuthHook(), ["conf_hook_excluded", "keep"]) + + +def _authored_task(): + _register_stems(AuthHook(), ["conf_hook_excluded", "authored_a"]) + _register_stems(QuietHook(), ["conf_excluded", "authored_b"]) + + +def _authored_clear(): + _register_stems(QuietHook(), ["conf_excluded"]) + _register_stems(AuthHook(), ["conf_hook_excluded"]) + + +def _hook_lineage_off(): + _register_stems(QuietHook(), ["hook_ignored"]) + + +def _fullmatch(): + _register_stems(QuietHook(), ["authored_partial", "authored_full"]) + + +def _identity(): + _register_uris( + QuietHook(), + inputs=[ + "s3://ol-exclude-bucket/staging/part_1.csv", + "s3://ol-exclude-bucket/final/part_1.csv", + "file://host1/slash_in.txt", + "file://host1/keep_in.txt", + ], + outputs=[ + "s3://ol-exclude-bucket/staging/part_1_out.csv", + "s3://ol-exclude-bucket/final/part_1_out.csv", + "file://host1/slash_out.txt", + "file://host1/keep_out.txt", + ], + ) + + +def _locked(): + _register_stems(QuietHook(), ["locked_excluded_a", "authored_a"]) + _register_stems(AuthHook(), ["specific_excluded", "locked_excluded_b", "authored_b"]) + + +def _locked_hook(): + _register_stems(LockedHook(), ["locked_hook_excluded", "authored_a"]) + _register_stems(QuietHook(), ["conf_excluded", "authored_b"]) + + +def _authored_over_dropped(): + _register_stems(AuthDroppedHook(), ["authored_a", "dropped_kept"]) + + +def _locked_dropped(): + _register_stems(AuthLockedDroppedHook(), ["dropped"]) + _register_stems(QuietHook(), ["conf_excluded", "authored_a", "keep"]) + + +def _do_nothing(): + pass + + +DAG_ID = "openlineage_policy_exclude_datasets_authoring_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + # Conf: the Dag list excludes conf_excluded, the AuthHook list conf_hook_excluded. The Dag-level authored + # list (dag_authored) replaces both: COMPLETE drops dag_authored and keeps conf_excluded and + # conf_hook_excluded. + t_authored_dag = PythonOperator(task_id="t_authored_dag", python_callable=_authored_dag) + # Gets the Dag-level authored list (dag_authored), then its own task-level list (task_authored) replaces + # it: COMPLETE drops task_authored and keeps dag_authored and conf_excluded. + t_dag_then_task = PythonOperator(task_id="t_dag_then_task", python_callable=_dag_then_task) + # Dag-level flags only reach tasks that already exist, so every task defined below is left out of this call. + extend_global_openlineage_emission_policy(dag, exclude_datasets=[".*/dag_authored_.*"]) + extend_global_openlineage_emission_policy(t_dag_then_task, exclude_datasets=[".*/task_authored_.*"]) + + # Defined after the Dag-level call, so the conf lists apply: COMPLETE drops conf_excluded and AuthHook's + # conf_hook_excluded, and keeps dag_authored and keep. + t_after_dag_call = PythonOperator(task_id="t_after_dag_call", python_callable=_after_dag_call) + # The authored list (authored_) replaces the Dag and AuthHook conf lists: COMPLETE drops authored_a and + # authored_b, and keeps conf_excluded and conf_hook_excluded. + t_authored_task = extend_global_openlineage_emission_policy( + PythonOperator(task_id="t_authored_task", python_callable=_authored_task), + exclude_datasets=[".*/authored_.*"], + ) + # The authored [] clears the conf lists: COMPLETE keeps conf_excluded and conf_hook_excluded. + t_authored_clear = extend_global_openlineage_emission_policy( + PythonOperator(task_id="t_authored_clear", python_callable=_authored_clear), + exclude_datasets=[], + ) + # The authored list (authored_) filters the declared inlets and outlets too: COMPLETE drops authored_a and + # keeps conf_excluded. + t_authored_inlets = extend_global_openlineage_emission_policy( + PythonOperator( + task_id="t_authored_inlets", + python_callable=_do_nothing, + inlets=_inlets(["authored_a", "conf_excluded"]), + outlets=_outlets(["authored_a", "conf_excluded"]), + ), + exclude_datasets=[".*/authored_.*"], + ) + # Authored hook_lineage=False drops the hook dataset (hook_ignored), so the declared inlets and outlets Review Comment: done ! -- 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]
