nifi add data provenance

noharm-ai · Jan 30, 2025 · ff8e42f · ff8e42f
1 parent aa67fd7
commit ff8e42f
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 14 deletions.
diff --git a/models/enums.py b/models/enums.py
@@ -127,6 +127,7 @@ class NifiQueueActionTypeEnum(Enum):
     CUSTOM_CALLBACK = "CUSTOM_CALLBACK"
     REFRESH_TEMPLATE = "REFRESH_TEMPLATE"
     UPDATE_PROPERTY = "UPDATE_PROPERTY"
+    VIEW_PROVENANCE = "VIEW_PROVENANCE"
 
 
 class DrugAlertTypeEnum(Enum):

diff --git a/services/admin/admin_integration_remote_service.py b/services/admin/admin_integration_remote_service.py
@@ -292,6 +292,21 @@ def _get_new_queue(id_processor: str, action_type: str, data: dict):
                     "config": {"properties": data["properties"]},
                 }
             }
+    elif NifiQueueActionTypeEnum.VIEW_PROVENANCE.value == action_type:
+        queue.url = f"nifi-api/provenance"
+        queue.method = "POST"
+        queue.body = {
+            "provenance": {
+                "request": {
+                    "maxResults": 100,
+                    "summarize": True,
+                    "incrementalResults": False,
+                    "searchTerms": {
+                        "ProcessorID": {"value": escape(id_processor), "inverse": False}
+                    },
+                }
+            }
+        }
 
     return queue
 
@@ -370,24 +385,26 @@ def get_errors(user_context: User):
 
 
 def _validate_custom_endpoint(endpoint: str):
-    pattern1 = re.compile(
-        "^nifi-api\/flowfile-queues\/[\w-]{36}\/flowfiles\/[\w-]{36}\/content$"
+    patterns = []
+    patterns.append(
+        re.compile(
+            "^nifi-api\/flowfile-queues\/[\w-]{36}\/flowfiles\/[\w-]{36}\/content$"
+        )
     )
-    pattern2 = re.compile(
-        "^nifi-api\/flowfile-queues\/[\w-]{36}\/listing-requests\/[\w-]{36}$"
+    patterns.append(
+        re.compile(
+            "^nifi-api\/flowfile-queues\/[\w-]{36}\/listing-requests\/[\w-]{36}$"
+        )
     )
-    pattern3 = re.compile(
-        "^nifi-api\/flowfile-queues\/[\w-]{36}\/flowfiles\/[\w-]{36}$"
+    patterns.append(
+        re.compile("^nifi-api\/flowfile-queues\/[\w-]{36}\/flowfiles\/[\w-]{36}$")
     )
+    patterns.append(re.compile("^nifi-api\/provenance-events\/\d*$"))
+    patterns.append(re.compile("^nifi-api\/provenance-events\/\d*\/content\/output$"))
 
-    if pattern1.match(endpoint):
-        return True
-
-    if pattern2.match(endpoint):
-        return True
-
-    if pattern3.match(endpoint):
-        return True
+    for p in patterns:
+        if p.match(endpoint):
+            return True
 
     raise ValidationError(
         "Endpoint custom inválido",