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

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

@@ -15,8 +15,10 @@ import traceback
 from copy import deepcopy
 from typing import Iterable, Optional, Tuple, Union
 
+from pydantic import EmailStr
+from pydantic_core import PydanticCustomError
 from pyhive.sqlalchemy_hive import _type_map
-from sqlalchemy import types, util
+from sqlalchemy import exc, types, util
 from sqlalchemy.engine import reflection
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.exc import DatabaseError
@@ -35,6 +37,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -43,6 +46,7 @@ from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
 from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_DDL,
     DATABRICKS_GET_CATALOGS,
     DATABRICKS_GET_CATALOGS_TAGS,
     DATABRICKS_GET_COLUMN_TAGS,
@@ -60,17 +64,15 @@ from metadata.utils.constants import DEFAULT_DATABASE
 from metadata.utils.filters import filter_by_database
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
-    get_all_table_ddls,
     get_all_view_definitions,
-    get_table_ddl,
     get_view_definition_wrapper,
 )
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 logger = ingestion_logger()
 
-DATABRICKS_TAG = "DATABRICK TAG"
-DATABRICKS_TAG_CLASSIFICATION = "DATABRICK TAG CLASSIFICATION"
+DATABRICKS_TAG = "DATABRICKS TAG"
+DATABRICKS_TAG_CLASSIFICATION = "DATABRICKS TAG CLASSIFICATION"
 DEFAULT_TAG_VALUE = "NONE"
 
 
@@ -162,16 +164,20 @@ def get_columns(self, connection, table_name, schema=None, **kw):
             "system_data_type": raw_col_type,
         }
         if col_type in {"array", "struct", "map"}:
-            rows = dict(
-                connection.execute(
-                    f"DESCRIBE {schema}.{table_name} {col_name}"
-                    if schema
-                    else f"DESCRIBE {table_name} {col_name}"
-                ).fetchall()
-            )
-
-            col_info["system_data_type"] = rows["data_type"]
-            col_info["is_complex"] = True
+            col_name = f"`{col_name}`" if "." in col_name else col_name
+            try:
+                rows = dict(
+                    connection.execute(
+                        f"DESCRIBE TABLE {kw.get('db_name')}.{schema}.{table_name} {col_name}"
+                    ).fetchall()
+                )
+                col_info["system_data_type"] = rows["data_type"]
+                col_info["is_complex"] = True
+            except DatabaseError as err:
+                logger.error(
+                    f"Failed to fetch column details for column {col_name} in table {table_name} due to: {err}"
+                )
+                logger.debug(traceback.format_exc())
         result.append(col_info)
     return result
 
@@ -221,7 +227,9 @@ def get_table_comment(  # pylint: disable=unused-argument
     """
     cursor = connection.execute(
         DATABRICKS_GET_TABLE_COMMENTS.format(
-            schema_name=schema_name, table_name=table_name
+            database_name=self.context.get().database,
+            schema_name=schema_name,
+            table_name=table_name,
         )
     )
     try:
@@ -250,6 +258,25 @@ def get_view_definition(
     return None
 
 
+@reflection.cache
+def get_table_ddl(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    Gets the Table DDL
+    """
+    schema = schema or self.default_schema_name
+    table_name = f"{schema}.{table_name}" if schema else table_name
+    cursor = connection.execute(DATABRICKS_DDL.format(table_name=table_name))
+    try:
+        result = cursor.fetchone()
+        if result:
+            return result[0]
+    except Exception:
+        pass
+    return None
+
+
 @reflection.cache
 def get_table_names(
     self, connection, schema=None, **kw
@@ -264,7 +291,7 @@ def get_table_names(
         # if it is > 1, we use spark thrift server with 3 columns in the result (schema, table, is_temporary)
         # else it is hive with 1 column in the result
         if len(row) > 1:
-            tables.append(row[1])
+            table_name = row[1]
         else:
             table_name = row[0]
         if schema:
@@ -312,7 +339,6 @@ DatabricksDialect.get_view_definition = get_view_definition
 DatabricksDialect.get_table_names = get_table_names
 DatabricksDialect.get_all_view_definitions = get_all_view_definitions
 reflection.Inspector.get_schema_names = get_schema_names_reflection
-reflection.Inspector.get_all_table_ddls = get_all_table_ddls
 reflection.Inspector.get_table_ddl = get_table_ddl
 
 
@@ -345,8 +371,8 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatabricksConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: DatabricksConnection = config.serviceConnection.root.config
         if not isinstance(connection, DatabricksConnection):
             raise InvalidSourceException(
                 f"Expected DatabricksConnection, but got {connection}"
@@ -489,9 +515,11 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                 )
                 if filter_by_database(
                     self.source_config.databaseFilterPattern,
-                    database_fqn
-                    if self.source_config.useFqnForFiltering
-                    else new_catalog,
+                    (
+                        database_fqn
+                        if self.source_config.useFqnForFiltering
+                        else new_catalog
+                    ),
                 ):
                     self.status.filter(database_fqn, "Database Filtered Out")
                     continue
@@ -651,7 +679,9 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
         try:
             cursor = self.connection.execute(
                 DATABRICKS_GET_TABLE_COMMENTS.format(
-                    schema_name=schema_name, table_name=table_name
+                    database_name=self.context.get().database,
+                    schema_name=schema_name,
+                    table_name=table_name,
                 )
             )
             for result in list(cursor):
@@ -674,3 +704,42 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                 f"Table description error for table [{schema_name}.{table_name}]: {exc}"
             )
         return description
+
+    def _filter_owner_name(self, owner_name: str) -> str:
+        """remove unnecessary keyword from name"""
+        pattern = r"\(Unknown\)"
+        filtered_name = re.sub(pattern, "", owner_name).strip()
+        return filtered_name
+
+    def get_owner_ref(self, table_name: str) -> Optional[EntityReferenceList]:
+        """
+        Method to process the table owners
+        """
+        try:
+            query = DATABRICKS_GET_TABLE_COMMENTS.format(
+                database_name=self.context.get().database,
+                schema_name=self.context.get().database_schema,
+                table_name=table_name,
+            )
+            result = self.connection.engine.execute(query)
+            owner = None
+            for row in result:
+                row_dict = dict(row)
+                if row_dict.get("col_name") == "Owner":
+                    owner = row_dict.get("data_type")
+                    break
+            if not owner:
+                return
+
+            owner = self._filter_owner_name(owner)
+            owner_ref = None
+            try:
+                owner_email = EmailStr._validate(owner)
+                owner_ref = self.metadata.get_reference_by_email(email=owner_email)
+            except PydanticCustomError:
+                owner_ref = self.metadata.get_reference_by_name(name=owner)
+            return owner_ref
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error processing owner for table {table_name}: {exc}")
+        return

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

@@ -24,7 +24,9 @@ DATABRICKS_VIEW_DEFINITIONS = textwrap.dedent(
     """
 )
 
-DATABRICKS_GET_TABLE_COMMENTS = "DESCRIBE TABLE EXTENDED {schema_name}.{table_name}"
+DATABRICKS_GET_TABLE_COMMENTS = (
+    "DESCRIBE TABLE EXTENDED {database_name}.{schema_name}.{table_name}"
+)
 
 DATABRICKS_GET_CATALOGS = "SHOW CATALOGS"
 
@@ -54,3 +56,5 @@ DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
     FROM {database_name}.information_schema.column_tags 
     """
 )
+
+DATABRICKS_DDL = "SHOW CREATE TABLE {table_name}"

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

@@ -53,8 +53,8 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
         """Create class instance"""
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatabricksConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: DatabricksConnection = config.serviceConnection.root.config
         if not isinstance(connection, DatabricksConnection):
             raise InvalidSourceException(
                 f"Expected DatabricksConnection, but got {connection}"

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

@@ -15,6 +15,7 @@ import traceback
 from datetime import datetime
 from typing import Iterable
 
+from metadata.generated.schema.type.basic import DateTime
 from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.database.databricks.query_parser import (
     DatabricksQueryParserSource,
@@ -46,9 +47,9 @@ class DatabricksUsageSource(DatabricksQueryParserSource, UsageSource):
                         TableQuery(
                             query=row.get("query_text"),
                             userName=row.get("user_name"),
-                            startTime=row.get("query_start_time_ms"),
-                            endTime=row.get("execution_end_time_ms"),
-                            analysisDate=datetime.now(),
+                            startTime=str(row.get("query_start_time_ms")),
+                            endTime=str(row.get("execution_end_time_ms")),
+                            analysisDate=DateTime(datetime.now()),
                             serviceName=self.config.serviceName,
                             duration=row.get("duration")
                             if row.get("duration")

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

@@ -0,0 +1,72 @@
+#  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.
+
+"""
+Datalake Azure Blob Client
+"""
+from functools import partial
+from typing import Callable, Iterable, Optional
+
+from azure.storage.blob import BlobServiceClient
+
+from metadata.clients.azure_client import AzureClient
+from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
+    AzureConfig,
+)
+from metadata.ingestion.source.database.datalake.clients.base import DatalakeBaseClient
+from metadata.utils.constants import DEFAULT_DATABASE
+
+
+class DatalakeAzureBlobClient(DatalakeBaseClient):
+    def __init__(self, client: BlobServiceClient):
+        self._client = client
+
+    @classmethod
+    def from_config(cls, config: AzureConfig) -> "DatalakeAzureBlobClient":
+        try:
+            if not config.securityConfig:
+                raise RuntimeError("AzureConfig securityConfig can't be None.")
+
+            client = AzureClient(config.securityConfig).create_blob_client()
+            return cls(client=client)
+        except Exception as exc:
+            raise RuntimeError(
+                f"Unknown error connecting with {config.securityConfig}: {exc}."
+            )
+
+    def update_client_database(self, config, database_name):
+        # For the AzureBlob Client we don't need to do anything when changing the database
+        pass
+
+    def get_database_names(self, service_connection) -> Iterable[str]:
+        yield service_connection.databaseName or DEFAULT_DATABASE
+
+    def get_database_schema_names(self, bucket_name: Optional[str]) -> Iterable[str]:
+        prefix = bucket_name or ""
+
+        for schema in self._client.list_containers(name_starts_with=prefix):
+            yield schema["name"]
+
+    def get_table_names(self, bucket_name: str, prefix: Optional[str]) -> Iterable[str]:
+        container_client = self._client.get_container_client(bucket_name)
+
+        for file in container_client.list_blobs(name_starts_with=prefix or None):
+            yield file.name
+
+    def close(self, service_connection):
+        self._client.close()
+
+    def get_test_list_buckets_fn(self, bucket_name: Optional[str]) -> Callable:
+        def list_buckets(client: BlobServiceClient):
+            conn = client.list_containers(name_starts_with="")
+            list(conn)
+
+        return partial(list_buckets, self._client)

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

@@ -0,0 +1,47 @@
+#  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.
+
+"""
+Datalake Base Client
+"""
+from abc import ABC, abstractmethod
+from typing import Callable, Iterable, Optional
+
+
+class DatalakeBaseClient(ABC):
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config) -> "DatalakeBaseClient":
+        """Returns a Datalake Client based on the DatalakeConfig passed."""
+
+    @abstractmethod
+    def update_client_database(self, config, database_name: str):
+        """Updates the Client when changing the Database."""
+
+    @abstractmethod
+    def get_database_names(self, service_connection) -> Iterable[str]:
+        """Returns the Database Names, based on the underlying client."""
+
+    @abstractmethod
+    def get_database_schema_names(self, bucket_name: Optional[str]) -> Iterable[str]:
+        """Returns the RAW database schema names, based on the underlying client."""
+
+    @abstractmethod
+    def get_table_names(self, bucket_name: str, prefix: Optional[str]) -> Iterable[str]:
+        """Returns the Table names, based on the underlying client."""
+
+    @abstractmethod
+    def close(self, service_connection):
+        """Closes the Client connection."""
+
+    @abstractmethod
+    def get_test_list_buckets_fn(self, bucket_name: Optional[str]) -> Callable:
+        """Returns a Callable used to test the ListBuckets condition."""

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

@@ -0,0 +1,136 @@
+#  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.
+
+"""
+Datalake GCS Client
+"""
+import os
+from copy import deepcopy
+from functools import partial
+from typing import Callable, Iterable, List, Optional
+
+from google.cloud import storage
+
+from metadata.generated.schema.entity.services.connections.database.datalake.gcsConfig import (
+    GCSConfig,
+)
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+    MultipleProjectId,
+    SingleProjectId,
+)
+from metadata.ingestion.source.database.datalake.clients.base import DatalakeBaseClient
+from metadata.utils.credentials import GOOGLE_CREDENTIALS, set_google_credentials
+
+
+class DatalakeGcsClient(DatalakeBaseClient):
+    def __init__(
+        self, client: storage.Client, temp_credentials_file_path_list: List[str]
+    ):
+        self._client = client
+        self._temp_credentials_file_path_list = temp_credentials_file_path_list
+
+    @property
+    def project(self):
+        return self._client.project
+
+    @staticmethod
+    def get_gcs_client(config: GCSConfig) -> storage.Client:
+        gcs_config = deepcopy(config)
+
+        if hasattr(config.securityConfig, "gcpConfig") and isinstance(
+            config.securityConfig.gcpConfig.projectId, MultipleProjectId
+        ):
+            gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
+                gcs_config.securityConfig.gcpConfig.projectId.root[0]
+            )
+
+        if not gcs_config.securityConfig:
+            raise RuntimeError("GCSConfig securityConfig can't be None.")
+
+        set_google_credentials(gcp_credentials=gcs_config.securityConfig)
+        return storage.Client()
+
+    def update_temp_credentials_file_path_list(self):
+        credentials = os.environ.get(GOOGLE_CREDENTIALS)
+
+        if credentials:
+            self._temp_credentials_file_path_list.append(credentials)
+
+    @classmethod
+    def from_config(cls, config: GCSConfig) -> "DatalakeGcsClient":
+        gcs_client = cls.get_gcs_client(config)
+
+        client = cls(client=gcs_client, temp_credentials_file_path_list=[])
+        client.update_temp_credentials_file_path_list()
+
+        return client
+
+    def get_database_names(self, service_connection):
+        project_id_list = (
+            service_connection.configSource.securityConfig.gcpConfig.projectId.root
+        )
+
+        if not isinstance(project_id_list, list):
+            project_id_list = [project_id_list]
+
+        for project_id in project_id_list:
+            yield project_id
+
+    def update_client_database(self, config: GCSConfig, database_name: str):
+        gcs_config = deepcopy(config)
+
+        if hasattr(gcs_config.securityConfig, "gcpConfig"):
+            gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
+                database_name
+            )
+
+        self._client = self.get_gcs_client(gcs_config)
+        self.update_temp_credentials_file_path_list()
+
+    def get_database_schema_names(self, bucket_name: Optional[str]) -> Iterable[str]:
+        if bucket_name:
+            yield bucket_name
+        else:
+            for bucket in self._client.list_buckets():
+                yield bucket.name
+
+    def get_table_names(self, bucket_name: str, prefix: Optional[str]) -> Iterable[str]:
+        bucket = self._client.get_bucket(bucket_name)
+
+        for key in bucket.list_blobs(prefix=prefix):
+            yield key.name
+
+    def close(self, service_connection):
+        os.environ.pop("GOOGLE_CLOUD_PROJECT", "")
+
+        if (
+            isinstance(service_connection, GcpCredentialsValues)
+            and GOOGLE_CREDENTIALS in os.environ
+        ):
+            del os.environ[GOOGLE_CREDENTIALS]
+            for temp_file_path in self._temp_credentials_file_path_list:
+                if os.path.exists(temp_file_path):
+                    os.remove(temp_file_path)
+
+    def get_test_list_buckets_fn(self, bucket_name: Optional[str]) -> Callable:
+
+        if bucket_name:
+            fn = partial(self._client.get_bucket, bucket_name)
+        else:
+            fn = self._client.list_buckets
+
+        os.environ.pop("GOOGLE_CLOUD_PROJECT", "")
+        if GOOGLE_CREDENTIALS in os.environ:
+            os.remove(os.environ[GOOGLE_CREDENTIALS])
+            del os.environ[GOOGLE_CREDENTIALS]
+
+        return fn

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

@@ -0,0 +1,79 @@
+#  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.
+
+"""
+Datalake S3 Client
+"""
+from functools import partial
+from typing import Callable, Iterable, Optional
+
+from metadata.clients.aws_client import AWSClient
+from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
+    S3Config,
+)
+from metadata.ingestion.source.database.datalake.clients.base import DatalakeBaseClient
+from metadata.utils.constants import DEFAULT_DATABASE
+from metadata.utils.s3_utils import list_s3_objects
+
+
+class DatalakeS3Client(DatalakeBaseClient):
+    def __init__(self, client):
+        self._client = client
+
+    @classmethod
+    def from_config(cls, config: S3Config) -> "DatalakeS3Client":
+        if not config.securityConfig:
+            raise RuntimeError("S3Config securityConfig can't be None.")
+
+        s3_client = AWSClient(config.securityConfig).get_client(service_name="s3")
+        return cls(client=s3_client)
+
+    def update_client_database(self, config, database_name: str):
+        # For the S3 Client we don't need to do anything when changing the database
+        pass
+
+    def get_database_names(self, service_connection) -> Iterable[str]:
+        yield service_connection.databaseName or DEFAULT_DATABASE
+
+    def get_database_schema_names(self, bucket_name: Optional[str]) -> Iterable[str]:
+        if bucket_name:
+            yield bucket_name
+        else:
+            for bucket in self._client.list_buckets()["Buckets"]:
+                yield bucket["Name"]
+
+    def get_table_names(self, bucket_name: str, prefix: Optional[str]) -> Iterable[str]:
+        kwargs = {"Bucket": bucket_name}
+
+        if prefix:
+            kwargs["Prefix"] = prefix if prefix.endswith("/") else f"{prefix}/"
+
+        for key in list_s3_objects(self._client, **kwargs):
+            yield key["Key"]
+
+    def get_folders_prefix(
+        self, bucket_name: str, prefix: Optional[str]
+    ) -> Iterable[str]:
+        for page in self._client.get_paginator("list_objects_v2").paginate(
+            Bucket=bucket_name, Prefix=prefix or "", Delimiter="/"
+        ):
+            for common_prefix in page.get("CommonPrefixes", []):
+                yield common_prefix.get("Prefix")
+
+    def close(self, service_connection):
+        # For the S3 Client we don't need to do anything when closing the connection
+        pass
+
+    def get_test_list_buckets_fn(self, bucket_name: Optional[str]) -> Callable:
+
+        if bucket_name:
+            return partial(self._client.list_objects, Bucket=bucket_name)
+        return self._client.list_buckets