acryl-datahub 1.0.0.3rc9__py3-none-any.whl → 1.0.0.3rc11__py3-none-any.whl

datahub/ingestion/source/mode.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import lru_cache
 from json import JSONDecodeError
-from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
 
 import dateutil.parser as dp
 import pydantic
@@ -203,6 +203,10 @@ class HTTPError429(HTTPError):
     pass
 
 
+class HTTPError504(HTTPError):
+    pass
+
+
 ModeRequestError = (HTTPError, JSONDecodeError)
 
 
@@ -217,6 +221,9 @@ class ModeSourceReport(StaleEntityRemovalSourceReport):
     num_query_template_render: int = 0
     num_query_template_render_failures: int = 0
     num_query_template_render_success: int = 0
+    num_requests_exceeding_rate_limit: int = 0
+    num_requests_retried_on_timeout: int = 0
+    num_spaces_retrieved: int = 0
 
     def report_dropped_space(self, ent_name: str) -> None:
         self.filtered_spaces.append(ent_name)
@@ -456,9 +463,23 @@ class ModeSource(StatefulIngestionSourceBase):
         # Datasets
         datasets = []
         for imported_dataset_name in report_info.get("imported_datasets", {}):
-            mode_dataset = self._get_request_json(
-                f"{self.workspace_uri}/reports/{imported_dataset_name.get('token')}"
-            )
+            try:
+                mode_dataset = self._get_request_json(
+                    f"{self.workspace_uri}/reports/{imported_dataset_name.get('token')}"
+                )
+            except HTTPError as http_error:
+                status_code = http_error.response.status_code
+                if status_code == 404:
+                    self.report.report_warning(
+                        title="Report Not Found",
+                        message="Referenced report for reusable dataset was not found.",
+                        context=f"Report: {report_info.get('id')}, "
+                        f"Imported Dataset Report: {imported_dataset_name.get('token')}",
+                    )
+                    continue
+                else:
+                    raise http_error
+
             dataset_urn = builder.make_dataset_urn_with_platform_instance(
                 self.platform,
                 str(mode_dataset.get("id")),
@@ -562,29 +583,34 @@ class ModeSource(StatefulIngestionSourceBase):
         space_info = {}
         try:
             logger.debug(f"Retrieving spaces for {self.workspace_uri}")
-            payload = self._get_request_json(f"{self.workspace_uri}/spaces?filter=all")
-            spaces = payload.get("_embedded", {}).get("spaces", {})
-            logger.debug(
-                f"Got {len(spaces)} spaces from workspace {self.workspace_uri}"
-            )
-            for s in spaces:
-                logger.debug(f"Space: {s.get('name')}")
-                space_name = s.get("name", "")
-                # Using both restricted and default_access_level because
-                # there is a current bug with restricted returning False everytime
-                # which has been reported to Mode team
-                if self.config.exclude_restricted and (
-                    s.get("restricted") or s.get("default_access_level") == "restricted"
-                ):
-                    logging.debug(
-                        f"Skipping space {space_name} due to exclude restricted"
-                    )
-                    continue
-                if not self.config.space_pattern.allowed(space_name):
-                    self.report.report_dropped_space(space_name)
-                    logging.debug(f"Skipping space {space_name} due to space pattern")
-                    continue
-                space_info[s.get("token", "")] = s.get("name", "")
+            for spaces_page in self._get_paged_request_json(
+                f"{self.workspace_uri}/spaces?filter=all", "spaces", 30
+            ):
+                logger.debug(
+                    f"Read {len(spaces_page)} spaces records from workspace {self.workspace_uri}"
+                )
+                self.report.num_spaces_retrieved += len(spaces_page)
+                for s in spaces_page:
+                    logger.debug(f"Space: {s.get('name')}")
+                    space_name = s.get("name", "")
+                    # Using both restricted and default_access_level because
+                    # there is a current bug with restricted returning False everytime
+                    # which has been reported to Mode team
+                    if self.config.exclude_restricted and (
+                        s.get("restricted")
+                        or s.get("default_access_level") == "restricted"
+                    ):
+                        logging.debug(
+                            f"Skipping space {space_name} due to exclude restricted"
+                        )
+                        continue
+                    if not self.config.space_pattern.allowed(space_name):
+                        self.report.report_dropped_space(space_name)
+                        logging.debug(
+                            f"Skipping space {space_name} due to space pattern"
+                        )
+                        continue
+                    space_info[s.get("token", "")] = s.get("name", "")
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Spaces",
@@ -1475,13 +1501,28 @@ class ModeSource(StatefulIngestionSourceBase):
             )
         return charts
 
+    def _get_paged_request_json(
+        self, url: str, key: str, per_page: int
+    ) -> Iterator[List[Dict]]:
+        page: int = 1
+        while True:
+            page_url = f"{url}&per_page={per_page}&page={page}"
+            response = self._get_request_json(page_url)
+            data: List[Dict] = response.get("_embedded", {}).get(key, [])
+            if not data:
+                break
+            yield data
+            page += 1
+
     def _get_request_json(self, url: str) -> Dict:
         r = tenacity.Retrying(
             wait=wait_exponential(
                 multiplier=self.config.api_options.retry_backoff_multiplier,
                 max=self.config.api_options.max_retry_interval,
             ),
-            retry=retry_if_exception_type((HTTPError429, ConnectionError)),
+            retry=retry_if_exception_type(
+                (HTTPError429, HTTPError504, ConnectionError)
+            ),
             stop=stop_after_attempt(self.config.api_options.max_attempts),
         )
 
@@ -1502,11 +1543,16 @@ class ModeSource(StatefulIngestionSourceBase):
             except HTTPError as http_error:
                 error_response = http_error.response
                 if error_response.status_code == 429:
+                    self.report.num_requests_exceeding_rate_limit += 1
                     # respect Retry-After
                     sleep_time = error_response.headers.get("retry-after")
                     if sleep_time is not None:
                         time.sleep(float(sleep_time))
                     raise HTTPError429 from None
+                elif error_response.status_code == 504:
+                    self.report.num_requests_retried_on_timeout += 1
+                    time.sleep(0.1)
+                    raise HTTPError504 from None
 
                 logger.debug(
                     f"Error response ({error_response.status_code}): {error_response.text}"

datahub/ingestion/source/neo4j/neo4j_source.py

@@ -5,27 +5,35 @@ from typing import Any, Dict, Iterable, List, Optional, Type, Union
 
 import pandas as pd
 from neo4j import GraphDatabase
-from pydantic.fields import Field
+from pydantic import Field
 
 from datahub.configuration.source_common import (
     EnvConfigMixin,
+    PlatformInstanceConfigMixin,
+)
+from datahub.emitter.mce_builder import (
+    make_data_platform_urn,
+    make_dataset_urn_with_platform_instance,
 )
-from datahub.emitter.mce_builder import make_data_platform_urn, make_dataset_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
 )
 from datahub.ingestion.api.source import (
     MetadataWorkUnitProcessor,
+    SourceCapability,
 )
+from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
+    StatefulStaleMetadataRemovalConfig,
 )
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
@@ -64,12 +72,16 @@ _type_mapping: Dict[Union[Type, str], Type] = {
 }
 
 
-class Neo4jConfig(EnvConfigMixin, StatefulIngestionConfigBase):
+class Neo4jConfig(
+    StatefulIngestionConfigBase, EnvConfigMixin, PlatformInstanceConfigMixin
+):
     username: str = Field(description="Neo4j Username")
     password: str = Field(description="Neo4j Password")
     uri: str = Field(description="The URI for the Neo4j server")
     env: str = Field(description="Neo4j env")
 
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
+
 
 @dataclass
 class Neo4jSourceReport(StatefulIngestionReport):
@@ -79,21 +91,27 @@ class Neo4jSourceReport(StatefulIngestionReport):
 
 @platform_name("Neo4j", id="neo4j")
 @config_class(Neo4jConfig)
+@capability(
+    SourceCapability.PLATFORM_INSTANCE, "Supported via the `platform_instance` config"
+)
 @support_status(SupportStatus.CERTIFIED)
 class Neo4jSource(StatefulIngestionSourceBase):
     NODE = "node"
     RELATIONSHIP = "relationship"
-    PLATFORM = "neo4j"
+    config: Neo4jConfig
+    report: Neo4jSourceReport
 
-    def __init__(self, ctx: PipelineContext, config: Neo4jConfig):
+    def __init__(self, config: Neo4jConfig, ctx: PipelineContext):
+        super().__init__(config, ctx)
         self.ctx = ctx
         self.config = config
+        self.platform = "neo4j"
         self.report: Neo4jSourceReport = Neo4jSourceReport()
 
     @classmethod
-    def create(cls, config_dict, ctx):
+    def create(cls, config_dict: Dict, ctx: PipelineContext) -> "Neo4jSource":
         config = Neo4jConfig.parse_obj(config_dict)
-        return cls(ctx, config)
+        return cls(config, ctx)
 
     def get_field_type(self, attribute_type: Union[type, str]) -> SchemaFieldDataType:
         type_class: type = _type_mapping.get(attribute_type, NullTypeClass)
@@ -123,34 +141,40 @@ class Neo4jSource(StatefulIngestionSourceBase):
         dataset: str,
         description: Optional[str] = None,
         custom_properties: Optional[Dict[str, str]] = None,
-    ) -> MetadataChangeProposalWrapper:
+    ) -> Iterable[MetadataWorkUnit]:
         dataset_properties = DatasetPropertiesClass(
             description=description,
             customProperties=custom_properties,
         )
-        return MetadataChangeProposalWrapper(
-            entityUrn=make_dataset_urn(
-                platform=self.PLATFORM, name=dataset, env=self.config.env
+        yield MetadataChangeProposalWrapper(
+            entityUrn=make_dataset_urn_with_platform_instance(
+                platform=self.platform,
+                name=dataset,
+                platform_instance=self.config.platform_instance,
+                env=self.config.env,
             ),
             aspect=dataset_properties,
-        )
+        ).as_workunit()
 
     def generate_neo4j_object(
         self, dataset: str, columns: list, obj_type: Optional[str] = None
-    ) -> MetadataChangeProposalWrapper:
+    ) -> Optional[MetadataChangeProposalWrapper]:
         try:
             fields = [
                 self.get_schema_field_class(key, value.lower(), obj_type=obj_type)
                 for d in columns
                 for key, value in d.items()
             ]
-            mcp = MetadataChangeProposalWrapper(
-                entityUrn=make_dataset_urn(
-                    platform=self.PLATFORM, name=dataset, env=self.config.env
+            return MetadataChangeProposalWrapper(
+                entityUrn=make_dataset_urn_with_platform_instance(
+                    platform=self.platform,
+                    name=dataset,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
                 ),
                 aspect=SchemaMetadataClass(
                     schemaName=dataset,
-                    platform=make_data_platform_urn(self.PLATFORM),
+                    platform=make_data_platform_urn(self.platform),
                     version=0,
                     hash="",
                     platformSchema=OtherSchemaClass(rawSchema=""),
@@ -161,13 +185,16 @@ class Neo4jSource(StatefulIngestionSourceBase):
                     fields=fields,
                 ),
             )
-            self.report.obj_created += 1
         except Exception as e:
             log.error(e)
-            self.report.obj_failures += 1
-        return mcp
+            self.report.report_failure(
+                message="Failed to process dataset",
+                context=dataset,
+                exc=e,
+            )
+            return None
 
-    def get_neo4j_metadata(self, query: str) -> pd.DataFrame:
+    def get_neo4j_metadata(self, query: str) -> Optional[pd.DataFrame]:
         driver = GraphDatabase.driver(
             self.config.uri, auth=(self.config.username, self.config.password)
         )
@@ -201,13 +228,14 @@ class Neo4jSource(StatefulIngestionSourceBase):
 
             union_cols = ["key", "obj_type", "property_data_types", "description"]
             df = pd.concat([node_df[union_cols], rel_df[union_cols]])
+            return df
         except Exception as e:
             self.report.failure(
                 message="Failed to get neo4j metadata",
                 exc=e,
             )
 
-        return df
+        return None
 
     def process_nodes(self, data: list) -> pd.DataFrame:
         nodes = [record for record in data if record["value"]["type"] == self.NODE]
@@ -306,46 +334,48 @@ class Neo4jSource(StatefulIngestionSourceBase):
         df = self.get_neo4j_metadata(
             "CALL apoc.meta.schema() YIELD value UNWIND keys(value) AS key RETURN key, value[key] AS value;"
         )
+        if df is None:
+            log.warning("No metadata retrieved from Neo4j")
+            return
+
         for _, row in df.iterrows():
             try:
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=self.generate_neo4j_object(
-                        columns=row["property_data_types"],
-                        dataset=row["key"],
-                    ),
-                    is_primary_source=True,
+                neo4j_obj = self.generate_neo4j_object(
+                    columns=row["property_data_types"],
+                    dataset=row["key"],
                 )
-
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=MetadataChangeProposalWrapper(
-                        entityUrn=make_dataset_urn(
-                            platform=self.PLATFORM,
-                            name=row["key"],
-                            env=self.config.env,
-                        ),
-                        aspect=SubTypesClass(
-                            typeNames=[
-                                DatasetSubTypes.NEO4J_NODE
-                                if row["obj_type"] == self.NODE
-                                else DatasetSubTypes.NEO4J_RELATIONSHIP
-                            ]
-                        ),
+                if neo4j_obj:
+                    yield from auto_workunit([neo4j_obj])
+
+                yield MetadataChangeProposalWrapper(
+                    entityUrn=make_dataset_urn_with_platform_instance(
+                        platform=self.platform,
+                        name=row["key"],
+                        platform_instance=self.config.platform_instance,
+                        env=self.config.env,
                     ),
-                )
-
-                yield MetadataWorkUnit(
-                    id=row["key"],
-                    mcp=self.add_properties(
-                        dataset=row["key"],
-                        custom_properties=None,
-                        description=row["description"],
+                    aspect=SubTypesClass(
+                        typeNames=[
+                            DatasetSubTypes.NEO4J_NODE
+                            if row["obj_type"] == self.NODE
+                            else DatasetSubTypes.NEO4J_RELATIONSHIP
+                        ]
                     ),
+                ).as_workunit()
+
+                yield from self.add_properties(
+                    dataset=row["key"],
+                    custom_properties=None,
+                    description=row["description"],
                 )
 
             except Exception as e:
-                raise e
+                log.error(f"Failed to process row {row['key']}: {str(e)}")
+                self.report.report_failure(
+                    message="Error processing Neo4j metadata",
+                    context=row["key"],
+                    exc=e,
+                )
 
-    def get_report(self):
+    def get_report(self) -> "Neo4jSourceReport":
         return self.report

datahub/ingestion/source/powerbi/config.py

@@ -513,7 +513,7 @@ class PowerBiDashboardSourceConfig(
     include_workspace_name_in_dataset_urn: bool = pydantic.Field(
         default=False,
         description="It is recommended to set this to true, as it helps prevent the overwriting of datasets."
-        "Read section #11560 at https://datahubproject.io/docs/how/updating-datahub/ before enabling this option."
+        "Read section #11560 at https://docs.datahub.com/docs/how/updating-datahub/ before enabling this option."
         "To maintain backward compatibility, this is set to False.",
     )
 

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

@@ -63,10 +63,10 @@ class SessionWithTimeout(requests.Session):
         super().__init__(*args, **kwargs)
         self.timeout = timeout
 
-    def request(self, method, url, **kwargs):
+    def request(self, method, url, *args, **kwargs):
         # Set the default timeout if none is provided
         kwargs.setdefault("timeout", self.timeout)
-        return super().request(method, url, **kwargs)
+        return super().request(method, url, *args, **kwargs)
 
 
 class DataResolverBase(ABC):

datahub/ingestion/source/redshift/usage.py

@@ -182,15 +182,16 @@ class RedshiftUsageExtractor:
         self.report.num_operational_stats_filtered = 0
 
         if self.config.include_operational_stats:
-            with self.report.new_stage(USAGE_EXTRACTION_OPERATIONAL_STATS):
-                with PerfTimer() as timer:
-                    # Generate operation aspect workunits
-                    yield from self._gen_operation_aspect_workunits(
-                        self.connection, all_tables
-                    )
-                    self.report.operational_metadata_extraction_sec[
-                        self.config.database
-                    ] = timer.elapsed_seconds(digits=2)
+            with self.report.new_stage(
+                USAGE_EXTRACTION_OPERATIONAL_STATS
+            ), PerfTimer() as timer:
+                # Generate operation aspect workunits
+                yield from self._gen_operation_aspect_workunits(
+                    self.connection, all_tables
+                )
+                self.report.operational_metadata_extraction_sec[
+                    self.config.database
+                ] = timer.elapsed_seconds(digits=2)
 
         # Generate aggregate events
         with self.report.new_stage(USAGE_EXTRACTION_USAGE_AGGREGATION):

datahub/ingestion/source/sql/clickhouse.py

@@ -145,7 +145,11 @@ class ClickHouseConfig(
     )
     include_materialized_views: Optional[bool] = Field(default=True, description="")
 
-    def get_sql_alchemy_url(self, current_db=None):
+    def get_sql_alchemy_url(
+        self,
+        uri_opts: Optional[Dict[str, Any]] = None,
+        current_db: Optional[str] = None,
+    ) -> str:
         url = make_url(
             super().get_sql_alchemy_url(uri_opts=self.uri_opts, current_db=current_db)
         )

datahub/ingestion/source/sql/druid.py

@@ -1,4 +1,6 @@
 # This import verifies that the dependencies are available.
+from typing import Any, Dict, Optional
+
 import pydruid  # noqa: F401
 from pydantic.fields import Field
 from pydruid.db.sqlalchemy import DruidDialect
@@ -38,8 +40,11 @@ class DruidConfig(BasicSQLAlchemyConfig):
         description="regex patterns for schemas to filter in ingestion.",
     )
 
-    def get_sql_alchemy_url(self):
-        return f"{super().get_sql_alchemy_url()}/druid/v2/sql/"
+    def get_sql_alchemy_url(
+        self, uri_opts: Optional[Dict[str, Any]] = None, database: Optional[str] = None
+    ) -> str:
+        base_url = super().get_sql_alchemy_url(uri_opts=uri_opts, database=database)
+        return f"{base_url}/druid/v2/sql/"
 
     """
     The pydruid library already formats the table name correctly, so we do not

datahub/ingestion/source/sql/oracle.py

@@ -127,11 +127,15 @@ class OracleConfig(BasicSQLAlchemyConfig):
             )
         return v
 
-    def get_sql_alchemy_url(self):
-        url = super().get_sql_alchemy_url()
+    def get_sql_alchemy_url(
+        self, uri_opts: Optional[Dict[str, Any]] = None, database: Optional[str] = None
+    ) -> str:
+        url = super().get_sql_alchemy_url(uri_opts=uri_opts, database=database)
+
         if self.service_name:
             assert not self.database
             url = f"{url}/?service_name={self.service_name}"
+
         return url
 
     def get_identifier(self, schema: str, table: str) -> str:

datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py

@@ -10,7 +10,8 @@ from datahub.ingestion.api.ingestion_job_checkpointing_provider_base import (
     IngestionCheckpointingProviderConfig,
     JobId,
 )
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.metadata.schema_classes import DatahubIngestionCheckpointClass
 
 logger = logging.getLogger(__name__)

datahub/ingestion/source/tableau/tableau_validation.py

@@ -24,7 +24,7 @@ def check_user_role(
     mitigation_message_prefix: str = (
         "Assign `Site Administrator Explorer` role to the user"
     )
-    mitigation_message_suffix: str = "Refer to the setup guide: https://datahubproject.io/docs/quick-ingestion-guides/tableau/setup"
+    mitigation_message_suffix: str = "Refer to the setup guide: https://docs.datahub.com/docs/quick-ingestion-guides/tableau/setup"
 
     try:
         # TODO: Add check for `Enable Derived Permissions`

datahub/ingestion/source/usage/clickhouse_usage.py

@@ -2,7 +2,7 @@ import collections
 import dataclasses
 import logging
 from datetime import datetime
-from typing import Dict, Iterable, List
+from typing import Any, Dict, Iterable, List, Optional
 
 from dateutil import parser
 from pydantic.fields import Field
@@ -74,8 +74,12 @@ class ClickHouseUsageConfig(ClickHouseConfig, BaseUsageConfig, EnvConfigMixin):
     options: dict = Field(default={}, description="")
     query_log_table: str = Field(default="system.query_log", exclude=True)
 
-    def get_sql_alchemy_url(self):
-        return super().get_sql_alchemy_url()
+    def get_sql_alchemy_url(
+        self,
+        uri_opts: Optional[Dict[str, Any]] = None,
+        current_db: Optional[str] = None,
+    ) -> str:
+        return super().get_sql_alchemy_url(uri_opts=uri_opts, current_db=current_db)
 
 
 @platform_name("ClickHouse")

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -4,7 +4,7 @@ import json
 import logging
 from datetime import datetime
 from email.utils import parseaddr
-from typing import Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional
 
 from dateutil import parser
 from pydantic.fields import Field
@@ -98,8 +98,10 @@ class TrinoUsageConfig(TrinoConfig, BaseUsageConfig, EnvBasedSourceBaseConfig):
     options: dict = Field(default={}, description="")
     database: str = Field(description="The name of the catalog from getting the usage")
 
-    def get_sql_alchemy_url(self):
-        return super().get_sql_alchemy_url()
+    def get_sql_alchemy_url(
+        self, uri_opts: Optional[Dict[str, Any]] = None, database: Optional[str] = None
+    ) -> str:
+        return super().get_sql_alchemy_url(uri_opts=uri_opts, database=database)
 
 
 @dataclasses.dataclass

datahub/integrations/assertion/common.py

@@ -3,6 +3,7 @@ from typing import List, Optional, Tuple, TypedDict
 
 from datahub.api.entities.assertion.assertion import BaseEntityAssertion
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import DatasetProperties
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaMetadata
 from datahub.utilities.urns.urn import Urn
@@ -15,7 +16,7 @@ class ColumnDict(TypedDict):
 
 @lru_cache
 def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         props: Optional[DatasetProperties] = graph.get_aspect(urn, DatasetProperties)
         if props is not None:
             return props.qualifiedName
@@ -24,7 +25,7 @@ def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
 
 @lru_cache
 def get_schema_from_datahub(urn: str) -> Optional[List[ColumnDict]]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.INGESTION) as graph:
         schema: Optional[SchemaMetadata] = graph.get_aspect(urn, SchemaMetadata)
         if schema is not None:
             return [