[feat] add with optimistic locking in BaseLocalDAO (#543)

* [feat] add with optimistic locking in BaseLocalDAO

* fix

* update unit test

Author: ybz1013

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

@@ -76,6 +76,8 @@
 import lombok.NonNull;
 import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
+import pegasus.com.linkedin.metadata.events.IngestionAspectETag;
+import pegasus.com.linkedin.metadata.events.IngestionAspectETagArray;
 
 import static com.linkedin.metadata.dao.utils.IngestionUtils.*;
 import static com.linkedin.metadata.dao.utils.ModelUtils.*;
@@ -484,7 +486,6 @@ public void setEmitAuditEvent(boolean emitAuditEvent) {
     _emitAuditEvent = emitAuditEvent;
   }
 
-
   /**
    * Logic common to both {@link #add(Urn, Class, Function, AuditStamp)} and {@link #delete(Urn, Class, AuditStamp, int)} methods.
    *
@@ -549,17 +550,44 @@ private <ASPECT extends RecordTemplate> AddResult<ASPECT> addCommon(@Nonnull URN
       return new AddResult<>(oldValue, oldValue, aspectClass);
     }
 
+    final AuditStamp optimisticLockAuditStamp = extractOptimisticLockForAspectFromIngestionParamsIfPossible(ingestionParams, aspectClass);
+
     // Save the newValue as the latest version
     long largestVersion =
-        saveLatest(urn, aspectClass, oldValue, oldAuditStamp, newValue, auditStamp, latest.isSoftDeleted,
-            trackingContext, ingestionParams.isTestMode());
+        saveLatest(urn, aspectClass, oldValue,
+            optimisticLockAuditStamp != null ? optimisticLockAuditStamp : oldAuditStamp,
+            newValue, auditStamp, latest.isSoftDeleted, trackingContext, ingestionParams.isTestMode());
 
     // Apply retention policy
     applyRetention(urn, aspectClass, getRetention(aspectClass), largestVersion);
 
     return new AddResult<>(oldValue, newValue, aspectClass);
   }
 
+  @VisibleForTesting
+  protected  <ASPECT extends RecordTemplate> AuditStamp extractOptimisticLockForAspectFromIngestionParamsIfPossible(
+      @Nullable IngestionParams ingestionParams, @Nonnull Class<ASPECT> aspectClass) {
+    if (ingestionParams == null) {
+      return null;
+    }
+
+    AuditStamp optimisticLockAuditStamp = null;
+
+    final IngestionAspectETagArray ingestionAspectETags = ingestionParams.getIngestionETags();
+
+    if (ingestionAspectETags != null) {
+      for (IngestionAspectETag ingestionAspectETag: ingestionAspectETags) {
+        if (aspectClass.getSimpleName().equalsIgnoreCase(ingestionAspectETag.getAspect_name())
+            && ingestionAspectETag.getETag() != null) {
+          optimisticLockAuditStamp = new AuditStamp();
+          optimisticLockAuditStamp.setTime(ingestionAspectETag.getETag());
+          break;
+        }
+      }
+    }
+    return optimisticLockAuditStamp;
+  }
+
   /**
    * Adds a new version of several aspects for an entity.
    *
@@ -1313,15 +1341,15 @@ private <ASPECT extends RecordTemplate> void applyRetention(@Nonnull URN urn, @N
    * @param urn the URN for the entity the aspect is attached to
    * @param aspectClass the aspectClass of the aspect being saved
    * @param oldEntry {@link RecordTemplate} of the previous latest value of aspect, null if new value is the first version
-   * @param oldAuditStamp the audit stamp of the previous latest aspect, null if new value is the first version
+   * @param optimisticLockAuditStamp the audit stamp of the previous latest aspect, null if new value is the first version. Used for optimistic locking.
    * @param newEntry {@link RecordTemplate} of the new latest value of aspect
    * @param newAuditStamp the audit stamp for the operation
    * @param isSoftDeleted flag to indicate if the previous latest value of aspect was soft deleted
    * @param isTestMode whether the test mode is enabled or not
    * @return the largest version
    */
   protected abstract <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn,
-      @Nonnull Class<ASPECT> aspectClass, @Nullable ASPECT oldEntry, @Nullable AuditStamp oldAuditStamp,
+      @Nonnull Class<ASPECT> aspectClass, @Nullable ASPECT oldEntry, @Nullable AuditStamp optimisticLockAuditStamp,
       @Nullable ASPECT newEntry, @Nonnull AuditStamp newAuditStamp, boolean isSoftDeleted,
       @Nullable IngestionTrackingContext trackingContext, boolean isTestMode);
 

dao-api/src/main/pegasus/pegasus/com/linkedin/metadata/events/IngestionAspectETag.pdl

@@ -6,7 +6,7 @@ namespace pegasus.com.linkedin.metadata.events
 record IngestionAspectETag {
 
   /**
-   * aspect FQCN
+   * aspect field name, e.g. "status"
    */
   aspect_name: optional string = ""
 

dao-api/src/test/java/com/linkedin/metadata/dao/BaseLocalDAOTest.java

@@ -52,6 +52,8 @@
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
+import pegasus.com.linkedin.metadata.events.IngestionAspectETag;
+import pegasus.com.linkedin.metadata.events.IngestionAspectETagArray;
 
 import static com.linkedin.common.AuditStamps.*;
 import static org.mockito.Mockito.*;
@@ -87,7 +89,7 @@ public DummyLocalDAO(Class<ENTITY_ASPECT_UNION> aspectClass, BiFunction<FooUrn,
 
     @Override
     protected <ASPECT extends RecordTemplate> long saveLatest(FooUrn urn, Class<ASPECT> aspectClass, ASPECT oldEntry,
-        AuditStamp oldAuditStamp, ASPECT newEntry, AuditStamp newAuditStamp, boolean isSoftDeleted,
+        AuditStamp optimisticLockAuditStamp, ASPECT newEntry, AuditStamp newAuditStamp, boolean isSoftDeleted,
         @Nullable IngestionTrackingContext trackingContext, boolean isTestMode) {
       return 0;
     }
@@ -802,4 +804,39 @@ public void testAspectCallbackHelperWithUnregisteredAspect() throws URISyntaxExc
     // Verify that the result is the same as the input aspect since it's not registered
     assertEquals(result.getUpdatedAspect(), foo);
   }
+
+  @Test
+  public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossible() {
+    IngestionAspectETag ingestionAspectETag = new IngestionAspectETag();
+    ingestionAspectETag.setAspect_name("aspectFoo");
+    ingestionAspectETag.setETag(1234L);
+
+    IngestionParams ingestionParams = new IngestionParams();
+    ingestionParams.setIngestionETags(new IngestionAspectETagArray(ingestionAspectETag));
+
+    AuditStamp result = _dummyLocalDAO.extractOptimisticLockForAspectFromIngestionParamsIfPossible(ingestionParams, AspectFoo.class);
+
+    assertEquals(result.getTime(), Long.valueOf(1234L));
+  }
+
+  @Test
+  public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossibleIngestionParamsIsNull() {
+    AuditStamp result = _dummyLocalDAO.extractOptimisticLockForAspectFromIngestionParamsIfPossible(null, AspectFoo.class);
+
+    assertNull(result);
+  }
+
+  @Test
+  public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossibleAspectNameDoesntMatch() {
+    IngestionAspectETag ingestionAspectETag = new IngestionAspectETag();
+    ingestionAspectETag.setAspect_name("aspectBar");
+    ingestionAspectETag.setETag(1234L);
+
+    IngestionParams ingestionParams = new IngestionParams();
+    ingestionParams.setIngestionETags(new IngestionAspectETagArray(ingestionAspectETag));
+
+    AuditStamp result = _dummyLocalDAO.extractOptimisticLockForAspectFromIngestionParamsIfPossible(ingestionParams, AspectFoo.class);
+
+    assertNull(result);
+  }
 }

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java

@@ -598,12 +598,12 @@ protected <T> T runInTransactionWithRetry(@Nonnull Supplier<T> block, int maxTra
 
   @Override
   protected <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn, @Nonnull Class<ASPECT> aspectClass,
-      @Nullable ASPECT oldValue, @Nullable AuditStamp oldAuditStamp, @Nullable ASPECT newValue,
+      @Nullable ASPECT oldValue, @Nullable AuditStamp optimisticLockAuditStamp, @Nullable ASPECT newValue,
       @Nonnull AuditStamp newAuditStamp, boolean isSoftDeleted, @Nullable IngestionTrackingContext trackingContext,
       boolean isTestMode) {
     // Save oldValue as the largest version + 1
     long largestVersion = 0;
-    if ((isSoftDeleted || oldValue != null) && oldAuditStamp != null && _changeLogEnabled) {
+    if ((isSoftDeleted || oldValue != null) && optimisticLockAuditStamp != null && _changeLogEnabled) {
       // When saving on entity which has history version (including being soft deleted), and changeLog is enabled,
       // the saveLatest will process the following steps:
 
@@ -620,11 +620,11 @@ protected <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn, @Non
       }
       // Move latest version to historical version by insert a new record only if we are not overwriting the latest version.
       if (!_overwriteLatestVersionEnabled) {
-        insert(urn, oldValue, aspectClass, oldAuditStamp, largestVersion, trackingContext, isTestMode);
+        insert(urn, oldValue, aspectClass, optimisticLockAuditStamp, largestVersion, trackingContext, isTestMode);
       }
       // update latest version
       updateWithOptimisticLocking(urn, newValue, aspectClass, newAuditStamp, LATEST_VERSION,
-          new Timestamp(oldAuditStamp.getTime()), trackingContext, isTestMode);
+          new Timestamp(optimisticLockAuditStamp.getTime()), trackingContext, isTestMode);
     } else {
       // When for fresh ingestion or with changeLog disabled
       // TODO(yanyang) added for job-gms duplicity debug, throwaway afterwards