Introduce MutationPoint-based filtering for EntityMutator

XJDKC · XJDKC · commit b5bc2b59dcb7 · 2025-06-05T21:53:46.000-07:00
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/identity/mutation/EntityMutationEngine.java b/polaris-core/src/main/java/org/apache/polaris/core/identity/mutation/EntityMutationEngine.java
@@ -33,9 +33,10 @@ public interface EntityMutationEngine {
   /**
    * Applies all registered entity mutators to the provided entity, in order.
    *
+   * @param mutationPoint The point in the entity lifecycle where mutations should be applied.
    * @param entity The original Polaris entity to mutate.
    * @return A new or modified instance of {@link PolarisBaseEntity} after all mutations are
    *     applied.
    */
-  PolarisBaseEntity applyMutations(PolarisBaseEntity entity);
+  PolarisBaseEntity applyMutations(MutationPoint mutationPoint, PolarisBaseEntity entity);
 }
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/identity/mutation/MutationPoint.java b/polaris-core/src/main/java/org/apache/polaris/core/identity/mutation/MutationPoint.java
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.polaris.core.identity.mutation;
+
+/**
+ * Defines points in the entity lifecycle where {@link EntityMutator} can be applied.
+ *
+ * <p>Each mutation point corresponds to a specific hook where mutations may be executed. Mutators
+ * can declare which points they support, allowing the engine to invoke only the relevant ones.
+ */
+public enum MutationPoint {
+
+  /** Applied before a catalog entity is persisted. */
+  CATALOG_PRE_PERSIST(0),
+  ;
+
+  private final int id;
+
+  MutationPoint(int id) {
+    this.id = id;
+  }
+
+  public int id() {
+    return id;
+  }
+}
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java
@@ -47,6 +47,8 @@
 import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
 import org.apache.polaris.core.entity.PolarisPrivilege;
 import org.apache.polaris.core.entity.PolarisTaskConstants;
+import org.apache.polaris.core.identity.mutation.EntityMutationEngine;
+import org.apache.polaris.core.identity.mutation.MutationPoint;
 import org.apache.polaris.core.persistence.*;
 import org.apache.polaris.core.persistence.dao.entity.BaseResult;
 import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult;
@@ -473,8 +475,10 @@ private void revokeGrantRecord(
 
     ms.persistStorageIntegrationIfNeededInCurrentTxn(callCtx, catalog, integration);
 
-    if (callCtx.getEntityMutationEngine() != null) {
-      catalog = callCtx.getEntityMutationEngine().applyMutations(catalog);
+    // CATALOG_PRE_PERSIST mutation point
+    EntityMutationEngine entityMutationEngine = callCtx.getEntityMutationEngine();
+    if (entityMutationEngine != null) {
+      entityMutationEngine.applyMutations(MutationPoint.CATALOG_PRE_PERSIST, catalog);
     }
 
     // now create and persist new catalog entity
diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/identity/mutation/NoOpEntityMutationEngine.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/identity/mutation/NoOpEntityMutationEngine.java
@@ -29,7 +29,7 @@
 public class NoOpEntityMutationEngine implements EntityMutationEngine {
 
   @Override
-  public PolarisBaseEntity applyMutations(PolarisBaseEntity entity) {
+  public PolarisBaseEntity applyMutations(MutationPoint mutationPoint, PolarisBaseEntity entity) {
     return entity;
   }
 }
diff --git a/quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusProducers.java b/quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusProducers.java
@@ -75,7 +75,7 @@ PolarisStorageIntegration<T> getStorageIntegrationForConfig(
   @Produces
   public EntityMutationEngine entityMutationEngine() {
     // An entity mutation engine is not required when running the admin tool.
-    return entity -> entity;
+    return ((mutationPoint, entity) -> entity);
   }
 
   @Produces
diff --git a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/identity/mutation/CatalogEntityConnectionConfigMutator.java b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/identity/mutation/CatalogEntityConnectionConfigMutator.java
@@ -29,7 +29,9 @@
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.core.identity.ServiceIdentityType;
 import org.apache.polaris.core.identity.dpo.ServiceIdentityInfoDpo;
+import org.apache.polaris.service.identity.mutation.AppliesTo;
 import org.apache.polaris.core.identity.mutation.EntityMutator;
+import org.apache.polaris.core.identity.mutation.MutationPoint;
 import org.apache.polaris.core.identity.registry.ServiceIdentityRegistry;
 
 /**
@@ -41,6 +43,7 @@
  */
 @RequestScoped
 @Identifier("catalog-connection-config")
+@AppliesTo(MutationPoint.CATALOG_PRE_PERSIST)
 public class CatalogEntityConnectionConfigMutator implements EntityMutator {
 
   private final ServiceIdentityRegistry serviceIdentityRegistry;
diff --git a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/identity/mutation/NoOpEntityMutator.java b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/identity/mutation/NoOpEntityMutator.java
@@ -23,6 +23,7 @@
 import jakarta.enterprise.context.ApplicationScoped;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.core.identity.mutation.EntityMutator;
+import org.apache.polaris.service.identity.mutation.AppliesTo;
 
 /**
  * A no-op implementation of {@link EntityMutator} that returns the entity unchanged.
diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java
@@ -45,7 +45,8 @@ public class JWTSymmetricKeyGeneratorTest {
   /** Sanity test to verify that we can generate a token */
   @Test
   public void testJWTSymmetricKeyGenerator() {
-    PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, null, null, null, null);
+    PolarisCallContext polarisCallContext =
+        new PolarisCallContext(null, null, null, null, null, null);
     PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class);
     String mainSecret = "test_secret";
     String clientId = "test_client_id";
diff --git a/service/common/src/main/java/org/apache/polaris/service/identity/mutation/AppliesTo.java b/service/common/src/main/java/org/apache/polaris/service/identity/mutation/AppliesTo.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.polaris.service.identity.mutation;
+
+import jakarta.enterprise.util.AnnotationLiteral;
+import jakarta.inject.Qualifier;
+import java.lang.annotation.*;
+import org.apache.polaris.core.identity.mutation.EntityMutationEngine;
+import org.apache.polaris.core.identity.mutation.EntityMutator;
+import org.apache.polaris.core.identity.mutation.MutationPoint;
+
+/**
+ * Qualifier to mark an {@link EntityMutator} as applicable to a specific {@link MutationPoint}.
+ *
+ * <p>This is used by the {@link EntityMutationEngine} to apply only relevant mutators based on
+ * context.
+ *
+ * <p>Supports being repeated on the same class to handle multiple mutation points.
+ */
+@Qualifier
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+@Repeatable(AppliesTos.class)
+public @interface AppliesTo {
+
+  /** The mutation point this mutator applies to. */
+  MutationPoint value();
+
+  /** Helper for creating {@link AppliesTo} qualifiers programmatically. */
+  final class Literal extends AnnotationLiteral<AppliesTo> implements AppliesTo {
+    private final MutationPoint value;
+
+    public static Literal of(MutationPoint value) {
+      return new Literal(value);
+    }
+
+    private Literal(MutationPoint value) {
+      this.value = value;
+    }
+
+    @Override
+    public MutationPoint value() {
+      return value;
+    }
+  }
+}
diff --git a/service/common/src/main/java/org/apache/polaris/service/identity/mutation/AppliesTos.java b/service/common/src/main/java/org/apache/polaris/service/identity/mutation/AppliesTos.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.polaris.service.identity.mutation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import org.apache.polaris.core.identity.mutation.EntityMutator;
+import org.apache.polaris.core.identity.mutation.MutationPoint;
+
+/**
+ * Container annotation for repeating {@link AppliesTo}.
+ *
+ * <p>Allows an {@link EntityMutator} to declare support for multiple {@link MutationPoint}s.
+ */
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface AppliesTos {
+  AppliesTo[] value();
+}
diff --git a/service/common/src/main/java/org/apache/polaris/service/identity/mutation/EntityMutationEngineImpl.java b/service/common/src/main/java/org/apache/polaris/service/identity/mutation/EntityMutationEngineImpl.java
@@ -25,9 +25,11 @@
 import jakarta.enterprise.inject.Instance;
 import jakarta.inject.Inject;
 import java.util.List;
+import java.util.Objects;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.core.identity.mutation.EntityMutationEngine;
 import org.apache.polaris.core.identity.mutation.EntityMutator;
+import org.apache.polaris.core.identity.mutation.MutationPoint;
 
 @ApplicationScoped
 public class EntityMutationEngineImpl implements EntityMutationEngine {
@@ -43,26 +45,28 @@ public EntityMutationEngineImpl(
   }
 
   @Override
-  public PolarisBaseEntity applyMutations(PolarisBaseEntity entity) {
+  public PolarisBaseEntity applyMutations(MutationPoint mutationPoint, PolarisBaseEntity entity) {
     PolarisBaseEntity result = entity;
 
+    // Collect mutators in configured order, filtering only those applicable to the mutation point
     List<EntityMutator> orderedMutators =
-        // config.map(EntityMutationConfiguration::mutators).orElse(List.of()).stream()
         config.mutators().orElse(List.of()).stream()
             .map(
                 id -> {
+                  // Resolve the mutator instance by ID
                   Instance<EntityMutator> matched =
                       mutatorInstances.select(Identifier.Literal.of(id));
-                  if (matched.isResolvable()) {
-                    return matched.get();
-                  } else {
+                  if (!matched.isResolvable()) {
                     throw new IllegalStateException("No EntityMutator found for ID: " + id);
                   }
+                  // Filter by MutationPoint via @AppliesTo
+                  Instance<EntityMutator> filtered =
+                      matched.select(AppliesTo.Literal.of(mutationPoint));
+                  return filtered.isResolvable() ? filtered.get() : null;
                 })
+            .filter(Objects::nonNull)
             .toList();
 
-    // Apply all mutators to the entity
-    // TODO: Add a way to control the order of mutators
     for (EntityMutator mutator : orderedMutators) {
       result = mutator.apply(result);
     }

Original file line number	Diff line number	Diff line change
`@@ -33,9 +33,10 @@ public interface EntityMutationEngine {`
`33`	`33`	`/**`
`34`	`34`	`* Applies all registered entity mutators to the provided entity, in order.`
`35`	`35`	`*`
	`36`	`+ * @param mutationPoint The point in the entity lifecycle where mutations should be applied.`
`36`	`37`	`* @param entity The original Polaris entity to mutate.`
`37`	`38`	`* @return A new or modified instance of {@link PolarisBaseEntity} after all mutations are`
`38`	`39`	`* applied.`
`39`	`40`	`*/`
`40`		`- PolarisBaseEntity applyMutations(PolarisBaseEntity entity);`
	`41`	`+ PolarisBaseEntity applyMutations(MutationPoint mutationPoint, PolarisBaseEntity entity);`
`41`	`42`	`}`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`	`public class NoOpEntityMutationEngine implements EntityMutationEngine {`
`30`	`30`
`31`	`31`	`@Override`
`32`		`- public PolarisBaseEntity applyMutations(PolarisBaseEntity entity) {`
	`32`	`+ public PolarisBaseEntity applyMutations(MutationPoint mutationPoint, PolarisBaseEntity entity) {`
`33`	`33`	`return entity;`
`34`	`34`	`}`
`35`	`35`	`}`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ PolarisStorageIntegration<T> getStorageIntegrationForConfig(`
`75`	`75`	`@Produces`
`76`	`76`	`public EntityMutationEngine entityMutationEngine() {`
`77`	`77`	`// An entity mutation engine is not required when running the admin tool.`
`78`		`- return entity -> entity;`
	`78`	`+ return ((mutationPoint, entity) -> entity);`
`79`	`79`	`}`
`80`	`80`
`81`	`81`	`@Produces`