This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 34c72c5e53a Allow excluding datasets from OpenLineage events by
pattern and hook (#71312)
34c72c5e53a is described below
commit 34c72c5e53a4a38bc57be3b702337d08dd80dc11
Author: Ramit Kataria <[email protected]>
AuthorDate: Fri Oct 23 11:58:44 2026 -0700
Allow excluding datasets from OpenLineage events by pattern and hook
(#71312)
A task that writes many objects through a hook reports one dataset per
object,
and the only existing control, hook_lineage, is all-or-nothing: silencing
the
noisy datasets also loses the ones worth tracking from the same hook.
Dataset patterns let a Dag keep the lineage that matters while dropping the
noise, and a hook scope lets that be narrowed to the hook producing it
without
touching each provider's hooks. Applying the patterns to every lineage
source
keeps one pattern behaving the same regardless of where a dataset came from.
closes: #63371
---
.../unit/google/cloud/openlineage/test_mixins.py | 3 +-
providers/openlineage/docs/emission_policy.rst | 147 +++++++++-
.../airflow/providers/openlineage/api/datasets.py | 11 +
.../providers/openlineage/api/emission_policy.py | 26 +-
.../src/airflow/providers/openlineage/api/sql.py | 4 +-
.../providers/openlineage/extractors/manager.py | 69 +++--
.../providers/openlineage/utils/emission_policy.py | 312 +++++++++++++++++++--
.../openlineage/utils/sql_hook_lineage.py | 16 +-
.../tests/unit/openlineage/api/test_datasets.py | 28 ++
.../unit/openlineage/api/test_emission_policy.py | 29 ++
.../tests/unit/openlineage/api/test_sql.py | 25 +-
.../unit/openlineage/extractors/test_manager.py | 167 ++++++++++-
.../unit/openlineage/utils/test_emission_policy.py | 253 +++++++++++++++++
.../openlineage/utils/test_sql_hook_lineage.py | 28 ++
14 files changed, 1044 insertions(+), 74 deletions(-)
diff --git
a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
index 2a44c648ef9..3cdc868edfc 100644
--- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
+++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
@@ -43,6 +43,7 @@ from airflow.providers.common.compat.openlineage.facet import
(
from airflow.providers.google.cloud.openlineage.facets import
BigQueryJobRunFacet
from airflow.providers.google.cloud.openlineage.mixins import
_BigQueryInsertJobOperatorOpenLineageMixin
from airflow.providers.openlineage.sqlparser import SQLParser
+from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy
QUERY_JOB_PROPERTIES = {
"configuration": {
@@ -645,7 +646,7 @@ class TestBigQueryOpenLineageMixin:
return [input_table2], [output_table2]
mock_get_inputs_and_outputs.side_effect = get_inputs_and_outputs
- mock_resolve_task_emission_policy.return_value = MagicMock(emit=True)
+ mock_resolve_task_emission_policy.return_value =
EmissionPolicy.defaults()
lineage =
self.operator.get_openlineage_facets_on_complete(make_task_instance())
diff --git a/providers/openlineage/docs/emission_policy.rst
b/providers/openlineage/docs/emission_policy.rst
index 5fb49148601..06488dfcafd 100644
--- a/providers/openlineage/docs/emission_policy.rst
+++ b/providers/openlineage/docs/emission_policy.rst
@@ -100,15 +100,19 @@ and the rule is skipped with a WARNING. The same applies
to unknown keys inside
- ``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.
+- ``hook``: fully-qualified class name (``module.ClassName``) of the object
that reported a
+ hook-collected dataset, usually a hook such as
+ ``airflow.providers.amazon.aws.hooks.s3.S3Hook``. It can be added to any of
the other scopes
+ and narrows the rule to the datasets that hook reports. See
:ref:`emission_policy_datasets:openlineage`.
- *(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.
+``scope`` must be one of: global (empty), operator-only, dag-only, or
dag+task, each optionally
+with ``hook``. ``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.
+value inside ``scope`` (``dag_id``, ``task_id``, ``operator``, ``hook``) is
treated as a
+``re.fullmatch`` pattern.
**Control flag keys** (all optional inside ``controls``; the dict must be
non-empty):
@@ -150,6 +154,54 @@ pattern.
``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.
+ * - ``exclude_datasets``
+ - ``[]``
+ - List of regex patterns; datasets whose ``<namespace>/<name>`` matches
any of them are
+ dropped. Task events only. See
:ref:`emission_policy_datasets:openlineage`.
+
+**Valid scope and control combinations.** A rule that uses a control its scope
does not accept
+is skipped with a WARNING that names the offending keys.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 10 14 10 18 18
+
+ * - Control
+ - Global (``{}``)
+ - ``operator``
+ - ``dag_id``
+ - ``dag_id`` + ``task_id``
+ - Any of these + ``hook``
+ * - ``emit``, ``emit_task_events``
+ - yes
+ - yes
+ - yes
+ - yes
+ - no
+ * - ``emit_dag_events``
+ - yes
+ - no
+ - yes
+ - no
+ - no
+ * - ``extract_operator_metadata``, ``include_source_code``,
``include_full_task_info``
+ - yes
+ - yes
+ - yes
+ - yes
+ - no
+ * - ``hook_lineage``
+ - yes
+ - yes
+ - yes
+ - yes
+ - yes, keeps or drops the hook's datasets; not with ``exclude_datasets``
+ * - ``exclude_datasets``
+ - yes
+ - yes
+ - yes
+ - yes
+ - yes; not with ``hook_lineage``
``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 /
@@ -220,6 +272,85 @@ In the example above:
- 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).
+.. _emission_policy_datasets:openlineage:
+
+Excluding datasets (``exclude_datasets``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``exclude_datasets`` drops individual datasets while keeping the rest of a
task's lineage, for
+example the per-chunk objects of a multipart upload that hook lineage reports
one by one.
+
+- Patterns are always matched with ``re.fullmatch``, independent of
``match_mode``, against
+ ``<namespace>/<name>`` of the OpenLineage dataset:
``s3://my-bucket/staging/part_1`` for a
+ dataset with namespace ``s3://my-bucket`` and name ``staging/part_1``.
Airflow assets are
+ matched after translation to OpenLineage, so assets that cannot be
translated are never
+ reported, whatever the patterns.
+- They apply to every source of task lineage: extractor results,
hook-collected assets,
+ manually declared ``inlets`` / ``outlets``, runtime outlet events, the
per-query events emitted
+ from SQL hook lineage, and the ``emit_dataset_lineage`` /
``emit_query_lineage`` APIs.
+- Only datasets are dropped, never events. A SQL hook lineage query event, or
an
+ ``emit_query_lineage`` event, is still emitted when every dataset of the
query is excluded,
+ because it also carries the query text and id. ``emit_dataset_lineage``
skips its event in
+ that case, since the event exists only to carry the datasets.
+- Excluding datasets never changes which source is used. If an extractor
returns datasets that
+ are all excluded, hook lineage is not consulted as a fallback; if hook
lineage collected assets
+ that are all excluded, manually declared ``inlets`` / ``outlets`` are not
used either.
+- A rule whose ``scope`` has a ``hook`` key only applies to datasets reported
by that hook, and
+ only to hook lineage and SQL hook lineage (the other sources have no
reporting hook). Such a
+ rule may only carry ``exclude_datasets`` or ``hook_lineage``; with a
``hook`` scope,
+ ``hook_lineage: false`` drops every dataset the hook reports and
``hook_lineage: true`` keeps
+ them all. A hook-scoped ``hook_lineage`` does not affect whether hook
lineage is collected.
+
+The list is resolved per dataset, and the most specific matching rule that
sets it **replaces**
+broader ones instead of merging with them, so a narrower rule can shorten a
broader exclusion or
+clear it with ``[]``. For a dataset reported by a hook, each tier is preceded
by its
+``hook`` variant:
+
+1. ``dag_id`` + ``task_id`` + ``hook``
+2. ``dag_id`` + ``task_id``
+3. ``dag_id`` + ``hook``
+4. ``dag_id``
+5. ``operator`` + ``hook``
+6. ``operator``
+7. ``hook``
+8. Global (empty ``scope``)
+
+Every replacement is logged at INFO level, naming the winning rule and the
broader rules it
+replaced. ``locked`` keeps its usual meaning: a locked rule blocks authoring
overrides of
+``exclude_datasets`` for the datasets it applies to, but a more specific conf
rule still replaces
+its list.
+
+.. code-block:: ini
+
+ [openlineage]
+ emission_policy = [
+ {"scope": {"hook": ".*\\.S3Hook"}, "match_mode": "regex", "controls":
{"exclude_datasets": ["s3://.*/_temporary/.*"]}},
+ {"scope": {"dag_id": "chunked_dag", "task_id": "upload", "hook":
"airflow.providers.amazon.aws.hooks.s3.S3Hook"}, "controls":
{"exclude_datasets": ["s3://my-bucket/staging/part_.*"]}},
+ {"scope": {"dag_id": "noisy_dag", "hook":
"airflow.providers.google.cloud.hooks.gcs.GCSHook"}, "controls":
{"hook_lineage": false}}
+ ]
+
+In the example above:
+
+- ``_temporary`` objects reported by ``S3Hook`` are dropped in every Dag.
+- ``upload`` in ``chunked_dag`` drops the staging parts its ``S3Hook`` writes.
Its list replaces
+ the first rule's, so ``_temporary`` objects are kept for that task.
+- Every dataset ``GCSHook`` reports in ``noisy_dag`` is dropped; other hooks
in that Dag are
+ unaffected.
+
+.. note::
+
+ **Exclusion happens after the collector cap.** Hook-collected assets are
filtered when
+ OpenLineage reads the hook lineage collector, after
+ :ref:`[lineage]
max_assets_per_collector<config:lineage__max_assets_per_collector>` has
+ limited what the collector kept. Excluding assets reduces what gets
reported, but it does not
+ make room for assets the cap already discarded.
+
+.. note::
+
+ Only the event ``inputs`` / ``outputs`` lists are filtered. Facets can
still refer to an
+ excluded dataset: ``columnLineage`` on the remaining datasets, and the
task's ``inlets``,
+ ``outlets`` and runtime outlet assets in the ``airflow`` run facet.
+
**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.
@@ -253,6 +384,9 @@ wins **unless** the matching conf rule is marked with
``"locked": true``.
# Suppress DAG-run events but keep task events
extend_global_openlineage_emission_policy(dag, emit_dag_events=False)
+ # Drop the staging objects a task writes, keep the rest of its lineage
+ extend_global_openlineage_emission_policy(extract,
exclude_datasets=["s3://my-bucket/staging/.*"])
+
Key semantics:
- **Resolution priority**: authoring flags > unlocked conf rules > built-in
defaults.
@@ -269,6 +403,9 @@ Key semantics:
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.
+- **``exclude_datasets`` has no hook dimension**: an authored list applies to
every dataset of
+ the task and replaces the conf lists, including hook-scoped ones, except
where a locked rule
+ applies.
- **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.
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py
b/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py
index d2c65043b89..c3e548e93ce 100644
--- a/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py
+++ b/providers/openlineage/src/airflow/providers/openlineage/api/datasets.py
@@ -146,6 +146,17 @@ def emit_dataset_lineage(
)
return
+ inputs = controls.dataset_filter.exclude(inputs)
+ outputs = controls.dataset_filter.exclude(outputs)
+ if not inputs and not outputs:
+ log.info(
+ "Skipping OpenLineage RUNNING event emission for task `%s` in
dag `%s`: every dataset "
+ "passed to emit_dataset_lineage is excluded by the emission
policy.",
+ task_instance.task_id,
+ task_instance.dag_id,
+ )
+ return
+
run_facets = build_task_event_run_facets(
task_instance=task_instance,
dag_run=dag_run,
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py
b/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py
index 898aae41a65..652b89d4e3f 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py
@@ -55,6 +55,10 @@ Quick-start examples
extend_global_openlineage_emission_policy(dag, emit_dag_events=False)
+**Drop the per-chunk objects a task writes, keeping the rest of its lineage**::
+
+ extend_global_openlineage_emission_policy(upload,
exclude_datasets=["s3://my-bucket/staging/part_.*"])
+
**Works with XComArg** — flags are applied to the underlying operator::
result = extend_global_openlineage_emission_policy(
@@ -91,12 +95,15 @@ from airflow.providers.openlineage.utils.emission_policy
import (
EMIT,
EMIT_DAG_EVENTS,
EMIT_TASK_EVENTS,
+ EXCLUDE_DATASETS,
EXTRACT_OPERATOR_METADATA,
HOOK_LINEAGE,
INCLUDE_FULL_TASK_INFO,
INCLUDE_SOURCE_CODE,
OL_EMISSION_POLICY_PARAM,
+ ControlValue,
_merge_param,
+ find_invalid_pattern,
)
if TYPE_CHECKING:
@@ -121,6 +128,7 @@ def extend_global_openlineage_emission_policy(
include_source_code: bool | None = None,
hook_lineage: bool | None = None,
include_full_task_info: bool | None = None,
+ exclude_datasets: list[str] | None = None,
) -> T:
"""
Extend the global OpenLineage emission policy with per-task / per-DAG
overrides.
@@ -132,8 +140,8 @@ def extend_global_openlineage_emission_policy(
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.
+ ``include_source_code``, ``hook_lineage``, ``include_full_task_info``,
+ ``exclude_datasets``) 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.
@@ -180,6 +188,11 @@ def extend_global_openlineage_emission_policy(
: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.
+ :param exclude_datasets: Regex patterns matched with ``re.fullmatch``
against
+ ``"<namespace>/<name>"`` of each OpenLineage dataset of the task;
matching datasets are
+ dropped. Replaces, instead of extending, any list from the
``emission_policy`` config,
+ including lists set for specific hooks.
+ :raises ValueError: If ``exclude_datasets`` is not a list of valid regex
patterns.
:return: The same *obj* — allows use as a decorator or in chained calls.
"""
if isinstance(obj, XComArg):
@@ -192,6 +205,7 @@ def extend_global_openlineage_emission_policy(
include_source_code=include_source_code,
hook_lineage=hook_lineage,
include_full_task_info=include_full_task_info,
+ exclude_datasets=exclude_datasets,
)
return obj
@@ -206,7 +220,12 @@ def extend_global_openlineage_emission_policy(
"configuration with a global rule ('scope': {}) instead."
)
- provided: dict[str, bool] = {
+ # A bad pattern stored here would only fail during lineage extraction,
where the error is
+ # swallowed and the task silently loses its lineage.
+ if exclude_datasets is not None and (invalid :=
find_invalid_pattern(EXCLUDE_DATASETS, exclude_datasets)):
+ raise ValueError(f"extend_global_openlineage_emission_policy():
{invalid}")
+
+ provided: dict[str, ControlValue] = {
k: v
for k, v in {
EMIT: emit,
@@ -216,6 +235,7 @@ def extend_global_openlineage_emission_policy(
INCLUDE_SOURCE_CODE: include_source_code,
HOOK_LINEAGE: hook_lineage,
INCLUDE_FULL_TASK_INFO: include_full_task_info,
+ EXCLUDE_DATASETS: exclude_datasets,
}.items()
if v is not None
}
diff --git a/providers/openlineage/src/airflow/providers/openlineage/api/sql.py
b/providers/openlineage/src/airflow/providers/openlineage/api/sql.py
index bf8e17a9e7b..14f0c9f421b 100644
--- a/providers/openlineage/src/airflow/providers/openlineage/api/sql.py
+++ b/providers/openlineage/src/airflow/providers/openlineage/api/sql.py
@@ -183,8 +183,8 @@ def emit_query_lineage(
task_instance=task_instance,
job_name=job_name,
is_successful=is_successful,
- inputs=all_inputs,
- outputs=all_outputs,
+ inputs=controls.dataset_filter.exclude(all_inputs),
+ outputs=controls.dataset_filter.exclude(all_outputs),
run_facets=run_facets,
job_facets=job_facets,
start_event_time=start_time,
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py
b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py
index c6b936970a1..b3af8de1c05 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py
@@ -31,7 +31,7 @@ from airflow.providers.openlineage.extractors.base import (
)
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.emission_policy import
EmissionPolicy, get_hook_class_name
from airflow.providers.openlineage.utils.utils import (
get_runtime_outlet_assets,
get_unknown_source_attribute_run_facet,
@@ -166,17 +166,20 @@ class ExtractorManager(LoggingMixin):
# If no inputs and outputs are present - check Hook Lineage if
enabled
if (not task_metadata.inputs) and (not task_metadata.outputs):
if controls.hook_lineage:
- hook_lineage = self.get_hook_lineage(task_instance,
task_instance_state)
+ hook_lineage = self.get_hook_lineage(task_instance,
task_instance_state, controls)
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)
+ self.extract_inlets_and_outlets(task_metadata,
task, task_instance, controls)
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)
+ self.extract_inlets_and_outlets(task_metadata, task,
task_instance, controls)
+ else:
+ task_metadata.inputs =
controls.dataset_filter.exclude(task_metadata.inputs)
+ task_metadata.outputs =
controls.dataset_filter.exclude(task_metadata.outputs)
return task_metadata
except Exception as e:
@@ -195,7 +198,7 @@ class ExtractorManager(LoggingMixin):
# 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)
+ hook_lineage = self.get_hook_lineage(task_instance,
task_instance_state, controls)
except Exception as e:
self.log.warning(
"Failed to extract OpenLineage hook lineage %s: %s.
Task event will be emitted without lineage.",
@@ -219,7 +222,7 @@ class ExtractorManager(LoggingMixin):
task_metadata = OperatorLineage(
run_facets=get_unknown_source_attribute_run_facet(task=task),
)
- self.extract_inlets_and_outlets(task_metadata, task, task_instance)
+ self.extract_inlets_and_outlets(task_metadata, task,
task_instance, controls)
return task_metadata
return OperatorLineage()
@@ -253,29 +256,32 @@ class ExtractorManager(LoggingMixin):
task_metadata: OperatorLineage,
task,
task_instance=None,
+ controls: EmissionPolicy | None = None,
) -> None:
+ if controls is None:
+ controls = EmissionPolicy.defaults()
if task.inlets or task.outlets:
self.log.debug("Manually extracting lineage metadata from inlets
and outlets")
- for i in task.inlets:
- if d := self.convert_to_ol_dataset(i):
- task_metadata.inputs.append(d)
- for o in task.outlets:
- if d := self.convert_to_ol_dataset(o):
- task_metadata.outputs.append(d)
+ inputs = [d for i in task.inlets if (d :=
self.convert_to_ol_dataset(i))]
+ outputs = [d for o in task.outlets if (d :=
self.convert_to_ol_dataset(o))]
# Add runtime-emitted outlets (alias resolutions + dynamic asset
events), deduped
# by (namespace, name) against both static outputs and each other.
- seen = {(d.namespace, d.name) for d in task_metadata.outputs}
+ seen = {(d.namespace, d.name) for d in task_metadata.outputs + outputs}
for asset, _alias in get_runtime_outlet_assets(task_instance):
ol = translate_airflow_asset(asset, None)
if ol is not None and (ol.namespace, ol.name) not in seen:
- task_metadata.outputs.append(ol)
+ outputs.append(ol)
seen.add((ol.namespace, ol.name))
+ task_metadata.inputs.extend(controls.dataset_filter.exclude(inputs))
+ task_metadata.outputs.extend(controls.dataset_filter.exclude(outputs))
+
def get_hook_lineage(
self,
task_instance=None,
task_instance_state: TaskInstanceState | None = None,
+ controls: EmissionPolicy | None = None,
) -> OperatorLineage | None:
"""
Extract lineage from the Hook Lineage Collector.
@@ -288,8 +294,13 @@ class ExtractorManager(LoggingMixin):
When ``task_instance`` is provided, each extra is parsed and
separate per-query
OpenLineage events are emitted.
- Returns ``None`` when nothing was collected.
+ Datasets matching the ``exclude_datasets`` patterns that *controls*
resolves for the
+ reporting hook are dropped. Returns ``None`` when nothing was
collected; an empty
+ :class:`OperatorLineage` when everything collected was excluded, so
exclusion never
+ makes the caller fall back to another lineage source.
"""
+ if controls is None:
+ controls = EmissionPolicy.defaults()
try:
from airflow.providers.common.compat.lineage.hook import
get_hook_lineage_collector
from airflow.providers.common.sql.hooks.lineage import
SqlJobHookLineageExtra
@@ -308,16 +319,8 @@ class ExtractorManager(LoggingMixin):
collected = collector.collected_assets
# Asset-based inputs/outputs - keep only assets that can be translated
to OL datasets
- inputs = [
- asset
- for asset_info in collected.inputs
- if (asset := translate_airflow_asset(asset_info.asset,
asset_info.context)) is not None
- ]
- outputs = [
- asset
- for asset_info in collected.outputs
- if (asset := translate_airflow_asset(asset_info.asset,
asset_info.context)) is not None
- ]
+ inputs = self._translate_hook_assets(collected.inputs)
+ outputs = self._translate_hook_assets(collected.outputs)
# SQL-based lineage - keep only SQL extra with query_text or job_id.
sql_extras = [
@@ -338,12 +341,26 @@ class ExtractorManager(LoggingMixin):
task_instance=task_instance,
sql_extras=sql_extras,
is_successful=task_instance_state != TaskInstanceState.FAILED,
+ dataset_filter=controls.dataset_filter,
)
if not inputs and not outputs:
return None
- return OperatorLineage(inputs=inputs, outputs=outputs)
+ dataset_filter = controls.dataset_filter
+ return OperatorLineage(
+ inputs=[d for hook, d in inputs if not
dataset_filter.is_excluded(d, hook)],
+ outputs=[d for hook, d in outputs if not
dataset_filter.is_excluded(d, hook)],
+ )
+
+ @staticmethod
+ def _translate_hook_assets(asset_infos) -> list[tuple[str | None,
Dataset]]:
+ """Translate collected assets to OL datasets, paired with the class
name of their reporter."""
+ return [
+ (get_hook_class_name(asset_info.context), dataset)
+ for asset_info in asset_infos
+ if (dataset := translate_airflow_asset(asset_info.asset,
asset_info.context)) is not None
+ ]
@staticmethod
def convert_to_ol_dataset_from_object_storage_uri(uri: str) -> Dataset |
None:
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py
index de386fce438..89a93b94035 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py
@@ -39,16 +39,20 @@ Scope keys (all optional inside ``scope``):
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.
+- ``hook``: fully-qualified class name of the object that reported a
hook-collected
+ dataset (usually a hook). Can be added to any of the scopes above, and
narrows the
+ rule to datasets reported by that hook. Only ``exclude_datasets`` and
``hook_lineage``
+ are valid controls for a hook-scoped rule.
- *(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.
+global, operator-only, dag-only, dag+task, each optionally with ``hook``).
``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.
+``operator``, ``hook``) is treated as a ``re.fullmatch`` pattern.
Control flag keys (all optional inside ``controls``; the dict must be
non-empty):
@@ -69,10 +73,20 @@ Control flag keys (all optional inside ``controls``; the
dict must be non-empty)
- ``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``.
+ Only meaningful for task events. Default: ``true``. In a hook-scoped rule
it instead
+ keeps (``true``) or drops (``false``) every dataset reported by the matching
hook, and it
+ cannot be combined with ``exclude_datasets`` in the same rule.
- ``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``.
+- ``exclude_datasets``: list of ``re.fullmatch`` patterns (independent of
``match_mode``)
+ matched against ``"<namespace>/<name>"`` of each OpenLineage dataset;
matching datasets
+ are dropped from task events and from the SQL hook lineage events of the
task. Applies
+ to extractor, hook-collected and manually declared lineage. Default: ``[]``.
+
+Unlike the boolean controls, ``exclude_datasets`` is resolved per dataset, and
the most
+specific matching rule's list **replaces** broader lists instead of merging
with them, so a
+narrower rule can shorten or clear (``[]``) a broader exclusion.
``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.
@@ -102,6 +116,11 @@ rule wins):
For ``emit`` at the task level: ``emit_task_events`` beats ``emit`` within the
same
rule.
+For ``exclude_datasets`` on a dataset reported by a hook, each tier above is
preceded by
+its hook-scoped variant: ``dag_id`` + ``task_id`` + ``hook``, ``dag_id`` +
``task_id``,
+``dag_id`` + ``hook``, ``dag_id``, ``operator`` + ``hook``, ``operator``,
``hook``,
+global. Datasets from other sources only see the tiers without ``hook``.
+
Priority for **DAG run events**:
1. ``dag_id`` rule (using ``emit_dag_events`` or ``emit``)
@@ -138,18 +157,27 @@ import logging
import re
import warnings
from collections.abc import Callable, Sequence
-from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING
+from dataclasses import dataclass, field as dataclass_field, replace
+from typing import TYPE_CHECKING, TypeVar
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 openlineage.client.event_v2 import Dataset
+
from airflow.providers.openlineage.utils.utils import AnyOperator
+ DatasetT = TypeVar("DatasetT", bound=Dataset)
+
log = logging.getLogger(__name__)
+V = TypeVar("V")
+
+# A validated ``controls`` value: bool for the flags, list of patterns for
``exclude_datasets``.
+ControlValue = bool | list[str]
+
# Control flag names — the keys that may appear inside a rule's ``controls``
dict.
EMIT = "emit"
EMIT_TASK_EVENTS = "emit_task_events"
@@ -158,11 +186,13 @@ EXTRACT_OPERATOR_METADATA = "extract_operator_metadata"
INCLUDE_SOURCE_CODE = "include_source_code"
HOOK_LINEAGE = "hook_lineage"
INCLUDE_FULL_TASK_INFO = "include_full_task_info"
+EXCLUDE_DATASETS = "exclude_datasets"
# 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"
+SCOPE_HOOK = "hook"
# Top-level rule keys.
RULE_SCOPE = "scope"
@@ -178,7 +208,7 @@ 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_SCOPE_KEYS: frozenset[str] = frozenset({SCOPE_DAG_ID, SCOPE_TASK_ID,
SCOPE_OPERATOR, SCOPE_HOOK})
_ALLOWED_CONTROL_KEYS: frozenset[str] = frozenset(
{
EMIT,
@@ -188,8 +218,12 @@ _ALLOWED_CONTROL_KEYS: frozenset[str] = frozenset(
INCLUDE_SOURCE_CODE,
HOOK_LINEAGE,
INCLUDE_FULL_TASK_INFO,
+ EXCLUDE_DATASETS,
}
)
+# The hook is only known per dataset, so a hook-scoped rule may only carry
controls that
+# act on individual datasets.
+_HOOK_SCOPE_CONTROL_KEYS: frozenset[str] = frozenset({EXCLUDE_DATASETS,
HOOK_LINEAGE})
# Fields scanned for ``locked: true`` on task-event resolution.
# ``emit`` is intentionally absent: emit-locking is driven by either ``emit``
or
@@ -213,6 +247,7 @@ _TASK_FLAG_KEYS: frozenset[str] = frozenset(
INCLUDE_SOURCE_CODE,
HOOK_LINEAGE,
INCLUDE_FULL_TASK_INFO,
+ EXCLUDE_DATASETS,
}
)
_DAG_FLAG_KEYS: frozenset[str] = frozenset({EMIT, EMIT_DAG_EVENTS})
@@ -227,6 +262,8 @@ class EmissionPolicy:
include_source_code: bool
hook_lineage: bool
include_full_task_info: bool
+ # Rules hold dicts, so hashing the filter would fail; equal policies still
hash equally without it.
+ dataset_filter: DatasetFilter = dataclass_field(default_factory=lambda:
DatasetFilter(), hash=False)
@classmethod
def defaults(cls) -> EmissionPolicy:
@@ -251,7 +288,7 @@ class Rule:
"""
scope: dict[str, str]
- controls: dict[str, bool]
+ controls: dict[str, ControlValue]
match_mode: str = "exact"
locked: bool = False
@@ -267,6 +304,10 @@ class Rule:
def has_op(self) -> bool:
return SCOPE_OPERATOR in self.scope
+ @property
+ def has_hook(self) -> bool:
+ return SCOPE_HOOK in self.scope
+
def _matches(pattern: str, value: str, match_mode: str) -> bool:
"""Return ``True`` if *value* matches *pattern* according to
*match_mode*."""
@@ -275,7 +316,133 @@ def _matches(pattern: str, value: str, match_mode: str)
-> bool:
return pattern == value
-def _read_param(obj: object, param_name: str) -> dict[str, bool]:
+def get_hook_class_name(hook: object | None) -> str | None:
+ """Return the fully-qualified class name that the ``hook`` scope key
matches, if any."""
+ if hook is None:
+ return None
+ hook_type = type(hook)
+ return f"{hook_type.__module__}.{hook_type.__name__}"
+
+
+def _get_dataset_identity(dataset: Dataset) -> str:
+ """Return the ``"<namespace>/<name>"`` string that ``exclude_datasets``
patterns match."""
+ return f"{dataset.namespace}/{dataset.name}"
+
+
+def _get_rule_exclude_patterns(rule: Rule) -> tuple[str, ...] | None:
+ """Return the ``exclude_datasets`` patterns *rule* sets, or ``None`` if it
sets none."""
+ patterns = rule.controls.get(EXCLUDE_DATASETS)
+ if isinstance(patterns, list):
+ return tuple(patterns)
+ if rule.has_hook and HOOK_LINEAGE in rule.controls:
+ return () if rule.controls[HOOK_LINEAGE] else (".*",)
+ return None
+
+
+@dataclass(frozen=True)
+class DatasetFilter:
+ """
+ ``exclude_datasets`` resolution for the datasets of one task.
+
+ Hook-scoped rules can only be matched once the reporting hook is known, so
the winning
+ patterns are resolved lazily per hook (``None`` for datasets not reported
by a hook)
+ and cached. *tiers* holds only rules that set patterns, most specific
first:
+ dag+task+hook, dag+task, dag+hook, dag, operator+hook, operator, hook,
global.
+ """
+
+ tiers: tuple[tuple[Rule, ...], ...] = ()
+ authoring_patterns: tuple[str, ...] | None = None
+ context: str = dataclass_field(default="", compare=False)
+ _cache: dict[str | None, tuple[tuple[str, ...], tuple[re.Pattern[str],
...]]] = dataclass_field(
+ default_factory=dict, init=False, compare=False, repr=False
+ )
+
+ def _get_resolved(self, hook: str | None) -> tuple[tuple[str, ...],
tuple[re.Pattern[str], ...]]:
+ if hook not in self._cache:
+ patterns = self._resolve_patterns(hook)
+ self._cache[hook] = (patterns, tuple(re.compile(p) for p in
patterns))
+ return self._cache[hook]
+
+ def get_patterns(self, hook: str | None = None) -> tuple[str, ...]:
+ """Return the exclusion patterns that apply to datasets reported by
*hook*."""
+ return self._get_resolved(hook)[0]
+
+ def is_excluded(self, dataset: Dataset, hook: str | None = None) -> bool:
+ """Return ``True`` if *dataset*, reported by *hook* (if any), matches
an exclusion pattern."""
+ compiled = self._get_resolved(hook)[1]
+ if not compiled:
+ return False
+ identity = _get_dataset_identity(dataset)
+ if any(p.fullmatch(identity) for p in compiled):
+ log.debug("OpenLineage emission policy: excluding dataset %r for
%s.", identity, self.context)
+ return True
+ return False
+
+ def exclude(self, datasets: Sequence[DatasetT], hook: str | None = None)
-> list[DatasetT]:
+ """Return *datasets* without the ones excluded for *hook*."""
+ return [d for d in datasets if not self.is_excluded(d, hook)]
+
+ def _resolve_patterns(self, hook: str | None) -> tuple[str, ...]:
+ if not self.tiers and self.authoring_patterns is None:
+ return ()
+ tiers = [
+ [
+ rule
+ for rule in tier
+ if not rule.has_hook
+ or (hook is not None and _matches(rule.scope[SCOPE_HOOK],
hook, rule.match_mode))
+ ]
+ for tier in self.tiers
+ ]
+ context = self.context if hook is None else f"datasets reported by
'{hook}' in {self.context}"
+ no_patterns: tuple[str, ...] = ()
+ patterns, winning_rule = _walk_tiers(tiers,
_get_rule_exclude_patterns, EXCLUDE_DATASETS, no_patterns)
+ if winning_rule is not None:
+ winning_tier = next(i for i, tier in enumerate(tiers) if
winning_rule in tier)
+ replaced = [rule for tier in tiers[winning_tier + 1 :] for rule in
tier]
+ _audit_log_conf_patterns(patterns, context, winning_rule, replaced)
+
+ if self.authoring_patterns is None or self.authoring_patterns ==
patterns:
+ return patterns
+ if any(rule.locked for tier in tiers for rule in tier):
+ 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",
+ EXCLUDE_DATASETS,
+ context,
+ list(patterns),
+ )
+ return patterns
+ log.info(
+ "OpenLineage emission policy: '%s' set to %r for %s "
+ "by manual `extend_global_openlineage_emission_policy` call.",
+ EXCLUDE_DATASETS,
+ list(self.authoring_patterns),
+ context,
+ )
+ return self.authoring_patterns
+
+
+def find_invalid_pattern(field_name: str, value: object) -> str | None:
+ """
+ Return an error message if *value* is not a list of valid regex patterns,
else ``None``.
+
+ Shared by the conf parser (which skips the offending rule) and the
authoring API (which
+ raises), so every pattern that reaches dataset matching has compiled at
least once.
+ """
+ if not isinstance(value, list):
+ return f"'{field_name}' must be a list of strings, got {value!r}"
+ for pattern in value:
+ if not isinstance(pattern, str):
+ return f"'{field_name}' must contain only strings, got {pattern!r}"
+ try:
+ re.compile(pattern)
+ except re.error as exc:
+ return f"'{field_name}' pattern {pattern!r} is not a valid regex
({exc})"
+ return None
+
+
+def _read_param(obj: object, param_name: str) -> dict[str, ControlValue]:
"""Read the flags dict stored at *param_name* on *obj*, or ``{}`` if
absent."""
val = getattr(obj, "params", {}).get(param_name)
if val is None:
@@ -285,7 +452,7 @@ def _read_param(obj: object, param_name: str) -> dict[str,
bool]:
return val if isinstance(val, dict) else {}
-def _merge_param(obj: object, param_name: str, new_flags: dict[str, bool]) ->
None:
+def _merge_param(obj: object, param_name: str, new_flags: dict[str,
ControlValue]) -> None:
"""
Merge *new_flags* into the ``param_name`` param on *obj*, creating it if
absent.
@@ -321,6 +488,29 @@ def _audit_log_conf_field(field: str, value: bool,
context: str, source: Rule) -
)
+def _audit_log_conf_patterns(
+ patterns: tuple[str, ...], context: str, source: Rule, replaced:
Sequence[Rule]
+) -> None:
+ """Log the winning ``exclude_datasets`` list, and the broader rules it
replaced, if any."""
+ if replaced:
+ log.info(
+ "OpenLineage emission policy: '%s' set to %r for %s by %r,
replacing broader rule(s) %r",
+ EXCLUDE_DATASETS,
+ list(patterns),
+ context,
+ source,
+ list(replaced),
+ )
+ else:
+ log.info(
+ "OpenLineage emission policy: '%s' set to %r for %s by %r",
+ EXCLUDE_DATASETS,
+ list(patterns),
+ context,
+ source,
+ )
+
+
def _audit_log_authoring_updates(
field_changes: dict[str, bool],
context: str,
@@ -433,7 +623,12 @@ def _parse_rule(rule: object) -> Rule | None:
)
return None
for k, v in controls.items():
- if not isinstance(v, bool):
+ if k == EXCLUDE_DATASETS:
+ invalid = find_invalid_pattern(k, v)
+ if invalid:
+ log.warning("OpenLineage emission_policy rule controls.%s;
ignoring: %r", invalid, rule)
+ return None
+ elif not isinstance(v, bool):
log.warning(
"OpenLineage emission_policy rule 'controls.%s' must be bool,
got %r; ignoring: %r",
k,
@@ -456,6 +651,24 @@ def _parse_rule(rule: object) -> Rule | None:
rule,
)
return None
+ if SCOPE_HOOK in scope:
+ unsupported = set(controls) - _HOOK_SCOPE_CONTROL_KEYS
+ if unsupported:
+ log.warning(
+ "OpenLineage emission_policy rule has scope 'hook' with
controls %s, but hook-scoped rules "
+ "only support %s; ignoring: %r",
+ sorted(unsupported),
+ sorted(_HOOK_SCOPE_CONTROL_KEYS),
+ rule,
+ )
+ return None
+ if EXCLUDE_DATASETS in controls and HOOK_LINEAGE in controls:
+ log.warning(
+ "OpenLineage emission_policy rule has scope 'hook' with both
'exclude_datasets' and "
+ "'hook_lineage', which both set the datasets dropped for that
hook; ignoring: %r",
+ rule,
+ )
+ return None
match_mode = rule.get(RULE_MATCH_MODE, "exact")
if match_mode not in ("exact", "regex"):
@@ -466,7 +679,7 @@ def _parse_rule(rule: object) -> Rule | None:
)
return None
if match_mode == "regex":
- for key in (SCOPE_DAG_ID, SCOPE_TASK_ID, SCOPE_OPERATOR):
+ for key in (SCOPE_DAG_ID, SCOPE_TASK_ID, SCOPE_OPERATOR, SCOPE_HOOK):
if key in scope:
try:
re.compile(scope[key])
@@ -516,10 +729,16 @@ def _classify_task_rules(
fqcn: str,
dag_id: str,
task_id: str,
+ *,
+ hook_scoped: bool = False,
) -> tuple[list[Rule], list[Rule], list[Rule], list[Rule]]:
"""
Classify pre-validated rules into the four priority tiers for task
resolution.
+ Only rules whose scope has a ``hook`` key (``hook_scoped=True``) or has
none
+ (``hook_scoped=False``) are considered; the ``hook`` key itself is left for
+ :class:`DatasetFilter` to match per dataset.
+
Returns ``(task_rules, dag_rules, operator_rules, global_rules)``.
"""
task_rules: list[Rule] = []
@@ -528,6 +747,8 @@ def _classify_task_rules(
global_rules: list[Rule] = []
for rule in rules:
+ if rule.has_hook != hook_scoped:
+ continue
scope = rule.scope
match_mode = rule.match_mode
@@ -575,24 +796,24 @@ def _classify_dag_rules(
def _walk_tiers(
- tiers: list[list[Rule]],
- extract_value: Callable[[dict[str, bool]], bool | None],
+ tiers: Sequence[Sequence[Rule]],
+ extract_value: Callable[[Rule], V | None],
field_label: str,
- default: bool,
-) -> tuple[bool, Rule | None]:
+ default: V,
+) -> tuple[V, 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
+ *extract_value* maps a rule to its value for the field; ``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
+ value: V | None = None
winning_rule: Rule | None = None
for rule in tier_rules:
- new_v = extract_value(rule.controls)
+ new_v = extract_value(rule)
if new_v is not None:
if winning_rule is not None and value != new_v:
log.warning(
@@ -616,7 +837,12 @@ 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 _extract(rule: Rule) -> bool | None:
+ value = rule.controls.get(field)
+ return value if isinstance(value, bool) else None
+
+ return _walk_tiers(tiers, _extract, field, default)
def _resolve_emit_with_source(
@@ -633,10 +859,9 @@ def _resolve_emit_with_source(
"""
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)
+ def _extract(rule: Rule) -> bool | None:
+ value = rule.controls.get(specific_key, rule.controls.get(EMIT))
+ return value if isinstance(value, bool) else None
return _walk_tiers(tiers, _extract, f"emit ({scope})", default)
@@ -860,7 +1085,8 @@ def _apply_authoring_overrides(
return config
_audit_log_authoring_updates(changed, context)
- return replace(config, **changed)
+ # ``changed`` only holds the boolean fields, never ``dataset_filter``.
+ return replace(config, **changed) # type: ignore[arg-type]
def _extend_policy_with_task_authoring(
@@ -882,14 +1108,26 @@ def _extend_policy_with_task_authoring(
flags = _read_param(operator, OL_EMISSION_POLICY_PARAM)
if not flags:
return config
- return _apply_authoring_overrides(
+ config = _apply_authoring_overrides(
config,
locked_fields,
- flags,
+ {k: v for k, v in flags.items() if isinstance(v, bool)},
EMIT_TASK_EVENTS,
context,
extra_fields=(EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE,
HOOK_LINEAGE, INCLUDE_FULL_TASK_INFO),
)
+ # Locks on exclude_datasets depend on the hook reporting each dataset, so
the filter
+ # checks them when it resolves patterns.
+ patterns = flags.get(EXCLUDE_DATASETS)
+ if patterns is None:
+ return config
+ # Params written without the authoring API skip its validation, and a bad
pattern would
+ # otherwise raise during extraction and drop the task event.
+ invalid = find_invalid_pattern(EXCLUDE_DATASETS, patterns)
+ if invalid or not isinstance(patterns, list):
+ log.warning("OpenLineage emission_policy: ignoring authored %s for
%s", invalid, context)
+ return config
+ return replace(config, dataset_filter=replace(config.dataset_filter,
authoring_patterns=tuple(patterns)))
def _extend_policy_with_dag_authoring(
@@ -907,7 +1145,13 @@ def _extend_policy_with_dag_authoring(
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)
+ return _apply_authoring_overrides(
+ config,
+ locked_fields,
+ {k: v for k, v in flags.items() if isinstance(v, bool)},
+ EMIT_DAG_EVENTS,
+ context,
+ )
def _resolve_task_policy_from_conf_only(
@@ -954,6 +1198,15 @@ def _resolve_task_policy_from_conf_only(
locked_fields = _compute_locked_task_fields(task_rules, dag_rules,
operator_rules, global_rules)
+ hook_tiers = _classify_task_rules(rules, fqcn, dag_id, task_id,
hook_scoped=True)
+ dataset_filter_tiers = tuple(
+ tuple(rule for rule in tier if _get_rule_exclude_patterns(rule) is not
None)
+ for pair in zip(hook_tiers, tiers)
+ for tier in pair
+ )
+ if not any(dataset_filter_tiers):
+ dataset_filter_tiers = ()
+
return (
EmissionPolicy(
emit=emit,
@@ -961,6 +1214,7 @@ def _resolve_task_policy_from_conf_only(
include_source_code=include_source_code,
hook_lineage=hook_lineage,
include_full_task_info=include_full_task_info,
+ dataset_filter=DatasetFilter(tiers=dataset_filter_tiers,
context=context),
),
locked_fields,
)
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
index f673971ac1a..1726943750a 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
@@ -46,6 +46,7 @@ from airflow.providers.openlineage.sqlparser import (
from_table_meta,
get_openlineage_facets_with_sql,
)
+from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
get_hook_class_name
from airflow.providers.openlineage.utils.utils import _get_parent_run_facet
if TYPE_CHECKING:
@@ -95,7 +96,12 @@ def _parse_query_into_datasets(
return inputs, outputs
-def emit_lineage_from_sql_extras(task_instance, sql_extras: list,
is_successful: bool = True) -> None:
+def emit_lineage_from_sql_extras(
+ task_instance,
+ sql_extras: list,
+ is_successful: bool = True,
+ dataset_filter: DatasetFilter | None = None,
+) -> None:
"""
Process ``sql_job`` extras and emit per-query OpenLineage events.
@@ -103,8 +109,11 @@ def emit_lineage_from_sql_extras(task_instance,
sql_extras: list, is_successful:
* Parse SQL via :func:`get_openlineage_facets_with_sql` to obtain inputs,
outputs and facets (schema enrichment, column lineage, etc.).
+ * Drop the inputs/outputs that *dataset_filter* excludes for the reporting
hook.
* Emit a separate START + COMPLETE/FAIL event pair (child job of the task).
"""
+ if dataset_filter is None:
+ dataset_filter = DatasetFilter()
if not sql_extras:
return None
@@ -191,13 +200,14 @@ def emit_lineage_from_sql_extras(task_instance,
sql_extras: list, is_successful:
),
)
+ hook_class_name = get_hook_class_name(hook)
events.extend(
_create_ol_event_pair(
task_instance=task_instance,
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.hll_query.{query_count}",
is_successful=is_successful,
- inputs=query_lineage.inputs,
- outputs=query_lineage.outputs,
+ inputs=dataset_filter.exclude(query_lineage.inputs,
hook_class_name),
+ outputs=dataset_filter.exclude(query_lineage.outputs,
hook_class_name),
run_facets=query_lineage.run_facets,
job_facets=query_lineage.job_facets,
)
diff --git a/providers/openlineage/tests/unit/openlineage/api/test_datasets.py
b/providers/openlineage/tests/unit/openlineage/api/test_datasets.py
index b02cc7fafc4..f8d06b0ee05 100644
--- a/providers/openlineage/tests/unit/openlineage/api/test_datasets.py
+++ b/providers/openlineage/tests/unit/openlineage/api/test_datasets.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import dataclasses
from datetime import UTC, datetime
from unittest import mock
@@ -24,6 +25,7 @@ from openlineage.client.event_v2 import Dataset, RunEvent,
RunState
from airflow.providers.openlineage.api.datasets import emit_dataset_lineage
from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter
+from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
EmissionPolicy
_MODULE = "airflow.providers.openlineage.api.datasets"
_CORE = "airflow.providers.openlineage.api.core"
@@ -276,3 +278,29 @@ def
test_passes_include_full_task_info_to_run_facets(patched_emit):
mock_build_run_facets.assert_called_once()
assert mock_build_run_facets.call_args.kwargs["include_full_task_info"] is
True
+
+
+def _policy_excluding(patterns: tuple[str, ...]) -> EmissionPolicy:
+ return dataclasses.replace(
+ EmissionPolicy.defaults(),
dataset_filter=DatasetFilter(authoring_patterns=patterns)
+ )
+
+
+def test_excludes_datasets_from_event(patched_emit):
+ ti = _make_task_instance()
+ with mock.patch(f"{_MODULE}.resolve_task_emission_policy",
return_value=_policy_excluding(("ns/drop",))):
+ emit_dataset_lineage(
+ inputs=[Dataset(namespace="ns", name="drop"),
Dataset(namespace="ns", name="keep")],
+ task_instance=ti,
+ )
+
+ (event,) = patched_emit.call_args.args
+ assert event.inputs == [Dataset(namespace="ns", name="keep")]
+
+
+def test_noop_when_every_dataset_excluded(patched_emit):
+ ti = _make_task_instance()
+ with mock.patch(f"{_MODULE}.resolve_task_emission_policy",
return_value=_policy_excluding(("ns/.*",))):
+ emit_dataset_lineage(inputs=[Dataset(namespace="ns", name="a")],
task_instance=ti)
+
+ patched_emit.assert_not_called()
diff --git
a/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py
b/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py
index 57ff1c124a9..4eb77714b28 100644
--- a/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py
+++ b/providers/openlineage/tests/unit/openlineage/api/test_emission_policy.py
@@ -279,3 +279,32 @@ class TestUnsupportedObjects:
"""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())
+
+
+class TestExcludeDatasets:
+ def test_stored_on_task(self):
+ _, task = _make_dag_and_task()
+ extend_global_openlineage_emission_policy(task,
exclude_datasets=["gs://b/.*"])
+ assert _task_flags(task) == {"exclude_datasets": ["gs://b/.*"]}
+
+ def test_dag_call_propagates_to_tasks_only(self):
+ dag_obj, task = _make_dag_and_task()
+ extend_global_openlineage_emission_policy(dag_obj,
exclude_datasets=["gs://b/.*"])
+ assert _task_flags(task) == {"exclude_datasets": ["gs://b/.*"]}
+ assert _dag_flags(dag_obj) == {}
+
+ def test_xcomarg_delegates_to_operator(self):
+ with DAG(dag_id="test_dag", schedule=None, start_date=now()):
+ op = PythonOperator(task_id="t", python_callable=lambda: 1)
+ extend_global_openlineage_emission_policy(XComArg(op),
exclude_datasets=["x"])
+ assert _task_flags(op) == {"exclude_datasets": ["x"]}
+
+ @pytest.mark.parametrize(
+ "patterns",
+ [pytest.param("gs://b/.*", id="not-a-list"), pytest.param(["("],
id="invalid-regex")],
+ )
+ def test_invalid_patterns_raise(self, patterns):
+ _, task = _make_dag_and_task()
+ with pytest.raises(ValueError, match="exclude_datasets"):
+ extend_global_openlineage_emission_policy(task,
exclude_datasets=patterns)
+ assert _task_flags(task) == {}
diff --git a/providers/openlineage/tests/unit/openlineage/api/test_sql.py
b/providers/openlineage/tests/unit/openlineage/api/test_sql.py
index 9c04a104eec..b04efeb8efe 100644
--- a/providers/openlineage/tests/unit/openlineage/api/test_sql.py
+++ b/providers/openlineage/tests/unit/openlineage/api/test_sql.py
@@ -16,13 +16,15 @@
# under the License.
from __future__ import annotations
+import dataclasses
from datetime import UTC, datetime
from unittest import mock
import pytest
-from openlineage.client.event_v2 import RunState
+from openlineage.client.event_v2 import Dataset, RunState
from airflow.providers.openlineage.api.sql import emit_query_lineage
+from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
EmissionPolicy
_MODULE = "airflow.providers.openlineage.api.sql"
_CORE = "airflow.providers.openlineage.api.core"
@@ -333,3 +335,24 @@ def
test_noop_when_emission_policy_blocks_emit(patched_emit):
)
patched_emit.assert_not_called()
+
+
+def _policy_excluding(patterns: tuple[str, ...]) -> EmissionPolicy:
+ return dataclasses.replace(
+ EmissionPolicy.defaults(),
dataset_filter=DatasetFilter(authoring_patterns=patterns)
+ )
+
+
+def test_excludes_datasets_from_events(patched_emit):
+ ti = _make_task_instance()
+ with mock.patch(f"{_MODULE}.resolve_task_emission_policy",
return_value=_policy_excluding(("ns/drop",))):
+ emit_query_lineage(
+ query_id="qid-1",
+ query_source_namespace="snowflake://ACCT",
+ inputs=[Dataset(namespace="ns", name="drop"),
Dataset(namespace="ns", name="keep")],
+ task_instance=ti,
+ )
+
+ start_event, end_event = (c.args[0] for c in patched_emit.call_args_list)
+ assert [d.name for d in start_event.inputs] == ["keep"]
+ assert [d.name for d in end_event.inputs] == ["keep"]
diff --git
a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
index d153b5b7895..7361a7be49d 100644
--- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
+++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import dataclasses
import functools
import tempfile
from typing import TYPE_CHECKING, Any
@@ -40,6 +41,7 @@ from airflow.providers.openlineage.extractors.manager import (
ExtractorManager,
_is_hook_lineage_collector_created,
)
+from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
EmissionPolicy, Rule
from airflow.providers.openlineage.utils.utils import Asset
from airflow.utils.state import State, TaskInstanceState
@@ -536,7 +538,9 @@ def
test_get_hook_lineage_with_sql_extras_only(hook_lineage_collector):
)
assert result is None
- mock_sql_fn.assert_called_once_with(task_instance=mock_ti,
sql_extras=mock.ANY, is_successful=True)
+ mock_sql_fn.assert_called_once_with(
+ task_instance=mock_ti, sql_extras=mock.ANY, is_successful=True,
dataset_filter=DatasetFilter()
+ )
sql_extras = mock_sql_fn.call_args.kwargs["sql_extras"]
assert len(sql_extras) == 1
assert
sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] ==
"SELECT 1"
@@ -564,7 +568,9 @@ def
test_get_hook_lineage_with_assets_and_sql_extras(hook_lineage_collector):
task_instance_state=TaskInstanceState.SUCCESS,
)
- mock_sql_fn.assert_called_once_with(task_instance=mock_ti,
sql_extras=mock.ANY, is_successful=True)
+ mock_sql_fn.assert_called_once_with(
+ task_instance=mock_ti, sql_extras=mock.ANY, is_successful=True,
dataset_filter=DatasetFilter()
+ )
sql_extras = mock_sql_fn.call_args.kwargs["sql_extras"]
assert len(sql_extras) == 1
assert (
@@ -599,7 +605,9 @@ def
test_get_hook_lineage_sql_extras_multiple_queries(hook_lineage_collector):
task_instance_state=TaskInstanceState.SUCCESS,
)
- mock_sql_fn.assert_called_once_with(task_instance=mock_ti,
sql_extras=mock.ANY, is_successful=True)
+ mock_sql_fn.assert_called_once_with(
+ task_instance=mock_ti, sql_extras=mock.ANY, is_successful=True,
dataset_filter=DatasetFilter()
+ )
sql_extras = mock_sql_fn.call_args.kwargs["sql_extras"]
assert len(sql_extras) == 2
assert
sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] ==
"SELECT a from src1"
@@ -700,7 +708,9 @@ def
test_get_hook_lineage_passes_failed_state(hook_lineage_collector):
task_instance_state=TaskInstanceState.FAILED,
)
- mock_sql_fn.assert_called_once_with(task_instance=mock_ti,
sql_extras=mock.ANY, is_successful=False)
+ mock_sql_fn.assert_called_once_with(
+ task_instance=mock_ti, sql_extras=mock.ANY, is_successful=False,
dataset_filter=DatasetFilter()
+ )
sql_extras = mock_sql_fn.call_args.kwargs["sql_extras"]
assert len(sql_extras) == 1
assert
sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] ==
"SELECT 1"
@@ -907,3 +917,152 @@ def
test_convert_to_ol_dataset_asset_returns_none_when_translation_fails():
return_value=None,
):
assert ExtractorManager.convert_to_ol_dataset(asset) is None
+
+
+class _FakeStorageHook:
+ pass
+
+
+class _FakeWarehouseHook:
+ pass
+
+
+def _hook_fqcn(hook_cls) -> str:
+ return f"{hook_cls.__module__}.{hook_cls.__name__}"
+
+
+def _controls(dataset_filter: DatasetFilter) -> EmissionPolicy:
+ return dataclasses.replace(EmissionPolicy.defaults(),
dataset_filter=dataset_filter)
+
+
+def _hook_exclusion(hook_cls, patterns: list[str]) -> DatasetFilter:
+ return DatasetFilter(
+ tiers=((Rule(scope={"hook": _hook_fqcn(hook_cls)},
controls={"exclude_datasets": patterns}),),)
+ )
+
+
+# gs:// because translating an asset needs the converter of the provider
owning the scheme, and
+# google is a dev dependency of this provider while amazon is not.
+def _collect_gs_parts(collector):
+ storage_hook, warehouse_hook = _FakeStorageHook(), _FakeWarehouseHook()
+ collector.add_output_asset(context=storage_hook,
uri="gs://bucket/staging/part_0")
+ collector.add_output_asset(context=storage_hook,
uri="gs://bucket/final/summary")
+ collector.add_output_asset(context=warehouse_hook,
uri="gs://bucket/staging/part_1")
+
+
+def
test_get_hook_lineage_excludes_datasets_per_reporting_hook(hook_lineage_collector):
+ _collect_gs_parts(hook_lineage_collector)
+
+ result = ExtractorManager().get_hook_lineage(
+ controls=_controls(_hook_exclusion(_FakeStorageHook,
["gs://bucket/staging/.*"]))
+ )
+
+ assert result == OperatorLineage(
+ outputs=[
+ OpenLineageDataset(namespace="gs://bucket", name="final/summary"),
+ OpenLineageDataset(namespace="gs://bucket", name="staging/part_1"),
+ ]
+ )
+
+
+def
test_get_hook_lineage_returns_empty_lineage_when_everything_excluded(hook_lineage_collector):
+ _collect_gs_parts(hook_lineage_collector)
+
+ result = ExtractorManager().get_hook_lineage(
+ controls=_controls(DatasetFilter(authoring_patterns=(".*",)))
+ )
+
+ assert result == OperatorLineage()
+
+
+def
test_get_hook_lineage_passes_dataset_filter_to_sql_extras(hook_lineage_collector):
+ hook_lineage_collector.add_extra(
+ context=MagicMock(),
+ key=SqlJobHookLineageExtra.KEY.value,
+ value={SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value: "SELECT 1"},
+ )
+ dataset_filter = DatasetFilter(authoring_patterns=("x",))
+
+ with patch(_SQL_FN_PATH, return_value=None) as mock_sql_fn:
+ ExtractorManager().get_hook_lineage(
+ task_instance=MagicMock(),
+ task_instance_state=TaskInstanceState.SUCCESS,
+ controls=_controls(dataset_filter),
+ )
+
+ assert mock_sql_fn.call_args.kwargs["dataset_filter"] is dataset_filter
+
+
+def
test_extract_metadata_excluded_hook_lineage_does_not_fall_back_to_inlets(hook_lineage_collector):
+ class FakeUnsupportedOperator(BaseOperator):
+ def execute(self, context: Context) -> Any:
+ pass
+
+ _collect_gs_parts(hook_lineage_collector)
+ task = FakeUnsupportedOperator(
+ task_id="unsupported_task", inlets=[OpenLineageDataset(namespace="ns",
name="inlet")]
+ )
+
+ metadata = ExtractorManager().extract_metadata(
+ dagrun=MagicMock(),
+ task=task,
+ task_instance_state=None,
+ task_instance=MagicMock(),
+ controls=_controls(DatasetFilter(authoring_patterns=("gs://.*",))),
+ )
+
+ assert metadata.inputs == []
+ assert metadata.outputs == []
+
+
+def test_extract_metadata_excludes_extractor_datasets():
+ class FakeSupportedOperator(BaseOperator):
+ def execute(self, context: Context) -> Any:
+ pass
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ return OperatorLineage(
+ inputs=[OpenLineageDataset(namespace="ns", name="keep")],
+ outputs=[OpenLineageDataset(namespace="ns", name="drop")],
+ )
+
+ metadata = ExtractorManager().extract_metadata(
+ dagrun=MagicMock(),
+ task=FakeSupportedOperator(task_id="supported_task"),
+ task_instance_state=None,
+ task_instance=MagicMock(),
+ controls=_controls(DatasetFilter(authoring_patterns=("ns/drop",))),
+ )
+
+ assert metadata.inputs == [OpenLineageDataset(namespace="ns", name="keep")]
+ assert metadata.outputs == []
+
+
+def test_extract_inlets_and_outlets_excludes_static_and_runtime_datasets():
+ task = PythonOperator(
+ task_id="t",
+ python_callable=lambda: 1,
+ inlets=[
+ OpenLineageDataset(namespace="ns", name="drop_in"),
+ OpenLineageDataset(namespace="ns", name="in"),
+ ],
+ outlets=[OpenLineageDataset(namespace="ns", name="out")],
+ )
+ lineage = OperatorLineage()
+
+ with (
+ patch(
+
"airflow.providers.openlineage.extractors.manager.get_runtime_outlet_assets",
+ return_value=[(Asset(uri="s3://bucket/runtime.txt", extra={}),
None)],
+ ),
+ patch(
+
"airflow.providers.openlineage.extractors.manager.translate_airflow_asset",
+ return_value=OpenLineageDataset(namespace="ns",
name="drop_runtime"),
+ ),
+ ):
+ ExtractorManager().extract_inlets_and_outlets(
+ lineage, task, MagicMock(),
_controls(DatasetFilter(authoring_patterns=("ns/drop_.*",)))
+ )
+
+ assert lineage.inputs == [OpenLineageDataset(namespace="ns", name="in")]
+ assert lineage.outputs == [OpenLineageDataset(namespace="ns", name="out")]
diff --git
a/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py
b/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py
index ee501d4a2b8..4784bcede1d 100644
--- a/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py
+++ b/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py
@@ -22,10 +22,12 @@ import warnings
from unittest import mock
import pytest
+from openlineage.client.event_v2 import Dataset
from airflow import DAG
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.openlineage.utils.emission_policy import (
+ DatasetFilter,
EmissionPolicy,
Rule,
resolve_dag_emission_policy,
@@ -2412,3 +2414,254 @@ class TestSelectiveEnableWithLockedRule:
# 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
+
+
+_HOOK = "airflow.providers.google.cloud.hooks.gcs.GCSHook"
+_OTHER_HOOK = "airflow.providers.postgres.hooks.postgres.PostgresHook"
+
+
+def _exclude_rule(scope: dict, patterns: list[str], **extra) -> dict:
+ return {"scope": scope, "controls": {"exclude_datasets": patterns},
**extra}
+
+
+class TestExcludeDatasetsValidation:
+ @pytest.mark.parametrize(
+ "rule",
+ [
+ pytest.param({"scope": {}, "controls": {"exclude_datasets":
"gs://b/.*"}}, id="not-a-list"),
+ pytest.param({"scope": {}, "controls": {"exclude_datasets": [1]}},
id="non-string-pattern"),
+ pytest.param({"scope": {}, "controls": {"exclude_datasets":
["("]}}, id="invalid-regex"),
+ pytest.param({"scope": {"hook": _HOOK}, "controls": {"emit":
False}}, id="hook-scope-emit"),
+ pytest.param(
+ {"scope": {"hook": _HOOK}, "controls": {"exclude_datasets":
[".*"], "hook_lineage": False}},
+ id="hook-scope-both-controls",
+ ),
+ pytest.param(
+ {"scope": {"hook": "("}, "match_mode": "regex", "controls":
{"exclude_datasets": [".*"]}},
+ id="hook-scope-invalid-regex",
+ ),
+ ],
+ )
+ def test_invalid_rule_is_skipped(self, rule):
+ cfg = _resolve_task_controls([rule], operator=MockOperator(),
dag_id="dag", task_id="task")
+ assert cfg == EmissionPolicy.defaults()
+
+ def test_hook_scoped_hook_lineage_does_not_disable_task_hook_lineage(self):
+ cfg = _resolve_task_controls(
+ [{"scope": {"hook": _HOOK}, "controls": {"hook_lineage": False}}],
+ operator=MockOperator(),
+ dag_id="dag",
+ task_id="task",
+ )
+ assert cfg.hook_lineage is True
+ assert cfg.dataset_filter.get_patterns(_HOOK) == (".*",)
+
+
+class TestExcludeDatasetsResolution:
+ @pytest.mark.parametrize(
+ ("rules", "hook", "expected"),
+ [
+ pytest.param([_exclude_rule({}, ["a"])], None, ("a",),
id="global"),
+ pytest.param([_exclude_rule({}, ["a"])], _HOOK, ("a",),
id="global-applies-to-hooks"),
+ pytest.param(
+ [_exclude_rule({}, ["a"]), _exclude_rule({"dag_id": "dag",
"task_id": "task"}, ["b"])],
+ None,
+ ("b",),
+ id="task-replaces-global",
+ ),
+ pytest.param(
+ [_exclude_rule({}, ["a"]), _exclude_rule({"dag_id": "dag",
"task_id": "task"}, [])],
+ None,
+ (),
+ id="task-clears-global",
+ ),
+ pytest.param([_exclude_rule({"hook": _HOOK}, ["a"])], _HOOK,
("a",), id="hook-matches"),
+ pytest.param([_exclude_rule({"hook": _HOOK}, ["a"])], _OTHER_HOOK,
(), id="hook-other-hook"),
+ pytest.param([_exclude_rule({"hook": _HOOK}, ["a"])], None, (),
id="hook-no-hook"),
+ pytest.param(
+ [_exclude_rule({"hook": ".*GCSHook"}, ["a"],
match_mode="regex")],
+ _HOOK,
+ ("a",),
+ id="hook-regex",
+ ),
+ pytest.param(
+ [_exclude_rule({"dag_id": "other", "hook": _HOOK}, ["a"])],
+ _HOOK,
+ (),
+ id="dag-hook-other-dag",
+ ),
+ pytest.param(
+ [
+ _exclude_rule({"dag_id": "dag", "task_id": "task"},
["task"]),
+ _exclude_rule({"dag_id": "dag", "task_id": "task", "hook":
_HOOK}, ["task+hook"]),
+ ],
+ _HOOK,
+ ("task+hook",),
+ id="task-hook-beats-task",
+ ),
+ pytest.param(
+ [
+ _exclude_rule({"dag_id": "dag", "hook": _HOOK},
["dag+hook"]),
+ _exclude_rule({"dag_id": "dag", "task_id": "task"},
["task"]),
+ ],
+ _HOOK,
+ ("task",),
+ id="task-beats-dag-hook",
+ ),
+ pytest.param(
+ [
+ _exclude_rule({"dag_id": "dag"}, ["dag"]),
+ _exclude_rule({"dag_id": "dag", "hook": _HOOK},
["dag+hook"]),
+ ],
+ _HOOK,
+ ("dag+hook",),
+ id="dag-hook-beats-dag",
+ ),
+ pytest.param(
+ [
+ _exclude_rule({"operator":
"tests.mock_module.MockOperator"}, ["op"]),
+ _exclude_rule({"operator":
"tests.mock_module.MockOperator", "hook": _HOOK}, ["op+hook"]),
+ ],
+ _HOOK,
+ ("op+hook",),
+ id="operator-hook-beats-operator",
+ ),
+ pytest.param(
+ [_exclude_rule({"hook": _HOOK}, ["hook"]),
_exclude_rule({"operator": "x.Y"}, ["op"])],
+ _HOOK,
+ ("hook",),
+ id="operator-rule-not-matching",
+ ),
+ pytest.param(
+ [_exclude_rule({}, ["global"]), _exclude_rule({"hook": _HOOK},
["hook"])],
+ _HOOK,
+ ("hook",),
+ id="hook-beats-global",
+ ),
+ pytest.param(
+ [
+ {"scope": {"hook": _HOOK}, "controls": {"hook_lineage":
False}},
+ {"scope": {"dag_id": "dag", "hook": _HOOK}, "controls":
{"hook_lineage": True}},
+ ],
+ _HOOK,
+ (),
+ id="hook-lineage-true-clears-broader-hook-rule",
+ ),
+ ],
+ )
+ def test_get_patterns(self, rules, hook, expected):
+ cfg = _resolve_task_controls(rules, operator=MockOperator(),
dag_id="dag", task_id="task")
+ assert cfg.dataset_filter.get_patterns(hook) == expected
+
+ @pytest.mark.parametrize(
+ ("patterns", "dataset", "excluded"),
+ [
+ pytest.param(
+ ["gs://bucket/staging/.*"], Dataset("gs://bucket",
"staging/part_1"), True, id="match"
+ ),
+ pytest.param(
+ ["staging/.*"], Dataset("gs://bucket", "staging/part_1"),
False, id="fullmatch-only"
+ ),
+ pytest.param(
+ ["gs://bucket/staging/.*"], Dataset("gs://bucket",
"final/out"), False, id="no-match"
+ ),
+ ],
+ )
+ def test_is_excluded_matches_namespace_and_name(self, patterns, dataset,
excluded):
+ cfg = _resolve_task_controls(
+ [_exclude_rule({}, patterns)], operator=MockOperator(),
dag_id="dag", task_id="task"
+ )
+ assert cfg.dataset_filter.is_excluded(dataset) is excluded
+
+ def test_exclude_keeps_order_of_remaining_datasets(self):
+ datasets = [Dataset("gs://b", "keep_1"), Dataset("gs://b", "drop"),
Dataset("gs://b", "keep_2")]
+ cfg = _resolve_task_controls(
+ [_exclude_rule({}, ["gs://b/drop"])], operator=MockOperator(),
dag_id="dag", task_id="task"
+ )
+ assert cfg.dataset_filter.exclude(datasets) == [datasets[0],
datasets[2]]
+
+ def test_audit_log_names_replaced_broader_rules(self):
+ global_rule = _exclude_rule({}, ["a"])
+ task_rule = _exclude_rule({"dag_id": "dag", "task_id": "task"}, ["b"])
+ cfg = _resolve_task_controls(
+ [global_rule, task_rule], operator=MockOperator(), dag_id="dag",
task_id="task"
+ )
+ with
mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as
mock_log:
+ cfg.dataset_filter.get_patterns(None)
+ cfg.dataset_filter.get_patterns(None)
+ mock_log.info.assert_called_once_with(
+ "OpenLineage emission policy: '%s' set to %r for %s by %r,
replacing broader rule(s) %r",
+ "exclude_datasets",
+ ["b"],
+ "task 'task' in dag 'dag'",
+ Rule(scope=task_rule["scope"], controls=task_rule["controls"]),
+ [Rule(scope={}, controls=global_rule["controls"])],
+ )
+
+
+class TestExcludeDatasetsAuthoring:
+ def test_authoring_replaces_conf_patterns_for_every_source(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,
exclude_datasets=["authored"])
+ rules = [_exclude_rule({}, ["global"]), _exclude_rule({"hook": _HOOK},
["hook"])]
+ cfg = _resolve_task_controls(rules, operator=task, dag_id="test_dag",
task_id="test_task")
+
+ assert cfg.dataset_filter.get_patterns(None) == ("authored",)
+ assert cfg.dataset_filter.get_patterns(_HOOK) == ("authored",)
+
+ def test_authoring_without_conf_rules(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,
exclude_datasets=["authored"])
+ cfg = _resolve_task_controls([], operator=task, dag_id="test_dag",
task_id="test_task")
+
+ assert cfg.dataset_filter ==
DatasetFilter(authoring_patterns=("authored",))
+ assert cfg.dataset_filter.get_patterns(None) == ("authored",)
+
+ def test_locked_rule_blocks_authoring_only_where_it_applies(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,
exclude_datasets=["authored"])
+ rules = [_exclude_rule({"hook": _HOOK}, ["locked"], locked=True)]
+ cfg = _resolve_task_controls(rules, operator=task, dag_id="test_dag",
task_id="test_task")
+
+ with
mock.patch("airflow.providers.openlineage.utils.emission_policy.log") as
mock_log:
+ assert cfg.dataset_filter.get_patterns(_HOOK) == ("locked",)
+ mock_log.warning.assert_called_once()
+ assert cfg.dataset_filter.get_patterns(_OTHER_HOOK) == ("authored",)
+ assert cfg.dataset_filter.get_patterns(None) == ("authored",)
+
+ def test_locked_rule_does_not_block_more_specific_conf_rule(self):
+ rules = [
+ _exclude_rule({}, ["locked"], locked=True),
+ _exclude_rule({"dag_id": "dag", "task_id": "task"}, ["task"]),
+ ]
+ cfg = _resolve_task_controls(rules, operator=MockOperator(),
dag_id="dag", task_id="task")
+ assert cfg.dataset_filter.get_patterns(None) == ("task",)
+
+ @pytest.mark.parametrize(
+ "patterns", [pytest.param(["("], id="invalid-regex"),
pytest.param([1], id="non-string")]
+ )
+ def test_invalid_authored_patterns_bypassing_api_are_ignored(self,
patterns):
+ _, task = _make_dag_and_task()
+ task.params["_openlineage_emission_policy"] = {"exclude_datasets":
patterns}
+ cfg = _resolve_task_controls(
+ [_exclude_rule({}, ["global"])], operator=task, dag_id="test_dag",
task_id="test_task"
+ )
+ assert cfg.dataset_filter.get_patterns(None) == ("global",)
+
+ def test_policy_with_exclude_rules_is_hashable(self):
+ cfg = _resolve_task_controls(
+ [_exclude_rule({}, ["a"])], operator=MockOperator(), dag_id="dag",
task_id="task"
+ )
+ assert hash(cfg) == hash(EmissionPolicy.defaults())
diff --git
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
index b2a716c97e7..bbe23dbb439 100644
---
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
+++
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
@@ -26,6 +26,7 @@ from openlineage.client.facet_v2 import external_query_run,
sql_job
from airflow.providers.common.sql.hooks.lineage import SqlJobHookLineageExtra
from airflow.providers.openlineage.extractors.base import OperatorLineage
from airflow.providers.openlineage.sqlparser import SQLParser
+from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
Rule
from airflow.providers.openlineage.utils.sql_hook_lineage import (
_get_hook_conn_id,
emit_lineage_from_sql_extras,
@@ -142,6 +143,33 @@ class TestEmitLineageFromSqlExtras:
mock.call(mock.sentinel.end_event),
]
+ def test_datasets_excluded_for_reporting_hook_only(self):
+ class FakeHook:
+ def get_connection(self, conn_id):
+ return mock.MagicMock()
+
+ def get_openlineage_database_info(self, connection):
+ return mock.MagicMock()
+
+ self.mock_ns.return_value = "postgres://host/db"
+ in_table = OpenLineageDataset(namespace="ns", name="in_table")
+ out_table = OpenLineageDataset(namespace="ns", name="out_table")
+ self.mock_facets_fn.side_effect = lambda **_:
OperatorLineage(inputs=[in_table], outputs=[out_table])
+ rule = Rule(
+ scope={"hook": f"{FakeHook.__module__}.{FakeHook.__name__}"},
+ controls={"exclude_datasets": ["ns/in_table"]},
+ )
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock.MagicMock(dag_id="dag_id", task_id="task_id"),
+ sql_extras=[_make_extra(sql="SELECT 1", hook=FakeHook()),
_make_extra(sql="SELECT 2")],
+ dataset_filter=DatasetFilter(tiers=((rule,),)),
+ )
+
+ first, second = self.mock_build.call_args_list
+ assert (first.kwargs["inputs"], first.kwargs["outputs"]) == ([],
[out_table])
+ assert (second.kwargs["inputs"], second.kwargs["outputs"]) ==
([in_table], [out_table])
+
def test_multiple_queries_increment_counter(self):
self.mock_ns.return_value = "postgres://host/db"
mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
