acryl-datahub 1.0.0.3rc8__py3-none-any.whl → 1.0.0.3rc10__py3-none-any.whl

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/slack/slack.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import textwrap
 from dataclasses import dataclass
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
@@ -613,6 +612,10 @@ class SlackSource(StatefulIngestionSourceBase):
             ),
         )
 
+    @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[Tuple[CorpUser, Optional[CorpUserEditableInfoClass]]]:
@@ -634,56 +637,5 @@ class SlackSource(StatefulIngestionSourceBase):
             if user_obj.email is not None:
                 yield (user_obj, editable_properties)
 
-    @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_oss(self) -> Iterable[CorpUser]:
-        graphql_query = textwrap.dedent(
-            """
-            query listUsers($input: ListUsersInput!) {
-                listUsers(input: $input) {
-                    total
-                    users {
-                        urn
-                        editableProperties {
-                            email
-                            slack
-                        }
-                    }
-                }
-            }
-        """
-        )
-        start = 0
-        count = 10
-        total = count
-
-        assert self.ctx.graph is not None
-
-        while start < total:
-            variables = {"input": {"start": start, "count": count}}
-            response = self.ctx.graph.execute_graphql(
-                query=graphql_query, variables=variables
-            )
-            list_users = response.get("listUsers", {})
-            total = list_users.get("total", 0)
-            users = list_users.get("users", [])
-            for user in users:
-                user_obj = CorpUser()
-                editable_properties = user.get("editableProperties", {})
-                user_obj.urn = user.get("urn")
-                if user_obj.urn is None:
-                    continue
-                if editable_properties is not None:
-                    user_obj.email = editable_properties.get("email")
-                if user_obj.email is None:
-                    urn_id = Urn.from_string(user_obj.urn).get_entity_id_as_string()
-                    if "@" in urn_id:
-                        user_obj.email = urn_id
-                if user_obj.email is not None:
-                    yield user_obj
-            start += count
-
     def get_report(self) -> SourceReport:
         return self.report

datahub/ingestion/source/snowflake/snowflake_connection.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 from typing import Any, Dict, Optional
 
 import pydantic
@@ -402,13 +403,30 @@ class SnowflakeConnection(Closeable):
     def __init__(self, connection: NativeSnowflakeConnection):
         self._connection = connection
 
+        self._query_num_lock = threading.Lock()
+        self._query_num = 1
+
     def native_connection(self) -> NativeSnowflakeConnection:
         return self._connection
 
+    def get_query_no(self) -> int:
+        with self._query_num_lock:
+            no = self._query_num
+            self._query_num += 1
+            return no
+
     def query(self, query: str) -> Any:
         try:
-            logger.info(f"Query: {query}", stacklevel=2)
+            # We often run multiple queries in parallel across multiple threads,
+            # so we need to number them to help with log readability.
+            query_num = self.get_query_no()
+            logger.info(f"Query #{query_num}: {query}", stacklevel=2)
             resp = self._connection.cursor(DictCursor).execute(query)
+            if resp is not None and resp.rowcount is not None:
+                logger.info(
+                    f"Query #{query_num} got {resp.rowcount} row(s) back from Snowflake",
+                    stacklevel=2,
+                )
             return resp
 
         except Exception as e:

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/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