refactor(ingest): Call source_helpers via new WorkUnitProcessors in base Source #8101

asikowitz · 2023-05-23T00:11:25Z

Somewhat serious refactor that provides a default implementation of get_workunits and changes get_workunits_internal to be the main overriden method. Adds get_workunit_processors() which returns a list of callables that take in a WorkUnit stream and alter that stream (i.e. source helpers). Provides a default list of workunit processors (source helpers) that can be changed by overriding the method.

Checklist

The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
Links to related issues (if applicable)
Tests for the changes have been added/updated (if applicable)
Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

…ase Source

asikowitz · 2023-05-23T00:12:01Z

metadata-ingestion/src/datahub/ingestion/api/source_helpers.py

@@ -119,14 +118,9 @@ def auto_workunit_reporter(report: SourceReport, stream: Iterable[T]) -> Iterabl

 def auto_materialize_referenced_tags(
    stream: Iterable[MetadataWorkUnit],
-    active: bool = True,


Now handled by passing in None for the workunit processor, although I don't think this was ever passed as False

asikowitz · 2023-05-23T00:12:35Z

metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py

-            state_type_class=GenericCheckpointState,
-            pipeline_name=self.ctx.pipeline_name,
-            run_id=self.ctx.run_id,
+        self.stale_entity_removal_handler = StaleEntityRemovalHandler.create(


This gets used to call self.stale_entity_removal_handler.add_urn_to_skip(node_datahub_urn) at some point

asikowitz · 2023-05-23T00:14:00Z

metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py

@@ -1197,6 +1185,17 @@ def get_workspace_workunit(
                for workunit in dataset_workunits:
                    yield workunit

+    def get_workunit_processors(self) -> Sequence[Optional[MetadataWorkUnitProcessor]]:


Pay attention to this file, some pretty non-standard behavior. Not sure if there's a cleaner way to do this... could directly overwrite get_workunits as one option

hsheth2 · 2023-05-23T15:09:06Z

metadata-ingestion/src/datahub/ingestion/source/file.py

-    def get_workunits(self) -> Iterable[MetadataWorkUnit]:
-        return auto_workunit_reporter(self.report, self.get_workunits_internal())
+    def get_workunit_processors(self) -> Sequence[Optional[MetadataWorkUnitProcessor]]:
+        return [partial(auto_workunit_reporter, self.report)]


add a comment that not calling super() is intentional here

hsheth2 · 2023-05-23T15:11:31Z

metadata-ingestion/src/datahub/ingestion/api/source.py

@@ -155,9 +176,35 @@ def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
        # can't make this method abstract.
        raise NotImplementedError('sources must implement "create"')

-    @abstractmethod
+    def get_workunit_processors(self) -> Sequence[Optional[MetadataWorkUnitProcessor]]:


I think we should change this return type to List[Optional[MetadataWorkUnitProcessor]]

So we can append to it? I guess it's more specific.

not for append - we shouldn't be using append here

Sequence is a bit weird to work with in some places with mypy. We can always make it more general in the future, but let's not restrain ourselves unnecessarily

Oh really, what issues have you seen with Sequence? In general, I think we want to use Sequence (if possible) for parameters but List / other specific types for return types

Actually i think sequences are fine. iirc the issue i've seen is that Iterable is not covariant (so passing Iterable[Square] to a function that takes Iterable[Shape] would throw an error), but sequence works fine there

In general, I think we want to use Sequence (if possible) for parameters but List / other specific types for return types

yup this seems like a good rule of thumb - might be worth adding to the list here https://datahubproject.io/docs/metadata-ingestion/developing/#code-style

hsheth2 · 2023-05-23T15:12:49Z

metadata-ingestion/tests/integration/integration_helpers.py

+)
+
+
+def get_current_checkpoint_from_pipeline(


we're gonna get merge conflicts on this with my PR

hsheth2

LGTM

mayurinehate

LGTM

mayurinehate · 2023-05-24T08:20:39Z

metadata-ingestion/src/datahub/ingestion/api/source.py

@@ -155,9 +177,35 @@ def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
        # can't make this method abstract.
        raise NotImplementedError('sources must implement "create"')

-    @abstractmethod
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        """A list of functions that transforms the workunits produced by this source.


So the returned list can have None ?
Not clear to me why we need Optional[MetadataWorkUnitProcessor] as opposed to simply MetadataWorkUnitProcessor

Yeah, the idea is we can do something like return [*super().get_workunit_processors(), other_workunit_processor if self.config.flag else None]

mayurinehate · 2023-05-24T08:38:36Z

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py

+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]


Stale Entity Removal would emit some workunits, which won't be reported I believe, as workunit_processor for stale entity removal comes after auto_workunit_reporter. The same was the case earlier, so guessing thats okay.

Yeah, was trying to match existing behavior. If we want it the other way we can always make the change

refactor(ingest): Call source_helpers via new WorkUnitProcessors in b…

f8ff3f1

…ase Source

asikowitz requested review from treff7es, hsheth2 and mayurinehate May 23, 2023 00:11

asikowitz commented May 23, 2023

View reviewed changes

github-actions bot added the ingestion PR or Issue related to the ingestion of metadata label May 23, 2023

vercel bot had a problem deploying to Preview May 23, 2023 00:27 Failure

hsheth2 reviewed May 23, 2023

View reviewed changes

asikowitz added 2 commits May 23, 2023 11:30

pr feedback

358be3e

Merge branch 'master' into workunit-processors

691d7e4

asikowitz requested a review from hsheth2 May 23, 2023 22:46

vercel bot had a problem deploying to Preview May 23, 2023 22:53 Failure

asikowitz added 2 commits May 23, 2023 17:35

remove cyclical dependencies

6428e77

fix tests

2c0d113

vercel bot deployed to Preview May 24, 2023 01:03 View deployment

update goldens

4a61ef9

vercel bot deployed to Preview May 24, 2023 02:42 View deployment

hsheth2 approved these changes May 24, 2023

View reviewed changes

fix tests

5079627

vercel bot deployed to Preview May 24, 2023 08:14 View deployment

mayurinehate approved these changes May 24, 2023

View reviewed changes

Merge branch 'master' into workunit-processors

f1349ac

vercel bot deployed to Preview May 24, 2023 18:16 View deployment

asikowitz merged commit fdbc4de into datahub-project:master May 24, 2023

asikowitz deleted the workunit-processors branch May 24, 2023 20:36

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor(ingest): Call source_helpers via new WorkUnitProcessors in base Source #8101

refactor(ingest): Call source_helpers via new WorkUnitProcessors in base Source #8101

asikowitz commented May 23, 2023

asikowitz May 23, 2023

asikowitz May 23, 2023

asikowitz May 23, 2023

hsheth2 May 23, 2023

hsheth2 May 23, 2023

asikowitz May 23, 2023

hsheth2 May 23, 2023

asikowitz May 23, 2023

hsheth2 May 23, 2023

hsheth2 May 23, 2023

hsheth2 May 23, 2023

hsheth2 left a comment

mayurinehate left a comment

mayurinehate May 24, 2023

asikowitz May 24, 2023

mayurinehate May 24, 2023

asikowitz May 24, 2023

refactor(ingest): Call source_helpers via new WorkUnitProcessors in base Source #8101

refactor(ingest): Call source_helpers via new WorkUnitProcessors in base Source #8101

Conversation

asikowitz commented May 23, 2023

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hsheth2 left a comment

Choose a reason for hiding this comment

mayurinehate left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment