Enhance purview registry error messages (#709)

blrchen · web-flow · commit 059f2b4b6311 · 2022-09-28T22:26:21.000+08:00
* Enhance purview registry error messages

* Update doc for REGISTRY_DEBUGGING
diff --git a/docs/dev_guide/build-and-push-feathr-registry-docker-image.md b/docs/dev_guide/build-and-push-feathr-registry-docker-image.md
@@ -76,4 +76,8 @@ docker push feathrfeaturestore/feathr-registry
 
 ## Published Feathr Registry Image
 
-The published feathr feature registry is located in [DockerHub here](https://hub.docker.com/r/feathrfeaturestore/feathr-registry).
+The published feathr feature registry is located in [DockerHub here](https://hub.docker.com/r/feathrfeaturestore/feathr-registry).
+
+## Include the detailed track back info in registry api HTTP error response
+
+Set environment REGISTRY_DEBUGGING to any non empty string will enable the detailed track back info in registry api http response. This variable is helpful for python client debugging and should only be used for debugging purposes.
diff --git a/registry/purview-registry/main.py b/registry/purview-registry/main.py
@@ -1,11 +1,12 @@
 import os
+import traceback
 from re import sub
 from typing import Optional
 from uuid import UUID
 from fastapi import APIRouter, FastAPI, HTTPException
+from fastapi.responses import JSONResponse
 from starlette.middleware.cors import CORSMiddleware
-from registry import *
-from registry.purview_registry import PurviewRegistry
+from registry.purview_registry import PurviewRegistry, ConflictError
 from registry.models import AnchorDef, AnchorFeatureDef, DerivedFeatureDef, EntityType, ProjectDef, SourceDef, to_snake
 
 rp = "/v1"
@@ -43,6 +44,48 @@ def to_camel(s):
                    allow_headers=["*"],
                    )
 
+def exc_to_content(e: Exception) -> dict:
+    content={"message": str(e)}
+    if os.environ.get("REGISTRY_DEBUGGING"):
+        content["traceback"] = "".join(traceback.TracebackException.from_exception(e).format())
+    return content
+
+@app.exception_handler(ConflictError)
+async def conflict_error_handler(_, exc: ConflictError):
+    return JSONResponse(
+        status_code=409,
+        content=exc_to_content(exc),
+    )
+
+
+@app.exception_handler(ValueError)
+async def value_error_handler(_, exc: ValueError):
+    return JSONResponse(
+        status_code=400,
+        content=exc_to_content(exc),
+    )
+
+@app.exception_handler(TypeError)
+async def type_error_handler(_, exc: ValueError):
+    return JSONResponse(
+        status_code=400,
+        content=exc_to_content(exc),
+    )
+
+
+@app.exception_handler(KeyError)
+async def key_error_handler(_, exc: KeyError):
+    return JSONResponse(
+        status_code=404,
+        content=exc_to_content(exc),
+    )
+
+@app.exception_handler(IndexError)
+async def index_error_handler(_, exc: IndexError):
+    return JSONResponse(
+        status_code=404,
+        content=exc_to_content(exc),
+    )
 
 @router.get("/projects",tags=["Project"])
 def get_projects() -> list[str]:
diff --git a/registry/purview-registry/registry/purview_registry.py b/registry/purview-registry/registry/purview_registry.py
@@ -1,8 +1,6 @@
 import copy
-from http.client import CONFLICT, HTTPException
 import itertools
-from typing import Any, Optional, Tuple, Union
-from urllib.error import HTTPError
+from typing import Optional, Tuple, Union
 from uuid import UUID
 
 from azure.identity import DefaultAzureCredential
@@ -11,7 +9,7 @@
 from pyapacheatlas.core import (AtlasEntity, AtlasProcess,
                                 PurviewClient)
 from pyapacheatlas.core.typedef import (AtlasAttributeDef,Cardinality,EntityTypeDef)
-from pyapacheatlas.core.util import GuidTracker
+from pyapacheatlas.core.util import GuidTracker, AtlasException
 from pyhocon import ConfigFactory
 
 from registry.interface import Registry
@@ -23,6 +21,10 @@
 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>"
+
+class ConflictError(Exception):
+    pass
+
 class PurviewRegistry(Registry):
     def __init__(self,azure_purview_name: str, registry_delimiter: str = "__", credential=None,register_types = True):
         self.registry_delimiter = registry_delimiter
@@ -568,18 +570,22 @@ 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
-            else:
-                raise RuntimeError("Feature registration failed.", results)            
-            
+            try:
+                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
+                else:
+                    raise RuntimeError("Feature registration failed.", results)
+            except AtlasException as e:
+                if "PreConditionCheckFailed" in e.args[0]:
+                    raise ConflictError(f"Entity {entity.guid}, {entity.typeName} -- {entity.qualifiedName} already exists in Purview. Please use a new name.")
+
     def _generate_fully_qualified_name(self, segments):
         return self.registry_delimiter.join(segments)
 
