feat(ingest/snowflake): use auto_workunit_reporter helper (#7568)

This commit is contained in:
Harshal Sheth 2023-03-14 16:54:57 -04:00 committed by GitHub
parent a18c93c028
commit cb6c3dc8dd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 8 additions and 8 deletions

View File

@ -378,7 +378,6 @@ class SnowflakeUsageExtractor(
id=f"{start_time.isoformat()}-operation-aspect-{resource}", id=f"{start_time.isoformat()}-operation-aspect-{resource}",
mcp=mcp, mcp=mcp,
) )
self.report.report_workunit(wu)
yield wu yield wu
def _process_snowflake_history_row( def _process_snowflake_history_row(

View File

@ -131,6 +131,7 @@ from datahub.utilities.registries.domain_registry import DomainRegistry
from datahub.utilities.source_helpers import ( from datahub.utilities.source_helpers import (
auto_stale_entity_removal, auto_stale_entity_removal,
auto_status_aspect, auto_status_aspect,
auto_workunit_reporter,
) )
from datahub.utilities.time import datetime_to_ts_millis from datahub.utilities.time import datetime_to_ts_millis
@ -567,7 +568,10 @@ class SnowflakeV2Source(
def get_workunits(self) -> Iterable[MetadataWorkUnit]: def get_workunits(self) -> Iterable[MetadataWorkUnit]:
return auto_stale_entity_removal( return auto_stale_entity_removal(
self.stale_entity_removal_handler, self.stale_entity_removal_handler,
auto_workunit_reporter(
self.report,
auto_status_aspect(self.get_workunits_internal()), auto_status_aspect(self.get_workunits_internal()),
),
) )
def report_warehouse_failure(self): def report_warehouse_failure(self):
@ -1005,7 +1009,6 @@ class SnowflakeV2Source(
yield from add_table_to_schema_container( yield from add_table_to_schema_container(
dataset_urn=dataset_urn, dataset_urn=dataset_urn,
parent_container_key=schema_container_key, parent_container_key=schema_container_key,
report=self.report,
) )
dpi_aspect = get_dataplatform_instance_aspect( dpi_aspect = get_dataplatform_instance_aspect(
dataset_urn=dataset_urn, dataset_urn=dataset_urn,
@ -1031,7 +1034,6 @@ class SnowflakeV2Source(
entity_urn=dataset_urn, entity_urn=dataset_urn,
domain_config=self.config.domain, domain_config=self.config.domain,
domain_registry=self.domain_registry, domain_registry=self.domain_registry,
report=self.report,
) )
if table.tags: if table.tags:
@ -1230,7 +1232,6 @@ class SnowflakeV2Source(
sub_types=[DatasetContainerSubTypes.DATABASE], sub_types=[DatasetContainerSubTypes.DATABASE],
domain_registry=self.domain_registry, domain_registry=self.domain_registry,
domain_config=self.config.domain, domain_config=self.config.domain,
report=self.report,
external_url=self.get_external_url_for_database(database.name) external_url=self.get_external_url_for_database(database.name)
if self.config.include_external_url if self.config.include_external_url
else None, else None,
@ -1275,7 +1276,6 @@ class SnowflakeV2Source(
domain_config=self.config.domain, domain_config=self.config.domain,
schema_container_key=schema_container_key, schema_container_key=schema_container_key,
sub_types=[DatasetContainerSubTypes.SCHEMA], sub_types=[DatasetContainerSubTypes.SCHEMA],
report=self.report,
domain_registry=self.domain_registry, domain_registry=self.domain_registry,
description=schema.comment, description=schema.comment,
external_url=self.get_external_url_for_schema(schema.name, db_name) external_url=self.get_external_url_for_schema(schema.name, db_name)

View File

@ -182,7 +182,7 @@ def get_domain_wu(
entity_urn: str, entity_urn: str,
domain_config: Dict[str, AllowDenyPattern], domain_config: Dict[str, AllowDenyPattern],
domain_registry: DomainRegistry, domain_registry: DomainRegistry,
report: SourceReport, report: Optional[SourceReport] = None,
) -> Iterable[MetadataWorkUnit]: ) -> Iterable[MetadataWorkUnit]:
domain_urn = gen_domain_urn(dataset_name, domain_config, domain_registry) domain_urn = gen_domain_urn(dataset_name, domain_config, domain_registry)
if domain_urn: if domain_urn:
@ -191,6 +191,7 @@ def get_domain_wu(
domain_urn=domain_urn, domain_urn=domain_urn,
) )
for wu in wus: for wu in wus:
if report:
report.report_workunit(wu) report.report_workunit(wu)
yield wu yield wu