Make entity drop and cleanup task creation atomic

persistence/nosql/persistence/metastore/src/main/java/org/apache/polaris/persistence/nosql/metastore/NonFunctionalBasePersistence.java

@@ -82,6 +82,14 @@ public void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBa
     throw unimplemented();
   }
 
+  @Override
+  public void deleteEntityAndCreateEntities(
+      @NonNull PolarisCallContext callCtx,
+      @NonNull PolarisBaseEntity entityToDelete,
+      @NonNull List<PolarisBaseEntity> entitiesToCreate) {
+    throw useMetaStoreManager("create/update/rename/delete");
+  }
+
   @Override
   public void deleteFromGrantRecords(
       @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) {

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

@@ -352,6 +352,49 @@ public void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBa
     }
   }
 
+  @Override
+  public void deleteEntityAndCreateEntities(
+      @NonNull PolarisCallContext callCtx,
+      @NonNull PolarisBaseEntity entityToDelete,
+      @NonNull List<PolarisBaseEntity> entitiesToCreate) {
+    ModelEntity modelEntity = ModelEntity.fromEntity(entityToDelete, schemaVersion);
+    Map<String, Object> params =
+        Map.of(
+            "id",
+            modelEntity.getId(),
+            "catalog_id",
+            modelEntity.getCatalogId(),
+            "realm_id",
+            realmId);
+    try {
+      datasourceOperations.runWithinTransaction(
+          connection -> {
+            datasourceOperations.execute(
+                connection,
+                QueryGenerator.generateDeleteQuery(
+                    ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params));
+            for (PolarisBaseEntity entityToCreate : entitiesToCreate) {
+              PolarisBaseEntity existingEntity =
+                  lookupEntity(
+                      connection,
+                      entityToCreate.getCatalogId(),
+                      entityToCreate.getId(),
+                      entityToCreate.getTypeCode());
+              if (existingEntity != null) {
+                continue;
+              }
+              persistEntity(
+                  callCtx, entityToCreate, null, connection, datasourceOperations::execute);
+            }
+            return true;
+          });
+    } catch (SQLException e) {
+      throw new RuntimeException(
+          String.format("Failed to delete entity and create entities due to %s", e.getMessage()),
+          e);
+    }
+  }
+
   @Override
   public void deleteFromGrantRecords(
       @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) {
@@ -421,11 +464,25 @@ public void deleteAll(@NonNull PolarisCallContext callCtx) {
   @Override
   public PolarisBaseEntity lookupEntity(
       @NonNull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode) {
+    return getPolarisBaseEntity(entityLookupQuery(catalogId, entityId, typeCode));
+  }
+
+  private PolarisBaseEntity lookupEntity(
+      @NonNull Connection connection, long catalogId, long entityId, int typeCode)
+      throws SQLException {
+    return getPolarisBaseEntity(
+        datasourceOperations.executeSelect(
+            connection,
+            entityLookupQuery(catalogId, entityId, typeCode),
+            new ModelEntity(schemaVersion)));
+  }
+
+  private QueryGenerator.PreparedQuery entityLookupQuery(
+      long catalogId, long entityId, int typeCode) {
     Map<String, Object> params =
         Map.of("catalog_id", catalogId, "id", entityId, "type_code", typeCode, "realm_id", realmId);
-    return getPolarisBaseEntity(
-        QueryGenerator.generateSelectQuery(
-            ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params));
+    return QueryGenerator.generateSelectQuery(
+        ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params);
   }
 
   @Override
@@ -455,23 +512,28 @@ public PolarisBaseEntity lookupEntityByName(
   @Nullable
   private PolarisBaseEntity getPolarisBaseEntity(QueryGenerator.PreparedQuery query) {
     try {
-      var results = datasourceOperations.executeSelect(query, new ModelEntity(schemaVersion));
-      if (results.isEmpty()) {
-        return null;
-      } else if (results.size() > 1) {
-        throw new IllegalStateException(
-            String.format(
-                "More than one(%s) entities were found for a given type code : %s",
-                results.size(), results.getFirst().getTypeCode()));
-      } else {
-        return results.getFirst();
-      }
+      return getPolarisBaseEntity(
+          datasourceOperations.executeSelect(query, new ModelEntity(schemaVersion)));
     } catch (SQLException e) {
       throw new RuntimeException(
           String.format("Failed to retrieve polaris entity due to %s", e.getMessage()), e);
     }
   }
 
+  @Nullable
+  private PolarisBaseEntity getPolarisBaseEntity(List<PolarisBaseEntity> results) {
+    if (results.isEmpty()) {
+      return null;
+    } else if (results.size() > 1) {
+      throw new IllegalStateException(
+          String.format(
+              "More than one(%s) entities were found for a given type code : %s",
+              results.size(), results.getFirst().getTypeCode()));
+    } else {
+      return results.getFirst();
+    }
+  }
+
   @NonNull
   @Override
   public List<PolarisBaseEntity> lookupEntities(

persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java

@@ -48,6 +48,16 @@
 public abstract class AtomicMetastoreManagerWithJdbcBasePersistenceImplTest
     extends BasePolarisMetaStoreManagerTest {
 
+  @Test
+  void testCleanupTaskCreationFailureRollsBackEntityDrop() {
+    assertCleanupTaskCreationFailureRollsBackEntityDrop();
+  }
+
+  @Test
+  void testCleanupTaskCreationRetryIsIdempotent() {
+    assertCleanupTaskCreationRetryIsIdempotent();
+  }
+
   protected DatabaseType databaseType() {
     return DatabaseType.H2;
   }

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -177,6 +177,14 @@ private void dropEntity(
       @NonNull PolarisCallContext callCtx,
       @NonNull BasePersistence ms,
       @NonNull PolarisBaseEntity entity) {
+    dropEntity(callCtx, ms, entity, List.of());
+  }
+
+  private void dropEntity(
+      @NonNull PolarisCallContext callCtx,
+      @NonNull BasePersistence ms,
+      @NonNull PolarisBaseEntity entity,
+      @NonNull List<PolarisBaseEntity> entitiesToCreateAtomically) {
 
     // validate the entity type and subtype
     getDiagnostics().checkNotNull(entity, "unexpected_null_dpo");
@@ -187,7 +195,11 @@ private void dropEntity(
 
     // Remove the main entity itself first-thing; once its id no longer resolves successfully
     // it will be pruned out of any grant-record lookups anyways.
-    ms.deleteEntity(callCtx, entity);
+    if (entitiesToCreateAtomically.isEmpty()) {
+      ms.deleteEntity(callCtx, entity);
+    } else {
+      ms.deleteEntityAndCreateEntities(callCtx, entity, entitiesToCreateAtomically);
+    }
 
     // Best-effort cleanup - drop grant records, update grantRecordVersions for affected
     // other entities.
@@ -1092,6 +1104,18 @@ public void deletePrincipalSecrets(
 
     // if this entity was not found, return failure
     if (refreshEntityToDrop == null) {
+      if (cleanup && entityToDrop.getType() != PolarisEntityType.POLICY) {
+        PolarisBaseEntity cleanupTask =
+            ms.lookupEntityByName(
+                callCtx,
+                PolarisEntityConstants.getNullId(),
+                PolarisEntityConstants.getNullId(),
+                PolarisEntityType.TASK.getCode(),
+                "entityCleanup_" + entityToDrop.getId());
+        if (cleanupTask != null) {
+          return new DropEntityResult(cleanupTask.getId());
+        }
+      }
       return new DropEntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null);
     }
 
@@ -1182,10 +1206,6 @@ public void deletePrincipalSecrets(
       }
     }
 
-    // simply delete that entity. Will be removed from entities_active, added to the
-    // entities_dropped and its version will be changed.
-    this.dropEntity(callCtx, ms, refreshEntityToDrop);
-
     // if cleanup, schedule a cleanup task for the entity. do this here, so that drop and scheduling
     // the cleanup task is transactional. Otherwise, we'll be unable to schedule the cleanup task
     // later
@@ -1201,22 +1221,33 @@ public void deletePrincipalSecrets(
               .propertiesAsMap(properties)
               .id(ms.generateNewId(callCtx))
               .catalogId(0L)
-              .name("entityCleanup_" + entityToDrop.getId())
+              .name("entityCleanup_" + refreshEntityToDrop.getId())
               .typeCode(PolarisEntityType.TASK.getCode())
               .subTypeCode(PolarisEntitySubType.NULL_SUBTYPE.getCode())
               .createTimestamp(clock.millis());
       if (cleanupProperties != null) {
         taskEntityBuilder.internalPropertiesAsMap(cleanupProperties);
       }
-      // TODO: Add a way to create the task entities atomically with dropping the entity;
-      // in the meantime, if the server fails partway through a dropEntity, it's possible that
-      // the entity is dropped but we don't have any persisted task records that will carry
-      // out the cleanup.
-      PolarisBaseEntity taskEntity = taskEntityBuilder.build();
-      createEntityIfNotExists(callCtx, null, taskEntity);
+      PolarisBaseEntity taskEntity =
+          prepareToPersistNewEntity(callCtx, ms, taskEntityBuilder.build());
+      try {
+        this.dropEntity(callCtx, ms, refreshEntityToDrop, List.of(taskEntity));
+      } catch (EntityAlreadyExistsException e) {
+        return new DropEntityResult(
+            BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS,
+            String.format(
+                "Existing entity id: '%s', type %s subtype %s",
+                e.getExistingEntity().getId(),
+                e.getExistingEntity().getTypeCode(),
+                e.getExistingEntity().getSubTypeCode()));
+      }
       return new DropEntityResult(taskEntity.getId());
     }
 
+    // simply delete that entity. Will be removed from entities_active, added to the
+    // entities_dropped and its version will be changed.
+    this.dropEntity(callCtx, ms, refreshEntityToDrop);
+
     // done, return success
     return new DropEntityResult();
   }

polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java

@@ -158,6 +158,23 @@ void writeToGrantRecords(
    */
   void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBaseEntity entity);
 
+  /**
+   * Delete one entity and create the supplied entities in one atomic persistence operation. If the
+   * operation succeeds, the deleted entity must be durably removed and every created entity must be
+   * durably visible; if it fails, none of those changes may be applied.
+   *
+   * <p>The created entities use the same create semantics as {@link
+   * #writeEntities(PolarisCallContext, List, List)} with {@code originalEntities == null}.
+   *
+   * @param callCtx call context
+   * @param entityToDelete entity to delete
+   * @param entitiesToCreate entities to create atomically with the delete
+   */
+  void deleteEntityAndCreateEntities(
+      @NonNull PolarisCallContext callCtx,
+      @NonNull PolarisBaseEntity entityToDelete,
+      @NonNull List<PolarisBaseEntity> entitiesToCreate);
+
   /**
    * Delete the specified grantRecord to the grant_records table.
    *

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java

@@ -285,6 +285,32 @@ public void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBa
     runActionInTransaction(callCtx, () -> this.deleteEntityInCurrentTxn(callCtx, entity));
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public void deleteEntityAndCreateEntities(
+      @NonNull PolarisCallContext callCtx,
+      @NonNull PolarisBaseEntity entityToDelete,
+      @NonNull List<PolarisBaseEntity> entitiesToCreate) {
+    runActionInTransaction(
+        callCtx,
+        () -> {
+          this.deleteEntityInCurrentTxn(callCtx, entityToDelete);
+          for (PolarisBaseEntity entityToCreate : entitiesToCreate) {
+            try {
+              this.checkConditionsForWriteEntityInCurrentTxn(callCtx, entityToCreate, null);
+            } catch (EntityAlreadyExistsException e) {
+              // Matching ids indicate a retried create whose id was already reserved for the same
+              // entity. Treat that as idempotent, matching writeEntities.
+              if (e.getExistingEntity().getId() != entityToCreate.getId()) {
+                throw e;
+              }
+              continue;
+            }
+            this.writeEntityInCurrentTxn(callCtx, entityToCreate, true, null);
+          }
+        });
+  }
+
   /** {@inheritDoc} */
   @Override
   public void deleteFromGrantRecords(

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java

@@ -51,6 +51,7 @@
 import org.apache.polaris.core.entity.PolarisTaskConstants;
 import org.apache.polaris.core.entity.PrincipalEntity;
 import org.apache.polaris.core.persistence.BaseMetaStoreManager;
+import org.apache.polaris.core.persistence.EntityAlreadyExistsException;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
 import org.apache.polaris.core.persistence.PolarisObjectMapperUtil;
 import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException;
@@ -1316,6 +1317,26 @@ public void deletePrincipalSecrets(
     // entity cannot be null
     getDiagnostics().checkNotNull(entityToDrop, "unexpected_null_entity");
 
+    // Resolve the entity by id before validating its path so a retry after a successful commit can
+    // recover the cleanup task even though the dropped entity no longer resolves.
+    PolarisBaseEntity refreshEntityToDrop =
+        ms.lookupEntityInCurrentTxn(
+            callCtx, entityToDrop.getCatalogId(), entityToDrop.getId(), entityToDrop.getTypeCode());
+    if (refreshEntityToDrop == null
+        && cleanup
+        && entityToDrop.getType() != PolarisEntityType.POLICY) {
+      PolarisBaseEntity cleanupTask =
+          ms.lookupEntityByNameInCurrentTxn(
+              callCtx,
+              PolarisEntityConstants.getNullId(),
+              PolarisEntityConstants.getNullId(),
+              PolarisEntityType.TASK.getCode(),
+              "entityCleanup_" + entityToDrop.getId());
+      if (cleanupTask != null) {
+        return new DropEntityResult(cleanupTask.getId());
+      }
+    }
+
     // re-resolve everything including that entity
     PolarisEntityResolver resolver =
         new PolarisEntityResolver(getDiagnostics(), callCtx, ms, catalogPath, entityToDrop);
@@ -1325,11 +1346,6 @@ public void deletePrincipalSecrets(
       return new DropEntityResult(BaseResult.ReturnStatus.CATALOG_PATH_CANNOT_BE_RESOLVED, null);
     }
 
-    // first find the entity to drop
-    PolarisBaseEntity refreshEntityToDrop =
-        ms.lookupEntityInCurrentTxn(
-            callCtx, entityToDrop.getCatalogId(), entityToDrop.getId(), entityToDrop.getTypeCode());
-
     // if this entity was not found, return failure
     if (refreshEntityToDrop == null) {
       return new DropEntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null);
@@ -1443,7 +1459,12 @@ public void deletePrincipalSecrets(
         taskEntityBuilder.internalPropertiesAsMap(cleanupProperties);
       }
       PolarisBaseEntity taskEntity = taskEntityBuilder.build();
-      createEntityIfNotExists(callCtx, ms, null, taskEntity);
+      EntityResult createTaskResult = createEntityIfNotExists(callCtx, ms, null, taskEntity);
+      if (!createTaskResult.isSuccess()) {
+        ms.rollback();
+        return new DropEntityResult(
+            createTaskResult.getReturnStatus(), createTaskResult.getExtraInformation());
+      }
       return new DropEntityResult(taskEntity.getId());
     }
 
@@ -1463,11 +1484,21 @@ public void deletePrincipalSecrets(
     TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore());
 
     // need to run inside a read/write transaction
-    return ms.runInTransaction(
-        callCtx,
-        () ->
-            this.dropEntityIfExists(
-                callCtx, ms, catalogPath, entityToDrop, cleanupProperties, cleanup));
+    try {
+      return ms.runInTransaction(
+          callCtx,
+          () ->
+              this.dropEntityIfExists(
+                  callCtx, ms, catalogPath, entityToDrop, cleanupProperties, cleanup));
+    } catch (EntityAlreadyExistsException e) {
+      return new DropEntityResult(
+          BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS,
+          String.format(
+              "Existing entity id: '%s', type %s subtype %s",
+              e.getExistingEntity().getId(),
+              e.getExistingEntity().getTypeCode(),
+              e.getExistingEntity().getSubTypeCode()));
+    }
   }
 
   /**