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

metadata/ingestion/ometa/auth_provider.py

@@ -11,51 +11,24 @@
 """
 Interface definition for an Auth provider
 """
-import json
 import os.path
-import sys
-import traceback
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Tuple
 
-import requests
 from dateutil.relativedelta import relativedelta
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
-from metadata.generated.schema.security.client.auth0SSOClientConfig import (
-    Auth0SSOClientConfig,
-)
-from metadata.generated.schema.security.client.azureSSOClientConfig import (
-    AzureSSOClientConfig,
-)
-from metadata.generated.schema.security.client.customOidcSSOClientConfig import (
-    CustomOIDCSSOClientConfig,
-)
-from metadata.generated.schema.security.client.googleSSOClientConfig import (
-    GoogleSSOClientConfig,
-)
-from metadata.generated.schema.security.client.oktaSSOClientConfig import (
-    OktaSSOClientConfig,
-)
 from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
     OpenMetadataJWTClientConfig,
 )
-from metadata.ingestion.ometa.client import APIError
 from metadata.utils.logger import ometa_logger
 
-# Only load security providers on call
-# pylint: disable=import-outside-toplevel
-
 logger = ometa_logger()
 
-ACCESS_TOKEN = "access_token"
-EXPIRY = "expires_in"
-
 
 class AuthenticationException(Exception):
     """
@@ -74,7 +47,6 @@ class AuthenticationProvider(metaclass=ABCMeta):
     def create(cls, config: ConfigModel) -> "AuthenticationProvider":
         """
         Create authentication
-
         Arguments:
             config (ConfigModel): configuration
         Returns:
@@ -85,7 +57,6 @@ class AuthenticationProvider(metaclass=ABCMeta):
     def auth_token(self) -> str:
         """
         Authentication token
-
         Returns:
             str
         """
@@ -94,331 +65,11 @@ class AuthenticationProvider(metaclass=ABCMeta):
     def get_access_token(self):
         """
         Authentication token
-
         Returns:
             str
         """
 
 
-class NoOpAuthenticationProvider(AuthenticationProvider):
-    """
-    Extends AuthenticationProvider class
-
-    Args:
-        config (MetadataServerConfig):
-
-    Attributes:
-        config (MetadataServerConfig)
-    """
-
-    def __init__(self, config: OpenMetadataConnection):
-        self.config = config
-
-    @classmethod
-    def create(cls, config: OpenMetadataConnection):
-        return cls(config)
-
-    def auth_token(self):
-        pass
-
-    def get_access_token(self):
-        return "no_token", None
-
-
-class GoogleAuthenticationProvider(AuthenticationProvider):
-    """
-    Google authentication implementation
-
-    Args:
-        config (MetadataServerConfig):
-
-    Attributes:
-        config (MetadataServerConfig)
-    """
-
-    def __init__(self, config: OpenMetadataConnection):
-        self.config = config
-        self.security_config: GoogleSSOClientConfig = self.config.securityConfig
-
-        self.generated_auth_token = None
-        self.expiry = None
-
-    @classmethod
-    def create(cls, config: OpenMetadataConnection):
-        return cls(config)
-
-    def auth_token(self) -> None:
-        import google.auth
-        import google.auth.transport.requests
-        from google.oauth2 import service_account
-
-        credentials = service_account.IDTokenCredentials.from_service_account_file(
-            self.security_config.secretKey.get_secret_value(),
-            target_audience=self.security_config.audience,
-        )
-        request = google.auth.transport.requests.Request()
-        credentials.refresh(request)
-        self.generated_auth_token = credentials.token
-        self.expiry = credentials.expiry
-
-    def get_access_token(self):
-        self.auth_token()
-        return self.generated_auth_token, self.expiry
-
-
-class OktaAuthenticationProvider(AuthenticationProvider):
-    """
-    Prepare the Json Web Token for Okta auth
-    """
-
-    def __init__(self, config: OpenMetadataConnection):
-        self.config = config
-        self.security_config: OktaSSOClientConfig = self.config.securityConfig
-
-        self.generated_auth_token = None
-        self.expiry = None
-
-    @classmethod
-    def create(cls, config: OpenMetadataConnection):
-        return cls(config)
-
-    async def auth_token(  # pylint: disable=invalid-overridden-method, too-many-locals
-        self,
-    ) -> None:
-        import time
-        import uuid
-        from urllib.parse import quote, urlencode
-
-        from okta.cache.okta_cache import OktaCache
-        from okta.jwt import JWT, jwt
-        from okta.request_executor import RequestExecutor
-
-        try:
-            _, my_jwk = JWT.get_PEM_JWK(
-                self.security_config.privateKey.get_secret_value()
-            )
-            issued_time = int(time.time())
-            expiry_time = issued_time + JWT.ONE_HOUR
-            generated_jwt_id = str(uuid.uuid4())
-            claims = {
-                "sub": self.security_config.clientId,
-                "iat": issued_time,
-                "exp": expiry_time,
-                "iss": self.security_config.clientId,
-                "aud": self.security_config.orgURL,
-                "jti": generated_jwt_id,
-            }
-            jwt_token = jwt.encode(claims, my_jwk.to_dict(), JWT.HASH_ALGORITHM)
-            config = {
-                "client": {
-                    "orgUrl": self.security_config.orgURL,
-                    "authorizationMode": "BEARER",
-                    "rateLimit": {},
-                    "privateKey": self.security_config.privateKey,
-                    "clientId": self.security_config.clientId,
-                    "token": jwt_token,
-                    "scopes": self.security_config.scopes,
-                }
-            }
-            request_exec = RequestExecutor(
-                config=config, cache=OktaCache(ttl=expiry_time, tti=issued_time)
-            )
-            parameters = {
-                "grant_type": "client_credentials",
-                "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
-                "client_assertion": jwt_token,
-            }
-            if config["client"].get("scopes"):
-                parameters["scope"] = " ".join(config["client"]["scopes"])
-            encoded_parameters = urlencode(parameters, quote_via=quote)
-            url = f"{self.security_config.orgURL}?" + encoded_parameters
-            token_request_object = await request_exec.create_request(
-                "POST",
-                url,
-                None,
-                {
-                    "Accept": "application/json",
-                    "Content-Type": "application/x-www-form-urlencoded",
-                },
-                oauth=True,
-            )
-            _, _, res_json, err = await request_exec.fire_request(
-                token_request_object[0]
-            )
-            if err:
-                raise APIError(f"{err}")
-
-            try:
-                response_dict = json.loads(res_json)
-            except ValueError:
-                raise AuthenticationException(
-                    "Could not fetch the access token please validate the orgURL & clientId in configuration"
-                )
-
-            token = response_dict.get(ACCESS_TOKEN)
-            if not token:
-                raise AuthenticationException(
-                    f"Error getting access token: {response_dict}"
-                )
-
-            self.generated_auth_token = token
-            self.expiry = response_dict.get(EXPIRY)
-
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.error(f"Unexpected error fetching Okta auth token: {exc}")
-            sys.exit()
-
-    def get_access_token(self):
-        import asyncio
-
-        asyncio.run(self.auth_token())
-        return self.generated_auth_token, self.expiry
-
-
-class Auth0AuthenticationProvider(AuthenticationProvider):
-    """
-    OAuth authentication implementation
-    Args:
-        config (MetadataServerConfig):
-    Attributes:
-        config (MetadataServerConfig)
-    """
-
-    def __init__(self, config: OpenMetadataConnection):
-        self.config = config
-        self.security_config: Auth0SSOClientConfig = self.config.securityConfig
-
-        self.generated_auth_token = None
-        self.expiry = None
-
-    @classmethod
-    def create(cls, config: OpenMetadataConnection):
-        return cls(config)
-
-    def auth_token(self) -> None:
-
-        payload = (
-            f"grant_type=client_credentials&client_id={self.security_config.clientId}"
-            f"&client_secret={self.security_config.secretKey.get_secret_value()}"
-            f"&audience=https://{self.security_config.domain}/api/v2/"
-        )
-        headers = {"content-type": "application/x-www-form-urlencoded"}
-        res = requests.post(
-            url=f"https://{self.security_config.domain}/oauth/token",
-            data=payload,
-            headers=headers,
-            timeout=60 * 5,
-        )
-
-        data = json.loads(res.read().decode("utf-8"))
-
-        token = data.get(ACCESS_TOKEN)
-        if not token:
-            raise AuthenticationException(f"Error getting access token: {data}")
-
-        self.generated_auth_token = token
-        self.expiry = data[EXPIRY]
-
-    def get_access_token(self):
-        self.auth_token()
-        return self.generated_auth_token, self.expiry
-
-
-class AzureAuthenticationProvider(AuthenticationProvider):
-    """
-    Prepare the Json Web Token for Azure auth
-    """
-
-    # TODO: Prepare JSON for Azure Auth
-    def __init__(self, config: OpenMetadataConnection):
-        self.config = config
-        self.security_config: AzureSSOClientConfig = self.config.securityConfig
-        self.generated_auth_token = None
-        self.expiry = None
-
-    @classmethod
-    def create(cls, config: OpenMetadataConnection):
-        return cls(config)
-
-    def auth_token(self) -> None:
-        from msal import ConfidentialClientApplication
-
-        app = ConfidentialClientApplication(
-            client_id=self.security_config.clientId,
-            client_credential=self.security_config.clientSecret.get_secret_value(),
-            authority=self.security_config.authority,
-        )
-        data = app.acquire_token_for_client(scopes=self.security_config.scopes)
-
-        token = data.get(ACCESS_TOKEN)
-        if not token:
-            raise AuthenticationException(f"Error getting access token: {data}")
-
-        self.generated_auth_token = token
-        self.expiry = data.get(EXPIRY)
-
-    def get_access_token(self):
-        self.auth_token()
-        return self.generated_auth_token, self.expiry
-
-
-class CustomOIDCAuthenticationProvider(AuthenticationProvider):
-    """
-    Custom OIDC authentication implementation
-
-    Args:
-        config (MetadataServerConfig):
-
-    Attributes:
-        config (MetadataServerConfig)
-    """
-
-    def __init__(self, config: OpenMetadataConnection) -> None:
-        self.config = config
-        self.security_config: CustomOIDCSSOClientConfig = self.config.securityConfig
-
-        self.generated_auth_token = None
-        self.expiry = None
-
-    @classmethod
-    def create(
-        cls, config: OpenMetadataConnection
-    ) -> "CustomOIDCAuthenticationProvider":
-        return cls(config)
-
-    def auth_token(self) -> None:
-        data = {
-            "grant_type": "client_credentials",
-            "client_id": self.security_config.clientId,
-            "client_secret": self.security_config.secretKey.get_secret_value(),
-        }
-        response = requests.post(
-            url=self.security_config.tokenEndpoint,
-            data=data,
-            timeout=60 * 5,
-        )
-        if response.ok:
-            response_json = response.json()
-
-            token = response_json.get(ACCESS_TOKEN)
-            if not token:
-                raise AuthenticationException(
-                    f"Error getting access token: {response_json}"
-                )
-
-            self.generated_auth_token = token
-            self.expiry = response_json.get(EXPIRY)
-        else:
-            raise APIError(
-                error={"message": response.text}, http_error=response.status_code
-            )
-
-    def get_access_token(self) -> Tuple[str, int]:
-        self.auth_token()
-        return self.generated_auth_token, self.expiry
-
-
 class OpenMetadataAuthenticationProvider(AuthenticationProvider):
     """
     OpenMetadata authentication implementation

metadata/ingestion/ometa/client.py

@@ -21,6 +21,7 @@ from requests.exceptions import HTTPError
 
 from metadata.config.common import ConfigModel
 from metadata.ingestion.ometa.credentials import URL, get_api_version
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -155,9 +156,12 @@ class REST:
                     self.config.expires_in = (
                         datetime.datetime.utcnow().timestamp() + expiry - 120
                     )
-        headers[
-            self.config.auth_header
-        ] = f"{self._auth_token_mode} {self.config.access_token}"
+
+        headers[self.config.auth_header] = (
+            f"{self._auth_token_mode} {self.config.access_token}"
+            if self._auth_token_mode
+            else self.config.access_token
+        )
 
         # Merge extra headers if provided.
         # If a header value is provided in modulo string format and matches an existing header,
@@ -248,6 +252,7 @@ class REST:
 
         return None
 
+    @calculate_execution_time(context="GET")
     def get(self, path, data=None):
         """
         GET method
@@ -261,6 +266,7 @@ class REST:
         """
         return self._request("GET", path, data)
 
+    @calculate_execution_time(context="POST")
     def post(self, path, data=None):
         """
         POST method
@@ -274,6 +280,7 @@ class REST:
         """
         return self._request("POST", path, data)
 
+    @calculate_execution_time(context="PUT")
     def put(self, path, data=None):
         """
         PUT method
@@ -287,6 +294,7 @@ class REST:
         """
         return self._request("PUT", path, data)
 
+    @calculate_execution_time(context="PATCH")
     def patch(self, path, data=None):
         """
         PATCH method
@@ -305,6 +313,7 @@ class REST:
             headers={"Content-type": "application/json-patch+json"},
         )
 
+    @calculate_execution_time(context="DELETE")
     def delete(self, path, data=None):
         """
         DELETE method

metadata/ingestion/ometa/mixins/custom_property_mixin.py

@@ -15,7 +15,8 @@ To be used by OpenMetadata class
 """
 from typing import Dict
 
-from metadata.generated.schema.api.data.createCustomProperty import PropertyType
+from metadata.generated.schema.type.customProperty import PropertyType
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.models.custom_properties import (
     CustomPropertyDataTypes,
     CustomPropertyType,
@@ -54,16 +55,6 @@ class OMetaCustomPropertyMixin:
             f"/metadata/types/name/{entity_type}?category=field"
         )
 
-        # Get the data type of the custom property
-        if not ometa_custom_property.createCustomPropertyRequest.propertyType:
-            custom_property_type = self.get_custom_property_type(
-                data_type=ometa_custom_property.custom_property_type
-            )
-            property_type = PropertyType(id=custom_property_type.id, type="type")
-            ometa_custom_property.createCustomPropertyRequest.propertyType = (
-                property_type
-            )
-
         resp = self.client.put(
             f"/metadata/types/{entity_schema.get('id')}",
             data=ometa_custom_property.createCustomPropertyRequest.json(),
@@ -78,3 +69,12 @@ class OMetaCustomPropertyMixin:
         """
         resp = self.client.get(f"/metadata/types/name/{data_type.value}?category=field")
         return CustomPropertyType(**resp)
+
+    def get_property_type_ref(self, data_type: CustomPropertyDataTypes) -> PropertyType:
+        """
+        Get the PropertyType for custom properties
+        """
+        custom_property_type = self.get_custom_property_type(data_type=data_type)
+        return PropertyType(
+            __root__=EntityReference(id=custom_property_type.id, type="type")
+        )

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -79,6 +79,23 @@ class ESMixin(Generic[T]):
 
         return None
 
+    def _get_entity_from_es(
+        self, entity: Type[T], query_string: str, fields: Optional[list] = None
+    ) -> Optional[T]:
+        """Fetch an entity instance from ES"""
+
+        try:
+            entity_list = self._search_es_entity(
+                entity_type=entity, query_string=query_string, fields=fields
+            )
+            for instance in entity_list or []:
+                return instance
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get {entity.__name__} info from ES due to {err}")
+
+        return None
+
     def es_search_from_fqn(
         self,
         entity_type: Type[T],

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -15,9 +15,9 @@ To be used by OpenMetadata class
 """
 import json
 import traceback
+from copy import deepcopy
 from typing import Dict, List, Optional, Type, TypeVar, Union
 
-import jsonpatch
 from pydantic import BaseModel
 
 from metadata.generated.schema.entity.automations.workflow import (
@@ -35,6 +35,7 @@ from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.generated.schema.type.lifeCycle import LifeCycle
 from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.ingestion.api.models import Entity
+from metadata.ingestion.models.patch_request import build_patch
 from metadata.ingestion.models.table_metadata import ColumnDescription, ColumnTag
 from metadata.ingestion.ometa.client import REST
 from metadata.ingestion.ometa.mixins.patch_mixin_utils import (
@@ -45,7 +46,7 @@ from metadata.ingestion.ometa.mixins.patch_mixin_utils import (
 )
 from metadata.ingestion.ometa.utils import model_str
 from metadata.utils.deprecation import deprecated
-from metadata.utils.logger import ometa_logger
+from metadata.utils.logger import get_log_name, ometa_logger
 
 logger = ometa_logger()
 
@@ -92,10 +93,8 @@ def update_column_description(
         desc_column = col_dict.get(col.fullyQualifiedName.__root__)
         if desc_column:
             if col.description and not force:
-                logger.warning(
-                    f"The entity with id [{model_str(col.fullyQualifiedName)}] already has a description."
-                    " To overwrite it, set `force` to True."
-                )
+                # If the description is already present and force is not passed,
+                # description will not be overridden
                 continue
 
             col.description = desc_column.__root__
@@ -120,6 +119,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         destination: T,
         allowed_fields: Optional[Dict] = None,
         restrict_update_fields: Optional[List] = None,
+        array_entity_fields: Optional[List] = None,
     ) -> Optional[T]:
         """
         Given an Entity type and Source entity and Destination entity,
@@ -136,47 +136,17 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             Updated Entity
         """
         try:
-            # Get the difference between source and destination
-            if allowed_fields:
-                patch = jsonpatch.make_patch(
-                    json.loads(
-                        source.json(
-                            exclude_unset=True,
-                            exclude_none=True,
-                            include=allowed_fields,
-                        )
-                    ),
-                    json.loads(
-                        destination.json(
-                            exclude_unset=True,
-                            exclude_none=True,
-                            include=allowed_fields,
-                        )
-                    ),
-                )
-            else:
-                patch = jsonpatch.make_patch(
-                    json.loads(source.json(exclude_unset=True, exclude_none=True)),
-                    json.loads(destination.json(exclude_unset=True, exclude_none=True)),
-                )
+            patch = build_patch(
+                source=source,
+                destination=destination,
+                allowed_fields=allowed_fields,
+                restrict_update_fields=restrict_update_fields,
+                array_entity_fields=array_entity_fields,
+            )
+
             if not patch:
-                logger.debug(
-                    "Nothing to update when running the patch. Are you passing `force=True`?"
-                )
                 return None
 
-            # for a user editable fields like descriptions, tags we only want to support "add" operation in patch
-            # we will remove the other operations for replace, remove from here
-            if restrict_update_fields:
-                patch.patch = [
-                    patch_ops
-                    for patch_ops in patch.patch
-                    if self._determine_restricted_operation(
-                        patch_ops=patch_ops,
-                        restrict_update_fields=restrict_update_fields,
-                    )
-                ]
-
             res = self.client.patch(
                 path=f"{self.get_suffix(entity)}/{model_str(source.id)}",
                 data=str(patch),
@@ -185,25 +155,10 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.error(
-                f"Error trying to PATCH {entity.__name__} [{source.id.__root__}]: {exc}"
-            )
+            logger.error(f"Error trying to PATCH {get_log_name(source)}: {exc}")
 
         return None
 
-    def _determine_restricted_operation(
-        self, patch_ops: Dict, restrict_update_fields: Optional[List] = None
-    ) -> bool:
-        """
-        Only retain add operation for restrict_update_fields fields
-        """
-        path = patch_ops.get("path")
-        op = patch_ops.get("op")
-        for field in restrict_update_fields or []:
-            if field in path and op != PatchOperation.ADD.value:
-                return False
-        return True
-
     def patch_description(
         self,
         entity: Type[T],
@@ -238,10 +193,8 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             return None
 
         if instance.description and not force:
-            logger.warning(
-                f"The entity with id [{model_str(source.id)}] already has a description."
-                " To overwrite it, set `force` to True."
-            )
+            # If the description is already present and force is not passed,
+            # description will not be overridden
             return None
 
         # https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
@@ -391,18 +344,14 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
 
         # Don't change existing data without force
         if instance.owner and not force:
-            logger.warning(
-                f"The entity with id [{model_str(source.id)}] already has an owner."
-                " To overwrite it, set `overrideOwner` to True."
-            )
+            # If a owner is already present and force is not passed,
+            # owner will not be overridden
             return None
 
-        source.owner = instance.owner
-
-        destination = source.copy(deep=True)
+        destination = deepcopy(instance)
         destination.owner = owner
 
-        return self.patch(entity=entity, source=source, destination=destination)
+        return self.patch(entity=entity, source=instance, destination=destination)
 
     def patch_column_tags(
         self,