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

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

@@ -0,0 +1,304 @@
+#  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 PySpark Client
+"""
+import re
+import traceback
+from enum import Enum
+from typing import Any, Callable, Dict, Iterable, List, Optional
+
+from pyspark.sql.utils import AnalysisException, ParseException
+
+from metadata.generated.schema.entity.data.table import Column, TableType
+from metadata.generated.schema.entity.services.connections.database.deltalake.metastoreConfig import (
+    MetastoreConfig,
+    MetastoreDbConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
+    DeltaLakeConnection,
+)
+from metadata.ingestion.connections.builders import get_connection_args_common
+from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
+from metadata.ingestion.source.database.deltalake.clients.base import (
+    DeltalakeBaseClient,
+    TableInfo,
+)
+from metadata.utils.constants import DEFAULT_DATABASE
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class SparkTableType(Enum):
+    MANAGED = "MANAGED"
+    TEMPORARY = "TEMPORARY"
+    VIEW = "VIEW"
+    EXTERNAL = "EXTERNAL"
+
+
+TABLE_TYPE_MAP = {
+    SparkTableType.MANAGED.value: TableType.Regular,
+    SparkTableType.VIEW.value: TableType.View,
+    SparkTableType.EXTERNAL.value: TableType.External,
+}
+
+ARRAY_CHILD_START_INDEX = 6
+ARRAY_CHILD_END_INDEX = -1
+
+ARRAY_DATATYPE_REPLACE_MAP = {"(": "<", ")": ">", "=": ":", "<>": ""}
+
+
+class DeltalakePySparkClient(DeltalakeBaseClient):
+    def __init__(self, spark_session):
+        self._spark = spark_session
+
+    @classmethod
+    def from_config(cls, config: DeltaLakeConnection) -> "DeltalakeBaseClient":
+        """Returns a Deltalake Client based on the DeltalakeConfig passed."""
+        import pyspark
+        from delta import configure_spark_with_delta_pip
+
+        builder = (
+            pyspark.sql.SparkSession.builder.appName(
+                config.configSource.appName or "OpenMetadata"
+            )
+            .enableHiveSupport()
+            .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+            .config(
+                "spark.sql.catalog.spark_catalog",
+                "org.apache.spark.sql.delta.catalog.DeltaCatalog",
+            )
+            # Download delta-core jars when creating the SparkSession
+            .config("spark.jars.packages", "io.delta:delta-core_2.12:2.0.0")
+        )
+
+        # Check that the attribute exists and is properly informed
+        if (
+            hasattr(config.configSource.connection, "metastoreHostPort")
+            and config.configSource.connection.metastoreHostPort
+        ):
+            builder.config(
+                "hive.metastore.uris",
+                f"thrift://{config.configSource.connection.metastoreHostPort}",
+            )
+
+        if isinstance(config.configSource.connection, MetastoreDbConnection):
+            if config.configSource.connection.metastoreDb:
+                builder.config(
+                    "spark.hadoop.javax.jdo.option.ConnectionURL",
+                    config.configSource.connection.metastoreDb,
+                )
+            if config.configSource.connection.jdbcDriverClassPath:
+                builder.config(
+                    "spark.driver.extraClassPath",
+                    config.configSource.connection.jdbcDriverClassPath,
+                )
+            if config.configSource.connection.username:
+                builder.config(
+                    "spark.hadoop.javax.jdo.option.ConnectionUserName",
+                    config.configSource.connection.username,
+                )
+            if config.configSource.connection.password:
+                builder.config(
+                    "spark.hadoop.javax.jdo.option.ConnectionPassword",
+                    config.configSource.connection.password.get_secret_value(),
+                )
+            if config.configSource.connection.driverName:
+                builder.config(
+                    "spark.hadoop.javax.jdo.option.ConnectionDriverName",
+                    config.configSource.connection.driverName,
+                )
+        if (
+            hasattr(config.configSource.connection, "metastoreFilePath")
+            and config.configSource.connection.metastoreFilePath
+        ):
+            # From https://stackoverflow.com/questions/38377188/how-to-get-rid-of-derby-log-metastore-db-from-spark-shell
+            # derby.system.home is the one in charge of the path for `metastore_db` dir and `derby.log`
+            # We can use this option to control testing, as well as to properly point to the right
+            # local database when ingesting data
+            builder.config(
+                "spark.driver.extraJavaOptions",
+                f"-Dderby.system.home={config.configSource.connection.metastoreFilePath}",
+            )
+
+        for key, value in get_connection_args_common(config).items():
+            builder.config(key, value)
+
+        return cls(spark_session=configure_spark_with_delta_pip(builder).getOrCreate())
+
+    def get_database_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Returns the Database Names, based on the underlying client."""
+        yield service_connection.databaseName or DEFAULT_DATABASE
+
+    def get_database_schema_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Returns the RAW database schema names, based on the underlying client."""
+        for schema in self._spark.catalog.listDatabases():
+            yield schema.name
+
+    def get_table_info(
+        self, service_connection: DeltaLakeConnection, schema_name: str
+    ) -> Iterable[TableInfo]:
+        """Returns the Tables name and type, based on the underlying client."""
+        for table in self._spark.catalog.listTables(dbName=schema_name):
+
+            if table.tableType == SparkTableType.TEMPORARY.value:
+                logger.debug(f"Skipping temporary table {table.name}")
+                continue
+
+            yield TableInfo(
+                schema=schema_name,
+                name=table.name,
+                description=table.description,
+                _type=TABLE_TYPE_MAP.get(table.tableType, TableType.Regular),
+            )
+
+    def update_table_info(self, table_info: TableInfo) -> TableInfo:
+        return TableInfo(
+            schema=table_info.schema,
+            name=table_info.name,
+            _type=table_info._type,
+            location=table_info.location,
+            description=table_info.description,
+            columns=self.get_columns(table_info.schema, table_info.name),
+            table_partitions=table_info.table_partitions,
+        )
+
+    def _check_col_length(self, datatype, col_raw_type):
+        if datatype and datatype.upper() in {"CHAR", "VARCHAR", "BINARYU", "VARBINARY"}:
+            try:
+                return col_raw_type.length if col_raw_type.length else 1
+            except AttributeError:
+                return 1
+        return None
+
+    def _get_display_data_type(self, row):
+        display_data_type = repr(row["data_type"]).lower()
+        for original, new in ARRAY_DATATYPE_REPLACE_MAP.items():
+            display_data_type = display_data_type.replace(original, new)
+        return display_data_type
+
+    def _get_col_info(self, row):
+        parsed_string = ColumnTypeParser._parse_datatype_string(row["data_type"])
+
+        if parsed_string:
+            parsed_string["dataLength"] = self._check_col_length(
+                parsed_string["dataType"], row["data_type"]
+            )
+            if row["data_type"] == "array":
+                array_data_type_display = self._get_display_data_type(row)
+                parsed_string["dataTypeDisplay"] = array_data_type_display
+                # Parse Primitive Datatype string
+                # if Datatype is Array(int) -> Parse int
+                parsed_string[
+                    "arrayDataType"
+                ] = ColumnTypeParser._parse_primitive_datatype_String(
+                    array_data_type_display[
+                        ARRAY_CHILD_START_INDEX:ARRAY_CHILD_END_INDEX
+                    ]
+                )[
+                    "dataType"
+                ]
+
+            column = Column(name=row["col_name"], **parsed_string)
+        else:
+            col_type = re.search(r"^\w+", row["data_type"]).group(0)
+            charlen = re.search(r"\(([\d]+)\)", row["data_type"])
+            if charlen:
+                charlen = int(charlen.group(1))
+            if (
+                col_type.upper() in {"CHAR", "VARCHAR", "VARBINARY", "BINARY"}
+                and charlen is None
+            ):
+                charlen = 1
+            column = Column(
+                name=row["col_name"],
+                description=row.get("comment"),
+                dataType=col_type,
+                dataLength=charlen,
+                displayName=row["data_type"],
+            )
+        return column
+
+    def fetch_view_schema(self, view_name: str) -> Optional[Dict]:
+        try:
+            describe_output = self._spark.sql(
+                f"describe extended {view_name}"
+            ).collect()
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unexpected exception to fetch view schema [{view_name}]: {exc}"
+            )
+            return None
+
+        view_detail = {}
+        col_details = False
+
+        for row in describe_output:
+            row_dict = row.asDict()
+            if col_details:
+                view_detail[row_dict["col_name"]] = row_dict["data_type"]
+            if "# Detailed Table" in row_dict["col_name"]:
+                col_details = True
+        return view_detail.get("View Text")
+
+    def get_columns(self, schema: str, table: str) -> List[Column]:
+        field_dict: Dict[str, Any] = {}
+        table_name = f"{schema}.{table}"
+
+        try:
+            raw_columns = self._spark.sql(f"describe {table_name}").collect()
+            for field in self._spark.table(f"{table_name}").schema:
+                field_dict[field.name] = field
+        except (AnalysisException, ParseException) as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unexpected exception getting columns for [{table_name}]: {exc}"
+            )
+            return []
+
+        parsed_columns: List[Column] = []
+        partition_cols = False
+        for row in raw_columns:
+            col_name = row["col_name"]
+            if col_name == "" or "#" in col_name:
+                partition_cols = True
+                continue
+            if not partition_cols:
+                column = self._get_col_info(row)
+                parsed_columns.append(column)
+
+        return parsed_columns
+
+    def close(self, service_connection: DeltaLakeConnection):
+        """Closes the Client connection."""
+        pass
+
+    def get_test_get_databases_fn(self, config: MetastoreConfig) -> Callable:
+        """Returns a Callable used to test the GetDatabases condition."""
+        return self._spark.catalog.listDatabases
+
+    def get_test_get_tables_fn(self, config: MetastoreConfig) -> Callable:
+        """Returns a Callable used to test the GetTables condition."""
+
+        def test_get_tables():
+            for database in self._spark.catalog.listDatabases():
+                if database:
+                    self._spark.catalog.listTables(database[0])
+                    break
+
+        return test_get_tables

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

@@ -0,0 +1,200 @@
+#  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 S3 Client
+"""
+import traceback
+from functools import partial
+from typing import Callable, Iterable, List, Optional
+
+from deltalake import DeltaTable
+from deltalake.exceptions import TableNotFoundError
+from urllib3.util import Url
+
+from metadata.clients.aws_client import AWSClient
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    PartitionColumnDetails,
+    TableType,
+)
+from metadata.generated.schema.entity.services.connections.database.deltalake.storageConfig import (
+    StorageConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
+    DeltaLakeConnection,
+)
+from metadata.ingestion.source.database.datalake.clients.s3 import DatalakeS3Client
+from metadata.ingestion.source.database.deltalake.clients.base import (
+    DeltalakeBaseClient,
+    TableInfo,
+)
+from metadata.utils.datalake.datalake_utils import ParquetDataFrameColumnParser
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class DeltalakeS3Client(DeltalakeBaseClient):
+    def __init__(self, client: DatalakeS3Client, storage_options: dict):
+        self._client = client
+        self._storage_options = storage_options
+
+    @classmethod
+    def from_config(
+        cls, service_connection: DeltaLakeConnection
+    ) -> "DeltalakeS3Client":
+        # Get the credentials to pass to the storage options
+        aws_client = AWSClient(
+            service_connection.configSource.connection.securityConfig
+        )
+        session = aws_client.create_session()
+
+        profile = session.profile_name
+        region = session.region_name
+
+        credentials = session.get_credentials().get_frozen_credentials()
+
+        endpoint_url = (
+            str(service_connection.configSource.connection.securityConfig.endPointURL)[
+                :-1
+            ]
+            if service_connection.configSource.connection.securityConfig.endPointURL
+            else None
+        )
+
+        storage_options = {
+            "AWS_ACCESS_KEY_ID": credentials.access_key,
+            "AWS_SECRET_ACCESS_KEY": credentials.secret_key,
+            "AWS_SESSION_TOKEN": credentials.token,
+            "AWS_PROFILE": profile,
+            "AWS_REGION": region,
+            "AWS_ENDPOINT_URL": endpoint_url,
+            # Allow to use MinIO
+            "AWS_ALLOW_HTTP": "true",
+        }
+
+        storage_options = {
+            key: value for key, value in storage_options.items() if value
+        }
+
+        return cls(
+            client=DatalakeS3Client.from_config(
+                config=service_connection.configSource.connection
+            ),
+            storage_options=storage_options,
+        )
+
+    @staticmethod
+    def _get_configured_bucket(
+        service_connection: DeltaLakeConnection,
+    ) -> Optional[str]:
+        return (
+            service_connection.configSource.bucketName
+            if service_connection.configSource.bucketName
+            else None
+        )
+
+    @staticmethod
+    def _get_configured_prefix(service_connection: DeltaLakeConnection) -> str:
+        return service_connection.configSource.prefix
+
+    def get_database_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Yields the Database name as set by the DatalakeS3Client."""
+        yield from self._client.get_database_names(service_connection)
+
+    def get_database_schema_names(
+        self, service_connection: DeltaLakeConnection
+    ) -> Iterable[str]:
+        """Yields the Bucket Names as Schema Names from the DatalakeS3Client."""
+        yield from self._client.get_database_schema_names(
+            self._get_configured_bucket(service_connection)
+        )
+
+    def _read_delta_table(self, schema_name: str, prefix: str) -> Optional[DeltaTable]:
+        url = Url(scheme="s3", host=schema_name, path=prefix)
+
+        try:
+            return DeltaTable(
+                url.url,
+                storage_options=self._storage_options,
+                without_files=True,  # Avoids tracking the files that belongs to the Delta Table, reducing memory footprint
+            )
+        except TableNotFoundError:
+            logger.debug(traceback.format_exc())
+            logger.warning("No Delta Table found at path '%s/%s'.", schema_name, prefix)
+            return None
+
+    def _get_columns(self, table) -> List[Column]:
+        return ParquetDataFrameColumnParser(data_frame=table.to_pandas()).get_columns()
+
+    def _get_partitions(self, table) -> Optional[List[PartitionColumnDetails]]:
+        return [
+            PartitionColumnDetails(columnName=column)
+            for column in table.metadata().partition_columns
+        ] or None
+
+    def _get_table_info(self, schema_name: str, prefix: str) -> Iterable[TableInfo]:
+        """Iterates on the s3 'folders' trying to find DeltaTables. If a DeltaTable is found, yield its information."""
+        table = self._read_delta_table(schema_name, prefix)
+
+        if table:
+            name = (
+                table.metadata().name
+                or [part for part in prefix.split("/") if part][-1]
+                or table.metadata().id
+            )
+            yield TableInfo(
+                schema=schema_name,
+                name=name,
+                location=prefix,
+                description=table.metadata().description,
+                _type=TableType.Regular,
+            )
+        else:
+            for folder in self._client.get_folders_prefix(schema_name, prefix):
+                yield from self._get_table_info(schema_name, folder)
+
+    def get_table_info(
+        self, service_connection: DeltaLakeConnection, schema_name: str
+    ) -> Iterable[TableInfo]:
+        """Yield TableInfo from found DeltaTables."""
+        yield from self._get_table_info(
+            schema_name, self._get_configured_prefix(service_connection)
+        )
+
+    def update_table_info(self, table_info: TableInfo) -> TableInfo:
+        table = self._read_delta_table(table_info.schema, table_info.location)
+        return TableInfo(
+            schema=table_info.schema,
+            name=table_info.name,
+            _type=table_info._type,
+            location=table_info.location,
+            description=table_info.description,
+            columns=self._get_columns(table),
+            table_partitions=self._get_partitions(table),
+        )
+
+    def close(self, service_connection: DeltaLakeConnection):
+        # For the S3 Client we don't need to do anything when closing the connection
+        pass
+
+    def get_test_get_databases_fn(self, config: StorageConfig) -> Callable:
+        """Returns a Callable used to test the GetDatabases condition."""
+        if config.bucketName:
+            return partial(self._client._client.list_objects, Bucket=config.bucketName)
+        return self._client._client.list_buckets
+
+    def get_test_get_tables_fn(self, config: StorageConfig) -> Callable:
+        """Returns a Callable used to test the GetTables condition."""
+        return lambda: True

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

@@ -12,100 +12,75 @@
 """
 Source connection handler
 """
+from dataclasses import dataclass
+from functools import singledispatch
 from typing import Optional
 
-import pyspark
-from delta import configure_spark_with_delta_pip
-from pyspark.sql import SparkSession
-
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
+    S3Config,
+)
+from metadata.generated.schema.entity.services.connections.database.deltalake.metastoreConfig import (
+    MetastoreConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.deltalake.storageConfig import (
+    StorageConfig,
+)
 from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
     DeltaLakeConnection,
-    MetastoreDbConnection,
 )
-from metadata.ingestion.connections.builders import get_connection_args_common
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 
 
-def get_connection(connection: DeltaLakeConnection) -> SparkSession:
-    """
-    Create connection
-    """
+@dataclass
+class DeltalakeClient:
+    def __init__(self, client, config) -> None:
+        self.client = client
+        self.config = config
+
+
+@singledispatch
+def get_deltalake_client(connection, config):
+    """Retrieve Deltalake Client from the config"""
+    msg = None
+    if config:
+        msg = f"Config not implemented for type {type(connection)}: {connection}"
+    raise NotImplementedError(msg)
+
 
-    builder = (
-        pyspark.sql.SparkSession.builder.appName(connection.appName or "OpenMetadata")
-        .enableHiveSupport()
-        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
-        .config(
-            "spark.sql.catalog.spark_catalog",
-            "org.apache.spark.sql.delta.catalog.DeltaCatalog",
-        )
-        # Download delta-core jars when creating the SparkSession
-        .config("spark.jars.packages", "io.delta:delta-core_2.12:2.0.0")
+@get_deltalake_client.register
+def _(connection: MetastoreConfig, config: DeltaLakeConnection):
+    from metadata.ingestion.source.database.deltalake.clients.pyspark import (
+        DeltalakePySparkClient,
     )
 
-    # Check that the attribute exists and is properly informed
-    if (
-        hasattr(connection.metastoreConnection, "metastoreHostPort")
-        and connection.metastoreConnection.metastoreHostPort
-    ):
-        builder.config(
-            "hive.metastore.uris",
-            f"thrift://{connection.metastoreConnection.metastoreHostPort}",
-        )
-
-    if isinstance(connection.metastoreConnection, MetastoreDbConnection):
-        if connection.metastoreConnection.metastoreDb:
-            builder.config(
-                "spark.hadoop.javax.jdo.option.ConnectionURL",
-                connection.metastoreConnection.metastoreDb,
-            )
-        if connection.metastoreConnection.jdbcDriverClassPath:
-            builder.config(
-                "spark.driver.extraClassPath",
-                connection.metastoreConnection.jdbcDriverClassPath,
-            )
-        if connection.metastoreConnection.username:
-            builder.config(
-                "spark.hadoop.javax.jdo.option.ConnectionUserName",
-                connection.metastoreConnection.username,
-            )
-        if connection.metastoreConnection.password:
-            builder.config(
-                "spark.hadoop.javax.jdo.option.ConnectionPassword",
-                connection.metastoreConnection.password.get_secret_value(),
-            )
-        if connection.metastoreConnection.driverName:
-            builder.config(
-                "spark.hadoop.javax.jdo.option.ConnectionDriverName",
-                connection.metastoreConnection.driverName,
-            )
-
-    if (
-        hasattr(connection.metastoreConnection, "metastoreFilePath")
-        and connection.metastoreConnection.metastoreFilePath
-    ):
-        # From https://stackoverflow.com/questions/38377188/how-to-get-rid-of-derby-log-metastore-db-from-spark-shell
-        # derby.system.home is the one in charge of the path for `metastore_db` dir and `derby.log`
-        # We can use this option to control testing, as well as to properly point to the right
-        # local database when ingesting data
-        builder.config(
-            "spark.driver.extraJavaOptions",
-            f"-Dderby.system.home={connection.metastoreConnection.metastoreFilePath}",
-        )
-
-    for key, value in get_connection_args_common(connection).items():
-        builder.config(key, value)
-
-    return configure_spark_with_delta_pip(builder).getOrCreate()
+    return DeltalakePySparkClient.from_config(config)
+
+
+@get_deltalake_client.register
+def _(connection: StorageConfig, config: DeltaLakeConnection):
+    from metadata.ingestion.source.database.deltalake.clients.s3 import (
+        DeltalakeS3Client,
+    )
+
+    if isinstance(connection.connection, S3Config):
+        return DeltalakeS3Client.from_config(config)
+
+
+def get_connection(connection: DeltaLakeConnection) -> DeltalakeClient:
+    """Create Deltalake Client"""
+    return DeltalakeClient(
+        client=get_deltalake_client(connection.configSource, connection),
+        config=connection,
+    )
 
 
 def test_connection(
     metadata: OpenMetadata,
-    spark: SparkSession,
+    connection: DeltalakeClient,
     service_connection: DeltaLakeConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
 ) -> None:
@@ -113,16 +88,13 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-
-    def custom_executor():
-        for database in spark.catalog.listDatabases():
-            if database:
-                spark.catalog.listTables(database[0])
-                break
-
     test_fn = {
-        "GetDatabases": spark.catalog.listDatabases,
-        "GetTables": custom_executor,
+        "GetDatabases": connection.client.get_test_get_databases_fn(
+            service_connection.configSource
+        ),
+        "GetTables": connection.client.get_test_get_tables_fn(
+            service_connection.configSource
+        ),
     }
 
     test_connection_steps(