openmetadata-ingestion 1.3.0.0.dev0__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

metadata/ingestion/source/database/databricks/metadata.py

@@ -13,7 +13,7 @@
 import re
 import traceback
 from copy import deepcopy
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Optional, Tuple, Union
 
 from pyhive.sqlalchemy_hive import _type_map
 from sqlalchemy import types, util
@@ -66,6 +66,7 @@ logger = ingestion_logger()
 
 DATABRICKS_TAG = "DATABRICK TAG"
 DATABRICKS_TAG_CLASSIFICATION = "DATABRICK TAG CLASSIFICATION"
+DEFAULT_TAG_VALUE = "NONE"
 
 
 class STRUCT(String):
@@ -268,6 +269,10 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         super().__init__(config, metadata)
         self.is_older_version = False
         self._init_version()
+        self.catalog_tags = {}
+        self.schema_tags = {}
+        self.table_tags = {}
+        self.column_tags = {}
 
     def _init_version(self):
         try:
@@ -314,10 +319,103 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         else:
             yield DEFAULT_DATABASE
 
+    def _clear_tag_cache(self) -> None:
+        """
+        Method to clean any existing tags available in memory
+        """
+        self.catalog_tags.clear()
+        self.table_tags.clear()
+        self.schema_tags.clear()
+        self.column_tags.clear()
+
+    def _add_to_tag_cache(
+        self, tag_dict: dict, key: Union[str, Tuple], value: Tuple[str, str]
+    ):
+        if tag_dict.get(key):
+            tag_dict.get(key).append(value)
+        else:
+            tag_dict[key] = [value]
+
+    def populate_tags_cache(self, database_name: str) -> None:
+        """
+        Method to fetch all the tags and populate the relevant caches
+        """
+        self._clear_tag_cache()
+        if self.source_config.includeTags is False:
+            return
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_CATALOGS_TAGS.format(database_name=database_name)
+            )
+
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.catalog_tags,
+                    tag.catalog_name,
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch catalog tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_SCHEMA_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.schema_tags,
+                    (tag.catalog_name, tag.schema_name),
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch schema tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_TABLE_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.table_tags,
+                    (tag.catalog_name, tag.schema_name, tag.table_name),
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch table tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_COLUMN_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                tag_table_id = (tag.catalog_name, tag.schema_name, tag.table_name)
+                if self.column_tags.get(tag_table_id):
+                    self._add_to_tag_cache(
+                        self.column_tags.get(tag_table_id),
+                        tag.column_name,
+                        # tag value is an optional field, if tag value is not available use default tag value
+                        (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                    )
+                else:
+                    self.column_tags[tag_table_id] = {
+                        tag.column_name: [
+                            (
+                                tag.tag_name,
+                                tag.tag_value or DEFAULT_TAG_VALUE,
+                            )
+                        ]
+                    }
+        except Exception as exc:
+            logger.debug(f"Failed to fetch column tags due to - {exc}")
+
     def get_database_names(self) -> Iterable[str]:
         configured_catalog = self.service_connection.catalog
         if configured_catalog:
             self.set_inspector(database_name=configured_catalog)
+            self.populate_tags_cache(database_name=configured_catalog)
             yield configured_catalog
         else:
             for new_catalog in self.get_database_names_raw():
@@ -337,6 +435,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                     continue
                 try:
                     self.set_inspector(database_name=new_catalog)
+                    self.populate_tags_cache(database_name=new_catalog)
                     yield new_catalog
                 except Exception as exc:
                     logger.error(traceback.format_exc())
@@ -361,10 +460,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         Method to yield database tags
         """
         try:
-            tags = self.connection.execute(
-                DATABRICKS_GET_CATALOGS_TAGS.format(database_name=database_name)
-            )
-            for tag in tags:
+            catalog_tags = self.catalog_tags.get(database_name, [])
+            for tag_name, tag_value in catalog_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -372,8 +469,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         service_name=self.context.database_service,
                         database_name=database_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
@@ -394,12 +491,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         Method to yield schema tags
         """
         try:
-            tags = self.connection.execute(
-                DATABRICKS_GET_SCHEMA_TAGS.format(
-                    database_name=self.context.database, schema_name=schema_name
-                )
-            )
-            for tag in tags:
+            schema_tags = self.schema_tags.get((self.context.database, schema_name), [])
+            for tag_name, tag_value in schema_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -408,8 +501,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         database_name=self.context.database,
                         schema_name=schema_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
@@ -428,14 +521,10 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
     ) -> Iterable[Either[OMetaTagAndClassification]]:
         table_name, _ = table_name_and_type
         try:
-            table_tags = self.connection.execute(
-                DATABRICKS_GET_TABLE_TAGS.format(
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
-                    table_name=table_name,
-                )
+            table_tags = self.table_tags.get(
+                (self.context.database, self.context.database_schema, table_name), []
             )
-            for tag in table_tags:
+            for tag_name, tag_value in table_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -445,35 +534,32 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         schema_name=self.context.database_schema,
                         table_name=table_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
 
-            column_tags = self.connection.execute(
-                DATABRICKS_GET_COLUMN_TAGS.format(
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
-                    table_name=table_name,
-                )
+            column_tags = self.column_tags.get(
+                (self.context.database, self.context.database_schema, table_name), {}
             )
-            for tag in column_tags:
-                yield from get_ometa_tag_and_classification(
-                    tag_fqn=fqn.build(
-                        self.metadata,
-                        Column,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
-                        table_name=table_name,
-                        column_name=tag.column_name,
-                    ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
-                    tag_description=DATABRICKS_TAG,
-                    classification_description=DATABRICKS_TAG_CLASSIFICATION,
-                )
+            for column_name, tags in column_tags.items():
+                for tag_name, tag_value in tags or []:
+                    yield from get_ometa_tag_and_classification(
+                        tag_fqn=fqn.build(
+                            self.metadata,
+                            Column,
+                            service_name=self.context.database_service,
+                            database_name=self.context.database,
+                            schema_name=self.context.database_schema,
+                            table_name=table_name,
+                            column_name=column_name,
+                        ),
+                        tags=[tag_value],
+                        classification_name=tag_name,
+                        tag_description=DATABRICKS_TAG,
+                        classification_description=DATABRICKS_TAG_CLASSIFICATION,
+                    )
 
         except Exception as exc:
             yield Either(

metadata/ingestion/source/database/databricks/queries.py

@@ -37,8 +37,7 @@ DATABRICKS_GET_SCHEMA_TAGS = textwrap.dedent(
     """
     SELECT 
         * 
-    FROM {database_name}.information_schema.schema_tags 
-    WHERE schema_name='{schema_name}'"""
+    FROM {database_name}.information_schema.schema_tags"""
 )
 
 DATABRICKS_GET_TABLE_TAGS = textwrap.dedent(
@@ -46,7 +45,7 @@ DATABRICKS_GET_TABLE_TAGS = textwrap.dedent(
     SELECT 
         * 
     FROM {database_name}.information_schema.table_tags 
-    WHERE schema_name = '{schema_name}' AND table_name = '{table_name}';"""
+    """
 )
 
 DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
@@ -54,5 +53,5 @@ DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
     SELECT 
         * 
     FROM {database_name}.information_schema.column_tags 
-    WHERE schema_name='{schema_name}' AND table_name='{table_name}';"""
+    """
 )

metadata/ingestion/source/database/datalake/metadata.py

@@ -77,8 +77,8 @@ from metadata.utils import fqn
 from metadata.utils.constants import DEFAULT_DATABASE
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
 from metadata.utils.datalake.datalake_utils import (
+    DataFrameColumnParser,
     fetch_dataframe,
-    get_columns,
     get_file_format_type,
 )
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
@@ -416,9 +416,14 @@ class DatalakeSource(DatabaseServiceSource):
                     file_extension=table_extension,
                 ),
             )
-
-            # If no data_frame (due to unsupported type), ignore
-            columns = get_columns(data_frame[0]) if data_frame else None
+            if data_frame:
+                column_parser = DataFrameColumnParser.create(
+                    data_frame[0], table_extension
+                )
+                columns = column_parser.get_columns()
+            else:
+                # If no data_frame (due to unsupported type), ignore
+                columns = None
             if columns:
                 table_request = CreateTableRequest(
                     name=table_name,

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

@@ -172,6 +172,11 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
             params_data["job_definition_id"] = job_id
 
         response = client.get(f"/accounts/{account_id}/runs", data=params_data)
+        if not response and not response.get("data"):
+            raise DBTConfigException(
+                "Unable to get the dbt job runs information.\n"
+                "Please check if the auth token is correct and has the necessary scopes to fetch dbt runs"
+            )
         runs_data = response.get("data")
         if runs_data:
             run_id = runs_data[0]["id"]

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

@@ -56,11 +56,6 @@ class DbtServiceTopology(ServiceTopology):
 
     root = TopologyNode(
         producer="get_dbt_files",
-        stages=[],
-        children=["process_dbt_files"],
-    )
-    process_dbt_files = TopologyNode(
-        producer="process_dbt_files",
         stages=[
             NodeStage(
                 type_=DbtFiles,
@@ -88,6 +83,7 @@ class DbtServiceTopology(ServiceTopology):
                 type_=DataModelLink,
                 processor="yield_data_models",
                 nullable=True,
+                consumer=["validate_dbt_files"],
             ),
         ],
     )
@@ -97,6 +93,7 @@ class DbtServiceTopology(ServiceTopology):
             NodeStage(
                 type_=AddLineageRequest,
                 processor="create_dbt_lineage",
+                consumer=["yield_data_models"],
             ),
             NodeStage(
                 type_=AddLineageRequest,
@@ -115,6 +112,7 @@ class DbtServiceTopology(ServiceTopology):
             NodeStage(
                 type_=CreateTestDefinitionRequest,
                 processor="create_dbt_tests_definition",
+                consumer=["yield_data_models"],
             ),
             NodeStage(
                 type_=CreateTestCaseRequest,
@@ -161,12 +159,6 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             }
         )
 
-    def process_dbt_files(self) -> Iterable[DbtFiles]:
-        """
-        Method return the dbt file from topology
-        """
-        yield self.context.dbt_file
-
     def get_dbt_files(self) -> Iterable[DbtFiles]:
         dbt_files = get_dbt_details(self.source_config.dbtConfigSource)
         for dbt_file in dbt_files:

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

@@ -14,6 +14,7 @@ DBT utils methods.
 import traceback
 from typing import Optional, Union
 
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -81,7 +82,8 @@ def generate_entity_link(dbt_test):
     manifest_node = dbt_test.get(DbtCommonEnum.MANIFEST_NODE.value)
     entity_link_list = [
         entity_link.get_entity_link(
-            table_fqn=table_fqn,
+            Table,
+            fqn=table_fqn,
             column_name=manifest_node.column_name
             if hasattr(manifest_node, "column_name")
             else None,

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

@@ -31,8 +31,6 @@ from metadata.generated.schema.entity.services.databaseService import DatabaseSe
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
-from metadata.generated.schema.entity.teams.team import Team
-from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -140,39 +138,29 @@ class DbtSource(DbtServiceSource):
 
     def get_dbt_owner(
         self, manifest_node: dict, catalog_node: Optional[dict]
-    ) -> Optional[str]:
+    ) -> Optional[EntityReference]:
         """
         Returns dbt owner
         """
-        owner = None
-        dbt_owner = None
-        if catalog_node:
-            dbt_owner = catalog_node.metadata.owner
-        if manifest_node:
-            dbt_owner = manifest_node.meta.get(DbtCommonEnum.OWNER.value)
-        if dbt_owner:
-            owner_name = dbt_owner
-            user_owner_fqn = fqn.build(
-                self.metadata, entity_type=User, user_name=owner_name
-            )
-            if user_owner_fqn:
-                owner = self.metadata.get_entity_reference(
-                    entity=User, fqn=user_owner_fqn
-                )
-            else:
-                team_owner_fqn = fqn.build(
-                    self.metadata, entity_type=Team, team_name=owner_name
-                )
-                if team_owner_fqn:
-                    owner = self.metadata.get_entity_reference(
-                        entity=Team, fqn=team_owner_fqn
-                    )
-                else:
+        try:
+            owner = None
+            dbt_owner = None
+            if catalog_node:
+                dbt_owner = catalog_node.metadata.owner
+            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)
+                if not owner:
                     logger.warning(
                         "Unable to ingest owner from DBT since no user or"
                         f" team was found with name {dbt_owner}"
                     )
-        return owner
+            return owner
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Unable to ingest owner from DBT due to: {exc}")
+        return None
 
     def check_columns(self, catalog_node):
         for catalog_key, catalog_column in catalog_node.get("columns").items():
@@ -692,7 +680,6 @@ class DbtSource(DbtServiceSource):
         )
         if table_entity:
             try:
-
                 service_name, database_name, schema_name, table_name = fqn.split(
                     table_entity.fullyQualifiedName.__root__
                 )