acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.0.0.1__py3-none-any.whl

datahub/ingestion/graph/filters.py

@@ -1,30 +1,58 @@
 import dataclasses
 import enum
-from typing import Any, Dict, List, Optional
+import warnings
+from typing import Dict, List, Literal, Optional, Union
+
+from typing_extensions import TypeAlias
 
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
 )
+from datahub.errors import SearchFilterWarning
 from datahub.utilities.urns.urn import guess_entity_type
 
-RawSearchFilterRule = Dict[str, Any]
+RawSearchFilterRule: TypeAlias = Dict[str, Union[str, bool, List[str]]]
+
+# This is a list of OR filters, each of which is a list of AND filters.
+# This can be put directly into the orFilters parameter in GraphQL.
+RawSearchFilter: TypeAlias = List[Dict[Literal["and"], List[RawSearchFilterRule]]]
+
+# Mirrors our GraphQL enum: https://datahubproject.io/docs/graphql/enums#filteroperator
+FilterOperator: TypeAlias = Literal[
+    "CONTAIN",
+    "EQUAL",
+    "IEQUAL",
+    "IN",
+    "EXISTS",
+    "GREATER_THAN",
+    "GREATER_THAN_OR_EQUAL_TO",
+    "LESS_THAN",
+    "LESS_THAN_OR_EQUAL_TO",
+    "START_WITH",
+    "END_WITH",
+    "DESCENDANTS_INCL",
+    "ANCESTORS_INCL",
+    "RELATED_INCL",
+]
 
 
 @dataclasses.dataclass
 class SearchFilterRule:
     field: str
-    condition: str  # TODO: convert to an enum
+    condition: FilterOperator
     values: List[str]
     negated: bool = False
 
     def to_raw(self) -> RawSearchFilterRule:
-        return {
+        rule: RawSearchFilterRule = {
             "field": self.field,
             "condition": self.condition,
             "values": self.values,
-            "negated": self.negated,
         }
+        if self.negated:
+            rule["negated"] = True
+        return rule
 
     def negate(self) -> "SearchFilterRule":
         return SearchFilterRule(
@@ -53,10 +81,10 @@ def generate_filter(
     platform_instance: Optional[str],
     env: Optional[str],
     container: Optional[str],
-    status: RemovedStatusFilter,
+    status: Optional[RemovedStatusFilter],
     extra_filters: Optional[List[RawSearchFilterRule]],
-    extra_or_filters: Optional[List[RawSearchFilterRule]] = None,
-) -> List[Dict[str, List[RawSearchFilterRule]]]:
+    extra_or_filters: Optional[RawSearchFilter] = None,
+) -> RawSearchFilter:
     """
     Generate a search filter based on the provided parameters.
     :param platform: The platform to filter by.
@@ -85,15 +113,16 @@ def generate_filter(
         and_filters.append(_get_container_filter(container).to_raw())
 
     # Status filter.
-    status_filter = _get_status_filter(status)
-    if status_filter:
-        and_filters.append(status_filter.to_raw())
+    if status:
+        status_filter = _get_status_filter(status)
+        if status_filter:
+            and_filters.append(status_filter.to_raw())
 
     # Extra filters.
     if extra_filters:
         and_filters += extra_filters
 
-    or_filters: List[Dict[str, List[RawSearchFilterRule]]] = [{"and": and_filters}]
+    or_filters: RawSearchFilter = [{"and": and_filters}]
 
     # Env filter
     if env:
@@ -107,11 +136,27 @@ def generate_filter(
 
     # Extra OR filters are distributed across the top level and lists.
     if extra_or_filters:
-        or_filters = [
-            {"and": and_filter["and"] + [extra_or_filter]}
-            for extra_or_filter in extra_or_filters
-            for and_filter in or_filters
-        ]
+        new_or_filters: RawSearchFilter = []
+        for and_filter in or_filters:
+            for extra_or_filter in extra_or_filters:
+                if isinstance(extra_or_filter, dict) and "and" in extra_or_filter:
+                    new_or_filters.append(
+                        {"and": and_filter["and"] + extra_or_filter["and"]}
+                    )
+                else:
+                    # Hack for backwards compatibility.
+                    # We have some code that erroneously passed a List[RawSearchFilterRule]
+                    # instead of a List[Dict["and", List[RawSearchFilterRule]]].
+                    warnings.warn(
+                        "Passing a List[RawSearchFilterRule] to extra_or_filters is deprecated. "
+                        "Please pass a List[Dict[str, List[RawSearchFilterRule]]] instead.",
+                        SearchFilterWarning,
+                        stacklevel=3,
+                    )
+                    new_or_filters.append(
+                        {"and": and_filter["and"] + [extra_or_filter]}  # type: ignore
+                    )
+        or_filters = new_or_filters
 
     return or_filters
 

datahub/ingestion/sink/datahub_rest.py

@@ -20,7 +20,7 @@ from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
-    DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_EMITTER_ENDPOINT,
     DEFAULT_REST_TRACE_MODE,
     DataHubRestEmitter,
     RestSinkEndpoint,
@@ -70,7 +70,7 @@ _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
 
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
-    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
+    endpoint: RestSinkEndpoint = DEFAULT_REST_EMITTER_ENDPOINT
     default_trace_mode: RestTraceMode = DEFAULT_REST_TRACE_MODE
 
     # These only apply in async modes.

datahub/ingestion/source/cassandra/cassandra.py

@@ -123,16 +123,7 @@ class CassandraSource(StatefulIngestionSourceBase):
             ).workunit_processor,
         ]
 
-    def get_workunits_internal(
-        self,
-    ) -> Iterable[MetadataWorkUnit]:
-        for metadata in self._get_metadata():
-            if isinstance(metadata, MetadataWorkUnit):
-                yield metadata
-            else:
-                yield from metadata.as_workunits()
-
-    def _get_metadata(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
+    def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         if not self.cassandra_api.authenticate():
             return
         keyspaces: List[CassandraKeyspace] = self.cassandra_api.get_keyspaces()

datahub/ingestion/source/common/data_platforms.py

@@ -0,0 +1,23 @@
+# This is a pretty limited list, and is not really complete yet. Right now it's only used to allow
+# automatic platform mapping when generating lineage and we have a manual override, so
+# it being incomplete is ok. This should not be used for urn validation.
+KNOWN_VALID_PLATFORM_NAMES = [
+    "bigquery",
+    "cassandra",
+    "databricks",
+    "delta-lake",
+    "dbt",
+    "feast",
+    "file",
+    "gcs",
+    "hdfs",
+    "hive",
+    "mssql",
+    "mysql",
+    "oracle",
+    "postgres",
+    "redshift",
+    "s3",
+    "sagemaker",
+    "snowflake",
+]

datahub/ingestion/source/common/gcp_credentials_config.py

@@ -51,3 +51,9 @@ class GCPCredential(ConfigModel):
             cred_json = json.dumps(configs, indent=4, separators=(",", ": "))
             fp.write(cred_json.encode())
             return fp.name
+
+    def to_dict(self, project_id: Optional[str] = None) -> Dict[str, str]:
+        configs = self.dict()
+        if project_id:
+            configs["project_id"] = project_id
+        return configs

datahub/ingestion/source/common/subtypes.py

@@ -25,6 +25,7 @@ class DatasetSubTypes(StrEnum):
     NEO4J_NODE = "Neo4j Node"
     NEO4J_RELATIONSHIP = "Neo4j Relationship"
     SNOWFLAKE_STREAM = "Snowflake Stream"
+    API_ENDPOINT = "API Endpoint"
 
     # TODO: Create separate entity...
     NOTEBOOK = "Notebook"
@@ -44,6 +45,7 @@ class DatasetContainerSubTypes(StrEnum):
     GCS_BUCKET = "GCS bucket"
     ABS_CONTAINER = "ABS container"
     KEYSPACE = "Keyspace"  # Cassandra
+    NAMESPACE = "Namespace"  # Iceberg
 
 
 class BIContainerSubTypes(StrEnum):
@@ -68,7 +70,7 @@ class FlowContainerSubTypes(StrEnum):
 class JobContainerSubTypes(StrEnum):
     NIFI_PROCESS_GROUP = "Process Group"
     MSSQL_JOBSTEP = "Job Step"
-    MSSQL_STORED_PROCEDURE = "Stored Procedure"
+    STORED_PROCEDURE = "Stored Procedure"
 
 
 class BIAssetSubTypes(StrEnum):
@@ -93,7 +95,21 @@ class BIAssetSubTypes(StrEnum):
     SAC_STORY = "Story"
     SAC_APPLICATION = "Application"
 
+    # Hex
+    HEX_PROJECT = "Project"
+    HEX_COMPONENT = "Component"
+
 
 class MLAssetSubTypes(StrEnum):
     MLFLOW_TRAINING_RUN = "ML Training Run"
     MLFLOW_EXPERIMENT = "ML Experiment"
+    VERTEX_EXPERIMENT = "Experiment"
+    VERTEX_EXPERIMENT_RUN = "Experiment Run"
+    VERTEX_EXECUTION = "Execution"
+
+    VERTEX_MODEL = "ML Model"
+    VERTEX_MODEL_GROUP = "ML Model Group"
+    VERTEX_TRAINING_JOB = "Training Job"
+    VERTEX_ENDPOINT = "Endpoint"
+    VERTEX_DATASET = "Dataset"
+    VERTEX_PROJECT = "Project"

datahub/ingestion/source/data_lake_common/path_spec.py

@@ -11,7 +11,7 @@ from cached_property import cached_property
 from pydantic.fields import Field
 from wcmatch import pathlib
 
-from datahub.configuration.common import ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel
 from datahub.ingestion.source.aws.s3_util import is_s3_uri
 from datahub.ingestion.source.azure.abs_utils import is_abs_uri
 from datahub.ingestion.source.gcs.gcs_utils import is_gcs_uri
@@ -145,6 +145,11 @@ class PathSpec(ConfigModel):
         description="Include hidden folders in the traversal (folders starting with . or _",
     )
 
+    tables_filter_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="The tables_filter_pattern configuration field uses regular expressions to filter the tables part of the Pathspec for ingestion, allowing fine-grained control over which tables are included or excluded based on specified patterns. The default setting allows all tables.",
+    )
+
     def is_path_hidden(self, path: str) -> bool:
         # Split the path into directories and filename
         dirs, filename = os.path.split(path)
@@ -177,6 +182,12 @@ class PathSpec(ConfigModel):
                 ):
                     return False
         logger.debug(f"{path} is not excluded")
+
+        table_name, _ = self.extract_table_name_and_path(path)
+        if not self.tables_filter_pattern.allowed(table_name):
+            return False
+        logger.debug(f"{path} is passed table name check")
+
         ext = os.path.splitext(path)[1].strip(".")
 
         if not ignore_ext:
@@ -218,6 +229,15 @@ class PathSpec(ConfigModel):
                     exclude_path.rstrip("/"), flags=pathlib.GLOBSTAR
                 ):
                     return False
+
+        file_name_pattern = self.include.rsplit("/", 1)[1]
+        table_name, _ = self.extract_table_name_and_path(
+            os.path.join(path, file_name_pattern)
+        )
+        if not self.tables_filter_pattern.allowed(table_name):
+            return False
+        logger.debug(f"{path} is passed table name check")
+
         return True
 
     @classmethod

datahub/ingestion/source/dbt/dbt_common.py

@@ -4,7 +4,7 @@ from abc import abstractmethod
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import auto
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import more_itertools
 import pydantic
@@ -849,7 +849,7 @@ class DBTSourceBase(StatefulIngestionSourceBase):
         test_nodes: List[DBTNode],
         extra_custom_props: Dict[str, str],
         all_nodes_map: Dict[str, DBTNode],
-    ) -> Iterable[MetadataWorkUnit]:
+    ) -> Iterable[MetadataChangeProposalWrapper]:
         for node in sorted(test_nodes, key=lambda n: n.dbt_name):
             upstreams = get_upstreams_for_test(
                 test_node=node,
@@ -902,7 +902,7 @@ class DBTSourceBase(StatefulIngestionSourceBase):
                     yield MetadataChangeProposalWrapper(
                         entityUrn=assertion_urn,
                         aspect=self._make_data_platform_instance_aspect(),
-                    ).as_workunit()
+                    )
 
                     yield make_assertion_from_test(
                         custom_props,
@@ -949,7 +949,9 @@ class DBTSourceBase(StatefulIngestionSourceBase):
             ),
         )
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         if self.config.write_semantics == "PATCH":
             self.ctx.require_graph("Using dbt with write_semantics=PATCH")
 

datahub/ingestion/source/dbt/dbt_core.py

@@ -343,6 +343,9 @@ class DBTRunResult(BaseModel):
     def timing_map(self) -> Dict[str, DBTRunTiming]:
         return {x.name: x for x in self.timing if x.name}
 
+    def has_success_status(self) -> bool:
+        return self.status in ("pass", "success")
+
 
 class DBTRunMetadata(BaseModel):
     dbt_schema_version: str
@@ -355,12 +358,7 @@ def _parse_test_result(
     dbt_metadata: DBTRunMetadata,
     run_result: DBTRunResult,
 ) -> Optional[DBTTestResult]:
-    if run_result.status == "success":
-        # This was probably a docs generate run result, so this isn't actually
-        # a test result.
-        return None
-
-    if run_result.status != "pass":
+    if not run_result.has_success_status():
         native_results = {"message": run_result.message or ""}
         if run_result.failures:
             native_results.update({"failures": str(run_result.failures)})

datahub/ingestion/source/dbt/dbt_tests.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
 from datahub.emitter import mce_builder
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
     AssertionInfoClass,
     AssertionResultClass,
@@ -43,6 +42,9 @@ class DBTTestResult:
 
     native_results: Dict[str, str]
 
+    def has_success_status(self) -> bool:
+        return self.status in ("pass", "success")
+
 
 def _get_name_for_relationship_test(kw_args: Dict[str, str]) -> Optional[str]:
     """
@@ -157,7 +159,7 @@ def make_assertion_from_test(
     node: "DBTNode",
     assertion_urn: str,
     upstream_urn: str,
-) -> MetadataWorkUnit:
+) -> MetadataChangeProposalWrapper:
     assert node.test_info
     qualified_test_name = node.test_info.qualified_test_name
     column_name = node.test_info.column_name
@@ -231,7 +233,7 @@ def make_assertion_from_test(
     return MetadataChangeProposalWrapper(
         entityUrn=assertion_urn,
         aspect=assertion_info,
-    ).as_workunit()
+    )
 
 
 def make_assertion_result_from_test(
@@ -240,7 +242,7 @@ def make_assertion_result_from_test(
     assertion_urn: str,
     upstream_urn: str,
     test_warnings_are_errors: bool,
-) -> MetadataWorkUnit:
+) -> MetadataChangeProposalWrapper:
     assertionResult = AssertionRunEventClass(
         timestampMillis=int(test_result.execution_time.timestamp() * 1000.0),
         assertionUrn=assertion_urn,
@@ -249,7 +251,7 @@ def make_assertion_result_from_test(
         result=AssertionResultClass(
             type=(
                 AssertionResultTypeClass.SUCCESS
-                if test_result.status == "pass"
+                if test_result.has_success_status()
                 or (not test_warnings_are_errors and test_result.status == "warn")
                 else AssertionResultTypeClass.FAILURE
             ),
@@ -261,4 +263,4 @@ def make_assertion_result_from_test(
     return MetadataChangeProposalWrapper(
         entityUrn=assertion_urn,
         aspect=assertionResult,
-    ).as_workunit()
+    )

datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py

@@ -66,7 +66,7 @@ class DremioToDataHubSourceTypeMapping:
     }
 
     @staticmethod
-    def get_datahub_source_type(dremio_source_type: str) -> str:
+    def get_datahub_platform(dremio_source_type: str) -> str:
         """
         Return the DataHub source type.
         """

datahub/ingestion/source/dremio/dremio_entities.py

@@ -294,7 +294,7 @@ class DremioContainer:
         )
 
 
-class DremioSource(DremioContainer):
+class DremioSourceContainer(DremioContainer):
     subclass: str = "Dremio Source"
     dremio_source_type: str
     root_path: Optional[str]
@@ -337,7 +337,7 @@ class DremioCatalog:
         self.dremio_api = dremio_api
         self.edition = dremio_api.edition
         self.datasets: Deque[DremioDataset] = deque()
-        self.sources: Deque[DremioSource] = deque()
+        self.sources: Deque[DremioSourceContainer] = deque()
         self.spaces: Deque[DremioSpace] = deque()
         self.folders: Deque[DremioFolder] = deque()
         self.glossary_terms: Deque[DremioGlossaryTerm] = deque()
@@ -380,12 +380,13 @@ class DremioCatalog:
                 container_type = container.get("container_type")
                 if container_type == DremioEntityContainerType.SOURCE:
                     self.sources.append(
-                        DremioSource(
+                        DremioSourceContainer(
                             container_name=container.get("name"),
                             location_id=container.get("id"),
                             path=[],
                             api_operations=self.dremio_api,
-                            dremio_source_type=container.get("source_type"),
+                            dremio_source_type=container.get("source_type")
+                            or "unknown",
                             root_path=container.get("root_path"),
                             database_name=container.get("database_name"),
                         )
@@ -426,7 +427,7 @@ class DremioCatalog:
         self.set_containers()
         return deque(itertools.chain(self.sources, self.spaces, self.folders))
 
-    def get_sources(self) -> Deque[DremioSource]:
+    def get_sources(self) -> Deque[DremioSourceContainer]:
         self.set_containers()
         return self.sources