openmetadata-ingestion 1.2.5.1__py3-none-any.whl → 1.3.0.0__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/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 (
@@ -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__
@@ -136,47 +135,16 @@ 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,
+            )
+
             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),
@@ -191,19 +159,6 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
 
         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,

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -0,0 +1,41 @@
+#  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.
+"""
+Mixin class containing Suggestions specific methods
+
+To be used by OpenMetadata class
+"""
+from metadata.generated.schema.entity.feed.suggestion import Suggestion
+from metadata.ingestion.ometa.client import REST
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaSuggestionsMixin:
+    """
+    OpenMetadata API methods related to the Suggestion Entity
+
+    To be inherited by OpenMetadata
+    """
+
+    client: REST
+
+    def update_suggestion(self, suggestion: Suggestion) -> Suggestion:
+        """
+        Update an existing Suggestion with new fields
+        """
+        resp = self.client.put(
+            f"{self.get_suffix(Suggestion)}/{str(suggestion.id.__root__)}",
+            data=suggestion.json(),
+        )
+
+        return Suggestion(**resp)

metadata/ingestion/ometa/mixins/table_mixin.py

@@ -22,6 +22,9 @@ from requests.utils import quote
 from metadata.generated.schema.api.data.createTableProfile import (
     CreateTableProfileRequest,
 )
+from metadata.generated.schema.api.tests.createCustomMetric import (
+    CreateCustomMetricRequest,
+)
 from metadata.generated.schema.entity.data.table import (
     ColumnProfile,
     DataModel,
@@ -288,3 +291,18 @@ class OMetaTableMixin:
             Optional[Table]: OM table object
         """
         return self._get(Table, f"{quote(model_str(fqn))}/tableProfile/latest")
+
+    def create_or_update_custom_metric(
+        self, custom_metric: CreateCustomMetricRequest, table_id: str
+    ) -> Table:
+        """Create or update custom metric. If custom metric name matches an existing
+        one then it will be updated.
+
+        Args:
+            custom_metric (CreateCustomMetricRequest): custom metric to be create or updated
+        """
+        resp = self.client.put(
+            f"{self.get_suffix(Table)}/{table_id}/customMetric",
+            data=custom_metric.json(),
+        )
+        return Table(**resp)

metadata/ingestion/ometa/mixins/tests_mixin.py

@@ -23,6 +23,9 @@ from metadata.generated.schema.api.tests.createLogicalTestCases import (
     CreateLogicalTestCases,
 )
 from metadata.generated.schema.api.tests.createTestCase import CreateTestCaseRequest
+from metadata.generated.schema.api.tests.createTestCaseResolutionStatus import (
+    CreateTestCaseResolutionStatus,
+)
 from metadata.generated.schema.api.tests.createTestDefinition import (
     CreateTestDefinitionRequest,
 )
@@ -30,6 +33,9 @@ from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteR
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.tests.basic import TestCaseResult
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
+from metadata.generated.schema.tests.testCaseResolutionStatus import (
+    TestCaseResolutionStatus,
+)
 from metadata.generated.schema.tests.testDefinition import (
     EntityType,
     TestCaseParameterDefinition,
@@ -242,10 +248,9 @@ class OMetaTestsMixin:
             end_ts (int): timestamp
         """
 
-        # timestamp should be changed to milliseconds in https://github.com/open-metadata/OpenMetadata/issues/8930
         params = {
-            "startTs": start_ts // 1000,
-            "endTs": end_ts // 1000,
+            "startTs": start_ts,
+            "endTs": end_ts,
         }
 
         resp = self.client.get(
@@ -302,3 +307,19 @@ class OMetaTestsMixin:
         """
         path = self.get_suffix(TestCase) + "/logicalTestCases"
         self.client.put(path, data=data.json(encoder=show_secrets_encoder))
+
+    def create_test_case_resolution(
+        self, data: CreateTestCaseResolutionStatus
+    ) -> TestCaseResolutionStatus:
+        """Create a test case resolution
+
+        Args:
+            data (CreateTestCaseResolutionStatus): test case resolution
+
+        Returns:
+            TestCaseResolutionStatus
+        """
+        path = self.get_suffix(TestCase) + "/testCaseIncidentStatus"
+        response = self.client.post(path, data=data.json(encoder=show_secrets_encoder))
+
+        return TestCaseResolutionStatus(**response)