From b433104fe0789bfaee305b186778a101218c8ec8 Mon Sep 17 00:00:00 2001 From: Ramit Kataria Date: Fri, 7 Aug 2026 11:32:58 -0700 Subject: [PATCH 1/2] Allow filtering hook-collected assets in OpenLineage emission policy Hook-level lineage registers an output asset for every object a hook writes, so a task that writes many objects reports one asset per object. The existing hook_lineage control is all-or-nothing: suppressing the noisy assets also loses the ones worth tracking from the same hook, and the collector cannot tell the two apart once its assets are read as one set. Filtering by asset URI or reporting hook class lets a Dag keep the lineage that matters while dropping the noise, and doing it where hook lineage is consumed covers every instrumented hook at once rather than one provider at a time. Note this reduces what gets reported rather than what gets collected, so it does not recover assets already dropped by max_assets_per_collector. closes: #63371 --- .../openlineage/docs/emission_policy.rst | 34 ++++ .../openlineage/api/emission_policy.py | 36 +++- .../openlineage/extractors/manager.py | 51 ++++- .../openlineage/utils/emission_policy.py | 170 ++++++++++++++-- .../openlineage/api/test_emission_policy.py | 40 ++++ .../openlineage/extractors/test_manager.py | 183 ++++++++++++++++++ .../openlineage/utils/test_emission_policy.py | 151 +++++++++++++++ 7 files changed, 640 insertions(+), 25 deletions(-) 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..2ad96bc701e54 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,186 @@ def test_get_hook_lineage_passes_failed_state(hook_lineage_collector): assert sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] == "SELECT 1" +class _FakeS3Hook: + pass + + +class _FakeGCSHook: + pass + + +def _hook_fqcn(hook_cls) -> str: + return f"{hook_cls.__module__}.{hook_cls.__name__}" + + +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. + """ + s3_hook, gcs_hook = _FakeS3Hook(), _FakeGCSHook() + for idx in range(3): + collector.add_output_asset(context=s3_hook, uri=f"s3://bucket/prefix/part_{idx}.txt") + collector.add_output_asset(context=s3_hook, uri="s3://bucket/prefix/summary.json") + collector.add_output_asset(context=gcs_hook, uri="gs://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/out.csv", + "s3://bucket/prefix/part_0.txt", + "s3://bucket/prefix/part_1.txt", + "s3://bucket/prefix/part_2.txt", + "s3://bucket/prefix/summary.json", + ], + id="no-filter-keeps-everything", + ), + pytest.param( + {"exclude_hook_lineage_assets": ("s3://bucket/prefix/part_.*",)}, + ["gs://bucket/out.csv", "s3://bucket/prefix/summary.json"], + id="asset-pattern-drops-parts-keeps-others", + ), + pytest.param( + {"exclude_hook_lineage_hooks": (_hook_fqcn(_FakeS3Hook),)}, + ["gs://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": ("s3://bucket/prefix/part_0.txt",)}, + [ + "gs://bucket/out.csv", + "s3://bucket/prefix/part_1.txt", + "s3://bucket/prefix/part_2.txt", + "s3://bucket/prefix/summary.json", + ], + 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 = _FakeS3Hook() + hook_lineage_collector.add_input_asset(context=hook, uri="s3://bucket/staging/tmp.txt") + hook_lineage_collector.add_input_asset(context=hook, uri="s3://bucket/in.txt") + + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=("s3://bucket/staging/.*",)) + result = ExtractorManager().get_hook_lineage( + task_instance=MagicMock(), + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + assert result.inputs == [OpenLineageDataset(namespace="s3://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("_FakeS3Hook", (), {"__module__": other_module.__name__}) + + hook_lineage_collector.add_output_asset(context=_FakeS3Hook(), uri="s3://bucket/from_local.txt") + hook_lineage_collector.add_output_asset(context=twin(), uri="s3://bucket/from_other.txt") + + controls = replace( + EmissionPolicy.defaults(), + exclude_hook_lineage_hooks=(f"{other_module.__name__}._FakeS3Hook",), + ) + result = ExtractorManager().get_hook_lineage( + task_instance=MagicMock(), + task_instance_state=TaskInstanceState.SUCCESS, + controls=controls, + ) + + assert _output_uris(result) == ["s3://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="s3://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) == ["s3://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=_FakeS3Hook(), uri="s3://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=("s3://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/out.csv", "s3://bucket/prefix/summary.json"] + + 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://.*",) From dac2744a6f99515306c80605e285aba5bb8689ad Mon Sep 17 00:00:00 2001 From: Ramit Kataria Date: Fri, 7 Aug 2026 14:35:01 -0700 Subject: [PATCH 2/2] Use gs scheme in hook lineage exclusion tests Translating an Airflow asset to an OpenLineage dataset needs the scheme's converter from the provider that owns it. The amazon provider is not a dev dependency of the openlineage provider, so s3 URIs do not translate in the lowest-dependencies CI job and every assertion about collected s3 assets fails there. --- .../openlineage/extractors/test_manager.py | 74 ++++++++++--------- 1 file changed, 39 insertions(+), 35 deletions(-) diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index 2ad96bc701e54..70cf78f332938 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -642,11 +642,11 @@ def test_get_hook_lineage_passes_failed_state(hook_lineage_collector): assert sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] == "SELECT 1" -class _FakeS3Hook: +class _FakeStorageHook: pass -class _FakeGCSHook: +class _FakeWarehouseHook: pass @@ -654,6 +654,10 @@ 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. @@ -661,11 +665,11 @@ def _collect_assets(collector): 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. """ - s3_hook, gcs_hook = _FakeS3Hook(), _FakeGCSHook() + storage_hook, warehouse_hook = _FakeStorageHook(), _FakeWarehouseHook() for idx in range(3): - collector.add_output_asset(context=s3_hook, uri=f"s3://bucket/prefix/part_{idx}.txt") - collector.add_output_asset(context=s3_hook, uri="s3://bucket/prefix/summary.json") - collector.add_output_asset(context=gcs_hook, uri="gs://bucket/out.csv") + 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): @@ -678,22 +682,22 @@ def _output_uris(lineage): pytest.param( {}, [ - "gs://bucket/out.csv", - "s3://bucket/prefix/part_0.txt", - "s3://bucket/prefix/part_1.txt", - "s3://bucket/prefix/part_2.txt", - "s3://bucket/prefix/summary.json", + "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": ("s3://bucket/prefix/part_.*",)}, - ["gs://bucket/out.csv", "s3://bucket/prefix/summary.json"], + {"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(_FakeS3Hook),)}, - ["gs://bucket/out.csv"], + {"exclude_hook_lineage_hooks": (_hook_fqcn(_FakeStorageHook),)}, + ["gs://other-bucket/out.csv"], id="hook-pattern-drops-all-assets-from-that-hook", ), pytest.param( @@ -702,12 +706,12 @@ def _output_uris(lineage): id="catch-all-asset-pattern-drops-everything", ), pytest.param( - {"exclude_hook_lineage_assets": ("s3://bucket/prefix/part_0.txt",)}, + {"exclude_hook_lineage_assets": ("gs://bucket/prefix/part_0.txt",)}, [ - "gs://bucket/out.csv", - "s3://bucket/prefix/part_1.txt", - "s3://bucket/prefix/part_2.txt", - "s3://bucket/prefix/summary.json", + "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", ), @@ -730,31 +734,31 @@ def test_get_hook_lineage_applies_exclusion_patterns(hook_lineage_collector, con def test_get_hook_lineage_exclusion_applies_to_inputs(hook_lineage_collector): - hook = _FakeS3Hook() - hook_lineage_collector.add_input_asset(context=hook, uri="s3://bucket/staging/tmp.txt") - hook_lineage_collector.add_input_asset(context=hook, uri="s3://bucket/in.txt") + 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=("s3://bucket/staging/.*",)) + 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="s3://bucket", name="in.txt")] + 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("_FakeS3Hook", (), {"__module__": other_module.__name__}) + twin = type("_FakeStorageHook", (), {"__module__": other_module.__name__}) - hook_lineage_collector.add_output_asset(context=_FakeS3Hook(), uri="s3://bucket/from_local.txt") - hook_lineage_collector.add_output_asset(context=twin(), uri="s3://bucket/from_other.txt") + 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__}._FakeS3Hook",), + exclude_hook_lineage_hooks=(f"{other_module.__name__}._FakeStorageHook",), ) result = ExtractorManager().get_hook_lineage( task_instance=MagicMock(), @@ -762,13 +766,13 @@ def test_get_hook_lineage_hook_pattern_distinguishes_same_class_name(hook_lineag controls=controls, ) - assert _output_uris(result) == ["s3://bucket/from_local.txt"] + 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="s3://bucket/curated/table.parquet", extra={})] + 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=(".*",)) @@ -780,12 +784,12 @@ def test_extract_metadata_falls_back_to_manual_outlets_when_all_assets_excluded( controls=controls, ) - assert _output_uris(result) == ["s3://bucket/curated/table.parquet"] + 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=_FakeS3Hook(), uri="s3://bucket/in.txt") + hook_lineage_collector.add_input_asset(context=_FakeStorageHook(), uri="gs://bucket/in.txt") hook_lineage_collector.add_extra( context=MagicMock(), key=SqlJobHookLineageExtra.KEY.value, @@ -809,7 +813,7 @@ 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=("s3://bucket/prefix/part_.*",)) + controls = replace(EmissionPolicy.defaults(), exclude_hook_lineage_assets=("gs://bucket/prefix/part_.*",)) result = ExtractorManager().extract_metadata( dagrun=MagicMock(), @@ -819,7 +823,7 @@ def test_extract_metadata_passes_controls_to_hook_lineage(hook_lineage_collector controls=controls, ) - assert _output_uris(result) == ["gs://bucket/out.csv", "s3://bucket/prefix/summary.json"] + assert _output_uris(result) == ["gs://bucket/prefix/summary.json", "gs://other-bucket/out.csv"] def test_extract_inlets_and_outlets_converts_asset_inlet_outlet():