openmetadata-ingestion 1.4.6.0__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

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

@@ -43,6 +43,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 +61,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,6 +161,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
             "system_data_type": raw_col_type,
         }
         if col_type in {"array", "struct", "map"}:
+            col_name = f"`{col_name}`" if "." in col_name else col_name
             rows = dict(
                 connection.execute(
                     f"DESCRIBE {schema}.{table_name} {col_name}"
@@ -249,6 +249,26 @@ 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
@@ -280,7 +300,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
 
 
@@ -313,8 +332,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}"

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

@@ -54,3 +54,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,
@@ -48,7 +49,7 @@ class DatabricksUsageSource(DatabricksQueryParserSource, UsageSource):
                             userName=row.get("user_name"),
                             startTime=row.get("query_start_time_ms"),
                             endTime=row.get("execution_end_time_ms"),
-                            analysisDate=datetime.now(),
+                            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,132 @@
+#  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
+
+    @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

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

@@ -12,15 +12,10 @@
 """
 Source connection handler
 """
-import os
-from copy import deepcopy
 from dataclasses import dataclass
-from functools import partial, singledispatch
+from functools import singledispatch
 from typing import Optional
 
-from google.cloud import storage
-
-from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -36,13 +31,13 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
-from metadata.generated.schema.security.credentials.gcpValues import (
-    MultipleProjectId,
-    SingleProjectId,
-)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.credentials import GOOGLE_CREDENTIALS, set_google_credentials
+from metadata.ingestion.source.database.datalake.clients.azure_blob import (
+    DatalakeAzureBlobClient,
+)
+from metadata.ingestion.source.database.datalake.clients.gcs import DatalakeGcsClient
+from metadata.ingestion.source.database.datalake.clients.s3 import DatalakeS3Client
 
 
 # Only import specific datalake dependencies if necessary
@@ -66,44 +61,17 @@ def get_datalake_client(config):
 
 @get_datalake_client.register
 def _(config: S3Config):
-    from metadata.clients.aws_client import AWSClient
-
-    s3_client = AWSClient(config.securityConfig).get_client(service_name="s3")
-    return s3_client
+    return DatalakeS3Client.from_config(config)
 
 
 @get_datalake_client.register
 def _(config: GCSConfig):
-    gcs_config = deepcopy(config)
-    if hasattr(config.securityConfig, "gcpConfig") and isinstance(
-        config.securityConfig.gcpConfig.projectId, MultipleProjectId
-    ):
-        gcs_config: GCSConfig = deepcopy(config)
-        gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-            gcs_config.securityConfig.gcpConfig.projectId.__root__[0]
-        )
-    set_google_credentials(gcp_credentials=gcs_config.securityConfig)
-    gcs_client = storage.Client()
-    return gcs_client
+    return DatalakeGcsClient.from_config(config)
 
 
 @get_datalake_client.register
 def _(config: AzureConfig):
-    try:
-        return AzureClient(config.securityConfig).create_blob_client()
-    except Exception as exc:
-        raise RuntimeError(
-            f"Unknown error connecting with {config.securityConfig}: {exc}."
-        )
-
-
-def set_gcs_datalake_client(config: GCSConfig, project_id: str):
-    gcs_config = deepcopy(config)
-    if hasattr(gcs_config.securityConfig, "gcpConfig"):
-        gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-            project_id
-        )
-    return get_datalake_client(gcs_config)
+    return DatalakeAzureBlobClient.from_config(config)
 
 
 def get_connection(connection: DatalakeConnection) -> DatalakeClient:
@@ -128,36 +96,10 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    config = connection.config.configSource
-    func = None
-    if isinstance(config, GCSConfig):
-        if connection.config.bucketName:
-            func = partial(connection.client.get_bucket, connection.config.bucketName)
-        else:
-            func = connection.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]
-
-    if isinstance(config, S3Config):
-        if connection.config.bucketName:
-            func = partial(
-                connection.client.list_objects, Bucket=connection.config.bucketName
-            )
-        else:
-            func = connection.client.list_buckets
-
-    if isinstance(config, AzureConfig):
-
-        def list_connection(connection):
-            conn = connection.client.list_containers(name_starts_with="")
-            list(conn)
-
-        func = partial(list_connection, connection)
-
     test_fn = {
-        "ListBuckets": func,
+        "ListBuckets": connection.client.get_test_list_buckets_fn(
+            connection.config.bucketName
+        ),
     }
 
     test_connection_steps(