From ba5db0be932b1e172f4cfacc6005965bc8261ca8 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 10:31:17 +0530 Subject: [PATCH 01/18] feat(ingestion): emit run-started progress update and thread total assets ingested Pre-existing in-flight branch work: emit a DISCOVERY progress update at run start so short runs are visible before the reporting timer's first tick, and forward totalAssetsIngested through the SSE payload and server tracker. Co-Authored-By: Claude Fable 5 --- ingestion/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ingestion/Dockerfile.ci b/ingestion/Dockerfile.ci index ec763025fc20..41af777ff170 100644 --- a/ingestion/Dockerfile.ci +++ b/ingestion/Dockerfile.ci @@ -127,7 +127,7 @@ RUN python /home/airflow/scripts/datamodel_generation.py # Argument to provide for Ingestion Dependencies to install. Defaults to all ARG INGESTION_DEPENDENCY="all" -RUN pip install ".[${INGESTION_DEPENDENCY}]" +RUN pip install ".[powerbi,snowflake,pii-processor]" # Temporary workaround for https://github.com/open-metadata/OpenMetadata/issues/9593 RUN [ $(uname -m) = "x86_64" ] \ From e1f7c175d9676add276fec77b12ce924769aa5e7 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 10:36:58 +0530 Subject: [PATCH 02/18] refactor(ingestion): move progress registry and tracking mixin into metadata/ingestion/progress package Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/api/topology_runner.py | 2 +- .../metadata/ingestion/progress/__init__.py | 39 +++++++++++++++++++ .../progress/registry.py} | 0 .../progress/tracking.py} | 2 +- .../source/database/query_parser_source.py | 2 +- .../src/metadata/workflow/progress_render.py | 2 +- .../database/test_snowflake_progress_count.py | 4 +- .../database/test_snowflake_progress_path.py | 2 +- .../tests/unit/test_progress_registry.py | 6 +-- .../unit/test_progress_tracking_mixin.py | 4 +- .../topology/test_topology_runner_progress.py | 2 +- .../unit/workflow/test_progress_rendering.py | 10 ++--- .../workflow/test_status_mixin_progress.py | 4 +- 13 files changed, 59 insertions(+), 20 deletions(-) create mode 100644 ingestion/src/metadata/ingestion/progress/__init__.py rename ingestion/src/metadata/{utils/progress_registry.py => ingestion/progress/registry.py} (100%) rename ingestion/src/metadata/{utils/progress_tracking.py => ingestion/progress/tracking.py} (95%) diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index 70e88df96adb..560e1e0fb424 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -46,7 +46,7 @@ from metadata.utils.custom_thread_pool import CustomThreadPoolExecutor from metadata.utils.logger import ingestion_logger from metadata.utils.operation_metrics import OperationMetricsState -from metadata.utils.progress_tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.utils.source_hash import generate_source_hash logger = ingestion_logger() diff --git a/ingestion/src/metadata/ingestion/progress/__init__.py b/ingestion/src/metadata/ingestion/progress/__init__.py new file mode 100644 index 000000000000..7908288ecdc9 --- /dev/null +++ b/ingestion/src/metadata/ingestion/progress/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ingestion progress capturing — the single home for all progress code. + +Ownership contract: + +* The **connector** knows how to query the source and how to calculate totals + (denominators). It never counts processed items. +* The **topology runner** counts processed entities for every connector by + default. +* The **ProgressRegistry** stores the state, calculates the ETA, and reports + it (CLI tree + SSE payload). +""" + +from metadata.ingestion.progress.registry import ( + DEFAULT_ACTIVE_LEAF_CAP, + GlobalCounter, + ProgressNode, + ProgressNodeSnapshot, + ProgressRegistry, +) +from metadata.ingestion.progress.tracking import ProgressTrackingMixin + +__all__ = [ + "DEFAULT_ACTIVE_LEAF_CAP", + "GlobalCounter", + "ProgressNode", + "ProgressNodeSnapshot", + "ProgressRegistry", + "ProgressTrackingMixin", +] diff --git a/ingestion/src/metadata/utils/progress_registry.py b/ingestion/src/metadata/ingestion/progress/registry.py similarity index 100% rename from ingestion/src/metadata/utils/progress_registry.py rename to ingestion/src/metadata/ingestion/progress/registry.py diff --git a/ingestion/src/metadata/utils/progress_tracking.py b/ingestion/src/metadata/ingestion/progress/tracking.py similarity index 95% rename from ingestion/src/metadata/utils/progress_tracking.py rename to ingestion/src/metadata/ingestion/progress/tracking.py index 4bd69a6a394c..4ab12a2c19be 100644 --- a/ingestion/src/metadata/utils/progress_tracking.py +++ b/ingestion/src/metadata/ingestion/progress/tracking.py @@ -15,7 +15,7 @@ without participating in the topology runner. """ -from metadata.utils.progress_registry import ProgressRegistry +from metadata.ingestion.progress.registry import ProgressRegistry class ProgressTrackingMixin: diff --git a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py index 599ad4407e9d..def95ff327c3 100644 --- a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py +++ b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py @@ -30,7 +30,7 @@ from metadata.ingestion.source.connections import test_connection_common from metadata.utils.helpers import get_start_and_end, retry_with_docker_host from metadata.utils.logger import ingestion_logger -from metadata.utils.progress_tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.utils.ssl_manager import get_ssl_connection logger = ingestion_logger() diff --git a/ingestion/src/metadata/workflow/progress_render.py b/ingestion/src/metadata/workflow/progress_render.py index dcbb4f4ce56d..077a322ccb80 100644 --- a/ingestion/src/metadata/workflow/progress_render.py +++ b/ingestion/src/metadata/workflow/progress_render.py @@ -18,7 +18,7 @@ from typing import List, Optional, Tuple # noqa: UP035 -from metadata.utils.progress_registry import ProgressNodeSnapshot, ProgressRegistry +from metadata.ingestion.progress.registry import ProgressNodeSnapshot, ProgressRegistry def render_progress_tree(root: Optional[ProgressNodeSnapshot]) -> str: # noqa: UP045 diff --git a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py index 321dcaa1f573..1b8fbf7fb0de 100644 --- a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py +++ b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py @@ -61,7 +61,7 @@ def test_filtered_names_are_cached_so_status_is_not_double_emitted(self): source.status.filter.assert_called_once() def test_producer_setup_runs_per_kept_database_without_re_filtering(self): - from metadata.utils.progress_registry import ProgressRegistry + from metadata.ingestion.progress.registry import ProgressRegistry source = _source(["A", "B"]) for setup in ( @@ -87,7 +87,7 @@ def test_producer_setup_runs_per_kept_database_without_re_filtering(self): @pytest.fixture def snowflake_source(): """Minimal SnowflakeSource stub for push-totals tests.""" - from metadata.utils.progress_registry import ProgressRegistry + from metadata.ingestion.progress.registry import ProgressRegistry source = object.__new__(SnowflakeSource) source.config = SimpleNamespace( diff --git a/ingestion/tests/unit/source/database/test_snowflake_progress_path.py b/ingestion/tests/unit/source/database/test_snowflake_progress_path.py index db0aba64b405..c9323d00ae7f 100644 --- a/ingestion/tests/unit/source/database/test_snowflake_progress_path.py +++ b/ingestion/tests/unit/source/database/test_snowflake_progress_path.py @@ -16,7 +16,7 @@ from metadata.ingestion.source.database import database_service from metadata.ingestion.source.database.snowflake import metadata as snowflake_metadata -from metadata.utils.progress_registry import ProgressRegistry +from metadata.ingestion.progress.registry import ProgressRegistry class _ConcreteSource(database_service.DatabaseServiceSource): diff --git a/ingestion/tests/unit/test_progress_registry.py b/ingestion/tests/unit/test_progress_registry.py index fbec7ca92cc3..726d055f680b 100644 --- a/ingestion/tests/unit/test_progress_registry.py +++ b/ingestion/tests/unit/test_progress_registry.py @@ -13,7 +13,7 @@ import threading from unittest.mock import patch -from metadata.utils.progress_registry import ProgressRegistry +from metadata.ingestion.progress.registry import ProgressRegistry def _by_label(snapshot): @@ -353,7 +353,7 @@ def test_none_before_first_open(self): def test_positive_after_first_open(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 100.0 reg.open([], "Database", None) clock.return_value = 130.0 @@ -361,7 +361,7 @@ def test_positive_after_first_open(self): def test_marker_set_on_first_open_only(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 100.0 reg.open([], "Database", None) clock.return_value = 200.0 diff --git a/ingestion/tests/unit/test_progress_tracking_mixin.py b/ingestion/tests/unit/test_progress_tracking_mixin.py index 5b68a5bc9701..38d4018efbe6 100644 --- a/ingestion/tests/unit/test_progress_tracking_mixin.py +++ b/ingestion/tests/unit/test_progress_tracking_mixin.py @@ -10,8 +10,8 @@ # limitations under the License. """Unit tests for the shared ProgressTrackingMixin.""" -from metadata.utils.progress_registry import ProgressRegistry -from metadata.utils.progress_tracking import ProgressTrackingMixin +from metadata.ingestion.progress.registry import ProgressRegistry +from metadata.ingestion.progress.tracking import ProgressTrackingMixin class _Owner(ProgressTrackingMixin): diff --git a/ingestion/tests/unit/topology/test_topology_runner_progress.py b/ingestion/tests/unit/topology/test_topology_runner_progress.py index ad1377e0c4ce..112c47940a54 100644 --- a/ingestion/tests/unit/topology/test_topology_runner_progress.py +++ b/ingestion/tests/unit/topology/test_topology_runner_progress.py @@ -24,7 +24,7 @@ from metadata.ingestion.source.database.database_service import ( DatabaseServiceTopology, ) -from metadata.utils.progress_registry import ProgressRegistry +from metadata.ingestion.progress.registry import ProgressRegistry class TestRunnerProgressSurface: diff --git a/ingestion/tests/unit/workflow/test_progress_rendering.py b/ingestion/tests/unit/workflow/test_progress_rendering.py index 0333f831387f..8381b746a40f 100644 --- a/ingestion/tests/unit/workflow/test_progress_rendering.py +++ b/ingestion/tests/unit/workflow/test_progress_rendering.py @@ -16,7 +16,7 @@ ProgressNode, ProgressUpdate, ) -from metadata.utils.progress_registry import ProgressNodeSnapshot, ProgressRegistry +from metadata.ingestion.progress.registry import ProgressNodeSnapshot, ProgressRegistry from metadata.workflow.progress_render import ( ProgressReporter, format_eta, @@ -270,7 +270,7 @@ def test_complete_is_none(self): def test_computable_without_open_when_counter_active(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 0.0 reg.set_total("DatabaseSchema", 10) for _ in range(2): @@ -280,7 +280,7 @@ def test_computable_without_open_when_counter_active(self): def test_normal_case_uses_cumulative_rate(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 0.0 reg.open([], "Database", None) reg.set_total("DatabaseSchema", 10) @@ -291,7 +291,7 @@ def test_normal_case_uses_cumulative_rate(self): def test_driver_is_last_counter_with_total(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 0.0 reg.open([], "Database", None) reg.set_total("Database", 4) @@ -306,7 +306,7 @@ def test_driver_is_last_counter_with_total(self): class TestEtaInHeader: def test_eta_suffix_on_driver_line_only(self): reg = ProgressRegistry() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 0.0 reg.open([], "Database", None) reg.set_total("Database", 4) diff --git a/ingestion/tests/unit/workflow/test_status_mixin_progress.py b/ingestion/tests/unit/workflow/test_status_mixin_progress.py index 146a62d7ee75..0ea30bab1638 100644 --- a/ingestion/tests/unit/workflow/test_status_mixin_progress.py +++ b/ingestion/tests/unit/workflow/test_status_mixin_progress.py @@ -16,7 +16,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.progressUpdate import ( ProgressUpdateType, ) -from metadata.utils.progress_registry import ProgressRegistry +from metadata.ingestion.progress.registry import ProgressRegistry from metadata.workflow.workflow_status_mixin import WorkflowStatusMixin @@ -46,7 +46,7 @@ class TestEstimatedSecondsRemaining: def test_maps_eta_seconds_into_payload(self): reg = ProgressRegistry() metadata = MagicMock() - with patch("metadata.utils.progress_registry.time.monotonic") as clock: + with patch("metadata.ingestion.progress.registry.time.monotonic") as clock: clock.return_value = 0.0 reg.open([], "Database", None) reg.set_total("DatabaseSchema", 10) From 1a3e5db1429eb837b52dbbb6ed787298cf6c955e Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 10:46:03 +0530 Subject: [PATCH 03/18] feat(ingestion): add ProgressMode knob with TotalsDeclarer and ManualProgress facades Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/progress/__init__.py | 10 ++ .../src/metadata/ingestion/progress/modes.py | 98 ++++++++++++++++ .../metadata/ingestion/progress/tracking.py | 34 ++++++ ingestion/tests/unit/progress/__init__.py | 0 .../unit/progress/test_progress_modes.py | 109 ++++++++++++++++++ 5 files changed, 251 insertions(+) create mode 100644 ingestion/src/metadata/ingestion/progress/modes.py create mode 100644 ingestion/tests/unit/progress/__init__.py create mode 100644 ingestion/tests/unit/progress/test_progress_modes.py diff --git a/ingestion/src/metadata/ingestion/progress/__init__.py b/ingestion/src/metadata/ingestion/progress/__init__.py index 7908288ecdc9..056c9a46bd7d 100644 --- a/ingestion/src/metadata/ingestion/progress/__init__.py +++ b/ingestion/src/metadata/ingestion/progress/__init__.py @@ -20,6 +20,12 @@ it (CLI tree + SSE payload). """ +from metadata.ingestion.progress.modes import ( + ManualProgress, + ProgressMode, + ProgressModeError, + TotalsDeclarer, +) from metadata.ingestion.progress.registry import ( DEFAULT_ACTIVE_LEAF_CAP, GlobalCounter, @@ -32,8 +38,12 @@ __all__ = [ "DEFAULT_ACTIVE_LEAF_CAP", "GlobalCounter", + "ManualProgress", + "ProgressMode", + "ProgressModeError", "ProgressNode", "ProgressNodeSnapshot", "ProgressRegistry", "ProgressTrackingMixin", + "TotalsDeclarer", ] diff --git a/ingestion/src/metadata/ingestion/progress/modes.py b/ingestion/src/metadata/ingestion/progress/modes.py new file mode 100644 index 000000000000..29cb1c98d053 --- /dev/null +++ b/ingestion/src/metadata/ingestion/progress/modes.py @@ -0,0 +1,98 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Progress modes and the two connector-facing facades. + +``TotalsDeclarer`` is denominator-only: an AUTO connector's +``declare_progress_totals`` hook receives it and structurally cannot increment +processed counts. ``ManualProgress`` is the MANUAL-mode escape hatch: the +runner makes zero progress calls and the connector drives counting itself — +auto XOR manual, never both. +""" + +from enum import Enum +from typing import TYPE_CHECKING, Dict, Optional # noqa: UP035 + +if TYPE_CHECKING: + from metadata.ingestion.progress.registry import ProgressRegistry + + +class ProgressMode(Enum): + AUTO = "auto" + MANUAL = "manual" + OFF = "off" + + +class ProgressModeError(RuntimeError): + """A source used a progress API its declared ``progress_mode`` forbids.""" + + +class TotalsDeclarer: + """Denominator-only view of the registry, handed to + ``declare_progress_totals``. Declarations feed % and ETA; they never + increment ``processed``.""" + + def __init__(self, registry: "ProgressRegistry") -> None: + self._registry = registry + + def set_total(self, entity_type: str, total: Optional[int]) -> None: # noqa: UP045 + self._registry.set_total(entity_type, total) + + def seed_scope_total(self, entity_type: str, scope: str, n: int) -> None: + self._registry.seed_scope_total(entity_type, scope, n) + + def mark_reconcilable(self, entity_type: str) -> None: + """Declare that the connector cannot pre-count ``entity_type``; the + runner will build the total from observed scope counts instead.""" + self._registry.set_reconcilable(entity_type) + + +class ManualProgress: + """Counting facade for ``ProgressMode.MANUAL`` sources — connectors whose + enumeration the topology runner cannot count (grouped dashboard walks, + custom ``_iter`` query sources).""" + + def __init__(self, registry: "ProgressRegistry") -> None: + self._registry = registry + self._group_label: Optional[str] = None # noqa: UP045 + + def set_total(self, entity_type: str, total: Optional[int]) -> None: # noqa: UP045 + self._registry.set_total(entity_type, total) + + def track(self, entity_type: Optional[str], n: int = 1) -> None: # noqa: UP045 + self._registry.track(entity_type, n) + + def seed_scope_total(self, entity_type: str, scope: str, n: int) -> None: + self._registry.seed_scope_total(entity_type, scope, n) + + def reconcile_scope_total(self, entity_type: Optional[str], scope: str, observed: int) -> None: # noqa: UP045 + self._registry.reconcile_scope_total(entity_type, scope, observed) + + def declare_groups(self, label: str, total: Optional[int]) -> None: # noqa: UP045 + """Declare the grouping axis (e.g. workspaces) as a global counter and + remember its label so ``close_group`` can count completions on it.""" + self._group_label = label + self._registry.set_total(label, total) + self._registry.open([], label, total) + + def open_group(self, group: str, expected_by_type: Dict[str, Optional[int]]) -> None: # noqa: UP006,UP045 + """Open one child node per asset type under ``group``; ``expected`` may + be None for lazy (running) counts.""" + for asset_type, expected in expected_by_type.items(): + self._registry.open([group, asset_type], asset_type, expected) + + def advance(self, group: str, asset_type: str) -> None: + self._registry.advance([group, asset_type], asset_type) + + def close_group(self, group: str) -> None: + """Count the finished group on its global counter and prune its subtree.""" + if self._group_label is not None: + self._registry.track(self._group_label) + self._registry.close([group]) diff --git a/ingestion/src/metadata/ingestion/progress/tracking.py b/ingestion/src/metadata/ingestion/progress/tracking.py index 4ab12a2c19be..1b90ae4391f6 100644 --- a/ingestion/src/metadata/ingestion/progress/tracking.py +++ b/ingestion/src/metadata/ingestion/progress/tracking.py @@ -15,10 +15,23 @@ without participating in the topology runner. """ +from typing import ClassVar + +from metadata.ingestion.progress.modes import ( + ManualProgress, + ProgressMode, + ProgressModeError, + TotalsDeclarer, +) from metadata.ingestion.progress.registry import ProgressRegistry class ProgressTrackingMixin: + progress_mode: ClassVar[ProgressMode] = ProgressMode.AUTO + """AUTO (default): the topology runner counts processed entities. + MANUAL: the runner makes zero progress calls; the source drives + ``manual_progress`` itself. OFF: no progress at all.""" + @property def progress(self) -> ProgressRegistry: """Per-Source progress registry. First access is single-threaded @@ -28,3 +41,24 @@ def progress(self) -> ProgressRegistry: registry = ProgressRegistry() self.__dict__["_progress_registry"] = registry return registry + + @property + def manual_progress(self) -> ManualProgress: + """MANUAL-mode counting facade. Raises for AUTO/OFF sources — the + runner counts AUTO sources, so a manual increment would double-count.""" + if self.progress_mode is not ProgressMode.MANUAL: + raise ProgressModeError( + f"manual_progress requires progress_mode=MANUAL, but {type(self).__name__} " + f"declares {self.progress_mode.name}. AUTO sources are counted by the topology " + "runner and may only declare totals via declare_progress_totals()." + ) + facade = self.__dict__.get("_manual_progress") + if facade is None: + facade = ManualProgress(self.progress) + self.__dict__["_manual_progress"] = facade + return facade + + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: + """Optional connector hook: declare denominators (totals) so the run + renders % and ETA. Called exactly once by the topology runner, just + before the first non-root node is processed. Default: no totals.""" diff --git a/ingestion/tests/unit/progress/__init__.py b/ingestion/tests/unit/progress/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/ingestion/tests/unit/progress/test_progress_modes.py b/ingestion/tests/unit/progress/test_progress_modes.py new file mode 100644 index 000000000000..bf3b6fe5d691 --- /dev/null +++ b/ingestion/tests/unit/progress/test_progress_modes.py @@ -0,0 +1,109 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ProgressMode knob, TotalsDeclarer and ManualProgress facades.""" + +import pytest + +from metadata.ingestion.progress.modes import ( + ManualProgress, + ProgressMode, + ProgressModeError, + TotalsDeclarer, +) +from metadata.ingestion.progress.registry import ProgressRegistry +from metadata.ingestion.progress.tracking import ProgressTrackingMixin + + +class _AutoSource(ProgressTrackingMixin): + pass + + +class _ManualSource(ProgressTrackingMixin): + progress_mode = ProgressMode.MANUAL + + +class _OffSource(ProgressTrackingMixin): + progress_mode = ProgressMode.OFF + + +def test_default_mode_is_auto(): + assert _AutoSource().progress_mode is ProgressMode.AUTO + + +def test_manual_progress_raises_in_auto_mode(): + with pytest.raises(ProgressModeError): + _ = _AutoSource().manual_progress + + +def test_manual_progress_raises_in_off_mode(): + with pytest.raises(ProgressModeError): + _ = _OffSource().manual_progress + + +def test_manual_progress_is_cached_per_instance(): + source = _ManualSource() + assert source.manual_progress is source.manual_progress + + +def test_manual_progress_wraps_the_source_registry(): + source = _ManualSource() + source.manual_progress.set_total("Workspaces", 3) + assert ("Workspaces", 0, 3) in source.progress.global_counters() + + +def test_totals_declarer_exposes_no_counting_methods(): + declarer = TotalsDeclarer(ProgressRegistry()) + for name in ("track", "advance", "open", "close", "reconcile_scope_total"): + assert not hasattr(declarer, name) + + +def test_totals_declarer_sets_denominators(): + registry = ProgressRegistry() + totals = TotalsDeclarer(registry) + totals.set_total("Database", 4) + totals.seed_scope_total("DatabaseSchema", "db1", 3) + totals.mark_reconcilable("Table") + counters = {t: (done, total) for t, done, total in registry.global_counters()} + assert counters["Database"] == (0, 4) + assert counters["DatabaseSchema"] == (0, 3) + assert registry.is_reconcilable("DatabaseSchema") + assert registry.is_reconcilable("Table") + + +def test_declare_progress_totals_default_is_noop(): + registry = ProgressRegistry() + _AutoSource().declare_progress_totals(TotalsDeclarer(registry)) + assert registry.global_counters() == [] + + +def test_manual_group_flow_drives_registry_like_the_old_helpers(): + registry = ProgressRegistry() + manual = ManualProgress(registry) + manual.declare_groups("Workspaces", 2) + manual.open_group("Finance", {"Dashboard": 2, "Chart": None}) + manual.advance("Finance", "Dashboard") + manual.advance("Finance", "Dashboard") + manual.advance("Finance", "Chart") + manual.close_group("Finance") + counters = {t: (done, total) for t, done, total in registry.global_counters()} + assert counters["Workspaces"] == (1, 2) + assert registry.assets_ingested() == 3 + assert registry.snapshot() is not None + + +def test_manual_counter_api_passthrough(): + registry = ProgressRegistry() + manual = ManualProgress(registry) + manual.seed_scope_total("Queries", "run", 100) + manual.track("Queries", 40) + manual.reconcile_scope_total("Queries", "run", 40) + counters = {t: (done, total) for t, done, total in registry.global_counters()} + assert counters["Queries"] == (40, 40) From 6951990088d4a35976e5d3f9a5f007b16977c831 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 10:48:51 +0530 Subject: [PATCH 04/18] fix(ingestion): declare_groups only sets the global total, not the tree root The extra open([], label, total) made the active-scope snapshot non-None after all groups are pruned, contradicting the established snapshot()-is-None behavior and reintroducing a parallel count path. Restore parity with the original _declare_progress_groups helper. Co-Authored-By: Claude Fable 5 --- ingestion/src/metadata/ingestion/progress/modes.py | 1 - ingestion/tests/unit/progress/test_progress_modes.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ingestion/src/metadata/ingestion/progress/modes.py b/ingestion/src/metadata/ingestion/progress/modes.py index 29cb1c98d053..d24a1301e68f 100644 --- a/ingestion/src/metadata/ingestion/progress/modes.py +++ b/ingestion/src/metadata/ingestion/progress/modes.py @@ -80,7 +80,6 @@ def declare_groups(self, label: str, total: Optional[int]) -> None: # noqa: UP0 remember its label so ``close_group`` can count completions on it.""" self._group_label = label self._registry.set_total(label, total) - self._registry.open([], label, total) def open_group(self, group: str, expected_by_type: Dict[str, Optional[int]]) -> None: # noqa: UP006,UP045 """Open one child node per asset type under ``group``; ``expected`` may diff --git a/ingestion/tests/unit/progress/test_progress_modes.py b/ingestion/tests/unit/progress/test_progress_modes.py index bf3b6fe5d691..7a303cb53f65 100644 --- a/ingestion/tests/unit/progress/test_progress_modes.py +++ b/ingestion/tests/unit/progress/test_progress_modes.py @@ -96,7 +96,7 @@ def test_manual_group_flow_drives_registry_like_the_old_helpers(): counters = {t: (done, total) for t, done, total in registry.global_counters()} assert counters["Workspaces"] == (1, 2) assert registry.assets_ingested() == 3 - assert registry.snapshot() is not None + assert registry.snapshot() is None def test_manual_counter_api_passthrough(): From 18d7542d1969fa0fe194f37a0b67f3890199eafa Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 10:59:14 +0530 Subject: [PATCH 05/18] refactor(ingestion): registry owns ETA and reporting; dissolve ProgressReporter Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/progress/__init__.py | 8 + .../metadata/ingestion/progress/_render.py | 92 +++++++++ .../metadata/ingestion/progress/registry.py | 75 ++++++++ ingestion/src/metadata/workflow/base.py | 6 +- .../src/metadata/workflow/progress_render.py | 179 ------------------ .../workflow/workflow_status_mixin.py | 31 ++- .../test_snowflake_access_history_progress.py | 3 +- .../unit/topology/dashboard/test_powerbi.py | 8 +- .../unit/workflow/test_progress_rendering.py | 57 +++--- .../workflow/test_status_mixin_progress.py | 2 +- 10 files changed, 224 insertions(+), 237 deletions(-) create mode 100644 ingestion/src/metadata/ingestion/progress/_render.py delete mode 100644 ingestion/src/metadata/workflow/progress_render.py diff --git a/ingestion/src/metadata/ingestion/progress/__init__.py b/ingestion/src/metadata/ingestion/progress/__init__.py index 056c9a46bd7d..b6cef80a65a2 100644 --- a/ingestion/src/metadata/ingestion/progress/__init__.py +++ b/ingestion/src/metadata/ingestion/progress/__init__.py @@ -20,6 +20,11 @@ it (CLI tree + SSE payload). """ +from metadata.ingestion.progress._render import ( + format_eta, + render_progress_tree, + snapshot_to_progress_payload, +) from metadata.ingestion.progress.modes import ( ManualProgress, ProgressMode, @@ -46,4 +51,7 @@ "ProgressRegistry", "ProgressTrackingMixin", "TotalsDeclarer", + "format_eta", + "render_progress_tree", + "snapshot_to_progress_payload", ] diff --git a/ingestion/src/metadata/ingestion/progress/_render.py b/ingestion/src/metadata/ingestion/progress/_render.py new file mode 100644 index 000000000000..fef77c273f6a --- /dev/null +++ b/ingestion/src/metadata/ingestion/progress/_render.py @@ -0,0 +1,92 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Render a ProgressRegistry snapshot tree for humans (CLI) and for the server +(nested progressNode payload). + +Dependency-free on purpose (stdlib only) so both the workflow base and the +status mixin can import it without a circular dependency. +""" + +from typing import TYPE_CHECKING, List, Optional # noqa: UP035 + +if TYPE_CHECKING: + from metadata.ingestion.progress.registry import ProgressNodeSnapshot + + +def render_progress_tree(root: "Optional[ProgressNodeSnapshot]") -> str: # noqa: UP045 + """Active-scope tree: the run rollup, then one indented line per in-flight + branch. Known counts render ``processed/expected``; lazy counts render a + bare ``processed``.""" + lines: List[str] = [] # noqa: UP006 + if root is not None: + _render_node(root, 0, lines) + return "\n".join(lines) + + +def _render_node(node: "ProgressNodeSnapshot", depth: int, lines: List[str]) -> None: # noqa: UP006 + prefix = " " * depth + ("└ " if depth else "") + label = f"{node.label} " if node.label else "" + type_ = f"{node.child_type} " if node.child_type else "" + if node.child_type is None: + count = "" + else: + count = f"{node.processed}/{node.expected}" if node.expected is not None else str(node.processed) + overflow = f" (+{node.overflow} more)" if node.overflow else "" + lines.append(f"{prefix}{label}{type_}{count}{overflow}".rstrip()) + for child in node.children: + _render_node(child, depth + 1, lines) + + +def snapshot_to_progress_payload(root: "Optional[ProgressNodeSnapshot]") -> Optional[dict]: # noqa: UP045 + """Map a snapshot tree to the ProgressUpdate.progress (progressNode) shape.""" + return _node_to_payload(root) if root is not None else None + + +def format_eta(seconds: int) -> str: + """Human ``~45s`` / ``~6m`` / ``~1h 20m`` from a rounded seconds count.""" + if seconds < 60: + result = f"~{seconds}s" + elif seconds < 3600: + result = f"~{seconds // 60}m" + else: + result = f"~{seconds // 3600}h {(seconds % 3600) // 60}m" + return result + + +def _render_joined(node: "ProgressNodeSnapshot", ancestors: "List[str]", lines: "List[str]") -> None: # noqa: UP006 + """Render active nodes with ancestor labels joined by '.' so sibling-level + hierarchy collapses into a single readable label, e.g. ``sales.public Table 45/310``.""" + path = [*ancestors, node.label] if node.label else list(ancestors) + if node.children: + for child in node.children: + _render_joined(child, path, lines) + else: + joined_label = ".".join(path) if path else "" + type_ = f"{node.child_type} " if node.child_type else "" + if node.child_type is None: + count = "" + else: + count = f"{node.processed}/{node.expected}" if node.expected is not None else str(node.processed) + overflow = f" (+{node.overflow} more)" if node.overflow else "" + lines.append(f"{joined_label} {type_}{count}{overflow}".strip()) + + +def _node_to_payload(node: "ProgressNodeSnapshot") -> dict: + return { + "label": node.label, + "entityType": node.child_type, + "processed": node.processed, + "expected": node.expected, + "active": node.active, + "overflow": node.overflow, + "children": [_node_to_payload(child) for child in node.children], + } diff --git a/ingestion/src/metadata/ingestion/progress/registry.py b/ingestion/src/metadata/ingestion/progress/registry.py index 96cbbc9086ac..75d796357e30 100644 --- a/ingestion/src/metadata/ingestion/progress/registry.py +++ b/ingestion/src/metadata/ingestion/progress/registry.py @@ -31,6 +31,12 @@ from dataclasses import dataclass, field from typing import Dict, List, Mapping, Optional, Tuple # noqa: UP035 +from metadata.ingestion.progress._render import ( + _render_joined, + format_eta, + snapshot_to_progress_payload, +) + DEFAULT_ACTIVE_LEAF_CAP = 20 @@ -186,6 +192,75 @@ def global_counters(self) -> "List[Tuple[str, int, Optional[int]]]": # noqa: UP with self._lock: return [(type_, c.done, c.total) for type_, c in self._global.items()] + def eta_seconds(self) -> Optional[int]: # noqa: UP045 + """Overall run ETA in seconds from the driver counter's cumulative rate: + ``elapsed * (total - done) / done``. ``None`` during warm-up (``done == + 0``), when there is no driver, when the driver is complete (``done >= + total``), or before elapsed time is available.""" + driver = self._driver_counter() + result = None + if driver is not None: + _, done, total = driver + elapsed = self.elapsed_seconds() + if done > 0 and total is not None and done < total and elapsed is not None and elapsed > 0: + result = round(elapsed * (total - done) / done) + return result + + def _driver_counter( + self, + counters: "Optional[List[Tuple[str, int, Optional[int]]]]" = None, # noqa: UP006,UP045 + ) -> "Optional[Tuple[str, int, Optional[int]]]": # noqa: UP006,UP045 + """The ETA driver: the last-declared global counter that has a known + total (finest-grained real unit of work), or ``None`` when no counter + declares a total.""" + pool = self.global_counters() if counters is None else counters + driver = None + for counter in pool: + if counter[2] is not None: + driver = counter + return driver + + def render_cli(self) -> str: + """Human progress report: global-counter header (with ETA on the driver + line) plus the active-scope tree.""" + snapshot = self.snapshot() + counters = self.global_counters() + header = self._render_header(counters) + if snapshot is None: + return header if counters else "" + lines: List[str] = [] # noqa: UP006 + _render_joined(snapshot, [], lines) + tree = "\n".join(lines) + return f"{header}\n{tree}" if tree else header + + def _render_header(self, counters: "List[Tuple[str, int, Optional[int]]]") -> str: # noqa: UP006,UP045 + lines: List[str] = [] # noqa: UP006 + driver = self._driver_counter(counters) + eta = self.eta_seconds() + for type_, done, total in counters: + line = f"{type_} {done}/{total}" if total is not None else f"{type_} {done}" + if eta is not None and driver is not None and type_ == driver[0]: + line = f"{line} {format_eta(eta)}" + lines.append(line) + assets = self.assets_ingested() + if assets > 0: + lines.append(f"Ingested: {assets:,} assets") + return "\n".join(lines) + + def sse_payload(self) -> Optional[dict]: # noqa: UP045 + """Bare ``progressNode`` tree (validates against ``ProgressUpdate``, + which is ``extra='forbid'``). The run-root ``processed`` carries the + monotonic asset total; ``expected`` is unknown (no global count).""" + snapshot = self.snapshot() + result = None + if snapshot is not None: + tree = snapshot_to_progress_payload(snapshot) + if tree is not None: + tree["processed"] = self.assets_ingested() + tree["expected"] = None + result = tree + return result + def close(self, path: List[str]) -> None: # noqa: UP006 """Remove the node at ``path`` from its parent's children once its work is complete, so the tree retains only active scopes. No-op on an empty diff --git a/ingestion/src/metadata/workflow/base.py b/ingestion/src/metadata/workflow/base.py index 08a1004d111e..a2ccd03de673 100644 --- a/ingestion/src/metadata/workflow/base.py +++ b/ingestion/src/metadata/workflow/base.py @@ -461,9 +461,9 @@ def _report_ingestion_status(self): f"Processes: {metrics.active_processes}" ) - reporter = self._progress_reporter() - if reporter is not None: - text = reporter.cli() + registry = self._find_progress_registry() + if registry is not None: + text = registry.render_cli() if text: logger.info("Ingestion progress:\n%s", text) diff --git a/ingestion/src/metadata/workflow/progress_render.py b/ingestion/src/metadata/workflow/progress_render.py deleted file mode 100644 index 077a322ccb80..000000000000 --- a/ingestion/src/metadata/workflow/progress_render.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright 2025 Collate -# Licensed under the Collate Community License, Version 1.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Render a ProgressRegistry snapshot tree for humans (CLI) and for the server -(nested progressNode payload). - -Dependency-free on purpose (stdlib only) so both the workflow base and the -status mixin can import it without a circular dependency. -""" - -from typing import List, Optional, Tuple # noqa: UP035 - -from metadata.ingestion.progress.registry import ProgressNodeSnapshot, ProgressRegistry - - -def render_progress_tree(root: Optional[ProgressNodeSnapshot]) -> str: # noqa: UP045 - """Active-scope tree: the run rollup, then one indented line per in-flight - branch. Known counts render ``processed/expected``; lazy counts render a - bare ``processed``.""" - lines: List[str] = [] # noqa: UP006 - if root is not None: - _render_node(root, 0, lines) - return "\n".join(lines) - - -def _render_node(node: ProgressNodeSnapshot, depth: int, lines: List[str]) -> None: # noqa: UP006 - prefix = " " * depth + ("└ " if depth else "") - label = f"{node.label} " if node.label else "" - type_ = f"{node.child_type} " if node.child_type else "" - if node.child_type is None: - count = "" - else: - count = f"{node.processed}/{node.expected}" if node.expected is not None else str(node.processed) - overflow = f" (+{node.overflow} more)" if node.overflow else "" - lines.append(f"{prefix}{label}{type_}{count}{overflow}".rstrip()) - for child in node.children: - _render_node(child, depth + 1, lines) - - -def snapshot_to_progress_payload(root: Optional[ProgressNodeSnapshot]) -> Optional[dict]: # noqa: UP045 - """Map a snapshot tree to the ProgressUpdate.progress (progressNode) shape.""" - return _node_to_payload(root) if root is not None else None - - -def format_eta(seconds: int) -> str: - """Human ``~45s`` / ``~6m`` / ``~1h 20m`` from a rounded seconds count.""" - if seconds < 60: - result = f"~{seconds}s" - elif seconds < 3600: - result = f"~{seconds // 60}m" - else: - result = f"~{seconds // 3600}h {(seconds % 3600) // 60}m" - return result - - -class ProgressReporter: - """Presentation of a registry's state. The registry stays a pure state - object; this reporter knows how to render it for CLI and for the SSE - payload. Connector-agnostic — the rollup header is generic.""" - - def __init__(self, registry: ProgressRegistry) -> None: - self._registry = registry - - def cli(self) -> str: - snapshot = self._registry.snapshot() - counters = self._registry.global_counters() - header = self._header(counters) - if snapshot is None: - return header if counters else "" - lines: List[str] = [] # noqa: UP006 - _render_joined(snapshot, [], lines) - tree = "\n".join(lines) - return f"{header}\n{tree}" if tree else header - - def _header(self, counters: "List[Tuple[str, int, Optional[int]]]") -> str: # noqa: UP006,UP045 - lines: List[str] = [] # noqa: UP006 - driver = self._driver_counter(counters) - eta = self.eta_seconds() - for type_, done, total in counters: - line = f"{type_} {done}/{total}" if total is not None else f"{type_} {done}" - if eta is not None and driver is not None and type_ == driver[0]: - line = f"{line} {format_eta(eta)}" - lines.append(line) - assets = self._registry.assets_ingested() - if assets > 0: - lines.append(f"Ingested: {assets:,} assets") - return "\n".join(lines) - - def global_counters(self) -> "List[Tuple[str, int, Optional[int]]]": # noqa: UP006,UP045 - """Run-level counters ``(type, done, total)`` for the SSE - ``ProgressUpdate.globalCounters``. Independent of the progress tree, so - reported even when the active tree is momentarily empty.""" - return self._registry.global_counters() - - def assets_ingested(self) -> int: - """Monotonic run-total of leaf assets ingested (tables + stored - procedures + other leaf entities). Independent of the tree, so it stays - correct after scope pruning and is authoritative on the terminal event.""" - return self._registry.assets_ingested() - - def _driver_counter( - self, - counters: "Optional[List[Tuple[str, int, Optional[int]]]]" = None, # noqa: UP006,UP045 - ) -> "Optional[Tuple[str, int, Optional[int]]]": # noqa: UP006,UP045 - """The ETA driver: the last-declared global counter that has a known - total (finest-grained real unit of work), or ``None`` when no counter - declares a total.""" - pool = self._registry.global_counters() if counters is None else counters - driver = None - for counter in pool: - if counter[2] is not None: - driver = counter - return driver - - def eta_seconds(self) -> Optional[int]: # noqa: UP045 - """Overall run ETA in seconds from the driver counter's cumulative rate: - ``elapsed * (total - done) / done``. ``None`` during warm-up (``done == - 0``), when there is no driver, when the driver is complete (``done >= - total``), or before elapsed time is available.""" - driver = self._driver_counter() - result = None - if driver is not None: - _, done, total = driver - elapsed = self._registry.elapsed_seconds() - if done > 0 and total is not None and done < total and elapsed is not None and elapsed > 0: - result = round(elapsed * (total - done) / done) - return result - - def payload(self) -> Optional[dict]: # noqa: UP045 - """Bare ``progressNode`` tree (validates against ``ProgressUpdate``, - which is ``extra='forbid'``). The run-root ``processed`` carries the - monotonic asset total; ``expected`` is unknown (no global count).""" - snapshot = self._registry.snapshot() - result = None - if snapshot is not None: - tree = snapshot_to_progress_payload(snapshot) - if tree is not None: - tree["processed"] = self._registry.assets_ingested() - tree["expected"] = None - result = tree - return result - - -def _render_joined(node: ProgressNodeSnapshot, ancestors: "List[str]", lines: "List[str]") -> None: # noqa: UP006 - """Render active nodes with ancestor labels joined by '.' so sibling-level - hierarchy collapses into a single readable label, e.g. ``sales.public Table 45/310``.""" - path = [*ancestors, node.label] if node.label else list(ancestors) - if node.children: - for child in node.children: - _render_joined(child, path, lines) - else: - joined_label = ".".join(path) if path else "" - type_ = f"{node.child_type} " if node.child_type else "" - if node.child_type is None: - count = "" - else: - count = f"{node.processed}/{node.expected}" if node.expected is not None else str(node.processed) - overflow = f" (+{node.overflow} more)" if node.overflow else "" - lines.append(f"{joined_label} {type_}{count}{overflow}".strip()) - - -def _node_to_payload(node: ProgressNodeSnapshot) -> dict: - return { - "label": node.label, - "entityType": node.child_type, - "processed": node.processed, - "expected": node.expected, - "active": node.active, - "overflow": node.overflow, - "children": [_node_to_payload(child) for child in node.children], - } diff --git a/ingestion/src/metadata/workflow/workflow_status_mixin.py b/ingestion/src/metadata/workflow/workflow_status_mixin.py index 96a43e81ffbb..0660fee06c74 100644 --- a/ingestion/src/metadata/workflow/workflow_status_mixin.py +++ b/ingestion/src/metadata/workflow/workflow_status_mixin.py @@ -16,10 +16,7 @@ import uuid from datetime import datetime from enum import Enum -from typing import TYPE_CHECKING, Optional, Tuple # noqa: UP035 - -if TYPE_CHECKING: - from metadata.workflow.progress_render import ProgressReporter +from typing import Optional, Tuple # noqa: UP035 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import ( IngestionPipeline, @@ -40,9 +37,9 @@ from metadata.generated.schema.type.basic import Map, Timestamp from metadata.ingestion.api.step import Step, Summary from metadata.ingestion.ometa.ometa_api import OpenMetadata +from metadata.ingestion.progress.registry import ProgressRegistry from metadata.utils.logger import ometa_logger from metadata.workflow.context.context_manager import ContextManager -from metadata.workflow.progress_render import ProgressReporter logger = ometa_logger() @@ -176,17 +173,17 @@ def build_ingestion_status(self) -> Optional[IngestionStatus]: # noqa: UP045 ] ) - def _progress_reporter(self) -> Optional["ProgressReporter"]: - """Reporter for the first workflow step that ran a topology walk. Reads - the backing attribute so we never create an empty registry on a step - that never tracked progress.""" - reporter = None + def _find_progress_registry(self) -> Optional[ProgressRegistry]: # noqa: UP045 + """Registry of the first workflow step that tracked progress. Reads the + backing attribute so we never create an empty registry on a step that + never tracked progress.""" + result = None for step in self.workflow_steps(): # pyright: ignore[reportAttributeAccessIssue] registry = getattr(step, "_progress_registry", None) if registry is not None: - reporter = ProgressReporter(registry) + result = registry break - return reporter + return result def send_progress_update(self, update_type: ProgressUpdateType = ProgressUpdateType.PROCESSING) -> None: """ @@ -199,11 +196,11 @@ def send_progress_update(self, update_type: ProgressUpdateType = ProgressUpdateT and self.ingestion_pipeline and self.ingestion_pipeline.fullyQualifiedName ): - reporter = self._progress_reporter() - progress_data = reporter.payload() if reporter is not None else None - counters = reporter.global_counters() if reporter is not None else [] - eta_seconds = reporter.eta_seconds() if reporter is not None else None - total_assets = reporter.assets_ingested() if reporter is not None else None + registry = self._find_progress_registry() + progress_data = registry.sse_payload() if registry is not None else None + counters = registry.global_counters() if registry is not None else [] + eta_seconds = registry.eta_seconds() if registry is not None else None + total_assets = registry.assets_ingested() if registry is not None else None progress_update = ProgressUpdate( runId=self.run_id, diff --git a/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py b/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py index aae250b67a4d..70cb4a857f51 100644 --- a/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py +++ b/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py @@ -14,7 +14,6 @@ from metadata.ingestion.api.models import Either from metadata.ingestion.source.database.snowflake.lineage import SnowflakeLineageSource -from metadata.workflow.progress_render import ProgressReporter class TestAccessHistoryProgress: @@ -29,4 +28,4 @@ def test_tracks_lineage_records_without_total_or_eta(self): assert len(results) == 4 assert source.progress.global_counters() == [("LineageRecords", 4, None)] - assert ProgressReporter(source.progress).eta_seconds() is None + assert source.progress.eta_seconds() is None diff --git a/ingestion/tests/unit/topology/dashboard/test_powerbi.py b/ingestion/tests/unit/topology/dashboard/test_powerbi.py index 8f212a02d7ae..387dd1b45642 100644 --- a/ingestion/tests/unit/topology/dashboard/test_powerbi.py +++ b/ingestion/tests/unit/topology/dashboard/test_powerbi.py @@ -2421,8 +2421,6 @@ def test_upstream_datamart_lineage(self, *_): @pytest.mark.order(55) @patch.object(fqn, "build", side_effect=lambda *args, **kwargs: kwargs.get("chart_name")) def test_yield_dashboard_advances_workspace_progress(self, *_): - from metadata.workflow.progress_render import ProgressReporter - self.powerbi.state = WorkspaceState() self.powerbi.__dict__.pop("_progress_registry", None) @@ -2444,7 +2442,7 @@ def test_yield_dashboard_advances_workspace_progress(self, *_): list(self.powerbi.yield_dashboard(Group(id="ws-1", name="Sales"))) assert self.powerbi.progress.assets_ingested() == 2 - out = ProgressReporter(self.powerbi.progress).cli() + out = self.powerbi.progress.render_cli() assert "Sales.Dashboard" in out assert "Dashboard 2" in out @@ -2508,8 +2506,6 @@ def test_admin_workspace_progress_total_reconciles_to_active_scan_results(self): @pytest.mark.order(58) @patch.object(fqn, "build", side_effect=lambda *args, **kwargs: kwargs.get("chart_name")) def test_unnamed_workspace_keys_progress_on_id(self, *_): - from metadata.workflow.progress_render import ProgressReporter - self.powerbi.__dict__.pop("_progress_registry", None) self.powerbi.state.add_filtered_dashboard(PowerBIDashboard(id="dash-1", displayName="One", tiles=[])) @@ -2524,7 +2520,7 @@ def test_unnamed_workspace_keys_progress_on_id(self, *_): self.powerbi._open_group_progress("ws-x", {"Dashboard": None}) list(self.powerbi.yield_dashboard(Group(id="ws-x", name=None))) - out = ProgressReporter(self.powerbi.progress).cli() + out = self.powerbi.progress.render_cli() assert "ws-x.Dashboard" in out assert "None.Dashboard" not in out diff --git a/ingestion/tests/unit/workflow/test_progress_rendering.py b/ingestion/tests/unit/workflow/test_progress_rendering.py index 8381b746a40f..236271b005c7 100644 --- a/ingestion/tests/unit/workflow/test_progress_rendering.py +++ b/ingestion/tests/unit/workflow/test_progress_rendering.py @@ -16,13 +16,12 @@ ProgressNode, ProgressUpdate, ) -from metadata.ingestion.progress.registry import ProgressNodeSnapshot, ProgressRegistry -from metadata.workflow.progress_render import ( - ProgressReporter, +from metadata.ingestion.progress import ( format_eta, render_progress_tree, snapshot_to_progress_payload, ) +from metadata.ingestion.progress.registry import ProgressNodeSnapshot, ProgressRegistry def _tree_snapshot(): @@ -46,7 +45,7 @@ def test_counts_all_leaf_types(self): reg.advance(["db", "sch"], "Table") for _ in range(2): reg.advance(["db", "sch"], "StoredProcedure") - assert ProgressReporter(reg).assets_ingested() == 5 + assert reg.assets_ingested() == 5 def test_survives_scope_pruning(self): reg = ProgressRegistry() @@ -56,7 +55,7 @@ def test_survives_scope_pruning(self): reg.advance(["db", "sch"], "Table") reg.close(["db", "sch"]) reg.close(["db"]) - assert ProgressReporter(reg).assets_ingested() == 2 + assert reg.assets_ingested() == 2 class TestProgressRendering: @@ -122,20 +121,20 @@ def _snowflake_like_registry() -> ProgressRegistry: return registry -class TestProgressReporter: +class TestRegistryReporting: def test_cli_header_is_assets_ingested(self): registry = _snowflake_like_registry() - text = ProgressReporter(registry).cli() + text = registry.render_cli() assert text.splitlines()[0] == f"Ingested: {registry.assets_ingested():,} assets" def test_cli_still_shows_active_scope(self): - text = ProgressReporter(_snowflake_like_registry()).cli() + text = _snowflake_like_registry().render_cli() assert "sales.public" in text assert "45/310" in text def test_payload_root_carries_asset_total(self): registry = _snowflake_like_registry() - payload = ProgressReporter(registry).payload() + payload = registry.sse_payload() assert set(payload) == { "label", "entityType", @@ -153,13 +152,13 @@ def test_payload_validates_against_progress_update_model(self): runId="r", timestamp=1, updateType="PROCESSING", - progress=ProgressReporter(_snowflake_like_registry()).payload(), + progress=_snowflake_like_registry().sse_payload(), ) assert update.progress is not None assert update.progress.processed == 45 def test_payload_is_none_before_anything_starts(self): - assert ProgressReporter(ProgressRegistry()).payload() is None + assert ProgressRegistry().sse_payload() is None def test_progress_update_carries_global_counters(self): reg = ProgressRegistry() @@ -167,7 +166,7 @@ def test_progress_update_carries_global_counters(self): reg.track("Workspaces") reg.track("Workspaces") reg.track("Workspaces") - counters = ProgressReporter(reg).global_counters() + counters = reg.global_counters() update = ProgressUpdate( runId="r1", timestamp=1, @@ -189,7 +188,7 @@ def test_header_shows_counters_then_assets(self): for _ in range(12): reg.track("DatabaseSchema") reg.advance([], "Table") - lines = ProgressReporter(reg).cli().splitlines() + lines = reg.render_cli().splitlines() assert lines[0] == "Database 2/4" assert lines[1].startswith("DatabaseSchema 12/45") assert lines[2] == "Ingested: 1 assets" @@ -198,25 +197,25 @@ def test_header_unknown_total_renders_bare_count(self): reg = ProgressRegistry() reg.set_total("Workspaces", None) reg.track("Workspaces") - assert ProgressReporter(reg).cli() == "Workspaces 1" + assert reg.render_cli() == "Workspaces 1" def test_header_renders_with_empty_tree_when_counter_active(self): reg = ProgressRegistry() reg.set_total("Workspaces", 4) for _ in range(4): reg.track("Workspaces") - assert ProgressReporter(reg).cli() == "Workspaces 4/4" + assert reg.render_cli() == "Workspaces 4/4" def test_no_group_keeps_legacy_header(self): reg = ProgressRegistry() reg.open([], "Database", None) reg.open(["x"], "Table", 2) reg.advance(["x"], "Table") - assert ProgressReporter(reg).cli().splitlines()[0] == "Ingested: 1 assets" + assert reg.render_cli().splitlines()[0] == "Ingested: 1 assets" def test_cli_empty_when_no_progress(self): reg = ProgressRegistry() - assert ProgressReporter(reg).cli() == "" + assert reg.render_cli() == "" class TestGlobalCountersOnSseUpdate: @@ -226,7 +225,7 @@ def test_reporter_global_counters_passthrough(self): reg.track("Workspaces") reg.track("Workspaces") reg.track("Workspaces") - assert ProgressReporter(reg).global_counters() == [("Workspaces", 3, 10)] + assert reg.global_counters() == [("Workspaces", 3, 10)] class TestFormatEta: @@ -251,14 +250,14 @@ def test_warmup_done_zero_is_none(self): reg = ProgressRegistry() reg.open([], "Database", None) reg.set_total("DatabaseSchema", 45) - assert ProgressReporter(reg).eta_seconds() is None + assert reg.eta_seconds() is None def test_no_counter_with_total_is_none(self): reg = ProgressRegistry() reg.open([], "Database", None) reg.set_total("Workspaces", None) reg.track("Workspaces") - assert ProgressReporter(reg).eta_seconds() is None + assert reg.eta_seconds() is None def test_complete_is_none(self): reg = ProgressRegistry() @@ -266,7 +265,7 @@ def test_complete_is_none(self): reg.set_total("DatabaseSchema", 4) for _ in range(4): reg.track("DatabaseSchema") - assert ProgressReporter(reg).eta_seconds() is None + assert reg.eta_seconds() is None def test_computable_without_open_when_counter_active(self): reg = ProgressRegistry() @@ -276,7 +275,7 @@ def test_computable_without_open_when_counter_active(self): for _ in range(2): reg.track("DatabaseSchema") clock.return_value = 60.0 - assert ProgressReporter(reg).eta_seconds() == 240 + assert reg.eta_seconds() == 240 def test_normal_case_uses_cumulative_rate(self): reg = ProgressRegistry() @@ -287,7 +286,7 @@ def test_normal_case_uses_cumulative_rate(self): for _ in range(2): reg.track("DatabaseSchema") clock.return_value = 60.0 - assert ProgressReporter(reg).eta_seconds() == 240 + assert reg.eta_seconds() == 240 def test_driver_is_last_counter_with_total(self): reg = ProgressRegistry() @@ -300,7 +299,7 @@ def test_driver_is_last_counter_with_total(self): for _ in range(4): reg.track("DatabaseSchema") clock.return_value = 100.0 - assert ProgressReporter(reg).eta_seconds() == 400 + assert reg.eta_seconds() == 400 class TestEtaInHeader: @@ -316,7 +315,7 @@ def test_eta_suffix_on_driver_line_only(self): for _ in range(9): reg.track("DatabaseSchema") clock.return_value = 120.0 - lines = ProgressReporter(reg).cli().splitlines() + lines = reg.render_cli().splitlines() assert lines[0] == "Database 2/4" assert lines[1].startswith("DatabaseSchema 9/45") assert "~8m" in lines[1] @@ -326,7 +325,7 @@ def test_no_eta_suffix_when_not_computable(self): reg = ProgressRegistry() reg.open([], "Database", None) reg.set_total("DatabaseSchema", 45) - assert "~" not in ProgressReporter(reg).cli() + assert "~" not in reg.render_cli() class TestHeaderIngestedSuppression: @@ -334,7 +333,7 @@ def test_ingested_line_suppressed_when_no_assets(self): registry = ProgressRegistry() registry.set_total("Queries", 100) registry.track("Queries", 12) - text = ProgressReporter(registry).cli() + text = registry.render_cli() assert "Queries 12/100" in text assert "Ingested:" not in text @@ -342,7 +341,7 @@ def test_totalless_counter_renders_without_denominator_or_eta(self): registry = ProgressRegistry() registry.set_total("LineageRecords", None) registry.track("LineageRecords", 8210) - text = ProgressReporter(registry).cli() + text = registry.render_cli() assert "LineageRecords 8210" in text assert "/" not in text assert "Ingested:" not in text @@ -351,5 +350,5 @@ def test_ingested_line_kept_when_assets_present(self): registry = ProgressRegistry() registry.open(["db", "schema"], "Table", 3) registry.advance(["db", "schema"], "Table") - text = ProgressReporter(registry).cli() + text = registry.render_cli() assert "Ingested: 1 assets" in text diff --git a/ingestion/tests/unit/workflow/test_status_mixin_progress.py b/ingestion/tests/unit/workflow/test_status_mixin_progress.py index 0ea30bab1638..29c3230448f5 100644 --- a/ingestion/tests/unit/workflow/test_status_mixin_progress.py +++ b/ingestion/tests/unit/workflow/test_status_mixin_progress.py @@ -8,7 +8,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""send_progress_update maps ProgressReporter.eta_seconds into the SSE payload.""" +"""send_progress_update maps ProgressRegistry.eta_seconds into the SSE payload.""" from types import SimpleNamespace from unittest.mock import MagicMock, patch From f3c4c28ac63fb154d1beb332966669bd48d80303 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 11:06:34 +0530 Subject: [PATCH 06/18] refactor(ingestion): formalize MANUAL progress mode for query-parser and PowerBI sources Co-Authored-By: Claude Fable 5 --- .../source/dashboard/dashboard_service.py | 16 +++++----------- .../source/dashboard/powerbi/metadata.py | 7 +++++-- .../ingestion/source/database/lineage_source.py | 6 +++--- .../source/database/query_parser_source.py | 5 ++++- .../source/database/snowflake/lineage.py | 6 +++--- .../ingestion/source/database/usage_source.py | 6 +++--- 6 files changed, 23 insertions(+), 23 deletions(-) diff --git a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py index 031e4678d80a..21021f4a419d 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py @@ -216,27 +216,21 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC): chart_source_state: Set = set() # noqa: RUF012, UP006 def _declare_progress_groups(self, label: str, total: Optional[int]) -> None: # noqa: UP045 - """Declare the grouping axis (e.g. workspaces) as a global counter and - remember its label so completion can target it at scope close.""" - self.__dict__["_progress_counter_label"] = label # pyright: ignore[reportIndexIssue] - self.progress.set_total(label, total) + """Declare the grouping axis (e.g. workspaces) as a global counter.""" + self.manual_progress.declare_groups(label, total) def _open_group_progress(self, group: str, expected_by_type: Dict[str, Optional[int]]) -> None: # noqa: UP006, UP045 """Open one child node per asset type under ``group`` so each type renders as its own line; ``expected`` may be None for lazy (running) counts.""" - for asset_type, expected in expected_by_type.items(): - self.progress.open([group, asset_type], asset_type, expected) + self.manual_progress.open_group(group, expected_by_type) def _advance_group_progress(self, group: str, asset_type: str) -> None: """Record one processed asset of ``asset_type`` under ``group``.""" - self.progress.advance([group, asset_type], asset_type) + self.manual_progress.advance(group, asset_type) def _close_group_progress(self, group: str) -> None: """Count the finished group on its global counter and prune its subtree.""" - label = self.__dict__.get("_progress_counter_label") - if label is not None: - self.progress.track(label) - self.progress.close([group]) + self.manual_progress.close_group(group) @retry_with_docker_host() def __init__( diff --git a/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py b/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py index c659c0d5b499..9daa3859d2ae 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py @@ -65,6 +65,7 @@ from metadata.ingestion.models.barrier import Barrier from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.ometa.utils import model_str +from metadata.ingestion.progress.modes import ProgressMode from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource from metadata.ingestion.source.dashboard.powerbi.constants import ( BIGQUERY_QUERY_EXPRESSION_KW, @@ -133,6 +134,8 @@ class LineageTargetSpec: class PowerbiSource(DashboardServiceSource): """PowerBi Source Class""" + progress_mode = ProgressMode.MANUAL + config: WorkflowSource metadata_config: OpenMetadataConnection @@ -167,7 +170,7 @@ def get_org_workspace_data(self) -> Iterable[Optional[Group]]: # noqa: UP045 workspaces = self.client.api_client.fetch_all_workspaces(filter_pattern) if workspaces: workspace_total += len(workspaces) - self.progress.set_total("Workspaces", workspace_total) + self.manual_progress.set_total("Workspaces", workspace_total) for workspace in workspaces: # add the dashboards to the workspace workspace.dashboards.extend( @@ -307,7 +310,7 @@ def get_admin_workspace_data(self) -> Iterable[Optional[Group]]: # noqa: UP045 missing_workspaces, ) active_workspace_total += len(active_workspaces) - self.progress.set_total("Workspaces", active_workspace_total) + self.manual_progress.set_total("Workspaces", active_workspace_total) yield from active_workspaces count += 1 else: diff --git a/ingestion/src/metadata/ingestion/source/database/lineage_source.py b/ingestion/src/metadata/ingestion/source/database/lineage_source.py index 13ae5121338f..7b1c3ffe0e53 100644 --- a/ingestion/src/metadata/ingestion/source/database/lineage_source.py +++ b/ingestion/src/metadata/ingestion/source/database/lineage_source.py @@ -367,14 +367,14 @@ def yield_query_lineage( result_limit = self.source_config.resultLimit # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue] if result_limit is not None: - self.progress.seed_scope_total("Queries", "run", result_limit) + self.manual_progress.seed_scope_total("Queries", "run", result_limit) produced = 0 def producer_fn(): nonlocal produced for table_query in self.query_lineage_producer(): produced += 1 - self.progress.track("Queries") + self.manual_progress.track("Queries") yield table_query processor_fn = query_lineage_processor @@ -394,7 +394,7 @@ def producer_fn(): args, max_threads=self.source_config.threads, # pyright: ignore[reportAttributeAccessIssue] ) - self.progress.reconcile_scope_total("Queries", "run", produced) + self.manual_progress.reconcile_scope_total("Queries", "run", produced) def view_lineage_producer(self) -> Iterable[TableView]: """ diff --git a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py index def95ff327c3..c22941ed3912 100644 --- a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py +++ b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py @@ -27,10 +27,11 @@ from metadata.ingestion.lineage.masker import masked_query_cache from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper from metadata.ingestion.ometa.ometa_api import OpenMetadata +from metadata.ingestion.progress.modes import ProgressMode +from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.ingestion.source.connections import test_connection_common from metadata.utils.helpers import get_start_and_end, retry_with_docker_host from metadata.utils.logger import ingestion_logger -from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.utils.ssl_manager import get_ssl_connection logger = ingestion_logger() @@ -46,6 +47,8 @@ class QueryParserSource(ProgressTrackingMixin, Source, ABC): some utilities to be overwritten when necessary """ + progress_mode = ProgressMode.MANUAL + sql_stmt: str dialect: str filters: str diff --git a/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py b/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py index 468d6aef3ab1..b0fcad36ea25 100644 --- a/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py +++ b/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py @@ -245,17 +245,17 @@ def _yield_access_history_lineage(self) -> Iterable[Either[AddLineageRequest]]: The two phases are isolated: a catastrophic failure in the combined ACCESS_HISTORY phase must not stop the COPY_HISTORY phase, and vice versa. """ - self.progress.set_total("LineageRecords", None) + self.manual_progress.set_total("LineageRecords", None) try: for either in self._yield_combined_access_history(): - self.progress.track("LineageRecords") + self.manual_progress.track("LineageRecords") yield either except Exception as exc: logger.warning("ACCESS_HISTORY combined lineage phase failed: %s", exc) logger.debug(traceback.format_exc()) try: for either in self._yield_copy_history_lineage(): - self.progress.track("LineageRecords") + self.manual_progress.track("LineageRecords") yield either except Exception as exc: logger.warning("COPY_HISTORY lineage phase failed: %s", exc) diff --git a/ingestion/src/metadata/ingestion/source/database/usage_source.py b/ingestion/src/metadata/ingestion/source/database/usage_source.py index 653d5177ec5e..6b417d486005 100644 --- a/ingestion/src/metadata/ingestion/source/database/usage_source.py +++ b/ingestion/src/metadata/ingestion/source/database/usage_source.py @@ -166,12 +166,12 @@ def _iter(self, *_, **__) -> Iterable[Either[TableQuery]]: days = max(1, (self.end - self.start).days) result_limit = self.source_config.resultLimit # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue] if result_limit is not None: - self.progress.seed_scope_total("Queries", "run", result_limit * days) + self.manual_progress.seed_scope_total("Queries", "run", result_limit * days) processed = 0 for table_queries in self.get_table_query(): if table_queries: count = len(table_queries.queries) # pyright: ignore[reportAttributeAccessIssue] - self.progress.track("Queries", count) + self.manual_progress.track("Queries", count) processed += count yield Either(right=table_queries) - self.progress.reconcile_scope_total("Queries", "run", processed) + self.manual_progress.reconcile_scope_total("Queries", "run", processed) From f615854511d15437d247c39104163304c2ea3327 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 11:13:52 +0530 Subject: [PATCH 07/18] feat(ingestion): TopologyProgressTracker extracts runner progress policy behind NodeProgress handles Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/progress/__init__.py | 8 + .../ingestion/progress/runner_tracker.py | 243 ++++++++++++++++++ .../unit/progress/test_runner_tracker.py | 153 +++++++++++ 3 files changed, 404 insertions(+) create mode 100644 ingestion/src/metadata/ingestion/progress/runner_tracker.py create mode 100644 ingestion/tests/unit/progress/test_runner_tracker.py diff --git a/ingestion/src/metadata/ingestion/progress/__init__.py b/ingestion/src/metadata/ingestion/progress/__init__.py index b6cef80a65a2..986ef03ff1c5 100644 --- a/ingestion/src/metadata/ingestion/progress/__init__.py +++ b/ingestion/src/metadata/ingestion/progress/__init__.py @@ -38,18 +38,26 @@ ProgressNodeSnapshot, ProgressRegistry, ) +from metadata.ingestion.progress.runner_tracker import ( + NO_OP_NODE_PROGRESS, + NodeProgress, + TopologyProgressTracker, +) from metadata.ingestion.progress.tracking import ProgressTrackingMixin __all__ = [ "DEFAULT_ACTIVE_LEAF_CAP", "GlobalCounter", "ManualProgress", + "NO_OP_NODE_PROGRESS", + "NodeProgress", "ProgressMode", "ProgressModeError", "ProgressNode", "ProgressNodeSnapshot", "ProgressRegistry", "ProgressTrackingMixin", + "TopologyProgressTracker", "TotalsDeclarer", "format_eta", "render_progress_tree", diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py new file mode 100644 index 000000000000..4d232e4bfcc8 --- /dev/null +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -0,0 +1,243 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Topology-walk progress policy, extracted from the runner. + +The runner asks ``for_node`` for a per-node handle and calls it +unconditionally. When the source is not in AUTO mode, the node is the +structural service root, or the node has no typed stage, the handle is a +shared no-op — the walk code carries zero progress conditionals. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set # noqa: UP035 + +from metadata.ingestion.models.topology import ( + NodeStage, + TopologyNode, + get_topology_nodes, + get_topology_root, +) +from metadata.ingestion.ometa.utils import model_str +from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer + +if TYPE_CHECKING: + from metadata.ingestion.progress.registry import ProgressRegistry + + +class _NoOpScope: + """No container scope to retire.""" + + def exit(self) -> None: + """Nothing to close.""" + + +NO_OP_SCOPE = _NoOpScope() + + +class _Scope: + """An open container scope: exiting prunes the subtree from the tree and + counts the finished container on its global counter.""" + + def __init__(self, registry: "ProgressRegistry", path: List[str], entity_type_name: str) -> None: # noqa: UP006 + self._registry = registry + self._path = path + self._entity_type_name = entity_type_name + + def exit(self) -> None: + self._registry.close(self._path) + self._registry.track(self._entity_type_name) + + +class _NoOpNodeProgress: + """Shared do-nothing handle: zero registry calls, zero per-entity overhead.""" + + wants_eager_count = False + + def open_with_count(self, count: int) -> None: + """No-op.""" + + def open_lazy(self) -> None: + """No-op.""" + + def advance_leaf(self) -> None: + """No-op.""" + + def enter_scope(self) -> _NoOpScope: + return NO_OP_SCOPE + + +NO_OP_NODE_PROGRESS = _NoOpNodeProgress() + + +class NodeProgress: + """Progress recording for one topology node. The parent path is resolved at + creation time — before worker threads spawn — so the handle can be shared + across a node's worker threads; ``enter_scope`` reads the *current thread's* + context at call time.""" + + def __init__( + self, + tracker: "TopologyProgressTracker", + node: TopologyNode, + entity_type_name: str, + parent_path: List[str], # noqa: UP006 + is_leaf: bool, + ) -> None: + self._tracker = tracker + self._node = node + self._registry = tracker.registry + self._entity_type_name = entity_type_name + self._parent_path = parent_path + self._is_leaf = is_leaf + self._reconcilable = self._registry.is_reconcilable(entity_type_name) + + @property + def wants_eager_count(self) -> bool: + """Materialize the producer for an exact child count: always for a + leaf; for a container only when its counter is reconcilable.""" + return self._is_leaf or self._reconcilable + + def open_with_count(self, count: int) -> None: + self._registry.open(self._parent_path, self._entity_type_name, count) + if self._reconcilable and self._parent_path: + self._registry.reconcile_scope_total(self._entity_type_name, self._parent_path[-1], count) + + def open_lazy(self) -> None: + self._registry.open(self._parent_path, self._entity_type_name, None) + + def advance_leaf(self) -> None: + if self._is_leaf: + self._registry.advance(self._parent_path, self._entity_type_name) + + def enter_scope(self): + """Scope handle for the container entity currently in this thread's + context; NO_OP for a leaf node or when the context value is unset.""" + result = NO_OP_SCOPE + if not self._is_leaf: + scope_path = self._tracker.scope_path_for_node(self._node, self._parent_path) + if scope_path is not None: + result = _Scope(self._registry, scope_path, self._entity_type_name) + return result + + +class TopologyProgressTracker: + """All topology-walk progress policy for one source: mode gating, ancestor + path resolution, once-only totals declaration, and handle creation. + + ``source`` is the TopologyRunnerMixin instance; the tracker reads its + ``topology``, ``context``, ``progress``, ``progress_mode``, + ``declare_progress_totals`` and primary-stage helpers.""" + + def __init__(self, source: Any) -> None: + self._source = source + self._root_node_ids: Set[int] = set() # noqa: UP006 + self._root_ctx_keys: Optional[Set[str]] = None # noqa: UP006,UP045 + self._primary_stage_idx: Optional[Dict[str, NodeStage]] = None # noqa: UP006,UP045 + self._totals_declared = False + + @property + def registry(self) -> "ProgressRegistry": + return self._source.progress + + def on_walk_start(self, root_nodes) -> None: + self._root_node_ids = {id(node) for node in root_nodes} + + def is_root_node(self, node: TopologyNode) -> bool: + """The topology root (the Service node) is a structural wrapper, not a + progress level: it always holds a single entity and would otherwise + collide with the first real level (databases) at the registry root.""" + return id(node) in self._root_node_ids + + def for_node(self, node: TopologyNode, is_leaf: bool): + """A NodeProgress for AUTO sources on real entity nodes; the shared + no-op otherwise. Declares connector totals on the first real node — + after the service root has populated the context the totals code may + need.""" + result = NO_OP_NODE_PROGRESS + if self._source.progress_mode is ProgressMode.AUTO and not self.is_root_node(node): + entity_type_name = self._source._get_entity_type_for_node(node) + if entity_type_name: + self._declare_totals_once() + parent_path = self.current_path(entity_type_name) + result = NodeProgress(self, node, entity_type_name, parent_path, is_leaf) + return result + + def _declare_totals_once(self) -> None: + if not self._totals_declared: + self._totals_declared = True + self._source.declare_progress_totals(TotalsDeclarer(self.registry)) + + def current_path(self, entity_type_name: Optional[str] = None) -> List[str]: # noqa: UP006,UP045 + """Ancestor container labels from the node's primary-stage ``consumer`` + chain minus the service-root key, each remaining key resolved to its + current context value. + + ``consumer`` lists only *ancestors*; a node's own context key is never + in its own consumer list. In the depth-first walk every ancestor + container is mid-iteration, so its value is fresh — this is why the + ``database_schema``-not-cleared-between-siblings bug cannot occur here. + Returns ``[]`` for the run-grain default (no ``entity_type_name``) or a + root-level entity (empty consumer).""" + path: List[str] = [] # noqa: UP006 + stage = self._primary_entity_stage(entity_type_name) + if stage is not None: + root_keys = self._root_context_keys() + ctx = self._source.context.get() + for key in stage.consumer or []: + if key not in root_keys: + value = getattr(ctx, key, None) + if value is not None: + path.append(model_str(value)) + return path + + def scope_path_for_node(self, node: TopologyNode, parent_path: List[str]) -> Optional[List[str]]: # noqa: UP006,UP045 + """The path of the container entity currently in context: its parent + ancestors plus this node's own context value. Used to prune a scope + from the progress tree once its children finish. ``None`` when the + node's context value is not set.""" + result = None + stage = self._source._node_primary_stage(node) + if stage is not None and stage.context and parent_path is not None: + value = getattr(self._source.context.get(), stage.context, None) + if value is not None: + result = [*parent_path, model_str(value)] + return result + + def _root_context_keys(self) -> Set[str]: # noqa: UP006 + """Context keys owned by root topology nodes (the service level), + derived from the topology — never a string literal. Cached; the + topology is static.""" + if self._root_ctx_keys is None: + self._root_ctx_keys = { + stage.context + for node in get_topology_root(self._source.topology) + for stage in node.stages + if stage.context + } + return self._root_ctx_keys + + def _primary_stage_index(self) -> Dict[str, NodeStage]: # noqa: UP006 + """Map of entity type name → that node's primary stage, built once from + the static topology. Entity types are unique per node, so the first + match wins on the rare collision.""" + if self._primary_stage_idx is None: + index: Dict[str, NodeStage] = {} # noqa: UP006 + for node in get_topology_nodes(self._source.topology): + stage = self._source._node_primary_stage(node) + if stage is not None: + index.setdefault(stage.type_.__name__, stage) + self._primary_stage_idx = index + return self._primary_stage_idx + + def _primary_entity_stage(self, entity_type_name: Optional[str]) -> Optional[NodeStage]: # noqa: UP045 + result = None + if entity_type_name is not None: + result = self._primary_stage_index().get(entity_type_name) + return result diff --git a/ingestion/tests/unit/progress/test_runner_tracker.py b/ingestion/tests/unit/progress/test_runner_tracker.py new file mode 100644 index 000000000000..68425e66f070 --- /dev/null +++ b/ingestion/tests/unit/progress/test_runner_tracker.py @@ -0,0 +1,153 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""TopologyProgressTracker: gating, path resolution, handles, totals hook.""" + +from unittest.mock import MagicMock + +from metadata.ingestion.api.topology_runner import TopologyRunnerMixin +from metadata.ingestion.models.topology import TopologyContextManager, get_topology_node, get_topology_root +from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer +from metadata.ingestion.progress.runner_tracker import ( + NO_OP_NODE_PROGRESS, + NodeProgress, + TopologyProgressTracker, +) +from metadata.ingestion.source.database.database_service import DatabaseServiceTopology + + +class _FakeSource(TopologyRunnerMixin): + topology = DatabaseServiceTopology() + + def __init__(self, mode=ProgressMode.AUTO): + self.progress_mode = mode + self.context = TopologyContextManager(self.topology) + self.status = MagicMock() + self.declared = 0 + + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: + self.declared += 1 + totals.set_total("Database", 7) + + +def _table_node(source): + return get_topology_node("table", source.topology) + + +def _database_node(source): + return get_topology_node("database", source.topology) + + +def test_for_node_returns_noop_when_mode_is_not_auto(): + for mode in (ProgressMode.MANUAL, ProgressMode.OFF): + source = _FakeSource(mode) + tracker = TopologyProgressTracker(source) + assert tracker.for_node(_table_node(source), is_leaf=True) is NO_OP_NODE_PROGRESS + assert "_progress_registry" not in source.__dict__ + + +def test_for_node_returns_noop_for_root_node(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + root = get_topology_root(source.topology)[0] + tracker.on_walk_start([root]) + assert tracker.is_root_node(root) + assert tracker.for_node(root, is_leaf=False) is NO_OP_NODE_PROGRESS + + +def test_for_node_returns_real_handle_for_auto_leaf(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + handle = tracker.for_node(_table_node(source), is_leaf=True) + assert isinstance(handle, NodeProgress) + assert handle.wants_eager_count is True + + +def test_totals_hook_called_exactly_once_and_only_for_real_nodes(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + root = get_topology_root(source.topology)[0] + tracker.on_walk_start([root]) + tracker.for_node(root, is_leaf=False) + assert source.declared == 0 + tracker.for_node(_database_node(source), is_leaf=False) + tracker.for_node(_table_node(source), is_leaf=True) + assert source.declared == 1 + assert ("Database", 0, 7) in source.progress.global_counters() + + +def test_leaf_handle_counts_open_and_advance(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + handle = tracker.for_node(_table_node(source), is_leaf=True) + handle.open_with_count(2) + handle.advance_leaf() + handle.advance_leaf() + snapshot = source.progress.snapshot() + assert snapshot.child_type == "Table" + assert snapshot.processed == 2 + + +def test_container_handle_is_lazy_without_reconcilable_counter(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + handle = tracker.for_node(_database_node(source), is_leaf=False) + assert handle.wants_eager_count is False + handle.open_lazy() + assert source.progress.snapshot().expected_by_type.get("Database") is None + + +def test_container_handle_reconciles_when_counter_is_reconcilable(): + source = _FakeSource() + source.progress.seed_scope_total("DatabaseSchema", "salesdb", 1) + tracker = TopologyProgressTracker(source) + key = source._node_primary_stage(_database_node(source)).context + setattr(source.context.get(), key, "salesdb") + schema_node = get_topology_node("databaseSchema", source.topology) + handle = tracker.for_node(schema_node, is_leaf=False) + assert handle.wants_eager_count is True + handle.open_with_count(3) + counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + assert counters["DatabaseSchema"] == (0, 3) + + +def test_enter_scope_closes_and_tracks_on_exit(): + source = _FakeSource() + source.progress.set_total("Database", 7) + tracker = TopologyProgressTracker(source) + node = _database_node(source) + handle = tracker.for_node(node, is_leaf=False) + key = source._node_primary_stage(node).context + setattr(source.context.get(), key, "salesdb") + closed = [] + source.progress.close = lambda path: closed.append(list(path)) + scope = handle.enter_scope() + scope.exit() + assert closed == [["salesdb"]] + counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + assert counters["Database"][0] == 1 + + +def test_enter_scope_is_noop_without_context_value(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + handle = tracker.for_node(_database_node(source), is_leaf=False) + scope = handle.enter_scope() + scope.exit() + assert source.progress.snapshot() is None # nothing opened, nothing closed + + +def test_current_path_resolves_ancestor_context(): + source = _FakeSource() + tracker = TopologyProgressTracker(source) + assert tracker.current_path() == [] + db_key = source._node_primary_stage(_database_node(source)).context + setattr(source.context.get(), db_key, "salesdb") + assert tracker.current_path("DatabaseSchema") == ["salesdb"] From 5e06f18b60c6c82f5e37e91c215599f93cf0a47f Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 11:35:13 +0530 Subject: [PATCH 08/18] refactor(ingestion): topology runner counts via NodeProgress handles; progress default-on with AUTO/MANUAL/OFF Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/api/topology_runner.py | 163 ++++-------------- .../database/test_snowflake_progress_path.py | 19 +- .../topology/test_topology_progress_path.py | 52 +++--- .../topology/test_topology_runner_progress.py | 96 +++++++---- 4 files changed, 134 insertions(+), 196 deletions(-) diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index 560e1e0fb424..c3113a74f560 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -38,15 +38,15 @@ TopologyContextManager, TopologyNode, get_topology_node, - get_topology_nodes, get_topology_root, ) from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.ometa.utils import model_str +from metadata.ingestion.progress.runner_tracker import TopologyProgressTracker +from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.utils.custom_thread_pool import CustomThreadPoolExecutor from metadata.utils.logger import ingestion_logger from metadata.utils.operation_metrics import OperationMetricsState -from metadata.ingestion.progress.tracking import ProgressTrackingMixin from metadata.utils.source_hash import generate_source_hash logger = ingestion_logger() @@ -79,12 +79,6 @@ class TopologyRunnerMixin(ProgressTrackingMixin, Generic[C]): "AddLineageRequest", } - progress_tracking_enabled: ClassVar[bool] = False - """Master gate for hierarchical progress tracking. Off for every connector - by default: when False the runner makes zero progress calls, the registry - is never created, and there is no per-entity overhead. A connector opts in - by setting this True on its source class.""" - def _node_primary_stage(self, node: TopologyNode) -> Optional[NodeStage]: # noqa: UP045 """The node's primary, non-side-output stage — the stage whose entity is the node's real entity (Table, Database, ...). Falls back to the first @@ -119,95 +113,22 @@ def _run_node_producer(self, node: TopologyNode) -> Iterable[Entity]: ) ) - def _root_context_keys(self) -> set[str]: - """Context keys owned by root topology nodes (the service level). Root - nodes are those with no consumers — the existing definition of a - topology root — so the exclusion is derived, never a string literal. - Cached per source; the topology is static.""" - cached = self.__dict__.get("_root_ctx_keys") - if cached is None: - cached = { - stage.context for node in get_topology_root(self.topology) for stage in node.stages if stage.context - } - self.__dict__["_root_ctx_keys"] = cached - return cached - - def _primary_stage_index(self) -> "dict[str, NodeStage]": - """Map of entity type name → that node's primary stage, built once from - the static topology. Entity types are unique per node, so the first - match wins on the rare collision.""" - cached = self.__dict__.get("_primary_stage_idx") - if cached is None: - cached = {} - for node in get_topology_nodes(self.topology): - stage = self._node_primary_stage(node) - if stage is not None: - cached.setdefault(stage.type_.__name__, stage) - self.__dict__["_primary_stage_idx"] = cached - return cached - - def _primary_entity_stage(self, entity_type_name: Optional[str]) -> Optional[NodeStage]: # noqa: UP045 - """The primary stage whose entity type is the node's progress key — the - same stage ``_get_entity_type_for_node`` selects. ``None`` when no node - produces ``entity_type_name`` (or it is ``None``), so callers can skip - topology access entirely for the run-grain default.""" - result = None - if entity_type_name is not None: - result = self._primary_stage_index().get(entity_type_name) - return result - - def current_progress_path(self, entity_type_name: Optional[str] = None) -> List[str]: # noqa: UP006,UP045 - """Ancestor container labels from the node's primary-stage ``consumer`` - chain minus the service-root key, each remaining key resolved to its - current context value. - - ``consumer`` lists only *ancestors*; a node's own context key is never - in its own consumer list. In the depth-first walk every ancestor - container is mid-iteration, so its value is fresh — this is why the - ``database_schema``-not-cleared-between-siblings bug cannot occur here. - Returns ``[]`` for the run-grain default (no ``entity_type_name``) or a - root-level entity (empty consumer).""" - path: List[str] = [] # noqa: UP006 - stage = self._primary_entity_stage(entity_type_name) - if stage is not None: - root_keys = self._root_context_keys() - ctx = self.context.get() - for key in stage.consumer or []: - if key not in root_keys: - value = getattr(ctx, key, None) - if value is not None: - path.append(model_str(value)) - return path - - def _is_root_node(self, node: TopologyNode) -> bool: - """The topology root (the Service node) is a structural wrapper, not a - progress level: it always holds a single entity and would otherwise - collide with the first real level (databases) at the registry root.""" - return id(node) in self.__dict__.get("_root_node_ids", set()) - - def _should_track_progress(self, node: TopologyNode, entity_type_name: Optional[str]) -> bool: # noqa: UP045 - """Progress is recorded only for opted-in sources, for real entity - nodes, and never for the structural service root.""" - return self.progress_tracking_enabled and bool(entity_type_name) and not self._is_root_node(node) - - def _scope_path_for_node(self, node: TopologyNode, parent_path: List[str]) -> Optional[List[str]]: # noqa: UP006,UP045 - """The path of the container entity currently in context: its parent - ancestors plus this node's own context value. Used to prune a scope - from the progress tree once its children finish. ``None`` when the - node's context value is not set.""" - result = None - stage = self._node_primary_stage(node) - if stage is not None and stage.context and parent_path is not None: - value = getattr(self.context.get(), stage.context, None) - if value is not None: - result = [*parent_path, model_str(value)] - return result + @property + def progress_tracker(self) -> TopologyProgressTracker: + """Per-source topology progress policy. Lazy: first access is in + ``_iter`` before worker threads spawn.""" + tracker = self.__dict__.get("_topology_progress_tracker") + if tracker is None: + tracker = TopologyProgressTracker(self) + self.__dict__["_topology_progress_tracker"] = tracker + return tracker def _multithread_process_node(self, node: TopologyNode, threads: int) -> Iterable[Entity]: """Multithread Processing of a Node with progress tracking""" child_nodes = self._get_child_nodes(node) entity_type_name = self._get_entity_type_for_node(node) operation_metrics = OperationMetricsState() + node_progress = self.progress_tracker.for_node(node, is_leaf=not child_nodes) # Track SOURCE time - fetching entities from producer source_start = perf_counter() @@ -223,12 +144,7 @@ def _multithread_process_node(self, node: TopologyNode, threads: int) -> Iterabl entity_type=entity_type_name, ) - track_progress = self._should_track_progress(node, entity_type_name) - parent_path = self.current_progress_path(entity_type_name) if track_progress else None - if track_progress and parent_path is not None: - self.progress.open(parent_path, entity_type_name, node_entities_length) - if self.progress.is_reconcilable(entity_type_name) and parent_path: - self.progress.reconcile_scope_total(entity_type_name, parent_path[-1], node_entities_length) + node_progress.open_with_count(node_entities_length) if node_entities_length == 0: return @@ -246,8 +162,7 @@ def _multithread_process_node(self, node: TopologyNode, threads: int) -> Iterabl chunk, child_nodes, self.context.get_current_thread_id(), - entity_type_name, - parent_path, + node_progress, ) for chunk in chunks ] @@ -273,28 +188,21 @@ def _process_node(self, node: TopologyNode) -> Iterable[Entity]: Container producers (database, schema) are iterated lazily so connectors can set up per-yield inspector/session state before each child is processed. Leaf producers (table, stored procedure) are materialized - eagerly ONLY when progress tracking is enabled — to record an exact child - count via ``progress.open``; otherwise they are iterated lazily so - connectors whose stages depend on per-yield producer state (e.g. PowerBI's - per-workspace ``state.enter`` / ``finally: state.exit``) are not torn down - before their stages run. + eagerly only when the node's progress handle wants an eager count — to + record an exact child count via ``progress.open``; otherwise they are + iterated lazily so connectors whose stages depend on per-yield producer + state (e.g. PowerBI's per-workspace ``state.enter`` / ``finally: + state.exit``) are not torn down before their stages run. """ child_nodes = self._get_child_nodes(node) - entity_type_name = self._get_entity_type_for_node(node) - is_leaf = not child_nodes - track_progress = self._should_track_progress(node, entity_type_name) - parent_path = self.current_progress_path(entity_type_name) if track_progress else [] + node_progress = self.progress_tracker.for_node(node, is_leaf=not child_nodes) - reconcilable = track_progress and not is_leaf and self.progress.is_reconcilable(entity_type_name) - if track_progress and (is_leaf or reconcilable): + if node_progress.wants_eager_count: node_entities = list(self._run_node_producer(node) or []) - self.progress.open(parent_path, entity_type_name, len(node_entities)) - if reconcilable and parent_path: - self.progress.reconcile_scope_total(entity_type_name, parent_path[-1], len(node_entities)) + node_progress.open_with_count(len(node_entities)) else: node_entities = self._run_node_producer(node) or [] - if track_progress: - self.progress.open(parent_path, entity_type_name, None) + node_progress.open_lazy() for node_entity in node_entities: for stage in node.stages: @@ -304,14 +212,11 @@ def _process_node(self, node: TopologyNode) -> Iterable[Entity]: if stage.clear_context: self.context.get().clear_stage(stage=stage) - if track_progress and is_leaf: - self.progress.advance(parent_path, entity_type_name) + node_progress.advance_leaf() - scope_path = self._scope_path_for_node(node, parent_path) if track_progress and not is_leaf else None + scope = node_progress.enter_scope() yield from self.process_nodes(child_nodes) - if scope_path is not None: - self.progress.close(scope_path) - self.progress.track(entity_type_name) + scope.exit() def process_nodes(self, nodes: List[TopologyNode]) -> Iterable[Entity]: # noqa: UP006 """ @@ -358,8 +263,7 @@ def _multithread_process_entity( node_entities: List[Any], # noqa: UP006 child_nodes: List[TopologyNode], # noqa: UP006 parent_thread_id: int, - entity_type_name: Optional[str] = None, # noqa: UP045 - parent_path: Optional[List[str]] = None, # noqa: UP006,UP045 + node_progress, ): """Multithread processing of a Node Entity with progress tracking""" # Generates a new context based on the parent thread. @@ -378,17 +282,12 @@ def _multithread_process_entity( if stage.clear_context: self.context.get().clear_stage(stage=stage) - if parent_path is not None and entity_type_name and not child_nodes: - self.progress.advance(parent_path, entity_type_name) + node_progress.advance_leaf() - scope_path = ( - self._scope_path_for_node(node, parent_path) if parent_path is not None and child_nodes else None - ) + scope = node_progress.enter_scope() for child_result in self.process_nodes(child_nodes): self.queue.put(child_result) - if scope_path is not None: - self.progress.close(scope_path) - self.progress.track(entity_type_name) + scope.exit() # Merge thread-local metrics into global state before thread exits operation_metrics.merge_thread_metrics() @@ -471,7 +370,7 @@ def _iter(self) -> Iterable[Either]: :return: Iterable of the Entities yielded by all nodes in the topology """ root_nodes = get_topology_root(self.topology) - self.__dict__["_root_node_ids"] = {id(node) for node in root_nodes} # pyright: ignore[reportIndexIssue] + self.progress_tracker.on_walk_start(root_nodes) yield from self.process_nodes(root_nodes) # pyright: ignore[reportReturnType] def create_patch_request(self, original_entity: Entity, create_request: C) -> PatchRequest: diff --git a/ingestion/tests/unit/source/database/test_snowflake_progress_path.py b/ingestion/tests/unit/source/database/test_snowflake_progress_path.py index c9323d00ae7f..538cf3c88a08 100644 --- a/ingestion/tests/unit/source/database/test_snowflake_progress_path.py +++ b/ingestion/tests/unit/source/database/test_snowflake_progress_path.py @@ -14,9 +14,10 @@ from types import SimpleNamespace from unittest.mock import MagicMock +from metadata.ingestion.progress.modes import ProgressMode +from metadata.ingestion.progress.registry import ProgressRegistry from metadata.ingestion.source.database import database_service from metadata.ingestion.source.database.snowflake import metadata as snowflake_metadata -from metadata.ingestion.progress.registry import ProgressRegistry class _ConcreteSource(database_service.DatabaseServiceSource): @@ -32,7 +33,7 @@ def _path_for(database, schema, entity_type_name): ctx = SimpleNamespace(database=database, database_schema=schema) source.context = MagicMock() source.context.get.return_value = ctx - return source.current_progress_path(entity_type_name) + return source.progress_tracker.current_path(entity_type_name) class TestSnowflakeProgressPath: @@ -65,14 +66,14 @@ def test_connector_references_no_count_constants(self): class TestProgressOptIn: - def test_database_family_is_off_by_default(self): - assert database_service.DatabaseServiceSource.progress_tracking_enabled is False + def test_database_family_is_auto_by_default(self): + assert database_service.DatabaseServiceSource.progress_mode is ProgressMode.AUTO - def test_concrete_db_source_inherits_off(self): - assert _ConcreteSource.progress_tracking_enabled is False + def test_concrete_db_source_inherits_auto(self): + assert _ConcreteSource.progress_mode is ProgressMode.AUTO - def test_snowflake_is_opted_in(self): - assert snowflake_metadata.SnowflakeSource.progress_tracking_enabled is True + def test_snowflake_is_auto(self): + assert snowflake_metadata.SnowflakeSource.progress_mode is ProgressMode.AUTO class TestStaleContextDoesNotStrandDatabases: @@ -83,7 +84,7 @@ class TestStaleContextDoesNotStrandDatabases: def _path(self, source, ctx, entity_type_name): source.context.get.return_value = ctx - return source.current_progress_path(entity_type_name) + return source.progress_tracker.current_path(entity_type_name) def test_second_database_completes_and_prunes(self): source = object.__new__(_ConcreteSource) diff --git a/ingestion/tests/unit/topology/test_topology_progress_path.py b/ingestion/tests/unit/topology/test_topology_progress_path.py index 725070e769db..aed36e46d14f 100644 --- a/ingestion/tests/unit/topology/test_topology_progress_path.py +++ b/ingestion/tests/unit/topology/test_topology_progress_path.py @@ -19,6 +19,9 @@ from unittest.mock import MagicMock from metadata.ingestion.api.topology_runner import TopologyRunnerMixin +from metadata.ingestion.models.topology import get_topology_node +from metadata.ingestion.progress.modes import ProgressMode +from metadata.ingestion.progress.runner_tracker import NO_OP_NODE_PROGRESS from metadata.ingestion.source.database.database_service import ( DatabaseServiceTopology, ) @@ -37,58 +40,63 @@ def __init__(self, database=None, database_schema=None): class TestGenericProgressPath: def test_database_node_opens_at_root(self): - assert _DBLikeRunner(database="db").current_progress_path("Database") == [] + assert _DBLikeRunner(database="db").progress_tracker.current_path("Database") == [] def test_schema_node_opens_under_its_database(self): runner = _DBLikeRunner(database="db") - assert runner.current_progress_path("DatabaseSchema") == ["db"] + assert runner.progress_tracker.current_path("DatabaseSchema") == ["db"] def test_table_node_opens_under_database_and_schema(self): runner = _DBLikeRunner(database="db", database_schema="sales") - assert runner.current_progress_path("Table") == ["db", "sales"] + assert runner.progress_tracker.current_path("Table") == ["db", "sales"] def test_stored_procedure_opens_under_database_and_schema(self): runner = _DBLikeRunner(database="db", database_schema="sales") - assert runner.current_progress_path("StoredProcedure") == ["db", "sales"] + assert runner.progress_tracker.current_path("StoredProcedure") == ["db", "sales"] def test_schema_node_ignores_stale_sibling_schema_context(self): runner = _DBLikeRunner(database="db2", database_schema="stale_from_db1") - assert runner.current_progress_path("DatabaseSchema") == ["db2"] + assert runner.progress_tracker.current_path("DatabaseSchema") == ["db2"] def test_unknown_entity_type_yields_empty_path(self): - assert _DBLikeRunner(database="db").current_progress_path("Nope") == [] + assert _DBLikeRunner(database="db").progress_tracker.current_path("Nope") == [] def test_none_entity_type_does_not_touch_topology_or_context(self): # The default surface test: no entity type means no path, and the # method must not require topology/context to be set. - assert TopologyRunnerMixin().current_progress_path() == [] + assert TopologyRunnerMixin().progress_tracker.current_path() == [] def test_root_context_keys_derived_from_topology(self): - assert _DBLikeRunner().current_progress_path # smoke: method resolves - assert _DBLikeRunner()._root_context_keys() == {"database_service"} + assert _DBLikeRunner().progress_tracker.current_path # smoke: method resolves + assert _DBLikeRunner().progress_tracker._root_context_keys() == {"database_service"} class TestShouldTrackProgress: - def _runner(self, enabled): - runner = TopologyRunnerMixin() - runner.progress_tracking_enabled = enabled - runner.__dict__["_root_node_ids"] = set() + def _runner(self, mode): + runner = _DBLikeRunner() + runner.progress_mode = mode return runner def test_disabled_never_tracks(self): - assert self._runner(False)._should_track_progress(object(), "Table") is False + runner = self._runner(ProgressMode.OFF) + table_node = get_topology_node("table", runner.topology) + assert runner.progress_tracker.for_node(table_node, is_leaf=True) is NO_OP_NODE_PROGRESS def test_enabled_real_entity_tracks(self): - assert self._runner(True)._should_track_progress(object(), "Table") is True + runner = self._runner(ProgressMode.AUTO) + table_node = get_topology_node("table", runner.topology) + assert runner.progress_tracker.for_node(table_node, is_leaf=True) is not NO_OP_NODE_PROGRESS def test_enabled_but_no_entity_type_does_not_track(self): - assert self._runner(True)._should_track_progress(object(), None) is False + runner = self._runner(ProgressMode.AUTO) + node = SimpleNamespace(stages=[]) + assert runner.progress_tracker.for_node(node, is_leaf=True) is NO_OP_NODE_PROGRESS def test_enabled_root_node_does_not_track(self): - runner = self._runner(True) - root = object() - runner.__dict__["_root_node_ids"] = {id(root)} - assert runner._should_track_progress(root, "Database") is False + runner = self._runner(ProgressMode.AUTO) + root_node = get_topology_node("root", runner.topology) + runner.progress_tracker.on_walk_start([root_node]) + assert runner.progress_tracker.for_node(root_node, is_leaf=False) is NO_OP_NODE_PROGRESS - def test_default_flag_is_false(self): - assert TopologyRunnerMixin.progress_tracking_enabled is False + def test_default_mode_is_auto(self): + assert TopologyRunnerMixin.progress_mode is ProgressMode.AUTO diff --git a/ingestion/tests/unit/topology/test_topology_runner_progress.py b/ingestion/tests/unit/topology/test_topology_runner_progress.py index 112c47940a54..dd23b7170da7 100644 --- a/ingestion/tests/unit/topology/test_topology_runner_progress.py +++ b/ingestion/tests/unit/topology/test_topology_runner_progress.py @@ -21,10 +21,11 @@ TopologyContextManager, get_topology_node, ) +from metadata.ingestion.progress.modes import ProgressMode +from metadata.ingestion.progress.registry import ProgressRegistry from metadata.ingestion.source.database.database_service import ( DatabaseServiceTopology, ) -from metadata.ingestion.progress.registry import ProgressRegistry class TestRunnerProgressSurface: @@ -38,7 +39,7 @@ def test_distinct_instances_distinct_registries(self): assert TopologyRunnerMixin().progress is not TopologyRunnerMixin().progress def test_current_progress_path_default_is_empty(self): - assert TopologyRunnerMixin().current_progress_path() == [] + assert TopologyRunnerMixin().progress_tracker.current_path() == [] def test_declare_totals_is_gone(self): assert not hasattr(TopologyRunnerMixin, "declare_totals") @@ -50,13 +51,14 @@ def test_no_count_pass_in_iter(self): source = inspect.getsource(topology_runner) assert "declare_totals" not in source - def test_root_node_detection_uses_iter_captured_ids(self): - runner = TopologyRunnerMixin() - root, child = object(), object() - assert runner._is_root_node(root) is False # nothing captured yet - runner.__dict__["_root_node_ids"] = {id(root)} - assert runner._is_root_node(root) is True - assert runner._is_root_node(child) is False + +def test_root_node_detection_uses_walk_captured_ids(): + runner = TopologyRunnerMixin() + root, child = object(), object() + assert runner.progress_tracker.is_root_node(root) is False + runner.progress_tracker.on_walk_start([root]) + assert runner.progress_tracker.is_root_node(root) is True + assert runner.progress_tracker.is_root_node(child) is False class _WalkRunner(TopologyRunnerMixin): @@ -65,10 +67,9 @@ class _WalkRunner(TopologyRunnerMixin): topology = DatabaseServiceTopology() - def __init__(self, enabled): - self.progress_tracking_enabled = enabled + def __init__(self, mode=ProgressMode.AUTO): + self.progress_mode = mode self.context = TopologyContextManager(self.topology) - self.__dict__["_root_node_ids"] = set() self.status = MagicMock() def _run_node_producer(self, node): @@ -94,25 +95,37 @@ def _drive_container(runner): def test_disabled_walk_creates_no_registry(): - runner = _WalkRunner(enabled=False) + runner = _WalkRunner(mode=ProgressMode.OFF) + _drive_leaf(runner) + assert "_progress_registry" not in runner.__dict__ + + +def test_manual_walk_creates_no_registry(): + runner = _WalkRunner(mode=ProgressMode.MANUAL) _drive_leaf(runner) assert "_progress_registry" not in runner.__dict__ def test_enabled_walk_records_into_registry(): - runner = _WalkRunner(enabled=True) + runner = _WalkRunner() _drive_leaf(runner) snapshot = runner.progress.snapshot() assert snapshot.child_type == "Table" assert snapshot.processed == 2 +def test_plain_connector_gets_progress_by_default(): + runner = _WalkRunner() + _drive_leaf(runner) + assert runner.progress.snapshot().processed == 2 + + def test_container_open_none_through_runner(): """Driving _process_node through the *database* container node must call ``progress.open(parent_path, "Database", None)`` — the lazy branch where no connector pushes a total — so the registry records ``expected=None`` for the Database level.""" - runner = _WalkRunner(enabled=True) + runner = _WalkRunner() _drive_container(runner) snapshot = runner.progress.snapshot() assert snapshot is not None @@ -123,11 +136,11 @@ def test_container_open_none_through_runner(): @pytest.fixture def progress_runner(): """A _WalkRunner with progress tracking enabled, pre-seeded with a registry.""" - return _WalkRunner(enabled=True) + return _WalkRunner() def test_container_without_push_is_unknown(progress_runner): - # progress_runner: a TopologyRunnerMixin test double with progress_tracking_enabled=True + # progress_runner: a TopologyRunnerMixin test double with progress_mode=AUTO progress_runner.progress.open(["db1"], "DatabaseSchema", None) snap = progress_runner.progress.snapshot() assert snap.children[0].expected is None @@ -143,7 +156,7 @@ def test_runner_has_no_pull_hooks(): class _CtxWalkRunner(_WalkRunner): """Like _WalkRunner but writes each entity into the node's context key, so - consumer-derived paths and _scope_path_for_node resolve to real names.""" + consumer-derived paths and scope_path_for_node resolve to real names.""" def _process_stage(self, stage, node_entity): if getattr(stage, "context", None): @@ -152,21 +165,21 @@ def _process_stage(self, stage, node_entity): def test_scope_path_for_node_appends_context_value(): - runner = _WalkRunner(enabled=True) + runner = _WalkRunner() database_node = get_topology_node("database", runner.topology) key = runner._node_primary_stage(database_node).context setattr(runner.context.get(), key, "salesdb") - assert runner._scope_path_for_node(database_node, []) == ["salesdb"] + assert runner.progress_tracker.scope_path_for_node(database_node, []) == ["salesdb"] def test_scope_path_for_node_is_none_without_context_value(): - runner = _WalkRunner(enabled=True) + runner = _WalkRunner() database_node = get_topology_node("database", runner.topology) - assert runner._scope_path_for_node(database_node, []) is None + assert runner.progress_tracker.scope_path_for_node(database_node, []) is None def test_completed_container_is_closed_through_runner(): - runner = _CtxWalkRunner(enabled=True) + runner = _CtxWalkRunner() closed: list = [] runner.progress.close = lambda path: closed.append(list(path)) container = get_topology_node("database", runner.topology) @@ -183,10 +196,9 @@ class _StatefulLeafRunner(TopologyRunnerMixin): topology = DatabaseServiceTopology() - def __init__(self, enabled): - self.progress_tracking_enabled = enabled + def __init__(self, mode=ProgressMode.AUTO): + self.progress_mode = mode self.context = TopologyContextManager(self.topology) - self.__dict__["_root_node_ids"] = set() self.status = MagicMock() self.state_live_at_stage = [] self._active = None @@ -208,20 +220,20 @@ def _process_stage(self, stage, node_entity): def test_leaf_is_lazy_when_progress_off_preserving_per_yield_state(): - runner = _StatefulLeafRunner(enabled=False) + runner = _StatefulLeafRunner(mode=ProgressMode.OFF) list(runner._process_node(get_topology_node("table", runner.topology))) assert runner.state_live_at_stage == ["a", "b"] def test_leaf_is_eager_when_progress_on_recording_count(): - runner = _StatefulLeafRunner(enabled=True) + runner = _StatefulLeafRunner() list(runner._process_node(get_topology_node("table", runner.topology))) assert runner.state_live_at_stage == [None, None] assert runner.progress.snapshot().processed == 2 def test_closing_container_tracks_global_done(): - runner = _CtxWalkRunner(enabled=True) + runner = _CtxWalkRunner() runner.progress.set_total("Database", 5) runner.progress.set_total("DatabaseSchema", 9) list(runner._process_node(get_topology_node("database", runner.topology))) @@ -231,7 +243,7 @@ def test_closing_container_tracks_global_done(): def test_reconcilable_container_reconciles_total(): - runner = _CtxWalkRunner(enabled=True) + runner = _CtxWalkRunner() runner.progress.seed_scope_total("DatabaseSchema", "a", 1) runner.progress.seed_scope_total("DatabaseSchema", "b", 1) # upfront total = 2 (1 seeded per declared database) @@ -247,8 +259,8 @@ class _MultiThreadCtxWalkRunner(_CtxWalkRunner): instead of _process_node. Post-process methods are no-ops because the runner stub does not implement the real connector hooks.""" - def __init__(self, enabled): - super().__init__(enabled) + def __init__(self, mode=ProgressMode.AUTO): + super().__init__(mode) self.context.threads = 2 def _run_node_post_process(self, node): @@ -261,7 +273,7 @@ def test_multithread_schema_node_reconciles_and_tracks_done(): reconcile_scope_total and track are exercised identically to the single-thread reconcile test: seed 1 per declared db, observe 2 per db, final counters show (done=4, total=4).""" - runner = _MultiThreadCtxWalkRunner(enabled=True) + runner = _MultiThreadCtxWalkRunner() runner.progress.seed_scope_total("DatabaseSchema", "a", 1) runner.progress.seed_scope_total("DatabaseSchema", "b", 1) # upfront total = 2; driving _process_node(database) recurses into @@ -271,3 +283,21 @@ def test_multithread_schema_node_reconciles_and_tracks_done(): counters = {t: (d, total) for t, d, total in runner.progress.global_counters()} # each database actually walks 2 schemas -> reconciled total 2*2=4; done=4 assert counters["DatabaseSchema"] == (4, 4) + + +class _TotalsWalkRunner(_WalkRunner): + def __init__(self): + super().__init__() + self.declared = 0 + + def declare_progress_totals(self, totals): + self.declared += 1 + totals.set_total("Table", 99) + + +def test_declare_progress_totals_called_once_per_walk(): + runner = _TotalsWalkRunner() + _drive_leaf(runner) + _drive_leaf(runner) + assert runner.declared == 1 + assert ("Table", 0, 99) in runner.progress.global_counters() From df7b515237a782346d15efffa97d54258bafdd4f Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 12:04:30 +0530 Subject: [PATCH 09/18] fix(ingestion): dashboard group-progress test double declares MANUAL mode The group-progress helpers on DashboardServiceSource were migrated to the manual_progress facade, which requires progress_mode=MANUAL. The _BareSource test double still used the default AUTO mode and raised ProgressModeError. Declare MANUAL on the double, matching how real PowerBI drives these helpers. Co-Authored-By: Claude Fable 5 --- .../unit/topology/dashboard/test_dashboard_group_progress.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py b/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py index 2a89a069b5dc..6cf74d7a165a 100644 --- a/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py +++ b/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py @@ -10,6 +10,7 @@ # limitations under the License. """Reusable per-group progress helpers on DashboardServiceSource.""" +from metadata.ingestion.progress.modes import ProgressMode from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource @@ -19,6 +20,8 @@ class _BareSource(DashboardServiceSource): # Overriding the abstract methods clears __abstractmethods__; __new__ then # skips the heavy __init__. The helpers and the ``progress`` property only # touch ``self.__dict__``, so no further setup is needed. + progress_mode = ProgressMode.MANUAL + def create(self, *args, **kwargs): ... def get_dashboard_details(self, *args, **kwargs): ... def get_dashboard_name(self, *args, **kwargs): ... From 266f7a2f0fca24187bc6f636c0a37b20abc62b81 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 12:11:35 +0530 Subject: [PATCH 10/18] refactor(ingestion): snowflake declares progress totals via the declare_progress_totals hook Co-Authored-By: Claude Fable 5 --- .../ingestion/source/database/snowflake/metadata.py | 13 +++++-------- .../database/test_snowflake_progress_count.py | 7 ++++--- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/ingestion/src/metadata/ingestion/source/database/snowflake/metadata.py b/ingestion/src/metadata/ingestion/source/database/snowflake/metadata.py index 8bbdc086e0cf..8677d937580a 100644 --- a/ingestion/src/metadata/ingestion/source/database/snowflake/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/snowflake/metadata.py @@ -76,6 +76,7 @@ IncrementalConfig, ) from metadata.ingestion.source.database.multi_db_source import MultiDBSource +from metadata.ingestion.progress.modes import TotalsDeclarer from metadata.ingestion.source.database.snowflake.constants import ( DEFAULT_STREAM_COLUMNS, PROCEDURE_TYPE_URL_MAP, @@ -217,8 +218,6 @@ class SnowflakeSource( service_connection: SnowflakeConnection - progress_tracking_enabled = True - def __init__( self, config, @@ -482,17 +481,17 @@ def _is_schema_filtered(self, database_name: str, schema_name: str) -> bool: filter_name = schema_fqn if self.source_config.useFqnForFiltering and schema_fqn else schema_name return filter_by_schema(self.source_config.schemaFilterPattern, filter_name) - def _declare_progress_totals(self) -> None: + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: """Seed the run-level ``Database`` and ``DatabaseSchema`` global counters upfront. ``Database`` is the filtered DB count; ``DatabaseSchema`` is the post-filter schema count per database from the account-wide SHOW. When that SHOW is unavailable, the schema counter is marked reconcilable so the walk fills its total instead.""" database_names = self._filtered_database_names() - self.progress.set_total(Database.__name__, len(database_names)) + totals.set_total(Database.__name__, len(database_names)) schemas_by_database = self._schema_names_by_database() if schemas_by_database is None: - self.progress.set_reconcilable(DatabaseSchema.__name__) + totals.mark_reconcilable(DatabaseSchema.__name__) else: for database_name in database_names: kept = [ @@ -500,11 +499,9 @@ def _declare_progress_totals(self) -> None: for schema_name in schemas_by_database.get(database_name, []) if not self._is_schema_filtered(database_name, schema_name) ] - self.progress.seed_scope_total(DatabaseSchema.__name__, database_name, len(kept)) + totals.seed_scope_total(DatabaseSchema.__name__, database_name, len(kept)) def get_database_names(self) -> Iterable[str]: - if self.progress_tracking_enabled: - self._declare_progress_totals() for database_name in self._filtered_database_names(): try: self.set_inspector(database_name=database_name) diff --git a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py index 1b8fbf7fb0de..baec5f745036 100644 --- a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py +++ b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py @@ -16,6 +16,7 @@ import pytest +from metadata.ingestion.progress.modes import TotalsDeclarer from metadata.ingestion.source.database.snowflake import metadata as snowflake_metadata SnowflakeSource = snowflake_metadata.SnowflakeSource @@ -110,7 +111,7 @@ def test_declare_progress_totals_seeds_database_and_schema(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1", "db2"] snowflake_source._schema_names_by_database = lambda: {"db1": ["s1", "s2"], "db2": ["s3"]} snowflake_source._is_schema_filtered = lambda db, sch: False - snowflake_source._declare_progress_totals() + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} assert counters["Database"] == (0, 2) assert counters["DatabaseSchema"] == (0, 3) @@ -120,7 +121,7 @@ def test_declare_progress_totals_applies_schema_filter(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1"] snowflake_source._schema_names_by_database = lambda: {"db1": ["keep", "drop"]} snowflake_source._is_schema_filtered = lambda db, sch: sch == "drop" - snowflake_source._declare_progress_totals() + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} assert counters["DatabaseSchema"] == (0, 1) @@ -128,7 +129,7 @@ def test_declare_progress_totals_applies_schema_filter(snowflake_source): def test_declare_progress_totals_falls_back_when_account_show_unavailable(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1"] snowflake_source._schema_names_by_database = lambda: None - snowflake_source._declare_progress_totals() + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) assert snowflake_source.progress.is_reconcilable("DatabaseSchema") is True counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} assert counters["Database"] == (0, 1) From 6930d181c2dd91510293c8270d71f3f0e8ee7da2 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 12:34:12 +0530 Subject: [PATCH 11/18] docs(ingestion): warn that dashboard group-progress helpers require MANUAL mode Guards a future dashboard-connector author from calling the group helpers on a default AUTO source, which would raise ProgressModeError and otherwise double-count against the runner's own tracking. Addresses the one latent trap from final review. Co-Authored-By: Claude Fable 5 --- .../ingestion/source/dashboard/dashboard_service.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py index 21021f4a419d..09ba665b60e5 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py @@ -216,7 +216,12 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC): chart_source_state: Set = set() # noqa: RUF012, UP006 def _declare_progress_groups(self, label: str, total: Optional[int]) -> None: # noqa: UP045 - """Declare the grouping axis (e.g. workspaces) as a global counter.""" + """Declare the grouping axis (e.g. workspaces) as a global counter. + + These group helpers drive ``self.manual_progress`` and therefore require + the connector to set ``progress_mode = ProgressMode.MANUAL`` (as PowerBI + does); calling them from a default AUTO source raises ``ProgressModeError`` + and would double-count against the runner's own tracking.""" self.manual_progress.declare_groups(label, total) def _open_group_progress(self, group: str, expected_by_type: Dict[str, Optional[int]]) -> None: # noqa: UP006, UP045 From 8f9566989ec1cc7074c1cd26a20d5cf9de2aad91 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 15:06:46 +0530 Subject: [PATCH 12/18] chore(ingestion): restore Dockerfile.ci to base (drop local dev extras) Dockerfile.ci was accidentally swept into this branch with a local build hack that pinned pip extras to [powerbi,snowflake,pii-processor]. Restore the canonical ARG-driven install so this file carries no branch-local change. Co-Authored-By: Claude Fable 5 --- ingestion/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ingestion/Dockerfile.ci b/ingestion/Dockerfile.ci index 41af777ff170..ec763025fc20 100644 --- a/ingestion/Dockerfile.ci +++ b/ingestion/Dockerfile.ci @@ -127,7 +127,7 @@ RUN python /home/airflow/scripts/datamodel_generation.py # Argument to provide for Ingestion Dependencies to install. Defaults to all ARG INGESTION_DEPENDENCY="all" -RUN pip install ".[powerbi,snowflake,pii-processor]" +RUN pip install ".[${INGESTION_DEPENDENCY}]" # Temporary workaround for https://github.com/open-metadata/OpenMetadata/issues/9593 RUN [ $(uname -m) = "x86_64" ] \ From 38371900171e86c1cb8fc65358b13cdbccd50ba5 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Mon, 6 Jul 2026 16:43:25 +0530 Subject: [PATCH 13/18] feat(ingestion): mysql declares progress totals via the declare_progress_totals hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MySQL already inherits AUTO processed-counting from the topology runner; add the Tier-1 declare_progress_totals hook so the run renders %/ETA. MySQL is effectively single-database, so declare the Database denominator directly and mark DatabaseSchema reconcilable — the runner fills the schema total from the live inspector it queries anyway, and DatabaseSchema drives the ETA. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../source/database/mysql/metadata.py | 11 ++++ .../database/test_mysql_progress_count.py | 52 +++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 ingestion/tests/unit/source/database/test_mysql_progress_count.py diff --git a/ingestion/src/metadata/ingestion/source/database/mysql/metadata.py b/ingestion/src/metadata/ingestion/source/database/mysql/metadata.py index cfc2eaef3289..467ac9e2983b 100644 --- a/ingestion/src/metadata/ingestion/source/database/mysql/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/mysql/metadata.py @@ -21,6 +21,7 @@ from metadata.generated.schema.api.data.createStoredProcedure import ( CreateStoredProcedureRequest, ) +from metadata.generated.schema.entity.data.database import Database from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode from metadata.generated.schema.entity.services.connections.database.mysqlConnection import ( @@ -36,6 +37,7 @@ from metadata.ingestion.api.models import Either from metadata.ingestion.api.steps import InvalidSourceException from metadata.ingestion.ometa.ometa_api import OpenMetadata +from metadata.ingestion.progress.modes import TotalsDeclarer from metadata.ingestion.source.database.common_db_source import CommonDbSourceService from metadata.ingestion.source.database.mysql.models import ( DEFAULT_STORED_PROC_LANGUAGE, @@ -76,6 +78,15 @@ def create(cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str raise InvalidSourceException(f"Expected MysqlConnection, but got {connection}") return cls(config, metadata) + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: + """Seed the run-level ``Database`` counter with MySQL's single database and + let the runner reconcile the ``DatabaseSchema`` total from the schemas it + observes. MySQL exposes schemas only through the live inspector, which the + walk queries anyway, so pre-counting them here would just duplicate that + query for no gain.""" + totals.set_total(Database.__name__, len(list(self.get_database_names()))) + totals.mark_reconcilable(DatabaseSchema.__name__) + def get_stored_procedures(self) -> Iterable[MysqlRoutine]: """List stored procedures and functions""" if self.source_config.includeStoredProcedures: diff --git a/ingestion/tests/unit/source/database/test_mysql_progress_count.py b/ingestion/tests/unit/source/database/test_mysql_progress_count.py new file mode 100644 index 000000000000..d68ba0975ea4 --- /dev/null +++ b/ingestion/tests/unit/source/database/test_mysql_progress_count.py @@ -0,0 +1,52 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MySQL declares the single-database ``Database`` denominator and defers the +``DatabaseSchema`` total to runner reconciliation.""" + +from types import SimpleNamespace + +import pytest + +from metadata.ingestion.progress.modes import TotalsDeclarer +from metadata.ingestion.progress.registry import ProgressRegistry +from metadata.ingestion.source.database.mysql.metadata import MysqlSource + + +def _source(service_connection): + source = object.__new__(MysqlSource) + source.service_connection = service_connection + source.__dict__["_progress_registry"] = ProgressRegistry() + return source + + +@pytest.fixture +def mysql_source(): + return _source(SimpleNamespace(databaseName=None, database="mydb")) + + +def test_declare_progress_totals_seeds_single_database(mysql_source): + mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress)) + counters = {t: (done, total) for t, done, total in mysql_source.progress.global_counters()} + assert counters["Database"] == (0, 1) + + +def test_declare_progress_totals_marks_schema_reconcilable(mysql_source): + mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress)) + assert mysql_source.progress.is_reconcilable("DatabaseSchema") is True + counters = {t: (done, total) for t, done, total in mysql_source.progress.global_counters()} + assert counters["DatabaseSchema"] == (0, None) + + +def test_declare_progress_totals_defaults_database_name_when_unset(mysql_source): + source = _source(SimpleNamespace(databaseName=None, database=None)) + source.declare_progress_totals(TotalsDeclarer(source.progress)) + counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + assert counters["Database"] == (0, 1) From 6d490bd138b8376de3f8ce4e0e25166f0566f9d1 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Tue, 7 Jul 2026 14:45:35 +0530 Subject: [PATCH 14/18] fix(ingestion): iterate leaf producers lazily under AUTO progress AUTO progress made every leaf node eager-count its producer, list()-draining it before any entity was sunk. Connectors whose producer reads context a stage of the just-yielded entity populates (e.g. S3 storage's get_containers reading context.container) then saw None, crashing with NoneType fqn / container.id errors and failing all S3 integration tests in shard-2. Leaf producers now always iterate lazily: their per-leaf advance already yields an accurate processed count, so eager counting is limited to reconcilable container counters. Only tradeoff is leaf bars show N/? instead of N/total. Adds a regression test mirroring the S3 pattern (leaf producer reads context written by the just-yielded entity's stage) and flips the two unit tests that had encoded the eager-drain behavior. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../metadata/ingestion/api/topology_runner.py | 17 ++++---- .../ingestion/progress/runner_tracker.py | 12 ++++-- .../unit/progress/test_runner_tracker.py | 4 +- .../topology/test_topology_runner_progress.py | 42 ++++++++++++++++++- 4 files changed, 60 insertions(+), 15 deletions(-) diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index c3113a74f560..3eb2156626ee 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -185,14 +185,15 @@ def _multithread_process_node(self, node: TopologyNode, threads: int) -> Iterabl def _process_node(self, node: TopologyNode) -> Iterable[Entity]: """Single-threaded processing of a Node. - Container producers (database, schema) are iterated lazily so connectors - can set up per-yield inspector/session state before each child is - processed. Leaf producers (table, stored procedure) are materialized - eagerly only when the node's progress handle wants an eager count — to - record an exact child count via ``progress.open``; otherwise they are - iterated lazily so connectors whose stages depend on per-yield producer - state (e.g. PowerBI's per-workspace ``state.enter`` / ``finally: - state.exit``) are not torn down before their stages run. + Producers are iterated lazily by default so connectors can set up + per-yield inspector/session state — and read context a stage of the + just-yielded entity populated (e.g. S3 storage's ``get_containers``) — + before the next child is produced. A producer is materialized eagerly + only when its node's progress handle wants an exact upfront child count, + which today is limited to reconcilable container counters (schema totals + nudged toward the observed count). Leaf producers are never drained + eagerly: their per-leaf ``advance`` already yields an accurate processed + count. """ child_nodes = self._get_child_nodes(node) node_progress = self.progress_tracker.for_node(node, is_leaf=not child_nodes) diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py index 4d232e4bfcc8..25b37908b6c4 100644 --- a/ingestion/src/metadata/ingestion/progress/runner_tracker.py +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -100,9 +100,15 @@ def __init__( @property def wants_eager_count(self) -> bool: - """Materialize the producer for an exact child count: always for a - leaf; for a container only when its counter is reconcilable.""" - return self._is_leaf or self._reconcilable + """Materialize the producer for an exact child count only when the node's + counter is reconcilable (a container whose scope total is nudged toward + the observed child count). Leaf producers are always iterated lazily: + their per-leaf ``advance`` already yields an accurate processed count, and + an eager ``list()``-drain would run the producer to completion before any + stage sinks an entity — breaking connectors (e.g. S3 storage's + ``get_containers``) whose producer reads context a stage of the + just-yielded entity populated.""" + return self._reconcilable def open_with_count(self, count: int) -> None: self._registry.open(self._parent_path, self._entity_type_name, count) diff --git a/ingestion/tests/unit/progress/test_runner_tracker.py b/ingestion/tests/unit/progress/test_runner_tracker.py index 68425e66f070..3e1c23f80dc8 100644 --- a/ingestion/tests/unit/progress/test_runner_tracker.py +++ b/ingestion/tests/unit/progress/test_runner_tracker.py @@ -62,12 +62,12 @@ def test_for_node_returns_noop_for_root_node(): assert tracker.for_node(root, is_leaf=False) is NO_OP_NODE_PROGRESS -def test_for_node_returns_real_handle_for_auto_leaf(): +def test_for_node_returns_lazy_handle_for_auto_leaf(): source = _FakeSource() tracker = TopologyProgressTracker(source) handle = tracker.for_node(_table_node(source), is_leaf=True) assert isinstance(handle, NodeProgress) - assert handle.wants_eager_count is True + assert handle.wants_eager_count is False def test_totals_hook_called_exactly_once_and_only_for_real_nodes(): diff --git a/ingestion/tests/unit/topology/test_topology_runner_progress.py b/ingestion/tests/unit/topology/test_topology_runner_progress.py index dd23b7170da7..12133f82100c 100644 --- a/ingestion/tests/unit/topology/test_topology_runner_progress.py +++ b/ingestion/tests/unit/topology/test_topology_runner_progress.py @@ -225,13 +225,51 @@ def test_leaf_is_lazy_when_progress_off_preserving_per_yield_state(): assert runner.state_live_at_stage == ["a", "b"] -def test_leaf_is_eager_when_progress_on_recording_count(): +def test_leaf_is_lazy_when_progress_on_preserving_per_yield_state(): + """AUTO-mode leaf producers must still iterate lazily: per-yield producer + state stays live when each stage runs (an eager list()-drain would tear it + down first), while the per-leaf advance keeps the processed count accurate.""" runner = _StatefulLeafRunner() list(runner._process_node(get_topology_node("table", runner.topology))) - assert runner.state_live_at_stage == [None, None] + assert runner.state_live_at_stage == ["a", "b"] assert runner.progress.snapshot().processed == 2 +class _CtxReadingLeafRunner(TopologyRunnerMixin): + """Mirrors S3 ``get_containers``: the leaf producer reads a context value that + a *stage* of the just-yielded entity populates. Correct only under lazy + iteration — an eager list()-drain runs the producer to completion before any + stage writes the context, so the producer reads ``None``.""" + + topology = DatabaseServiceTopology() + + def __init__(self, mode=ProgressMode.AUTO): + self.progress_mode = mode + self.context = TopologyContextManager(self.topology) + self.status = MagicMock() + self.observed_context = [] + + def _run_node_producer(self, node): + for name in ["t1", "t2"]: + yield name + self.observed_context.append(getattr(self.context.get(), "table", None)) + + def _process_stage(self, stage, node_entity): + if getattr(stage, "context", None): + setattr(self.context.get(), stage.context, node_entity) + return iter(()) + + +def test_leaf_producer_reads_context_written_by_prior_yield_stage(): + """Regression for the S3 ``get_containers`` crash (NoneType fqn / + ``container.id``): a leaf producer that reads context populated by the + just-yielded entity's stage must observe it. An eager list()-drain leaves the + context unset, so the producer would read ``None`` for every entity.""" + runner = _CtxReadingLeafRunner() + list(runner._process_node(get_topology_node("table", runner.topology))) + assert runner.observed_context == ["t1", "t2"] + + def test_closing_container_tracks_global_done(): runner = _CtxWalkRunner() runner.progress.set_total("Database", 5) From 8560e7e3d240e7feb348b16f4fc1ce0b54ff419a Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Tue, 7 Jul 2026 23:41:55 +0530 Subject: [PATCH 15/18] refactor(ingestion): apply progress review feedback (facade dedup, one open(), with-scopes) - ManualProgress extends TotalsDeclarer instead of re-implementing its denominator methods; AUTO hooks still receive a plain TotalsDeclarer, so the denominator-only guarantee is unchanged. - NodeProgress.open(count) replaces the open_with_count/open_lazy pair; reconciliation fires only when an exact count exists, which is the answer to why only one of them reconciled. - Container scopes are context managers and the runner uses with-blocks: an exception or early generator close now prunes the subtree without counting the container complete (previously the scope leaked in the tree on failure). Co-Authored-By: Claude Fable 5 --- .../metadata/ingestion/api/topology_runner.py | 18 +++++----- .../src/metadata/ingestion/progress/modes.py | 14 +++----- .../ingestion/progress/runner_tracker.py | 35 +++++++++++-------- .../unit/progress/test_runner_tracker.py | 33 +++++++++++++---- 4 files changed, 60 insertions(+), 40 deletions(-) diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index 3eb2156626ee..8f41034ea69a 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -144,7 +144,7 @@ def _multithread_process_node(self, node: TopologyNode, threads: int) -> Iterabl entity_type=entity_type_name, ) - node_progress.open_with_count(node_entities_length) + node_progress.open(node_entities_length) if node_entities_length == 0: return @@ -200,10 +200,10 @@ def _process_node(self, node: TopologyNode) -> Iterable[Entity]: if node_progress.wants_eager_count: node_entities = list(self._run_node_producer(node) or []) - node_progress.open_with_count(len(node_entities)) + node_progress.open(len(node_entities)) else: node_entities = self._run_node_producer(node) or [] - node_progress.open_lazy() + node_progress.open(None) for node_entity in node_entities: for stage in node.stages: @@ -215,9 +215,8 @@ def _process_node(self, node: TopologyNode) -> Iterable[Entity]: node_progress.advance_leaf() - scope = node_progress.enter_scope() - yield from self.process_nodes(child_nodes) - scope.exit() + with node_progress.enter_scope(): + yield from self.process_nodes(child_nodes) def process_nodes(self, nodes: List[TopologyNode]) -> Iterable[Entity]: # noqa: UP006 """ @@ -285,10 +284,9 @@ def _multithread_process_entity( node_progress.advance_leaf() - scope = node_progress.enter_scope() - for child_result in self.process_nodes(child_nodes): - self.queue.put(child_result) - scope.exit() + with node_progress.enter_scope(): + for child_result in self.process_nodes(child_nodes): + self.queue.put(child_result) # Merge thread-local metrics into global state before thread exits operation_metrics.merge_thread_metrics() diff --git a/ingestion/src/metadata/ingestion/progress/modes.py b/ingestion/src/metadata/ingestion/progress/modes.py index d24a1301e68f..8f299fc5fe06 100644 --- a/ingestion/src/metadata/ingestion/progress/modes.py +++ b/ingestion/src/metadata/ingestion/progress/modes.py @@ -54,24 +54,20 @@ def mark_reconcilable(self, entity_type: str) -> None: self._registry.set_reconcilable(entity_type) -class ManualProgress: +class ManualProgress(TotalsDeclarer): """Counting facade for ``ProgressMode.MANUAL`` sources — connectors whose enumeration the topology runner cannot count (grouped dashboard walks, - custom ``_iter`` query sources).""" + custom ``_iter`` query sources). Extends the denominator-only declarer + with the counting side; AUTO hooks still receive a plain + ``TotalsDeclarer``, never this subclass.""" def __init__(self, registry: "ProgressRegistry") -> None: - self._registry = registry + super().__init__(registry) self._group_label: Optional[str] = None # noqa: UP045 - def set_total(self, entity_type: str, total: Optional[int]) -> None: # noqa: UP045 - self._registry.set_total(entity_type, total) - def track(self, entity_type: Optional[str], n: int = 1) -> None: # noqa: UP045 self._registry.track(entity_type, n) - def seed_scope_total(self, entity_type: str, scope: str, n: int) -> None: - self._registry.seed_scope_total(entity_type, scope, n) - def reconcile_scope_total(self, entity_type: Optional[str], scope: str, observed: int) -> None: # noqa: UP045 self._registry.reconcile_scope_total(entity_type, scope, observed) diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py index 25b37908b6c4..13d031d5c1e3 100644 --- a/ingestion/src/metadata/ingestion/progress/runner_tracker.py +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -34,7 +34,10 @@ class _NoOpScope: """No container scope to retire.""" - def exit(self) -> None: + def __enter__(self) -> "_NoOpScope": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: """Nothing to close.""" @@ -42,17 +45,23 @@ def exit(self) -> None: class _Scope: - """An open container scope: exiting prunes the subtree from the tree and - counts the finished container on its global counter.""" + """An open container scope: exiting prunes the subtree from the tree. + The finished container is counted on its global counter only on a clean + exit — a failure (or early generator close) mid-subtree prunes without + claiming the container completed.""" def __init__(self, registry: "ProgressRegistry", path: List[str], entity_type_name: str) -> None: # noqa: UP006 self._registry = registry self._path = path self._entity_type_name = entity_type_name - def exit(self) -> None: + def __enter__(self) -> "_Scope": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: self._registry.close(self._path) - self._registry.track(self._entity_type_name) + if exc_type is None: + self._registry.track(self._entity_type_name) class _NoOpNodeProgress: @@ -60,10 +69,7 @@ class _NoOpNodeProgress: wants_eager_count = False - def open_with_count(self, count: int) -> None: - """No-op.""" - - def open_lazy(self) -> None: + def open(self, count: Optional[int]) -> None: # noqa: UP045 """No-op.""" def advance_leaf(self) -> None: @@ -110,14 +116,15 @@ def wants_eager_count(self) -> bool: just-yielded entity populated.""" return self._reconcilable - def open_with_count(self, count: int) -> None: + def open(self, count: Optional[int]) -> None: # noqa: UP045 + """Open this node's counter under its parent scope. An exact ``count`` + (the runner materialized the producer) additionally reconciles a + reconcilable container's scope total toward the observed child count; + a lazy open (``None``) has nothing to reconcile with.""" self._registry.open(self._parent_path, self._entity_type_name, count) - if self._reconcilable and self._parent_path: + if count is not None and self._reconcilable and self._parent_path: self._registry.reconcile_scope_total(self._entity_type_name, self._parent_path[-1], count) - def open_lazy(self) -> None: - self._registry.open(self._parent_path, self._entity_type_name, None) - def advance_leaf(self) -> None: if self._is_leaf: self._registry.advance(self._parent_path, self._entity_type_name) diff --git a/ingestion/tests/unit/progress/test_runner_tracker.py b/ingestion/tests/unit/progress/test_runner_tracker.py index 3e1c23f80dc8..d6d53dc0b499 100644 --- a/ingestion/tests/unit/progress/test_runner_tracker.py +++ b/ingestion/tests/unit/progress/test_runner_tracker.py @@ -12,6 +12,8 @@ from unittest.mock import MagicMock +import pytest + from metadata.ingestion.api.topology_runner import TopologyRunnerMixin from metadata.ingestion.models.topology import TopologyContextManager, get_topology_node, get_topology_root from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer @@ -87,7 +89,7 @@ def test_leaf_handle_counts_open_and_advance(): source = _FakeSource() tracker = TopologyProgressTracker(source) handle = tracker.for_node(_table_node(source), is_leaf=True) - handle.open_with_count(2) + handle.open(2) handle.advance_leaf() handle.advance_leaf() snapshot = source.progress.snapshot() @@ -100,7 +102,7 @@ def test_container_handle_is_lazy_without_reconcilable_counter(): tracker = TopologyProgressTracker(source) handle = tracker.for_node(_database_node(source), is_leaf=False) assert handle.wants_eager_count is False - handle.open_lazy() + handle.open(None) assert source.progress.snapshot().expected_by_type.get("Database") is None @@ -113,7 +115,7 @@ def test_container_handle_reconciles_when_counter_is_reconcilable(): schema_node = get_topology_node("databaseSchema", source.topology) handle = tracker.for_node(schema_node, is_leaf=False) assert handle.wants_eager_count is True - handle.open_with_count(3) + handle.open(3) counters = {t: (done, total) for t, done, total in source.progress.global_counters()} assert counters["DatabaseSchema"] == (0, 3) @@ -128,19 +130,36 @@ def test_enter_scope_closes_and_tracks_on_exit(): setattr(source.context.get(), key, "salesdb") closed = [] source.progress.close = lambda path: closed.append(list(path)) - scope = handle.enter_scope() - scope.exit() + with handle.enter_scope(): + pass assert closed == [["salesdb"]] counters = {t: (done, total) for t, done, total in source.progress.global_counters()} assert counters["Database"][0] == 1 +def test_enter_scope_prunes_without_counting_on_failure(): + source = _FakeSource() + source.progress.set_total("Database", 7) + tracker = TopologyProgressTracker(source) + node = _database_node(source) + handle = tracker.for_node(node, is_leaf=False) + key = source._node_primary_stage(node).context + setattr(source.context.get(), key, "salesdb") + closed = [] + source.progress.close = lambda path: closed.append(list(path)) + with pytest.raises(RuntimeError), handle.enter_scope(): + raise RuntimeError("boom") + assert closed == [["salesdb"]] + counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + assert counters["Database"][0] == 0 + + def test_enter_scope_is_noop_without_context_value(): source = _FakeSource() tracker = TopologyProgressTracker(source) handle = tracker.for_node(_database_node(source), is_leaf=False) - scope = handle.enter_scope() - scope.exit() + with handle.enter_scope(): + pass assert source.progress.snapshot() is None # nothing opened, nothing closed From e55834a7ccf8f0245ca00282d44f7acbca02666f Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Wed, 8 Jul 2026 13:37:14 +0530 Subject: [PATCH 16/18] refactor(ingestion): compose progress tracking instead of a mixin Address review feedback to prefer composition over the ProgressTrackingMixin. - Replace ProgressTrackingMixin with a plain ProgressTracking object that owns the registry, the declared mode, and the guarded manual facade. A source now *has a* ProgressTracking (via the progress_tracking property its base defines) instead of *being* a progress mixin. - attach_progress_tracking(source) builds it lazily and caches it under _progress_tracking, preserving the 'no empty registry on untracked steps' discovery contract the workflow reporter relies on. - TopologyRunnerMixin and QueryParserSource drop the mixin base and expose a two-line progress_tracking property; the runner's declare_progress_totals hook default stays on TopologyRunnerMixin. - Production counting routes through self.progress_tracking.manual / .registry, making the composition explicit at the call sites the reviewer flagged as implicit attributes. - Reporter reads _progress_tracking.registry; behavior otherwise unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../metadata/ingestion/api/topology_runner.py | 24 +++++- .../metadata/ingestion/progress/__init__.py | 8 +- .../ingestion/progress/runner_tracker.py | 2 +- .../metadata/ingestion/progress/tracking.py | 76 +++++++++++-------- .../source/dashboard/dashboard_service.py | 10 +-- .../source/dashboard/powerbi/metadata.py | 4 +- .../source/database/lineage_source.py | 6 +- .../source/database/query_parser_source.py | 13 +++- .../source/database/snowflake/lineage.py | 6 +- .../ingestion/source/database/usage_source.py | 6 +- .../workflow/workflow_status_mixin.py | 6 +- .../unit/lineage/test_lineage_progress.py | 2 +- .../test_snowflake_access_history_progress.py | 4 +- .../unit/progress/test_progress_modes.py | 47 +++++------- .../unit/progress/test_runner_tracker.py | 26 +++---- .../database/test_mysql_progress_count.py | 20 ++--- .../database/test_snowflake_progress_count.py | 24 +++--- .../tests/unit/test_progress_tracking.py | 40 ++++++++++ .../unit/test_progress_tracking_mixin.py | 31 -------- ingestion/tests/unit/test_usage_progress.py | 4 +- .../test_dashboard_group_progress.py | 8 +- .../unit/topology/dashboard/test_powerbi.py | 24 +++--- .../topology/test_topology_runner_progress.py | 44 +++++------ .../workflow/test_status_mixin_progress.py | 2 +- 24 files changed, 240 insertions(+), 197 deletions(-) create mode 100644 ingestion/tests/unit/test_progress_tracking.py delete mode 100644 ingestion/tests/unit/test_progress_tracking_mixin.py diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index 8f41034ea69a..5911371572cd 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -42,8 +42,12 @@ ) from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.ometa.utils import model_str +from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer from metadata.ingestion.progress.runner_tracker import TopologyProgressTracker -from metadata.ingestion.progress.tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ( + ProgressTracking, + attach_progress_tracking, +) from metadata.utils.custom_thread_pool import CustomThreadPoolExecutor from metadata.utils.logger import ingestion_logger from metadata.utils.operation_metrics import OperationMetricsState @@ -61,7 +65,7 @@ class MissingExpectedEntityAckException(Exception): # noqa: N818 """ -class TopologyRunnerMixin(ProgressTrackingMixin, Generic[C]): +class TopologyRunnerMixin(Generic[C]): """ Prepares the _run function dynamically based on the source topology @@ -73,12 +77,28 @@ class TopologyRunnerMixin(ProgressTrackingMixin, Generic[C]): queue = Queue() + progress_mode: ClassVar[ProgressMode] = ProgressMode.AUTO + """AUTO (default): the topology runner counts processed entities. + MANUAL: the runner makes zero progress calls; the source drives + ``progress_tracking.manual`` itself. OFF: no progress at all.""" + _SIDE_OUTPUT_STAGE_TYPES: ClassVar[set[str]] = { "OMetaTagAndClassification", "OMetaLifeCycleData", "AddLineageRequest", } + @property + def progress_tracking(self) -> ProgressTracking: + """Composed per-source progress state (registry, mode, manual facade). + Lazy: a source that never tracks progress never builds a registry.""" + return attach_progress_tracking(self) + + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: + """Optional connector hook: declare denominators (totals) so the run + renders % and ETA. Called exactly once by the topology runner, just + before the first non-root node is processed. Default: no totals.""" + def _node_primary_stage(self, node: TopologyNode) -> Optional[NodeStage]: # noqa: UP045 """The node's primary, non-side-output stage — the stage whose entity is the node's real entity (Table, Database, ...). Falls back to the first diff --git a/ingestion/src/metadata/ingestion/progress/__init__.py b/ingestion/src/metadata/ingestion/progress/__init__.py index 986ef03ff1c5..790bdb828e75 100644 --- a/ingestion/src/metadata/ingestion/progress/__init__.py +++ b/ingestion/src/metadata/ingestion/progress/__init__.py @@ -43,7 +43,10 @@ NodeProgress, TopologyProgressTracker, ) -from metadata.ingestion.progress.tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ( + ProgressTracking, + attach_progress_tracking, +) __all__ = [ "DEFAULT_ACTIVE_LEAF_CAP", @@ -56,9 +59,10 @@ "ProgressNode", "ProgressNodeSnapshot", "ProgressRegistry", - "ProgressTrackingMixin", + "ProgressTracking", "TopologyProgressTracker", "TotalsDeclarer", + "attach_progress_tracking", "format_eta", "render_progress_tree", "snapshot_to_progress_payload", diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py index 13d031d5c1e3..641493855eb5 100644 --- a/ingestion/src/metadata/ingestion/progress/runner_tracker.py +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -157,7 +157,7 @@ def __init__(self, source: Any) -> None: @property def registry(self) -> "ProgressRegistry": - return self._source.progress + return self._source.progress_tracking.registry def on_walk_start(self, root_nodes) -> None: self._root_node_ids = {id(node) for node in root_nodes} diff --git a/ingestion/src/metadata/ingestion/progress/tracking.py b/ingestion/src/metadata/ingestion/progress/tracking.py index 1b90ae4391f6..c41b822efe86 100644 --- a/ingestion/src/metadata/ingestion/progress/tracking.py +++ b/ingestion/src/metadata/ingestion/progress/tracking.py @@ -8,57 +8,67 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Shared mixin that lazily owns a per-source ProgressRegistry. +"""Per-source progress state, composed into a source rather than mixed in. -Any Source that mixes this in gets a ``progress`` registry discoverable by the -workflow reporter (which scans steps for the ``_progress_registry`` attribute), -without participating in the topology runner. +A source *has a* ``ProgressTracking`` object (reached through the +``progress_tracking`` property its base defines) instead of *being* a progress +mixin. The object owns the registry, the declared mode, and the guarded manual +facade; ``attach_progress_tracking`` builds it lazily so a step that never +tracks progress has no registry and the workflow reporter (which scans steps +for the ``_progress_tracking`` attribute) skips it. """ -from typing import ClassVar +from typing import Any, Optional from metadata.ingestion.progress.modes import ( ManualProgress, ProgressMode, ProgressModeError, - TotalsDeclarer, ) from metadata.ingestion.progress.registry import ProgressRegistry -class ProgressTrackingMixin: - progress_mode: ClassVar[ProgressMode] = ProgressMode.AUTO - """AUTO (default): the topology runner counts processed entities. - MANUAL: the runner makes zero progress calls; the source drives - ``manual_progress`` itself. OFF: no progress at all.""" +class ProgressTracking: + """One source's progress state: the registry, the declared mode, and the + guarded manual facade. Composed into a source, not inherited.""" + + def __init__(self, mode: ProgressMode, source_name: str) -> None: + self._mode = mode + self._source_name = source_name + self._registry = ProgressRegistry() + self._manual: Optional[ManualProgress] = None # noqa: UP045 + + @property + def mode(self) -> ProgressMode: + return self._mode @property - def progress(self) -> ProgressRegistry: - """Per-Source progress registry. First access is single-threaded - (in _iter, before worker threads spawn), so lazy init is safe.""" - registry = self.__dict__.get("_progress_registry") - if registry is None: - registry = ProgressRegistry() - self.__dict__["_progress_registry"] = registry - return registry + def registry(self) -> ProgressRegistry: + return self._registry @property - def manual_progress(self) -> ManualProgress: + def manual(self) -> ManualProgress: """MANUAL-mode counting facade. Raises for AUTO/OFF sources — the runner counts AUTO sources, so a manual increment would double-count.""" - if self.progress_mode is not ProgressMode.MANUAL: + if self._mode is not ProgressMode.MANUAL: raise ProgressModeError( - f"manual_progress requires progress_mode=MANUAL, but {type(self).__name__} " - f"declares {self.progress_mode.name}. AUTO sources are counted by the topology " - "runner and may only declare totals via declare_progress_totals()." + f"progress_tracking.manual requires progress_mode=MANUAL, but {self._source_name} " + f"declares {self._mode.name}. AUTO sources are counted by the topology runner " + "and may only declare totals via declare_progress_totals()." ) - facade = self.__dict__.get("_manual_progress") - if facade is None: - facade = ManualProgress(self.progress) - self.__dict__["_manual_progress"] = facade - return facade + if self._manual is None: + self._manual = ManualProgress(self._registry) + return self._manual + - def declare_progress_totals(self, totals: TotalsDeclarer) -> None: - """Optional connector hook: declare denominators (totals) so the run - renders % and ETA. Called exactly once by the topology runner, just - before the first non-root node is processed. Default: no totals.""" +def attach_progress_tracking(source: Any) -> ProgressTracking: + """Lazily build and cache the source's ``ProgressTracking``. Kept a + module-level helper (not a base class) so any host composes progress + without inheriting a progress mixin. First access is single-threaded (in + ``_iter``, before worker threads spawn), so the lazy init is safe.""" + tracking = source.__dict__.get("_progress_tracking") + if tracking is None: + mode = getattr(source, "progress_mode", ProgressMode.AUTO) + tracking = ProgressTracking(mode, type(source).__name__) + source.__dict__["_progress_tracking"] = tracking + return tracking diff --git a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py index 09ba665b60e5..506dc9b4a6b6 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py @@ -218,24 +218,24 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC): def _declare_progress_groups(self, label: str, total: Optional[int]) -> None: # noqa: UP045 """Declare the grouping axis (e.g. workspaces) as a global counter. - These group helpers drive ``self.manual_progress`` and therefore require + These group helpers drive ``self.progress_tracking.manual`` and therefore require the connector to set ``progress_mode = ProgressMode.MANUAL`` (as PowerBI does); calling them from a default AUTO source raises ``ProgressModeError`` and would double-count against the runner's own tracking.""" - self.manual_progress.declare_groups(label, total) + self.progress_tracking.manual.declare_groups(label, total) def _open_group_progress(self, group: str, expected_by_type: Dict[str, Optional[int]]) -> None: # noqa: UP006, UP045 """Open one child node per asset type under ``group`` so each type renders as its own line; ``expected`` may be None for lazy (running) counts.""" - self.manual_progress.open_group(group, expected_by_type) + self.progress_tracking.manual.open_group(group, expected_by_type) def _advance_group_progress(self, group: str, asset_type: str) -> None: """Record one processed asset of ``asset_type`` under ``group``.""" - self.manual_progress.advance(group, asset_type) + self.progress_tracking.manual.advance(group, asset_type) def _close_group_progress(self, group: str) -> None: """Count the finished group on its global counter and prune its subtree.""" - self.manual_progress.close_group(group) + self.progress_tracking.manual.close_group(group) @retry_with_docker_host() def __init__( diff --git a/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py b/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py index 9daa3859d2ae..a91ab36e99c2 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/powerbi/metadata.py @@ -170,7 +170,7 @@ def get_org_workspace_data(self) -> Iterable[Optional[Group]]: # noqa: UP045 workspaces = self.client.api_client.fetch_all_workspaces(filter_pattern) if workspaces: workspace_total += len(workspaces) - self.manual_progress.set_total("Workspaces", workspace_total) + self.progress_tracking.manual.set_total("Workspaces", workspace_total) for workspace in workspaces: # add the dashboards to the workspace workspace.dashboards.extend( @@ -310,7 +310,7 @@ def get_admin_workspace_data(self) -> Iterable[Optional[Group]]: # noqa: UP045 missing_workspaces, ) active_workspace_total += len(active_workspaces) - self.manual_progress.set_total("Workspaces", active_workspace_total) + self.progress_tracking.manual.set_total("Workspaces", active_workspace_total) yield from active_workspaces count += 1 else: diff --git a/ingestion/src/metadata/ingestion/source/database/lineage_source.py b/ingestion/src/metadata/ingestion/source/database/lineage_source.py index 7b1c3ffe0e53..974410382204 100644 --- a/ingestion/src/metadata/ingestion/source/database/lineage_source.py +++ b/ingestion/src/metadata/ingestion/source/database/lineage_source.py @@ -367,14 +367,14 @@ def yield_query_lineage( result_limit = self.source_config.resultLimit # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue] if result_limit is not None: - self.manual_progress.seed_scope_total("Queries", "run", result_limit) + self.progress_tracking.manual.seed_scope_total("Queries", "run", result_limit) produced = 0 def producer_fn(): nonlocal produced for table_query in self.query_lineage_producer(): produced += 1 - self.manual_progress.track("Queries") + self.progress_tracking.manual.track("Queries") yield table_query processor_fn = query_lineage_processor @@ -394,7 +394,7 @@ def producer_fn(): args, max_threads=self.source_config.threads, # pyright: ignore[reportAttributeAccessIssue] ) - self.manual_progress.reconcile_scope_total("Queries", "run", produced) + self.progress_tracking.manual.reconcile_scope_total("Queries", "run", produced) def view_lineage_producer(self) -> Iterable[TableView]: """ diff --git a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py index c22941ed3912..20585d7be5ff 100644 --- a/ingestion/src/metadata/ingestion/source/database/query_parser_source.py +++ b/ingestion/src/metadata/ingestion/source/database/query_parser_source.py @@ -28,7 +28,10 @@ from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.progress.modes import ProgressMode -from metadata.ingestion.progress.tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ( + ProgressTracking, + attach_progress_tracking, +) from metadata.ingestion.source.connections import test_connection_common from metadata.utils.helpers import get_start_and_end, retry_with_docker_host from metadata.utils.logger import ingestion_logger @@ -37,7 +40,7 @@ logger = ingestion_logger() -class QueryParserSource(ProgressTrackingMixin, Source, ABC): +class QueryParserSource(Source, ABC): """ Core class to be inherited for sources that parse query logs, be it for usage or lineage. @@ -49,6 +52,12 @@ class QueryParserSource(ProgressTrackingMixin, Source, ABC): progress_mode = ProgressMode.MANUAL + @property + def progress_tracking(self) -> ProgressTracking: + """Composed per-source progress state; these query sources are MANUAL, + so they drive ``progress_tracking.manual`` directly.""" + return attach_progress_tracking(self) + sql_stmt: str dialect: str filters: str diff --git a/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py b/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py index b0fcad36ea25..a987b138355f 100644 --- a/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py +++ b/ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py @@ -245,17 +245,17 @@ def _yield_access_history_lineage(self) -> Iterable[Either[AddLineageRequest]]: The two phases are isolated: a catastrophic failure in the combined ACCESS_HISTORY phase must not stop the COPY_HISTORY phase, and vice versa. """ - self.manual_progress.set_total("LineageRecords", None) + self.progress_tracking.manual.set_total("LineageRecords", None) try: for either in self._yield_combined_access_history(): - self.manual_progress.track("LineageRecords") + self.progress_tracking.manual.track("LineageRecords") yield either except Exception as exc: logger.warning("ACCESS_HISTORY combined lineage phase failed: %s", exc) logger.debug(traceback.format_exc()) try: for either in self._yield_copy_history_lineage(): - self.manual_progress.track("LineageRecords") + self.progress_tracking.manual.track("LineageRecords") yield either except Exception as exc: logger.warning("COPY_HISTORY lineage phase failed: %s", exc) diff --git a/ingestion/src/metadata/ingestion/source/database/usage_source.py b/ingestion/src/metadata/ingestion/source/database/usage_source.py index 6b417d486005..48d64190aefd 100644 --- a/ingestion/src/metadata/ingestion/source/database/usage_source.py +++ b/ingestion/src/metadata/ingestion/source/database/usage_source.py @@ -166,12 +166,12 @@ def _iter(self, *_, **__) -> Iterable[Either[TableQuery]]: days = max(1, (self.end - self.start).days) result_limit = self.source_config.resultLimit # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue] if result_limit is not None: - self.manual_progress.seed_scope_total("Queries", "run", result_limit * days) + self.progress_tracking.manual.seed_scope_total("Queries", "run", result_limit * days) processed = 0 for table_queries in self.get_table_query(): if table_queries: count = len(table_queries.queries) # pyright: ignore[reportAttributeAccessIssue] - self.manual_progress.track("Queries", count) + self.progress_tracking.manual.track("Queries", count) processed += count yield Either(right=table_queries) - self.manual_progress.reconcile_scope_total("Queries", "run", processed) + self.progress_tracking.manual.reconcile_scope_total("Queries", "run", processed) diff --git a/ingestion/src/metadata/workflow/workflow_status_mixin.py b/ingestion/src/metadata/workflow/workflow_status_mixin.py index 0660fee06c74..b9444d2ba309 100644 --- a/ingestion/src/metadata/workflow/workflow_status_mixin.py +++ b/ingestion/src/metadata/workflow/workflow_status_mixin.py @@ -179,9 +179,9 @@ def _find_progress_registry(self) -> Optional[ProgressRegistry]: # noqa: UP045 never tracked progress.""" result = None for step in self.workflow_steps(): # pyright: ignore[reportAttributeAccessIssue] - registry = getattr(step, "_progress_registry", None) - if registry is not None: - result = registry + tracking = getattr(step, "_progress_tracking", None) + if tracking is not None: + result = tracking.registry break return result diff --git a/ingestion/tests/unit/lineage/test_lineage_progress.py b/ingestion/tests/unit/lineage/test_lineage_progress.py index 3048bc7a95e2..392a9112e76c 100644 --- a/ingestion/tests/unit/lineage/test_lineage_progress.py +++ b/ingestion/tests/unit/lineage/test_lineage_progress.py @@ -60,4 +60,4 @@ def fake_generate(producer_fn, processor_fn, args, **kwargs): with patch.object(LineageSource, "generate_lineage_with_processes", staticmethod(fake_generate)): list(LineageSource.yield_query_lineage(source)) - assert source.progress.global_counters() == [("Queries", 4, 4)] + assert source.progress_tracking.registry.global_counters() == [("Queries", 4, 4)] diff --git a/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py b/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py index 70cb4a857f51..d31c3eb07783 100644 --- a/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py +++ b/ingestion/tests/unit/lineage/test_snowflake_access_history_progress.py @@ -27,5 +27,5 @@ def test_tracks_lineage_records_without_total_or_eta(self): results = list(SnowflakeLineageSource._yield_access_history_lineage(source)) assert len(results) == 4 - assert source.progress.global_counters() == [("LineageRecords", 4, None)] - assert source.progress.eta_seconds() is None + assert source.progress_tracking.registry.global_counters() == [("LineageRecords", 4, None)] + assert source.progress_tracking.registry.eta_seconds() is None diff --git a/ingestion/tests/unit/progress/test_progress_modes.py b/ingestion/tests/unit/progress/test_progress_modes.py index 7a303cb53f65..88ae28bdff8d 100644 --- a/ingestion/tests/unit/progress/test_progress_modes.py +++ b/ingestion/tests/unit/progress/test_progress_modes.py @@ -19,44 +19,28 @@ TotalsDeclarer, ) from metadata.ingestion.progress.registry import ProgressRegistry -from metadata.ingestion.progress.tracking import ProgressTrackingMixin +from metadata.ingestion.progress.tracking import ProgressTracking -class _AutoSource(ProgressTrackingMixin): - pass - - -class _ManualSource(ProgressTrackingMixin): - progress_mode = ProgressMode.MANUAL - - -class _OffSource(ProgressTrackingMixin): - progress_mode = ProgressMode.OFF - - -def test_default_mode_is_auto(): - assert _AutoSource().progress_mode is ProgressMode.AUTO - - -def test_manual_progress_raises_in_auto_mode(): +def test_manual_raises_in_auto_mode(): with pytest.raises(ProgressModeError): - _ = _AutoSource().manual_progress + _ = ProgressTracking(ProgressMode.AUTO, "AutoSource").manual -def test_manual_progress_raises_in_off_mode(): +def test_manual_raises_in_off_mode(): with pytest.raises(ProgressModeError): - _ = _OffSource().manual_progress + _ = ProgressTracking(ProgressMode.OFF, "OffSource").manual -def test_manual_progress_is_cached_per_instance(): - source = _ManualSource() - assert source.manual_progress is source.manual_progress +def test_manual_is_cached_per_tracking(): + tracking = ProgressTracking(ProgressMode.MANUAL, "ManualSource") + assert tracking.manual is tracking.manual -def test_manual_progress_wraps_the_source_registry(): - source = _ManualSource() - source.manual_progress.set_total("Workspaces", 3) - assert ("Workspaces", 0, 3) in source.progress.global_counters() +def test_manual_wraps_the_tracking_registry(): + tracking = ProgressTracking(ProgressMode.MANUAL, "ManualSource") + tracking.manual.set_total("Workspaces", 3) + assert ("Workspaces", 0, 3) in tracking.registry.global_counters() def test_totals_declarer_exposes_no_counting_methods(): @@ -79,8 +63,13 @@ def test_totals_declarer_sets_denominators(): def test_declare_progress_totals_default_is_noop(): + from metadata.ingestion.api.topology_runner import TopologyRunnerMixin + + class _Plain(TopologyRunnerMixin): + pass + registry = ProgressRegistry() - _AutoSource().declare_progress_totals(TotalsDeclarer(registry)) + _Plain().declare_progress_totals(TotalsDeclarer(registry)) assert registry.global_counters() == [] diff --git a/ingestion/tests/unit/progress/test_runner_tracker.py b/ingestion/tests/unit/progress/test_runner_tracker.py index d6d53dc0b499..9ea07a568920 100644 --- a/ingestion/tests/unit/progress/test_runner_tracker.py +++ b/ingestion/tests/unit/progress/test_runner_tracker.py @@ -52,7 +52,7 @@ def test_for_node_returns_noop_when_mode_is_not_auto(): source = _FakeSource(mode) tracker = TopologyProgressTracker(source) assert tracker.for_node(_table_node(source), is_leaf=True) is NO_OP_NODE_PROGRESS - assert "_progress_registry" not in source.__dict__ + assert "_progress_tracking" not in source.__dict__ def test_for_node_returns_noop_for_root_node(): @@ -82,7 +82,7 @@ def test_totals_hook_called_exactly_once_and_only_for_real_nodes(): tracker.for_node(_database_node(source), is_leaf=False) tracker.for_node(_table_node(source), is_leaf=True) assert source.declared == 1 - assert ("Database", 0, 7) in source.progress.global_counters() + assert ("Database", 0, 7) in source.progress_tracking.registry.global_counters() def test_leaf_handle_counts_open_and_advance(): @@ -92,7 +92,7 @@ def test_leaf_handle_counts_open_and_advance(): handle.open(2) handle.advance_leaf() handle.advance_leaf() - snapshot = source.progress.snapshot() + snapshot = source.progress_tracking.registry.snapshot() assert snapshot.child_type == "Table" assert snapshot.processed == 2 @@ -103,12 +103,12 @@ def test_container_handle_is_lazy_without_reconcilable_counter(): handle = tracker.for_node(_database_node(source), is_leaf=False) assert handle.wants_eager_count is False handle.open(None) - assert source.progress.snapshot().expected_by_type.get("Database") is None + assert source.progress_tracking.registry.snapshot().expected_by_type.get("Database") is None def test_container_handle_reconciles_when_counter_is_reconcilable(): source = _FakeSource() - source.progress.seed_scope_total("DatabaseSchema", "salesdb", 1) + source.progress_tracking.registry.seed_scope_total("DatabaseSchema", "salesdb", 1) tracker = TopologyProgressTracker(source) key = source._node_primary_stage(_database_node(source)).context setattr(source.context.get(), key, "salesdb") @@ -116,41 +116,41 @@ def test_container_handle_reconciles_when_counter_is_reconcilable(): handle = tracker.for_node(schema_node, is_leaf=False) assert handle.wants_eager_count is True handle.open(3) - counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + counters = {t: (done, total) for t, done, total in source.progress_tracking.registry.global_counters()} assert counters["DatabaseSchema"] == (0, 3) def test_enter_scope_closes_and_tracks_on_exit(): source = _FakeSource() - source.progress.set_total("Database", 7) + source.progress_tracking.registry.set_total("Database", 7) tracker = TopologyProgressTracker(source) node = _database_node(source) handle = tracker.for_node(node, is_leaf=False) key = source._node_primary_stage(node).context setattr(source.context.get(), key, "salesdb") closed = [] - source.progress.close = lambda path: closed.append(list(path)) + source.progress_tracking.registry.close = lambda path: closed.append(list(path)) with handle.enter_scope(): pass assert closed == [["salesdb"]] - counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + counters = {t: (done, total) for t, done, total in source.progress_tracking.registry.global_counters()} assert counters["Database"][0] == 1 def test_enter_scope_prunes_without_counting_on_failure(): source = _FakeSource() - source.progress.set_total("Database", 7) + source.progress_tracking.registry.set_total("Database", 7) tracker = TopologyProgressTracker(source) node = _database_node(source) handle = tracker.for_node(node, is_leaf=False) key = source._node_primary_stage(node).context setattr(source.context.get(), key, "salesdb") closed = [] - source.progress.close = lambda path: closed.append(list(path)) + source.progress_tracking.registry.close = lambda path: closed.append(list(path)) with pytest.raises(RuntimeError), handle.enter_scope(): raise RuntimeError("boom") assert closed == [["salesdb"]] - counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + counters = {t: (done, total) for t, done, total in source.progress_tracking.registry.global_counters()} assert counters["Database"][0] == 0 @@ -160,7 +160,7 @@ def test_enter_scope_is_noop_without_context_value(): handle = tracker.for_node(_database_node(source), is_leaf=False) with handle.enter_scope(): pass - assert source.progress.snapshot() is None # nothing opened, nothing closed + assert source.progress_tracking.registry.snapshot() is None # nothing opened, nothing closed def test_current_path_resolves_ancestor_context(): diff --git a/ingestion/tests/unit/source/database/test_mysql_progress_count.py b/ingestion/tests/unit/source/database/test_mysql_progress_count.py index d68ba0975ea4..b57cd0dacf3b 100644 --- a/ingestion/tests/unit/source/database/test_mysql_progress_count.py +++ b/ingestion/tests/unit/source/database/test_mysql_progress_count.py @@ -15,15 +15,15 @@ import pytest -from metadata.ingestion.progress.modes import TotalsDeclarer -from metadata.ingestion.progress.registry import ProgressRegistry +from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer +from metadata.ingestion.progress.tracking import ProgressTracking from metadata.ingestion.source.database.mysql.metadata import MysqlSource def _source(service_connection): source = object.__new__(MysqlSource) source.service_connection = service_connection - source.__dict__["_progress_registry"] = ProgressRegistry() + source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test") return source @@ -33,20 +33,20 @@ def mysql_source(): def test_declare_progress_totals_seeds_single_database(mysql_source): - mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress)) - counters = {t: (done, total) for t, done, total in mysql_source.progress.global_counters()} + mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress_tracking.registry)) + counters = {t: (done, total) for t, done, total in mysql_source.progress_tracking.registry.global_counters()} assert counters["Database"] == (0, 1) def test_declare_progress_totals_marks_schema_reconcilable(mysql_source): - mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress)) - assert mysql_source.progress.is_reconcilable("DatabaseSchema") is True - counters = {t: (done, total) for t, done, total in mysql_source.progress.global_counters()} + mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress_tracking.registry)) + assert mysql_source.progress_tracking.registry.is_reconcilable("DatabaseSchema") is True + counters = {t: (done, total) for t, done, total in mysql_source.progress_tracking.registry.global_counters()} assert counters["DatabaseSchema"] == (0, None) def test_declare_progress_totals_defaults_database_name_when_unset(mysql_source): source = _source(SimpleNamespace(databaseName=None, database=None)) - source.declare_progress_totals(TotalsDeclarer(source.progress)) - counters = {t: (done, total) for t, done, total in source.progress.global_counters()} + source.declare_progress_totals(TotalsDeclarer(source.progress_tracking.registry)) + counters = {t: (done, total) for t, done, total in source.progress_tracking.registry.global_counters()} assert counters["Database"] == (0, 1) diff --git a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py index baec5f745036..068be8965822 100644 --- a/ingestion/tests/unit/source/database/test_snowflake_progress_count.py +++ b/ingestion/tests/unit/source/database/test_snowflake_progress_count.py @@ -62,7 +62,8 @@ def test_filtered_names_are_cached_so_status_is_not_double_emitted(self): source.status.filter.assert_called_once() def test_producer_setup_runs_per_kept_database_without_re_filtering(self): - from metadata.ingestion.progress.registry import ProgressRegistry + from metadata.ingestion.progress.modes import ProgressMode + from metadata.ingestion.progress.tracking import ProgressTracking source = _source(["A", "B"]) for setup in ( @@ -76,7 +77,7 @@ def test_producer_setup_runs_per_kept_database_without_re_filtering(self): "set_database_tags_map", ): setattr(source, setup, MagicMock()) - source.__dict__["_progress_registry"] = ProgressRegistry() + source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test") with patch.object(snowflake_metadata, "filter_by_database", return_value=False): source._filtered_database_names() # warms the database names cache yielded = list(source.get_database_names()) @@ -88,7 +89,8 @@ def test_producer_setup_runs_per_kept_database_without_re_filtering(self): @pytest.fixture def snowflake_source(): """Minimal SnowflakeSource stub for push-totals tests.""" - from metadata.ingestion.progress.registry import ProgressRegistry + from metadata.ingestion.progress.modes import ProgressMode + from metadata.ingestion.progress.tracking import ProgressTracking source = object.__new__(SnowflakeSource) source.config = SimpleNamespace( @@ -103,7 +105,7 @@ def snowflake_source(): source.status = MagicMock() source.context = MagicMock() source.context.get.return_value = SimpleNamespace(database=None, database_service="svc") - source.__dict__["_progress_registry"] = ProgressRegistry() + source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test") return source @@ -111,8 +113,8 @@ def test_declare_progress_totals_seeds_database_and_schema(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1", "db2"] snowflake_source._schema_names_by_database = lambda: {"db1": ["s1", "s2"], "db2": ["s3"]} snowflake_source._is_schema_filtered = lambda db, sch: False - snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) - counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry)) + counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()} assert counters["Database"] == (0, 2) assert counters["DatabaseSchema"] == (0, 3) @@ -121,16 +123,16 @@ def test_declare_progress_totals_applies_schema_filter(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1"] snowflake_source._schema_names_by_database = lambda: {"db1": ["keep", "drop"]} snowflake_source._is_schema_filtered = lambda db, sch: sch == "drop" - snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) - counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry)) + counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()} assert counters["DatabaseSchema"] == (0, 1) def test_declare_progress_totals_falls_back_when_account_show_unavailable(snowflake_source): snowflake_source._filtered_database_names = lambda: ["db1"] snowflake_source._schema_names_by_database = lambda: None - snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress)) - assert snowflake_source.progress.is_reconcilable("DatabaseSchema") is True - counters = {t: (d, total) for t, d, total in snowflake_source.progress.global_counters()} + snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry)) + assert snowflake_source.progress_tracking.registry.is_reconcilable("DatabaseSchema") is True + counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()} assert counters["Database"] == (0, 1) assert counters["DatabaseSchema"] == (0, None) diff --git a/ingestion/tests/unit/test_progress_tracking.py b/ingestion/tests/unit/test_progress_tracking.py new file mode 100644 index 000000000000..a8ed829369fd --- /dev/null +++ b/ingestion/tests/unit/test_progress_tracking.py @@ -0,0 +1,40 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the composed ProgressTracking and its attach helper.""" + +from metadata.ingestion.progress.modes import ProgressMode +from metadata.ingestion.progress.registry import ProgressRegistry +from metadata.ingestion.progress.tracking import ProgressTracking, attach_progress_tracking + + +class _Owner: + progress_mode = ProgressMode.AUTO + + +class TestAttachProgressTracking: + def test_lazily_creates_one_tracking(self): + owner = _Owner() + tracking = attach_progress_tracking(owner) + assert isinstance(tracking, ProgressTracking) + assert isinstance(tracking.registry, ProgressRegistry) + assert attach_progress_tracking(owner) is tracking + + def test_caches_under_the_scanned_attribute(self): + owner = _Owner() + tracking = attach_progress_tracking(owner) + assert owner.__dict__.get("_progress_tracking") is tracking + + def test_defaults_to_auto_when_owner_declares_no_mode(self): + class _NoMode: + pass + + tracking = attach_progress_tracking(_NoMode()) + assert tracking.mode is ProgressMode.AUTO diff --git a/ingestion/tests/unit/test_progress_tracking_mixin.py b/ingestion/tests/unit/test_progress_tracking_mixin.py deleted file mode 100644 index 38d4018efbe6..000000000000 --- a/ingestion/tests/unit/test_progress_tracking_mixin.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright 2025 Collate -# Licensed under the Collate Community License, Version 1.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for the shared ProgressTrackingMixin.""" - -from metadata.ingestion.progress.registry import ProgressRegistry -from metadata.ingestion.progress.tracking import ProgressTrackingMixin - - -class _Owner(ProgressTrackingMixin): - pass - - -class TestProgressTrackingMixin: - def test_lazily_creates_one_registry(self): - owner = _Owner() - registry = owner.progress - assert isinstance(registry, ProgressRegistry) - assert owner.progress is registry - - def test_exposes_the_scanned_attribute(self): - owner = _Owner() - _ = owner.progress - assert owner.__dict__.get("_progress_registry") is owner.progress diff --git a/ingestion/tests/unit/test_usage_progress.py b/ingestion/tests/unit/test_usage_progress.py index 19b64f04a8fb..d2aa7403e19c 100644 --- a/ingestion/tests/unit/test_usage_progress.py +++ b/ingestion/tests/unit/test_usage_progress.py @@ -54,14 +54,14 @@ def test_seeds_ceiling_and_reconciles_to_real_count(self): list(UsageSource._iter(source)) - assert source.progress.global_counters() == [("Queries", 3, 3)] + assert source.progress_tracking.registry.global_counters() == [("Queries", 3, 3)] def test_ceiling_is_result_limit_times_days_before_reconcile(self): captured = {} def batches(): # peek the seeded total after the first batch, before completion - captured["mid"] = source.progress.global_counters() + captured["mid"] = source.progress_tracking.registry.global_counters() yield TableQueries(queries=[_query()]) source = _make_usage_source(batches(), result_limit=1000, start=datetime(2026, 1, 1), end=datetime(2026, 1, 3)) diff --git a/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py b/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py index 6cf74d7a165a..da0b0ab19a0d 100644 --- a/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py +++ b/ingestion/tests/unit/topology/dashboard/test_dashboard_group_progress.py @@ -42,10 +42,10 @@ def test_helpers_drive_registry_and_prune(self): src._open_group_progress("Sales", {"Dashboard": 3, "Chart": None}) src._advance_group_progress("Sales", "Dashboard") src._advance_group_progress("Sales", "Chart") - assert src.progress.global_counters() == [("Workspaces", 0, 2)] - assert src.progress.assets_ingested() == 2 + assert src.progress_tracking.registry.global_counters() == [("Workspaces", 0, 2)] + assert src.progress_tracking.registry.assets_ingested() == 2 src._close_group_progress("Sales") - assert src.progress.global_counters() == [("Workspaces", 1, 2)] - snapshot = src.progress.snapshot() + assert src.progress_tracking.registry.global_counters() == [("Workspaces", 1, 2)] + snapshot = src.progress_tracking.registry.snapshot() assert snapshot is None or all(child.label != "Sales" for child in snapshot.children) diff --git a/ingestion/tests/unit/topology/dashboard/test_powerbi.py b/ingestion/tests/unit/topology/dashboard/test_powerbi.py index 387dd1b45642..20d69f35671b 100644 --- a/ingestion/tests/unit/topology/dashboard/test_powerbi.py +++ b/ingestion/tests/unit/topology/dashboard/test_powerbi.py @@ -2422,7 +2422,7 @@ def test_upstream_datamart_lineage(self, *_): @patch.object(fqn, "build", side_effect=lambda *args, **kwargs: kwargs.get("chart_name")) def test_yield_dashboard_advances_workspace_progress(self, *_): self.powerbi.state = WorkspaceState() - self.powerbi.__dict__.pop("_progress_registry", None) + self.powerbi.__dict__.pop("_progress_tracking", None) for dash in ( PowerBIDashboard(id="dash-1", displayName="One", tiles=[]), @@ -2441,15 +2441,15 @@ def test_yield_dashboard_advances_workspace_progress(self, *_): ) list(self.powerbi.yield_dashboard(Group(id="ws-1", name="Sales"))) - assert self.powerbi.progress.assets_ingested() == 2 - out = self.powerbi.progress.render_cli() + assert self.powerbi.progress_tracking.registry.assets_ingested() == 2 + out = self.powerbi.progress_tracking.registry.render_cli() assert "Sales.Dashboard" in out assert "Dashboard 2" in out @pytest.mark.order(56) def test_get_dashboard_tracks_workspace_group(self): self.powerbi.state = WorkspaceState() - self.powerbi.__dict__.pop("_progress_registry", None) + self.powerbi.__dict__.pop("_progress_tracking", None) ws_a = Group(id="wa", name="Alpha") ws_b = Group(id="wb", name="Beta") @@ -2460,13 +2460,13 @@ def test_get_dashboard_tracks_workspace_group(self): produced = list(self.powerbi.get_dashboard()) assert [w.id for w in produced] == ["wa", "wb"] - assert self.powerbi.progress.global_counters() == [("Workspaces", 2, None)] - assert self.powerbi.progress.snapshot() is None + assert self.powerbi.progress_tracking.registry.global_counters() == [("Workspaces", 2, None)] + assert self.powerbi.progress_tracking.registry.snapshot() is None @pytest.mark.order(56) def test_get_dashboard_does_not_count_workspace_failing_before_open(self): self.powerbi.state = WorkspaceState() - self.powerbi.__dict__.pop("_progress_registry", None) + self.powerbi.__dict__.pop("_progress_tracking", None) ws_ok = Group(id="ok", name="Ok") ws_bad = Group(id="bad", name="Bad") @@ -2477,11 +2477,11 @@ def test_get_dashboard_does_not_count_workspace_failing_before_open(self): produced = list(self.powerbi.get_dashboard()) assert [w.id for w in produced] == ["ok"] - assert self.powerbi.progress.global_counters() == [("Workspaces", 1, None)] + assert self.powerbi.progress_tracking.registry.global_counters() == [("Workspaces", 1, None)] @pytest.mark.order(57) def test_admin_workspace_progress_total_reconciles_to_active_scan_results(self): - self.powerbi.__dict__.pop("_progress_registry", None) + self.powerbi.__dict__.pop("_progress_tracking", None) self.powerbi.pagination_entity_per_page = 100 candidate_workspaces = [Group(id=f"ws-{index}", name=f"Workspace {index}") for index in range(46)] @@ -2501,12 +2501,12 @@ def test_admin_workspace_progress_total_reconciles_to_active_scan_results(self): produced = list(self.powerbi.get_admin_workspace_data()) assert [workspace.id for workspace in produced] == [f"ws-{index}" for index in range(40)] - assert self.powerbi.progress.global_counters() == [("Workspaces", 0, 40)] + assert self.powerbi.progress_tracking.registry.global_counters() == [("Workspaces", 0, 40)] @pytest.mark.order(58) @patch.object(fqn, "build", side_effect=lambda *args, **kwargs: kwargs.get("chart_name")) def test_unnamed_workspace_keys_progress_on_id(self, *_): - self.powerbi.__dict__.pop("_progress_registry", None) + self.powerbi.__dict__.pop("_progress_tracking", None) self.powerbi.state.add_filtered_dashboard(PowerBIDashboard(id="dash-1", displayName="One", tiles=[])) @@ -2520,7 +2520,7 @@ def test_unnamed_workspace_keys_progress_on_id(self, *_): self.powerbi._open_group_progress("ws-x", {"Dashboard": None}) list(self.powerbi.yield_dashboard(Group(id="ws-x", name=None))) - out = self.powerbi.progress.render_cli() + out = self.powerbi.progress_tracking.registry.render_cli() assert "ws-x.Dashboard" in out assert "None.Dashboard" not in out diff --git a/ingestion/tests/unit/topology/test_topology_runner_progress.py b/ingestion/tests/unit/topology/test_topology_runner_progress.py index 12133f82100c..d0ee1bed79f7 100644 --- a/ingestion/tests/unit/topology/test_topology_runner_progress.py +++ b/ingestion/tests/unit/topology/test_topology_runner_progress.py @@ -31,12 +31,12 @@ class TestRunnerProgressSurface: def test_progress_is_lazy_per_instance(self): runner = TopologyRunnerMixin() - registry = runner.progress + registry = runner.progress_tracking.registry assert isinstance(registry, ProgressRegistry) - assert runner.progress is registry + assert runner.progress_tracking.registry is registry def test_distinct_instances_distinct_registries(self): - assert TopologyRunnerMixin().progress is not TopologyRunnerMixin().progress + assert TopologyRunnerMixin().progress_tracking.registry is not TopologyRunnerMixin().progress_tracking.registry def test_current_progress_path_default_is_empty(self): assert TopologyRunnerMixin().progress_tracker.current_path() == [] @@ -97,19 +97,19 @@ def _drive_container(runner): def test_disabled_walk_creates_no_registry(): runner = _WalkRunner(mode=ProgressMode.OFF) _drive_leaf(runner) - assert "_progress_registry" not in runner.__dict__ + assert "_progress_tracking" not in runner.__dict__ def test_manual_walk_creates_no_registry(): runner = _WalkRunner(mode=ProgressMode.MANUAL) _drive_leaf(runner) - assert "_progress_registry" not in runner.__dict__ + assert "_progress_tracking" not in runner.__dict__ def test_enabled_walk_records_into_registry(): runner = _WalkRunner() _drive_leaf(runner) - snapshot = runner.progress.snapshot() + snapshot = runner.progress_tracking.registry.snapshot() assert snapshot.child_type == "Table" assert snapshot.processed == 2 @@ -117,7 +117,7 @@ def test_enabled_walk_records_into_registry(): def test_plain_connector_gets_progress_by_default(): runner = _WalkRunner() _drive_leaf(runner) - assert runner.progress.snapshot().processed == 2 + assert runner.progress_tracking.registry.snapshot().processed == 2 def test_container_open_none_through_runner(): @@ -127,7 +127,7 @@ def test_container_open_none_through_runner(): the Database level.""" runner = _WalkRunner() _drive_container(runner) - snapshot = runner.progress.snapshot() + snapshot = runner.progress_tracking.registry.snapshot() assert snapshot is not None assert snapshot.child_type == "Database" assert snapshot.expected_by_type.get("Database") is None @@ -141,8 +141,8 @@ def progress_runner(): def test_container_without_push_is_unknown(progress_runner): # progress_runner: a TopologyRunnerMixin test double with progress_mode=AUTO - progress_runner.progress.open(["db1"], "DatabaseSchema", None) - snap = progress_runner.progress.snapshot() + progress_runner.progress_tracking.registry.open(["db1"], "DatabaseSchema", None) + snap = progress_runner.progress_tracking.registry.snapshot() assert snap.children[0].expected is None @@ -181,7 +181,7 @@ def test_scope_path_for_node_is_none_without_context_value(): def test_completed_container_is_closed_through_runner(): runner = _CtxWalkRunner() closed: list = [] - runner.progress.close = lambda path: closed.append(list(path)) + runner.progress_tracking.registry.close = lambda path: closed.append(list(path)) container = get_topology_node("database", runner.topology) list(runner._process_node(container)) # producer yields "a","b"; each database entity is closed at its own path @@ -232,7 +232,7 @@ def test_leaf_is_lazy_when_progress_on_preserving_per_yield_state(): runner = _StatefulLeafRunner() list(runner._process_node(get_topology_node("table", runner.topology))) assert runner.state_live_at_stage == ["a", "b"] - assert runner.progress.snapshot().processed == 2 + assert runner.progress_tracking.registry.snapshot().processed == 2 class _CtxReadingLeafRunner(TopologyRunnerMixin): @@ -272,21 +272,21 @@ def test_leaf_producer_reads_context_written_by_prior_yield_stage(): def test_closing_container_tracks_global_done(): runner = _CtxWalkRunner() - runner.progress.set_total("Database", 5) - runner.progress.set_total("DatabaseSchema", 9) + runner.progress_tracking.registry.set_total("Database", 5) + runner.progress_tracking.registry.set_total("DatabaseSchema", 9) list(runner._process_node(get_topology_node("database", runner.topology))) - counters = {t: (d, total) for t, d, total in runner.progress.global_counters()} + counters = {t: (d, total) for t, d, total in runner.progress_tracking.registry.global_counters()} assert counters["Database"][0] == 2 # 2 databases ("a","b") closed assert counters["DatabaseSchema"][0] == 4 # 2 schemas per database closed def test_reconcilable_container_reconciles_total(): runner = _CtxWalkRunner() - runner.progress.seed_scope_total("DatabaseSchema", "a", 1) - runner.progress.seed_scope_total("DatabaseSchema", "b", 1) + runner.progress_tracking.registry.seed_scope_total("DatabaseSchema", "a", 1) + runner.progress_tracking.registry.seed_scope_total("DatabaseSchema", "b", 1) # upfront total = 2 (1 seeded per declared database) list(runner._process_node(get_topology_node("database", runner.topology))) - counters = {t: (d, total) for t, d, total in runner.progress.global_counters()} + counters = {t: (d, total) for t, d, total in runner.progress_tracking.registry.global_counters()} # each database actually walks 2 schemas -> reconciled total 2*2=4; done=4 assert counters["DatabaseSchema"] == (4, 4) @@ -312,13 +312,13 @@ def test_multithread_schema_node_reconciles_and_tracks_done(): single-thread reconcile test: seed 1 per declared db, observe 2 per db, final counters show (done=4, total=4).""" runner = _MultiThreadCtxWalkRunner() - runner.progress.seed_scope_total("DatabaseSchema", "a", 1) - runner.progress.seed_scope_total("DatabaseSchema", "b", 1) + runner.progress_tracking.registry.seed_scope_total("DatabaseSchema", "a", 1) + runner.progress_tracking.registry.seed_scope_total("DatabaseSchema", "b", 1) # upfront total = 2; driving _process_node(database) recurses into # process_nodes([databaseSchema]) which routes to _multithread_process_node # because databaseSchema.threads=True and context.threads=2 list(runner._process_node(get_topology_node("database", runner.topology))) - counters = {t: (d, total) for t, d, total in runner.progress.global_counters()} + counters = {t: (d, total) for t, d, total in runner.progress_tracking.registry.global_counters()} # each database actually walks 2 schemas -> reconciled total 2*2=4; done=4 assert counters["DatabaseSchema"] == (4, 4) @@ -338,4 +338,4 @@ def test_declare_progress_totals_called_once_per_walk(): _drive_leaf(runner) _drive_leaf(runner) assert runner.declared == 1 - assert ("Table", 0, 99) in runner.progress.global_counters() + assert ("Table", 0, 99) in runner.progress_tracking.registry.global_counters() diff --git a/ingestion/tests/unit/workflow/test_status_mixin_progress.py b/ingestion/tests/unit/workflow/test_status_mixin_progress.py index 29c3230448f5..06b6e1a5ac77 100644 --- a/ingestion/tests/unit/workflow/test_status_mixin_progress.py +++ b/ingestion/tests/unit/workflow/test_status_mixin_progress.py @@ -22,7 +22,7 @@ class _FakeStep: def __init__(self, registry: ProgressRegistry) -> None: - self._progress_registry = registry + self._progress_tracking = SimpleNamespace(registry=registry) class _Harness(WorkflowStatusMixin): From f4d03567c576381f941d7ee58e11fd8fecc140b7 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Wed, 8 Jul 2026 17:10:31 +0530 Subject: [PATCH 17/18] fix(ingestion): totals-hook failure degrades progress instead of aborting the run declare_progress_totals runs in for_node, before _run_node_producer's own try/except, so a raising totals hook (e.g. Snowflake SHOW DATABASES or MySQL get_database_names failing) propagated out of the topology walk and killed the whole ingestion run for a progress-only feature. Wrap the hook in _declare_totals_once: warn and continue with denominator-less progress. Flagged by greptile (P1) and Copilot. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../metadata/ingestion/progress/runner_tracker.py | 14 +++++++++++++- .../tests/unit/progress/test_runner_tracker.py | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py index 641493855eb5..e622236af750 100644 --- a/ingestion/src/metadata/ingestion/progress/runner_tracker.py +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -16,6 +16,7 @@ shared no-op — the walk code carries zero progress conditionals. """ +import traceback from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set # noqa: UP035 from metadata.ingestion.models.topology import ( @@ -26,6 +27,9 @@ ) from metadata.ingestion.ometa.utils import model_str from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer +from metadata.utils.logger import ingestion_logger + +logger = ingestion_logger() if TYPE_CHECKING: from metadata.ingestion.progress.registry import ProgressRegistry @@ -183,9 +187,17 @@ def for_node(self, node: TopologyNode, is_leaf: bool): return result def _declare_totals_once(self) -> None: + """Run the connector's totals hook once. Totals are a progress-only + enhancement (they drive % and ETA), so a failure here — e.g. a totals + query raising — must not abort the ingestion walk: warn and continue + with denominator-less progress.""" if not self._totals_declared: self._totals_declared = True - self._source.declare_progress_totals(TotalsDeclarer(self.registry)) + try: + self._source.declare_progress_totals(TotalsDeclarer(self.registry)) + except Exception as exc: + logger.warning("Progress totals declaration failed; continuing without totals: %s", exc) + logger.debug(traceback.format_exc()) def current_path(self, entity_type_name: Optional[str] = None) -> List[str]: # noqa: UP006,UP045 """Ancestor container labels from the node's primary-stage ``consumer`` diff --git a/ingestion/tests/unit/progress/test_runner_tracker.py b/ingestion/tests/unit/progress/test_runner_tracker.py index 9ea07a568920..bb7fda6a56ca 100644 --- a/ingestion/tests/unit/progress/test_runner_tracker.py +++ b/ingestion/tests/unit/progress/test_runner_tracker.py @@ -85,6 +85,21 @@ def test_totals_hook_called_exactly_once_and_only_for_real_nodes(): assert ("Database", 0, 7) in source.progress_tracking.registry.global_counters() +class _RaisingTotalsSource(_FakeSource): + def declare_progress_totals(self, totals: TotalsDeclarer) -> None: + self.declared += 1 + raise RuntimeError("totals query blew up") + + +def test_totals_hook_failure_does_not_abort_the_walk(): + source = _RaisingTotalsSource() + tracker = TopologyProgressTracker(source) + handle = tracker.for_node(_database_node(source), is_leaf=False) + assert isinstance(handle, NodeProgress) # walk continues, no exception propagates + tracker.for_node(_table_node(source), is_leaf=True) + assert source.declared == 1 # marked declared up front, so it is not retried per node + + def test_leaf_handle_counts_open_and_advance(): source = _FakeSource() tracker = TopologyProgressTracker(source) From a3d166326524a378efc0d5f7046509ae07d06188 Mon Sep 17 00:00:00 2001 From: ulixius9 Date: Wed, 8 Jul 2026 17:16:04 +0530 Subject: [PATCH 18/18] docs(ingestion): flag the eager-count memory risk on wants_eager_count Make it explicit that returning True materializes the node's entire producer output in memory: safe only for bounded container nodes (schemas per db), never for high-cardinality leaves (tables/dashboards). Cross-reference the list() site in the runner. --- .../metadata/ingestion/api/topology_runner.py | 2 ++ .../ingestion/progress/runner_tracker.py | 29 ++++++++++++++----- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/ingestion/src/metadata/ingestion/api/topology_runner.py b/ingestion/src/metadata/ingestion/api/topology_runner.py index 5911371572cd..7d6bc6e3c0b1 100644 --- a/ingestion/src/metadata/ingestion/api/topology_runner.py +++ b/ingestion/src/metadata/ingestion/api/topology_runner.py @@ -219,6 +219,8 @@ def _process_node(self, node: TopologyNode) -> Iterable[Entity]: node_progress = self.progress_tracker.for_node(node, is_leaf=not child_nodes) if node_progress.wants_eager_count: + # Holds the node's whole producer output in memory — safe only for + # bounded container nodes; see NodeProgress.wants_eager_count. node_entities = list(self._run_node_producer(node) or []) node_progress.open(len(node_entities)) else: diff --git a/ingestion/src/metadata/ingestion/progress/runner_tracker.py b/ingestion/src/metadata/ingestion/progress/runner_tracker.py index e622236af750..a0d083c0e2e5 100644 --- a/ingestion/src/metadata/ingestion/progress/runner_tracker.py +++ b/ingestion/src/metadata/ingestion/progress/runner_tracker.py @@ -110,14 +110,27 @@ def __init__( @property def wants_eager_count(self) -> bool: - """Materialize the producer for an exact child count only when the node's - counter is reconcilable (a container whose scope total is nudged toward - the observed child count). Leaf producers are always iterated lazily: - their per-leaf ``advance`` already yields an accurate processed count, and - an eager ``list()``-drain would run the producer to completion before any - stage sinks an entity — breaking connectors (e.g. S3 storage's - ``get_containers``) whose producer reads context a stage of the - just-yielded entity populated.""" + """Whether the runner should drain this node's producer into a ``list()`` + up front to learn its exact child count. + + MEMORY WARNING: returning ``True`` makes the topology runner hold the + node's ENTIRE producer output in memory at once (see + ``TopologyRunnerMixin._process_node``). For a high-cardinality node — + tables in a large schema, dashboards in a big tenant — that is a real + OOM risk. Keep this ``True`` ONLY for low-cardinality *container* nodes + whose count is bounded (e.g. schemas per database, typically tens), and + NEVER for leaf entity nodes, which can number in the hundreds of + thousands. Any new reconcilable declaration must respect that bound; if + a container can itself be huge, add a cap or count it lazily instead. + + So: eager only when the node's counter is reconcilable — a bounded + container whose scope total is nudged toward the observed child count. + Leaf producers are always iterated lazily: their per-leaf ``advance`` + already yields an accurate processed count without materializing + anything, and an eager ``list()``-drain would additionally run the + producer to completion before any stage sinks an entity — breaking + connectors (e.g. S3 storage's ``get_containers``) whose producer reads + context a stage of the just-yielded entity populated.""" return self._reconcilable def open(self, count: Optional[int]) -> None: # noqa: UP045