Added RPC invocation for PreUpdate

dao-api/build.gradle

@@ -10,6 +10,7 @@ dependencies {
   compile externalDependency.reflections
   compile externalDependency.commonsLang
   implementation 'com.google.protobuf:protobuf-java:3.21.1'
+  implementation 'io.grpc:grpc-stub:1.57.0'
   dataModel project(':core-models')
   dataModel project(':validators')
 

dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java

@@ -1,6 +1,7 @@
 package com.linkedin.metadata.dao;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.Descriptors;
 import com.google.protobuf.Message;
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.Urn;
@@ -47,6 +48,9 @@
 import com.linkedin.metadata.query.IndexFilter;
 import com.linkedin.metadata.query.IndexGroupByCriterion;
 import com.linkedin.metadata.query.IndexSortCriterion;
+import io.grpc.stub.AbstractStub;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.sql.Timestamp;
 import java.time.Clock;
 import java.util.Collections;
@@ -215,6 +219,8 @@ public static class AspectUpdateLambda<ASPECT extends RecordTemplate> {
 
   private Clock _clock = Clock.systemUTC();
 
+  private Map<Class<? extends RecordTemplate>, Object[]> _grpcPreUpdateRoutingMap;
+
   /**
    * Constructor for BaseLocalDAO.
    *
@@ -414,6 +420,9 @@ public RestliPreUpdateAspectRegistry getRestliPreUpdateAspectRegistry() {
     return _restliPreUpdateAspectRegistry;
   }
 
+  public void setGrpcPreUpdateRoutingMap(Map<Class<? extends RecordTemplate>, Object[]> grpcPreUpdateRoutingMap) {
+    _grpcPreUpdateRoutingMap = grpcPreUpdateRoutingMap;
+  }
 
   /**
    * Enables or disables atomic updates of multiple aspects.
@@ -854,6 +863,9 @@ public <ASPECT extends RecordTemplate> ASPECT add(@Nonnull URN urn, @Nonnull ASP
       @Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext,
       @Nullable IngestionParams ingestionParams) {
     ASPECT updatedAspect = preUpdateRouting(urn, newValue);
+    if (_grpcPreUpdateRoutingMap != null && _grpcPreUpdateRoutingMap.containsKey(newValue.getClass())) {
+      updatedAspect = grpcPreUpdateRouting(urn, updatedAspect);
+    }
     return rawAdd(urn, updatedAspect, auditStamp, trackingContext, ingestionParams);
   }
 
@@ -1675,4 +1687,40 @@ protected <ASPECT extends RecordTemplate> ASPECT preUpdateRouting(URN urn, ASPEC
     }
     return newValue;
   }
-}
+
+  protected <ASPECT extends RecordTemplate> ASPECT grpcPreUpdateRouting(URN urn, ASPECT aspect) {
+    Object[] routingData = _grpcPreUpdateRoutingMap.get(aspect.getClass());
+    // Extract stored information
+    AbstractStub<?> stub = (AbstractStub<?>) routingData[0];
+    Method newBuilderMethod = (Method) routingData[1];  // Method for newBuilder()
+    Method preUpdateMethod = (Method) routingData[2];   // Method for gRPC call
+
+    // Dynamically create the request using newBuilder method
+    try {
+      Message.Builder requestBuilder = (Message.Builder) newBuilderMethod.invoke(null);
+
+      // Get the descriptor for the request type
+      Descriptors.Descriptor requestDescriptor = requestBuilder.getDescriptorForType();
+
+      // Find the "urn" field descriptor and set its value
+      Descriptors.FieldDescriptor urnField = requestDescriptor.findFieldByName("urn");
+      requestBuilder.setField(urnField, urn);
+
+      // Find the "value" field descriptor and set its value (this is the aspect in your case)
+      Descriptors.FieldDescriptor valueField = requestDescriptor.findFieldByName("value");
+      requestBuilder.setField(valueField, aspect);
+
+      // Build the request object
+      Message request = requestBuilder.build();
+
+      // Invoke the gRPC preUpdate method
+      Object response = preUpdateMethod.invoke(stub, request);
+      // Extract and return the updated aspect from the response
+      Descriptors.Descriptor responseDescriptor = ((Message) response).getDescriptorForType();
+      log.info("PreUpdateRouting completed in BaseLocalDao, urn: {}, updated aspect: {}", urn, response);
+      return (ASPECT) ((Message) response).getField(responseDescriptor.findFieldByName("value"));
+    } catch (IllegalAccessException | InvocationTargetException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}