diff --git a/providers/openlineage/docs/configurations-ref.rst b/providers/openlineage/docs/configurations-ref.rst index fb78ad2424fca..61c3d6c6be6c4 100644 --- a/providers/openlineage/docs/configurations-ref.rst +++ b/providers/openlineage/docs/configurations-ref.rst @@ -25,6 +25,20 @@ Highlighted configurations =========================== + +Configuration precedence +------------------------- + +Primary, and recommended method of configuring OpenLineage Airflow Provider is Airflow configuration. +As there are multiple possible ways of configuring OpenLineage, it's important to keep in mind the precedence of different configurations. +OpenLineage Airflow Provider looks for the configuration in the following order: + +1. Check ``config_conn_id`` in ``airflow.cfg`` under ``openlineage`` section. +2. Check ``config_path`` in ``airflow.cfg`` under ``openlineage`` section (or AIRFLOW__OPENLINEAGE__CONFIG_PATH environment variable) +3. Check ``transport`` in ``airflow.cfg`` under ``openlineage`` section (or AIRFLOW__OPENLINEAGE__TRANSPORT environment variable) +4. If all the above options are missing, the OpenLineage Python client used underneath looks for configuration in the order described in `this `_ documentation. Please note that **using Airflow configuration is encouraged** and is the only future proof solution. + + .. _configuration_transport:openlineage: Transport setup @@ -123,83 +137,6 @@ Example content of config YAML file: Detailed description, together with example config files, can be found `in Python client documentation `_. -Configuration precedence -^^^^^^^^^^^^^^^^^^^^^^^^^ - -Primary, and recommended method of configuring OpenLineage Airflow Provider is Airflow configuration. -As there are multiple possible ways of configuring OpenLineage, it's important to keep in mind the precedence of different configurations. -OpenLineage Airflow Provider looks for the configuration in the following order: - -1. Check ``config_conn_id`` in ``airflow.cfg`` under ``openlineage`` section. -2. Check ``config_path`` in ``airflow.cfg`` under ``openlineage`` section (or AIRFLOW__OPENLINEAGE__CONFIG_PATH environment variable) -3. Check ``transport`` in ``airflow.cfg`` under ``openlineage`` section (or AIRFLOW__OPENLINEAGE__TRANSPORT environment variable) -4. If all the above options are missing, the OpenLineage Python client used underneath looks for configuration in the order described in `this `_ documentation. Please note that **using Airflow configuration is encouraged** and is the only future proof solution. - - -.. _configuration_selective_enable:openlineage: - -Enabling OpenLineage on Dag/task level ---------------------------------------- - -One can selectively enable OpenLineage for specific Dags and tasks by using the ``selective_enable`` policy. -To enable this policy, set the ``selective_enable`` option to True in the [openlineage] section of your Airflow configuration file: - -.. code-block:: ini - - [openlineage] - selective_enable = True - -``AIRFLOW__OPENLINEAGE__SELECTIVE_ENABLE`` environment variable is an equivalent. - -.. code-block:: ini - - AIRFLOW__OPENLINEAGE__SELECTIVE_ENABLE=true - - -While ``selective_enable`` enables selective control, the ``disabled`` option still has precedence. -If you set ``disabled`` to True in the configuration, OpenLineage will be disabled for all Dags and tasks regardless of the ``selective_enable`` setting. - -Once the ``selective_enable`` policy is enabled, you can choose to enable OpenLineage -for individual Dags and tasks using the ``enable_lineage`` and ``disable_lineage`` functions. - -1. Enabling Lineage on a Dag: - -.. code-block:: python - - from airflow.providers.openlineage.utils.selective_enable import disable_lineage, enable_lineage - - with enable_lineage(Dag(...)): - # Tasks within this Dag will have lineage tracking enabled - MyOperator(...) - - AnotherOperator(...) - -2. Enabling Lineage on a Task: - -While enabling lineage on a Dag implicitly enables it for all tasks within that Dag, you can still selectively disable it for specific tasks: - -.. code-block:: python - - from airflow.providers.openlineage.utils.selective_enable import disable_lineage, enable_lineage - - with DAG(...) as dag: - t1 = MyOperator(...) - t2 = AnotherOperator(...) - - # Enable lineage for the entire Dag - enable_lineage(dag) - - # Disable lineage for task t1 - disable_lineage(t1) - -Enabling lineage on the Dag level automatically enables it for all tasks within that Dag unless explicitly disabled per task. - -Enabling lineage on the task level implicitly enables lineage on its Dag. -This is because each emitting task sends a `ParentRunFacet `_, -which requires the Dag-level lineage to be enabled in some OpenLineage backend systems. -Disabling Dag-level lineage while enabling task-level lineage might cause errors or inconsistencies. - - .. _configuration_custom_facets:openlineage: Custom Facets diff --git a/providers/openlineage/docs/emission_policy.rst b/providers/openlineage/docs/emission_policy.rst new file mode 100644 index 0000000000000..726f91cfe2a49 --- /dev/null +++ b/providers/openlineage/docs/emission_policy.rst @@ -0,0 +1,382 @@ + .. 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. + + +.. _emission_policy:openlineage: + +Controlling what OpenLineage emits +================================== + +This page is the one place to look when you want to control **which OpenLineage events are +emitted**, **how much metadata they contain**, and **who is allowed to override what**. + +.. _emission_policy_kill_switch:openlineage: + +Kill switch: ``disabled`` (turn OpenLineage off entirely) +--------------------------------------------------------- + +If you want to turn the OpenLineage integration off **completely** — no plugin loaded, no listener registered, +no events generated, no overhead — without uninstalling the provider, set: + +.. code-block:: ini + + [openlineage] + disabled = true + +or the equivalent environment variable: + +.. code-block:: bash + + AIRFLOW__OPENLINEAGE__DISABLED=true + + +.. note:: + + Changing this variable may require Airflow restart, as the plugins (including OpenLineage) are loaded at startup. + + +This is the all-or-nothing kill switch — useful for incident response, for testing without OL +side-effects, or to keep the provider installed but inactive. When ``disabled = true`` no +``emission_policy`` rule is evaluated and no authoring override has any effect: the OL plugin +is not loaded. + +For anything **less than all-or-nothing** — disabling specific DAGs, tasks, operator classes, +or specific facets; controlling source-code capture or hook-lineage extraction; locking +admin policy against per-Dag author overrides — use ``emission_policy`` instead, documented +below. + + +Two layers of control +--------------------- + +When OL is enabled, there are two layers of control, sitting on different sides of the wire: + +1. **The Airflow OpenLineage provider** (this provider) controls what the *producer* generates + in the first place — whether to run the extractor, whether to include source code, whether + to emit task or DAG-run events at all, whether to inject the full ``AirflowRunFacet``, and + so on. This is the ``emission_policy`` Airflow configuration, documented below. +2. **The OpenLineage Python client** controls what happens to events *after* they are built — + dropping events by name pattern, applying transport-level rules. See + `OpenLineage Python client – filters `_. + +The provider runs first (skipping work) and the client runs second (reshaping what survives). + +.. _emission_policy_unified:openlineage: + +Unified ``emission_policy`` configuration +----------------------------------------- + +The ``emission_policy`` option (under the ``[openlineage]`` section) accepts a JSON array of +rule objects. Each rule has this shape:: + + { + "scope": {}, # required; use {} for "global" + "match_mode": "exact" | "regex", # optional; default "exact" + "controls": {}, # required; must be non-empty + "locked": true | false # optional; default false + } + +**Top-level keys** outside ``scope`` / ``match_mode`` / ``controls`` / ``locked`` are rejected +and the rule is skipped with a WARNING. The same applies to unknown keys inside ``scope`` or +``controls``. This catches typos like ``"dgg_id"`` immediately at config load. + +**Scope keys** (all optional inside ``scope``): + +- ``operator`` — fully-qualified operator class name; applies to task events for that operator type. +- ``dag_id`` — applies to task and/or DAG-level events for the named DAG (use ``emit_task_events`` + / ``emit_dag_events`` in ``controls`` to target one type selectively). +- ``task_id`` — only valid alongside ``dag_id``; targets a specific task. +- *(empty ``scope: {}``)* — global default; applies to every task and DAG event. + +``scope`` must be one of: global (empty), operator-only, dag-only, or dag+task. ``task_id`` +without ``dag_id``, or ``operator`` combined with ``dag_id`` / ``task_id``, is rejected with +a WARNING. + +**``match_mode``** (top-level): ``"exact"`` (default) or ``"regex"``. When ``"regex"``, every +value inside ``scope`` (``dag_id``, ``task_id``, ``operator``) is treated as a ``re.fullmatch`` +pattern. + +**Control flag keys** (all optional inside ``controls``; the dict must be non-empty): + +.. list-table:: + :header-rows: 1 + :widths: 25 10 65 + + * - Flag + - Default + - Description + * - ``emit`` + - ``true`` + - Shorthand: disable **all** OpenLineage events in scope (both task and DAG-level). + * - ``emit_task_events`` + - ``true`` + - Disable task-level events only; takes precedence over ``emit`` for task events within the + same rule. + * - ``emit_dag_events`` + - ``true`` + - Disable DAG-run-level events only; takes precedence over ``emit`` for DAG events within + the same rule. + * - ``extract_operator_metadata`` + - ``true`` + - Whether to run operator-specific extractor-based metadata collection (inputs/outputs). Task events only. + * - ``include_source_code`` + - ``true`` + - Whether to include source code in Python/Bash operator events. Task events only. + **No effect when ``extract_operator_metadata`` is ``false``** — the entire extraction + pipeline is skipped. + * - ``hook_lineage`` + - ``true`` + - Whether to use ``HookLineageCollector`` as a fallback when the extractor finds no + inputs/outputs. Task events only. + **No effect when ``extract_operator_metadata`` is ``false``** — the entire extraction + pipeline (including hook lineage) is skipped. + * - ``include_full_task_info`` + - ``false`` + - Whether to include the full serialized operator state in ``AirflowRunFacet``. When + ``false``, only a curated subset of task attributes is sent. When ``true``, all + serializable task parameters are included, which may significantly increase event size. + Task events only. + +``locked`` (top-level, default ``false``) is an admin-only floor lock: when ``true``, the +control fields carried by this rule's ``controls`` dict cannot be overridden by per-Dag / +per-task authoring flags (see :ref:`emission_policy_authoring:openlineage`). The rule still +participates in normal tier resolution against other conf rules — locking only blocks +authoring overrides. + +``emit_task_events`` / ``emit_dag_events`` take precedence over ``emit`` within the same rule. +For example, ``{"scope": {"dag_id": "X"}, "controls": {"emit": false, "emit_task_events": true}}`` +suppresses DAG-level events for X while leaving task events enabled. + +**Priority** for task events (most specific tier wins; within a tier, the last matching rule wins): + +1. ``dag_id`` + ``task_id`` +2. ``dag_id`` +3. ``operator`` +4. Global (empty ``scope``) +5. Built-in defaults + +**Priority** for DAG-level events: + +1. ``dag_id`` rule (using ``emit_dag_events`` or ``emit``) +2. Global (empty ``scope``) +3. Built-in defaults + +.. note:: + + **Contradictory rules within the same tier.** When two rules in the same priority tier set the + same flag to different values, the *last* rule in the list wins. A ``WARNING`` is logged to help + you spot accidental conflicts. For example, + ``[{"scope": {}, "controls": {"emit": false}}, {"scope": {}, "controls": {"emit": true}}]`` + (both global) results in ``emit=true`` with a warning. Ordering rules intentionally is cleaner + than relying on last-wins. + +.. code-block:: ini + + [openlineage] + emission_policy = [ + {"scope": {"operator": "airflow.providers.http.operators.http.HttpOperator"}, "controls": {"emit": false}}, + {"scope": {"operator": "airflow.providers.standard.operators.python.PythonOperator"}, "controls": {"include_source_code": false}}, + {"scope": {"dag_id": "expensive_dag"}, "controls": {"extract_operator_metadata": false}}, + {"scope": {"dag_id": "expensive_dag", "task_id": "send_report"}, "controls": {"extract_operator_metadata": true}}, + {"scope": {"dag_id": "my_dag", "task_id": "sensitive_task"}, "controls": {"hook_lineage": false}}, + {"scope": {"dag_id": "reporting_dag"}, "controls": {"emit_dag_events": false}}, + {"scope": {"dag_id": "full_control_dag"}, "controls": {"emit": false}}, + {"scope": {"dag_id": "full_control_dag", "task_id": "critical_task"}, "controls": {"emit": true}}, + {"scope": {"dag_id": "^staging_.*"}, "match_mode": "regex", "controls": {"emit": false}}, + {"scope": {"dag_id": "audit_dag"}, "controls": {"include_full_task_info": true}}, + {"scope": {}, "controls": {"include_source_code": false}} + ] + + +``AIRFLOW__OPENLINEAGE__EMISSION_POLICY`` environment variable is an equivalent. + +In the example above: + +- ``HttpOperator`` task events are suppressed entirely for all DAGs. +- ``PythonOperator`` tasks emit events but without source code (all DAGs). +- Tasks in ``expensive_dag`` skip dataset extraction — except ``send_report``, which overrides this + with a more-specific task-scoped rule and extracts normally. +- ``sensitive_task`` in ``my_dag`` runs its extractor normally but skips the + ``HookLineageCollector`` fallback. +- DAG-level events for ``reporting_dag`` are suppressed; task events are still emitted. +- Both task events and DAG-level events for ``full_control_dag`` are suppressed — except + ``critical_task``, which re-enables task event emission via a task-scoped rule. +- All events for any DAG whose ID starts with ``staging_`` are suppressed (regex match). +- Full task info is included in all task events for ``audit_dag``. +- Source code is disabled globally for all operators (``PythonOperator`` was already covered above, + but the global rule applies to ``BashOperator`` and any other operator with source code). + +**Audit logging**: whenever a control flag is non-default (suppressed or enabled), an INFO-level log +is emitted identifying the field, the event context, and the exact rule that caused the change. +This provides a clear audit trail for operational troubleshooting. + +.. _emission_policy_authoring:openlineage: + +Per-Dag / per-task authoring overrides (``extend_global_openlineage_emission_policy``) +-------------------------------------------------------------------------------------- + +Dag authors can override most ``emission_policy`` flags directly in Dag code, without touching +Airflow configuration, via +:func:`~airflow.providers.openlineage.api.emission_policy.extend_global_openlineage_emission_policy`. This sits +**above** ``emission_policy`` conf rules in the resolution stack — the value provided in code +wins **unless** the matching conf rule is marked with ``"locked": true``. + +.. code-block:: python + + from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy + + with DAG("my_dag", ...) as dag: + extract = PythonOperator(task_id="extract", ...) + sensitive = PythonOperator(task_id="sensitive", ...) + + # Disable all events for one task + extend_global_openlineage_emission_policy(sensitive, emit=False) + + # Disable source-code capture for every task in the Dag, re-enable for one + extend_global_openlineage_emission_policy(dag, include_source_code=False) + extend_global_openlineage_emission_policy(extract, include_source_code=True) + + # Suppress DAG-run events but keep task events + extend_global_openlineage_emission_policy(dag, emit_dag_events=False) + +Key semantics: + +- **Resolution priority**: authoring flags > unlocked conf rules > built-in defaults. +- **DAG-or-task only — no global authoring scope.** ``extend_global_openlineage_emission_policy`` only accepts + a single DAG, operator, or :class:`XComArg`. There is no API for "apply to every DAG in the + deployment." Deployment-wide changes are an admin concern and belong in ``emission_policy`` + with an empty ``"scope": {}`` (see the :ref:`unified configuration` above). + Passing any other object type raises :class:`TypeError`. +- **Locked conf rules win**: a rule with ``"locked": true`` blocks authoring overrides for the + field(s) it carries. The override is silently ignored and an INFO-level log records the lock + hit. This is a floor lock — *any* matching conf rule that locks a field is honored, even if + another (more specific) conf rule wins the value resolution. +- **DAG-level calls only propagate to existing tasks** — flags set on a Dag are pushed to the + tasks present on the Dag at the moment the call is made. Tasks defined later (e.g. inside + ``with DAG(...) as dag:`` before the operators are declared) will not inherit them. Call + ``extend_global_openlineage_emission_policy(dag, ...)`` *after* defining the tasks, or set flags per task. +- **No unset API**: pass an explicit boolean to override; passing ``None`` means "not provided" + and leaves any previously-stored value intact. +- **Legacy options cannot be locked**. ``"locked": true`` is an ``emission_policy`` feature. + If you rely only on legacy options (see :ref:`emission_policy_legacy:openlineage` below), + Dag authors can override them via ``extend_global_openlineage_emission_policy``. To enforce a real lock, + move the constraint into ``emission_policy`` and add ``"locked": true``. + +See the docstring of +:func:`~airflow.providers.openlineage.api.emission_policy.extend_global_openlineage_emission_policy` +for the full list of supported flags and examples. + +.. _emission_policy_legacy:openlineage: + +Legacy configuration options (deprecated) +----------------------------------------- + +.. deprecated:: 2.18.0 + + The legacy options ``disabled_for_operators``, ``disable_source_code``, + ``include_full_task_info``, and ``selective_enable`` are superseded by + ``emission_policy``. They continue to work for now but will be removed in a future version + — migrate by translating them into the equivalent rules shown below. + +Any active legacy options are **automatically translated** into equivalent ``emission_policy`` rules. +The translated rules are **prepended** before your explicit rules, so your explicit rules win within each priority tier +(last-wins within a tier). A ``DeprecationWarning`` is issued listing every translated option — migrate +your configuration to ``emission_policy`` exclusively to silence it. + +The translation mapping is: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Legacy option + - Translated to + * - ``disabled_for_operators = some.Operator`` + - ``{"scope": {"operator": "some.Operator"}, "controls": {"emit": false}}`` (operator-scoped, one per entry) + * - ``disable_source_code = True`` + - ``{"scope": {}, "controls": {"include_source_code": false}}`` (global rule) + * - ``include_full_task_info = True`` + - ``{"scope": {}, "controls": {"include_full_task_info": true}}`` (global rule) + * - ``selective_enable = True`` + - ``{"scope": {}, "controls": {"emit": false}}`` (global baseline) + + ``{"scope": {"dag_id": "...", "task_id": "..."}, "controls": {"emit": true}}`` + injected at runtime for each task that called ``enable_lineage()`` + +**Full translation example.** Given this configuration: + +.. code-block:: ini + + [openlineage] + emission_policy = [{"scope": {"dag_id": "my_dag"}, "controls": {"include_source_code": true}}] + disabled_for_operators = airflow.providers.standard.operators.python.PythonOperator + disable_source_code = True + include_full_task_info = True + selective_enable = True + +The effective rule set evaluated at runtime becomes (legacy rules first, your rules last): + +.. code-block:: text + + Tier: operator + {"scope": {"operator": "airflow.providers.standard.operators.python.PythonOperator"}, + "controls": {"emit": false}} + ↳ from disabled_for_operators + + Tier: global + {"scope": {}, "controls": {"include_source_code": false}} + ↳ from disable_source_code + + {"scope": {}, "controls": {"include_full_task_info": true}} + ↳ from include_full_task_info + + {"scope": {}, "controls": {"emit": false}} + ↳ from selective_enable (global baseline — all tasks off by default) + + Tier: task (injected at runtime for each opted-in task) + {"scope": {"dag_id": "my_dag", "task_id": "opted_in_task"}, "controls": {"emit": true}} + ↳ from selective_enable + enable_lineage(task) + + Your explicit emission_policy rules (appended last, win within tier): + Tier: dag + {"scope": {"dag_id": "my_dag"}, "controls": {"include_source_code": true}} + ↳ from emission_policy + +Result for ``PythonOperator`` task in ``my_dag`` (task opted in via ``enable_lineage``): + +- ``emit`` — task-tier opt-in rule sets ``true``, overriding the global ``false`` baseline. +- ``include_source_code`` — your explicit dag-tier rule overrides the global + ``include_source_code=false`` from the translated legacy rule, because dag tier beats global tier. +- ``include_full_task_info`` — global ``true`` from the translated legacy rule. + +.. note:: + + Because ``disabled_for_operators`` translates to an **operator-tier** rule, overriding it from + ``emission_policy`` requires an operator-tier rule with a higher list position (i.e., appearing + later in the same ``emission_policy`` list). A global + ``{"scope": {}, "controls": {"emit": true}}`` rule will *not* override it — operator tier beats + global tier. To re-enable a specific operator that is listed in ``disabled_for_operators``, add + an explicit operator rule to ``emission_policy``: + + .. code-block:: ini + + [openlineage] + emission_policy = [{"scope": {"operator": "my.pkg.MyOperator"}, "controls": {"emit": true}}] + disabled_for_operators = my.pkg.MyOperator + + The translated legacy rule ``{"scope": {"operator": "my.pkg.MyOperator"}, "controls": {"emit": false}}`` + is prepended, but the explicit ``{"scope": {"operator": "my.pkg.MyOperator"}, "controls": {"emit": true}}`` + follows it in the same operator tier — last-wins, so the operator is re-enabled. diff --git a/providers/openlineage/docs/index.rst b/providers/openlineage/docs/index.rst index fd006817c5680..106ad00a808f7 100644 --- a/providers/openlineage/docs/index.rst +++ b/providers/openlineage/docs/index.rst @@ -35,6 +35,7 @@ :caption: Guides Intro + Emission policy Supported classes Custom Operators Job Hierarchy & Macros diff --git a/providers/openlineage/docs/troubleshooting.rst b/providers/openlineage/docs/troubleshooting.rst index 612367c376122..f679e6ff015d9 100644 --- a/providers/openlineage/docs/troubleshooting.rst +++ b/providers/openlineage/docs/troubleshooting.rst @@ -61,8 +61,8 @@ facets are not emitted outside task execution and thus are not applicable in thi To determine which operators will emit OpenLineage events ahead of time, DagRun START events contain AirflowJobFacet with a list of tasks, where each task contains an ``emits_ol_events`` boolean. This checks if the operator is empty, -has callbacks or outlets, and whether task lineage has not been :ref:`selectively disabled ` -or :ref:`disabled for operator `. +has callbacks or outlets, and whether task lineage has not been suppressed by an +:ref:`emission_policy ` rule. Limited lineage from PythonOperator @@ -208,7 +208,7 @@ You can also use some simple transport like the ``ConsoleTransport`` to print ev - Verify the documentation of provider and `client `_, maybe something has changed. - Configuration present: Ensure a working transport is configured. See :ref:`Transport `. -- Disabled settings: Verify you did not disable the integration globally via :ref:`Disabled ` or selectively via :ref:`Disabled for operators ` or :ref:`Selective Enable ` policy. +- Disabled settings: Verify you did not disable the integration globally via :ref:`Disabled ` or selectively via an :ref:`emission_policy ` rule (or the deprecated :ref:`Disabled for operators ` / :ref:`Selective Enable ` options). - Extraction precedence: If inputs/outputs are missing, remember the order described in :ref:`extraction_precedence:openlineage`. - Custom extractors registration: If using custom extractors, confirm they are registered via :ref:`Extractors ` and importable by both Scheduler and Workers. - Environment variables: For legacy environments, note the backwards-compatibility env vars in :ref:`Backwards Compatibility ` (e.g., ``OPENLINEAGE_URL``) but prefer Airflow config. @@ -221,7 +221,7 @@ No events emitted at all: - Ensure the provider is installed and at a supported Airflow version (see provider "Requirements"). - Check :ref:`Disabled ` is not set to ``true``. - - If using selective enablement, verify :ref:`Selective Enable ` and that the DAG/task is enabled via ``enable_lineage``. + - Check your :ref:`emission_policy `: a global ``{"scope": {}, "controls": {"emit": false}}`` rule (or a matching dag / operator / task-scoped rule) silences events. Scheduler / worker logs INFO-level ``emission_policy: 'emit' disabled for ...`` lines whenever a rule suppresses an event. - Confirm the OpenLineage plugin/listener is loaded in Scheduler/Worker logs. Events emitted but not received by backend @@ -256,8 +256,11 @@ Spark jobs missing parent linkage or transport settings Very large event payloads or serialization failures - - If :ref:`Include Full Task Info ` is enabled, events may become large; consider disabling or trimming task parameters. - - :ref:`Disable Source Code ` can reduce payloads for Python/Bash operators that include source code by default. + - Use :ref:`emission_policy ` to trim what gets emitted. Per-control rules let you scope the change (global, per-dag, per-task, or per-operator-class) rather than flipping a deployment-wide toggle: + + - ``{"scope": {}, "controls": {"include_full_task_info": false}}`` keeps the ``AirflowRunFacet`` slim instead of including the full serialized task. + - ``{"scope": {}, "controls": {"include_source_code": false}}`` drops the ``SourceCodeJobFacet`` for Python / Bash operators. + - ``{"scope": {}, "controls": {"extract_operator_metadata": false}}`` skips the entire extraction pipeline, producing a minimal event. 6. Check for open bugs and issues in the provider and the client diff --git a/providers/openlineage/provider.yaml b/providers/openlineage/provider.yaml index 69dae74992eb9..1943fa9a7cd46 100644 --- a/providers/openlineage/provider.yaml +++ b/providers/openlineage/provider.yaml @@ -149,6 +149,8 @@ config: version_added: 1.11.0 disable_source_code: description: | + Deprecated. Use ``emission_policy`` instead. + Disable the inclusion of source code in OpenLineage events by setting this to `true`. By default, several Operators (e.g. Python, Bash) will include their source code in the events unless disabled. @@ -165,6 +167,8 @@ config: version_added: ~ disabled_for_operators: description: | + Deprecated. Use ``emission_policy`` instead. + Exclude some Operators from emitting OpenLineage events by passing a string of semicolon separated full import paths of Operators to disable. type: string @@ -172,6 +176,15 @@ config: airflow.providers.standard.operators.python.PythonOperator" default: "" version_added: 1.1.0 + emission_policy: + description: | + Unified per-scope control over what OpenLineage emits. See the `Emission policy documentation page + `_ + for the full schema, examples, the dag authoring API, and migration from the legacy options. + type: string + example: '[{"scope": {"dag_id": "expensive_dag"}, "controls": {"extract_operator_metadata": false}}]' + default: "[]" + version_added: 2.18.0 execution_timeout: description: | Maximum amount of time (in seconds) that OpenLineage can spend executing metadata extraction for @@ -192,6 +205,8 @@ config: version_added: ~ include_full_task_info: description: | + Deprecated. Use ``emission_policy`` instead. + If true, OpenLineage task events include full serialized task (operator) information. By default, the AirflowRunFacet attached to task events contains only a selected subset of task parameters. With this flag on, all serializable task parameters are sent @@ -213,6 +228,8 @@ config: default: ~ selective_enable: description: | + Deprecated. Use ``emission_policy`` instead. + If this setting is enabled, OpenLineage integration won't collect and emit metadata, unless you explicitly enable it per `DAG` or `Task` using `enable_lineage` method. type: boolean diff --git a/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py b/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py index 6e83e66fb6f66..e8e9f5b68ccc7 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py +++ b/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py @@ -34,6 +34,7 @@ lineage_job_namespace, lineage_run_id, ) +from airflow.providers.openlineage.utils.emission_policy import resolve_task_emission_policy from airflow.providers.openlineage.utils.utils import ( build_task_event_job_facets, build_task_event_run_facets, @@ -131,6 +132,20 @@ def my_task(): dag_run, dag, task = get_dag_run_dag_and_task_from_ti(task_instance) task_uuid = lineage_run_id(task_instance) + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + log.info( + "Skipping OpenLineage RUNNING event emission for task `%s` in dag `%s` " + "due to emission policy. emit_dataset_lineage will have no effect.", + task_instance.task_id, + task_instance.dag_id, + ) + return + run_facets = build_task_event_run_facets( task_instance=task_instance, dag_run=dag_run, @@ -146,6 +161,7 @@ def my_task(): parent_job_name=dag.dag_id, dr_conf=_get_dag_run_conf(task_instance), additional_run_facets=additional_run_facets, + include_full_task_info=controls.include_full_task_info, ) job_facets = build_task_event_job_facets( task=task, dag=dag, additional_job_facets=additional_job_facets diff --git a/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py b/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py new file mode 100644 index 0000000000000..898aae41a6500 --- /dev/null +++ b/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py @@ -0,0 +1,252 @@ +# 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. + +""" +Public authoring-time API for per-task / per-DAG OpenLineage emission control. + +The central function is :func:`extend_global_openlineage_emission_policy`. It lets Dag authors +override the global ``emission_policy`` Airflow configuration on individual tasks or entire +Dags at authoring time, *extending* the deployment-wide policy with a per-Dag / per-task delta. + +Quick-start examples +-------------------- + +**Disable all OpenLineage events for a sensitive task**:: + + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + with DAG("my_dag", ...) as dag: + extract = PythonOperator(task_id="extract", ...) + sensitive = PythonOperator(task_id="sensitive", ...) + + extend_global_openlineage_emission_policy(sensitive, emit=False) + +**Disable source-code capture for an entire Dag, then re-enable for one task**:: + + extend_global_openlineage_emission_policy(dag, include_source_code=False) # all tasks in dag + extend_global_openlineage_emission_policy(extract, include_source_code=True) # override + +**Use as the return value (fluent / inline style)**:: + + with DAG("my_dag", ...) as dag: + task = extend_global_openlineage_emission_policy( + PythonOperator(task_id="my_task", python_callable=my_fn), + include_source_code=False, + hook_lineage=False, + ) + +**Suppress DAG-run events while keeping task events**:: + + extend_global_openlineage_emission_policy(dag, emit_dag_events=False) + +**Works with XComArg** — flags are applied to the underlying operator:: + + result = extend_global_openlineage_emission_policy( + my_python_task_function(), # returns XComArg + extract_operator_metadata=False, + ) + +Flags and locked conf rules +---------------------------- + +These flags sit **above** the ``emission_policy`` Airflow configuration in the resolution stack. +However, if an Airflow admin marks a conf rule with ``locked: true``, that field is protected +and cannot be overridden here — the attempt is silently ignored and an INFO log is emitted. + +Example conf rule that locks include_source_code for all tasks:: + + [openlineage] + emission_policy = [{"scope": {}, "controls": {"include_source_code": false}, "locked": true}] + +Even if a Dag author calls +``extend_global_openlineage_emission_policy(task, include_source_code=True)`` after that, +the lock wins and source code is NOT included. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, TypeVar + +from airflow.providers.common.compat.sdk import DAG, XComArg +from airflow.providers.openlineage.utils.emission_policy import ( + _DAG_FLAG_KEYS, + _TASK_FLAG_KEYS, + EMIT, + EMIT_DAG_EVENTS, + EMIT_TASK_EVENTS, + EXTRACT_OPERATOR_METADATA, + HOOK_LINEAGE, + INCLUDE_FULL_TASK_INFO, + INCLUDE_SOURCE_CODE, + OL_EMISSION_POLICY_PARAM, + _merge_param, +) + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import BaseOperator, MappedOperator + + T = TypeVar("T", bound="DAG | BaseOperator | MappedOperator") + +log = logging.getLogger(__name__) + +__all__ = [ + "extend_global_openlineage_emission_policy", +] + + +def extend_global_openlineage_emission_policy( + obj: T, + *, + emit: bool | None = None, + emit_task_events: bool | None = None, + emit_dag_events: bool | None = None, + extract_operator_metadata: bool | None = None, + include_source_code: bool | None = None, + hook_lineage: bool | None = None, + include_full_task_info: bool | None = None, +) -> T: + """ + Extend the global OpenLineage emission policy with per-task / per-DAG overrides. + + Flags left as ``None`` are not set and will fall through to the ``emission_policy`` + Airflow configuration (or its built-in defaults). Only explicitly provided flags are + stored — successive calls **merge** into any previously stored flags. + + When called on a **DAG**, flags are applied as follows: + + - Task-relevant flags (``emit``, ``emit_task_events``, ``extract_operator_metadata``, + ``include_source_code``, ``hook_lineage``, ``include_full_task_info``) are + **propagated to all tasks** in the Dag at the time of the call. + - DAG-run-level flags (``emit``, ``emit_dag_events``) are stored on the Dag itself. + - ``emit_dag_events`` is meaningless on a task and logs a warning if provided. + + .. warning:: + + **DAG-level calls only propagate to tasks that already exist on the DAG when the call + runs.** Tasks added later (for example, when + ``extend_global_openlineage_emission_policy`` is called inside ``with DAG(...) as dag:`` + before the operators are defined) will **not** inherit the flags. Call + ``extend_global_openlineage_emission_policy(dag, ...)`` **after** all task definitions, + or set flags per-task. + + .. note:: + + There is no "unset" API. Passing ``None`` for a flag is treated as "not provided" — + it does not remove a previously-stored value. To override, pass the explicit boolean + you want, or call again with the new value (successive calls merge, later wins per key). + + ``emit`` is a shorthand that affects both task and Dag events: + + - ``emit=False`` — disables **all** OpenLineage events (task + Dag). + - ``emit=True, emit_task_events=False`` — Dag-run events only (task events off). + - ``emit=True, emit_dag_events=False`` — task events only (Dag-run events off). + + These flags sit **above** the ``emission_policy`` Airflow configuration in priority. + If an admin marks a conf rule with ``locked: true``, that field is protected and cannot + be overridden by this function. + + .. note:: + + **No global authoring scope.** ``extend_global_openlineage_emission_policy`` only + accepts a single DAG, operator, or :class:`XComArg` — there is no equivalent of + "apply globally to every DAG in the deployment." Deployment-wide changes are an + admin concern and belong in the ``emission_policy`` Airflow configuration with an + empty ``"scope": {}``. Passing any other object type raises :class:`TypeError`. + + :param obj: An Airflow Dag, operator, or XComArg. + :param emit: Enable/disable all OpenLineage events (shorthand for both scopes). + :param emit_task_events: Enable/disable task-level events only; takes precedence over + ``emit`` for task events. + :param emit_dag_events: Enable/disable Dag-run-level events only; takes precedence over + ``emit`` for Dag events. Ignored (with a warning) when called on a task. + :param extract_operator_metadata: Whether to run operator-specific extractor-based metadata collection. + :param include_source_code: Whether to include operator source code in Python/Bash operator events. + :param hook_lineage: Whether to use ``HookLineageCollector`` as a fallback. + :param include_full_task_info: Whether to include the full serialized operator state. + :return: The same *obj* — allows use as a decorator or in chained calls. + """ + if isinstance(obj, XComArg): + extend_global_openlineage_emission_policy( + obj.operator, # type: ignore[arg-type] + emit=emit, + emit_task_events=emit_task_events, + emit_dag_events=emit_dag_events, + extract_operator_metadata=extract_operator_metadata, + include_source_code=include_source_code, + hook_lineage=hook_lineage, + include_full_task_info=include_full_task_info, + ) + return obj + + # Type guard: only DAGs and task/operator objects are supported here. + # There is no global authoring scope — see the function docstring and + # the 'emission_policy' Airflow configuration for deployment-wide changes. + if not isinstance(obj, DAG) and not hasattr(obj, "task_id"): + raise TypeError( + "extend_global_openlineage_emission_policy() must be called on a DAG, an Operator, " + f"or an XComArg; got {type(obj).__name__!s}. There is no global authoring scope — " + "for deployment-wide changes, use the [openlineage] 'emission_policy' Airflow " + "configuration with a global rule ('scope': {}) instead." + ) + + provided: dict[str, bool] = { + k: v + for k, v in { + EMIT: emit, + EMIT_TASK_EVENTS: emit_task_events, + EMIT_DAG_EVENTS: emit_dag_events, + EXTRACT_OPERATOR_METADATA: extract_operator_metadata, + INCLUDE_SOURCE_CODE: include_source_code, + HOOK_LINEAGE: hook_lineage, + INCLUDE_FULL_TASK_INFO: include_full_task_info, + }.items() + if v is not None + } + + if not provided: + log.warning( + "OpenLineage extend_global_openlineage_emission_policy(): no emission-control flags were " + "provided for %r — the call has no effect. Pass at least one flag (e.g. emit=False) to " + "store an override.", + getattr(obj, "task_id", None) or getattr(obj, "dag_id", None) or repr(obj), + ) + return obj + + if isinstance(obj, DAG): + dag_flags = {k: v for k, v in provided.items() if k in _DAG_FLAG_KEYS} + if dag_flags: + _merge_param(obj, OL_EMISSION_POLICY_PARAM, dag_flags) + + task_flags = {k: v for k, v in provided.items() if k in _TASK_FLAG_KEYS} + if task_flags: + for task in obj.task_dict.values(): + _merge_param(task, OL_EMISSION_POLICY_PARAM, task_flags) + else: # Task / operator call + if EMIT_DAG_EVENTS in provided: + log.warning( + "OpenLineage extend_global_openlineage_emission_policy(): 'emit_dag_events' has no effect " + "on a task (task_id=%r) — set it on the Dag instead.", + getattr(obj, "task_id", repr(obj)), + ) + task_flags = {k: v for k, v in provided.items() if k in _TASK_FLAG_KEYS} + if task_flags: + _merge_param(obj, OL_EMISSION_POLICY_PARAM, task_flags) + + return obj diff --git a/providers/openlineage/src/airflow/providers/openlineage/api/sql.py b/providers/openlineage/src/airflow/providers/openlineage/api/sql.py index b59b1f2056f6a..bf8e17a9e7b60 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/api/sql.py +++ b/providers/openlineage/src/airflow/providers/openlineage/api/sql.py @@ -26,11 +26,13 @@ from airflow.providers.openlineage.api.core import emit, is_openlineage_active from airflow.providers.openlineage.plugins.adapter import _PRODUCER from airflow.providers.openlineage.plugins.macros import lineage_job_name +from airflow.providers.openlineage.utils.emission_policy import resolve_task_emission_policy from airflow.providers.openlineage.utils.sql_hook_lineage import ( _create_ol_event_pair, _parse_query_into_datasets, ) from airflow.providers.openlineage.utils.utils import ( + get_dag_run_dag_and_task_from_ti, get_task_instance_from_context, next_query_counter_from_context, ) @@ -129,6 +131,21 @@ def my_task(): log.debug("TaskInstance not provided, retrieving it from context.") task_instance = get_task_instance_from_context() + _, _, task = get_dag_run_dag_and_task_from_ti(task_instance) + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + log.info( + "Skipping OpenLineage QUERY event emission for task `%s` in dag `%s` " + "due to emission policy. emit_query_lineage will have no effect.", + task_instance.task_id, + task_instance.dag_id, + ) + return + # Copy caller-supplied lists so we never mutate user inputs. all_inputs = list(inputs) if inputs else [] all_outputs = list(outputs) if outputs else [] @@ -174,7 +191,9 @@ def my_task(): end_event_time=end_time, ) - log.info("emit_query_lineage will emit 2 OpenLineage events for job `%s`.", start_event.job.name) + log.info( + "emit_query_lineage will emit 2 OpenLineage QUERY events for job `%s`.", start_event.job.name + ) emit(start_event) emit(end_event) except Exception as err: diff --git a/providers/openlineage/src/airflow/providers/openlineage/conf.py b/providers/openlineage/src/airflow/providers/openlineage/conf.py index 3a2ca201e53e4..d46dcc67cd1b9 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/conf.py +++ b/providers/openlineage/src/airflow/providers/openlineage/conf.py @@ -179,3 +179,12 @@ def include_full_task_info() -> bool: def debug_mode() -> bool: """[openlineage] debug_mode.""" return conf.getboolean(_CONFIG_SECTION, "debug_mode", fallback="False") + + +@cache +def emission_policy() -> list[dict]: + """[openlineage] emission_policy.""" + option = conf.getjson(_CONFIG_SECTION, "emission_policy", fallback=[]) + if not isinstance(option, list): + raise ValueError(f"[openlineage] emission_policy must be a JSON array, got: {type(option).__name__}") + return option diff --git a/providers/openlineage/src/airflow/providers/openlineage/extractors/base.py b/providers/openlineage/src/airflow/providers/openlineage/extractors/base.py index f8d4eac2b49a3..c4115c83b66c6 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/extractors/base.py +++ b/providers/openlineage/src/airflow/providers/openlineage/extractors/base.py @@ -67,9 +67,10 @@ class BaseExtractor(ABC, LoggingMixin): _allowed_query_params: list[str] = [] - def __init__(self, operator): + def __init__(self, operator, source_code_enabled: bool = True): super().__init__() self.operator = operator + self.source_code_enabled = source_code_enabled @classmethod @abstractmethod diff --git a/providers/openlineage/src/airflow/providers/openlineage/extractors/bash.py b/providers/openlineage/src/airflow/providers/openlineage/extractors/bash.py index 6e1b3f28eefe0..45d5df7102cad 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/extractors/bash.py +++ b/providers/openlineage/src/airflow/providers/openlineage/extractors/bash.py @@ -19,7 +19,6 @@ from openlineage.client.facet_v2 import source_code_job -from airflow.providers.openlineage import conf from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage from airflow.providers.openlineage.utils.utils import get_unknown_source_attribute_run_facet @@ -45,7 +44,7 @@ def get_operator_classnames(cls) -> list[str]: def _execute_extraction(self) -> OperatorLineage | None: job_facets: dict = {} - if conf.is_source_enabled(): + if self.source_code_enabled: job_facets = { "sourceCode": source_code_job.SourceCodeJobFacet( language="bash", @@ -55,7 +54,7 @@ def _execute_extraction(self) -> OperatorLineage | None: } else: self.log.debug( - "OpenLineage disable_source_code option is on - no source code is extracted.", + "OpenLineage source_code is disabled - no source code is extracted.", ) return OperatorLineage( diff --git a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py index 276636e129e23..7509044958072 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py +++ b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py @@ -29,6 +29,7 @@ ) from airflow.providers.openlineage.extractors.bash import BashExtractor from airflow.providers.openlineage.extractors.python import PythonExtractor +from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy from airflow.providers.openlineage.utils.utils import ( get_runtime_outlet_assets, get_unknown_source_attribute_run_facet, @@ -92,9 +93,16 @@ def add_extractor(self, operator_class: str, extractor: type[BaseExtractor]): self.extractors[operator_class] = extractor def extract_metadata( - self, dagrun, task, task_instance_state: TaskInstanceState, task_instance + self, + dagrun, + task, + task_instance_state: TaskInstanceState, + task_instance, + controls: EmissionPolicy | None = None, ) -> OperatorLineage: - extractor = self._get_extractor(task) + if controls is None: + controls = EmissionPolicy.defaults() + extractor = self._get_extractor(task, source_code_enabled=controls.include_source_code) task_info = ( f"task_type={task.task_type} " f"airflow_dag_id={task.dag_id} " @@ -126,45 +134,59 @@ def extract_metadata( str(task_metadata), ) task_metadata = self.validate_task_metadata(task_metadata) or OperatorLineage() - # If no inputs and outputs are present - check Hook Lineage + # If no inputs and outputs are present - check Hook Lineage if enabled if (not task_metadata.inputs) and (not task_metadata.outputs): - hook_lineage = self.get_hook_lineage(task_instance, task_instance_state) - if hook_lineage is not None: - task_metadata = task_metadata.merge(hook_lineage) - else: # Last resort - check manual annotations + if controls.hook_lineage: + hook_lineage = self.get_hook_lineage(task_instance, task_instance_state) + if hook_lineage is not None: + task_metadata = task_metadata.merge(hook_lineage) + else: # Last resort - check manual annotations + self.extract_inlets_and_outlets(task_metadata, task, task_instance) + else: + self.log.info( + "Skipping OpenLineage hook lineage collection for task '%s' due to emission_policy.", + task.task_id, + ) self.extract_inlets_and_outlets(task_metadata, task, task_instance) return task_metadata except Exception as e: self.log.warning( - "Failed to extract metadata using found extractor %s - %s %s", + "Failed to extract OpenLineage metadata using found extractor %s - %s %s", extractor, e, task_info, ) self.log.debug("OpenLineage extraction failure details:", exc_info=True) else: - # No extractor found — fall back to hook lineage. This call must be wrapped in - # try/except: it runs emit_lineage_from_sql_extras → _create_ol_event_pair which - # is not guarded internally. An uncaught exception here would propagate up to the - # listener's @print_warning decorator, silently suppressing the task-level event. - try: - hook_lineage = self.get_hook_lineage(task_instance, task_instance_state) - except Exception as e: - self.log.warning( - "Failed to extract hook lineage %s: %s. Task event will be emitted without lineage.", - task_info, - e, + # No extractor found. + if controls.hook_lineage: + # Fall back to hook lineage. This call must be wrapped in try/except: it runs + # emit_lineage_from_sql_extras → _create_ol_event_pair which is not guarded + # internally. An uncaught exception here would propagate up to the listener's + # @print_warning decorator, silently suppressing the task-level event. + try: + hook_lineage = self.get_hook_lineage(task_instance, task_instance_state) + except Exception as e: + self.log.warning( + "Failed to extract OpenLineage hook lineage %s: %s. Task event will be emitted without lineage.", + task_info, + e, + ) + self.log.debug("OpenLineage hook lineage failure details:", exc_info=True) + hook_lineage = None + + if hook_lineage is not None: + return hook_lineage + else: + self.log.info( + "Skipping OpenLineage hook lineage collection for task '%s' due to emission_policy.", + task.task_id, ) - self.log.debug("OpenLineage hook lineage failure details:", exc_info=True) - hook_lineage = None - - if hook_lineage is not None: - return hook_lineage - self.log.debug("Unable to find an extractor %s", task_info) - - # Only include the unknownSourceAttribute facet if there is no extractor + # No extractor and no hook lineage (or hook lineage disabled) — always emit the + # unknownSourceAttribute facet and fall back to manually-declared inlets/outlets. + self.log.debug("Unable to find an OpenLineage extractor %s", task_info) task_metadata = OperatorLineage( run_facets=get_unknown_source_attribute_run_facet(task=task), ) @@ -184,12 +206,18 @@ def method_exists(method_name): return self.default_extractor return None - def _get_extractor(self, task: BaseOperator) -> BaseExtractor | None: - extractor = self.get_extractor_class(task) - self.log.debug("extractor for %s is %s", task.task_type, extractor) - if extractor: - return extractor(task) - return None + def _get_extractor(self, task: BaseOperator, source_code_enabled: bool = True) -> BaseExtractor | None: + extractor_cls = self.get_extractor_class(task) + self.log.debug("extractor for %s is %s", task.task_type, extractor_cls) + if extractor_cls is None: + return None + # Set the flag after construction rather than passing it as a constructor kwarg: + # custom extractors may still use the historically-public `__init__(self, operator)` + # signature, which would raise TypeError on an unexpected keyword and drop the whole + # task event. Built-in extractors read `source_code_enabled` at extraction time. + extractor = extractor_cls(task) + extractor.source_code_enabled = source_code_enabled + return extractor def extract_inlets_and_outlets( self, diff --git a/providers/openlineage/src/airflow/providers/openlineage/extractors/python.py b/providers/openlineage/src/airflow/providers/openlineage/extractors/python.py index 37054d25d2942..c6e27d67f8b85 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/extractors/python.py +++ b/providers/openlineage/src/airflow/providers/openlineage/extractors/python.py @@ -22,7 +22,6 @@ from openlineage.client.facet_v2 import source_code_job -from airflow.providers.openlineage import conf from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage from airflow.providers.openlineage.utils.utils import get_unknown_source_attribute_run_facet @@ -49,7 +48,7 @@ def get_operator_classnames(cls) -> list[str]: def _execute_extraction(self) -> OperatorLineage | None: source_code = self.get_source_code(self.operator.python_callable) job_facet: dict = {} - if conf.is_source_enabled() and source_code: + if self.source_code_enabled and source_code: job_facet = { "sourceCode": source_code_job.SourceCodeJobFacet( language="python", @@ -59,7 +58,7 @@ def _execute_extraction(self) -> OperatorLineage | None: } else: self.log.debug( - "OpenLineage disable_source_code option is on - no source code is extracted.", + "OpenLineage source_code is disabled - no source code is extracted.", ) return OperatorLineage( diff --git a/providers/openlineage/src/airflow/providers/openlineage/get_provider_info.py b/providers/openlineage/src/airflow/providers/openlineage/get_provider_info.py index 60cf1a981a6cf..117981894b400 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/get_provider_info.py +++ b/providers/openlineage/src/airflow/providers/openlineage/get_provider_info.py @@ -86,7 +86,7 @@ def get_provider_info(): "version_added": "1.11.0", }, "disable_source_code": { - "description": "Disable the inclusion of source code in OpenLineage events by setting this to `true`.\nBy default, several Operators (e.g. Python, Bash) will include their source code in the events\nunless disabled.\n", + "description": "Deprecated. Use ``emission_policy`` instead.\n\nDisable the inclusion of source code in OpenLineage events by setting this to `true`.\nBy default, several Operators (e.g. Python, Bash) will include their source code in the events\nunless disabled.\n", "default": "False", "example": None, "type": "boolean", @@ -100,12 +100,19 @@ def get_provider_info(): "version_added": None, }, "disabled_for_operators": { - "description": "Exclude some Operators from emitting OpenLineage events by passing a string of semicolon separated\nfull import paths of Operators to disable.\n", + "description": "Deprecated. Use ``emission_policy`` instead.\n\nExclude some Operators from emitting OpenLineage events by passing a string of semicolon separated\nfull import paths of Operators to disable.\n", "type": "string", "example": "airflow.providers.standard.operators.bash.BashOperator; airflow.providers.standard.operators.python.PythonOperator", "default": "", "version_added": "1.1.0", }, + "emission_policy": { + "description": "Unified per-scope control over what OpenLineage emits. See the `Emission policy documentation page\n`_\nfor the full schema, examples, the dag authoring API, and migration from the legacy options.\n", + "type": "string", + "example": '[{"scope": {"dag_id": "expensive_dag"}, "controls": {"extract_operator_metadata": false}}]', + "default": "[]", + "version_added": "2.18.0", + }, "execution_timeout": { "description": "Maximum amount of time (in seconds) that OpenLineage can spend executing metadata extraction for\ntask (on worker). Note that other configurations, sometimes with higher priority, such as\n`[core] task_success_overtime\n`_,\nmay also affect how much time OpenLineage has for execution.\n", "default": "10", @@ -121,7 +128,7 @@ def get_provider_info(): "version_added": None, }, "include_full_task_info": { - "description": "If true, OpenLineage task events include full serialized task (operator) information.\nBy default, the AirflowRunFacet attached to task events contains only a selected subset\nof task parameters. With this flag on, all serializable task parameters are sent\n(excluding known non-serializable elements), which may significantly increase event size.\n\nWarning: By setting this variable to true, OpenLineage event can potentially include elements that\nare megabytes in size or larger, depending on the size of data you pass to the task.\n", + "description": "Deprecated. Use ``emission_policy`` instead.\n\nIf true, OpenLineage task events include full serialized task (operator) information.\nBy default, the AirflowRunFacet attached to task events contains only a selected subset\nof task parameters. With this flag on, all serializable task parameters are sent\n(excluding known non-serializable elements), which may significantly increase event size.\n\nWarning: By setting this variable to true, OpenLineage event can potentially include elements that\nare megabytes in size or larger, depending on the size of data you pass to the task.\n", "default": "False", "example": None, "type": "boolean", @@ -135,7 +142,7 @@ def get_provider_info(): "default": None, }, "selective_enable": { - "description": "If this setting is enabled, OpenLineage integration won't collect and emit metadata,\nunless you explicitly enable it per `DAG` or `Task` using `enable_lineage` method.\n", + "description": "Deprecated. Use ``emission_policy`` instead.\n\nIf this setting is enabled, OpenLineage integration won't collect and emit metadata,\nunless you explicitly enable it per `DAG` or `Task` using `enable_lineage` method.\n", "type": "boolean", "default": "False", "example": None, diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index 2d4d74e828f09..ec24c37128f89 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -34,6 +34,10 @@ from airflow.providers.openlineage import conf from airflow.providers.openlineage.extractors import ExtractorManager, OperatorLineage from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter, RunState +from airflow.providers.openlineage.utils.emission_policy import ( + resolve_dag_emission_policy, + resolve_task_emission_policy, +) from airflow.providers.openlineage.utils.utils import ( AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, @@ -50,8 +54,6 @@ get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, - is_operator_disabled, - is_selective_lineage_enabled, print_warning, ) from airflow.settings import configure_orm @@ -181,20 +183,16 @@ def on_task_instance_running( # type: ignore[misc] def _on_task_instance_running( self, task_instance: RuntimeTaskInstance | TaskInstance, dag, dagrun, task, start_date: datetime ): - if is_operator_disabled(task): - self.log.debug( - "Skipping OpenLineage event emission for operator `%s` " - "due to its presence in [openlineage] disabled_for_operators.", - task.task_type, - ) - return - - if not is_selective_lineage_enabled(task): - self.log.debug( - "Skipping OpenLineage event emission for task `%s` " - "due to lack of explicit lineage enablement for task or DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + self.log.info( + "Skipping OpenLineage event emission for task `%s` in dag `%s`.", task_instance.task_id, + task_instance.dag_id, ) return @@ -243,13 +241,23 @@ def on_running(): if not doc: doc, doc_type = get_dag_documentation(dag) - with Stats.timer("ol.extract", tags={"event_type": event_type, "operator_name": operator_name}): - task_metadata = self.extractor_manager.extract_metadata( - dagrun=dagrun, - task=task, - task_instance_state=TaskInstanceState.RUNNING, - task_instance=task_instance, + if controls.extract_operator_metadata: + with Stats.timer( + "ol.extract", tags={"event_type": event_type, "operator_name": operator_name} + ): + task_metadata = self.extractor_manager.extract_metadata( + dagrun=dagrun, + task=task, + task_instance_state=TaskInstanceState.RUNNING, + task_instance=task_instance, + controls=controls, + ) + else: + self.log.info( + "Skipping OpenLineage operator metadata extraction for task `%s` due to emission_policy.", + task_instance.task_id, ) + task_metadata = OperatorLineage() redacted_event = self.adapter.start_task( run_id=task_uuid, @@ -271,7 +279,14 @@ def on_running(): dr_conf=getattr(dagrun, "conf", {}), ), **get_airflow_mapped_task_facet(task_instance), - **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid), + **get_airflow_run_facet( + dagrun, + dag, + task_instance, + task, + task_uuid, + include_full_task_info=controls.include_full_task_info, + ), **debug_facet, }, ) @@ -325,20 +340,16 @@ def on_task_instance_success( # type: ignore[misc] def _on_task_instance_success(self, task_instance: RuntimeTaskInstance, dag, dagrun, task): end_date = timezone.utcnow() - if is_operator_disabled(task): - self.log.debug( - "Skipping OpenLineage event emission for operator `%s` " - "due to its presence in [openlineage] disabled_for_operators.", - task.task_type, - ) - return - - if not is_selective_lineage_enabled(task): - self.log.debug( - "Skipping OpenLineage event emission for task `%s` " - "due to lack of explicit lineage enablement for task or DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + self.log.info( + "Skipping OpenLineage event emission for task `%s` in dag `%s`.", task_instance.task_id, + task_instance.dag_id, ) return @@ -375,13 +386,23 @@ def on_success(): if not doc: doc, doc_type = get_dag_documentation(dag) - with Stats.timer("ol.extract", tags={"event_type": event_type, "operator_name": operator_name}): - task_metadata = self.extractor_manager.extract_metadata( - dagrun=dagrun, - task=task, - task_instance_state=TaskInstanceState.SUCCESS, - task_instance=task_instance, + if controls.extract_operator_metadata: + with Stats.timer( + "ol.extract", tags={"event_type": event_type, "operator_name": operator_name} + ): + task_metadata = self.extractor_manager.extract_metadata( + dagrun=dagrun, + task=task, + task_instance_state=TaskInstanceState.SUCCESS, + task_instance=task_instance, + controls=controls, + ) + else: + self.log.info( + "Skipping OpenLineage operator metadata extraction for task `%s` due to emission_policy.", + task_instance.task_id, ) + task_metadata = OperatorLineage() redacted_event = self.adapter.complete_task( run_id=task_uuid, @@ -402,7 +423,14 @@ def on_success(): parent_job_name=dag.dag_id, dr_conf=getattr(dagrun, "conf", {}), ), - **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid), + **get_airflow_run_facet( + dagrun, + dag, + task_instance, + task, + task_uuid, + include_full_task_info=controls.include_full_task_info, + ), **get_airflow_debug_facet(), }, ) @@ -471,20 +499,16 @@ def _on_task_instance_failed( ) -> None: end_date = timezone.utcnow() - if is_operator_disabled(task): - self.log.debug( - "Skipping OpenLineage event emission for operator `%s` " - "due to its presence in [openlineage] disabled_for_operators.", - task.task_type, - ) - return - - if not is_selective_lineage_enabled(task): - self.log.debug( - "Skipping OpenLineage event emission for task `%s` " - "due to lack of explicit lineage enablement for task or DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + self.log.info( + "Skipping OpenLineage event emission for task `%s` in dag `%s`.", task_instance.task_id, + task_instance.dag_id, ) return @@ -521,13 +545,23 @@ def on_failure(): if not doc: doc, doc_type = get_dag_documentation(dag) - with Stats.timer("ol.extract", tags={"event_type": event_type, "operator_name": operator_name}): - task_metadata = self.extractor_manager.extract_metadata( - dagrun=dagrun, - task=task, - task_instance_state=TaskInstanceState.FAILED, - task_instance=task_instance, + if controls.extract_operator_metadata: + with Stats.timer( + "ol.extract", tags={"event_type": event_type, "operator_name": operator_name} + ): + task_metadata = self.extractor_manager.extract_metadata( + dagrun=dagrun, + task=task, + task_instance_state=TaskInstanceState.FAILED, + task_instance=task_instance, + controls=controls, + ) + else: + self.log.info( + "Skipping OpenLineage operator metadata extraction for task `%s` due to emission_policy.", + task_instance.task_id, ) + task_metadata = OperatorLineage() redacted_event = self.adapter.fail_task( run_id=task_uuid, @@ -549,7 +583,14 @@ def on_failure(): parent_job_name=dag.dag_id, dr_conf=getattr(dagrun, "conf", {}), ), - **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid), + **get_airflow_run_facet( + dagrun, + dag, + task_instance, + task, + task_uuid, + include_full_task_info=controls.include_full_task_info, + ), **get_airflow_debug_facet(), }, ) @@ -594,20 +635,16 @@ def _on_task_instance_skipped( ) -> None: end_date = timezone.utcnow() - if is_operator_disabled(task): - self.log.debug( - "Skipping OpenLineage event emission for operator `%s` " - "due to its presence in [openlineage] disabled_for_operators.", - task.task_type, - ) - return - - if not is_selective_lineage_enabled(task): - self.log.debug( - "Skipping OpenLineage event emission for task `%s` " - "due to lack of explicit lineage enablement for task or DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_task_emission_policy( + operator=task, + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + ) + if not controls.emit: + self.log.info( + "Skipping OpenLineage event emission for task `%s` in dag `%s`.", task_instance.task_id, + task_instance.dag_id, ) return @@ -644,13 +681,23 @@ def on_skipped(): if not doc: doc, doc_type = get_dag_documentation(dag) - with Stats.timer("ol.extract", tags={"event_type": event_type, "operator_name": operator_name}): - task_metadata = self.extractor_manager.extract_metadata( - dagrun=dagrun, - task=task, - task_instance_state=TaskInstanceState.SKIPPED, - task_instance=task_instance, + if controls.extract_operator_metadata: + with Stats.timer( + "ol.extract", tags={"event_type": event_type, "operator_name": operator_name} + ): + task_metadata = self.extractor_manager.extract_metadata( + dagrun=dagrun, + task=task, + task_instance_state=TaskInstanceState.SKIPPED, + task_instance=task_instance, + controls=controls, + ) + else: + self.log.info( + "Skipping OpenLineage operator metadata extraction for task `%s` due to emission_policy.", + task_instance.task_id, ) + task_metadata = OperatorLineage() redacted_event = self.adapter.complete_task( run_id=task_uuid, @@ -671,7 +718,14 @@ def on_skipped(): parent_job_name=dag.dag_id, dr_conf=getattr(dagrun, "conf", {}), ), - **get_airflow_run_facet(dagrun, dag, task_instance, task, task_uuid), + **get_airflow_run_facet( + dagrun, + dag, + task_instance, + task, + task_uuid, + include_full_task_info=controls.include_full_task_info, + ), **get_airflow_debug_facet(), }, ) @@ -705,23 +759,22 @@ def _on_task_instance_manual_state_change( self.log.debug("`_on_task_instance_manual_state_change` was called with state: `%s`.", ti_state) end_date = timezone.utcnow() + include_full_task_info = False task = getattr(ti, "task") # on scheduler, we should have access to task - if task and is_operator_disabled(task): - self.log.debug( - "Skipping OpenLineage event emission for operator `%s` " - "due to its presence in [openlineage] disabled_for_operators.", - task.task_type, - ) - return - - if task and not is_selective_lineage_enabled(task): - self.log.debug( - "Skipping OpenLineage event emission for task `%s` " - "due to lack of explicit lineage enablement for task or DAG while " - "[openlineage] selective_enable is on.", - ti.task_id, + if task: + controls = resolve_task_emission_policy( + operator=task, + dag_id=ti.dag_id, + task_id=ti.task_id, ) - return + if not controls.emit: + self.log.info( + "Skipping OpenLineage event emission for task `%s` in dag `%s`.", + ti.task_id, + ti.dag_id, + ) + return + include_full_task_info = controls.include_full_task_info try: if not self.executor: @@ -778,7 +831,14 @@ def _on_task_instance_manual_state_change( doc, doc_type = get_dag_documentation(dag) dag_tags = dag.tags owners = [x.strip() for x in (task if task.owner != "airflow" else dag).owner.split(",")] - airflow_run_facet = get_airflow_run_facet(dagrun, dag, ti, task, task_uuid) + airflow_run_facet = get_airflow_run_facet( + dagrun, + dag, + ti, + task, + task_uuid, + include_full_task_info=include_full_task_info, + ) adapter_kwargs: dict = { "run_id": task_uuid, @@ -895,11 +955,10 @@ def before_stopping(self, component) -> None: @hookimpl def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None: try: - if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag): - self.log.debug( - "Skipping OpenLineage event emission for DAG `%s` " - "due to lack of explicit lineage enablement for DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_dag_emission_policy(dag_run.dag_id, dag=dag_run.dag) + if not controls.emit: + self.log.info( + "Skipping OpenLineage dag event emission for DAG `%s`.", dag_run.dag_id, ) return @@ -947,11 +1006,10 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None: @hookimpl def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None: try: - if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag): - self.log.debug( - "Skipping OpenLineage event emission for DAG `%s` " - "due to lack of explicit lineage enablement for DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_dag_emission_policy(dag_run.dag_id, dag=dag_run.dag) + if not controls.emit: + self.log.info( + "Skipping OpenLineage dag event emission for DAG `%s`.", dag_run.dag_id, ) return @@ -999,11 +1057,10 @@ def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None: @hookimpl def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None: try: - if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag): - self.log.debug( - "Skipping OpenLineage event emission for DAG `%s` " - "due to lack of explicit lineage enablement for DAG while " - "[openlineage] selective_enable is on.", + controls = resolve_dag_emission_policy(dag_run.dag_id, dag=dag_run.dag) + if not controls.emit: + self.log.info( + "Skipping OpenLineage dag event emission for DAG `%s`.", dag_run.dag_id, ) return diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py b/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py new file mode 100644 index 0000000000000..de386fce438f5 --- /dev/null +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py @@ -0,0 +1,1112 @@ +# 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. + +""" +Per-scope emission policy resolution for the OpenLineage provider. + +The ``emission_policy`` Airflow configuration option (``[openlineage]`` section) +accepts a JSON array of rule objects. Each rule has the shape:: + + { + "scope": {}, # required, may be empty for "global" + "match_mode": "exact" | "regex", # optional, default "exact" + "controls": {}, # required, must be non-empty + "locked": true | false # optional, default false + } + +Top-level keys outside the four above are rejected with a WARNING and the rule +is skipped. + +Scope keys (all optional inside ``scope``): + +- ``operator`` — fully-qualified operator class name; matches task events for that + operator type. +- ``dag_id`` — matches task and/or dag run events for all tasks / dag runs belonging + to the named DAG. Use ``emit_task_events`` / ``emit_dag_events`` to target one + event type selectively. +- ``task_id`` — only valid alongside ``dag_id``; targets a specific task. +- *(empty ``scope: {}``)* — global default override; applies to every task and DAG event. + +``operator`` cannot be combined with ``dag_id`` or ``task_id`` (scope is one of: +global, operator-only, dag-only, dag+task). ``task_id`` without ``dag_id`` and +``operator`` combined with ``emit_dag_events`` (or ``task_id`` combined with +``emit_dag_events``) are also rejected with a WARNING. + +``match_mode`` lives at the top level: ``"exact"`` (default) or ``"regex"``. +When set to ``"regex"``, every value inside ``scope`` (``dag_id``, ``task_id``, +``operator``) is treated as a ``re.fullmatch`` pattern. + +Control flag keys (all optional inside ``controls``; the dict must be non-empty): + +- ``emit`` — shorthand: disable ALL OpenLineage events in scope (both task and dag + run events). Default: ``true``. +- ``emit_task_events`` — disable task-level events only; takes precedence over + ``emit`` for task event decisions. Default: ``true``. +- ``emit_dag_events`` — disable dag-run-level events only; takes precedence over + ``emit`` for dag event decisions. Default: ``true``. +- ``extract_operator_metadata`` — whether to run operator-specific extractor-based metadata + collection. When ``true``, the extractor manager calls the operator's OpenLineage extractor + (if registered), which may produce dataset inputs/outputs, job facets, run facets, and + other operator-specific metadata. When ``false``, the entire extraction pipeline is skipped + and a minimal event is emitted. Only meaningful for task events. Default: ``true``. +- ``include_source_code`` — whether to include operator source code in the + ``SourceCodeJobFacet`` for Python and Bash operators. Only meaningful for task events + when ``extract_operator_metadata`` is also ``true``. Default: ``true``. +- ``hook_lineage`` — whether to use ``HookLineageCollector`` as a fallback when the + extractor finds no inputs/outputs. **Has no effect when ``extract_operator_metadata`` + is ``false``** — the entire extraction pipeline (including hook lineage) is skipped. + Only meaningful for task events. Default: ``true``. +- ``include_full_task_info`` — whether to include the full serialized operator state + in the ``AirflowRunFacet``. When ``false`` (default), only a curated subset of + task attributes is sent. Only meaningful for task events. Default: ``false``. + +``locked: true`` is an admin floor lock that prevents per-Dag / per-task authoring +overrides from changing the field(s) carried by this rule's ``controls`` dict. + +Flag hierarchy +~~~~~~~~~~~~~~ + +Flags are not fully independent — some only take effect when a higher-level flag is +enabled: + +- ``extract_operator_metadata: false`` skips the **entire** operator extraction pipeline. + When this is ``false``, both ``include_source_code`` and ``hook_lineage`` have **no + effect** regardless of their values, because the code paths they control are never + reached. +- ``include_source_code`` only applies inside Python and Bash operator extractors; setting + it to ``false`` on other operator types is a no-op. + +Priority for **task events** (most specific tier wins; within a tier, last matching +rule wins): + +1. ``dag_id`` + ``task_id`` +2. ``dag_id`` +3. ``operator`` +4. Global (no match keys) +5. Built-in defaults + +For ``emit`` at the task level: ``emit_task_events`` beats ``emit`` within the same +rule. + +Priority for **DAG run events**: + +1. ``dag_id`` rule (using ``emit_dag_events`` or ``emit``) +2. Global (no match keys) +3. Built-in defaults + +For ``emit`` at the dag level: ``emit_dag_events`` beats ``emit`` within the same +rule. + +Legacy config translation +------------------------- + +Any active legacy config option (``disabled_for_operators``, ``disable_source_code``, +``include_full_task_info``, ``selective_enable``) is **always** translated into equivalent +``emission_policy`` rules, regardless of whether ``emission_policy`` itself is set. The +translated rules are prepended before any user-provided rules, so user rules win within +each priority tier (last-wins). A ``DeprecationWarning`` is issued listing every translated +option — silence it by migrating those options into ``emission_policy`` exclusively. + +When no legacy option is active and ``emission_policy`` is empty, the resolver returns +the built-in defaults with no warnings. + +Audit logging +------------- + +Every resolved field that is non-default — whether from a rule or from a legacy +translation — is logged at INFO level, identifying the field, the event context, +and the exact rule that caused the change. +""" + +from __future__ import annotations + +import logging +import re +import warnings +from collections.abc import Callable, Sequence +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.providers.common.compat.sdk import Param +from airflow.providers.openlineage import conf as _ol_conf + +if TYPE_CHECKING: + from airflow.providers.openlineage.utils.utils import AnyOperator + +log = logging.getLogger(__name__) + +# Control flag names — the keys that may appear inside a rule's ``controls`` dict. +EMIT = "emit" +EMIT_TASK_EVENTS = "emit_task_events" +EMIT_DAG_EVENTS = "emit_dag_events" +EXTRACT_OPERATOR_METADATA = "extract_operator_metadata" +INCLUDE_SOURCE_CODE = "include_source_code" +HOOK_LINEAGE = "hook_lineage" +INCLUDE_FULL_TASK_INFO = "include_full_task_info" + +# Scope keys — the keys that may appear inside a rule's ``scope`` dict. +SCOPE_DAG_ID = "dag_id" +SCOPE_TASK_ID = "task_id" +SCOPE_OPERATOR = "operator" + +# Top-level rule keys. +RULE_SCOPE = "scope" +RULE_CONTROLS = "controls" +RULE_MATCH_MODE = "match_mode" +RULE_LOCKED = "locked" + +# Param names used by the authoring API to store flow-control flags on operator / DAG +# objects. Kept in this module (rather than the api module) so the resolver can read +# them without a circular import. +OL_EMISSION_POLICY_PARAM = "_openlineage_emission_policy" + +# Schema enforcement: only these keys may appear at each level. Unknown keys cause +# the rule to be skipped with a WARNING (catches typos like {"scope": {"dgg_id": "x"}}). +_ALLOWED_TOP_LEVEL_KEYS: frozenset[str] = frozenset({RULE_SCOPE, RULE_MATCH_MODE, RULE_CONTROLS, RULE_LOCKED}) +_ALLOWED_SCOPE_KEYS: frozenset[str] = frozenset({SCOPE_DAG_ID, SCOPE_TASK_ID, SCOPE_OPERATOR}) +_ALLOWED_CONTROL_KEYS: frozenset[str] = frozenset( + { + EMIT, + EMIT_TASK_EVENTS, + EMIT_DAG_EVENTS, + EXTRACT_OPERATOR_METADATA, + INCLUDE_SOURCE_CODE, + HOOK_LINEAGE, + INCLUDE_FULL_TASK_INFO, + } +) + +# Fields scanned for ``locked: true`` on task-event resolution. +# ``emit`` is intentionally absent: emit-locking is driven by either ``emit`` or +# ``emit_task_events`` appearing in a locked rule (handled separately in +# :func:`_compute_locked_task_fields`), so listing it here would double-count. +_LOCKABLE_TASK_FIELDS: tuple[str, ...] = ( + EXTRACT_OPERATOR_METADATA, + INCLUDE_SOURCE_CODE, + HOOK_LINEAGE, + INCLUDE_FULL_TASK_INFO, +) + +# Authoring flag classification: which keys are relevant for task-level vs DAG-run +# resolution. Used by the authoring API to split a single call into the two stored +# param dicts. +_TASK_FLAG_KEYS: frozenset[str] = frozenset( + { + EMIT, + EMIT_TASK_EVENTS, + EXTRACT_OPERATOR_METADATA, + INCLUDE_SOURCE_CODE, + HOOK_LINEAGE, + INCLUDE_FULL_TASK_INFO, + } +) +_DAG_FLAG_KEYS: frozenset[str] = frozenset({EMIT, EMIT_DAG_EVENTS}) + + +@dataclass(frozen=True) +class EmissionPolicy: + """Resolved emission policy for a specific event context.""" + + emit: bool + extract_operator_metadata: bool + include_source_code: bool + hook_lineage: bool + include_full_task_info: bool + + @classmethod + def defaults(cls) -> EmissionPolicy: + """Return the default policy (all controls at their built-in defaults).""" + return cls( + emit=True, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) + + +@dataclass(frozen=True) +class Rule: + """ + A single, fully-validated ``emission_policy`` rule. + + Instances are produced exclusively by :func:`_parse_rule`, so holding a ``Rule`` + is itself the "this passed schema validation" guarantee — downstream resolution + never re-validates. ``scope`` / ``controls`` are the raw (validated) sub-dicts. + """ + + scope: dict[str, str] + controls: dict[str, bool] + match_mode: str = "exact" + locked: bool = False + + @property + def has_dag(self) -> bool: + return SCOPE_DAG_ID in self.scope + + @property + def has_task(self) -> bool: + return SCOPE_TASK_ID in self.scope + + @property + def has_op(self) -> bool: + return SCOPE_OPERATOR in self.scope + + +def _matches(pattern: str, value: str, match_mode: str) -> bool: + """Return ``True`` if *value* matches *pattern* according to *match_mode*.""" + if match_mode == "regex": + return bool(re.fullmatch(pattern, value)) + return pattern == value + + +def _read_param(obj: object, param_name: str) -> dict[str, bool]: + """Read the flags dict stored at *param_name* on *obj*, or ``{}`` if absent.""" + val = getattr(obj, "params", {}).get(param_name) + if val is None: + return {} + if isinstance(val, Param): + val = val.value + return val if isinstance(val, dict) else {} + + +def _merge_param(obj: object, param_name: str, new_flags: dict[str, bool]) -> None: + """ + Merge *new_flags* into the ``param_name`` param on *obj*, creating it if absent. + + The public entry point + (:func:`~airflow.providers.openlineage.api.emission_policy.extend_global_openlineage_emission_policy`) + validates that *obj* is a DAG or task-like object before reaching here, so a missing + ``params`` attribute indicates a programmer error (e.g. an incorrectly-mocked object + in tests) — we raise loud rather than silently dropping flags. + """ + params = getattr(obj, "params", None) + if params is None: + raise TypeError( + f"extend_global_openlineage_emission_policy: {type(obj).__name__} instance has no " + "'params' attribute — cannot store flow-control flags. " + "Pass a DAG, an Operator, or an XComArg." + ) + existing = params.get(param_name) + if existing is not None: + existing_val = existing.value if isinstance(existing, Param) else existing + if isinstance(existing_val, dict): + new_flags = {**existing_val, **new_flags} + params[param_name] = Param(new_flags) + + +def _audit_log_conf_field(field: str, value: bool, context: str, source: Rule) -> None: + """Log the winning conf value for a field whenever any rule touched it.""" + log.info( + "OpenLineage emission policy: '%s' %s for %s by %r", + field, + "enabled" if value else "disabled", + context, + source, + ) + + +def _audit_log_authoring_updates( + field_changes: dict[str, bool], + context: str, +) -> None: + """Audit-log authoring overrides that actually change the resolved policy value.""" + for field, new_value in field_changes.items(): + log.info( + "OpenLineage emission policy: '%s' %s for %s " + "by manual `extend_global_openlineage_emission_policy` call.", + field, + "enabled" if new_value else "disabled", + context, + ) + + +def _parse_rule(rule: object) -> Rule | None: + """ + Parse and validate a raw rule dict into a :class:`Rule`; warn and return ``None`` otherwise. + + A valid rule has exactly these top-level keys: ``scope`` (dict, possibly empty), + ``controls`` (non-empty dict), and the optionals ``match_mode`` / ``locked``. + Unknown keys at any level cause the rule to be skipped — that catches user + typos like ``{"scope": {"dgg_id": "x"}}``. + + This is the single validation gate: it runs once at config-read time (see + :func:`_parse_user_rules`), so the resolution hot path never re-validates. + """ + if not isinstance(rule, dict): + log.warning("OpenLineage emission_policy entry is not a dict: %r; ignoring.", rule) + return None + + unknown_top = set(rule) - _ALLOWED_TOP_LEVEL_KEYS + if unknown_top: + log.warning( + "OpenLineage emission_policy rule has unknown top-level key(s) %s (allowed: %s); ignoring: %r", + sorted(unknown_top), + sorted(_ALLOWED_TOP_LEVEL_KEYS), + rule, + ) + return None + + if RULE_SCOPE not in rule: + log.warning( + "OpenLineage emission_policy rule is missing required 'scope' key (use 'scope': {} for global); ignoring: %r", + rule, + ) + return None + scope = rule[RULE_SCOPE] + if not isinstance(scope, dict): + log.warning( + "OpenLineage emission_policy rule 'scope' must be a dict (use {} for global), got %r; ignoring: %r", + type(scope).__name__, + rule, + ) + return None + unknown_scope = set(scope) - _ALLOWED_SCOPE_KEYS + if unknown_scope: + log.warning( + "OpenLineage emission_policy rule 'scope' has unknown key(s) %s (allowed: %s); ignoring: %r", + sorted(unknown_scope), + sorted(_ALLOWED_SCOPE_KEYS), + rule, + ) + return None + has_dag = SCOPE_DAG_ID in scope + has_task = SCOPE_TASK_ID in scope + has_op = SCOPE_OPERATOR in scope + if has_task and not has_dag: + log.warning( + "OpenLineage emission_policy rule scope has 'task_id' without 'dag_id'; ignoring: %r", + rule, + ) + return None + if has_op and (has_dag or has_task): + log.warning( + "OpenLineage emission_policy rule scope combines 'operator' with 'dag_id'/'task_id' " + "(must be exactly one of: global, operator-only, dag-only, dag+task); ignoring: %r", + rule, + ) + return None + + if RULE_CONTROLS not in rule: + log.warning( + "OpenLineage emission_policy rule is missing required 'controls' dict; ignoring: %r", + rule, + ) + return None + controls = rule[RULE_CONTROLS] + if not isinstance(controls, dict): + log.warning( + "OpenLineage emission_policy rule 'controls' must be a dict, got %r; ignoring: %r", + type(controls).__name__, + rule, + ) + return None + if not controls: + log.warning( + "OpenLineage emission_policy rule has empty 'controls' dict " + "(at least one control flag is required); ignoring: %r", + rule, + ) + return None + unknown_controls = set(controls) - _ALLOWED_CONTROL_KEYS + if unknown_controls: + log.warning( + "OpenLineage emission_policy rule 'controls' has unknown key(s) %s (allowed: %s); ignoring: %r", + sorted(unknown_controls), + sorted(_ALLOWED_CONTROL_KEYS), + rule, + ) + return None + for k, v in controls.items(): + if not isinstance(v, bool): + log.warning( + "OpenLineage emission_policy rule 'controls.%s' must be bool, got %r; ignoring: %r", + k, + v, + rule, + ) + return None + + if has_op and EMIT_DAG_EVENTS in controls: + log.warning( + "OpenLineage emission_policy rule has scope 'operator' with controls 'emit_dag_events' which is meaningless " + "(operators are not associated with DAG run events); ignoring: %r", + rule, + ) + return None + if has_task and EMIT_DAG_EVENTS in controls: + log.warning( + "OpenLineage emission_policy rule has scope 'task_id' with controls 'emit_dag_events' which is meaningless " + "(task-specific rules do not target DAG run events); ignoring: %r", + rule, + ) + return None + + match_mode = rule.get(RULE_MATCH_MODE, "exact") + if match_mode not in ("exact", "regex"): + log.warning( + "OpenLineage emission_policy rule has invalid 'match_mode' %r (must be 'exact' or 'regex'); ignoring: %r", + match_mode, + rule, + ) + return None + if match_mode == "regex": + for key in (SCOPE_DAG_ID, SCOPE_TASK_ID, SCOPE_OPERATOR): + if key in scope: + try: + re.compile(scope[key]) + except re.error as exc: + log.warning( + "OpenLineage emission_policy rule scope.'%s' pattern %r is not a valid regex (%s); ignoring: %r", + key, + scope[key], + exc, + rule, + ) + return None + + if RULE_LOCKED in rule and not isinstance(rule[RULE_LOCKED], bool): + log.warning( + "OpenLineage emission_policy rule 'locked' must be bool, got %r; ignoring: %r", + rule[RULE_LOCKED], + rule, + ) + return None + # Under the nested schema every allowed control key is lockable (validation already + # requires non-empty controls with keys from the allowed set), so there is no + # "lock with no lockable field" warning case to raise. + + return Rule( + scope=rule[RULE_SCOPE], + controls=rule[RULE_CONTROLS], + match_mode=match_mode, + locked=rule.get(RULE_LOCKED, False), + ) + + +@_ol_conf.cache +def _parse_user_rules() -> tuple[Rule, ...]: + """ + Parse and cache the user-configured ``emission_policy`` rules. + + Validation runs here exactly once per config value (``conf.emission_policy()`` is + itself cached), so task / dag resolution operates on pre-validated :class:`Rule` + objects without re-validating on every event. + """ + return tuple(parsed for raw in _ol_conf.emission_policy() if (parsed := _parse_rule(raw)) is not None) + + +def _classify_task_rules( + rules: Sequence[Rule], + fqcn: str, + dag_id: str, + task_id: str, +) -> tuple[list[Rule], list[Rule], list[Rule], list[Rule]]: + """ + Classify pre-validated rules into the four priority tiers for task resolution. + + Returns ``(task_rules, dag_rules, operator_rules, global_rules)``. + """ + task_rules: list[Rule] = [] + dag_rules: list[Rule] = [] + operator_rules: list[Rule] = [] + global_rules: list[Rule] = [] + + for rule in rules: + scope = rule.scope + match_mode = rule.match_mode + + if rule.has_dag and rule.has_task: + if _matches(scope[SCOPE_DAG_ID], dag_id, match_mode) and _matches( + scope[SCOPE_TASK_ID], task_id, match_mode + ): + task_rules.append(rule) + elif rule.has_dag: + if _matches(scope[SCOPE_DAG_ID], dag_id, match_mode): + dag_rules.append(rule) + elif rule.has_op: + if _matches(scope[SCOPE_OPERATOR], fqcn, match_mode): + operator_rules.append(rule) + else: + global_rules.append(rule) + + return task_rules, dag_rules, operator_rules, global_rules + + +def _classify_dag_rules( + rules: Sequence[Rule], + dag_id: str, +) -> tuple[list[Rule], list[Rule]]: + """ + Classify pre-validated rules into the two priority tiers for dag event resolution. + + Only rules with empty ``scope`` (global) and rules with exactly ``{dag_id}`` + contribute to dag-event resolution. Operator-scoped and task-scoped rules + never affect dag-run events. + + Returns ``(dag_rules, global_rules)``. + """ + dag_rules: list[Rule] = [] + global_rules: list[Rule] = [] + + for rule in rules: + if rule.has_dag and not rule.has_task and not rule.has_op: + if _matches(rule.scope[SCOPE_DAG_ID], dag_id, rule.match_mode): + dag_rules.append(rule) + elif not (rule.has_dag or rule.has_task or rule.has_op): # global = empty scope dict + global_rules.append(rule) + + return dag_rules, global_rules + + +def _walk_tiers( + tiers: list[list[Rule]], + extract_value: Callable[[dict[str, bool]], bool | None], + field_label: str, + default: bool, +) -> tuple[bool, Rule | None]: + """ + Walk priority tiers from most-specific to least-specific, last-wins within each tier. + + *extract_value* maps a rule's ``controls`` dict to ``bool | None``; ``None`` means + "this rule does not cover this field". A contradiction WARNING is emitted when two + rules in the same tier produce different non-``None`` values. Returns + ``(default, None)`` when no rule in any tier matches. + """ + for tier_rules in tiers: + value: bool | None = None + winning_rule: Rule | None = None + for rule in tier_rules: + new_v = extract_value(rule.controls) + if new_v is not None: + if winning_rule is not None and value != new_v: + log.warning( + "OpenLineage emission_policy: field '%s' has contradictory values in the same " + "priority tier (using last value %r); conflicting rules:" + "\n - %r" + "\n - %r", + field_label, + new_v, + winning_rule, + rule, + ) + value = new_v + winning_rule = rule + if value is not None: + return value, winning_rule + return default, None + + +def _resolve_field_with_source( + tiers: list[list[Rule]], field: str, default: bool +) -> tuple[bool, Rule | None]: + """Walk tiers for a single boolean *field* in ``controls``.""" + return _walk_tiers(tiers, lambda c: c.get(field), field, default) + + +def _resolve_emit_with_source( + tiers: list[list[Rule]], + scope: str, + default: bool, +) -> tuple[bool, Rule | None]: + """ + Resolve the effective ``emit`` decision for the given *scope*. + + ``scope`` must be ``"task"`` or ``"dag"``. Within each rule, the + scope-specific key (``emit_task_events`` / ``emit_dag_events``) takes + precedence over the generic ``emit`` shorthand. + """ + specific_key = EMIT_TASK_EVENTS if scope == "task" else EMIT_DAG_EVENTS + + def _extract(c: dict[str, bool]) -> bool | None: + if specific_key in c: + return c[specific_key] + return c.get(EMIT) + + return _walk_tiers(tiers, _extract, f"emit ({scope})", default) + + +def _synthesize_legacy_task_rules( + operator: AnyOperator, + dag_id: str, + task_id: str, +) -> list[Rule]: + """ + Translate legacy config options into equivalent ``emission_policy`` rules for task resolution. + + The returned rules are intended to be **prepended** to the user-provided rules so that + user rules (appended later) win within each tier (last-wins within a tier). They are + constructed directly (trusted, no validation path). + """ + rules: list[Rule] = [] + + # disabled_for_operators -> operator-scope emit:false rules + for fqcn in _ol_conf.disabled_operators(): + rules.append(Rule(scope={SCOPE_OPERATOR: fqcn}, controls={EMIT: False})) + + # disable_source_code -> global include_source_code:false + if not _ol_conf.is_source_enabled(): + rules.append(Rule(scope={}, controls={INCLUDE_SOURCE_CODE: False})) + + # include_full_task_info -> global include_full_task_info:true (only when non-default True) + if _ol_conf.include_full_task_info(): + rules.append(Rule(scope={}, controls={INCLUDE_FULL_TASK_INFO: True})) + + # selective_enable -> global emit:false baseline + per-task opt-in rule + if _ol_conf.selective_enable(): + rules.append(Rule(scope={}, controls={EMIT: False})) # global baseline + try: + from airflow.providers.openlineage.utils.selective_enable import is_task_lineage_enabled + + if is_task_lineage_enabled(operator): + # Task is explicitly opted in -> task-tier rule that overrides the global baseline. + rules.append( + Rule(scope={SCOPE_DAG_ID: dag_id, SCOPE_TASK_ID: task_id}, controls={EMIT: True}) + ) + except Exception: + # Non-standard operator types may not expose params correctly. Surface at debug + # level so the failure is discoverable without spamming production logs. + log.debug( + "OpenLineage selective_enable translation: is_task_lineage_enabled() failed for task '%s'", + task_id, + exc_info=True, + ) + + return rules + + +def _synthesize_legacy_dag_rules( + dag_id: str, + dag: object | None, +) -> list[Rule]: + """Translate legacy config options into equivalent ``emission_policy`` rules for dag event resolution.""" + rules: list[Rule] = [] + + if _ol_conf.selective_enable() and dag is not None: + rules.append(Rule(scope={}, controls={EMIT: False})) # global baseline + try: + from airflow.providers.openlineage.utils.selective_enable import is_dag_lineage_enabled + + if is_dag_lineage_enabled(dag): # type: ignore[arg-type] + rules.append(Rule(scope={SCOPE_DAG_ID: dag_id}, controls={EMIT: True})) + except Exception: + log.debug( + "OpenLineage selective_enable translation: is_dag_lineage_enabled() failed for dag '%s'", + dag_id, + exc_info=True, + ) + + return rules + + +def _warn_legacy_with_emission_policy(legacy_rules: list[Rule], scope: str) -> None: + """ + Issue a DeprecationWarning listing legacy options that were translated to rules. + + *scope* is ``"task"`` or ``"dag"``: a task-scope warning enumerates every legacy + option (``disabled_for_operators``, ``disable_source_code``, + ``include_full_task_info``, ``selective_enable``); a dag-scope warning only + mentions ``selective_enable`` because the other three are task-only. + """ + if not legacy_rules: + return + + parts: list[str] = [] + if scope == "task": + if _ol_conf.disabled_operators(): + translated = [ + {RULE_SCOPE: {SCOPE_OPERATOR: f}, RULE_CONTROLS: {EMIT: False}} + for f in _ol_conf.disabled_operators() + ] + parts.append(f" - disabled_for_operators -> {translated}") + if not _ol_conf.is_source_enabled(): + parts.append( + ' - disable_source_code -> [{"scope": {}, "controls": {"include_source_code": false}}]' + ) + if _ol_conf.include_full_task_info(): + parts.append( + ' - include_full_task_info -> [{"scope": {}, "controls": {"include_full_task_info": true}}]' + ) + if _ol_conf.selective_enable(): + parts.append( + " - selective_enable -> " + '[{"scope": {}, "controls": {"emit": false}}] (global baseline) ' + "+ per-task opt-in rules injected at runtime from enable_lineage()/disable_lineage()" + ) + else: # dag scope: only selective_enable contributes legacy dag-event rules + if not _ol_conf.selective_enable(): + return + parts.append( + " - selective_enable -> " + '[{"scope": {}, "controls": {"emit": false}}] (global baseline) ' + "+ dag opt-in rules injected at runtime from enable_lineage()/disable_lineage()" + ) + + if not parts: + return + + warnings.warn( + "[openlineage] One or more legacy config options are set and have been translated " + "into equivalent 'emission_policy' rules:\n" + "\n".join(parts) + "\n" + "Migrate to 'emission_policy' exclusively to silence this warning. " + "These legacy options will be removed in a future version.", + AirflowProviderDeprecationWarning, + stacklevel=4, + ) + + +def _compute_locked_fields( + rules: list[Rule], + emit_key: str, + extra_lockable: Sequence[str] = (), +) -> frozenset[str]: + """ + Return the set of fields locked by *any* rule in *rules* that carries ``locked: true``. + + **Floor-lock semantics**: a field is locked if *any* matching rule marks it locked, + regardless of whether that rule wins the value race. + + *emit_key* is the scope-specific emit control (``emit_task_events`` or + ``emit_dag_events``); both it and the generic ``emit`` shorthand lock the ``"emit"`` + policy field. *extra_lockable* lists additional fields to scan (task-event only). + """ + locked: set[str] = set() + for rule in rules: + if not rule.locked: + continue + controls = rule.controls + if emit_key in controls or EMIT in controls: + locked.add(EMIT) + for field in extra_lockable: + if field in controls: + locked.add(field) + return frozenset(locked) + + +def _compute_locked_task_fields( + task_rules: list[Rule], + dag_rules: list[Rule], + operator_rules: list[Rule], + global_rules: list[Rule], +) -> frozenset[str]: + """Compute the set of task-event fields locked by any matching conf rule.""" + return _compute_locked_fields( + task_rules + dag_rules + operator_rules + global_rules, + EMIT_TASK_EVENTS, + _LOCKABLE_TASK_FIELDS, + ) + + +def _apply_authoring_overrides( + config: EmissionPolicy, + locked_fields: frozenset[str], + flags: dict[str, bool], + emit_key: str, + context: str, + extra_fields: tuple[str, ...] = (), +) -> EmissionPolicy: + """ + Core of the authoring-override merge: apply *flags* onto *config*, skipping locked fields. + + *emit_key* is the scope-specific emit control (``emit_task_events`` for tasks, + ``emit_dag_events`` for DAG events). *extra_fields* lists additional lockable + fields that the task path carries but the DAG path does not. + """ + updates: dict[str, bool] = {} + + effective_emit: bool | None = flags.get(emit_key, flags.get(EMIT)) + if effective_emit is not None: + if EMIT in locked_fields: + log.warning( + "OpenLineage emission_policy: extend_global_openlineage_emission_policy call for 'emit' on %s" + " has no effect — locked by conf rule at value %r", + context, + config.emit, + ) + else: + updates[EMIT] = effective_emit + + for field in extra_fields: + if field not in flags: + continue + if field in locked_fields: + log.warning( + "OpenLineage emission_policy: extend_global_openlineage_emission_policy call for '%s' on %s" + " has no effect — locked by conf rule at value %r", + field, + context, + getattr(config, field), + ) + else: + updates[field] = flags[field] + + changed = {k: v for k, v in updates.items() if getattr(config, k) != v} + if not changed: + return config + + _audit_log_authoring_updates(changed, context) + return replace(config, **changed) + + +def _extend_policy_with_task_authoring( + config: EmissionPolicy, + locked_fields: frozenset[str], + operator: AnyOperator, + context: str, +) -> EmissionPolicy: + """ + Apply per-task authoring flags on top of *config*. + + Authoring flags are stored on operator objects by + :func:`~airflow.providers.openlineage.api.emission_policy.extend_global_openlineage_emission_policy`. + Fields listed in *locked_fields* are protected — attempts to override them are + logged at WARNING level and silently ignored. ``emit_task_events`` in the authoring + flags takes precedence over ``emit`` for the task ``emit`` field, mirroring the + same precedence used in conf rule resolution. + """ + flags = _read_param(operator, OL_EMISSION_POLICY_PARAM) + if not flags: + return config + return _apply_authoring_overrides( + config, + locked_fields, + flags, + EMIT_TASK_EVENTS, + context, + extra_fields=(EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE, HOOK_LINEAGE, INCLUDE_FULL_TASK_INFO), + ) + + +def _extend_policy_with_dag_authoring( + config: EmissionPolicy, + locked_fields: frozenset[str], + dag: object, + context: str, +) -> EmissionPolicy: + """ + Apply per-DAG authoring flags on top of *config* for dag-run events. + + ``emit_dag_events`` in the authoring flags takes precedence over ``emit``, + mirroring conf rule resolution. + """ + flags = _read_param(dag, OL_EMISSION_POLICY_PARAM) + if not flags: + return config + return _apply_authoring_overrides(config, locked_fields, flags, EMIT_DAG_EVENTS, context) + + +def _resolve_task_policy_from_conf_only( + rules: Sequence[Rule], + fqcn: str, + dag_id: str, + task_id: str, +) -> tuple[EmissionPolicy, frozenset[str]]: + """ + Pure, deterministic conf resolution for task events. + + Takes the already-parsed, pre-validated *rules* and a pre-computed *fqcn* string. + Rules are validated once at config-read time (:func:`_parse_user_rules`), so this + path performs no re-validation. + """ + task_rules, dag_rules, operator_rules, global_rules = _classify_task_rules(rules, fqcn, dag_id, task_id) + + defaults = EmissionPolicy.defaults() + tiers = [task_rules, dag_rules, operator_rules, global_rules] + context = f"task '{task_id}' in dag '{dag_id}'" + + emit, emit_rule = _resolve_emit_with_source(tiers, "task", defaults.emit) + extract_operator_metadata, em_rule = _resolve_field_with_source( + tiers, EXTRACT_OPERATOR_METADATA, defaults.extract_operator_metadata + ) + include_source_code, isc_rule = _resolve_field_with_source( + tiers, INCLUDE_SOURCE_CODE, defaults.include_source_code + ) + hook_lineage, hl_rule = _resolve_field_with_source(tiers, HOOK_LINEAGE, defaults.hook_lineage) + include_full_task_info, ift_rule = _resolve_field_with_source( + tiers, INCLUDE_FULL_TASK_INFO, defaults.include_full_task_info + ) + + if emit_rule is not None: + _audit_log_conf_field(EMIT, emit, context, emit_rule) + if em_rule is not None: + _audit_log_conf_field(EXTRACT_OPERATOR_METADATA, extract_operator_metadata, context, em_rule) + if isc_rule is not None: + _audit_log_conf_field(INCLUDE_SOURCE_CODE, include_source_code, context, isc_rule) + if hl_rule is not None: + _audit_log_conf_field(HOOK_LINEAGE, hook_lineage, context, hl_rule) + if ift_rule is not None: + _audit_log_conf_field(INCLUDE_FULL_TASK_INFO, include_full_task_info, context, ift_rule) + + locked_fields = _compute_locked_task_fields(task_rules, dag_rules, operator_rules, global_rules) + + return ( + EmissionPolicy( + emit=emit, + extract_operator_metadata=extract_operator_metadata, + include_source_code=include_source_code, + hook_lineage=hook_lineage, + include_full_task_info=include_full_task_info, + ), + locked_fields, + ) + + +def _resolve_dag_policy_from_conf_only( + rules: Sequence[Rule], + dag_id: str, +) -> tuple[EmissionPolicy, frozenset[str]]: + """ + Pure, deterministic conf resolution for dag-run events. + + Returns ``(config, locked_fields)``. Only ``emit`` is meaningful for dag events; + all other fields in :class:`EmissionPolicy` carry their built-in defaults. + + **Floor-lock semantics** (see :func:`_compute_locked_task_fields`): a field + is locked if *any* matching conf rule carries ``locked: true`` for that field. + + ``emission_policy()`` is cached at the conf layer, so calling this directly + is cheap without any additional memoisation. + """ + defaults = EmissionPolicy.defaults() + dag_rules, global_rules = _classify_dag_rules(rules, dag_id) + tiers = [dag_rules, global_rules] + context = f"dag event '{dag_id}'" + + emit, emit_rule = _resolve_emit_with_source(tiers, "dag", defaults.emit) + + if emit_rule is not None: + _audit_log_conf_field(EMIT, emit, context, emit_rule) + + locked_fields = _compute_locked_fields(dag_rules + global_rules, EMIT_DAG_EVENTS) + + return ( + EmissionPolicy( + emit=emit, + extract_operator_metadata=defaults.extract_operator_metadata, + include_source_code=defaults.include_source_code, + hook_lineage=defaults.hook_lineage, + include_full_task_info=defaults.include_full_task_info, + ), + locked_fields, + ) + + +def resolve_task_emission_policy( + operator: AnyOperator, + dag_id: str, + task_id: str, +) -> EmissionPolicy: + """ + Resolve the emission policy for a task-level event. + + This is the **single authoritative entry point** for task event emission decisions. + + Any active legacy options (``disabled_for_operators``, ``disable_source_code``, + ``include_full_task_info``, ``selective_enable``) are *always* translated into equivalent + ``emission_policy`` rules, prepended to the user-provided rules so user rules win within + each tier (last-wins). A ``DeprecationWarning`` is issued whenever a legacy option + contributes a rule — silence it by migrating those options into ``emission_policy`` + exclusively. + + Never raises: any failure (e.g. a misconfigured ``emission_policy``) is logged and the + built-in defaults are returned, so a bad config degrades to "emit with defaults" rather + than crashing the listener's task-notification path. + + :param operator: The Airflow operator/task object. + :param dag_id: The DAG ID for this task instance. + :param task_id: The task ID for this task instance. + """ + try: + context = f"task '{task_id}' in dag '{dag_id}'" + user_rules = _parse_user_rules() + + legacy_rules = _synthesize_legacy_task_rules(operator, dag_id, task_id) + _warn_legacy_with_emission_policy(legacy_rules, scope="task") + + all_rules = list(legacy_rules) + list(user_rules) + from airflow.providers.openlineage.utils.utils import get_fully_qualified_class_name + + fqcn = get_fully_qualified_class_name(operator) + config, locked_fields = _resolve_task_policy_from_conf_only(all_rules, fqcn, dag_id, task_id) + return _extend_policy_with_task_authoring(config, locked_fields, operator, context) + except Exception as err: + # This runs in the scheduler/listener hot path and must never raise — any failure + # (a misconfigured ``emission_policy``, or even a deprecation warning promoted to an + # error via ``-W error``) degrades to "emit with built-in defaults" rather than + # breaking task-event emission. + log.warning( + "Failed to resolve OpenLineage emission policy for task `%s` in dag `%s`; " + "emitting with default controls. Error: %s", + task_id, + dag_id, + err, + ) + log.debug("Exception details:", exc_info=True) + return EmissionPolicy.defaults() + + +def resolve_dag_emission_policy(dag_id: str, dag: object | None = None) -> EmissionPolicy: + """ + Resolve the emission policy for a DAG-level event. + + Returns a full :class:`EmissionPolicy` for a uniform resolver API across task and + DAG scopes. Today only ``emit`` governs DAG-run events; the remaining fields + (``extract_operator_metadata``, ``include_source_code``, ``hook_lineage``, + ``include_full_task_info``) are reserved — they carry their built-in defaults and + are not yet meaningful at DAG scope (no extraction happens at the DAG level). The + unified return type keeps room for future DAG-level controls (e.g. terminal-event + filtering) without re-widening the API. + + Any active legacy options (currently only ``selective_enable`` affects DAG events) + are *always* translated into equivalent ``emission_policy`` rules, prepended to + the user-provided rules so user rules win within each tier (last-wins). A + ``DeprecationWarning`` is issued whenever a legacy option contributes a rule. + + Never raises: any failure (e.g. a misconfigured ``emission_policy``) is logged and the + built-in defaults are returned, so a bad config degrades to "emit with defaults" rather + than crashing the listener's dag-run-notification path. + + :param dag_id: The DAG ID for the dag run event. + :param dag: Optional DAG object used for the ``selective_enable`` check / translation. + """ + try: + context = f"dag event '{dag_id}'" + user_rules = _parse_user_rules() + + legacy_rules = _synthesize_legacy_dag_rules(dag_id, dag) + _warn_legacy_with_emission_policy(legacy_rules, scope="dag") + + all_rules = list(legacy_rules) + list(user_rules) + config, locked_fields = _resolve_dag_policy_from_conf_only(all_rules, dag_id) + + if dag is not None: + config = _extend_policy_with_dag_authoring(config, locked_fields, dag, context) + return config + except Exception as err: + # This runs in the scheduler/listener hot path and must never raise — any failure + # (a misconfigured ``emission_policy``, or even a deprecation warning promoted to an + # error via ``-W error``) degrades to "emit with built-in defaults" rather than + # breaking dag-run-event emission. + log.warning( + "Failed to resolve OpenLineage emission policy for dag `%s`;" + " emitting with default controls. Error: %s", + dag_id, + err, + ) + log.debug("Exception details:", exc_info=True) + return EmissionPolicy.defaults() diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index afc33804cd761..6bf7658754615 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -257,6 +257,7 @@ def build_task_event_run_facets( parent_job_name: str | None = None, dr_conf: dict | None = None, additional_run_facets: dict[str, RunFacet] | None = None, + include_full_task_info: bool = False, ) -> dict[str, RunFacet]: """Build the task-event run-facet dict.""" if dr_conf is None: @@ -287,7 +288,12 @@ def build_task_event_run_facets( dr_conf=dr_conf, ), **get_airflow_run_facet( - dag_run=dag_run, dag=dag, task_instance=task_instance, task=task, task_uuid=task_uuid + dag_run=dag_run, + dag=dag, + task_instance=task_instance, + task=task, + task_uuid=task_uuid, + include_full_task_info=include_full_task_info, ), **get_airflow_debug_facet(), **get_processing_engine_facet(), @@ -1244,6 +1250,7 @@ class TaskInfoComplete(TaskInfo): includes = [] excludes = [ "_BaseOperator__instantiated", + "_BaseOperator__init_kwargs", # Causes recursion error on AF3, nothing useful there "_dag", "_hook", "_log", @@ -1324,6 +1331,7 @@ def get_airflow_run_facet( task_instance: TaskInstance, task: BaseOperator, task_uuid: str, + include_full_task_info: bool = False, ) -> dict[str, AirflowRunFacet]: runtime_assets = get_runtime_outlet_assets(task_instance) return { @@ -1333,7 +1341,7 @@ def get_airflow_run_facet( taskInstance=TaskInstanceInfo(task_instance), task=( TaskInfoComplete(task, runtime_assets=runtime_assets) - if conf.include_full_task_info() + if include_full_task_info else TaskInfo(task, runtime_assets=runtime_assets) ), taskUuid=task_uuid, @@ -1839,8 +1847,14 @@ def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple def _emits_ol_events(task: AnyOperator) -> bool: - config_selective_enabled = is_selective_lineage_enabled(task) - config_disabled_for_operators = is_operator_disabled(task) + from airflow.providers.openlineage.utils.emission_policy import resolve_task_emission_policy + + # resolve_task_emission_policy already incorporates the selective_enable check. + controls = resolve_task_emission_policy( + operator=task, + dag_id=task.dag_id, + task_id=task.task_id, + ) is_task_schedulable_method = getattr(TaskInstance, "is_task_schedulable", None) # Added in 3.2.0 #56039 if is_task_schedulable_method and callable(is_task_schedulable_method): @@ -1870,8 +1884,7 @@ def _emits_ol_events(task: AnyOperator) -> bool: emits_ol_events = all( ( - config_selective_enabled, - not config_disabled_for_operators, + controls.emit, not is_skipped_as_empty_operator, ) ) diff --git a/providers/openlineage/tests/unit/openlineage/api/test_datasets.py b/providers/openlineage/tests/unit/openlineage/api/test_datasets.py index 734b7e5d842b4..3f67c554ea58b 100644 --- a/providers/openlineage/tests/unit/openlineage/api/test_datasets.py +++ b/providers/openlineage/tests/unit/openlineage/api/test_datasets.py @@ -234,3 +234,45 @@ def test_resolves_task_instance_from_context(patched_emit): emit_dataset_lineage(inputs=[Dataset(namespace="ns", name="a")]) get_ti.assert_called_once() patched_emit.assert_called_once() + + +def test_noop_when_emission_policy_blocks_emit(patched_emit): + from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy + + ti = _make_task_instance() + with mock.patch( + f"{_MODULE}.resolve_task_emission_policy", + return_value=EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ), + ): + emit_dataset_lineage(inputs=[Dataset(namespace="ns", name="a")], task_instance=ti) + + patched_emit.assert_not_called() + + +def test_passes_include_full_task_info_to_run_facets(patched_emit): + from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy + + ti = _make_task_instance() + with ( + mock.patch( + f"{_MODULE}.resolve_task_emission_policy", + return_value=EmissionPolicy( + emit=True, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=True, + ), + ), + mock.patch(f"{_MODULE}.build_task_event_run_facets", return_value={}) as mock_build_run_facets, + ): + emit_dataset_lineage(inputs=[Dataset(namespace="ns", name="a")], task_instance=ti) + + mock_build_run_facets.assert_called_once() + assert mock_build_run_facets.call_args.kwargs["include_full_task_info"] is True diff --git a/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py b/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py new file mode 100644 index 0000000000000..57ff1c124a985 --- /dev/null +++ b/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py @@ -0,0 +1,281 @@ +# 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. + +from __future__ import annotations + +import pytest +from pendulum import now + +from airflow.providers.common.compat.sdk import DAG, XComArg, dag, task +from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, +) +from airflow.providers.openlineage.utils.emission_policy import ( + OL_EMISSION_POLICY_PARAM, + _read_param, +) +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.providers.standard.operators.python import PythonOperator + + +def _make_dag_and_task(dag_id: str = "test_dag"): + """Create a minimal DAG with one EmptyOperator task.""" + with DAG(dag_id=dag_id, schedule=None, start_date=now()) as d: + t = EmptyOperator(task_id="test_task") + return d, t + + +def _task_flags(task) -> dict: + """Internal-access helper for tests — read the raw stored task flags.""" + return _read_param(task, OL_EMISSION_POLICY_PARAM) + + +def _dag_flags(dag) -> dict: + """Internal-access helper for tests — read the raw stored dag flags.""" + return _read_param(dag, OL_EMISSION_POLICY_PARAM) + + +class TestExtendOnTask: + def test_emit_false_stored_on_task(self): + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=False) + assert _task_flags(task) == {"emit": False} + + def test_multiple_flags_stored_together(self): + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy( + task, + include_source_code=False, + hook_lineage=False, + extract_operator_metadata=True, + ) + assert _task_flags(task) == { + "include_source_code": False, + "hook_lineage": False, + "extract_operator_metadata": True, + } + + def test_no_flags_provided_returns_obj_unchanged(self): + _, task = _make_dag_and_task() + result = extend_global_openlineage_emission_policy(task) + assert result is task + assert _task_flags(task) == {} + + def test_returns_same_object(self): + _, task = _make_dag_and_task() + result = extend_global_openlineage_emission_policy(task, emit=False) + assert result is task + + def test_successive_calls_merge_flags(self): + """Multiple calls accumulate — later call wins for the same key.""" + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=False, include_source_code=False) + extend_global_openlineage_emission_policy(task, include_source_code=True) + flags = _task_flags(task) + assert flags["emit"] is False + assert flags["include_source_code"] is True + + def test_emit_dag_events_on_task_not_stored(self): + """emit_dag_events on a task is silently discarded (not a task-level flag).""" + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit_dag_events=False) + assert "emit_dag_events" not in _task_flags(task) + + +class TestExtendOnDag: + def test_dag_emit_false_stored_on_dag(self): + dag, _ = _make_dag_and_task() + extend_global_openlineage_emission_policy(dag, emit=False) + assert _dag_flags(dag) == {"emit": False} + + def test_emit_dag_events_stored_on_dag(self): + dag, _ = _make_dag_and_task() + extend_global_openlineage_emission_policy(dag, emit_dag_events=False) + assert _dag_flags(dag) == {"emit_dag_events": False} + + def test_dag_call_propagates_task_flags_to_all_tasks(self): + with DAG(dag_id="prop_dag", schedule=None, start_date=now()) as dag: + t1 = EmptyOperator(task_id="t1") + t2 = EmptyOperator(task_id="t2") + t3 = EmptyOperator(task_id="t3") + + extend_global_openlineage_emission_policy(dag, include_source_code=False) + + for t in (t1, t2, t3): + assert _task_flags(t).get("include_source_code") is False + + def test_dag_emit_propagates_to_all_tasks_and_dag(self): + with DAG(dag_id="emit_dag", schedule=None, start_date=now()) as dag: + t1 = EmptyOperator(task_id="t1") + t2 = EmptyOperator(task_id="t2") + + extend_global_openlineage_emission_policy(dag, emit=False) + + for t in (t1, t2): + assert _task_flags(t).get("emit") is False + assert _dag_flags(dag).get("emit") is False + + def test_emit_dag_events_only_affects_dag_not_tasks(self): + with DAG(dag_id="ede_dag", schedule=None, start_date=now()) as dag: + t1 = EmptyOperator(task_id="t1") + + extend_global_openlineage_emission_policy(dag, emit_dag_events=False) + + assert _dag_flags(dag).get("emit_dag_events") is False + assert "emit_dag_events" not in _task_flags(t1) + + def test_dag_returns_dag_object(self): + dag, _ = _make_dag_and_task() + result = extend_global_openlineage_emission_policy(dag, emit=False) + assert result is dag + + def test_dag_no_tasks_does_not_fail(self): + """Calling on a DAG with no tasks should not raise.""" + with DAG(dag_id="empty_dag", schedule=None, start_date=now()) as dag: + pass + extend_global_openlineage_emission_policy(dag, include_source_code=False) + + +class TestExtendOnXComArg: + def test_xcomarg_delegates_to_operator(self): + with DAG(dag_id="xcom_dag", schedule=None, start_date=now()): + op = PythonOperator(task_id="xcom_task", python_callable=lambda: None) + + xarg = XComArg(op) + result = extend_global_openlineage_emission_policy(xarg, emit=False) + assert result is xarg + assert _task_flags(op).get("emit") is False + + +class TestExtendOnTaskFlow: + """TaskFlow API: ``@task``-decorated calls return an XComArg, ``@dag`` returns a DAG.""" + + def test_task_decorator_call_stores_flags_on_operator(self): + @task + def my_task(): + return None + + with DAG(dag_id="tf_task_dag", schedule=None, start_date=now()): + xarg = my_task() + + result = extend_global_openlineage_emission_policy(xarg, emit=False) + assert result is xarg + assert _task_flags(xarg.operator).get("emit") is False + + def test_task_decorator_multiple_flags(self): + @task + def my_task(): + return None + + with DAG(dag_id="tf_task_multi_dag", schedule=None, start_date=now()): + xarg = my_task() + + extend_global_openlineage_emission_policy( + xarg, + include_source_code=False, + hook_lineage=False, + ) + assert _task_flags(xarg.operator) == { + "include_source_code": False, + "hook_lineage": False, + } + + def test_dag_decorator_propagates_task_flags(self): + @task + def t1(): + return None + + @task + def t2(): + return None + + @dag(schedule=None, start_date=now()) + def my_pipeline(): + t1() + t2() + + dag_obj = my_pipeline() + + extend_global_openlineage_emission_policy(dag_obj, include_source_code=False) + + for task_obj in dag_obj.task_dict.values(): + assert _task_flags(task_obj).get("include_source_code") is False + + def test_dag_decorator_emit_dag_events_stored_on_dag(self): + @task + def t1(): + return None + + @dag(schedule=None, start_date=now()) + def my_pipeline(): + t1() + + dag_obj = my_pipeline() + + extend_global_openlineage_emission_policy(dag_obj, emit_dag_events=False) + + assert _dag_flags(dag_obj).get("emit_dag_events") is False + for task_obj in dag_obj.task_dict.values(): + assert "emit_dag_events" not in _task_flags(task_obj) + + def test_dag_decorator_emit_propagates_to_tasks_and_dag(self): + @task + def t1(): + return None + + @task + def t2(): + return None + + @dag(schedule=None, start_date=now()) + def my_pipeline(): + t1() + t2() + + dag_obj = my_pipeline() + + extend_global_openlineage_emission_policy(dag_obj, emit=False) + + assert _dag_flags(dag_obj).get("emit") is False + for task_obj in dag_obj.task_dict.values(): + assert _task_flags(task_obj).get("emit") is False + + +class TestUnsupportedObjects: + """extend_global_openlineage_emission_policy() rejects anything else. + + There is intentionally no "global authoring scope": deployment-wide changes belong in the + ``emission_policy`` Airflow configuration with an empty ``scope: {}`` rule. + """ + + @pytest.mark.parametrize( + "obj", + [None, 42, "a_string", object(), {"some": "dict"}, [1, 2, 3]], + ids=["None", "int", "str", "object", "dict", "list"], + ) + def test_rejects_non_dag_non_task_object(self, obj): + with pytest.raises(TypeError, match="DAG, an Operator, or an XComArg"): + extend_global_openlineage_emission_policy(obj, emit=False) + + def test_error_message_mentions_global_scope_workaround(self): + with pytest.raises(TypeError, match="emission_policy"): + extend_global_openlineage_emission_policy(object(), emit=False) + + def test_no_flags_provided_still_raises_for_unsupported_type(self): + """Even with no flags, a wrong argument type fails loud rather than silently no-op.""" + with pytest.raises(TypeError): + extend_global_openlineage_emission_policy(object()) diff --git a/providers/openlineage/tests/unit/openlineage/api/test_sql.py b/providers/openlineage/tests/unit/openlineage/api/test_sql.py index 3331555ca313c..565db5d8440f0 100644 --- a/providers/openlineage/tests/unit/openlineage/api/test_sql.py +++ b/providers/openlineage/tests/unit/openlineage/api/test_sql.py @@ -310,3 +310,26 @@ def test_raises_on_context_resolution_failure_when_flag_set(patched_emit): ) patched_emit.assert_not_called() + + +def test_noop_when_emission_policy_blocks_emit(patched_emit): + from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy + + ti = _make_task_instance() + with mock.patch( + f"{_MODULE}.resolve_task_emission_policy", + return_value=EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ), + ): + emit_query_lineage( + query_id="qid-1", + query_source_namespace="snowflake://ACCT", + task_instance=ti, + ) + + patched_emit.assert_not_called() diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_bash.py b/providers/openlineage/tests/unit/openlineage/extractors/test_bash.py index 5e788805cad89..e3ed6f929aff0 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_bash.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_bash.py @@ -19,7 +19,6 @@ import warnings from datetime import datetime -from unittest.mock import patch from openlineage.client.facet_v2 import source_code_job @@ -40,13 +39,11 @@ bash_task = BashOperator(task_id="bash-task", bash_command="ls -halt && exit 0", dag=dag) -@patch("airflow.providers.openlineage.conf.is_source_enabled") -def test_extract_operator_bash_command_disabled(mocked_source_enabled): - mocked_source_enabled.return_value = False +def test_extract_operator_bash_command_disabled(): operator = BashOperator(task_id="taskid", bash_command="exit 0;", env={"A": "1"}, append_env=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", AirflowProviderDeprecationWarning) - result = BashExtractor(operator).extract() + result = BashExtractor(operator, source_code_enabled=False).extract() assert "sourceCode" not in result.job_facets assert "unknownSourceAttribute" in result.run_facets unknown_items = result.run_facets["unknownSourceAttribute"]["unknownItems"] @@ -58,13 +55,11 @@ def test_extract_operator_bash_command_disabled(mocked_source_enabled): assert "task_id" in unknown_items[0]["properties"] -@patch("airflow.providers.openlineage.conf.is_source_enabled") -def test_extract_operator_bash_command_enabled(mocked_source_enabled): - mocked_source_enabled.return_value = True +def test_extract_operator_bash_command_enabled(): operator = BashOperator(task_id="taskid", bash_command="exit 0;", env={"A": "1"}, append_env=True) with warnings.catch_warnings(): warnings.simplefilter("ignore", AirflowProviderDeprecationWarning) - result = BashExtractor(operator).extract() + result = BashExtractor(operator, source_code_enabled=True).extract() assert result.job_facets["sourceCode"] == source_code_job.SourceCodeJobFacet("bash", "exit 0;") assert "unknownSourceAttribute" in result.run_facets unknown_items = result.run_facets["unknownSourceAttribute"]["unknownItems"] diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index 2c02f5709fc6f..d026c1b5039e6 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -316,6 +316,120 @@ def get_openlineage_facets_on_start(self): assert metadata.outputs == [] +@pytest.mark.parametrize("hook_lineage", [True, False]) +def test_extract_metadata_extractor_empty_result_uses_hook_lineage_when_enabled(hook_lineage): + """ + Extractor found but returned empty inputs/outputs: hook lineage should be merged when + ``hook_lineage=True`` and skipped (falling back to inlets/outlets) when ``hook_lineage=False``. + """ + from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy + + class FakeSupportedOperator(BaseOperator): + def execute(self, context: Context) -> Any: + pass + + def get_openlineage_facets_on_complete(self, task_instance): + return OperatorLineage() + + hook_input = OpenLineageDataset(namespace="s3://bucket", name="hook_input") + dagrun = MagicMock() + task = FakeSupportedOperator(task_id="test_task_hook_lineage_policy") + ti = MagicMock() + + controls = EmissionPolicy( + emit=True, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=hook_lineage, + include_full_task_info=False, + ) + + with mock.patch.object( + ExtractorManager, "get_hook_lineage", return_value=OperatorLineage(inputs=[hook_input]) + ): + extractor_manager = ExtractorManager() + metadata = extractor_manager.extract_metadata( + dagrun=dagrun, task=task, task_instance_state=None, task_instance=ti, controls=controls + ) + + if hook_lineage: + assert hook_input in metadata.inputs + else: + assert hook_input not in metadata.inputs + + +@pytest.mark.parametrize("hook_lineage", [True, False]) +def test_extract_metadata_no_extractor_emits_unknown_source_and_inlets_outlets( + hook_lineage, hook_lineage_collector +): + """ + Regression: when no extractor matches and no hook lineage is collected, the no-extractor + fallback must still emit the ``unknownSourceAttribute`` run facet AND extract manually + declared inlets/outlets — for both ``hook_lineage=True`` (default) and ``hook_lineage=False``. + """ + from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy + + inlets = [OpenLineageDataset(namespace="namespace1", name="name1")] + outlets = [OpenLineageDataset(namespace="namespace2", name="name2")] + + class FakeUnsupportedOperator(BaseOperator): + def execute(self, context: Context) -> Any: + pass + + dagrun = MagicMock() + task = FakeUnsupportedOperator(task_id="unsupported_task", inlets=inlets, outlets=outlets) + ti = MagicMock() + + controls = EmissionPolicy( + emit=True, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=hook_lineage, + include_full_task_info=False, + ) + + extractor_manager = ExtractorManager() + metadata = extractor_manager.extract_metadata( + dagrun=dagrun, task=task, task_instance_state=None, task_instance=ti, controls=controls + ) + + assert "unknownSourceAttribute" in metadata.run_facets + assert metadata.inputs == inlets + assert metadata.outputs == outlets + + +def test_get_extractor_supports_legacy_custom_extractor_signature(): + """ + Regression: custom extractors may use the historically-public ``__init__(self, operator)`` + signature. ``_get_extractor`` must construct them without passing ``source_code_enabled`` as + a constructor kwarg (which would raise ``TypeError`` and drop the whole task event), and set + the flag afterward. + """ + from airflow.providers.openlineage.extractors.base import BaseExtractor + + class LegacyExtractor(BaseExtractor): + def __init__(self, operator): # legacy signature - no source_code_enabled kwarg + self.operator = operator + + @classmethod + def get_operator_classnames(cls): + return ["LegacyOperator"] + + def _execute_extraction(self): + return OperatorLineage() + + task = MagicMock() + task.task_type = "LegacyOperator" + + extractor_manager = ExtractorManager() + extractor_manager.add_extractor("LegacyOperator", LegacyExtractor) + + extractor = extractor_manager._get_extractor(task, source_code_enabled=False) + assert isinstance(extractor, LegacyExtractor) + assert extractor.operator is task + assert extractor.source_code_enabled is False + + @pytest.mark.skipif( AIRFLOW_V_3_0_PLUS, reason="Test for hook level lineage in Airflow < 3.0", diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_python.py b/providers/openlineage/tests/unit/openlineage/extractors/test_python.py index abfe10dbe0989..f35dd3e307111 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_python.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_python.py @@ -21,7 +21,6 @@ import os import warnings from datetime import datetime -from unittest.mock import patch from openlineage.client.facet_v2 import source_code_job @@ -60,13 +59,11 @@ def test_extract_source_code(): assert code == CODE -@patch("airflow.providers.openlineage.conf.is_source_enabled") -def test_extract_operator_code_disabled(mocked_source_enabled): - mocked_source_enabled.return_value = False +def test_extract_operator_code_disabled(): operator = PythonOperator(task_id="taskid", python_callable=callable, op_args=(1, 2), op_kwargs={"a": 1}) with warnings.catch_warnings(): warnings.simplefilter("ignore", AirflowProviderDeprecationWarning) - result = PythonExtractor(operator).extract() + result = PythonExtractor(operator, source_code_enabled=False).extract() assert "sourceCode" not in result.job_facets assert "unknownSourceAttribute" in result.run_facets unknown_items = result.run_facets["unknownSourceAttribute"]["unknownItems"] @@ -78,13 +75,11 @@ def test_extract_operator_code_disabled(mocked_source_enabled): assert "task_id" in unknown_items[0]["properties"] -@patch("airflow.providers.openlineage.conf.is_source_enabled") -def test_extract_operator_code_enabled(mocked_source_enabled): - mocked_source_enabled.return_value = True +def test_extract_operator_code_enabled(): operator = PythonOperator(task_id="taskid", python_callable=callable, op_args=(1, 2), op_kwargs={"a": 1}) with warnings.catch_warnings(): warnings.simplefilter("ignore", AirflowProviderDeprecationWarning) - result = PythonExtractor(operator).extract() + result = PythonExtractor(operator, source_code_enabled=True).extract() assert result.job_facets["sourceCode"] == source_code_job.SourceCodeJobFacet("python", CODE) assert "unknownSourceAttribute" in result.run_facets unknown_items = result.run_facets["unknownSourceAttribute"]["unknownItems"] diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index edec29202c0db..f3132746cd660 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json import uuid from collections import defaultdict from collections.abc import Callable @@ -38,6 +39,7 @@ from airflow.providers.openlineage.extractors.base import OperatorLineage from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter from airflow.providers.openlineage.plugins.listener import OpenLineageListener +from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy from airflow.providers.openlineage.utils.selective_enable import disable_lineage, enable_lineage from airflow.utils import types from airflow.utils.state import DagRunState, State @@ -422,7 +424,7 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -454,7 +456,7 @@ def test_adapter_start_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() listener.on_task_instance_running(None, task_instance, None) listener.adapter.start_task.assert_called_once_with( @@ -479,7 +481,7 @@ def test_adapter_start_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -503,7 +505,7 @@ def test_adapter_start_task_is_called_with_dag_owners_when_task_owner_is_default mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner listener.on_task_instance_running(None, task_instance, None) @@ -511,7 +513,7 @@ def test_adapter_start_task_is_called_with_dag_owners_when_task_owner_is_default @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -538,7 +540,7 @@ def test_adapter_start_task_is_called_with_dag_description_when_task_doc_is_empt mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description listener.on_task_instance_running(None, task_instance, None) @@ -547,7 +549,7 @@ def test_adapter_start_task_is_called_with_dag_description_when_task_doc_is_empt @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") @@ -579,7 +581,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() err = ValueError("test") listener.on_task_instance_failed( @@ -607,7 +609,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -631,7 +633,7 @@ def test_adapter_fail_task_is_called_with_dag_owners_when_task_owner_is_default( mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner err = ValueError("test") @@ -642,7 +644,7 @@ def test_adapter_fail_task_is_called_with_dag_owners_when_task_owner_is_default( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -666,7 +668,7 @@ def test_adapter_fail_task_is_called_with_dag_description_when_task_doc_is_empty mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description err = ValueError("test") @@ -678,7 +680,7 @@ def test_adapter_fail_task_is_called_with_dag_description_when_task_doc_is_empty @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") @@ -710,7 +712,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() listener.on_task_instance_success(None, task_instance, None) # This run_id will be different as we did NOT simulate increase of the try_number attribute, @@ -738,7 +740,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -762,7 +764,7 @@ def test_adapter_complete_task_is_called_with_dag_owners_when_task_owner_is_defa mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner listener.on_task_instance_success(None, task_instance, None) @@ -770,7 +772,7 @@ def test_adapter_complete_task_is_called_with_dag_owners_when_task_owner_is_defa @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -794,7 +796,7 @@ def test_adapter_complete_task_is_called_with_dag_description_when_task_doc_is_e mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description listener.on_task_instance_success(None, task_instance, None) @@ -924,7 +926,7 @@ def success_callable(**kwargs): # try_number after task has been executed assert task_instance.try_number == TRY_NUMBER_AFTER_EXECUTION - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_operator( @@ -933,46 +935,64 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) listener.on_task_instance_running(None, task_instance, None) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_operator( self, mock_get_user_provided_run_facets, mock_disabled ): listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) on_task_failed_kwargs = {"error": ValueError("test")} listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs, session=None ) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_operator( self, mock_get_user_provided_run_facets, mock_disabled ): listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) listener.on_task_instance_success(None, task_instance, None) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() @@ -1326,7 +1346,7 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1359,7 +1379,7 @@ def test_adapter_start_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() listener.on_task_instance_running(None, task_instance) listener.adapter.start_task.assert_called_once_with( @@ -1384,7 +1404,180 @@ def test_adapter_start_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_task_hook_emits_with_defaults_when_emission_policy_misconfigured( + self, + mock_get_user_provided_run_facets, + mock_get_airflow_mapped_task_facet, + mock_get_airflow_run_facet, + mock_get_task_parent_run_facet, + mock_debug_facet, + mock_debug_mode, + ): + """A non-list ``emission_policy`` config raises ``ValueError`` deep in policy resolution. + + The task hook must guard the resolve call so the misconfiguration degrades to "emit with + default controls" instead of crashing the task-notification path (i.e. ``start_task`` is + still invoked). + """ + listener, task_instance = self._create_listener_and_task_instance() + mock_get_airflow_mapped_task_facet.return_value = {} + mock_get_user_provided_run_facets.return_value = {} + mock_get_airflow_run_facet.return_value = {} + mock_get_task_parent_run_facet.return_value = {} + mock_debug_facet.return_value = {} + + with conf_vars({("openlineage", "emission_policy"): json.dumps({"not": "a list"})}): + listener.on_task_instance_running(None, task_instance) + + listener.adapter.start_task.assert_called_once() + + @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_task_hook_failed_emits_with_defaults_when_emission_policy_misconfigured( + self, + mock_get_user_provided_run_facets, + mock_get_airflow_mapped_task_facet, + mock_get_airflow_run_facet, + mock_get_task_parent_run_facet, + mock_debug_facet, + mock_debug_mode, + ): + """on_task_instance_failed must also degrade gracefully on a non-list emission_policy.""" + listener, task_instance = self._create_listener_and_task_instance() + mock_get_airflow_mapped_task_facet.return_value = {} + mock_get_user_provided_run_facets.return_value = {} + mock_get_airflow_run_facet.return_value = {} + mock_get_task_parent_run_facet.return_value = {} + mock_debug_facet.return_value = {} + + with conf_vars({("openlineage", "emission_policy"): json.dumps({"not": "a list"})}): + listener.on_task_instance_failed( + previous_state=None, task_instance=task_instance, error=ValueError("boom") + ) + + listener.adapter.fail_task.assert_called_once() + + @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_task_hook_success_emits_with_defaults_when_emission_policy_misconfigured( + self, + mock_get_user_provided_run_facets, + mock_get_airflow_mapped_task_facet, + mock_get_airflow_run_facet, + mock_get_task_parent_run_facet, + mock_debug_facet, + mock_debug_mode, + ): + """on_task_instance_success must also degrade gracefully on a non-list emission_policy.""" + listener, task_instance = self._create_listener_and_task_instance() + mock_get_airflow_mapped_task_facet.return_value = {} + mock_get_user_provided_run_facets.return_value = {} + mock_get_airflow_run_facet.return_value = {} + mock_get_task_parent_run_facet.return_value = {} + mock_debug_facet.return_value = {} + + with conf_vars({("openlineage", "emission_policy"): json.dumps({"not": "a list"})}): + listener.on_task_instance_success(None, task_instance) + + listener.adapter.complete_task.assert_called_once() + + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_extract_operator_metadata_false_skips_extraction( + self, + mock_get_user_provided_run_facets, + mock_get_airflow_mapped_task_facet, + mock_get_airflow_run_facet, + mock_get_task_parent_run_facet, + mock_policy, + ): + """When extract_operator_metadata=False the listener must skip ExtractorManager and emit bare OperatorLineage.""" + listener, task_instance = self._create_listener_and_task_instance() + mock_get_airflow_mapped_task_facet.return_value = {} + mock_get_user_provided_run_facets.return_value = {} + mock_get_airflow_run_facet.return_value = {} + mock_get_task_parent_run_facet.return_value = {} + mock_policy.return_value = EmissionPolicy( + emit=True, + extract_operator_metadata=False, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) + + listener.on_task_instance_running(None, task_instance) + + listener.extractor_manager.extract_metadata.assert_not_called() + listener.adapter.start_task.assert_called_once() + call_task = listener.adapter.start_task.call_args.kwargs["task"] + assert call_task == OperatorLineage() + + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_include_full_task_info_true_forwarded_to_get_airflow_run_facet( + self, + mock_get_user_provided_run_facets, + mock_get_airflow_mapped_task_facet, + mock_get_airflow_run_facet, + mock_get_task_parent_run_facet, + mock_policy, + ): + """When include_full_task_info=True the listener must forward that flag to get_airflow_run_facet.""" + listener, task_instance = self._create_listener_and_task_instance() + mock_get_airflow_mapped_task_facet.return_value = {} + mock_get_user_provided_run_facets.return_value = {} + mock_get_airflow_run_facet.return_value = {} + mock_get_task_parent_run_facet.return_value = {} + mock_policy.return_value = EmissionPolicy( + emit=True, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=True, + ) + + listener.on_task_instance_running(None, task_instance) + + mock_get_airflow_run_facet.assert_called_once() + _, call_kwargs = mock_get_airflow_run_facet.call_args + assert call_kwargs.get("include_full_task_info") is True + + @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) + @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1408,7 +1601,7 @@ def test_adapter_start_task_is_called_with_dag_owners_when_task_owner_is_default mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner listener.on_task_instance_running(None, task_instance) @@ -1417,7 +1610,7 @@ def test_adapter_start_task_is_called_with_dag_owners_when_task_owner_is_default @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1444,7 +1637,7 @@ def test_adapter_start_task_is_called_with_dag_description_when_task_doc_is_empt mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description listener.on_task_instance_running(None, task_instance) @@ -1453,7 +1646,7 @@ def test_adapter_start_task_is_called_with_dag_description_when_task_doc_is_empt @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") @@ -1485,7 +1678,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() err = ValueError("test") listener.on_task_instance_failed(previous_state=None, task_instance=task_instance, error=err) @@ -1511,7 +1704,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1535,7 +1728,7 @@ def test_adapter_fail_task_is_called_with_dag_owners_when_task_owner_is_default( mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner err = ValueError("test") @@ -1546,7 +1739,7 @@ def test_adapter_fail_task_is_called_with_dag_owners_when_task_owner_is_default( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1570,7 +1763,7 @@ def test_adapter_fail_task_is_called_with_dag_description_when_task_doc_is_empty mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description err = ValueError("test") @@ -1648,7 +1841,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments_for_db_task_instance_ @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") @@ -1680,7 +1873,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments( mock_get_airflow_run_facet.return_value = {"airflow": {"task": "..."}} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() listener.on_task_instance_success(None, task_instance) calls = listener.adapter.complete_task.call_args_list @@ -1706,7 +1899,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments( @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1730,7 +1923,7 @@ def test_adapter_complete_task_is_called_with_dag_owners_when_task_owner_is_defa mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.owner = "airflow" # Simulate default owner on task to force fallback to DAG owner listener.on_task_instance_success(None, task_instance) @@ -1740,7 +1933,7 @@ def test_adapter_complete_task_is_called_with_dag_owners_when_task_owner_is_defa @mock.patch("airflow.providers.openlineage.conf.debug_mode", return_value=True) @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet") - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_mapped_task_facet") @@ -1764,7 +1957,7 @@ def test_adapter_complete_task_is_called_with_dag_description_when_task_doc_is_e mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} mock_get_task_parent_run_facet.return_value = {"parent": 4} mock_debug_facet.return_value = {"debug": "packages"} - mock_disabled.return_value = False + mock_disabled.return_value = EmissionPolicy.defaults() task_instance.task.doc_md = None # Simulate lack of task doc to force fallback to DAG description listener.on_task_instance_success(None, task_instance) @@ -1899,7 +2092,7 @@ def test_on_task_instance_success_correctly_calls_openlineage_adapter_run_id_met map_index=-1, ) - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_operator( @@ -1908,46 +2101,64 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} mock_get_airflow_run_facet.return_value = {"airflow_run_facet": 3} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) listener.on_task_instance_running(None, task_instance) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_operator( self, mock_get_user_provided_run_facets, mock_disabled ): listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) on_task_failed_kwargs = {"error": ValueError("test")} listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs ) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_operator( self, mock_get_user_provided_run_facets, mock_disabled ): listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) listener.on_task_instance_success(None, task_instance) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() @@ -1973,17 +2184,23 @@ def test_on_task_instance_skipped_correctly_calls_openlineage_adapter_run_id_met map_index=-1, ) - @mock.patch("airflow.providers.openlineage.plugins.listener.is_operator_disabled") + @mock.patch("airflow.providers.openlineage.plugins.listener.resolve_task_emission_policy") @mock.patch("airflow.providers.openlineage.plugins.listener.get_user_provided_run_facets") def test_listener_on_task_instance_skipped_do_not_call_adapter_when_disabled_operator( self, mock_get_user_provided_run_facets, mock_disabled ): listener, task_instance = self._create_listener_and_task_instance() mock_get_user_provided_run_facets.return_value = {"custom_facet": 2} - mock_disabled.return_value = True + mock_disabled.return_value = EmissionPolicy( + emit=False, + extract_operator_metadata=True, + include_source_code=True, + hook_lineage=True, + include_full_task_info=False, + ) listener.on_task_instance_skipped(previous_state=None, task_instance=task_instance) - mock_disabled.assert_called_once_with(task_instance.task) + mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() listener.adapter.build_task_instance_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() @@ -2162,6 +2379,7 @@ def dummy_callable(): @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests") +@pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning") class TestOpenLineageSelectiveEnableAirflow2: def setup_method(self): date = timezone.datetime(2022, 1, 1) diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py index 4bfc5090230fe..159ac9d316f06 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py @@ -325,9 +325,7 @@ def test_is_operator_disabled(mock_disabled_operators): assert is_operator_disabled(op) is True -@patch("airflow.providers.openlineage.conf.include_full_task_info") -def test_includes_full_task_info(mock_include_full_task_info): - mock_include_full_task_info.return_value = True +def test_includes_full_task_info(): # There should be no 'bash_command' in excludes and it's not in includes - so # it's a good choice for checking TaskInfo vs TaskInfoComplete assert ( @@ -338,13 +336,12 @@ def test_includes_full_task_info(mock_include_full_task_info): MagicMock(), BashOperator(task_id="bash_op", bash_command="sleep 1"), MagicMock(), + include_full_task_info=True, )["airflow"].task ) -@patch("airflow.providers.openlineage.conf.include_full_task_info") -def test_does_not_include_full_task_info(mock_include_full_task_info): - mock_include_full_task_info.return_value = False +def test_does_not_include_full_task_info(): # There should be no 'bash_command' in excludes and it's not in includes - so # it's a good choice for checking TaskInfo vs TaskInfoComplete assert ( @@ -355,10 +352,76 @@ def test_does_not_include_full_task_info(mock_include_full_task_info): MagicMock(), BashOperator(task_id="bash_op", bash_command="sleep 1"), MagicMock(), + include_full_task_info=False, )["airflow"].task ) +@pytest.mark.skipif( + not AIRFLOW_V_3_0_PLUS, reason="__init_kwargs is an Airflow 3 dataclass implementation detail" +) +def test_full_task_info_excludes_init_kwargs(): + """TaskInfoComplete must not expose _BaseOperator__init_kwargs. + + That field holds the raw DAG and TaskGroup objects passed to __init__, which are attrs + classes with circular references (DAG -> TaskGroup -> DAG). Including them causes + RecursionError in attrs.asdict during OpenLineage event emission. + """ + with DAG("test_dag", start_date=datetime.datetime(2025, 1, 1)): + op = BashOperator(task_id="bash_op", bash_command="sleep 1") + task_dict = get_airflow_run_facet( + MagicMock(), + MagicMock(), + MagicMock(), + op, + MagicMock(), + include_full_task_info=True, + )["airflow"].task + assert "_BaseOperator__init_kwargs" not in task_dict + + +@pytest.mark.skipif( + not AIRFLOW_V_3_0_PLUS, reason="DAG and TaskGroup are attrs classes with circular refs in Airflow 3" +) +def test_full_task_info_no_raw_attrs_objects_in_dag_context(): + """No raw attrs objects may appear (recursively) in TaskInfoComplete when the operator is part of a DAG. + + Attrs objects with circular back-references (DAG.task_group -> TaskGroup.dag -> ...) cause + RecursionError in attrs.asdict when the OL client serialises the RunEvent. + """ + import attrs as _attrs + + with DAG("test_dag", start_date=datetime.datetime(2025, 1, 1)): + op = BashOperator(task_id="bash_op", bash_command="sleep 1") + task_dict = get_airflow_run_facet( + MagicMock(), + MagicMock(), + MagicMock(), + op, + MagicMock(), + include_full_task_info=True, + )["airflow"].task + + def _find_attrs_objs(obj, depth=0): + if depth > 10: + return [] + if _attrs.has(type(obj)): + return [obj] + found = [] + if isinstance(obj, dict): + for v in obj.values(): + found.extend(_find_attrs_objs(v, depth + 1)) + elif isinstance(obj, (list, tuple)): + for v in obj: + found.extend(_find_attrs_objs(v, depth + 1)) + return found + + raw_attrs = _find_attrs_objs(task_dict) + assert not raw_attrs, ( + f"TaskInfoComplete contains raw attrs objects that would cause RecursionError: {raw_attrs}" + ) + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks serialization only in 3.0 conditions") def test_serialize_timetable_complex_with_alias(): from airflow.providers.common.compat.assets import AssetAlias, AssetAll, AssetAny diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py b/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py new file mode 100644 index 0000000000000..ee501d4a2b87b --- /dev/null +++ b/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py @@ -0,0 +1,2414 @@ +# 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. + +from __future__ import annotations + +import json +import warnings +from unittest import mock + +import pytest + +from airflow import DAG +from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.providers.openlineage.utils.emission_policy import ( + EmissionPolicy, + Rule, + resolve_dag_emission_policy, + resolve_task_emission_policy, +) +from airflow.providers.openlineage.utils.selective_enable import enable_lineage + +from tests_common.test_utils.compat import EmptyOperator +from tests_common.test_utils.config import conf_vars + + +class MockOperator: + """Minimal stand-in for an Airflow operator.""" + + def __init__(self, module: str = "tests.mock_module", class_name: str = "MockOperator"): + self.__class__.__module__ = module + self.__class__.__qualname__ = class_name + + +def _operator_fqcn(operator) -> str: + return f"{operator.__class__.__module__}.{operator.__class__.__qualname__}" + + +def _resolve_task_controls(rules: list[dict], operator, dag_id: str, task_id: str): + """Resolve task-level controls for ``operator`` (``dag_id``/``task_id``) under ``rules``. + + Every argument is required so each test states exactly which operator/dag/task it + resolves against — there are no hidden defaults that the expected result depends on. + """ + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + return resolve_task_emission_policy(operator, dag_id, task_id) + + +def _resolve_dag_controls(rules: list[dict], dag_id: str): + """Resolve dag-level controls for ``dag_id`` under ``rules`` (no hidden defaults).""" + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + return resolve_dag_emission_policy(dag_id) + + +class TestEmissionPolicyDefaults: + def test_defaults_all_true(self): + cfg = EmissionPolicy.defaults() + assert cfg.emit is True + assert cfg.extract_operator_metadata is True + assert cfg.include_source_code is True + assert cfg.hook_lineage is True + assert cfg.include_full_task_info is False + + def test_frozen(self): + cfg = EmissionPolicy.defaults() + with pytest.raises((TypeError, AttributeError)): + cfg.emit = False # type: ignore[misc] + + +class TestResolveLineageControlsEmpty: + def test_defaults_when_no_rules(self): + with conf_vars({("openlineage", "emission_policy"): "[]"}): + cfg = resolve_task_emission_policy(MockOperator(), "dag", "task") + assert cfg == EmissionPolicy.defaults() + + def test_defaults_when_config_absent(self): + with conf_vars({("openlineage", "emission_policy"): ""}): + cfg = resolve_task_emission_policy(MockOperator(), "dag", "task") + assert cfg == EmissionPolicy.defaults() + + +class TestResolveLineageControlsGlobal: + def test_global_emit_false(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + assert cfg.extract_operator_metadata is True + + def test_global_source_code_false(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"include_source_code": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_source_code is False + + def test_global_hook_lineage_false(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"hook_lineage": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.hook_lineage is False + + def test_global_applies_when_no_match(self): + """Global rule applies to any task when no specific scope matches.""" + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"extract_operator_metadata": False}}], + dag_id="unrelated_dag", + task_id="t1", + operator=MockOperator(), + ) + assert cfg.extract_operator_metadata is False + + def test_last_global_rule_wins(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}, {"scope": {}, "controls": {"emit": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + +class TestResolveLineageControlsOperator: + def test_operator_emit_false(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [{"scope": {"operator": fqcn}, "controls": {"emit": False}}], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_operator_no_match_uses_default(self): + op = MockOperator() + cfg = _resolve_task_controls( + [{"scope": {"operator": "some.other.Operator"}, "controls": {"emit": False}}], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_operator_last_matching_rule_wins(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [ + {"scope": {"operator": fqcn}, "controls": {"emit": False}}, + {"scope": {"operator": fqcn}, "controls": {"emit": True}}, + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + +class TestResolveLineageControlsDag: + def test_dag_extract_metadata_false(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"extract_operator_metadata": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.extract_operator_metadata is False + + def test_dag_no_match(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "other_dag"}, "controls": {"extract_operator_metadata": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.extract_operator_metadata is True + + def test_dag_last_rule_wins_within_tier(self): + cfg = _resolve_task_controls( + [ + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": True}}, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + +class TestResolveLineageControlsTask: + def test_task_emit_false(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag", "task_id": "my_task"}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_task_no_match_different_task(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag", "task_id": "other_task"}, "controls": {"emit": False}}], + task_id="my_task", + operator=MockOperator(), + dag_id="my_dag", + ) + assert cfg.emit is True + + def test_task_no_match_different_dag(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "other_dag", "task_id": "my_task"}, "controls": {"emit": False}}], + dag_id="my_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.emit is True + + +class TestResolveLineageControlsPriority: + def test_task_overrides_dag(self): + cfg = _resolve_task_controls( + [ + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}, + {"scope": {"dag_id": "my_dag", "task_id": "my_task"}, "controls": {"emit": True}}, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_dag_overrides_operator(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [ + {"scope": {"operator": fqcn}, "controls": {"emit": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": True}}, + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_operator_overrides_global(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"emit": False}}, # global + {"scope": {"operator": fqcn}, "controls": {"emit": True}}, + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_global_applied_when_no_specific_match(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_independent_field_resolution(self): + """emit from task tier, source_code from operator tier.""" + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [ + {"scope": {"operator": fqcn}, "controls": {"include_source_code": False}}, + {"scope": {"dag_id": "my_dag", "task_id": "my_task"}, "controls": {"emit": False}}, + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + assert cfg.include_source_code is False + assert cfg.extract_operator_metadata is True + assert cfg.hook_lineage is True + + def test_dag_does_not_override_task(self): + """Dag-level rule does NOT override task-level resolution for the same field.""" + cfg = _resolve_task_controls( + [ + {"scope": {"dag_id": "my_dag", "task_id": "my_task"}, "controls": {"emit": True}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}, # lower priority + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_global_source_code_false_but_dag_enables_it(self): + """More specific dag rule re-enables source_code that global rule disabled.""" + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"include_source_code": False}}, # global + {"scope": {"dag_id": "my_dag"}, "controls": {"include_source_code": True}}, # dag tier wins + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_source_code is True + + +class TestEmitTaskDagEventFlags: + """Tests for emit_task_events / emit_dag_events flags and emit shorthand.""" + + def test_emit_task_events_false_disables_task_emit(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_task_events": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_emit_dag_events_false_does_not_affect_task_resolution(self): + """`emit_dag_events: false` on a dag rule must NOT suppress task events.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_emit_shorthand_disables_task_events(self): + """`emit: false` disables task events (shorthand).""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_emit_task_events_true_overrides_emit_false(self): + """`emit_task_events: true` restores task emission even when `emit: false`.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit": False, "emit_task_events": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_emit_task_events_false_overrides_emit_true(self): + """`emit_task_events: false` suppresses task emission even when `emit: true`.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit": True, "emit_task_events": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + + def test_emit_dag_events_false_disables_dag_event_emit(self): + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": False}}], dag_id="my_dag" + ) + assert cfg.emit is False + + def test_emit_shorthand_disables_dag_events(self): + """`emit: false` on a dag_id rule also disables dag run events (shorthand).""" + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}], dag_id="my_dag" + ) + assert cfg.emit is False + + def test_emit_task_events_false_does_not_affect_dag_event_resolution(self): + """`emit_task_events: false` must NOT suppress dag run events.""" + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_task_events": False}}], dag_id="my_dag" + ) + assert cfg.emit is True + + def test_emit_dag_events_true_overrides_emit_false_for_dag_events(self): + """`emit_dag_events: true` restores dag event emission when `emit: false`.""" + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit": False, "emit_dag_events": True}}], + dag_id="my_dag", + ) + assert cfg.emit is True + + def test_global_emit_false_disables_both_task_and_dag(self): + """Global `emit: false` disables both task events and dag run events.""" + task_cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + dag_cfg = _resolve_dag_controls([{"scope": {}, "controls": {"emit": False}}], dag_id="my_dag") + assert task_cfg.emit is False + assert dag_cfg.emit is False + + def test_global_emit_dag_events_false_disables_dag_events_only(self): + """Global `emit_dag_events: false` leaves task events unaffected.""" + task_cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit_dag_events": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + dag_cfg = _resolve_dag_controls( + [{"scope": {}, "controls": {"emit_dag_events": False}}], dag_id="my_dag" + ) + assert task_cfg.emit is True + assert dag_cfg.emit is False + + def test_global_emit_task_events_false_disables_task_events_only(self): + """Global `emit_task_events: false` leaves dag events unaffected.""" + task_cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit_task_events": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + dag_cfg = _resolve_dag_controls( + [{"scope": {}, "controls": {"emit_task_events": False}}], dag_id="my_dag" + ) + assert task_cfg.emit is False + assert dag_cfg.emit is True + + +class TestResolveDagEventControls: + def test_defaults_when_no_rules(self): + with conf_vars({("openlineage", "emission_policy"): "[]"}): + cfg = resolve_dag_emission_policy("my_dag") + assert cfg == EmissionPolicy.defaults() + + def test_dag_id_emit_dag_events_false(self): + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": False}}], dag_id="my_dag" + ) + assert cfg.emit is False + + def test_dag_id_no_match(self): + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "other_dag"}, "controls": {"emit_dag_events": False}}], dag_id="my_dag" + ) + assert cfg.emit is True + + def test_global_emit_false_applies_to_dag_event(self): + cfg = _resolve_dag_controls([{"scope": {}, "controls": {"emit": False}}], dag_id="my_dag") + assert cfg.emit is False + + def test_dag_id_wins_over_global(self): + cfg = _resolve_dag_controls( + [ + {"scope": {}, "controls": {"emit": False}}, # global + { + "scope": {"dag_id": "my_dag"}, + "controls": {"emit_dag_events": True}, + }, # more specific — wins + ], + dag_id="my_dag", + ) + assert cfg.emit is True + + def test_task_specific_rules_do_not_affect_dag_event_resolution(self): + """dag_id + task_id rules are task-specific and must NOT match dag event resolution.""" + cfg = _resolve_dag_controls( + [{"scope": {"dag_id": "my_dag", "task_id": "t1"}, "controls": {"emit": False}}], dag_id="my_dag" + ) + assert cfg.emit is True + + def test_operator_rules_do_not_affect_dag_event_resolution(self): + cfg = _resolve_dag_controls( + [{"scope": {"operator": "some.Operator"}, "controls": {"emit": False}}], dag_id="my_dag" + ) + assert cfg.emit is True + + def test_extract_metadata_source_code_hook_lineage_always_true(self): + """Non-emit fields are N/A for dag events and always return defaults (True).""" + cfg = _resolve_dag_controls( + [ + { + "scope": {"dag_id": "my_dag"}, + "controls": {"emit_dag_events": False, "extract_operator_metadata": False}, + } + ], + dag_id="my_dag", + ) + assert cfg.extract_operator_metadata is True + assert cfg.include_source_code is True + assert cfg.hook_lineage is True + + +class TestMatchModeRegex: + def test_match_mode_regex_dag_id(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "^my_.*"}, "match_mode": "regex", "controls": {"emit": False}}], + dag_id="my_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.emit is False + + def test_match_mode_regex_dag_id_no_match(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "^other_.*"}, "match_mode": "regex", "controls": {"emit": False}}], + dag_id="my_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.emit is True + + def test_match_mode_regex_operator(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + prefix = fqcn.rsplit(".", 1)[0] + cfg = _resolve_task_controls( + [ + { + "scope": {"operator": f"{prefix}\\..*"}, + "match_mode": "regex", + "controls": {"include_source_code": False}, + } + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_source_code is False + + def test_match_mode_regex_task_id(self): + cfg = _resolve_task_controls( + [ + { + "scope": {"dag_id": "my_dag", "task_id": "my_.*"}, + "match_mode": "regex", + "controls": {"hook_lineage": False}, + } + ], + dag_id="my_dag", + task_id="my_task", + operator=MockOperator(), + ) + assert cfg.hook_lineage is False + + def test_match_mode_regex_task_id_no_match(self): + cfg = _resolve_task_controls( + [ + { + "scope": {"dag_id": "my_dag", "task_id": "other_.*"}, + "match_mode": "regex", + "controls": {"hook_lineage": False}, + } + ], + dag_id="my_dag", + task_id="my_task", + operator=MockOperator(), + ) + assert cfg.hook_lineage is True + + def test_match_mode_regex_dag_event_resolution(self): + cfg = _resolve_dag_controls( + [ + { + "scope": {"dag_id": "^prod_.*"}, + "match_mode": "regex", + "controls": {"emit_dag_events": False}, + } + ], + dag_id="prod_daily", + ) + assert cfg.emit is False + + def test_match_mode_exact_is_default(self): + """Without match_mode, exact matching is used.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "^my_.*"}, "controls": {"emit": False}}], # regex pattern, but exact mode + dag_id="my_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.emit is True # "^my_.*" != "my_dag" under exact match + + def test_match_mode_invalid_warns_and_skips(self): + """Invalid match_mode value → rule is skipped.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "match_mode": "glob", "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_match_mode_invalid_regex_warns_and_skips(self): + """Malformed regex pattern → rule is skipped.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "[invalid"}, "match_mode": "regex", "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + +class TestResolveLineageControlsValidation: + def test_task_id_without_dag_id_ignored(self): + cfg = _resolve_task_controls( + [{"scope": {"task_id": "my_task"}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_operator_with_emit_dag_events_invalid(self): + """operator + emit_dag_events is meaningless → rule ignored.""" + cfg = _resolve_task_controls( + [{"scope": {"operator": "some.Operator"}, "controls": {"emit_dag_events": False, "emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_task_id_with_emit_dag_events_invalid(self): + """task_id + emit_dag_events is meaningless → rule ignored.""" + cfg = _resolve_task_controls( + [ + { + "scope": {"dag_id": "my_dag", "task_id": "my_task"}, + "controls": {"emit_dag_events": False, "emit": False}, + } + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_non_bool_control_value_skipped(self): + """Non-bool value inside controls → rule is dropped entirely.""" + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": "yes"}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True # rule rejected, defaults remain + + def test_non_dict_rule_ignored(self): + cfg = _resolve_task_controls( + ["not_a_dict", {"scope": {}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False # valid rule still applies + + +class TestAuditLogging: + """Audit-log assertions pin the exact format string, field, context, and source for every call.""" + + _CONF_FMT = "OpenLineage emission policy: '%s' %s for %s by %r" + _AUTH_FMT = ( + "OpenLineage emission policy: '%s' %s for %s " + "by manual `extend_global_openlineage_emission_policy` call." + ) + + def test_audit_log_info_when_emit_disabled_by_rule(self): + op = MockOperator() + rules = [{"scope": {}, "controls": {"emit": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(op, "dag", "task") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "emit", + "disabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"emit": False}, match_mode="exact", locked=False), + ) + + def test_audit_log_info_when_source_code_disabled(self): + op = MockOperator() + rule = {"scope": {}, "controls": {"include_source_code": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(op, "dag", "task") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"include_source_code": False}, match_mode="exact", locked=False), + ) + + def test_no_audit_log_when_all_defaults(self): + """No INFO log when no rule touches any field.""" + op = MockOperator() + with conf_vars({("openlineage", "emission_policy"): "[]"}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(op, "dag", "task") + mock_log.info.assert_not_called() + + def test_audit_log_includes_winning_rule(self): + """The parsed Rule object (with scope/controls) is the source argument to log.info.""" + op = MockOperator() + rule = {"scope": {"dag_id": "audit_dag"}, "controls": {"extract_operator_metadata": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(op, "audit_dag", "t") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "extract_operator_metadata", + "disabled", + "task 't' in dag 'audit_dag'", + Rule( + scope={"dag_id": "audit_dag"}, + controls={"extract_operator_metadata": False}, + match_mode="exact", + locked=False, + ), + ) + + def test_audit_log_dag_event_emit_disabled(self): + """INFO log emitted when dag event emit is disabled by a rule.""" + rules = [{"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_dag_emission_policy("my_dag") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "emit", + "disabled", + "dag event 'my_dag'", + Rule( + scope={"dag_id": "my_dag"}, + controls={"emit_dag_events": False}, + match_mode="exact", + locked=False, + ), + ) + + def test_audit_log_selective_enable_suppression(self): + """The translated emit:false baseline from selective_enable is logged as disabled.""" + dag = DAG(dag_id="test_se_dag", schedule=None) + task = EmptyOperator(task_id="t", dag=dag) + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + resolve_task_emission_policy(task, "test_se_dag", "t") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "emit", + "disabled", + "task 't' in dag 'test_se_dag'", + Rule(scope={}, controls={"emit": False}, match_mode="exact", locked=False), + ) + + def test_audit_log_dag_event_selective_enable_suppression(self): + """The translated emit:false rule is logged for dag event suppression.""" + dag = DAG(dag_id="se_dag2", schedule=None) + EmptyOperator(task_id="t", dag=dag) + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + resolve_dag_emission_policy("se_dag2", dag=dag) + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "emit", + "disabled", + "dag event 'se_dag2'", + Rule(scope={}, controls={"emit": False}, match_mode="exact", locked=False), + ) + + def test_audit_log_conf_reenable_by_higher_tier(self): + """Dag-tier rule re-enabling a globally-disabled field is logged as 'enabled'. + + The dag-tier rule is the winner; the global rule is not the winner so it is + not logged. Previously only the non-default final value triggered logging, + producing zero INFO logs for this combination. + """ + rules = [ + {"scope": {}, "controls": {"include_source_code": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"include_source_code": True}}, + ] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(MockOperator(), "my_dag", "task") + assert cfg.include_source_code is True + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "enabled", + "task 'task' in dag 'my_dag'", + Rule( + scope={"dag_id": "my_dag"}, + controls={"include_source_code": True}, + match_mode="exact", + locked=False, + ), + ) + + def test_audit_log_authoring_reenable_logs_enabled(self): + """Authoring re-enabling a conf-disabled field produces a second 'enabled' log in order. + + Previously the authoring layer only logged non-default values, so restoring + a field to its default via authoring was silent — the last visible log said + 'disabled', which was misleading. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + rules = [{"scope": {}, "controls": {"include_source_code": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is True + assert mock_log.info.call_count == 2 + mock_log.info.assert_has_calls( + [ + mock.call( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + Rule(scope={}, controls={"include_source_code": False}, match_mode="exact", locked=False), + ), + mock.call( + self._AUTH_FMT, + "include_source_code", + "enabled", + "task 'test_task' in dag 'test_dag'", + ), + ] + ) + + def test_no_duplicate_audit_log_when_authoring_matches_conf(self): + """Authoring that agrees with conf does not produce a second log for the same field.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) + + rules = [{"scope": {}, "controls": {"include_source_code": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + Rule(scope={}, controls={"include_source_code": False}, match_mode="exact", locked=False), + ) + + def test_conf_explicit_default_value_is_logged(self): + """A conf rule that explicitly sets a field to its default value is still logged. + + Any rule that touches a field is now logged — even when the resulting value + equals the built-in default — so the operator can see that a rule is active. + """ + rule = {"scope": {}, "controls": {"include_source_code": True}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(MockOperator(), "dag", "task") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "enabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"include_source_code": True}, match_mode="exact", locked=False), + ) + + def test_conf_two_same_tier_same_value_one_log_from_last_rule(self): + """Two global rules setting the same field to the same value produce exactly one log. + + The second (last-in-tier) rule is the winner and is the source in the log. + No duplication occurs even though both rules match. + """ + rule1 = {"scope": {}, "controls": {"include_source_code": False}} + rule2 = {"scope": {}, "controls": {"include_source_code": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule1, rule2])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(MockOperator(), "dag", "task") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"include_source_code": False}, match_mode="exact", locked=False), + ) + + def test_conf_contradictory_same_tier_winner_logged_loser_silent(self): + """A contradictory pair in the same tier fires a WARNING and logs only the winner. + + Last-in-tier wins (the True rule), which is logged as 'enabled'. + The losing rule (False) is not logged — only the winner's effect is audited. + """ + rule_false = {"scope": {}, "controls": {"include_source_code": False}} + rule_true = {"scope": {}, "controls": {"include_source_code": True}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule_false, rule_true])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(MockOperator(), "dag", "task") + assert cfg.include_source_code is True + assert mock_log.info.call_count == 1 + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "enabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"include_source_code": True}, match_mode="exact", locked=False), + ) + assert any("contradictory" in str(c) for c in mock_log.warning.call_args_list) + + def test_conf_single_rule_multiple_fields_logs_each_in_fixed_order(self): + """A rule with two control flags produces two INFO logs in resolution order. + + Resolution order is fixed by the code: emit first, then include_source_code. + """ + rule = {"scope": {}, "controls": {"emit": False, "include_source_code": False}} + winning_rule = Rule( + scope={}, controls={"emit": False, "include_source_code": False}, match_mode="exact", locked=False + ) + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(MockOperator(), "dag", "task") + assert mock_log.info.call_count == 2 + mock_log.info.assert_has_calls( + [ + mock.call(self._CONF_FMT, "emit", "disabled", "task 'task' in dag 'dag'", winning_rule), + mock.call( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'task' in dag 'dag'", + winning_rule, + ), + ] + ) + + def test_conf_include_full_task_info_enabled_logged(self): + """A conf rule enabling include_full_task_info (default False) is logged as 'enabled'.""" + rule = {"scope": {}, "controls": {"include_full_task_info": True}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(MockOperator(), "dag", "task") + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_full_task_info", + "enabled", + "task 'task' in dag 'dag'", + Rule(scope={}, controls={"include_full_task_info": True}, match_mode="exact", locked=False), + ) + + def test_authoring_only_non_default_produces_log(self): + """With no conf rule, authoring that changes a field from its default is logged.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) + + with conf_vars({("openlineage", "emission_policy"): "[]"}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + mock_log.info.assert_called_once_with( + self._AUTH_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + ) + + def test_authoring_only_same_as_default_no_log(self): + """Authoring that sets a field to the built-in default (no conf rule) produces no log. + + include_source_code defaults to True; authoring it to True is a no-op against + the conf-resolved state, so no change and no log. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + with conf_vars({("openlineage", "emission_policy"): "[]"}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + resolve_task_emission_policy(task, "test_dag", "test_task") + + mock_log.info.assert_not_called() + + def test_authoring_two_calls_net_matches_conf_no_extra_log(self): + """Two authoring calls are merged; the NET result is compared against conf, not each call. + + First call: include_source_code=True (enables back to default). + Second call: include_source_code=False (disables again). + Net merged authoring: False — which matches the conf rule's False — so no authoring log. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + extend_global_openlineage_emission_policy(task, include_source_code=False) + + rules = [{"scope": {}, "controls": {"include_source_code": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + mock_log.info.assert_called_once_with( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + Rule(scope={}, controls={"include_source_code": False}, match_mode="exact", locked=False), + ) + + def test_authoring_emit_task_events_key_logged_as_emit_field(self): + """Authoring via emit_task_events=False is logged under the resolved field name 'emit'. + + The authoring layer maps emit_task_events to the emit field in EmissionPolicy; + the log therefore records 'emit', not the raw key name 'emit_task_events'. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit_task_events=False) + + with conf_vars({("openlineage", "emission_policy"): "[]"}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False + mock_log.info.assert_called_once_with( + self._AUTH_FMT, + "emit", + "disabled", + "task 'test_task' in dag 'test_dag'", + ) + + def test_conf_enables_explicit_then_authoring_disables(self): + """Conf logs 'enabled' (explicit default rule), authoring then logs 'disabled'. + + Two logs in order: conf 'enabled', authoring 'disabled'. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) + + rules = [{"scope": {}, "controls": {"include_source_code": True}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + assert mock_log.info.call_count == 2 + mock_log.info.assert_has_calls( + [ + mock.call( + self._CONF_FMT, + "include_source_code", + "enabled", + "task 'test_task' in dag 'test_dag'", + Rule(scope={}, controls={"include_source_code": True}, match_mode="exact", locked=False), + ), + mock.call( + self._AUTH_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + ), + ] + ) + + def test_no_conf_for_field_authoring_same_as_default_no_log(self): + """When conf has no rule for a field and authoring sets it to the default, no log is emitted. + + include_source_code resolves to True by default; authoring True is a no-op. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + rules = [{"scope": {}, "controls": {"emit": False}}] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is True + emit_logs = [c for c in mock_log.info.call_args_list if c.args[1] == "include_source_code"] + assert len(emit_logs) == 0 + + def test_conf_higher_tier_re_enables_then_authoring_disables(self): + """Global conf disables, dag-tier conf re-enables (logs 'enabled'), authoring then disables. + + Three-step trace: global disable (not logged — not the winner), dag-tier enabled (logged), + authoring disabled (logged). Two INFO logs total in order. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) + + rules = [ + {"scope": {}, "controls": {"include_source_code": False}}, + {"scope": {"dag_id": "test_dag"}, "controls": {"include_source_code": True}}, + ] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + assert mock_log.info.call_count == 2 + mock_log.info.assert_has_calls( + [ + mock.call( + self._CONF_FMT, + "include_source_code", + "enabled", + "task 'test_task' in dag 'test_dag'", + Rule( + scope={"dag_id": "test_dag"}, + controls={"include_source_code": True}, + match_mode="exact", + locked=False, + ), + ), + mock.call( + self._AUTH_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + ), + ] + ) + + def test_two_conf_fields_authoring_re_enables_one_three_logs_in_order(self): + """Conf disables two fields; authoring re-enables one of them. + + Resolution order for the conf layer is fixed: emit first, then include_source_code. + Authoring log follows after both conf logs. + Total: 3 INFO logs in order. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + rule = {"scope": {}, "controls": {"emit": False, "include_source_code": False}} + winning_rule = Rule( + scope={}, + controls={"emit": False, "include_source_code": False}, + match_mode="exact", + locked=False, + ) + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False + assert cfg.include_source_code is True + assert mock_log.info.call_count == 3 + mock_log.info.assert_has_calls( + [ + mock.call( + self._CONF_FMT, "emit", "disabled", "task 'test_task' in dag 'test_dag'", winning_rule + ), + mock.call( + self._CONF_FMT, + "include_source_code", + "disabled", + "task 'test_task' in dag 'test_dag'", + winning_rule, + ), + mock.call( + self._AUTH_FMT, + "include_source_code", + "enabled", + "task 'test_task' in dag 'test_dag'", + ), + ] + ) + + +class TestContradictoryRules: + def test_contradictory_global_emit_warns_and_uses_last(self): + """Two global rules set emit to opposite values → warning, last wins.""" + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}, {"scope": {}, "controls": {"emit": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True # last rule wins + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("contradictory" in c for c in warn_calls) + + def test_contradictory_dag_scoped_emit_warns_and_uses_last(self): + """Two dag-scoped rules set emit to opposite values in the same tier.""" + rules = [ + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": True}}, + ] + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls(rules, operator=MockOperator(), dag_id="my_dag", task_id="my_task") + assert cfg.emit is True + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("contradictory" in c for c in warn_calls) + + def test_same_value_twice_does_not_warn(self): + """Redundant rules with the same value should not produce a contradictory warning.""" + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"emit": False}}, {"scope": {}, "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is False + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert not any("contradictory" in c for c in warn_calls) + + def test_contradictory_in_different_tiers_does_not_warn(self): + """Rules in different tiers are resolved by priority, not considered contradictory.""" + # global emit=false, dag emit=true — different tiers, dag wins, no contradiction warning + rules = [ + {"scope": {}, "controls": {"emit": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit": True}}, + ] + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls(rules, operator=MockOperator(), dag_id="my_dag", task_id="my_task") + assert cfg.emit is True # dag tier beats global tier + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert not any("contradictory" in c for c in warn_calls) + + def test_contradictory_field_warns_for_non_emit_fields(self): + """Contradiction warning also applies to non-emit fields like source_code.""" + rules = [ + {"scope": {}, "controls": {"include_source_code": True}}, + {"scope": {}, "controls": {"include_source_code": False}}, + ] + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls(rules, operator=MockOperator(), dag_id="my_dag", task_id="my_task") + assert cfg.include_source_code is False + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("contradictory" in c and "include_source_code" in c for c in warn_calls) + + def test_contradictory_dag_event_emit_warns(self): + """Contradiction in dag-event resolution triggers the warning.""" + rules = [ + {"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": False}}, + {"scope": {"dag_id": "my_dag"}, "controls": {"emit_dag_events": True}}, + ] + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_dag_controls(rules, dag_id="my_dag") + assert cfg.emit is True + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("contradictory" in c for c in warn_calls) + + +class TestResolveLineageControlsWithLegacy: + def test_uses_legacy_disabled_operators_when_no_emission_policy(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "disabled_for_operators"): fqcn, + } + ): + # Mixing legacy options with the resolver emits a deprecation warning; capture and + # assert it here so it does not propagate (the test env promotes it to an error). + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg.emit is False + + def test_uses_legacy_source_code_when_no_emission_policy(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "disable_source_code"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg.include_source_code is False + + def test_emission_policy_takes_precedence_over_disabled_operators(self): + """An explicit operator-tier emission_policy rule wins over the translated legacy rule. + + Legacy disabled_for_operators becomes an operator-tier emit:false rule. Within the same + (operator) tier, last-wins applies, so the user's explicit emission_policy rule overrides + the translated legacy rule. + """ + op = MockOperator() + fqcn = _operator_fqcn(op) + with conf_vars( + { + # Explicit operator-tier rule re-enables this operator + ("openlineage", "emission_policy"): json.dumps( + [{"scope": {"operator": fqcn}, "controls": {"emit": True}}] + ), + ("openlineage", "disabled_for_operators"): fqcn, + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg.emit is True + assert any(issubclass(w.category, DeprecationWarning) for w in caught) + + def test_emission_policy_takes_precedence_over_disable_source_code(self): + op = MockOperator() + with conf_vars( + { + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"include_source_code": true}}]', + ("openlineage", "disable_source_code"): "True", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg.include_source_code is True + assert any(issubclass(w.category, DeprecationWarning) for w in caught) + + def test_no_deprecation_warning_when_legacy_configs_at_default(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): '[{"scope": {}, "controls": {"emit": false}}]', + ("openlineage", "disabled_for_operators"): "", + ("openlineage", "disable_source_code"): "False", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg.emit is False + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert not dep_warnings + + def test_defaults_returned_when_both_absent(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "disabled_for_operators"): "", + ("openlineage", "disable_source_code"): "False", + } + ): + cfg = resolve_task_emission_policy(op, "dag", "task") + assert cfg == EmissionPolicy.defaults() + + +def _make_dag_and_task(dag_id: str = "test_dag", task_id: str = "test_task"): + """Helper to create a real DAG + EmptyOperator for selective_enable tests.""" + dag = DAG(dag_id=dag_id, schedule=None) + task = EmptyOperator(task_id=task_id, dag=dag) + return dag, task + + +class TestSelectiveEnableInTaskControls: + """selective_enable is folded into resolve_task_emission_policy.""" + + def test_selective_enable_off_task_resolution(self): + _, task = _make_dag_and_task() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "False", + } + ): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is True + + def test_selective_enable_on_task_not_opted_in(self): + _, task = _make_dag_and_task() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is False + + def test_selective_enable_on_task_opted_in(self): + _, task = _make_dag_and_task() + enable_lineage(task) + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is True + + def test_selective_enable_does_not_override_already_false_emit(self): + """If emit is already False from emission_policy, selective_enable check is not applied.""" + _, task = _make_dag_and_task() + enable_lineage(task) # task IS opted in + with conf_vars( + { + ("openlineage", "emission_policy"): json.dumps( + [{"scope": {"dag_id": "test_dag", "task_id": "test_task"}, "controls": {"emit": False}}] + ), + ("openlineage", "selective_enable"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is False + + def test_other_fields_preserved_when_selective_enable_forces_emit_false(self): + _, task = _make_dag_and_task() + with conf_vars( + { + ("openlineage", "emission_policy"): json.dumps( + [{"scope": {}, "controls": {"include_source_code": False, "hook_lineage": False}}] + ), + ("openlineage", "selective_enable"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is False + assert cfg.include_source_code is False + assert cfg.hook_lineage is False + assert cfg.extract_operator_metadata is True + + +class TestSelectiveEnableInDagEventControls: + """selective_enable is folded into resolve_dag_emission_policy via the dag= parameter.""" + + def test_no_dag_object_skips_selective_check(self): + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + cfg = resolve_dag_emission_policy("any_dag") + assert cfg.emit is True + + def test_selective_enable_off_dag_event(self): + dag, _ = _make_dag_and_task() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "False", + } + ): + cfg = resolve_dag_emission_policy("test_dag", dag=dag) + assert cfg.emit is True + + def test_selective_enable_on_dag_not_opted_in(self): + dag, _ = _make_dag_and_task() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_dag_emission_policy("test_dag", dag=dag) + assert cfg.emit is False + + def test_selective_enable_on_dag_opted_in(self): + dag, _ = _make_dag_and_task() + enable_lineage(dag) + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "selective_enable"): "True", + } + ): + cfg = resolve_dag_emission_policy("test_dag", dag=dag) + assert cfg.emit is True + + def test_selective_enable_does_not_override_emission_policy_emit_false(self): + """If emission_policy suppresses dag event, selective opt-in doesn't restore it.""" + dag, _ = _make_dag_and_task() + enable_lineage(dag) + with conf_vars( + { + ("openlineage", "emission_policy"): json.dumps( + [{"scope": {"dag_id": "test_dag"}, "controls": {"emit_dag_events": False}}] + ), + ("openlineage", "selective_enable"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_dag_emission_policy("test_dag", dag=dag) + assert cfg.emit is False + + +class TestIncludeFullTaskInfo: + def test_default_is_false(self): + cfg = _resolve_task_controls([], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + assert cfg.include_full_task_info is False + + def test_global_include_full_task_info_true(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"include_full_task_info": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is True + + def test_dag_scope_include_full_task_info_true(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "my_dag"}, "controls": {"include_full_task_info": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is True + + def test_dag_scope_does_not_affect_other_dag(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": "other_dag"}, "controls": {"include_full_task_info": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is False + + def test_task_scope_include_full_task_info_true(self): + cfg = _resolve_task_controls( + [ + { + "scope": {"dag_id": "my_dag", "task_id": "my_task"}, + "controls": {"include_full_task_info": True}, + } + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is True + + def test_task_scope_overrides_global_false(self): + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"include_full_task_info": False}}, + { + "scope": {"dag_id": "my_dag", "task_id": "my_task"}, + "controls": {"include_full_task_info": True}, + }, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is True + + def test_operator_scope_include_full_task_info_true(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [{"scope": {"operator": fqcn}, "controls": {"include_full_task_info": True}}], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_full_task_info is True + + def test_include_full_task_info_does_not_affect_emit(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"include_full_task_info": True}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + assert cfg.extract_operator_metadata is True + + +class TestLegacyTranslation: + """When emission_policy is set, legacy options are translated to rules.""" + + def test_disabled_for_operators_translated_when_emission_policy_set(self): + op = MockOperator() + fqcn = _operator_fqcn(op) + with conf_vars( + { + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"extract_operator_metadata": false}}]', + ("openlineage", "disabled_for_operators"): fqcn, + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + # Operator is in disabled_for_operators → translated to emit:false rule → emit should be False + assert cfg.emit is False + # The emission_policy rule also applies + assert cfg.extract_operator_metadata is False + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert dep_warnings, "Expected DeprecationWarning for mixed legacy+emission_policy" + + def test_disable_source_code_translated_when_emission_policy_set(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): '[{"scope": {}, "controls": {"hook_lineage": false}}]', + ("openlineage", "disable_source_code"): "True", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + assert cfg.include_source_code is False # from translated legacy + assert cfg.hook_lineage is False # from emission_policy + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert dep_warnings + + def test_include_full_task_info_translated_when_emission_policy_set(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): '[{"scope": {}, "controls": {"emit": true}}]', + ("openlineage", "include_full_task_info"): "True", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + assert cfg.include_full_task_info is True # from translated legacy + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert dep_warnings + + def test_emission_policy_rule_overrides_translated_legacy_rule(self): + """User's explicit emission_policy rule wins over translated legacy (last-wins in tier).""" + op = MockOperator() + with conf_vars( + { + # emission_policy explicitly sets include_source_code=True + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"include_source_code": true}}]', + # legacy sets disable_source_code=True (translated to source_code=False global rule) + ("openlineage", "disable_source_code"): "True", + } + ): + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + # User rule (include_source_code=True) appears AFTER the translated legacy rule in the combined list + # and is in the same global tier → user wins (last-wins within tier) + assert cfg.include_source_code is True + + def test_no_deprecation_warning_when_only_emission_policy_set(self): + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): '[{"scope": {}, "controls": {"emit": false}}]', + ("openlineage", "disabled_for_operators"): "", + ("openlineage", "disable_source_code"): "False", + ("openlineage", "include_full_task_info"): "False", + ("openlineage", "selective_enable"): "False", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + assert cfg.emit is False + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert not dep_warnings + + def test_selective_enable_translated_with_emission_policy_set(self): + """selective_enable=True translates to global emit:false + task opt-in rule.""" + _, task = _make_dag_and_task() + enable_lineage(task) + with conf_vars( + { + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"extract_operator_metadata": false}}]', + ("openlineage", "selective_enable"): "True", + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + # Task is opted in → emit should be True (task-tier rule overrides global emit:false) + assert cfg.emit is True + # emission_policy rule still applies + assert cfg.extract_operator_metadata is False + dep_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)] + assert dep_warnings + + def test_selective_enable_translated_not_opted_in_with_emission_policy_set(self): + """Task not opted in under selective_enable + emission_policy → emit False.""" + _, task = _make_dag_and_task() + # task is NOT opted in + with conf_vars( + { + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"include_source_code": false}}]', + ("openlineage", "selective_enable"): "True", + } + ): + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + assert cfg.emit is False + assert cfg.include_source_code is False + + def test_include_full_task_info_legacy_path_unchanged(self): + """When emission_policy is empty, include_full_task_info= uses legacy conf directly.""" + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "include_full_task_info"): "True", + } + ): + with pytest.warns(AirflowProviderDeprecationWarning): + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + assert cfg.include_full_task_info is True + + def test_include_full_task_info_legacy_path_default_false(self): + """When emission_policy is empty and include_full_task_info not set, default False.""" + op = MockOperator() + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "include_full_task_info"): "False", + } + ): + cfg = resolve_task_emission_policy(op, "my_dag", "my_task") + assert cfg.include_full_task_info is False + + +class TestLockedField: + """``locked: true`` on a conf rule prevents per-task authoring from overriding that field.""" + + def test_locked_emit_blocks_authoring_override(self): + """A task-level emit=True from extend_global_openlineage_emission_policy cannot override a locked emit=False rule.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) # authoring says True + + rule = {"scope": {}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False # locked conf rule wins + + def test_unlocked_emit_allows_authoring_override(self): + """Without locked, an authoring emit=True overrides a conf emit=False rule.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) # authoring says True + + rule = {"scope": {}, "controls": {"emit": False}} # no locked + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is True # authoring overrides + + def test_locked_source_code_blocks_authoring(self): + """Locked include_source_code=False in conf cannot be re-enabled by extend_global_openlineage_emission_policy.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + rule = {"scope": {}, "locked": True, "controls": {"include_source_code": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False # locked + + def test_locked_field_does_not_block_other_fields(self): + """A locked rule only protects the field(s) it carries, not unrelated fields.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True, hook_lineage=False) + + # Only emit is locked; include_source_code authoring should still work + rule = {"scope": {}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False # locked + assert cfg.hook_lineage is False # authoring applied (not locked) + + def test_locked_dag_emit_blocks_dag_authoring(self): + """locked emit on dag rule blocks extend_global_openlineage_emission_policy(dag, emit=True) for dag events.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + dag, _ = _make_dag_and_task("locked_dag_test") + extend_global_openlineage_emission_policy(dag, emit=True) # authoring says True for dag events + + rule = {"scope": {"dag_id": "locked_dag_test"}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_dag_emission_policy("locked_dag_test", dag) + + assert cfg.emit is False # locked conf rule wins + + def test_locked_invalid_type_rule_ignored(self): + """A rule with locked of non-bool type is invalid and ignored.""" + cfg = _resolve_task_controls( + [{"scope": {}, "locked": "yes", "controls": {"emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + # Rule is invalid → falls back to default + assert cfg.emit is True + + def test_locked_no_authoring_has_no_effect(self): + """locked with no authoring flags in play still resolves correctly.""" + _, task = _make_dag_and_task() + # No extend_global_openlineage_emission_policy call + + rule = {"scope": {}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False # locked rule still applies normally + + def test_floor_lock_lower_tier_locked_rule_blocks_authoring_despite_higher_tier_override(self): + """Floor lock: a global locked rule blocks authoring even when a dag-tier conf rule wins. + + Scenario: + - Global conf rule: include_source_code=False, locked=True (admin mandate, lower tier) + - Dag-tier conf rule: include_source_code=True (more specific, no lock — wins value) + - Authoring: extend_global_openlineage_emission_policy(task, include_source_code=False) — should be blocked + + The dag-tier rule wins for VALUE (include_source_code=True), but the global locked rule adds + "include_source_code" to locked_fields, so the authoring override cannot touch it. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) # authoring wants False + + rules = [ + { + "scope": {}, + "locked": True, + "controls": {"include_source_code": False}, + }, # global, locked — lower tier + { + "scope": {"dag_id": "test_dag"}, + "controls": {"include_source_code": True}, + }, # dag-tier, no lock — wins value + ] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + # Dag-tier rule wins: include_source_code=True (from conf tier resolution) + assert cfg.include_source_code is True + # But authoring is blocked by floor lock — include_source_code stays at conf-resolved value + + def test_floor_lock_regex_rule_blocks_authoring_for_matching_dag(self): + """Floor lock: a regex-based locked rule blocks authoring for all matching dags.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task("prod_reporting") + extend_global_openlineage_emission_policy(task, include_source_code=True) # authoring wants True + + rules = [ + { + "scope": {"dag_id": "^prod_.*"}, + "match_mode": "regex", + "locked": True, + "controls": {"include_source_code": False}, + } + ] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + cfg = resolve_task_emission_policy(task, "prod_reporting", "test_task") + + assert cfg.include_source_code is False # locked regex rule wins + + def test_floor_lock_regex_rule_does_not_block_non_matching_dag(self): + """Floor lock: a regex-based locked rule does NOT affect non-matching dags.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task("dev_reporting") + extend_global_openlineage_emission_policy(task, include_source_code=False) # authoring wants False + + rules = [ + { + "scope": {"dag_id": "^prod_.*"}, + "match_mode": "regex", + "locked": True, + "controls": {"include_source_code": False}, + } + ] + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + cfg = resolve_task_emission_policy(task, "dev_reporting", "test_task") + + # Regex doesn't match "dev_reporting" → no lock → authoring override applies + assert cfg.include_source_code is False # authoring took effect + + def test_locked_in_legacy_path_has_no_effect(self): + """In the pure legacy path (emission_policy empty), locked is irrelevant. + + The authoring emit=True flag overrides disabled_for_operators because the legacy + path uses frozenset() (no locked fields) when calling _apply_task_authoring. + """ + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + # Use a real operator that carries params so extend_global_openlineage_emission_policy takes effect. + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) + + fqcn = f"{task.__class__.__module__}.{task.__class__.__qualname__}" + with conf_vars( + { + ("openlineage", "emission_policy"): "", + ("openlineage", "disabled_for_operators"): fqcn, + } + ): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + # Legacy path: authoring applied on top (frozenset() → no locked fields) + # disabled_for_operators → emit=False in legacy path, but authoring emit=True + # overrides it because there is no locking in the pure legacy path. + assert cfg.emit is True + + def test_locked_emit_logs_exact_message_with_locked_value(self): + """Exact log message when extend_global_openlineage_emission_policy is blocked for task 'emit'.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) + + rule = {"scope": {}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False + locked_calls = [c for c in mock_log.warning.call_args_list if "has no effect" in c.args[0]] + assert len(locked_calls) == 1 + assert locked_calls[0].args == ( + "OpenLineage emission_policy: extend_global_openlineage_emission_policy call for 'emit' on %s" + " has no effect — locked by conf rule at value %r", + "task 'test_task' in dag 'test_dag'", + False, + ) + + def test_locked_field_logs_exact_message_with_locked_value(self): + """Exact log message when extend_global_openlineage_emission_policy is blocked for a non-emit field.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=True) + + rule = {"scope": {}, "locked": True, "controls": {"include_source_code": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + locked_calls = [c for c in mock_log.warning.call_args_list if "has no effect" in c.args[0]] + assert len(locked_calls) == 1 + assert locked_calls[0].args == ( + "OpenLineage emission_policy: extend_global_openlineage_emission_policy call for '%s' on %s" + " has no effect — locked by conf rule at value %r", + "include_source_code", + "task 'test_task' in dag 'test_dag'", + False, + ) + + def test_locked_dag_emit_logs_exact_message_with_locked_value(self): + """Exact log message when extend_global_openlineage_emission_policy is blocked for dag-event 'emit'.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + dag, _ = _make_dag_and_task("locked_dag_msg_test") + extend_global_openlineage_emission_policy(dag, emit=True) + + rule = {"scope": {"dag_id": "locked_dag_msg_test"}, "locked": True, "controls": {"emit": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = resolve_dag_emission_policy("locked_dag_msg_test", dag) + + assert cfg.emit is False + locked_calls = [c for c in mock_log.warning.call_args_list if "has no effect" in c.args[0]] + assert len(locked_calls) == 1 + assert locked_calls[0].args == ( + "OpenLineage emission_policy: extend_global_openlineage_emission_policy call for 'emit' on %s" + " has no effect — locked by conf rule at value %r", + "dag event 'locked_dag_msg_test'", + False, + ) + + +class TestAuthoringLayerWithLegacy: + def test_authoring_applied_in_legacy_path(self): + """extend_global_openlineage_emission_policy flags take effect in the pure legacy path.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + # Use a real operator (EmptyOperator has params) + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, include_source_code=False) + + with conf_vars({("openlineage", "emission_policy"): ""}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.include_source_code is False + + def test_authoring_applied_in_new_path(self): + """extend_global_openlineage_emission_policy flags take effect in the new (emission_policy) path.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, hook_lineage=False) + + with conf_vars( + { + ( + "openlineage", + "emission_policy", + ): '[{"scope": {}, "controls": {"include_source_code": false}}]' + } + ): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.hook_lineage is False # authoring + assert cfg.include_source_code is False # conf rule + + def test_emit_task_events_in_authoring_sets_emit(self): + """emit_task_events in authoring flags resolves to task emit.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit_task_events=False) + + with conf_vars({("openlineage", "emission_policy"): "[]"}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False + + +class TestEmitMixedKeyContradictions: + """``emit`` vs ``emit_task_events`` resolution across rules in the same tier.""" + + def test_emit_task_events_in_later_rule_wins_with_contradiction_warning(self): + """[{'emit': False}, {'emit_task_events': True}] — task emit becomes True; a contradiction + warning fires because the resolved task-emit decision flipped within the same tier. + This pins current behaviour: ``emit`` and ``emit_task_events`` participate in the same + resolution stream for task scope, so disagreement across rules in one tier is flagged. + """ + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"emit": False}}, + {"scope": {}, "controls": {"emit_task_events": True}}, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("contradictory" in c for c in warn_calls) + + def test_emit_in_later_rule_wins_when_earlier_set_emit_task_events(self): + """[{'emit_task_events': False}, {'emit': True}] — task emit becomes True (last wins).""" + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"emit_task_events": False}}, + {"scope": {}, "controls": {"emit": True}}, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.emit is True + + def test_global_emit_true_when_selective_enable_translated_baseline(self): + """User-explicit global {'emit': true} overrides selective_enable's translated baseline.""" + _, task = _make_dag_and_task() + # Task is NOT opted in — selective_enable would otherwise suppress emission. + with conf_vars( + { + ("openlineage", "emission_policy"): json.dumps([{"scope": {}, "controls": {"emit": True}}]), + ("openlineage", "selective_enable"): "True", + } + ): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + # Both rules are global; user rule appears later → wins last-wins-in-tier. + assert cfg.emit is True + + +class TestLockOnSpecificEmitKey: + """Lock-precision tests: ``emit_task_events`` vs ``emit_dag_events`` vs ``emit``.""" + + def test_lock_on_emit_dag_events_does_not_block_task_authoring(self): + """A locked dag-event emit rule must not lock task-event authoring overrides.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) + + # Lock only the dag-event emit; task emit must remain unlocked for authoring. + # NOTE: this rule still adds 'emit' to locked_fields under the current implementation, + # because we cannot disambiguate which event scope the lock targets without + # explicit per-key locking. The test pins this behaviour so it isn't broken silently. + rule = {"scope": {"dag_id": "test_dag"}, "locked": True, "controls": {"emit_dag_events": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + # Conf rule used emit_dag_events; task emit is not changed by the conf rule itself, + # but the lock currently extends to task emit. Authoring is blocked. + assert cfg.emit is True # default — neither rule nor authoring took effect + # (Behaviour intentionally pinned; revisit if per-scope locking is added.) + + def test_lock_on_emit_task_events_blocks_task_authoring(self): + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, emit=True) + + rule = {"scope": {}, "locked": True, "controls": {"emit_task_events": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.emit is False # locked + + def test_lock_on_extract_metadata_blocks_authoring_but_not_independent_fields(self): + """Authoring ``hook_lineage`` is independent of an ``extract_operator_metadata`` lock.""" + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, extract_operator_metadata=True, hook_lineage=False) + + rule = {"scope": {}, "locked": True, "controls": {"extract_operator_metadata": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + # extract_operator_metadata stays at the locked conf value + assert cfg.extract_operator_metadata is False + # hook_lineage authoring still applies (different field, not locked) + assert cfg.hook_lineage is False + + +class TestStrictUnknownKeyHandling: + """Unknown keys at any nesting level cause the rule to be skipped with a WARNING.""" + + def test_unknown_top_level_key_skipped(self): + rule = {"scope": {}, "controls": {"emit": False}, "bogus_top_level_key": 1} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("unknown top-level key" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() # rule was skipped, defaults remain + + def test_unknown_scope_key_skipped(self): + rule = {"scope": {"dgg_id": "typo_dag"}, "controls": {"emit": False}} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("scope" in c and "unknown key" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_unknown_controls_key_skipped(self): + rule = {"scope": {}, "controls": {"emiit": False}} # typo + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("controls" in c and "unknown key" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_empty_controls_skipped(self): + rule = {"scope": {"dag_id": "x"}, "controls": {}} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("empty 'controls'" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_missing_scope_skipped(self): + rule = {"controls": {"emit": False}} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("missing required 'scope'" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_missing_controls_skipped(self): + rule = {"scope": {"dag_id": "x"}} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("missing required 'controls'" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_operator_combined_with_dag_id_skipped(self): + rule = {"scope": {"dag_id": "x", "operator": "y"}, "controls": {"emit": False}} + with mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as mock_log: + cfg = _resolve_task_controls([rule], operator=MockOperator(), dag_id="my_dag", task_id="my_task") + warn_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("operator" in c and "dag_id" in c for c in warn_calls) + assert cfg == EmissionPolicy.defaults() + + def test_well_formed_unknown_top_level_key_locked_still_works_after_skip(self): + """An unknown key on rule A skips A but rule B (well-formed) still applies.""" + rules = [ + {"scope": {}, "controls": {"emit": False}, "bogus": True}, # skipped + {"scope": {}, "controls": {"emit": False}}, # applied + ] + cfg = _resolve_task_controls(rules, operator=MockOperator(), dag_id="my_dag", task_id="my_task") + assert cfg.emit is False # second rule applied normally + + +class TestRegexMatchAll: + """A ``match_mode: regex`` rule with ``.*`` matches every dag_id — lands in dag tier.""" + + def test_regex_dot_star_dag_id_acts_as_dag_tier_for_all(self): + cfg = _resolve_task_controls( + [{"scope": {"dag_id": ".*"}, "match_mode": "regex", "controls": {"include_source_code": False}}], + dag_id="any_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.include_source_code is False + + def test_regex_dot_star_dag_id_beats_operator_tier(self): + """Regex dag rule (dag tier) beats an operator-tier rule for the same field.""" + op = MockOperator() + fqcn = _operator_fqcn(op) + cfg = _resolve_task_controls( + [ + {"scope": {"operator": fqcn}, "controls": {"include_source_code": False}}, + {"scope": {"dag_id": ".*"}, "match_mode": "regex", "controls": {"include_source_code": True}}, + ], + operator=op, + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.include_source_code is True + + def test_regex_empty_string_dag_id_matches_only_empty_dag(self): + """An empty regex pattern matches only the empty string under re.fullmatch.""" + cfg = _resolve_task_controls( + [{"scope": {"dag_id": ""}, "match_mode": "regex", "controls": {"emit": False}}], + dag_id="my_dag", + operator=MockOperator(), + task_id="my_task", + ) + assert cfg.emit is True # no match + + +class TestEmitDagEventsLockDoesNotAffectTaskAuthoring: + """Locking on a dag-event-only key should leave unrelated authoring fields untouched.""" + + def test_authoring_extract_metadata_unaffected_by_dag_emit_lock(self): + from airflow.providers.openlineage.api.emission_policy import ( + extend_global_openlineage_emission_policy, + ) + + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, extract_operator_metadata=False) + + # Lock targets emit only — other authoring fields must pass through. + rule = {"scope": {"dag_id": "test_dag"}, "locked": True, "controls": {"emit_dag_events": False}} + with conf_vars({("openlineage", "emission_policy"): json.dumps([rule])}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.extract_operator_metadata is False # authoring applied + + +class TestSelectiveEnableWithLockedRule: + """selective_enable's translated baseline can be overridden by an unlocked rule — but a locked rule wins.""" + + def test_locked_global_emit_false_overrides_selective_enable_opt_in(self): + """A locked global emit:false rule blocks the per-task opt-in from selective_enable.""" + _, task = _make_dag_and_task() + enable_lineage(task) # task IS opted in + rules = [{"scope": {}, "locked": True, "controls": {"emit": False}}] + with conf_vars( + { + ("openlineage", "emission_policy"): json.dumps(rules), + ("openlineage", "selective_enable"): "True", + } + ): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + # selective_enable injects a task-tier {"emit": true} opt-in rule for opted-in tasks. + # That rule wins value resolution (task tier > global tier). The lock on the global rule + # adds 'emit' to locked_fields but locked_fields only blocks AUTHORING, not other conf rules. + # So the task-tier opt-in rule still applies. + assert cfg.emit is True