acryl-datahub 0.15.0.3rc1__py3-none-any.whl → 0.15.0.4__py3-none-any.whl

datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py

@@ -71,13 +71,13 @@ class Workspace:
     id: str
     name: str
     type: str  # This is used as a subtype of the Container entity.
-    dashboards: List["Dashboard"]
-    reports: List["Report"]
-    datasets: Dict[str, "PowerBIDataset"]
-    report_endorsements: Dict[str, List[str]]
-    dashboard_endorsements: Dict[str, List[str]]
+    dashboards: Dict[str, "Dashboard"]  # key = dashboard id
+    reports: Dict[str, "Report"]  # key = report id
+    datasets: Dict[str, "PowerBIDataset"]  # key = dataset id
+    report_endorsements: Dict[str, List[str]]  # key = report id
+    dashboard_endorsements: Dict[str, List[str]]  # key = dashboard id
     scan_result: dict
-    independent_datasets: List["PowerBIDataset"]
+    independent_datasets: Dict[str, "PowerBIDataset"]  # key = dataset id
     app: Optional["App"]
 
     def get_urn_part(self, workspace_id_as_urn_part: Optional[bool] = False) -> str:

datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py

@@ -193,15 +193,18 @@ class PowerBiAPI:
     def get_report_users(self, workspace_id: str, report_id: str) -> List[User]:
         return self._get_entity_users(workspace_id, Constant.REPORTS, report_id)
 
-    def get_reports(self, workspace: Workspace) -> List[Report]:
+    def get_reports(self, workspace: Workspace) -> Dict[str, Report]:
         """
         Fetch the report from PowerBi for the given Workspace
         """
-        reports: List[Report] = []
+        reports: Dict[str, Report] = {}
         try:
-            reports = self._get_resolver().get_reports(workspace)
+            reports = {
+                report.id: report
+                for report in self._get_resolver().get_reports(workspace)
+            }
             # Fill Report dataset
-            for report in reports:
+            for report in reports.values():
                 if report.dataset_id:
                     report.dataset = self.dataset_registry.get(report.dataset_id)
                     if report.dataset is None:
@@ -222,7 +225,7 @@ class PowerBiAPI:
                 )
                 return
 
-            for report in reports:
+            for report in reports.values():
                 report.users = self.get_report_users(
                     workspace_id=workspace.id, report_id=report.id
                 )
@@ -234,7 +237,7 @@ class PowerBiAPI:
                 )
                 return
 
-            for report in reports:
+            for report in reports.values():
                 report.tags = workspace.report_endorsements.get(report.id, [])
 
         fill_ownership()
@@ -270,12 +273,12 @@ class PowerBiAPI:
                 name=workspace[Constant.NAME],
                 type=workspace[Constant.TYPE],
                 datasets={},
-                dashboards=[],
-                reports=[],
+                dashboards={},
+                reports={},
                 report_endorsements={},
                 dashboard_endorsements={},
                 scan_result={},
-                independent_datasets=[],
+                independent_datasets={},
                 app=None,  # It will be populated in _fill_metadata_from_scan_result method
             )
             for workspace in groups
@@ -561,12 +564,12 @@ class PowerBiAPI:
                 name=workspace_metadata[Constant.NAME],
                 type=workspace_metadata[Constant.TYPE],
                 datasets={},
-                dashboards=[],
-                reports=[],
+                dashboards={},
+                reports={},
                 report_endorsements={},
                 dashboard_endorsements={},
                 scan_result={},
-                independent_datasets=[],
+                independent_datasets={},
                 app=None,  # It is getting set from scan-result
             )
             cur_workspace.scan_result = workspace_metadata
@@ -597,25 +600,28 @@ class PowerBiAPI:
     def _fill_independent_datasets(self, workspace: Workspace) -> None:
         reachable_datasets: List[str] = []
         # Find out reachable datasets
-        for dashboard in workspace.dashboards:
+        for dashboard in workspace.dashboards.values():
             for tile in dashboard.tiles:
                 if tile.dataset is not None:
                     reachable_datasets.append(tile.dataset.id)
 
-        for report in workspace.reports:
+        for report in workspace.reports.values():
             if report.dataset is not None:
                 reachable_datasets.append(report.dataset.id)
 
         # Set datasets not present in reachable_datasets
         for dataset in workspace.datasets.values():
             if dataset.id not in reachable_datasets:
-                workspace.independent_datasets.append(dataset)
+                workspace.independent_datasets[dataset.id] = dataset
 
     def _fill_regular_metadata_detail(self, workspace: Workspace) -> None:
         def fill_dashboards() -> None:
-            workspace.dashboards = self._get_resolver().get_dashboards(workspace)
+            workspace.dashboards = {
+                dashboard.id: dashboard
+                for dashboard in self._get_resolver().get_dashboards(workspace)
+            }
             # set tiles of Dashboard
-            for dashboard in workspace.dashboards:
+            for dashboard in workspace.dashboards.values():
                 dashboard.tiles = self._get_resolver().get_tiles(
                     workspace, dashboard=dashboard
                 )
@@ -644,7 +650,7 @@ class PowerBiAPI:
                     "Skipping tag retrieval for dashboard as extract_endorsements_to_tags is set to false"
                 )
                 return
-            for dashboard in workspace.dashboards:
+            for dashboard in workspace.dashboards.values():
                 dashboard.tags = workspace.dashboard_endorsements.get(dashboard.id, [])
 
         if self.__config.extract_dashboards:

datahub/ingestion/source/s3/source.py

@@ -40,7 +40,6 @@ from datahub.ingestion.source.aws.s3_util import (
     get_bucket_name,
     get_bucket_relative_path,
     get_key_prefix,
-    group_s3_objects_by_dirname,
     strip_s3_prefix,
 )
 from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
@@ -73,6 +72,7 @@ from datahub.metadata.schema_classes import (
     _Aspect,
 )
 from datahub.telemetry import stats, telemetry
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.perf_timer import PerfTimer
 
 if TYPE_CHECKING:
@@ -868,7 +868,11 @@ class S3Source(StatefulIngestionSourceBase):
         """
         partitions: List[Folder] = []
         s3_objects = bucket.objects.filter(Prefix=prefix).page_size(PAGE_SIZE)
-        for key, group in group_s3_objects_by_dirname(s3_objects).items():
+        grouped_s3_objects_by_dirname = groupby_unsorted(
+            s3_objects,
+            key=lambda obj: obj.key.rsplit("/", 1)[0],
+        )
+        for key, group in grouped_s3_objects_by_dirname:
             file_size = 0
             creation_time = None
             modification_time = None

datahub/ingestion/source/slack/slack.py

@@ -5,6 +5,8 @@ from typing import Iterable, List, Optional, Tuple
 
 from pydantic import Field, SecretStr
 from slack_sdk import WebClient
+from tenacity import retry, wait_exponential
+from tenacity.before_sleep import before_sleep_log
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import ConfigModel
@@ -294,6 +296,10 @@ class SlackSource(Source):
                 return
             raise e
 
+    @retry(
+        wait=wait_exponential(multiplier=2, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.ERROR, True),
+    )
     def get_user_to_be_updated(self) -> Iterable[CorpUser]:
         graphql_query = textwrap.dedent(
             """

datahub/ingestion/source/sql/hive_metastore.py

@@ -2,7 +2,6 @@ import base64
 import json
 import logging
 from collections import namedtuple
-from itertools import groupby
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from pydantic.dataclasses import dataclass
@@ -58,6 +57,7 @@ from datahub.metadata.schema_classes import (
     SubTypesClass,
     ViewPropertiesClass,
 )
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
 from datahub.utilities.str_enum import StrEnum
 
@@ -490,7 +490,7 @@ class HiveMetastoreSource(SQLAlchemySource):
 
         iter_res = self._alchemy_client.execute_query(statement)
 
-        for key, group in groupby(iter_res, self._get_table_key):
+        for key, group in groupby_unsorted(iter_res, self._get_table_key):
             schema_name = (
                 f"{db_name}.{key.schema}"
                 if self.config.include_catalog_name_in_ids
@@ -647,7 +647,7 @@ class HiveMetastoreSource(SQLAlchemySource):
         )
 
         iter_res = self._alchemy_client.execute_query(statement)
-        for key, group in groupby(iter_res, self._get_table_key):
+        for key, group in groupby_unsorted(iter_res, self._get_table_key):
             db_name = self.get_db_name(inspector)
 
             schema_name = (

datahub/ingestion/source/sql/mssql/job_models.py

@@ -156,7 +156,7 @@ class MSSQLDataJob:
     entity: Union[StoredProcedure, JobStep]
     type: str = "dataJob"
     source: str = "mssql"
-    external_url: str = ""
+    external_url: Optional[str] = None
     description: Optional[str] = None
     status: Optional[str] = None
     incoming: List[str] = field(default_factory=list)
@@ -228,7 +228,7 @@ class MSSQLDataFlow:
     entity: Union[MSSQLJob, MSSQLProceduresContainer]
     type: str = "dataFlow"
     source: str = "mssql"
-    external_url: str = ""
+    external_url: Optional[str] = None
     flow_properties: Dict[str, str] = field(default_factory=dict)
 
     def add_property(

datahub/ingestion/source/sql/mssql/source.py

@@ -11,6 +11,7 @@ from sqlalchemy.engine.base import Connection
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.exc import ProgrammingError, ResourceClosedError
 
+import datahub.metadata.schema_classes as models
 from datahub.configuration.common import AllowDenyPattern
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -49,21 +50,15 @@ from datahub.ingestion.source.sql.sql_config import (
     make_sqlalchemy_uri,
 )
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
-from datahub.metadata.schema_classes import (
-    BooleanTypeClass,
-    NumberTypeClass,
-    StringTypeClass,
-    UnionTypeClass,
-)
 from datahub.utilities.file_backed_collections import FileBackedList
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-register_custom_type(sqlalchemy.dialects.mssql.BIT, BooleanTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.MONEY, NumberTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.SMALLMONEY, NumberTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.SQL_VARIANT, UnionTypeClass)
-register_custom_type(sqlalchemy.dialects.mssql.UNIQUEIDENTIFIER, StringTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.BIT, models.BooleanTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.MONEY, models.NumberTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.SMALLMONEY, models.NumberTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.SQL_VARIANT, models.UnionTypeClass)
+register_custom_type(sqlalchemy.dialects.mssql.UNIQUEIDENTIFIER, models.StringTypeClass)
 
 
 class SQLServerConfig(BasicSQLAlchemyConfig):
@@ -651,6 +646,26 @@ class SQLServerSource(SQLAlchemySource):
                 entityUrn=data_job.urn,
                 aspect=data_job.as_datajob_input_output_aspect,
             ).as_workunit()
+
+        if (
+            self.config.include_stored_procedures_code
+            and isinstance(data_job.entity, StoredProcedure)
+            and data_job.entity.code is not None
+        ):
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_job.urn,
+                aspect=models.DataTransformLogicClass(
+                    transforms=[
+                        models.DataTransformClass(
+                            queryStatement=models.QueryStatementClass(
+                                value=data_job.entity.code,
+                                language=models.QueryLanguageClass.SQL,
+                            ),
+                        )
+                    ]
+                ),
+            ).as_workunit()
+
         # TODO: Add SubType when it appear
 
     def construct_flow_workunits(

datahub/ingestion/source/sql/teradata.py

@@ -3,7 +3,6 @@ from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime
 from functools import lru_cache
-from itertools import groupby
 from typing import (
     Any,
     Dict,
@@ -59,6 +58,7 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
 from datahub.metadata.schema_classes import SchemaMetadataClass
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.utilities.groupby import groupby_unsorted
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -286,7 +286,7 @@ def optimized_get_foreign_keys(self, connection, table_name, schema=None, **kw):
 
     # TODO: Check if there's a better way
     fk_dicts = list()
-    for constraint_info, constraint_cols in groupby(res, grouper):
+    for constraint_info, constraint_cols in groupby_unsorted(res, grouper):
         fk_dict = {
             "name": str(constraint_info["name"]),
             "constrained_columns": list(),

datahub/ingestion/source/tableau/tableau.py

@@ -1147,23 +1147,36 @@ class TableauSiteSource:
                 )
             # Set parent project name
             for _project_id, project in all_project_map.items():
-                if (
-                    project.parent_id is not None
-                    and project.parent_id in all_project_map
-                ):
+                if project.parent_id is None:
+                    continue
+
+                if project.parent_id in all_project_map:
                     project.parent_name = all_project_map[project.parent_id].name
+                else:
+                    self.report.warning(
+                        title="Incomplete project hierarchy",
+                        message="Project details missing. Child projects will be ingested without reference to their parent project. We generally need Site Administrator Explorer permissions to extract the complete project hierarchy.",
+                        context=f"Missing {project.parent_id}, referenced by {project.id} {project.project_name}",
+                    )
+                    project.parent_id = None
+
+            # Post-condition
+            assert all(
+                [
+                    ((project.parent_id is None) == (project.parent_name is None))
+                    and (
+                        project.parent_id is None
+                        or project.parent_id in all_project_map
+                    )
+                    for project in all_project_map.values()
+                ]
+            ), "Parent project id and name should be consistent"
 
         def set_project_path():
             def form_path(project_id: str) -> List[str]:
                 cur_proj = all_project_map[project_id]
                 ancestors = [cur_proj.name]
                 while cur_proj.parent_id is not None:
-                    if cur_proj.parent_id not in all_project_map:
-                        self.report.warning(
-                            "project-issue",
-                            f"Parent project {cur_proj.parent_id} not found. We need Site Administrator Explorer permissions.",
-                        )
-                        break
                     cur_proj = all_project_map[cur_proj.parent_id]
                     ancestors = [cur_proj.name, *ancestors]
                 return ancestors