openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/ingestion/source/database/iceberg/catalog/base.py

@@ -0,0 +1,40 @@
+#  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.
+
+"""
+Iceberg Catalog base module.
+"""
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from pyiceberg.catalog import Catalog
+
+from metadata.generated.schema.entity.services.connections.database.iceberg.icebergCatalog import (
+    IcebergCatalog,
+)
+from metadata.ingestion.source.database.iceberg.fs import (
+    FileSystemConfig,
+    IcebergFileSystemFactory,
+)
+
+
+class IcebergCatalogBase(ABC):
+    @classmethod
+    @abstractmethod
+    def get_catalog(cls, catalog: IcebergCatalog) -> Catalog:
+        """Returns a Catalog for given catalog configuration."""
+
+    @staticmethod
+    def get_fs_parameters(fs_config: Optional[FileSystemConfig]) -> dict:
+        """Gets the FileSystem parameters based on passed configuration."""
+        if not fs_config:
+            return {}
+        return IcebergFileSystemFactory.parse(fs_config)

metadata/ingestion/source/database/iceberg/catalog/dynamodb.py

@@ -0,0 +1,102 @@
+#  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.
+
+"""
+Iceberg DynamoDB Catalog
+"""
+import boto3
+from pyiceberg.catalog import Catalog
+from pyiceberg.catalog.dynamodb import DynamoDbCatalog
+
+from metadata.generated.schema.entity.services.connections.database.iceberg.dynamoDbCatalogConnection import (
+    DynamoDbCatalogConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.iceberg.icebergCatalog import (
+    IcebergCatalog,
+)
+from metadata.ingestion.source.database.iceberg.catalog.base import IcebergCatalogBase
+
+
+class CustomDynamoDbCatalog(DynamoDbCatalog):
+    """Custom DynamoDb Catalog implementation to override the PyIceberg one.
+    This is needed due to PyIceberg not handling the __init__ method correctly by
+    instantiating the Boto3 Client without the correct configuration and having side
+    effects on it.
+    """
+
+    @staticmethod
+    def get_boto3_client(parameters: dict):
+        """
+        Overrides the boto3 client created by PyIceberg.
+        PyIceberg doesn't handle the Boto3 Session.
+        """
+        boto_session_config_keys = [
+            "aws_access_key_id",
+            "aws_secret_access_key",
+            "aws_session_token",
+            "region_name",
+            "profile_name",
+        ]
+
+        session_config = {
+            k: v for k, v in parameters.items() if k in boto_session_config_keys
+        }
+
+        session = boto3.Session(**session_config)
+        return session.client("dynamodb")
+
+    def __init__(self, name: str, **properties: str):
+        # HACK: Runs Catalog.__init__ without running DynamoDbCatalog.__init__
+        super(DynamoDbCatalog, self).__init__(  # pylint: disable=E1003,I0021
+            name, **properties
+        )
+
+        self.dynamodb = self.get_boto3_client(properties)
+        self.dynamodb_table_name = self.properties.get("table-name", "iceberg")
+
+
+class IcebergDynamoDbCatalog(IcebergCatalogBase):
+    """Responsible for building a PyIceberg DynamoDB Catalog."""
+
+    @classmethod
+    def get_catalog(cls, catalog: IcebergCatalog) -> Catalog:
+        """Returns a DynamoDB Catalog for the given connection and file storage.
+
+        For more information, check the PyIceberg [docs](https://py.iceberg.apache.org/configuration/#dynamodb-catalog)
+        """
+        if not isinstance(catalog.connection, DynamoDbCatalogConnection):
+            raise RuntimeError(
+                "'connection' is not an instance of 'DynamoDbCatalogConnection'"
+            )
+
+        parameters = {"warehouse": catalog.warehouseLocation}
+
+        if catalog.connection.tableName:
+            parameters = {"table-name": catalog.connection.tableName}
+
+        if catalog.connection.awsConfig:
+            aws_config = catalog.connection.awsConfig
+
+            parameters = {
+                **parameters,
+                "aws_secret_key_id": aws_config.awsAccessKeyId,
+                "aws_secret_access_key": aws_config.awsSecretAccessKey
+                if aws_config.awsSecretAccessKey
+                else None,
+                "aws_session_token": aws_config.awsSessionToken,
+                "region_name": aws_config.awsRegion,
+                "profile_name": aws_config.profileName,
+                # Needed because the way PyIceberg instantiates the PyArrowFileIO
+                # is different from how they instantiate the Boto3 Client.
+                **cls.get_fs_parameters(aws_config),
+            }
+
+        return CustomDynamoDbCatalog(catalog.name, **parameters)

metadata/ingestion/source/database/iceberg/catalog/glue.py

@@ -0,0 +1,88 @@
+#  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.
+
+"""
+Iceberg Glue Catalog
+"""
+import boto3
+from pyiceberg.catalog import Catalog
+from pyiceberg.catalog.glue import GlueCatalog
+
+from metadata.generated.schema.entity.services.connections.database.iceberg.glueCatalogConnection import (
+    GlueCatalogConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.iceberg.icebergCatalog import (
+    IcebergCatalog,
+)
+from metadata.ingestion.source.database.iceberg.catalog.base import IcebergCatalogBase
+
+
+class IcebergGlueCatalog(IcebergCatalogBase):
+    """Responsible for building a PyIceberg Glue Catalog."""
+
+    @staticmethod
+    def override_boto3_glue_client(catalog: GlueCatalog, parameters: dict):
+        """
+        Overrides the boto3 client created by PyIceberg.
+        This is needed because PyIceberg 0.4.0 has a bug
+        where they try to pass the 'aws_access_key_id' as 'aws_secret_key_id' and
+        it breaks.
+        """
+        boto_session_config_keys = [
+            "aws_access_key_id",
+            "aws_secret_access_key",
+            "aws_session_token",
+            "region_name",
+            "profile_name",
+        ]
+
+        session_config = {
+            k: v for k, v in parameters.items() if k in boto_session_config_keys
+        }
+        session = boto3.Session(**session_config)
+        catalog.glue = session.client("glue")
+
+    @classmethod
+    def get_catalog(cls, catalog: IcebergCatalog) -> Catalog:
+        """Returns a Glue Catalog for the given connection and file storage.
+
+        For more information, check the PyIceberg [docs](https://py.iceberg.apache.org/configuration/#glue-catalog)
+        """
+        if not isinstance(catalog.connection, GlueCatalogConnection):
+            raise RuntimeError(
+                "'connection' is not an instance of 'GlueCatalogConnection'"
+            )
+
+        parameters = {"warehouse": catalog.warehouseLocation}
+
+        if catalog.connection.awsConfig:
+            aws_config = catalog.connection.awsConfig
+
+            parameters = {
+                **parameters,
+                "aws_access_key_id": aws_config.awsAccessKeyId,
+                "aws_secret_access_key": aws_config.awsSecretAccessKey.get_secret_value()
+                if aws_config.awsSecretAccessKey
+                else None,
+                "aws_session_token": aws_config.awsSessionToken,
+                "region_name": aws_config.awsRegion,
+                "profile_name": aws_config.profileName,
+                # Needed because the way PyIceberg instantiates the PyArrowFileIO
+                # is different from how they instantiate the Boto3 Client.
+                **cls.get_fs_parameters(aws_config),
+            }
+
+        glue_catalog: GlueCatalog = GlueCatalog(catalog.name, **parameters)
+
+        # HACK: Overriding the Boto3 Glue client due to PyIceberg 0.4.0 Bug.
+        cls.override_boto3_glue_client(glue_catalog, parameters)
+
+        return glue_catalog

metadata/ingestion/source/database/iceberg/catalog/hive.py

@@ -0,0 +1,51 @@
+#  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.
+
+"""
+Iceberg Hive Catalog
+"""
+from pyiceberg.catalog import Catalog, load_hive
+
+from metadata.generated.schema.entity.services.connections.database.iceberg.hiveCatalogConnection import (
+    HiveCatalogConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.iceberg.icebergCatalog import (
+    IcebergCatalog,
+)
+from metadata.ingestion.source.database.iceberg.catalog.base import IcebergCatalogBase
+
+
+class IcebergHiveCatalog(IcebergCatalogBase):
+    """Responsible for building a PyIceberg Hive Catalog."""
+
+    @classmethod
+    def get_catalog(cls, catalog: IcebergCatalog) -> Catalog:
+        """Returns a Hive Catalog for the given connection and file storage.
+
+        For more information, check the PyIceberg [docs](https://py.iceberg.apache.org/configuration/#hive-catalog)
+        """
+        if not isinstance(catalog.connection, HiveCatalogConnection):
+            raise RuntimeError(
+                "'connection' is not an instance of 'HiveCatalogConnection'"
+            )
+
+        parameters = {
+            "warehouse": catalog.warehouseLocation,
+            "uri": catalog.connection.uri,
+        }
+
+        if catalog.connection.fileSystem:
+            parameters = {
+                **parameters,
+                **cls.get_fs_parameters(catalog.connection.fileSystem.type),
+            }
+
+        return load_hive(catalog.name, parameters)

metadata/ingestion/source/database/iceberg/catalog/rest.py

@@ -0,0 +1,84 @@
+#  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.
+
+"""
+Iceberg Rest Catalog
+"""
+from pyiceberg.catalog import Catalog, load_rest
+
+from metadata.generated.schema.entity.services.connections.database.iceberg.icebergCatalog import (
+    IcebergCatalog,
+)
+from metadata.generated.schema.entity.services.connections.database.iceberg.restCatalogConnection import (
+    RestCatalogConnection,
+)
+from metadata.ingestion.source.database.iceberg.catalog.base import IcebergCatalogBase
+
+
+class IcebergRestCatalog(IcebergCatalogBase):
+    """Responsible for building a PyIceberg Rest Catalog."""
+
+    @classmethod
+    def get_catalog(cls, catalog: IcebergCatalog) -> Catalog:
+        """Returns a Rest Catalog for the given connection and file storage.
+
+        For more information, check the PyIceberg [docs](https://py.iceberg.apache.org/configuration/#rest-catalog)
+        """
+        if not isinstance(catalog.connection, RestCatalogConnection):
+            raise RuntimeError(
+                "'connection' is not an instance of 'RestCatalogConnection'"
+            )
+
+        if catalog.connection.credential:
+            client_id = catalog.connection.credential.clientId.get_secret_value()
+            client_secret = (
+                catalog.connection.credential.clientSecret.get_secret_value()
+            )
+
+            credential = f"{client_id}:{client_secret}"
+        else:
+            credential = None
+
+        parameters = {
+            "warehouse": catalog.warehouseLocation,
+            "uri": catalog.connection.uri,
+            "credential": credential,
+            "token": catalog.connection.token.get_secret_value()
+            if catalog.connection.token
+            else None,
+        }
+
+        if catalog.connection.fileSystem:
+            parameters = {
+                **parameters,
+                **cls.get_fs_parameters(catalog.connection.fileSystem.type),
+            }
+
+        if catalog.connection.ssl:
+            parameters = {
+                **parameters,
+                "ssl": {
+                    "client": {
+                        "cert": catalog.connection.ssl.clientCertPath,
+                        "key": catalog.connection.ssl.privateKeyPath,
+                    },
+                    "cabundle": catalog.connection.ssl.caCertPath,
+                },
+            }
+
+        if catalog.connection.sigv4:
+            parameters = {
+                **parameters,
+                "rest.sigv4": True,
+                "rest.signing_region": catalog.connection.sigv4.signingRegion,
+                "rest.signing_name": catalog.connection.sigv4.signingName,
+            }
+        return load_rest(catalog.name, parameters)

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

@@ -0,0 +1,68 @@
+#  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.
+
+"""
+Source connection handler
+"""
+
+from typing import Optional
+
+from pyiceberg.catalog import Catalog
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.icebergConnection import (
+    IcebergConnection,
+)
+from metadata.ingestion.connections.test_connections import test_connection_steps
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.iceberg.catalog import IcebergCatalogFactory
+
+
+def get_connection(connection: IcebergConnection) -> Catalog:
+    """
+    Create connection
+    """
+    return IcebergCatalogFactory.from_connection(
+        connection.catalog,
+    )
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    catalog: Catalog,
+    service_connection: IcebergConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+
+    def custom_executor_for_namespaces():
+        list(catalog.list_namespaces())
+
+    def custom_executor_for_tables():
+        for namespace in catalog.list_namespaces():
+            return list(catalog.list_tables(namespace))
+
+    test_fn = {
+        "GetNamespaces": custom_executor_for_namespaces,
+        "GetTables": custom_executor_for_tables,
+    }
+
+    test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+    )

metadata/ingestion/source/database/iceberg/fs/__init__.py

@@ -0,0 +1,52 @@
+#  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.
+
+"""
+Iceberg FileSystem Factory module.
+"""
+from __future__ import annotations
+
+from typing import Dict, Optional, Type
+
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+from metadata.ingestion.source.database.iceberg.fs.azure import AzureFileSystem
+from metadata.ingestion.source.database.iceberg.fs.base import (
+    FileSystemConfig,
+    IcebergFileSystemBase,
+)
+from metadata.ingestion.source.database.iceberg.fs.s3 import S3FileSystem
+
+
+class IcebergFileSystemFactory:
+    """Factory class to return any supported PyIceberg FileSystem."""
+
+    file_system_config_map: Dict[str, Type[IcebergFileSystemBase]] = {
+        AWSCredentials.__name__: S3FileSystem,
+        AzureCredentials.__name__: AzureFileSystem,
+    }
+
+    @classmethod
+    def parse(cls, fs_config: Optional[FileSystemConfig]) -> dict:
+        """Returns the Iceberg FileSystem parameters, depending on specific implementation."""
+        if not fs_config:
+            return {}
+
+        file_system_type = cls.file_system_config_map.get(fs_config.__class__.__name__)
+
+        if not file_system_type:
+            raise NotImplementedError(
+                f"Iceberg File System otype ['{fs_config.__class__.__name__}'] Not Implemented."
+            )
+
+        return file_system_type.get_fs_params(fs_config)

metadata/ingestion/source/database/iceberg/fs/azure.py

@@ -0,0 +1,44 @@
+#  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.
+
+"""
+Iceberg S3 File System.
+"""
+from __future__ import annotations
+
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+from metadata.ingestion.source.database.iceberg.fs.base import (
+    FileSystemConfig,
+    IcebergFileSystemBase,
+)
+
+
+class AzureFileSystem(IcebergFileSystemBase):
+    @classmethod
+    def get_fs_params(cls, fs_config: FileSystemConfig) -> dict:
+        """Returns the parameters expected by PyIceberg for Azure Data Lake.
+
+        For more information, check the
+            [PyIceberg documentation](https://py.iceberg.apache.org/configuration/#azure-data-lake).
+        """
+        if not isinstance(fs_config, AzureCredentials):
+            raise RuntimeError(
+                "FileSystem Configuration is not an instance of 'AzureCredentials'."
+            )
+
+        return {
+            "adlfs.account-name": fs_config.accountName,
+            "adlfs.tenant-id": fs_config.tenantId,
+            "adlfs.client-id": fs_config.clientId,
+            "adlfs.client-secret": fs_config.clientSecret.get_secret_value(),
+        }

metadata/ingestion/source/database/iceberg/fs/base.py

@@ -0,0 +1,30 @@
+#  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.
+
+"""
+Iceberg FileSystem base module.
+"""
+from abc import ABC, abstractmethod
+from typing import Union
+
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+
+FileSystemConfig = Union[AWSCredentials, AzureCredentials, None]
+
+
+class IcebergFileSystemBase(ABC):
+    @classmethod
+    @abstractmethod
+    def get_fs_params(cls, fs_config: FileSystemConfig) -> dict:
+        """Returns a Catalog for given catalog_config."""

metadata/ingestion/source/database/iceberg/fs/s3.py

@@ -0,0 +1,77 @@
+#  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.
+
+"""
+Iceberg S3 File System.
+"""
+from __future__ import annotations
+
+from pyiceberg.io import (
+    S3_ACCESS_KEY_ID,
+    S3_ENDPOINT,
+    S3_REGION,
+    S3_SECRET_ACCESS_KEY,
+    S3_SESSION_TOKEN,
+)
+
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.ingestion.source.database.iceberg.fs.base import (
+    FileSystemConfig,
+    IcebergFileSystemBase,
+)
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+SUPPORTED_KEYS = [
+    "endPointURL",
+    "awsAccessKeyId",
+    "awsSecretAccessKey",
+    "awsSessionToken",
+    "awsRegion",
+]
+
+
+class S3FileSystem(IcebergFileSystemBase):
+    """Responsible for returning a PyIceberg S3FileSystem compatible configuration."""
+
+    @classmethod
+    def get_fs_params(cls, fs_config: FileSystemConfig) -> dict:
+        """Returns the parameters expected by PyIceberg for AWS S3.
+
+        For more information, check the [PyIceberg documentation](https://py.iceberg.apache.org/configuration/#s3).
+        """
+        if not isinstance(fs_config, AWSCredentials):
+            raise RuntimeError(
+                "FileSystem Configuration is not an instance of 'AWSCredentials'."
+            )
+
+        unused_keys = []
+
+        for key, value in fs_config.dict().items():
+            if key not in SUPPORTED_KEYS and value:
+                unused_keys.append(key)
+
+        if unused_keys:
+            logger.warning(
+                "[%s] Not Supported by the Iceberg Connector. They won't be used.",
+                ",".join(unused_keys),
+            )
+
+        return {
+            S3_ENDPOINT: fs_config.endPointURL,
+            S3_ACCESS_KEY_ID: fs_config.awsAccessKeyId,
+            S3_SECRET_ACCESS_KEY: fs_config.awsSecretAccessKey.get_secret_value()
+            if fs_config.awsSecretAccessKey
+            else None,
+            S3_SESSION_TOKEN: fs_config.awsSessionToken,
+            S3_REGION: fs_config.awsRegion,
+        }