feat(gql): support operationName (datahub-project#10210)

sleeperdeep · Jun 25, 2024 · 72b1f99 · 72b1f99
1 parent 689e3b8
commit 72b1f99
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java
@@ -87,6 +87,7 @@ private GraphQLEngine(
 
   public ExecutionResult execute(
       @Nonnull final String query,
+      @Nullable final String operationName,
       @Nullable final Map<String, Object> variables,
       @Nonnull final QueryContext context) {
     /*
@@ -100,6 +101,7 @@ public ExecutionResult execute(
     ExecutionInput executionInput =
         ExecutionInput.newExecutionInput()
             .query(query)
+            .operationName(operationName)
             .variables(variables)
             .dataLoaderRegistry(register)
             .context(context)

diff --git a/metadata-ingestion/src/datahub/ingestion/graph/client.py b/metadata-ingestion/src/datahub/ingestion/graph/client.py
@@ -859,15 +859,21 @@ def get_aspect_counts(self, aspect: str, urn_like: Optional[str] = None) -> int:
         results = self._post_generic(self._aspect_count_endpoint, args)
         return results["value"]
 
-    def execute_graphql(self, query: str, variables: Optional[Dict] = None) -> Dict:
+    def execute_graphql(
+        self,
+        query: str,
+        variables: Optional[Dict] = None,
+        operation_name: Optional[str] = None,
+    ) -> Dict:
         url = f"{self.config.server}/api/graphql"
 
         body: Dict = {
             "query": query,
         }
-
         if variables:
             body["variables"] = variables
+        if operation_name:
+            body["operationName"] = operation_name
 
         logger.debug(
             f"Executing graphql query: {query} with variables: {json.dumps(variables)}"

diff --git a/...ata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java b/...ata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java
@@ -86,6 +86,15 @@ CompletableFuture<ResponseEntity<String>> postGraphQL(HttpEntity<String> httpEnt
       return CompletableFuture.completedFuture(new ResponseEntity<>(HttpStatus.BAD_REQUEST));
     }
 
+    /*
+     * Extract "operationName" field
+     */
+    JsonNode operationNameJson = bodyJson.get("operationName");
+    final String operationName =
+        (operationNameJson != null && !operationNameJson.isNull())
+            ? operationNameJson.asText()
+            : null;
+
     /*
      * Extract "variables" map
      */
@@ -118,7 +127,8 @@ CompletableFuture<ResponseEntity<String>> postGraphQL(HttpEntity<String> httpEnt
           /*
            * Execute GraphQL Query
            */
-          ExecutionResult executionResult = _engine.execute(queryJson.asText(), variables, context);
+          ExecutionResult executionResult =
+              _engine.execute(queryJson.asText(), operationName, variables, context);
 
           if (executionResult.getErrors().size() != 0) {
             // There were GraphQL errors. Report in error logs.