acryl-datahub 1.0.0.1rc6__py3-none-any.whl → 1.0.0.2__py3-none-any.whl

datahub/ingestion/source/iceberg/iceberg.py

@@ -2,6 +2,7 @@ import json
 import logging
 import threading
 import uuid
+from functools import partial
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 from dateutil import parser as dateutil_parser
@@ -47,6 +48,12 @@ from datahub.emitter.mce_builder import (
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import NamespaceKey
+from datahub.ingestion.api.auto_work_units.auto_dataset_properties_aspect import (
+    auto_patch_last_modified,
+)
+from datahub.ingestion.api.auto_work_units.auto_ensure_aspect_size import (
+    EnsureAspectSizeProcessor,
+)
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -57,6 +64,14 @@ from datahub.ingestion.api.decorators import (
     support_status,
 )
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
+from datahub.ingestion.api.source_helpers import (
+    AutoSystemMetadata,
+    auto_fix_duplicate_schema_field_paths,
+    auto_fix_empty_field_paths,
+    auto_lowercase_urns,
+    auto_materialize_referenced_tags_terms,
+    auto_workunit_reporter,
+)
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.extractor import schema_util
 from datahub.ingestion.source.common.subtypes import (
@@ -82,6 +97,8 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaMetadata,
 )
 from datahub.metadata.schema_classes import (
+    BrowsePathEntryClass,
+    BrowsePathsV2Class,
     ContainerClass,
     DataPlatformInstanceClass,
     DatasetPropertiesClass,
@@ -134,6 +151,7 @@ class IcebergSource(StatefulIngestionSourceBase):
         super().__init__(config, ctx)
         self.report: IcebergSourceReport = IcebergSourceReport()
         self.config: IcebergSourceConfig = config
+        self.ctx: PipelineContext = ctx
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "IcebergSource":
@@ -141,8 +159,47 @@ class IcebergSource(StatefulIngestionSourceBase):
         return cls(config, ctx)
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        # This source needs to overwrite standard `get_workunit_processor`, because it is unique in terms of usage
+        # of parallelism. Because of this, 2 processors won't work as expected:
+        # 1. browse_path_processor - it needs aspects for a single entity to be continuous - which is not guaranteed
+        #    in this source
+        # 2. automatic stamping with systemMetadata - in current implementation of the Source class this processor
+        #    would have been applied in a thread (single) shared between the source, processors and transformers.
+        #    Since the metadata scraping happens in separate threads, this could lead to difference between
+        #    time used by systemMetadata and actual time at which metadata was read
+        auto_lowercase_dataset_urns: Optional[MetadataWorkUnitProcessor] = None
+        if (
+            self.ctx.pipeline_config
+            and self.ctx.pipeline_config.source
+            and self.ctx.pipeline_config.source.config
+            and (
+                (
+                    hasattr(
+                        self.ctx.pipeline_config.source.config,
+                        "convert_urns_to_lowercase",
+                    )
+                    and self.ctx.pipeline_config.source.config.convert_urns_to_lowercase
+                )
+                or (
+                    hasattr(self.ctx.pipeline_config.source.config, "get")
+                    and self.ctx.pipeline_config.source.config.get(
+                        "convert_urns_to_lowercase"
+                    )
+                )
+            )
+        ):
+            auto_lowercase_dataset_urns = auto_lowercase_urns
+
         return [
-            *super().get_workunit_processors(),
+            auto_lowercase_dataset_urns,
+            auto_materialize_referenced_tags_terms,
+            partial(
+                auto_fix_duplicate_schema_field_paths, platform=self._infer_platform()
+            ),
+            partial(auto_fix_empty_field_paths, platform=self._infer_platform()),
+            partial(auto_workunit_reporter, self.get_report()),
+            auto_patch_last_modified,
+            EnsureAspectSizeProcessor(self.get_report()).ensure_aspect_size,
             StaleEntityRemovalHandler.create(
                 self, self.config, self.ctx
             ).workunit_processor,
@@ -208,6 +265,12 @@ class IcebergSource(StatefulIngestionSourceBase):
                     )
                     thread_local.local_catalog = self.config.get_catalog()
 
+                if not hasattr(thread_local, "stamping_processor"):
+                    LOGGER.debug(
+                        f"Didn't find stamping_processor in thread_local ({thread_local}), initializing new workunit processor"
+                    )
+                    thread_local.stamping_processor = AutoSystemMetadata(self.ctx)
+
                 with PerfTimer() as timer:
                     table = thread_local.local_catalog.load_table(dataset_path)
                     time_taken = timer.elapsed_seconds()
@@ -224,9 +287,11 @@ class IcebergSource(StatefulIngestionSourceBase):
                 for aspect in self._create_iceberg_table_aspects(
                     dataset_name, table, namespace_urn
                 ):
-                    yield MetadataChangeProposalWrapper(
-                        entityUrn=dataset_urn, aspect=aspect
-                    ).as_workunit()
+                    yield thread_local.stamping_processor.stamp_wu(
+                        MetadataChangeProposalWrapper(
+                            entityUrn=dataset_urn, aspect=aspect
+                        ).as_workunit()
+                    )
             except NoSuchPropertyException as e:
                 self.report.warning(
                     title="Unable to process table",
@@ -308,6 +373,7 @@ class IcebergSource(StatefulIngestionSourceBase):
             return
 
         try:
+            stamping_processor = AutoSystemMetadata(self.ctx)
             namespace_ids = self._get_namespaces(catalog)
             namespaces: List[Tuple[Identifier, str]] = []
             for namespace in namespace_ids:
@@ -323,9 +389,11 @@ class IcebergSource(StatefulIngestionSourceBase):
                 )
                 namespaces.append((namespace, namespace_urn))
                 for aspect in self._create_iceberg_namespace_aspects(namespace):
-                    yield MetadataChangeProposalWrapper(
-                        entityUrn=namespace_urn, aspect=aspect
-                    ).as_workunit()
+                    yield stamping_processor.stamp_wu(
+                        MetadataChangeProposalWrapper(
+                            entityUrn=namespace_urn, aspect=aspect
+                        ).as_workunit()
+                    )
             LOGGER.debug("Namespaces ingestion completed")
         except Exception as e:
             self.report.report_failure(
@@ -366,7 +434,9 @@ class IcebergSource(StatefulIngestionSourceBase):
                 yield dataset_ownership
 
             yield self._create_schema_metadata(dataset_name, table)
-            yield self._get_dataplatform_instance_aspect()
+            dpi = self._get_dataplatform_instance_aspect()
+            yield dpi
+            yield self._create_browse_paths_aspect(dpi.instance, str(namespace_urn))
             yield ContainerClass(container=str(namespace_urn))
 
         self.report.report_table_processing_time(
@@ -377,6 +447,22 @@ class IcebergSource(StatefulIngestionSourceBase):
             profiler = IcebergProfiler(self.report, self.config.profiling)
             yield from profiler.profile_table(dataset_name, table)
 
+    def _create_browse_paths_aspect(
+        self,
+        platform_instance_urn: Optional[str] = None,
+        container_urn: Optional[str] = None,
+    ) -> BrowsePathsV2Class:
+        path = []
+        if platform_instance_urn:
+            path.append(
+                BrowsePathEntryClass(
+                    id=platform_instance_urn, urn=platform_instance_urn
+                )
+            )
+        if container_urn:
+            path.append(BrowsePathEntryClass(id=container_urn, urn=container_urn))
+        return BrowsePathsV2Class(path=path)
+
     def _get_partition_aspect(self, table: Table) -> Optional[str]:
         """Extracts partition information from the provided table and returns a JSON array representing the [partition spec](https://iceberg.apache.org/spec/?#partition-specs) of the table.
         Each element of the returned array represents a field in the [partition spec](https://iceberg.apache.org/spec/?#partition-specs) that follows [Appendix-C](https://iceberg.apache.org/spec/?#appendix-c-json-serialization) of the Iceberg specification.
@@ -425,23 +511,21 @@ class IcebergSource(StatefulIngestionSourceBase):
     def _get_dataset_properties_aspect(
         self, dataset_name: str, table: Table
     ) -> DatasetPropertiesClass:
-        additional_properties = {}
+        created: Optional[TimeStampClass] = None
         custom_properties = table.metadata.properties.copy()
         custom_properties["location"] = table.metadata.location
         custom_properties["format-version"] = str(table.metadata.format_version)
         custom_properties["partition-spec"] = str(self._get_partition_aspect(table))
+        last_modified: Optional[int] = table.metadata.last_updated_ms
         if table.current_snapshot():
             custom_properties["snapshot-id"] = str(table.current_snapshot().snapshot_id)
             custom_properties["manifest-list"] = table.current_snapshot().manifest_list
-            additional_properties["lastModified"] = TimeStampClass(
-                int(table.current_snapshot().timestamp_ms)
-            )
+            if not last_modified:
+                last_modified = int(table.current_snapshot().timestamp_ms)
         if "created-at" in custom_properties:
             try:
                 dt = dateutil_parser.isoparse(custom_properties["created-at"])
-                additional_properties["created"] = TimeStampClass(
-                    int(dt.timestamp() * 1000)
-                )
+                created = TimeStampClass(int(dt.timestamp() * 1000))
             except Exception as ex:
                 LOGGER.warning(
                     f"Exception while trying to parse creation date {custom_properties['created-at']}, ignoring: {ex}"
@@ -451,8 +535,10 @@ class IcebergSource(StatefulIngestionSourceBase):
             name=table.name()[-1],
             description=table.metadata.properties.get("comment", None),
             customProperties=custom_properties,
-            lastModified=additional_properties.get("lastModified"),
-            created=additional_properties.get("created"),
+            lastModified=TimeStampClass(last_modified)
+            if last_modified is not None
+            else None,
+            created=created,
             qualifiedName=dataset_name,
         )
 
@@ -530,7 +616,9 @@ class IcebergSource(StatefulIngestionSourceBase):
             name=namespace_repr, qualifiedName=namespace_repr, env=self.config.env
         )
         yield SubTypes(typeNames=[DatasetContainerSubTypes.NAMESPACE])
-        yield self._get_dataplatform_instance_aspect()
+        dpi = self._get_dataplatform_instance_aspect()
+        yield dpi
+        yield self._create_browse_paths_aspect(dpi.instance)
 
 
 class ToAvroSchemaIcebergVisitor(SchemaVisitorPerPrimitiveType[Dict[str, Any]]):

datahub/ingestion/source/kafka/kafka.py

@@ -568,10 +568,7 @@ class KafkaSource(StatefulIngestionSourceBase, TestableSource):
 
         for config_key in KafkaTopicConfigKeys:
             try:
-                if (
-                    config_key in topic_config.keys()
-                    and topic_config[config_key] is not None
-                ):
+                if config_key in topic_config and topic_config[config_key] is not None:
                     config_value = topic_config[config_key].value
                     custom_props[config_key] = (
                         config_value

datahub/ingestion/source/kafka_connect/sink_connectors.py

@@ -197,7 +197,7 @@ class BigQuerySinkConnector(BaseConnector):
         for name in transform_names:
             transform = {"name": name}
             transforms.append(transform)
-            for key in self.connector_manifest.config.keys():
+            for key in self.connector_manifest.config:
                 if key.startswith(f"transforms.{name}."):
                     transform[key.replace(f"transforms.{name}.", "")] = (
                         self.connector_manifest.config[key]

datahub/ingestion/source/kafka_connect/source_connectors.py

@@ -121,7 +121,7 @@ class ConfluentJDBCSourceConnector(BaseConnector):
         for name in transform_names:
             transform = {"name": name}
             transforms.append(transform)
-            for key in self.connector_manifest.config.keys():
+            for key in self.connector_manifest.config:
                 if key.startswith(f"transforms.{name}."):
                     transform[key.replace(f"transforms.{name}.", "")] = (
                         self.connector_manifest.config[key]

datahub/ingestion/source/looker/looker_source.py

@@ -363,7 +363,7 @@ class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
         filters: MutableMapping[str, Any] = (
             query.filters if query.filters is not None else {}
         )
-        for field in filters.keys():
+        for field in filters:
             if field is None:
                 continue
 
@@ -877,8 +877,7 @@ class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
             # fine to set them to None.
             # TODO: Track project names for each explore.
             explores_to_fetch = [
-                (None, model, explore)
-                for (model, explore) in self.reachable_explores.keys()
+                (None, model, explore) for (model, explore) in self.reachable_explores
             ]
         explores_to_fetch.sort()
 

datahub/ingestion/source/mlflow.py

@@ -16,7 +16,7 @@ from datahub.api.entities.dataprocess.dataprocess_instance import (
 )
 from datahub.configuration.source_common import EnvConfigMixin
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.mcp_builder import ContainerKey
+from datahub.emitter.mcp_builder import ExperimentKey
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -36,6 +36,7 @@ from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
+    StatefulStaleMetadataRemovalConfig,
 )
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
@@ -77,10 +78,6 @@ from datahub.sdk.dataset import Dataset
 T = TypeVar("T")
 
 
-class ContainerKeyWithId(ContainerKey):
-    id: str
-
-
 class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
@@ -123,6 +120,8 @@ class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
         default=None, description="Password for MLflow authentication"
     )
 
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
+
 
 @dataclass
 class MLflowRegisteredModelStageInfo:
@@ -252,7 +251,7 @@ class MLflowSource(StatefulIngestionSourceBase):
         self, experiment: Experiment
     ) -> Iterable[MetadataWorkUnit]:
         experiment_container = Container(
-            container_key=ContainerKeyWithId(
+            container_key=ExperimentKey(
                 platform=str(DataPlatformUrn(platform_name=self.platform)),
                 id=experiment.name,
             ),
@@ -470,7 +469,7 @@ class MLflowSource(StatefulIngestionSourceBase):
     def _get_run_workunits(
         self, experiment: Experiment, run: Run
     ) -> Iterable[MetadataWorkUnit]:
-        experiment_key = ContainerKeyWithId(
+        experiment_key = ExperimentKey(
             platform=str(DataPlatformUrn(self.platform)), id=experiment.name
         )
 

datahub/ingestion/source/mode.py

@@ -899,7 +899,7 @@ class ModeSource(StatefulIngestionSourceBase):
                 for match in matches:
                     definition = Template(source=match).render()
                     parameters = yaml.safe_load(definition)
-                    for key in parameters.keys():
+                    for key in parameters:
                         jinja_params[key] = parameters[key].get("default", "")
 
                 normalized_query = re.sub(
@@ -1601,7 +1601,7 @@ class ModeSource(StatefulIngestionSourceBase):
 
     def emit_chart_mces(self) -> Iterable[MetadataWorkUnit]:
         # Space/collection -> report -> query -> Chart
-        for space_token in self.space_tokens.keys():
+        for space_token in self.space_tokens:
             reports = self._get_reports(space_token)
             for report in reports:
                 report_token = report.get("token", "")

datahub/ingestion/source/nifi.py

@@ -703,7 +703,7 @@ class NifiSource(StatefulIngestionSourceBase):
             if (
                 component.nifi_type is NifiType.PROCESSOR
                 and component.type
-                not in NifiProcessorProvenanceEventAnalyzer.KNOWN_INGRESS_EGRESS_PROCESORS.keys()
+                not in NifiProcessorProvenanceEventAnalyzer.KNOWN_INGRESS_EGRESS_PROCESORS
             ) or component.nifi_type not in [
                 NifiType.PROCESSOR,
                 NifiType.REMOTE_INPUT_PORT,
@@ -977,7 +977,7 @@ class NifiSource(StatefulIngestionSourceBase):
                     )
 
             for incoming_from in incoming:
-                if incoming_from in self.nifi_flow.remotely_accessible_ports.keys():
+                if incoming_from in self.nifi_flow.remotely_accessible_ports:
                     dataset_name = f"{self.config.site_name}.{self.nifi_flow.remotely_accessible_ports[incoming_from].name}"
                     dataset_urn = builder.make_dataset_urn(
                         NIFI, dataset_name, self.config.env
@@ -994,7 +994,7 @@ class NifiSource(StatefulIngestionSourceBase):
                     )
 
             for outgoing_to in outgoing:
-                if outgoing_to in self.nifi_flow.remotely_accessible_ports.keys():
+                if outgoing_to in self.nifi_flow.remotely_accessible_ports:
                     dataset_name = f"{self.config.site_name}.{self.nifi_flow.remotely_accessible_ports[outgoing_to].name}"
                     dataset_urn = builder.make_dataset_urn(
                         NIFI, dataset_name, self.config.env

datahub/ingestion/source/openapi.py

@@ -102,7 +102,7 @@ class OpenApiConfig(ConfigModel):
                 # details there once, and then use that session for all requests.
                 self.token = f"Bearer {self.bearer_token}"
             else:
-                assert "url_complement" in self.get_token.keys(), (
+                assert "url_complement" in self.get_token, (
                     "When 'request_type' is set to 'get', an url_complement is needed for the request."
                 )
                 if self.get_token["request_type"] == "get":
@@ -317,7 +317,7 @@ class APISource(Source, ABC):
                 yield wu
 
             # Handle schema metadata if available
-            if "data" in endpoint_dets.keys():
+            if "data" in endpoint_dets:
                 # we are lucky! data is defined in the swagger for this endpoint
                 schema_metadata = set_metadata(dataset_name, endpoint_dets["data"])
                 wu = MetadataWorkUnit(
@@ -371,7 +371,7 @@ class APISource(Source, ABC):
                 else:
                     self.report_bad_responses(response.status_code, type=endpoint_k)
             else:
-                if endpoint_k not in config.forced_examples.keys():
+                if endpoint_k not in config.forced_examples:
                     # start guessing...
                     url_guess = try_guessing(endpoint_k, root_dataset_samples)
                     tot_url = clean_url(config.url + self.url_basepath + url_guess)

datahub/ingestion/source/openapi_parser.py

@@ -128,18 +128,18 @@ def get_endpoints(sw_dict: dict) -> dict:
 
     for p_k, p_o in sw_dict["paths"].items():
         method = list(p_o)[0]
-        if "200" in p_o[method]["responses"].keys():
+        if "200" in p_o[method]["responses"]:
             base_res = p_o[method]["responses"]["200"]
-        elif 200 in p_o[method]["responses"].keys():
+        elif 200 in p_o[method]["responses"]:
             # if you read a plain yml file the 200 will be an integer
             base_res = p_o[method]["responses"][200]
         else:
             # the endpoint does not have a 200 response
             continue
 
-        if "description" in p_o[method].keys():
+        if "description" in p_o[method]:
             desc = p_o[method]["description"]
-        elif "summary" in p_o[method].keys():
+        elif "summary" in p_o[method]:
             desc = p_o[method]["summary"]
         else:  # still testing
             desc = ""
@@ -156,7 +156,7 @@ def get_endpoints(sw_dict: dict) -> dict:
             url_details[p_k]["data"] = example_data
 
         # checking whether there are defined parameters to execute the call...
-        if "parameters" in p_o[method].keys():
+        if "parameters" in p_o[method]:
             url_details[p_k]["parameters"] = p_o[method]["parameters"]
 
     return dict(sorted(url_details.items()))
@@ -169,7 +169,7 @@ def check_for_api_example_data(base_res: dict, key: str) -> dict:
     data = {}
     if "content" in base_res:
         res_cont = base_res["content"]
-        if "application/json" in res_cont.keys():
+        if "application/json" in res_cont:
             ex_field = None
             if "example" in res_cont["application/json"]:
                 ex_field = "example"
@@ -186,7 +186,7 @@ def check_for_api_example_data(base_res: dict, key: str) -> dict:
                 logger.warning(
                     f"Field in swagger file does not give consistent data --- {key}"
                 )
-        elif "text/csv" in res_cont.keys():
+        elif "text/csv" in res_cont:
             data = res_cont["text/csv"]["schema"]
     elif "examples" in base_res:
         data = base_res["examples"]["application/json"]
@@ -239,7 +239,7 @@ def guessing_url_name(url: str, examples: dict) -> str:
 
     # substituting the parameter's name w the value
     for name, clean_name in zip(needed_n, cleaned_needed_n):
-        if clean_name in examples[ex2use].keys():
+        if clean_name in examples[ex2use]:
             guessed_url = re.sub(name, str(examples[ex2use][clean_name]), guessed_url)
 
     return guessed_url

datahub/ingestion/source/powerbi/config.py

@@ -555,7 +555,7 @@ class PowerBiDashboardSourceConfig(
     def map_data_platform(cls, value):
         # For backward compatibility convert input PostgreSql to PostgreSQL
         # PostgreSQL is name of the data-platform in M-Query
-        if "PostgreSql" in value.keys():
+        if "PostgreSql" in value:
             platform_name = value["PostgreSql"]
             del value["PostgreSql"]
             value["PostgreSQL"] = platform_name

datahub/ingestion/source/powerbi/powerbi.py

@@ -94,7 +94,7 @@ from datahub.metadata.schema_classes import (
     UpstreamLineageClass,
     ViewPropertiesClass,
 )
-from datahub.metadata.urns import ChartUrn
+from datahub.metadata.urns import ChartUrn, DatasetUrn
 from datahub.sql_parsing.sqlglot_lineage import ColumnLineageInfo
 from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.urns.urn_iter import lowercase_dataset_urn
@@ -263,7 +263,7 @@ class Mapper:
             for upstream_dpt in lineage.upstreams:
                 if (
                     upstream_dpt.data_platform_pair.powerbi_data_platform_name
-                    not in self.__config.dataset_type_mapping.keys()
+                    not in self.__config.dataset_type_mapping
                 ):
                     logger.debug(
                         f"Skipping upstream table for {ds_urn}. The platform {upstream_dpt.data_platform_pair.powerbi_data_platform_name} is not part of dataset_type_mapping",
@@ -1083,6 +1083,7 @@ class Mapper:
         report: powerbi_data_classes.Report,
         chart_mcps: List[MetadataChangeProposalWrapper],
         user_mcps: List[MetadataChangeProposalWrapper],
+        dataset_edges: List[EdgeClass],
     ) -> List[MetadataChangeProposalWrapper]:
         """
         Map PowerBi report to Datahub dashboard
@@ -1104,6 +1105,7 @@ class Mapper:
             charts=chart_urn_list,
             lastModified=ChangeAuditStamps(),
             dashboardUrl=report.webUrl,
+            datasetEdges=dataset_edges,
         )
 
         info_mcp = self.new_mcp(
@@ -1197,12 +1199,23 @@ class Mapper:
         ds_mcps = self.to_datahub_dataset(report.dataset, workspace)
         chart_mcps = self.pages_to_chart(report.pages, workspace, ds_mcps)
 
+        # collect all upstream datasets; using a set to retain unique urns
+        dataset_urns = {
+            dataset.entityUrn
+            for dataset in ds_mcps
+            if dataset.entityType == DatasetUrn.ENTITY_TYPE and dataset.entityUrn
+        }
+        dataset_edges = [
+            EdgeClass(destinationUrn=dataset_urn) for dataset_urn in dataset_urns
+        ]
+
         # Let's convert report to datahub dashboard
         report_mcps = self.report_to_dashboard(
             workspace=workspace,
             report=report,
             chart_mcps=chart_mcps,
             user_mcps=user_mcps,
+            dataset_edges=dataset_edges,
         )
 
         # Now add MCPs in sequence
@@ -1340,7 +1353,7 @@ class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
             for data_platform in SupportedDataPlatform
         ]
 
-        for key in self.source_config.dataset_type_mapping.keys():
+        for key in self.source_config.dataset_type_mapping:
             if key not in powerbi_data_platforms:
                 raise ValueError(f"PowerBI DataPlatform {key} is not supported")
 

datahub/ingestion/source/redshift/profile.py

@@ -42,9 +42,9 @@ class RedshiftProfiler(GenericProfiler):
                 "max_overflow", self.config.profiling.max_workers
             )
 
-        for db in tables.keys():
+        for db in tables:
             profile_requests = []
-            for schema in tables.get(db, {}).keys():
+            for schema in tables.get(db, {}):
                 if not self.config.schema_pattern.allowed(schema):
                     continue
                 for table in tables[db].get(schema, {}):

datahub/ingestion/source/sigma/sigma.py

@@ -170,7 +170,9 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
             if self.config.workspace_pattern.allowed(workspace.name):
                 allowed_workspaces.append(workspace)
             else:
-                self.reporter.workspaces.dropped(workspace.workspaceId)
+                self.reporter.workspaces.dropped(
+                    f"{workspace.name} ({workspace.workspaceId})"
+                )
         logger.info(f"Number of allowed workspaces = {len(allowed_workspaces)}")
 
         return allowed_workspaces
@@ -661,7 +663,9 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
             yield from self._gen_workbook_workunit(workbook)
 
         for workspace in self._get_allowed_workspaces():
-            self.reporter.workspaces.processed(workspace.workspaceId)
+            self.reporter.workspaces.processed(
+                f"{workspace.name} ({workspace.workspaceId})"
+            )
             yield from self._gen_workspace_workunit(workspace)
         yield from self._gen_sigma_dataset_upstream_lineage_workunit()
 

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -77,7 +77,7 @@ class SnowsightUrlBuilder:
         region: str,
     ) -> Tuple[str, str]:
         cloud: str
-        if region in SNOWFLAKE_REGION_CLOUD_REGION_MAPPING.keys():
+        if region in SNOWFLAKE_REGION_CLOUD_REGION_MAPPING:
             cloud, cloud_region_id = SNOWFLAKE_REGION_CLOUD_REGION_MAPPING[region]
         elif region.startswith(("aws_", "gcp_", "azure_")):
             # e.g. aws_us_west_2, gcp_us_central1, azure_northeurope

datahub/ingestion/source/sql/stored_procedures/base.py

@@ -26,6 +26,7 @@ from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
     DataTransformClass,
     DataTransformLogicClass,
+    QueryLanguageClass,
     QueryStatementClass,
     SubTypesClass,
 )
@@ -176,7 +177,17 @@ def _generate_job_workunits(
                     DataTransformClass(
                         queryStatement=QueryStatementClass(
                             value=procedure.procedure_definition,
-                            language=procedure.language,
+                            language=(
+                                QueryLanguageClass.SQL
+                                if procedure.language == "SQL"
+                                # The language field uses a pretty limited enum.
+                                # The "UNKNOWN" enum value is pretty new, so we don't want to
+                                # emit it until it has broader server-side support. As a
+                                # short-term solution, we map all languages to "SQL".
+                                # TODO: Once we've released server 1.1.0, we should change
+                                # this to be "UNKNOWN" for all languages except "SQL".
+                                else QueryLanguageClass.SQL
+                            ),
                         ),
                     )
                 ]

datahub/ingestion/source/sql/trino.py

@@ -128,9 +128,10 @@ def get_table_comment(self, connection, table_name: str, schema: str = None, **k
         if catalog_name is None:
             raise exc.NoSuchTableError("catalog is required in connection")
         connector_name = get_catalog_connector_name(connection.engine, catalog_name)
-        if connector_name is None:
-            return {}
-        if connector_name in PROPERTIES_TABLE_SUPPORTED_CONNECTORS:
+        if (
+            connector_name is not None
+            and connector_name in PROPERTIES_TABLE_SUPPORTED_CONNECTORS
+        ):
             properties_table = self._get_full_table(f"{table_name}$properties", schema)
             query = f"SELECT * FROM {properties_table}"
             row = connection.execute(sql.text(query)).fetchone()

datahub/ingestion/source/state/stale_entity_removal_handler.py

@@ -45,7 +45,6 @@ class StatefulStaleMetadataRemovalConfig(StatefulIngestionConfig):
         description="Prevents large amount of soft deletes & the state from committing from accidental changes to the source configuration if the relative change percent in entities compared to the previous state is above the 'fail_safe_threshold'.",
         le=100.0,
         ge=0.0,
-        hidden_from_docs=True,
     )