Allow recreating entities for PurView registry

feathr_project/feathr/registry/_feature_registry_purview.py

@@ -3,6 +3,7 @@
 import inspect
 import itertools
 import os
+import re
 import sys
 import ast
 import types
@@ -44,6 +45,25 @@
 
 from feathr.constants import *
 
+def _to_snake(d, level: int = 0):
+    """
+    Convert `string`, `list[string]`, or all keys in a `dict` into snake case
+    The maximum length of input string or list is 100, or it will be truncated before being processed, for dict, the exception will be thrown if it has more than 100 keys.
+    the maximum nested level is 10, otherwise the exception will be thrown
+    """
+    if level >= 10:
+        raise ValueError("Too many nested levels")
+    if isinstance(d, str):
+        d = d[:100]
+        return re.sub(r'(?<!^)(?=[A-Z])', '_', d).lower()
+    if isinstance(d, list):
+        d = d[:100]
+        return [_to_snake(i, level + 1) if isinstance(i, (dict, list)) else i for i in d]
+    if len(d) > 100:
+        raise ValueError("Dict has too many keys")
+    return {_to_snake(a, level + 1): _to_snake(b, level + 1) if isinstance(b, (dict, list)) else b for a, b in d.items()}
+
+
 class _PurviewRegistry(FeathrRegistry):
     """
     Initializes the feature registry, doing the following:
@@ -720,6 +740,32 @@ def upload_single_entity_to_purview(self,entity:Union[AtlasEntity,AtlasProcess])
         The entity itself will also be modified, fill the GUID with real GUID in Purview.
         In order to avoid having concurrency issue, and provide clear guidance, this method only allows entity uploading once at a time.
         '''
+        try:
+            """
+            Try to find existing entity/process first, if found, return the existing entity's GUID
+            """
+            id = self.get_entity_id(entity.qualifiedName)
+            response =  self.purview_client.get_entity(id)['entities'][0]
+            j = entity.to_json()
+            if j["typeName"] == response["typeName"]:
+                if j["typeName"] == "Process":
+                    if response["attributes"]["qualifiedName"] != j["attributes"]["qualifiedName"]:
+                        raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+                else:
+                    if "type" in response['attributes'] and response["typeName"] in (TYPEDEF_ANCHOR_FEATURE, TYPEDEF_DERIVED_FEATURE):
+                        conf = ConfigFactory.parse_string(response['attributes']['type'])
+                        response['attributes']['type'] = dict(conf)
+                    keys = set([_to_snake(key) for key in j["attributes"].keys()]) - set(["qualified_name"])
+                    keys.add("qualifiedName")
+                    for k in keys:
+                        if response["attributes"][k] != j["attributes"][k]:
+                            raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+                return response["guid"]
+            else:
+                raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+        except AtlasException as e:
+            pass
+
         try:
             entity.lastModifiedTS="0"
             result = self.purview_client.upload_entities([entity])

registry/purview-registry/registry/purview_registry.py

@@ -5,6 +5,9 @@
 from urllib.error import HTTPError
 from uuid import UUID
 
+from registry.models import to_snake
+from pyapacheatlas.core.util import AtlasException
+
 from azure.identity import DefaultAzureCredential
 from loguru import logger
 from pyapacheatlas.auth.azcredential import AzCredentialWrapper
@@ -20,6 +23,9 @@
 Label_BelongsTo = "BELONGSTO"
 Label_Consumes = "CONSUMES"
 Label_Produces = "PRODUCES"
+TYPEDEF_DERIVED_FEATURE="feathr_derived_feature_v1"
+TYPEDEF_ANCHOR_FEATURE="feathr_anchor_feature_v1"
+
 TYPEDEF_ARRAY_ANCHOR=f"array<feathr_anchor_v1>"
 TYPEDEF_ARRAY_DERIVED_FEATURE=f"array<feathr_derived_feature_v1>"
 TYPEDEF_ARRAY_ANCHOR_FEATURE=f"array<feathr_anchor_feature_v1>"
@@ -568,17 +574,47 @@ def _register_feathr_feature_types(self):
     def _upload_entity_batch(self, entity_batch:list[AtlasEntity]):
         # we only support entity creation, update is not supported. 
         # setting lastModifiedTS ==0 will ensure this, if another entity with ts>=1 exist
-        # upload funtion will fail with 412 Precondition fail.
+        # upload function will fail with 412 Precondition fail.
         for entity in entity_batch:
-            entity.lastModifiedTS="0"
-            results = self.purview_client.upload_entities(
-                batch=entity)
-            if results:
-                dict = {x.guid: x for x in entity_batch}
-                for k, v in results['guidAssignments'].items():
-                    dict[k].guid = v
+            self._upload_single_entity(entity)
+
+    def _upload_single_entity(self, entity:AtlasEntity):
+        try:
+            """
+            Try to find existing entity/process first, if found, return the existing entity's GUID
+            """
+            id = self.get_entity_id(entity.qualifiedName)
+            response =  self.purview_client.get_entity(id)['entities'][0]
+            j = entity.to_json()
+            if j["typeName"] == response["typeName"]:
+                if j["typeName"] == "Process":
+                    if response["attributes"]["qualifiedName"] != j["attributes"]["qualifiedName"]:
+                        raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+                else:
+                    if "type" in response['attributes'] and response["typeName"] in (TYPEDEF_ANCHOR_FEATURE, TYPEDEF_DERIVED_FEATURE):
+                        conf = ConfigFactory.parse_string(response['attributes']['type'])
+                        response['attributes']['type'] = dict(conf)
+                    keys = set([to_snake(key) for key in j["attributes"].keys()]) - set(["qualified_name"])
+                    keys.add("qualifiedName")
+                    for k in keys:
+                        if response["attributes"][k] != j["attributes"][k]:
+                            raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+                entity.guid = response["guid"]
+                return
             else:
-                raise RuntimeError("Feature registration failed.", results)            
+                raise RuntimeError("The requested entity %s conflicts with the existing entity in PurView" % j["attributes"]["qualifiedName"])
+        except AtlasException as e:
+            pass
+
+        entity.lastModifiedTS="0"
+        results = self.purview_client.upload_entities(
+            batch=entity)
+        if results:
+            d = {x.guid: x for x in [entity]}
+            for k, v in results['guidAssignments'].items():
+                d[k].guid = v
+        else:
+            raise RuntimeError("Feature registration failed.", results)
 
     def _generate_fully_qualified_name(self, segments):
         return self.registry_delimiter.join(segments)