[#13983] feat(tms): Register policy timeTrigger on the expand scheduler

core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHook.java

@@ -0,0 +1,88 @@
+/*
+ * 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.gravitino.policy.hook;
+
+import org.apache.gravitino.meta.PolicyEntity;
+
+/**
+ * Storage-layer callbacks inside policy create / update / delete transactions, with JPA-style
+ * lifecycle points relative to the policy DML.
+ *
+ * <p>Distinct from {@code org.apache.gravitino.hook.PolicyHookDispatcher}, which decorates the
+ * policy API after a successful manager call (for example ownership). This SPI runs in {@code
+ * PolicyMetaService} on the open SQL session so extension writes commit or roll back with the
+ * policy row.
+ *
+ * <p>Implementations must use that connection, and must not commit, roll back, or close it. An
+ * exception rolls the policy write back together with the hook's own writes.
+ *
+ * <p>Defaults are no-ops so an extension only overrides the points it needs.
+ *
+ * <p>{@code insertPolicy(..., overwritten = true)} that replaces an existing row fires the update
+ * lifecycle ({@link #beforeUpdate} / {@link #afterUpdate}), not persist. A create, or an overwrite
+ * that inserts because no row was found, fires persist.
+ */
+public interface PolicyLifecycleHook {
+
+  /**
+   * About to insert the policy row in the current transaction.
+   *
+   * @param policy the policy about to be stored
+   */
+  default void beforePersist(PolicyEntity policy) {}
+
+  /**
+   * The policy row has been inserted in the current transaction.
+   *
+   * @param policy the policy as stored, including id, enabled flag, and content
+   */
+  default void afterPersist(PolicyEntity policy) {}
+
+  /**
+   * About to update the policy row in the current transaction.
+   *
+   * <p>Also used when {@code insertPolicy} overwrites an existing row by name or id.
+   *
+   * @param policy the policy about to be stored (id is the retained row id on overwrite)
+   */
+  default void beforeUpdate(PolicyEntity policy) {}
+
+  /**
+   * The policy row has been updated in the current transaction.
+   *
+   * <p>Also used when {@code insertPolicy} overwrites an existing row by name or id.
+   *
+   * @param policy the policy as stored, including the retained id, enabled flag, and content
+   */
+  default void afterUpdate(PolicyEntity policy) {}
+
+  /**
+   * The policy row is about to be deleted in the current transaction.
+   *
+   * @param policyId persisted policy id
+   */
+  default void beforeRemove(long policyId) {}
+
+  /**
+   * The policy row has been deleted in the current transaction.
+   *
+   * @param policyId persisted policy id
+   */
+  default void afterRemove(long policyId) {}
+}

core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHooks.java

@@ -0,0 +1,66 @@
+/*
+ * 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.gravitino.policy.hook;
+
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
+
+/**
+ * Holds the optional {@link PolicyLifecycleHook} registered by an extension.
+ *
+ * <p>This is a storage-layer SPI invoked from {@code PolicyMetaService} inside the open policy
+ * transaction. It is not the same as {@code org.apache.gravitino.hook.PolicyHookDispatcher}, which
+ * wraps the policy API for owner side effects after a successful manager call.
+ */
+public final class PolicyLifecycleHooks {
+
+  private static final AtomicReference<PolicyLifecycleHook> HOOK = new AtomicReference<>();
+
+  private PolicyLifecycleHooks() {}
+
+  /**
+   * Installs the hook invoked at policy create / update / delete lifecycle points inside the open
+   * transaction.
+   *
+   * @param hook hook to run during policy create, update, and delete transactions
+   * @throws IllegalStateException if a hook is already registered
+   */
+  public static void register(PolicyLifecycleHook hook) {
+    if (!HOOK.compareAndSet(null, hook)) {
+      throw new IllegalStateException("A policy lifecycle hook is already registered");
+    }
+  }
+
+  /**
+   * Removes the hook when it is the one currently registered.
+   *
+   * @param hook hook passed to {@link #register}
+   */
+  public static void unregister(PolicyLifecycleHook hook) {
+    HOOK.compareAndSet(hook, null);
+  }
+
+  /**
+   * @return the registered hook, or {@code null} when no extension is participating
+   */
+  @Nullable
+  public static PolicyLifecycleHook get() {
+    return HOOK.get();
+  }
+}

core/src/main/java/org/apache/gravitino/storage/relational/service/PolicyMetaService.java

@@ -40,6 +40,8 @@
 import org.apache.gravitino.exceptions.NoSuchEntityException;
 import org.apache.gravitino.meta.PolicyEntity;
 import org.apache.gravitino.metrics.Monitored;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHook;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHooks;
 import org.apache.gravitino.storage.relational.mapper.MetalakeMetaMapper;
 import org.apache.gravitino.storage.relational.mapper.OwnerMetaMapper;
 import org.apache.gravitino.storage.relational.mapper.PolicyMetaMapper;
@@ -172,8 +174,10 @@ public <E extends Entity & HasIdentifier> PolicyEntity updatePolicy(
       PolicyPO newPolicyPO =
           POConverters.updatePolicyPOWithVersion(oldPolicyPO, updatedPolicyEntity);
       SessionUtils.doMultipleWithCommit(
+          () -> beforeUpdate(updatedPolicyEntity),
           () -> updatePolicyRootWithVersion(ident, oldPolicyPO, newPolicyPO),
-          () -> insertPolicyVersionIfAllocated(oldPolicyPO, newPolicyPO));
+          () -> insertPolicyVersionIfAllocated(oldPolicyPO, newPolicyPO),
+          () -> afterUpdate(updatedPolicyEntity));
     } catch (RuntimeException re) {
       ExceptionUtils.checkSQLException(
           re, Entity.EntityType.POLICY, updatedPolicyEntity.nameIdentifier().toString());
@@ -212,6 +216,7 @@ boolean deletePolicy(NameIdentifier ident, PolicyPO policyPO) {
 
     try {
       SessionUtils.doMultipleWithCommit(
+          () -> beforeRemove(policyId),
           () -> deletePolicyWithVersion(ident, policyPO),
           () ->
               SessionUtils.doWithoutCommit(
@@ -237,7 +242,8 @@ boolean deletePolicy(NameIdentifier ident, PolicyPO policyPO) {
                   SecurableObjectMapper.class,
                   mapper ->
                       mapper.softDeleteObjectRelsByMetadataObject(
-                          policyId, MetadataObject.Type.POLICY.name())));
+                          policyId, MetadataObject.Type.POLICY.name())),
+          () -> afterRemove(policyId));
       return true;
     } catch (NoSuchEntityException e) {
       return false;
@@ -363,7 +369,12 @@ private void lockMetalakeForPolicyCreate(MetalakePO observedMetalakePO) {
   }
 
   /**
-   * Writes the policy row and its content snapshot.
+   * Writes the policy row and its content snapshot without committing, and fires the matching
+   * lifecycle hooks around the DML.
+   *
+   * <p>A plain create (or an overwrite that finds no row) uses persist hooks. An overwrite that
+   * replaces an existing row uses update hooks and passes the entity as stored — including the
+   * retained policy id when the caller supplied a different id for the same name.
    *
    * <p>An overwrite is not an upsert any more: the existing row is located and locked first, and
    * the replacement is written as the next version of that row, so the snapshot history survives
@@ -378,7 +389,9 @@ private void lockMetalakeForPolicyCreate(MetalakePO observedMetalakePO) {
   private void insertPolicyWithoutCommit(
       PolicyEntity policyEntity, PolicyPO initializedPolicyPO, boolean overwritten) {
     if (!overwritten) {
+      beforePersist(policyEntity);
       insertNewPolicyWithoutCommit(initializedPolicyPO);
+      afterPersist(policyEntity);
       return;
     }
 
@@ -392,16 +405,22 @@ private void insertPolicyWithoutCommit(
             "The policy ID %s is reserved by a deleted policy; use a new ID",
             initializedPolicyPO.getPolicyId());
       }
+      beforePersist(policyEntity);
       insertNewPolicyWithoutCommit(initializedPolicyPO);
+      afterPersist(policyEntity);
       return;
     }
 
     PolicyPO replacementPolicyPO =
         POConverters.updatePolicyPOWithVersion(existingPolicyPO, initializedPolicyPO);
+    PolicyEntity storedPolicy =
+        POConverters.fromPolicyPO(replacementPolicyPO, policyEntity.namespace());
     NameIdentifier observedIdentifier =
         NameIdentifier.of(policyEntity.namespace(), existingPolicyPO.getPolicyName());
+    beforeUpdate(storedPolicy);
     updatePolicyRootWithVersion(observedIdentifier, existingPolicyPO, replacementPolicyPO);
     insertPolicyVersionIfAllocated(existingPolicyPO, replacementPolicyPO);
+    afterUpdate(storedPolicy);
   }
 
   /**
@@ -570,4 +589,46 @@ private PolicyPO getPolicyPOByMetalakeAndName(String metalakeName, String policy
     }
     return policyPO;
   }
+
+  private static void beforePersist(PolicyEntity policy) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.beforePersist(policy);
+    }
+  }
+
+  private static void afterPersist(PolicyEntity policy) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.afterPersist(policy);
+    }
+  }
+
+  private static void beforeUpdate(PolicyEntity policy) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.beforeUpdate(policy);
+    }
+  }
+
+  private static void afterUpdate(PolicyEntity policy) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.afterUpdate(policy);
+    }
+  }
+
+  private static void beforeRemove(long policyId) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.beforeRemove(policyId);
+    }
+  }
+
+  private static void afterRemove(long policyId) {
+    PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+    if (hook != null) {
+      hook.afterRemove(policyId);
+    }
+  }
 }

core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessions.java

@@ -20,6 +20,8 @@
 package org.apache.gravitino.storage.relational.session;
 
 import com.google.common.annotations.VisibleForTesting;
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.ibatis.session.SqlSession;
 import org.apache.ibatis.session.TransactionIsolationLevel;
@@ -112,6 +114,23 @@ public static boolean hasActiveSession() {
     return sessions.get() != null;
   }
 
+  /**
+   * Returns the connection of the current thread's open session without changing the session count.
+   *
+   * <p>Callers must not commit, roll back, or close the connection. The session owner commits it.
+   *
+   * @return the connection bound to the current session
+   * @throws IllegalStateException if the current thread has no open session
+   * @throws SQLException if the session cannot provide its connection
+   */
+  public static Connection currentConnection() throws SQLException {
+    SqlSession sqlSession = sessions.get();
+    if (sqlSession == null) {
+      throw new IllegalStateException("No active SQL session on the current thread");
+    }
+    return sqlSession.getConnection();
+  }
+
   private static void handleSessionClose(boolean commit, boolean rollback) {
     SqlSession sqlSession = sessions.get();
     if (sqlSession == null) {
@@ -124,7 +143,11 @@ private static void handleSessionClose(boolean commit, boolean rollback) {
         if (commit) {
           sqlSession.commit();
         } else if (rollback) {
-          sqlSession.rollback();
+          // Force the JDBC rollback. A policy lifecycle hook can write on this connection without
+          // going
+          // through a mapped statement, and MyBatis skips rollback when the session looks clean.
+          // Closing that session then turns auto-commit back on, which would commit the hook write.
+          sqlSession.rollback(true);
         }
       } finally {
         try {

core/src/test/java/org/apache/gravitino/policy/hook/TestPolicyLifecycleHooks.java

@@ -0,0 +1,59 @@
+/*
+ * 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.gravitino.policy.hook;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+public class TestPolicyLifecycleHooks {
+
+  private final PolicyLifecycleHook hook = new PolicyLifecycleHook() {};
+
+  @AfterEach
+  void tearDown() {
+    PolicyLifecycleHooks.unregister(hook);
+  }
+
+  @Test
+  void testRegisterAndUnregister() {
+    PolicyLifecycleHooks.register(hook);
+    assertSame(hook, PolicyLifecycleHooks.get());
+    PolicyLifecycleHooks.unregister(hook);
+    assertNull(PolicyLifecycleHooks.get());
+  }
+
+  @Test
+  void testDoubleRegisterRejected() {
+    PolicyLifecycleHooks.register(hook);
+    PolicyLifecycleHook other = new PolicyLifecycleHook() {};
+    assertThrows(IllegalStateException.class, () -> PolicyLifecycleHooks.register(other));
+    assertSame(hook, PolicyLifecycleHooks.get());
+  }
+
+  @Test
+  void testUnregisterWrongHookIsNoOp() {
+    PolicyLifecycleHooks.register(hook);
+    PolicyLifecycleHooks.unregister(new PolicyLifecycleHook() {});
+    assertSame(hook, PolicyLifecycleHooks.get());
+  }
+}