feat(metadata-io): enable rollback transaction support (#12509)

Author: david-leifker

docs/how/updating-datahub.md

@@ -31,7 +31,11 @@ This file documents any backwards-incompatible changes in DataHub and assists pe
 ## 0.15.0
 
 - OpenAPI Update: PIT Keep Alive parameter added to scroll endpoints. NOTE: This parameter requires the `pointInTimeCreationEnabled` feature flag to be enabled and the `elasticSearch.implementation` configuration to be `elasticsearch`. This feature is not supported for OpenSearch at this time and the parameter will not be respected without both of these set.
-- OpenAPI Update 2: Previously there was an incorrectly marked parameter named `sort` on the generic list entities endpoint for v3. This parameter is deprecated and only supports a single string value while the documentation indicates it supports a list of strings. This documentation error has been fixed and the correct field, `sortCriteria`, is now documented which supports a list of strings. 
+- OpenAPI Update 2: Previously there was an incorrectly marked parameter named `sort` on the generic list entities endpoint for v3. This parameter is deprecated and only supports a single string value while the documentation indicates it supports a list of strings. This documentation error has been fixed and the correct field, `sortCriteria`, is now documented which supports a list of strings.
+
+### Known Issues
+
+- Persistence Exception: No Rows Updated may occur if a transaction does not change any aspect's data.
 
 ### Breaking Changes
 

docs/managed-datahub/release-notes/v_0_3_7.md

@@ -13,6 +13,8 @@ If you are using an older CLI/SDK version, then please upgrade it. This applies
 
 ## Known Issues
 
+- Persistence Exception: No Rows Updated may occur if a transaction does not change any aspect's data.
+
 ### v0.3.7.8
    * Notes Feature
       * Adding a Note to an entity will result in that note showing up in the Settings > Home Page list of announcements as well as the profile page of the entity.

docs/managed-datahub/release-notes/v_0_3_8.md

@@ -12,6 +12,7 @@ Recommended CLI/SDK
 ## Known Issues
 
 - Async APIs - DataHub's asynchronous APIs perform only basic schema validation when receiving MCP requests, similar to direct production to MCP Kafka topics. While requests must conform to the MCP schema to be accepted, actual processing happens later in the pipeline. Any processing failures that occur after the initial acceptance are captured in the Failed MCP topic, but these failures are not immediately surfaced to the API caller since they happen asynchronously.
+- Persistence Exception: No Rows Updated may occur if a transaction does not change any aspect's data.
 
 ## Release Changelog
 ---

metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java

@@ -10,6 +10,7 @@
 import java.sql.Timestamp;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Stream;
@@ -169,15 +170,16 @@ default Map<String, Long> getNextVersions(
   void setWritable(boolean canWrite);
 
   @Nonnull
-  <T> T runInTransactionWithRetry(
-      @Nonnull final Function<TransactionContext, T> block, final int maxTransactionRetry);
+  <T> Optional<T> runInTransactionWithRetry(
+      @Nonnull final Function<TransactionContext, TransactionResult<T>> block,
+      final int maxTransactionRetry);
 
   @Nonnull
-  default <T> List<T> runInTransactionWithRetry(
-      @Nonnull final Function<TransactionContext, T> block,
+  default <T> Optional<T> runInTransactionWithRetry(
+      @Nonnull final Function<TransactionContext, TransactionResult<T>> block,
       AspectsBatch batch,
       final int maxTransactionRetry) {
-    return List.of(runInTransactionWithRetry(block, maxTransactionRetry));
+    return runInTransactionWithRetry(block, maxTransactionRetry);
   }
 
   default void incrementWriteMetrics(String aspectName, long count, long bytes) {

metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java

@@ -72,4 +72,10 @@ public void flush() {
       tx.flush();
     }
   }
+
+  public void rollback() {
+    if (tx != null) {
+      tx.rollback();
+    }
+  }
 }

metadata-io/src/main/java/com/linkedin/metadata/entity/TransactionContext.java

@@ -30,6 +30,7 @@
 import com.linkedin.metadata.entity.EntityAspectIdentifier;
 import com.linkedin.metadata.entity.ListResult;
 import com.linkedin.metadata.entity.TransactionContext;
+import com.linkedin.metadata.entity.TransactionResult;
 import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
 import com.linkedin.metadata.entity.ebean.PartitionedStream;
 import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs;
@@ -44,6 +45,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -294,17 +296,17 @@ public ListResult<String> listAspectMetadata(
         aspectMetadatas, listResultMetadata, start, pageNumber, pageSize, totalCount);
   }
 
-  @Override
   @Nonnull
-  public <T> T runInTransactionWithRetry(
-      @Nonnull final Function<TransactionContext, T> block, final int maxTransactionRetry) {
+  @Override
+  public <T> Optional<T> runInTransactionWithRetry(
+      @Nonnull Function<TransactionContext, TransactionResult<T>> block, int maxTransactionRetry) {
     validateConnection();
     TransactionContext txContext = TransactionContext.empty(maxTransactionRetry);
     do {
       try {
         // TODO: Try to bend this code to make use of Cassandra batches. This method is called from
         // single-urn operations, so perf should not suffer much
-        return block.apply(txContext);
+        return block.apply(txContext).getResults();
       } catch (DriverException exception) {
         txContext.addException(exception);
       }

metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java

@@ -5,9 +5,6 @@
 import com.codahale.metrics.MetricRegistry;
 import com.datahub.util.exception.ModelConversionException;
 import com.datahub.util.exception.RetryLimitReached;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.aspect.RetrieverContext;
@@ -20,6 +17,7 @@
 import com.linkedin.metadata.entity.EntityAspectIdentifier;
 import com.linkedin.metadata.entity.ListResult;
 import com.linkedin.metadata.entity.TransactionContext;
+import com.linkedin.metadata.entity.TransactionResult;
 import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs;
 import com.linkedin.metadata.models.AspectSpec;
@@ -53,14 +51,10 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -89,12 +83,6 @@ public class EbeanAspectDao implements AspectDao, AspectMigrationsDao {
   // more testing.
   private int _queryKeysCount = 375; // 0 means no pagination on keys
 
-  /**
-   * Used to control write concurrency when an entity key aspect is present. If a batch contains an
-   * entity key aspect, only allow a single execution per URN
-   */
-  private final LoadingCache<String, Lock> locks;
-
   private final String batchGetMethod;
 
   public EbeanAspectDao(@Nonnull final Database server, EbeanConfiguration ebeanConfiguration) {
@@ -103,21 +91,6 @@ public EbeanAspectDao(@Nonnull final Database server, EbeanConfiguration ebeanCo
         ebeanConfiguration.getBatchGetMethod() != null
             ? ebeanConfiguration.getBatchGetMethod()
             : "IN";
-    if (ebeanConfiguration.getLocking().isEnabled()) {
-      this.locks =
-          CacheBuilder.newBuilder()
-              .maximumSize(ebeanConfiguration.getLocking().getMaximumLocks())
-              .expireAfterWrite(
-                  ebeanConfiguration.getLocking().getDurationSeconds(), TimeUnit.SECONDS)
-              .build(
-                  new CacheLoader<>() {
-                    public Lock load(String key) {
-                      return new ReentrantLock(true);
-                    }
-                  });
-    } else {
-      this.locks = null;
-    }
   }
 
   @Override
@@ -764,83 +737,42 @@ public ListResult<String> listLatestAspectMetadata(
 
   @Override
   @Nonnull
-  public <T> T runInTransactionWithRetry(
-      @Nonnull final Function<TransactionContext, T> block, final int maxTransactionRetry) {
-    return runInTransactionWithRetry(block, null, maxTransactionRetry).get(0);
+  public <T> Optional<T> runInTransactionWithRetry(
+      @Nonnull final Function<TransactionContext, TransactionResult<T>> block,
+      final int maxTransactionRetry) {
+    return runInTransactionWithRetry(block, null, maxTransactionRetry);
   }
 
   @Override
   @Nonnull
-  public <T> List<T> runInTransactionWithRetry(
-      @Nonnull final Function<TransactionContext, T> block,
+  public <T> Optional<T> runInTransactionWithRetry(
+      @Nonnull final Function<TransactionContext, TransactionResult<T>> block,
       @Nullable AspectsBatch batch,
       final int maxTransactionRetry) {
 
-    LinkedList<T> result = new LinkedList<>();
-
-    if (locks != null && batch != null) {
-      Set<Urn> urnsWithKeyAspects =
-          batch.getMCPItems().stream()
-              .filter(i -> i.getEntitySpec().getKeyAspectSpec().equals(i.getAspectSpec()))
-              .map(MCPItem::getUrn)
-              .collect(Collectors.toSet());
-
-      if (!urnsWithKeyAspects.isEmpty()) {
-
-        // Split into batches by urn with key aspect, remaining aspects in the pair's second
-        Pair<List<AspectsBatch>, AspectsBatch> splitBatches =
-            splitByUrn(batch, urnsWithKeyAspects, batch.getRetrieverContext());
-
-        // Run non-key aspect `other` batch per normal
-        if (!splitBatches.getSecond().getItems().isEmpty()) {
-          result.add(
-              runInTransactionWithRetryUnlocked(
-                  block, splitBatches.getSecond(), maxTransactionRetry));
-        }
-
-        // For each key aspect batch
-        for (AspectsBatch splitBatch : splitBatches.getFirst()) {
-          try {
-            Lock lock =
-                locks.get(splitBatch.getMCPItems().stream().findFirst().get().getUrn().toString());
-            lock.lock();
-            try {
-              result.add(runInTransactionWithRetryUnlocked(block, splitBatch, maxTransactionRetry));
-            } finally {
-              lock.unlock();
-            }
-          } catch (ExecutionException e) {
-            throw new RuntimeException(e);
-          }
-        }
-      } else {
-        // No key aspects found, run per normal
-        result.add(runInTransactionWithRetryUnlocked(block, batch, maxTransactionRetry));
-      }
-    } else {
-      // locks disabled or null batch
-      result.add(runInTransactionWithRetryUnlocked(block, batch, maxTransactionRetry));
-    }
-
-    return result;
+    return runInTransactionWithRetryUnlocked(block, batch, maxTransactionRetry).getResults();
   }
 
   @Nonnull
-  public <T> T runInTransactionWithRetryUnlocked(
-      @Nonnull final Function<TransactionContext, T> block,
+  public <T> TransactionResult<T> runInTransactionWithRetryUnlocked(
+      @Nonnull final Function<TransactionContext, TransactionResult<T>> block,
       @Nullable AspectsBatch batch,
       final int maxTransactionRetry) {
 
     validateConnection();
     TransactionContext transactionContext = TransactionContext.empty(maxTransactionRetry);
 
-    T result = null;
+    TransactionResult<T> result = null;
     do {
       try (Transaction transaction =
           _server.beginTransaction(TxScope.requiresNew().setIsolation(TX_ISOLATION))) {
         transaction.setBatchMode(true);
         result = block.apply(transactionContext.tx(transaction));
-        transaction.commit();
+        if (result.isCommitOrRollback()) {
+          transaction.commit();
+        } else {
+          transaction.rollback();
+        }
         break;
       } catch (PersistenceException exception) {
         if (exception instanceof DuplicateKeyException) {

metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java

@@ -41,6 +41,7 @@
 import java.sql.Timestamp;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import org.mockito.Mockito;
 import org.testng.annotations.Test;
 
@@ -129,7 +130,7 @@ public void testDeleteUniqueRefGeneratesValidMCP() {
             1);
 
     Mockito.when(_aspectDao.runInTransactionWithRetry(Mockito.any(), Mockito.anyInt()))
-        .thenReturn(result);
+        .thenReturn(Optional.of(result));
 
     final DeleteReferencesResponse response =
         _deleteEntityService.deleteReferencesTo(opContext, container, false);

metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java

@@ -12,6 +12,7 @@
 import com.linkedin.metadata.aspect.batch.AspectsBatch;
 import com.linkedin.metadata.config.EbeanConfiguration;
 import com.linkedin.metadata.entity.EntityAspectIdentifier;
+import com.linkedin.metadata.entity.TransactionResult;
 import io.ebean.Database;
 import io.ebean.test.LoggedSql;
 import java.util.List;
@@ -38,7 +39,7 @@ public void testGetNextVersionForUpdate() {
         (txContext) -> {
           testDao.getNextVersions(
               Map.of("urn:li:corpuser:testGetNextVersionForUpdate", Set.of("status")));
-          return "";
+          return TransactionResult.commit("");
         },
         mock(AspectsBatch.class),
         0);
@@ -61,7 +62,7 @@ public void testGetLatestAspectsForUpdate() throws JsonProcessingException {
         (txContext) -> {
           testDao.getLatestAspects(
               Map.of("urn:li:corpuser:testGetLatestAspectsForUpdate", Set.of("status")), true);
-          return "";
+          return TransactionResult.commit("");
         },
         mock(AspectsBatch.class),
         0);
@@ -94,7 +95,7 @@ public void testbatchGetForUpdate() throws JsonProcessingException {
                       DATA_PLATFORM_INSTANCE_ASPECT_NAME,
                       ASPECT_LATEST_VERSION)),
               true);
-          return "";
+          return TransactionResult.commit("");
         },
         mock(AspectsBatch.class),
         0);

metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java

@@ -22,26 +22,7 @@ public class EbeanConfiguration {
   private long waitTimeoutMillis;
   private boolean autoCreateDdl;
   private boolean postgresUseIamAuth;
-  private LockingConfiguration locking;
   private String batchGetMethod;
 
-  public static final EbeanConfiguration testDefault =
-      EbeanConfiguration.builder().locking(LockingConfiguration.testDefault).build();
-
-  @Data
-  @Builder
-  @AllArgsConstructor
-  @NoArgsConstructor
-  public static class LockingConfiguration {
-    private boolean enabled;
-    private long durationSeconds;
-    private long maximumLocks;
-
-    public static final LockingConfiguration testDefault =
-        LockingConfiguration.builder()
-            .enabled(true)
-            .durationSeconds(60)
-            .maximumLocks(10000)
-            .build();
-  }
+  public static final EbeanConfiguration testDefault = EbeanConfiguration.builder().build();
 }

metadata-service/configuration/src/main/java/com/linkedin/metadata/config/EbeanConfiguration.java

@@ -172,10 +172,6 @@ ebean:
   autoCreateDdl: ${EBEAN_AUTOCREATE:false}
   postgresUseIamAuth: ${EBEAN_POSTGRES_USE_AWS_IAM_AUTH:false}
   batchGetMethod: ${EBEAN_BATCH_GET_METHOD:IN} # Alternative UNION
-  locking:
-    enabled: ${EBEAN_LOCKING_ENABLED:false}
-    durationSeconds: ${EBEAN_LOCKING_DURATION_SECONDS:60}
-    maximumLocks: ${EBEAN_LOCKING_MAXIMUM_LOCKS:20000}
 
 # Only required if entityService.impl is cassandra
 cassandra:

metadata-service/configuration/src/main/resources/application.yaml

@@ -16,6 +16,7 @@
 import com.linkedin.metadata.entity.AspectDao;
 import com.linkedin.metadata.entity.IngestAspectsResult;
 import com.linkedin.metadata.entity.TransactionContext;
+import com.linkedin.metadata.entity.TransactionResult;
 import com.linkedin.metadata.entity.UpdateAspectResult;
 import com.linkedin.metadata.event.EventProducer;
 import io.datahubproject.metadata.context.OperationContext;
@@ -71,14 +72,16 @@ public void setup()
     OperationContext opContext = TestOperationContexts.systemContextNoSearchAuthorization();
     AspectDao aspectDao = Mockito.mock(AspectDao.class);
     when(aspectDao.runInTransactionWithRetry(
-            ArgumentMatchers.<Function<TransactionContext, List<UpdateAspectResult>>>any(),
+            ArgumentMatchers
+                .<Function<TransactionContext, TransactionResult<List<UpdateAspectResult>>>>any(),
             any(AspectsBatch.class),
             anyInt()))
         .thenAnswer(
             i ->
-                List.of(
-                    ((Function<TransactionContext, IngestAspectsResult>) i.getArgument(0))
-                        .apply(TransactionContext.empty(Mockito.mock(Transaction.class), 0))));
+                ((Function<TransactionContext, TransactionResult<IngestAspectsResult>>)
+                        i.getArgument(0))
+                    .apply(TransactionContext.empty(Mockito.mock(Transaction.class), 0))
+                    .getResults());
 
     EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class);
     PreProcessHooks preProcessHooks = new PreProcessHooks();