openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -16,6 +16,8 @@ from abc import ABC, abstractmethod
 from typing import Iterable
 
 from dbt_artifacts_parser.parser import parse_catalog, parse_manifest, parse_run_results
+from pydantic import Field
+from typing_extensions import Annotated
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.api.tests.createTestCase import CreateTestCaseRequest
@@ -54,7 +56,9 @@ class DbtServiceTopology(ServiceTopology):
     dbt files -> dbt tags -> data models -> descriptions -> lineage -> tests.
     """
 
-    root = TopologyNode(
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
         producer="get_dbt_files",
         stages=[
             NodeStage(
@@ -69,7 +73,9 @@ class DbtServiceTopology(ServiceTopology):
             "process_dbt_tests",
         ],
     )
-    process_dbt_data_model = TopologyNode(
+    process_dbt_data_model: Annotated[
+        TopologyNode, Field(description="Process dbt data models")
+    ] = TopologyNode(
         producer="get_dbt_objects",
         stages=[
             NodeStage(
@@ -87,7 +93,9 @@ class DbtServiceTopology(ServiceTopology):
             ),
         ],
     )
-    process_dbt_entities = TopologyNode(
+    process_dbt_entities: Annotated[
+        TopologyNode, Field(description="Process dbt entities")
+    ] = TopologyNode(
         producer="get_data_model",
         stages=[
             NodeStage(
@@ -106,7 +114,9 @@ class DbtServiceTopology(ServiceTopology):
             ),
         ],
     )
-    process_dbt_tests = TopologyNode(
+    process_dbt_tests: Annotated[
+        TopologyNode, Field(description="Process dbt tests")
+    ] = TopologyNode(
         producer="get_dbt_tests",
         stages=[
             NodeStage(
@@ -207,8 +217,7 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
         """
         Prepare the data models
         """
-        for data_model_link in self.context.get().data_model_links:
-            yield data_model_link
+        yield from self.context.get().data_model_links
 
     @abstractmethod
     def create_dbt_lineage(self, data_model_link: DataModelLink) -> AddLineageRequest:

metadata/ingestion/source/database/dbt/metadata.py

@@ -14,7 +14,7 @@ DBT source methods.
 """
 import traceback
 from datetime import datetime
-from typing import Iterable, List, Optional, Union
+from typing import Any, Iterable, List, Optional, Union
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.api.tests.createTestCase import CreateTestCaseRequest
@@ -46,10 +46,16 @@ from metadata.generated.schema.tests.testDefinition import (
     TestDefinition,
     TestPlatform,
 )
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    FullyQualifiedEntityName,
+    SqlQuery,
+    Timestamp,
+    Uuid,
+)
 from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
@@ -87,6 +93,7 @@ from metadata.ingestion.source.database.dbt.dbt_utils import (
 from metadata.ingestion.source.database.dbt.models import DbtMeta
 from metadata.utils import fqn
 from metadata.utils.elasticsearch import get_entity_from_es_result
+from metadata.utils.entity_link import get_table_fqn
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
 from metadata.utils.time_utils import convert_timestamp_to_milliseconds
@@ -120,7 +127,7 @@ class DbtSource(DbtServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
         return cls(config, metadata)
 
     def test_connection(self) -> None:
@@ -132,11 +139,10 @@ class DbtSource(DbtServiceSource):
         """
         By default for DBT nothing is required to be prepared
         """
-        pass
 
     def get_dbt_owner(
-        self, manifest_node: dict, catalog_node: Optional[dict]
-    ) -> Optional[EntityReference]:
+        self, manifest_node: Any, catalog_node: Optional[Any]
+    ) -> Optional[EntityReferenceList]:
         """
         Returns dbt owner
         """
@@ -148,7 +154,16 @@ class DbtSource(DbtServiceSource):
             if manifest_node:
                 dbt_owner = manifest_node.meta.get(DbtCommonEnum.OWNER.value)
             if dbt_owner:
-                owner = self.metadata.get_reference_by_name(name=dbt_owner)
+                owner = self.metadata.get_reference_by_name(
+                    name=dbt_owner, is_owner=True
+                )
+
+                if owner:
+                    return owner
+
+                # If owner is not found, try to find the owner in OMD using email
+                owner = self.metadata.get_reference_by_email(name=dbt_owner)
+
                 if not owner:
                     logger.warning(
                         "Unable to ingest owner from DBT since no user or"
@@ -185,6 +200,7 @@ class DbtSource(DbtServiceSource):
                 **dbt_files.dbt_manifest[DbtCommonEnum.NODES.value],
                 **dbt_files.dbt_manifest[DbtCommonEnum.SOURCES.value],
             }
+            catalog_entities = None
             if dbt_files.dbt_catalog:
                 catalog_entities = {
                     **dbt_files.dbt_catalog[DbtCommonEnum.NODES.value],
@@ -209,7 +225,7 @@ class DbtSource(DbtServiceSource):
                     )
 
                 # Validate the catalog file if it is passed
-                if dbt_files.dbt_catalog:
+                if catalog_entities:
                     catalog_node = catalog_entities.get(key)
                     if catalog_node and "columns" in catalog_node:
                         self.check_columns(catalog_node=catalog_node)
@@ -324,6 +340,7 @@ class DbtSource(DbtServiceSource):
                 **dbt_objects.dbt_manifest.sources,
                 **dbt_objects.dbt_manifest.nodes,
             }
+            catalog_entities = None
             if dbt_objects.dbt_catalog:
                 catalog_entities = {
                     **dbt_objects.dbt_catalog.sources,
@@ -387,7 +404,7 @@ class DbtSource(DbtServiceSource):
                     logger.debug(f"Processing DBT node: {model_name}")
 
                     catalog_node = None
-                    if dbt_objects.dbt_catalog:
+                    if catalog_entities:
                         catalog_node = catalog_entities.get(key)
 
                     dbt_table_tags_list = []
@@ -411,7 +428,6 @@ class DbtSource(DbtServiceSource):
                     dbt_raw_query = get_dbt_raw_query(manifest_node)
 
                     # Get the table entity from ES
-                    # TODO: Change to get_by_name once the postgres case sensitive calls is fixed
                     table_fqn = fqn.build(
                         self.metadata,
                         entity_type=Table,
@@ -442,15 +458,19 @@ class DbtSource(DbtServiceSource):
                                 if manifest_node.description
                                 else None,
                                 path=get_data_model_path(manifest_node=manifest_node),
-                                rawSql=dbt_raw_query if dbt_raw_query else "",
-                                sql=dbt_compiled_query if dbt_compiled_query else "",
+                                rawSql=SqlQuery(dbt_raw_query)
+                                if dbt_raw_query
+                                else None,
+                                sql=SqlQuery(dbt_compiled_query)
+                                if dbt_compiled_query
+                                else None,
                                 columns=self.parse_data_model_columns(
                                     manifest_node, catalog_node
                                 ),
                                 upstream=self.parse_upstream_nodes(
                                     manifest_entities, manifest_node
                                 ),
-                                owner=self.get_dbt_owner(
+                                owners=self.get_dbt_owner(
                                     manifest_node=manifest_node,
                                     catalog_node=catalog_node,
                                 ),
@@ -516,7 +536,6 @@ class DbtSource(DbtServiceSource):
                         )
 
                         # check if the parent table exists in OM before adding it to the upstream list
-                        # TODO: Change to get_by_name once the postgres case sensitive calls is fixed
                         parent_table_entity: Optional[
                             Union[Table, List[Table]]
                         ] = get_entity_from_es_result(
@@ -536,7 +555,7 @@ class DbtSource(DbtServiceSource):
         return upstream_nodes
 
     def parse_data_model_columns(
-        self, manifest_node: dict, catalog_node: dict
+        self, manifest_node: Any, catalog_node: Any
     ) -> List[Column]:
         """
         Method to parse the DBT columns
@@ -546,7 +565,7 @@ class DbtSource(DbtServiceSource):
         for key, manifest_column in manifest_columns.items():
             try:
                 logger.debug(f"Processing DBT column: {key}")
-                # If catalog file is passed pass the column information from catalog file
+                # If catalog file is passed, pass the column information from catalog file
                 catalog_column = None
                 if catalog_node and catalog_node.columns:
                     catalog_column = catalog_node.columns.get(key)
@@ -618,9 +637,7 @@ class DbtSource(DbtServiceSource):
         Method to process DBT lineage from upstream nodes
         """
         to_entity: Table = data_model_link.table_entity
-        logger.debug(
-            f"Processing DBT lineage for: {to_entity.fullyQualifiedName.__root__}"
-        )
+        logger.debug(f"Processing DBT lineage for: {to_entity.fullyQualifiedName.root}")
 
         for upstream_node in data_model_link.datamodel.upstream:
             try:
@@ -638,11 +655,11 @@ class DbtSource(DbtServiceSource):
                         right=AddLineageRequest(
                             edge=EntitiesEdge(
                                 fromEntity=EntityReference(
-                                    id=from_entity.id.__root__,
+                                    id=Uuid(from_entity.id.root),
                                     type="table",
                                 ),
                                 toEntity=EntityReference(
-                                    id=to_entity.id.__root__,
+                                    id=Uuid(to_entity.id.root),
                                     type="table",
                                 ),
                                 lineageDetails=LineageDetails(
@@ -664,48 +681,48 @@ class DbtSource(DbtServiceSource):
         """
         Method to process DBT lineage from queries
         """
-        to_entity: Table = data_model_link.table_entity
-        logger.debug(
-            f"Processing DBT Query lineage for: {to_entity.fullyQualifiedName.__root__}"
-        )
-
-        try:
-            source_elements = fqn.split(to_entity.fullyQualifiedName.__root__)
-            # remove service name from fqn to make it parseable in format db.schema.table
-            query_fqn = fqn._build(  # pylint: disable=protected-access
-                *source_elements[-3:]
-            )
-            query = (
-                f"create table {query_fqn} as {data_model_link.datamodel.sql.__root__}"
-            )
-            connection_type = str(
-                self.config.serviceConnection.__root__.config.type.value
+        if data_model_link.datamodel.sql:
+            to_entity: Table = data_model_link.table_entity
+            logger.debug(
+                f"Processing DBT Query lineage for: {to_entity.fullyQualifiedName.root}"
             )
-            dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
-            lineages = get_lineage_by_query(
-                self.metadata,
-                query=query,
-                service_name=source_elements[0],
-                database_name=source_elements[1],
-                schema_name=source_elements[2],
-                dialect=dialect,
-                timeout_seconds=self.source_config.parsingTimeoutLimit,
-                lineage_source=LineageSource.DbtLineage,
-            )
-            for lineage_request in lineages or []:
-                yield lineage_request
 
-        except Exception as exc:  # pylint: disable=broad-except
-            yield Either(
-                left=StackTraceError(
-                    name=data_model_link.datamodel.sql.__root__,
-                    error=(
-                        f"Failed to parse the query {data_model_link.datamodel.sql.__root__}"
-                        f" to capture lineage: {exc}"
-                    ),
-                    stackTrace=traceback.format_exc(),
+            try:
+                source_elements = fqn.split(to_entity.fullyQualifiedName.root)
+                # remove service name from fqn to make it parseable in format db.schema.table
+                query_fqn = fqn._build(  # pylint: disable=protected-access
+                    *source_elements[-3:]
+                )
+                query = (
+                    f"create table {query_fqn} as {data_model_link.datamodel.sql.root}"
+                )
+                connection_type = str(
+                    self.config.serviceConnection.root.config.type.value
+                )
+                dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
+                lineages = get_lineage_by_query(
+                    self.metadata,
+                    query=query,
+                    service_name=source_elements[0],
+                    database_name=source_elements[1],
+                    schema_name=source_elements[2],
+                    dialect=dialect,
+                    timeout_seconds=self.source_config.parsingTimeoutLimit,
+                    lineage_source=LineageSource.DbtLineage,
+                )
+                yield from lineages or []
+
+            except Exception as exc:  # pylint: disable=broad-except
+                yield Either(
+                    left=StackTraceError(
+                        name=data_model_link.datamodel.sql.root,
+                        error=(
+                            f"Failed to parse the query {data_model_link.datamodel.sql.root}"
+                            f" to capture lineage: {exc}"
+                        ),
+                        stackTrace=traceback.format_exc(),
+                    )
                 )
-            )
 
     def process_dbt_meta(self, manifest_meta):
         """
@@ -719,7 +736,7 @@ class DbtSource(DbtServiceSource):
                     get_tag_labels(
                         metadata=self.metadata,
                         tags=dbt_meta_info.openmetadata.glossary,
-                        include_tags=self.source_config.includeTags,
+                        include_tags=True,
                         tag_type=GlossaryTerm,
                     )
                     or []
@@ -732,7 +749,7 @@ class DbtSource(DbtServiceSource):
                         metadata=self.metadata,
                         tags=[tier_fqn.split(fqn.FQN_SEPARATOR)[-1]],
                         classification_name=tier_fqn.split(fqn.FQN_SEPARATOR)[0],
-                        include_tags=self.source_config.includeTags,
+                        include_tags=True,
                     )
                     or []
                 )
@@ -749,12 +766,12 @@ class DbtSource(DbtServiceSource):
         """
         table_entity: Table = data_model_link.table_entity
         logger.debug(
-            f"Processing DBT Descriptions for: {table_entity.fullyQualifiedName.__root__}"
+            f"Processing DBT Descriptions for: {table_entity.fullyQualifiedName.root}"
         )
         if table_entity:
             try:
                 service_name, database_name, schema_name, table_name = fqn.split(
-                    table_entity.fullyQualifiedName.__root__
+                    table_entity.fullyQualifiedName.root
                 )
                 data_model = data_model_link.datamodel
                 force_override = False
@@ -769,7 +786,7 @@ class DbtSource(DbtServiceSource):
                     self.metadata.patch_description(
                         entity=Table,
                         source=table_entity,
-                        description=data_model.description.__root__,
+                        description=data_model.description.root,
                         force=force_override,
                     )
 
@@ -786,7 +803,7 @@ class DbtSource(DbtServiceSource):
                                     database_name=database_name,
                                     schema_name=schema_name,
                                     table_name=table_name,
-                                    column_name=column.name.__root__,
+                                    column_name=column.name.root,
                                 ),
                                 description=column.description,
                             )
@@ -799,7 +816,7 @@ class DbtSource(DbtServiceSource):
             except Exception as exc:  # pylint: disable=broad-except
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Failed to parse the node {table_entity.fullyQualifiedName.__root__} "
+                    f"Failed to parse the node {table_entity.fullyQualifiedName.root} "
                     f"to update dbt description: {exc}"
                 )
 
@@ -836,7 +853,7 @@ class DbtSource(DbtServiceSource):
                                 manifest_node
                             ),
                             displayName=None,
-                            owner=None,
+                            owners=None,
                         )
                     )
         except Exception as err:  # pylint: disable=broad-except
@@ -863,20 +880,42 @@ class DbtSource(DbtServiceSource):
                     test_suite = check_or_create_test_suite(
                         self.metadata, entity_link_str
                     )
-                    yield Either(
-                        right=CreateTestCaseRequest(
-                            name=manifest_node.name,
-                            description=manifest_node.description,
-                            testDefinition=FullyQualifiedEntityName(
-                                __root__=manifest_node.name
-                            ),
-                            entityLink=entity_link_str,
-                            testSuite=test_suite.fullyQualifiedName,
-                            parameterValues=create_test_case_parameter_values(dbt_test),
-                            displayName=None,
-                            owner=None,
-                        )
+                    table_fqn = get_table_fqn(entity_link_str)
+                    source_elements = table_fqn.split(fqn.FQN_SEPARATOR)
+                    test_case_fqn = fqn.build(
+                        self.metadata,
+                        entity_type=TestCase,
+                        service_name=source_elements[0],
+                        database_name=source_elements[1],
+                        schema_name=source_elements[2],
+                        table_name=source_elements[3],
+                        column_name=manifest_node.column_name
+                        if hasattr(manifest_node, "column_name")
+                        else None,
+                        test_case_name=manifest_node.name,
+                    )
+
+                    test_case = self.metadata.get_by_name(
+                        TestCase, test_case_fqn, fields=["testDefinition,testSuite"]
                     )
+                    if test_case is None:
+                        # Create the test case only if it does not exist
+                        yield Either(
+                            right=CreateTestCaseRequest(
+                                name=manifest_node.name,
+                                description=manifest_node.description,
+                                testDefinition=FullyQualifiedEntityName(
+                                    manifest_node.name
+                                ),
+                                entityLink=entity_link_str,
+                                testSuite=test_suite.fullyQualifiedName,
+                                parameterValues=create_test_case_parameter_values(
+                                    dbt_test
+                                ),
+                                displayName=None,
+                                owners=None,
+                            )
+                        )
         except Exception as err:  # pylint: disable=broad-except
             yield Either(
                 left=StackTraceError(
@@ -898,6 +937,12 @@ class DbtSource(DbtServiceSource):
                     f"Adding DBT Test Case Results for node: {manifest_node.name}"
                 )
                 dbt_test_result = dbt_test.get(DbtCommonEnum.RESULTS.value)
+                if not dbt_test_result:
+                    logger.warning(
+                        f"DBT Test Case Results not found for node: {manifest_node.name}"
+                    )
+                    return
+
                 test_case_status = TestCaseStatus.Aborted
                 test_result_value = 0
                 if dbt_test_result.status.value in [
@@ -931,8 +976,8 @@ class DbtSource(DbtServiceSource):
 
                 # Create the test case result object
                 test_case_result = TestCaseResult(
-                    timestamp=convert_timestamp_to_milliseconds(
-                        dbt_timestamp.timestamp()
+                    timestamp=Timestamp(
+                        convert_timestamp_to_milliseconds(dbt_timestamp.timestamp())
                     ),
                     testCaseStatus=test_case_status,
                     testResultValue=[

metadata/ingestion/source/database/dbt/models.py

@@ -18,27 +18,27 @@ from pydantic import BaseModel
 
 
 class DbtFiles(BaseModel):
-    dbt_catalog: Optional[dict]
+    dbt_catalog: Optional[dict] = None
     dbt_manifest: dict
-    dbt_run_results: Optional[List[dict]]
+    dbt_run_results: Optional[List[dict]] = None
 
 
 class DbtObjects(BaseModel):
-    dbt_catalog: Optional[Any]
+    dbt_catalog: Optional[Any] = None
     dbt_manifest: Any
-    dbt_run_results: Optional[List[Any]]
+    dbt_run_results: Optional[List[Any]] = None
 
 
 class DbtFilteredModel(BaseModel):
     is_filtered: Optional[bool] = False
-    message: Optional[str]
-    model_fqn: Optional[str]
+    message: Optional[str] = None
+    model_fqn: Optional[str] = None
 
 
 class DbtMetaGlossaryTier(BaseModel):
-    tier: Optional[str]
-    glossary: Optional[List[str]]
+    tier: Optional[str] = None
+    glossary: Optional[List[str]] = None
 
 
 class DbtMeta(BaseModel):
-    openmetadata: Optional[DbtMetaGlossaryTier]
+    openmetadata: Optional[DbtMetaGlossaryTier] = None

metadata/ingestion/source/database/deltalake/clients/base.py

@@ -0,0 +1,82 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Deltalake Base Client
+"""
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Iterable, List, Optional
+
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    PartitionColumnDetails,
+    TableType,
+)
+from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
+    DeltaLakeConnection,
+)
+
+
+@dataclass
+class TableInfo:
+    """Small Class to store TableInfo from different Clients in the same format."""
+
+    schema: str
+    name: str
+    _type: TableType
+    location: Optional[str] = None
+    description: Optional[str] = None
+    columns: Optional[List[Column]] = None
+    table_partitions: Optional[List[PartitionColumnDetails]] = None
+
+
+class DeltalakeBaseClient(ABC):
+    @classmethod
+    @abstractmethod
+    def from_config(
+        cls, service_connection: DeltaLakeConnection
+    ) -> "DeltalakeBaseClient":
+        """Returns a Deltalake Client based on the DatalakeConfig passed."""
+
+    @abstractmethod
+    def get_database_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Returns the Database Names, based on the underlying client."""
+
+    @abstractmethod
+    def get_database_schema_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Returns the RAW database schema names, based on the underlying client."""
+
+    @abstractmethod
+    def get_table_info(
+        self, service_connection: DeltaLakeConnection, schema_name: str
+    ) -> Iterable[TableInfo]:
+        """Returns the TableInfo, based on the underlying client."""
+
+    @abstractmethod
+    def update_table_info(self, table_info: TableInfo) -> TableInfo:
+        """Updates TableInfo with extra Metadata that was defered."""
+
+    @abstractmethod
+    def close(self, service_connection: DeltaLakeConnection):
+        """Closes the Client connection."""
+
+    @abstractmethod
+    def get_test_get_databases_fn(self, config) -> Callable:
+        """Returns a Callable used to test the GetDatabases condition."""
+
+    @abstractmethod
+    def get_test_get_tables_fn(self, config) -> Callable:
+        """Returns a Callable used to test the GetTables condition."""