diff --git a/providers/openlineage/docs/emission_policy.rst b/providers/openlineage/docs/emission_policy.rst index 726f91cfe2a49..d1c35b57502ec 100644 --- a/providers/openlineage/docs/emission_policy.rst +++ b/providers/openlineage/docs/emission_policy.rst @@ -144,6 +144,16 @@ pattern. inputs/outputs. Task events only. **No effect when ``extract_operator_metadata`` is ``false``** — the entire extraction pipeline (including hook lineage) is skipped. + * - ``exclude_hook_lineage_assets`` + - ``[]`` + - List of regex patterns; hook-collected assets whose URI matches any of them are dropped. + Task events only. **No effect when ``hook_lineage`` is ``false``** — hook lineage is not + collected at all. + * - ``exclude_hook_lineage_hooks`` + - ``[]`` + - List of regex patterns; hook-collected assets reported by a hook whose **fully-qualified** + class name (``module.ClassName``, as with the ``operator`` scope key) matches any of them are + dropped. Task events only. **No effect when ``hook_lineage`` is ``false``**. * - ``include_full_task_info`` - ``false`` - Whether to include the full serialized operator state in ``AirflowRunFacet``. When @@ -151,6 +161,22 @@ pattern. serializable task parameters are included, which may significantly increase event size. Task events only. +Both exclude lists are always matched with ``re.fullmatch``, independent of the rule's +``match_mode`` (which governs ``scope`` only). Unlike the boolean controls they are **replaced** +rather than merged across tiers: the most specific matching rule that sets a list wins outright, +so a task-scoped rule can narrow — or clear, with ``[]`` — a broader global exclusion. Both +filters apply to hook-collected assets only; SQL-based hook lineage is unaffected. + +.. note:: + + **Exclusion happens after the collector cap.** These filters run when OpenLineage reads the + hook lineage collector, which is after + :ref:`[lineage] max_assets_per_collector` has + already limited what the collector retained. Excluding assets reduces what gets reported, but + it does not create room for assets the cap already discarded — a task that writes more objects + than the cap allows never collected the later ones in the first place. Reducing the number of + assets a hook registers requires changes in the hooks themselves. + ``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 @@ -193,6 +219,7 @@ suppresses DAG-level events for X while leaving task events enabled. {"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": "chunked_dag", "task_id": "upload"}, "controls": {"exclude_hook_lineage_assets": ["s3://my-bucket/staging/part_.*"]}}, {"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}}, @@ -212,6 +239,8 @@ In the example above: with a more-specific task-scoped rule and extracts normally. - ``sensitive_task`` in ``my_dag`` runs its extractor normally but skips the ``HookLineageCollector`` fallback. +- ``upload`` in ``chunked_dag`` keeps hook lineage but drops the staging part objects it writes, + so only the assets that are not staging parts are reported. - 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. @@ -253,6 +282,11 @@ 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) + # Keep hook lineage, but drop the per-chunk objects a loop uploads + extend_global_openlineage_emission_policy( + extract, exclude_hook_lineage_assets=["s3://my-bucket/staging/part_.*"] + ) + Key semantics: - **Resolution priority**: authoring flags > unlocked conf rules > built-in defaults. 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 898aae41a6500..a64f8bb4f1c99 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py +++ b/providers/openlineage/src/airflow/providers/openlineage/api/emission_policy.py @@ -91,12 +91,16 @@ EMIT, EMIT_DAG_EVENTS, EMIT_TASK_EVENTS, + EXCLUDE_HOOK_LINEAGE_ASSETS, + EXCLUDE_HOOK_LINEAGE_HOOKS, 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: @@ -120,6 +124,8 @@ def extend_global_openlineage_emission_policy( extract_operator_metadata: bool | None = None, include_source_code: bool | None = None, hook_lineage: bool | None = None, + exclude_hook_lineage_assets: list[str] | None = None, + exclude_hook_lineage_hooks: list[str] | None = None, include_full_task_info: bool | None = None, ) -> T: """ @@ -132,8 +138,9 @@ 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``, ``exclude_hook_lineage_assets``, + ``exclude_hook_lineage_hooks``, ``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. @@ -179,7 +186,13 @@ def extend_global_openlineage_emission_policy( :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 exclude_hook_lineage_assets: Regex patterns; hook-collected assets whose URI matches + any of them are dropped. Has no effect when ``hook_lineage`` is disabled. + :param exclude_hook_lineage_hooks: Regex patterns; hook-collected assets reported by a hook + whose fully-qualified class name (``module.ClassName``) matches any of them are dropped. + Has no effect when ``hook_lineage`` is disabled. :param include_full_task_info: Whether to include the full serialized operator state. + :raises ValueError: If either exclusion argument 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): @@ -191,6 +204,8 @@ def extend_global_openlineage_emission_policy( extract_operator_metadata=extract_operator_metadata, include_source_code=include_source_code, hook_lineage=hook_lineage, + exclude_hook_lineage_assets=exclude_hook_lineage_assets, + exclude_hook_lineage_hooks=exclude_hook_lineage_hooks, include_full_task_info=include_full_task_info, ) return obj @@ -206,7 +221,20 @@ def extend_global_openlineage_emission_policy( "configuration with a global rule ('scope': {}) instead." ) - provided: dict[str, bool] = { + # Validate eagerly: a bad pattern stored here would otherwise surface much later as an + # re.error inside hook lineage extraction, which the extractor manager swallows — the Dag + # would silently lose all lineage instead of reporting the typo. + for field, patterns in ( + (EXCLUDE_HOOK_LINEAGE_ASSETS, exclude_hook_lineage_assets), + (EXCLUDE_HOOK_LINEAGE_HOOKS, exclude_hook_lineage_hooks), + ): + if patterns is None: + continue + invalid = find_invalid_pattern(field, patterns) + if invalid: + raise ValueError(f"extend_global_openlineage_emission_policy(): {invalid}") + + provided: dict[str, ControlValue] = { k: v for k, v in { EMIT: emit, @@ -215,6 +243,8 @@ def extend_global_openlineage_emission_policy( EXTRACT_OPERATOR_METADATA: extract_operator_metadata, INCLUDE_SOURCE_CODE: include_source_code, HOOK_LINEAGE: hook_lineage, + EXCLUDE_HOOK_LINEAGE_ASSETS: exclude_hook_lineage_assets, + EXCLUDE_HOOK_LINEAGE_HOOKS: exclude_hook_lineage_hooks, INCLUDE_FULL_TASK_INFO: include_full_task_info, }.items() if v is not None diff --git a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py index 7509044958072..ef50200e70e0d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py +++ b/providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import re from collections.abc import Iterator from typing import TYPE_CHECKING @@ -137,7 +138,7 @@ def extract_metadata( # 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 @@ -166,7 +167,7 @@ def extract_metadata( # 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.", @@ -243,10 +244,39 @@ def extract_inlets_and_outlets( task_metadata.outputs.append(ol) seen.add((ol.namespace, ol.name)) + def _is_excluded_asset( + self, + asset_info, + exclude_assets: tuple[str, ...], + exclude_hooks: tuple[str, ...], + ) -> bool: + """ + Return ``True`` if *asset_info* matches any asset-URI or hook-class exclusion pattern. + + Hooks are matched on their fully-qualified class name, mirroring how the ``operator`` + scope key identifies operators, so two identically named classes from different modules + stay distinguishable. + """ + if exclude_assets: + uri = getattr(asset_info.asset, "uri", None) + if uri and any(re.fullmatch(p, uri) for p in exclude_assets): + self.log.debug("Excluding hook-collected asset %r by asset pattern.", uri) + return True + + if exclude_hooks: + hook_type = type(asset_info.context) + hook_name = f"{hook_type.__module__}.{hook_type.__name__}" + if any(re.fullmatch(p, hook_name) for p in exclude_hooks): + self.log.debug("Excluding hook-collected asset reported by hook %r.", hook_name) + return True + + return False + 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. @@ -259,8 +289,13 @@ def get_hook_lineage( When ``task_instance`` is provided, each extra is parsed and separate per-query OpenLineage events are emitted. + Assets matching the ``exclude_hook_lineage_assets`` / ``exclude_hook_lineage_hooks`` + patterns of *controls* are dropped. SQL-based lineage is not filtered. + Returns ``None`` when nothing was collected. """ + 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 @@ -276,16 +311,22 @@ def get_hook_lineage( self.log.debug("OpenLineage will extract lineage from Hook Lineage Collector.") collected = collector.collected_assets - # Asset-based inputs/outputs - keep only assets that can be translated to OL datasets + exclude_assets = controls.exclude_hook_lineage_assets + exclude_hooks = controls.exclude_hook_lineage_hooks + + # Asset-based inputs/outputs - keep only assets that are not excluded by policy and + # 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 + if not self._is_excluded_asset(asset_info, exclude_assets, exclude_hooks) + and (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 + if not self._is_excluded_asset(asset_info, exclude_assets, exclude_hooks) + and (asset := translate_airflow_asset(asset_info.asset, asset_info.context)) is not None ] # SQL-based lineage - keep only SQL extra with query_text or job_id. 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 de386fce438f5..1e13ed77d15ad 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/emission_policy.py @@ -70,10 +70,30 @@ 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``. +- ``exclude_hook_lineage_assets`` — list of regex patterns; hook-collected assets whose URI + matches any of them are dropped. Only meaningful for task events when ``hook_lineage`` + is ``true``. Default: ``[]`` (drop nothing). +- ``exclude_hook_lineage_hooks`` — list of regex patterns; hook-collected assets reported by + a hook whose **fully-qualified** class name (``module.ClassName``, as with the ``operator`` + scope key) matches any of them are dropped. Only meaningful for task events when + ``hook_lineage`` is ``true``. Default: ``[]`` (drop nothing). - ``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``. +Both exclude lists are always matched as ``re.fullmatch`` patterns, independent of the +rule's ``match_mode`` (which governs ``scope`` only) — a URI or hook name is rarely worth +excluding one exact value at a time. Unlike the boolean flags they are not merged across +tiers: the most specific rule that sets a list replaces any broader one, so a task-scoped +rule can narrow a global exclusion rather than only adding to it. + +Both filters run when OpenLineage **reads** the collector, which is *after* +``[lineage] max_assets_per_collector`` has already capped what the collector kept. Excluding +assets therefore reduces what is reported, but it does not free capacity for assets the cap +already dropped: if a task writes more objects than the cap allows, the ones written past the +limit were never collected and no exclusion pattern can bring them back. Lowering the volume +of collected assets requires producer-side changes in the hooks themselves. + ``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. @@ -87,6 +107,8 @@ 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. +- ``hook_lineage: false`` skips hook lineage collection entirely, so + ``exclude_hook_lineage_assets`` and ``exclude_hook_lineage_hooks`` have **no effect**. - ``include_source_code`` only applies inside Python and Bash operator extractors; setting it to ``false`` on other operator types is a no-op. @@ -150,6 +172,10 @@ log = logging.getLogger(__name__) +# A validated ``controls`` value: bool for the flags, list of patterns for the +# ``exclude_hook_lineage_*`` controls. +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" @@ -157,8 +183,14 @@ EXTRACT_OPERATOR_METADATA = "extract_operator_metadata" INCLUDE_SOURCE_CODE = "include_source_code" HOOK_LINEAGE = "hook_lineage" +EXCLUDE_HOOK_LINEAGE_ASSETS = "exclude_hook_lineage_assets" +EXCLUDE_HOOK_LINEAGE_HOOKS = "exclude_hook_lineage_hooks" INCLUDE_FULL_TASK_INFO = "include_full_task_info" +# Controls whose value is a list of patterns rather than a bool. They are validated and +# resolved separately from the boolean flags. +_PATTERN_LIST_CONTROLS: frozenset[str] = frozenset({EXCLUDE_HOOK_LINEAGE_ASSETS, EXCLUDE_HOOK_LINEAGE_HOOKS}) + # Scope keys — the keys that may appear inside a rule's ``scope`` dict. SCOPE_DAG_ID = "dag_id" SCOPE_TASK_ID = "task_id" @@ -187,6 +219,8 @@ EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE, HOOK_LINEAGE, + EXCLUDE_HOOK_LINEAGE_ASSETS, + EXCLUDE_HOOK_LINEAGE_HOOKS, INCLUDE_FULL_TASK_INFO, } ) @@ -199,6 +233,8 @@ EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE, HOOK_LINEAGE, + EXCLUDE_HOOK_LINEAGE_ASSETS, + EXCLUDE_HOOK_LINEAGE_HOOKS, INCLUDE_FULL_TASK_INFO, ) @@ -212,6 +248,8 @@ EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE, HOOK_LINEAGE, + EXCLUDE_HOOK_LINEAGE_ASSETS, + EXCLUDE_HOOK_LINEAGE_HOOKS, INCLUDE_FULL_TASK_INFO, } ) @@ -227,6 +265,8 @@ class EmissionPolicy: include_source_code: bool hook_lineage: bool include_full_task_info: bool + exclude_hook_lineage_assets: tuple[str, ...] = () + exclude_hook_lineage_hooks: tuple[str, ...] = () @classmethod def defaults(cls) -> EmissionPolicy: @@ -251,7 +291,7 @@ class Rule: """ scope: dict[str, str] - controls: dict[str, bool] + controls: dict[str, ControlValue] match_mode: str = "exact" locked: bool = False @@ -275,7 +315,26 @@ 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 find_invalid_pattern(field: str, value: object) -> str | None: + """ + Return an error message if *value* is not a usable list of 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 asset matching has already compiled once here. + """ + if not isinstance(value, list): + return f"'{field}' must be a list of strings, got {value!r}" + for pattern in value: + if not isinstance(pattern, str): + return f"'{field}' must contain only strings, got {pattern!r}" + try: + re.compile(pattern) + except re.error as exc: + return f"'{field}' 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 +344,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,17 +380,32 @@ def _audit_log_conf_field(field: str, value: bool, context: str, source: Rule) - ) +def _audit_log_conf_patterns(field: str, patterns: tuple[str, ...], context: str, source: Rule) -> None: + """Log the winning conf value for a pattern-list field whenever any rule set it.""" + log.info( + "OpenLineage emission policy: '%s' set to %r for %s by %r", + field, + list(patterns), + context, + source, + ) + + def _audit_log_authoring_updates( - field_changes: dict[str, bool], + field_changes: dict[str, bool | tuple[str, ...]], context: str, ) -> None: """Audit-log authoring overrides that actually change the resolved policy value.""" for field, new_value in field_changes.items(): + if isinstance(new_value, bool): + state = "enabled" if new_value else "disabled" + else: + state = f"set to {list(new_value)!r}" log.info( "OpenLineage emission policy: '%s' %s for %s " "by manual `extend_global_openlineage_emission_policy` call.", field, - "enabled" if new_value else "disabled", + state, context, ) @@ -433,7 +507,16 @@ def _parse_rule(rule: object) -> Rule | None: ) return None for k, v in controls.items(): - if not isinstance(v, bool): + if k in _PATTERN_LIST_CONTROLS: + 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, @@ -576,7 +659,7 @@ def _classify_dag_rules( def _walk_tiers( tiers: list[list[Rule]], - extract_value: Callable[[dict[str, bool]], bool | None], + extract_value: Callable[[dict[str, ControlValue]], bool | None], field_label: str, default: bool, ) -> tuple[bool, Rule | None]: @@ -612,11 +695,42 @@ def _walk_tiers( return default, None +def _read_bool_control(controls: dict[str, ControlValue], field: str) -> bool | None: + """ + Read a boolean control, or ``None`` when *field* is absent. + + Rule validation already guarantees boolean controls hold bools, so the isinstance check + only narrows the union for the type checker. + """ + value = controls.get(field) + return value if isinstance(value, bool) else 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) + return _walk_tiers(tiers, lambda c: _read_bool_control(c, field), field, default) + + +def _resolve_pattern_list(tiers: list[list[Rule]], field: str) -> tuple[tuple[str, ...], Rule | None]: + """ + Walk tiers for a pattern-list *field*, returning the winning rule's list verbatim. + + Lists replace rather than merge: the most specific tier that sets *field* wins outright, + so a task-scoped rule can narrow a broad global exclusion instead of adding to it. + Within a tier the last matching rule wins, mirroring the boolean resolution order. + """ + for tier_rules in tiers: + winning_rule: Rule | None = None + for rule in tier_rules: + if field in rule.controls: + winning_rule = rule + if winning_rule is not None: + patterns = winning_rule.controls[field] + # Validation guarantees a list here; the check only narrows the union. + return (tuple(patterns) if isinstance(patterns, list) else ()), winning_rule + return (), None def _resolve_emit_with_source( @@ -633,10 +747,10 @@ 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: + def _extract(c: dict[str, ControlValue]) -> bool | None: if specific_key in c: - return c[specific_key] - return c.get(EMIT) + return _read_bool_control(c, specific_key) + return _read_bool_control(c, EMIT) return _walk_tiers(tiers, _extract, f"emit ({scope})", default) @@ -815,7 +929,7 @@ def _compute_locked_task_fields( def _apply_authoring_overrides( config: EmissionPolicy, locked_fields: frozenset[str], - flags: dict[str, bool], + flags: dict[str, ControlValue], emit_key: str, context: str, extra_fields: tuple[str, ...] = (), @@ -827,9 +941,13 @@ def _apply_authoring_overrides( ``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] = {} + # Values here are already normalised to their EmissionPolicy field types, so pattern lists + # are tuples rather than the lists the authoring API accepts. + updates: dict[str, bool | tuple[str, ...]] = {} - effective_emit: bool | None = flags.get(emit_key, flags.get(EMIT)) + effective_emit = _read_bool_control(flags, emit_key) + if effective_emit is None: + effective_emit = _read_bool_control(flags, EMIT) if effective_emit is not None: if EMIT in locked_fields: log.warning( @@ -853,14 +971,17 @@ def _apply_authoring_overrides( getattr(config, field), ) else: - updates[field] = flags[field] + value = flags[field] + # Pattern lists are stored as tuples on the policy, so normalise before comparing + # to avoid a list-vs-tuple mismatch always reading as "changed". + updates[field] = tuple(value) if isinstance(value, list) else value 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) + return replace(config, **changed) # type: ignore[arg-type] def _extend_policy_with_task_authoring( @@ -888,7 +1009,14 @@ def _extend_policy_with_task_authoring( flags, EMIT_TASK_EVENTS, context, - extra_fields=(EXTRACT_OPERATOR_METADATA, INCLUDE_SOURCE_CODE, HOOK_LINEAGE, INCLUDE_FULL_TASK_INFO), + extra_fields=( + EXTRACT_OPERATOR_METADATA, + INCLUDE_SOURCE_CODE, + HOOK_LINEAGE, + EXCLUDE_HOOK_LINEAGE_ASSETS, + EXCLUDE_HOOK_LINEAGE_HOOKS, + INCLUDE_FULL_TASK_INFO, + ), ) @@ -937,6 +1065,8 @@ def _resolve_task_policy_from_conf_only( tiers, INCLUDE_SOURCE_CODE, defaults.include_source_code ) hook_lineage, hl_rule = _resolve_field_with_source(tiers, HOOK_LINEAGE, defaults.hook_lineage) + exclude_assets, ea_rule = _resolve_pattern_list(tiers, EXCLUDE_HOOK_LINEAGE_ASSETS) + exclude_hooks, eh_rule = _resolve_pattern_list(tiers, EXCLUDE_HOOK_LINEAGE_HOOKS) include_full_task_info, ift_rule = _resolve_field_with_source( tiers, INCLUDE_FULL_TASK_INFO, defaults.include_full_task_info ) @@ -949,6 +1079,10 @@ def _resolve_task_policy_from_conf_only( _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 ea_rule is not None: + _audit_log_conf_patterns(EXCLUDE_HOOK_LINEAGE_ASSETS, exclude_assets, context, ea_rule) + if eh_rule is not None: + _audit_log_conf_patterns(EXCLUDE_HOOK_LINEAGE_HOOKS, exclude_hooks, context, eh_rule) if ift_rule is not None: _audit_log_conf_field(INCLUDE_FULL_TASK_INFO, include_full_task_info, context, ift_rule) @@ -961,6 +1095,8 @@ 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, + exclude_hook_lineage_assets=exclude_assets, + exclude_hook_lineage_hooks=exclude_hooks, ), locked_fields, ) 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 57ff1c124a985..ab9b8600a1f45 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,43 @@ 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()) + + +class TestExcludePatternValidation: + """Malformed exclusion patterns are rejected at authoring time, not at extraction time.""" + + FIELDS = ("exclude_hook_lineage_assets", "exclude_hook_lineage_hooks") + + @pytest.mark.parametrize("field", FIELDS) + @pytest.mark.parametrize( + ("value", "message"), + [ + pytest.param("s3://bucket/.*", "must be a list of strings", id="string-instead-of-list"), + pytest.param([123], "must contain only strings", id="non-string-element"), + pytest.param(["["], "is not a valid regex", id="unterminated-character-set"), + pytest.param(["ok.*", "(unclosed"], "is not a valid regex", id="second-pattern-invalid"), + ], + ) + def test_invalid_patterns_raise_value_error(self, field, value, message): + _, task = _make_dag_and_task() + with pytest.raises(ValueError, match=message): + extend_global_openlineage_emission_policy(task, **{field: value}) + + @pytest.mark.parametrize("field", FIELDS) + def test_invalid_patterns_store_nothing(self, field): + _, task = _make_dag_and_task() + with pytest.raises(ValueError, match="is not a valid regex"): + extend_global_openlineage_emission_policy(task, **{field: ["["]}) + assert _task_flags(task) == {} + + @pytest.mark.parametrize("field", FIELDS) + def test_valid_patterns_are_stored(self, field): + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, **{field: ["s3://bucket/part_.*"]}) + assert _task_flags(task)[field] == ["s3://bucket/part_.*"] + + @pytest.mark.parametrize("field", FIELDS) + def test_empty_list_is_accepted(self, field): + _, task = _make_dag_and_task() + extend_global_openlineage_emission_policy(task, **{field: []}) + assert _task_flags(task)[field] == [] diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index d026c1b5039e6..70cf78f332938 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -18,6 +18,8 @@ from __future__ import annotations import tempfile +import types +from dataclasses import replace from typing import TYPE_CHECKING, Any from unittest import mock from unittest.mock import MagicMock, patch @@ -36,6 +38,7 @@ from airflow.providers.common.sql.hooks.lineage import SqlJobHookLineageExtra from airflow.providers.openlineage.extractors import OperatorLineage from airflow.providers.openlineage.extractors.manager import ExtractorManager +from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy from airflow.providers.openlineage.utils.utils import Asset from airflow.utils.state import State, TaskInstanceState @@ -639,6 +642,190 @@ def test_get_hook_lineage_passes_failed_state(hook_lineage_collector): assert sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] == "SELECT 1" +class _FakeStorageHook: + pass + + +class _FakeWarehouseHook: + pass + + +def _hook_fqcn(hook_cls) -> str: + return f"{hook_cls.__module__}.{hook_cls.__name__}" + + +# Assets are built on gs:// because translating an Airflow asset to an OL dataset needs the +# scheme's converter from the owning provider, and google is a dev dependency of this provider +# while amazon is not - s3:// URIs silently fail to translate when only the lowest set of +# dependencies is installed. +def _collect_assets(collector): + """ + Seed the collector with staging parts plus unrelated assets from two hooks. + + All of these fit well inside ``max_assets_per_collector``; exclusion runs after that cap, + so this fixture deliberately does not model a collection that overflows it. + """ + storage_hook, warehouse_hook = _FakeStorageHook(), _FakeWarehouseHook() + for idx in range(3): + collector.add_output_asset(context=storage_hook, uri=f"gs://bucket/prefix/part_{idx}.txt") + collector.add_output_asset(context=storage_hook, uri="gs://bucket/prefix/summary.json") + collector.add_output_asset(context=warehouse_hook, uri="gs://other-bucket/out.csv") + + +def _output_uris(lineage): + return sorted(f"{dataset.namespace}/{dataset.name}" for dataset in lineage.outputs) + + +@pytest.mark.parametrize( + ("controls_kwargs", "expected"), + [ + pytest.param( + {}, + [ + "gs://bucket/prefix/part_0.txt", + "gs://bucket/prefix/part_1.txt", + "gs://bucket/prefix/part_2.txt", + "gs://bucket/prefix/summary.json", + "gs://other-bucket/out.csv", + ], + id="no-filter-keeps-everything", + ), + pytest.param( + {"exclude_hook_lineage_assets": ("gs://bucket/prefix/part_.*",)}, + ["gs://bucket/prefix/summary.json", "gs://other-bucket/out.csv"], + id="asset-pattern-drops-parts-keeps-others", + ), + pytest.param( + {"exclude_hook_lineage_hooks": (_hook_fqcn(_FakeStorageHook),)}, + ["gs://other-bucket/out.csv"], + id="hook-pattern-drops-all-assets-from-that-hook", + ), + pytest.param( + {"exclude_hook_lineage_assets": (".*",)}, + [], + id="catch-all-asset-pattern-drops-everything", + ), + pytest.param( + {"exclude_hook_lineage_assets": ("gs://bucket/prefix/part_0.txt",)}, + [ + "gs://bucket/prefix/part_1.txt", + "gs://bucket/prefix/part_2.txt", + "gs://bucket/prefix/summary.json", + "gs://other-bucket/out.csv", + ], + id="pattern-is-fullmatch-not-substring", + ), + ], +) +def test_get_hook_lineage_applies_exclusion_patterns(hook_lineage_collector, controls_kwargs, expected): + _collect_assets(hook_lineage_collector) + controls = replace(EmissionPolicy.defaults(), **controls_kwargs) + + result = ExtractorManager().get_hook_lineage( + task_instance=MagicMock(), + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + if expected: + assert _output_uris(result) == expected + else: + assert result is None + + +def test_get_hook_lineage_exclusion_applies_to_inputs(hook_lineage_collector): + hook = _FakeStorageHook() + hook_lineage_collector.add_input_asset(context=hook, uri="gs://bucket/staging/tmp.txt") + hook_lineage_collector.add_input_asset(context=hook, uri="gs://bucket/in.txt") + + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=("gs://bucket/staging/.*",)) + result = ExtractorManager().get_hook_lineage( + task_instance=MagicMock(), + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + assert result.inputs == [OpenLineageDataset(namespace="gs://bucket", name="in.txt")] + + +def test_get_hook_lineage_hook_pattern_distinguishes_same_class_name(hook_lineage_collector): + """Two hooks with the same class name in different modules must be separately targetable.""" + other_module = types.ModuleType("other_provider.hooks") + twin = type("_FakeStorageHook", (), {"__module__": other_module.__name__}) + + hook_lineage_collector.add_output_asset(context=_FakeStorageHook(), uri="gs://bucket/from_local.txt") + hook_lineage_collector.add_output_asset(context=twin(), uri="gs://bucket/from_other.txt") + + controls = replace( + EmissionPolicy.defaults(), + exclude_hook_lineage_hooks=(f"{other_module.__name__}._FakeStorageHook",), + ) + result = ExtractorManager().get_hook_lineage( + task_instance=MagicMock(), + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + assert _output_uris(result) == ["gs://bucket/from_local.txt"] + + +def test_extract_metadata_falls_back_to_manual_outlets_when_all_assets_excluded(hook_lineage_collector): + """Excluding every hook-collected asset must not suppress manually declared outlets.""" + _collect_assets(hook_lineage_collector) + outlets = [Asset(uri="gs://bucket/curated/table.parquet", extra={})] + task = PythonOperator(task_id="task_id", python_callable=lambda: None, outlets=outlets) + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=(".*",)) + + result = ExtractorManager().extract_metadata( + dagrun=MagicMock(), + task=task, + task_instance_state=TaskInstanceState.SUCCESS, + task_instance=MagicMock(), + controls=controls, + ) + + assert _output_uris(result) == ["gs://bucket/curated/table.parquet"] + + +def test_get_hook_lineage_does_not_filter_sql_extras(hook_lineage_collector): + """Exclusion patterns target assets only; SQL-based lineage is unaffected.""" + hook_lineage_collector.add_input_asset(context=_FakeStorageHook(), uri="gs://bucket/in.txt") + hook_lineage_collector.add_extra( + context=MagicMock(), + key=SqlJobHookLineageExtra.KEY.value, + value={SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value: "SELECT 1"}, + ) + + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=(".*",)) + mock_ti = MagicMock() + with patch(_SQL_FN_PATH, return_value=None) as mock_sql_fn: + result = ExtractorManager().get_hook_lineage( + task_instance=mock_ti, + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + assert result is None + mock_sql_fn.assert_called_once_with(task_instance=mock_ti, sql_extras=mock.ANY, is_successful=True) + + +def test_extract_metadata_passes_controls_to_hook_lineage(hook_lineage_collector): + """The exclusion patterns resolved for a task reach the collector read path.""" + _collect_assets(hook_lineage_collector) + task = PythonOperator(task_id="task_id", python_callable=lambda: None) + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=("gs://bucket/prefix/part_.*",)) + + result = ExtractorManager().extract_metadata( + dagrun=MagicMock(), + task=task, + task_instance_state=TaskInstanceState.SUCCESS, + task_instance=MagicMock(), + controls=controls, + ) + + assert _output_uris(result) == ["gs://bucket/prefix/summary.json", "gs://other-bucket/out.csv"] + + def test_extract_inlets_and_outlets_converts_asset_inlet_outlet(): """An Airflow ``Asset`` entry goes through ``translate_airflow_asset`` into an OL Dataset.""" extractor_manager = ExtractorManager() 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 ee501d4a2b87b..53877cbe7a922 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_emission_policy.py @@ -2412,3 +2412,154 @@ def test_locked_global_emit_false_overrides_selective_enable_opt_in(self): # 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 + + +class TestHookLineageExclusionPatterns: + """Validation and tier resolution for the ``exclude_hook_lineage_*`` pattern lists.""" + + def test_defaults_are_empty(self): + cfg = EmissionPolicy.defaults() + assert cfg.exclude_hook_lineage_assets == () + assert cfg.exclude_hook_lineage_hooks == () + + @pytest.mark.parametrize( + "control", + ["exclude_hook_lineage_assets", "exclude_hook_lineage_hooks"], + ) + def test_patterns_resolved_as_tuple(self, control): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {control: ["a.*", "b.*"]}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert getattr(cfg, control) == ("a.*", "b.*") + + @pytest.mark.parametrize( + "value", + [ + pytest.param("s3://bucket/.*", id="string-instead-of-list"), + pytest.param([1], id="non-string-element"), + pytest.param(["["], id="invalid-regex"), + pytest.param(True, id="bool-instead-of-list"), + ], + ) + def test_invalid_pattern_list_skips_rule(self, value): + """An invalid exclude list drops the whole rule, leaving defaults in place.""" + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"exclude_hook_lineage_assets": value, "emit": False}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_assets == () + assert cfg.emit is True + + def test_empty_list_is_valid(self): + cfg = _resolve_task_controls( + [{"scope": {}, "controls": {"exclude_hook_lineage_assets": []}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_assets == () + + def test_patterns_validated_as_regex_under_exact_match_mode(self): + """Exclude patterns are always regex, so a bad one is rejected even in exact mode.""" + cfg = _resolve_task_controls( + [{"scope": {}, "match_mode": "exact", "controls": {"exclude_hook_lineage_assets": ["("]}}], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_assets == () + + def test_more_specific_tier_replaces_broader_list(self): + """Lists replace rather than merge, so a task rule fully overrides a global one.""" + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"exclude_hook_lineage_assets": ["s3://.*"]}}, + { + "scope": {"dag_id": "my_dag", "task_id": "my_task"}, + "controls": {"exclude_hook_lineage_assets": ["s3://bucket/part_.*"]}, + }, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_assets == ("s3://bucket/part_.*",) + + def test_task_rule_can_clear_a_global_exclusion(self): + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"exclude_hook_lineage_assets": ["s3://.*"]}}, + { + "scope": {"dag_id": "my_dag", "task_id": "my_task"}, + "controls": {"exclude_hook_lineage_assets": []}, + }, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_assets == () + + def test_last_rule_wins_within_a_tier(self): + cfg = _resolve_task_controls( + [ + {"scope": {}, "controls": {"exclude_hook_lineage_hooks": ["FirstHook"]}}, + {"scope": {}, "controls": {"exclude_hook_lineage_hooks": ["SecondHook"]}}, + ], + operator=MockOperator(), + dag_id="my_dag", + task_id="my_task", + ) + assert cfg.exclude_hook_lineage_hooks == ("SecondHook",) + + def test_authoring_api_sets_patterns(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_hook_lineage_assets=["s3://.*/part_.*"]) + + with conf_vars({("openlineage", "emission_policy"): "[]"}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.exclude_hook_lineage_assets == ("s3://.*/part_.*",) + + def test_authoring_api_overrides_conf_patterns(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_hook_lineage_assets=["gs://.*"]) + rules = [{"scope": {}, "controls": {"exclude_hook_lineage_assets": ["s3://.*"]}}] + + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.exclude_hook_lineage_assets == ("gs://.*",) + + def test_locked_conf_rule_blocks_authoring_override(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_hook_lineage_assets=[]) + rules = [ + { + "scope": {}, + "locked": True, + "controls": {"exclude_hook_lineage_assets": ["s3://.*"]}, + } + ] + + with conf_vars({("openlineage", "emission_policy"): json.dumps(rules)}): + cfg = resolve_task_emission_policy(task, "test_dag", "test_task") + + assert cfg.exclude_hook_lineage_assets == ("s3://.*",)