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

metadata/utils/secrets/azure_kv_secrets_manager.py

@@ -0,0 +1,132 @@
+#  Copyright 2022 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.
+
+"""
+Abstract class for AWS based secrets manager implementations
+"""
+import os
+import traceback
+from abc import ABC
+from typing import Optional
+
+from azure.keyvault.secrets import KeyVaultSecret
+
+from metadata.clients.azure_client import AzureClient
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
+from metadata.generated.schema.security.secrets.secretsManagerProvider import (
+    SecretsManagerProvider,
+)
+from metadata.utils.dispatch import enum_register
+from metadata.utils.logger import utils_logger
+from metadata.utils.secrets.external_secrets_manager import (
+    SECRET_MANAGER_AIRFLOW_CONF,
+    ExternalSecretsManager,
+    SecretsManagerConfigException,
+)
+
+logger = utils_logger()
+
+secrets_manager_client_loader = enum_register()
+
+
+# pylint: disable=import-outside-toplevel
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.noop.value)
+def _() -> None:
+    return None
+
+
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.airflow.value)
+def _() -> Optional["AzureCredentials"]:
+    from airflow.configuration import conf
+
+    from metadata.generated.schema.security.credentials.azureCredentials import (
+        AzureCredentials,
+    )
+
+    key_vault_name = conf.get(
+        SECRET_MANAGER_AIRFLOW_CONF, "azure_key_vault_name", fallback=None
+    )
+    if not key_vault_name:
+        raise ValueError(
+            "Missing `azure_key_vault_name` config for Azure Key Vault Secrets Manager Provider."
+        )
+
+    tenant_id = conf.get(SECRET_MANAGER_AIRFLOW_CONF, "azure_tenant_id", fallback=None)
+    client_id = conf.get(SECRET_MANAGER_AIRFLOW_CONF, "azure_client_id", fallback=None)
+    client_secret = conf.get(
+        SECRET_MANAGER_AIRFLOW_CONF, "azure_client_secret", fallback=None
+    )
+
+    return AzureCredentials(
+        clientId=client_id,
+        clientSecret=client_secret,
+        tenantId=tenant_id,
+        vaultName=key_vault_name,
+    )
+
+
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.env.value)
+def _() -> Optional["AzureCredentials"]:
+    from metadata.generated.schema.security.credentials.azureCredentials import (
+        AzureCredentials,
+    )
+
+    # Load only the AZURE_KEY_VAULT_NAME (required) variable and use the
+    # Default Auth chain
+    # https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential
+    key_vault_name = os.getenv("AZURE_KEY_VAULT_NAME")
+
+    if not key_vault_name:
+        raise ValueError(
+            "Missing `azure_key_vault_name` config for Azure Key Vault Secrets Manager Provider."
+        )
+
+    return AzureCredentials(vaultName=key_vault_name)
+
+
+class AzureKVSecretsManager(ExternalSecretsManager, ABC):
+    """
+    Azure Key Vault Secrets Manager class
+    """
+
+    def __init__(
+        self,
+        loader: SecretsManagerClientLoader,
+    ):
+        super().__init__(provider=SecretsManagerProvider.azure_kv, loader=loader)
+
+        self.client = AzureClient(self.credentials).create_secret_client()
+
+    def get_string_value(self, secret_id: str) -> str:
+        """
+        :param secret_id: The secret id to retrieve
+        :return: The value of the secret
+        """
+        try:
+            secret: KeyVaultSecret = self.client.get_secret(secret_id)
+            logger.debug(f"Got value for secret {secret_id}")
+            return secret.value
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Could not get the secret value of {secret_id} due to [{exc}]"
+            )
+            raise exc
+
+    def load_credentials(self) -> Optional["AzureCredentials"]:
+        """Load the provider credentials based on the loader type"""
+        try:
+            loader_fn = secrets_manager_client_loader.registry.get(self.loader.value)
+            return loader_fn()
+        except Exception as err:
+            raise SecretsManagerConfigException(f"Error loading credentials - [{err}]")

metadata/utils/secrets/{noop_secrets_manager.py → db_secrets_manager.py}

@@ -18,12 +18,13 @@ from metadata.generated.schema.security.secrets.secretsManagerProvider import (
 from metadata.utils.secrets.secrets_manager import SecretsManager
 
 
-class NoopSecretsManager(SecretsManager):
+class DBSecretsManager(SecretsManager):
     """
-    LocalSecretsManager is used when there is not a secrets' manager configured.
+    Will return the string as received. The client does not need to pick
+    it up from any external resource.
     """
 
-    provider: str = SecretsManagerProvider.noop.name
+    provider: str = SecretsManagerProvider.db.name
 
     def get_string_value(self, secret_id: str) -> str:
         return secret_id

metadata/utils/secrets/external_secrets_manager.py

@@ -12,18 +12,40 @@
 """
 Abstract class for third party secrets' manager implementations
 """
-from abc import ABC
+from abc import ABC, abstractmethod
+from typing import Any
 
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
 from metadata.generated.schema.security.secrets.secretsManagerProvider import (
     SecretsManagerProvider,
 )
 from metadata.utils.secrets.secrets_manager import SecretsManager
 
+SECRET_MANAGER_AIRFLOW_CONF = "openmetadata_secrets_manager"
+
+
+class SecretsManagerConfigException(Exception):
+    """
+    Invalid config that does not allow us to create
+    the SecretsManagerFactory
+    """
+
 
 class ExternalSecretsManager(SecretsManager, ABC):
     """
     Abstract class for third party secrets' manager implementations
     """
 
-    def __init__(self, provider: SecretsManagerProvider):
-        self.provider = provider.name
+    def __init__(
+        self, provider: SecretsManagerProvider, loader: SecretsManagerClientLoader
+    ):
+        self.provider = provider
+        self.loader = loader
+
+        self.credentials = self.load_credentials()
+
+    @abstractmethod
+    def load_credentials(self) -> Any:
+        """Load the provider credentials based on the loader type"""

metadata/utils/secrets/secrets_manager_factory.py

@@ -12,7 +12,7 @@
 """
 Secrets manager factory module
 """
-from typing import Any, Optional
+from typing import Optional
 
 from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
     SecretsManagerClientLoader,
@@ -22,19 +22,12 @@ from metadata.generated.schema.security.secrets.secretsManagerProvider import (
 )
 from metadata.utils.secrets.aws_secrets_manager import AWSSecretsManager
 from metadata.utils.secrets.aws_ssm_secrets_manager import AWSSSMSecretsManager
-from metadata.utils.secrets.client.loader import secrets_manager_client_loader
-from metadata.utils.secrets.noop_secrets_manager import NoopSecretsManager
+from metadata.utils.secrets.azure_kv_secrets_manager import AzureKVSecretsManager
+from metadata.utils.secrets.db_secrets_manager import DBSecretsManager
 from metadata.utils.secrets.secrets_manager import SecretsManager
 from metadata.utils.singleton import Singleton
 
 
-class SecretsManagerConfigException(Exception):
-    """
-    Invalid config that does not allow us to create
-    the SecretsManagerFactory
-    """
-
-
 class SecretsManagerFactory(metaclass=Singleton):
     """
     Singleton factory to initialize a secret manager. It will return always the same secret manager instance.
@@ -56,11 +49,9 @@ class SecretsManagerFactory(metaclass=Singleton):
         self._secrets_manager_provider = secrets_manager_provider
         self._secrets_manager_loader = secrets_manager_loader
 
-        credentials = self._load_secrets_manager_credentials()
-
         self.secrets_manager = self._get_secrets_manager(
             secrets_manager_provider,
-            credentials,
+            secrets_manager_loader,
         )
 
     @property
@@ -74,44 +65,35 @@ class SecretsManagerFactory(metaclass=Singleton):
     def _get_secrets_manager(
         self,
         secrets_manager_provider: SecretsManagerProvider,
-        credentials: Any = None,
+        secrets_manager_loader: SecretsManagerClientLoader,
     ) -> SecretsManager:
         """
         Method to get the secrets manager based on the arguments passed
         :param secrets_manager_provider: the secrets' manager provider
-        :param credentials: optional credentials that could be required by the clients of the secrets' manager
-                            implementations
+        :param secrets_manager_loader: how to retrieve the secrets manager keys from the environment
         :return: a secrets manager
         """
         if (
             secrets_manager_provider is None
-            or secrets_manager_provider == SecretsManagerProvider.noop
+            or secrets_manager_provider == SecretsManagerProvider.db
         ):
-            return NoopSecretsManager()
+            return DBSecretsManager()
         if secrets_manager_provider in (
             SecretsManagerProvider.aws,
             SecretsManagerProvider.managed_aws,
         ):
-            return AWSSecretsManager(credentials)
+            return AWSSecretsManager(secrets_manager_loader)
         if secrets_manager_provider in (
             SecretsManagerProvider.aws_ssm,
             SecretsManagerProvider.managed_aws_ssm,
         ):
-            return AWSSSMSecretsManager(credentials)
+            return AWSSSMSecretsManager(secrets_manager_loader)
+        if secrets_manager_provider in (
+            SecretsManagerProvider.azure_kv,
+            SecretsManagerProvider.managed_azure_kv,
+        ):
+            return AzureKVSecretsManager(secrets_manager_loader)
         raise NotImplementedError(f"[{secrets_manager_provider}] is not implemented.")
 
     def get_secrets_manager(self):
         return self.secrets_manager
-
-    def _load_secrets_manager_credentials(self) -> Optional["AWSCredentials"]:
-
-        if not self.secrets_manager_loader:
-            return None
-
-        try:
-            loader_fn = secrets_manager_client_loader.registry.get(
-                self.secrets_manager_loader.value
-            )
-            return loader_fn(self.secrets_manager_provider)
-        except Exception as err:
-            raise SecretsManagerConfigException(f"Error loading credentials - [{err}]")

metadata/utils/{source_hash_utils.py → source_hash.py}

@@ -23,11 +23,20 @@ SOURCE_HASH_EXCLUDE_FIELDS = {
 }
 
 
-def generate_source_hash(create_request: C, exclude_fields: Optional[Dict]) -> str:
+def generate_source_hash(
+    create_request: C, exclude_fields: Optional[Dict] = None
+) -> str:
     """
     Given a create_request model convert it to json string and generate a hash value
     """
 
+    # We always want to exclude the sourceHash when generating the fingerprint
+    exclude_fields = (
+        SOURCE_HASH_EXCLUDE_FIELDS.update(exclude_fields)
+        if exclude_fields
+        else SOURCE_HASH_EXCLUDE_FIELDS
+    )
+
     create_request_json = create_request.json(exclude=exclude_fields)
 
     json_bytes = create_request_json.encode("utf-8")

metadata/utils/sqlalchemy_utils.py

@@ -39,6 +39,27 @@ def get_table_comment_wrapper(self, connection, query, table_name, schema=None):
     return {"text": self.all_table_comments.get((table_name, schema))}
 
 
+@reflection.cache
+def get_all_table_owners(
+    self, connection, query, schema_name, **kw
+):  # pylint: disable=unused-argument
+    """
+    Method to fetch owners of all available tables
+    """
+    self.all_table_owners: Dict[Tuple[str, str], str] = {}
+    result = connection.execute(query)
+    for table in result:
+        self.all_table_owners[(table[0], table[1])] = table[2]
+
+
+def get_table_owner_wrapper(
+    self, connection, query, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    if not hasattr(self, "all_table_owners"):
+        self.get_all_table_owners(connection, query, schema)
+    return self.all_table_owners.get((schema, table_name), "")
+
+
 @reflection.cache
 def get_all_view_definitions(self, connection, query):
     """

metadata/utils/storage_metadata_config.py

@@ -17,9 +17,13 @@ from functools import singledispatch
 
 import requests
 
+from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
     AzureConfig,
 )
+from metadata.generated.schema.entity.services.connections.database.datalake.gcsConfig import (
+    GCSConfig,
+)
 from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
     S3Config,
 )
@@ -29,6 +33,9 @@ from metadata.generated.schema.metadataIngestion.storage.manifestMetadataConfig
 from metadata.generated.schema.metadataIngestion.storage.storageMetadataADLSConfig import (
     StorageMetadataAdlsConfig,
 )
+from metadata.generated.schema.metadataIngestion.storage.storageMetadataGCSConfig import (
+    StorageMetadataGcsConfig,
+)
 from metadata.generated.schema.metadataIngestion.storage.storageMetadataHttpConfig import (
     StorageMetadataHttpConfig,
 )
@@ -39,6 +46,7 @@ from metadata.generated.schema.metadataIngestion.storage.storageMetadataS3Config
     StorageMetadataS3Config,
 )
 from metadata.readers.file.config_source_factory import get_reader
+from metadata.utils.credentials import set_google_credentials
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -146,25 +154,45 @@ def _(config: StorageMetadataAdlsConfig) -> ManifestMetadataConfig:
             else STORAGE_METADATA_MANIFEST_FILE_NAME
         )
 
-        from azure.identity import (  # pylint: disable=import-outside-toplevel
-            ClientSecretCredential,
+        blob_client = AzureClient(config.securityConfig).create_blob_client()
+
+        reader = get_reader(
+            config_source=AzureConfig(securityConfig=config.securityConfig),
+            client=blob_client,
         )
-        from azure.storage.blob import (  # pylint: disable=import-outside-toplevel
-            BlobServiceClient,
+
+        manifest = reader.read(path=path, bucket_name=bucket_name)
+        return ManifestMetadataConfig.parse_obj(json.loads(manifest))
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        raise StorageMetadataConfigException(
+            f"Error fetching manifest file from adls: {exc}"
         )
 
-        blob_client = BlobServiceClient(
-            account_url=f"https://{config.securityConfig.accountName}.blob.core.windows.net/",
-            credential=ClientSecretCredential(
-                config.securityConfig.tenantId,
-                config.securityConfig.clientId,
-                config.securityConfig.clientSecret.get_secret_value(),
-            ),
+
+@get_manifest.register
+def _(config: StorageMetadataGcsConfig) -> ManifestMetadataConfig:
+    try:
+        bucket_name, prefix = (
+            config.prefixConfig.containerName,
+            config.prefixConfig.objectPrefix,
+        )
+
+        path = (
+            f"{prefix}/{STORAGE_METADATA_MANIFEST_FILE_NAME}"
+            if prefix
+            else STORAGE_METADATA_MANIFEST_FILE_NAME
+        )
+
+        from google.cloud.storage import (  # pylint: disable=import-outside-toplevel
+            Client,
         )
 
+        set_google_credentials(gcp_credentials=config.securityConfig)
+        gcs_client = Client()
         reader = get_reader(
-            config_source=AzureConfig(securityConfig=config.securityConfig),
-            client=blob_client,
+            config_source=GCSConfig(securityConfig=config.securityConfig),
+            client=gcs_client,
         )
 
         manifest = reader.read(path=path, bucket_name=bucket_name)
@@ -172,5 +200,5 @@ def _(config: StorageMetadataAdlsConfig) -> ManifestMetadataConfig:
     except Exception as exc:
         logger.debug(traceback.format_exc())
         raise StorageMetadataConfigException(
-            f"Error fetching manifest file from s3: {exc}"
+            f"Error fetching manifest file from gcs: {exc}"
         )

metadata/utils/tag_utils.py

@@ -21,6 +21,9 @@ from metadata.generated.schema.api.classification.createClassification import (
 )
 from metadata.generated.schema.api.classification.createTag import CreateTagRequest
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.generated.schema.type.tagLabel import (
     LabelType,
@@ -28,7 +31,7 @@ from metadata.generated.schema.type.tagLabel import (
     TagLabel,
     TagSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils import fqn
@@ -72,7 +75,7 @@ def get_ometa_tag_and_classification(
                     left=StackTraceError(
                         name=tag,
                         error=f"Error yielding tag [{tag}]: [{err}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 

metadata/workflow/application.py

@@ -0,0 +1,154 @@
+#  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.
+"""
+Generic Workflow entrypoint to execute Applications
+"""
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from metadata.config.common import WorkflowExecutionError
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.application import (
+    OpenMetadataApplicationConfig,
+)
+from metadata.generated.schema.metadataIngestion.workflow import LogLevels
+from metadata.ingestion.api.step import Step, Summary
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.importer import import_from_module
+from metadata.utils.logger import ingestion_logger
+from metadata.workflow.base import BaseWorkflow
+from metadata.workflow.workflow_status_mixin import SUCCESS_THRESHOLD_VALUE
+
+logger = ingestion_logger()
+
+
+class InvalidAppConfiguration(Exception):
+    """
+    To be raised if the config received by the App
+    is not the one expected
+    """
+
+
+class AppRunner(Step, ABC):
+    """Class that knows how to execute the Application logic."""
+
+    def __init__(
+        self,
+        config: OpenMetadataApplicationConfig,
+        metadata: OpenMetadata,
+    ):
+        self.app_config = config.appConfig.__root__ if config.appConfig else None
+        self.private_config = (
+            config.appPrivateConfig.__root__ if config.appPrivateConfig else None
+        )
+        self.metadata = metadata
+
+        super().__init__()
+
+    @property
+    def name(self) -> str:
+        return "AppRunner"
+
+    @abstractmethod
+    def run(self) -> None:
+        """App logic to execute"""
+
+    @classmethod
+    def create(cls, config_dict: dict, metadata: OpenMetadata) -> "Step":
+        config = OpenMetadataApplicationConfig.parse_obj(config_dict)
+        return cls(config=config, metadata=metadata)
+
+
+class ApplicationWorkflow(BaseWorkflow, ABC):
+    """Base Application Workflow implementation"""
+
+    config: OpenMetadataApplicationConfig
+    runner: Optional[AppRunner]
+
+    def __init__(self, config_dict: dict):
+        self.runner = None  # Will be passed in post-init
+        # TODO: Create a parse_gracefully method
+        self.config = OpenMetadataApplicationConfig.parse_obj(config_dict)
+
+        # Applications are associated to the OpenMetadata Service
+        self.service_type: ServiceType = ServiceType.Metadata
+
+        metadata_config: OpenMetadataConnection = (
+            self.config.workflowConfig.openMetadataServerConfig
+        )
+        log_level: LogLevels = self.config.workflowConfig.loggerLevel
+
+        super().__init__(
+            config=self.config,
+            log_level=log_level,
+            metadata_config=metadata_config,
+            service_type=self.service_type,
+        )
+
+    @classmethod
+    def create(cls, config_dict: dict):
+        return cls(config_dict)
+
+    def post_init(self) -> None:
+        """
+        Method to execute after we have initialized all the internals.
+        Here we will load the runner since it needs the `metadata` object
+        """
+        runner_class = import_from_module(self.config.sourcePythonClass)
+        if not issubclass(runner_class, AppRunner):
+            raise ValueError(
+                "We need a valid AppRunner to initialize the ApplicationWorkflow!"
+            )
+
+        try:
+            self.runner = runner_class(
+                config=self.config,
+                metadata=self.metadata,
+            )
+        except Exception as exc:
+            logger.error(
+                f"Error trying to init the AppRunner [{self.config.sourcePythonClass}] due to [{exc}]"
+            )
+            raise exc
+
+    def execute_internal(self) -> None:
+        """Workflow-specific logic to execute safely"""
+        self.runner.run()
+
+    def calculate_success(self) -> float:
+        return self.runner.get_status().calculate_success()
+
+    def get_failures(self) -> List[StackTraceError]:
+        return self.workflow_steps()[0].get_status().failures
+
+    def workflow_steps(self) -> List[Step]:
+        return [self.runner]
+
+    def raise_from_status_internal(self, raise_warnings=False):
+        """Check failed status in the runner"""
+        if (
+            self.runner.get_status().failures
+            and self.calculate_success() < SUCCESS_THRESHOLD_VALUE
+        ):
+            raise WorkflowExecutionError(
+                f"{self.runner.name} reported errors: {Summary.from_step(self.runner)}"
+            )
+
+        if raise_warnings and self.runner.get_status().warnings:
+            raise WorkflowExecutionError(
+                f"{self.runner.name} reported warning: {Summary.from_step(self.runner)}"
+            )

metadata/workflow/application_output_handler.py

@@ -0,0 +1,34 @@
+#  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.
+
+"""
+Module handles the output messages for applications
+"""
+import time
+
+from metadata.utils.helpers import pretty_print_time_duration
+from metadata.utils.logger import ANSI, log_ansi_encoded_string
+from metadata.workflow.output_handler import print_workflow_summary
+
+
+def print_status(workflow: "ApplicationWorkflow") -> None:
+    """
+    Print the workflow results
+    """
+    print_workflow_summary(workflow)
+
+    if workflow.runner.get_status().source_start_time:
+        log_ansi_encoded_string(
+            color=ANSI.BRIGHT_CYAN,
+            bold=True,
+            message="Workflow finished in time: "
+            f"{pretty_print_time_duration(time.time()-workflow.runner.get_status().source_start_time)}",
+        )