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

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

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

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

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

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