This is an automated email from the ASF dual-hosted git repository. roryqi pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit fb255f9885509d674926a46fddfcb6a1d1441417 Author: MaSai <[email protected]> AuthorDate: Mon Oct 26 10:20:44 2026 +0800 [#13984] feat(tms): Nearest-policy resolver and expand table enumerator (#14055) ### What changes were proposed in this pull request? First slice of the expand track. Targets `branch-tms-scheduler`. - Extend `EffectiveTagResolver` / `ObjectPolicyResolver` so callers can resolve matched policies with contributing metadata-object sources (`ResolvedTag` / `Match` nested on the resolvers). Existing `resolve()` / PolicyManager list shape is unchanged. - `NearestMaintenancePolicyResolver` applies maintenance nearest-level and same-level conflict rules on those matches (`WINNER` / `NONE` / `CONFLICT`). - `ExpandTableCandidate` carries the table id and metadata-object type of a winner. - `NearestPolicyTableEnumerator` walks policy → tags → metadata objects, expands catalog/schema attachments to tables, and keeps a table only when this policy is the nearest winner and `ExpandOccupancyGate` allows it. Cooldown / in-flight winners do not consume the batch budget. This PR does not write `policy-expand` or `cron-job` rows. ### Why are the changes needed? Expand has to pick one nearest maintenance policy per table, and list the tables that policy wins, before any cron-job row is written. Matching must reuse effective-tag inheritance so a child tag-value override (e.g. table `data_domain=risk`) does not keep selecting an ancestor `TAG_VALUE(finance)` policy. Fix: #13984 ### Does this PR introduce _any_ user-facing change? No. Core adds resolver APIs used internally; no REST / property-key changes. ### How was this patch tested? - `./gradlew :core:test -PskipITs --tests org.apache.gravitino.policy.TestObjectPolicyResolver --tests org.apache.gravitino.tag.TestEffectiveTagResolver` - `./gradlew :plugins:tms:test -PskipITs -PskipDockerTests=true --tests org.apache.gravitino.maintenance.scheduler.expand.TestNearestMaintenancePolicyResolver --tests org.apache.gravitino.maintenance.scheduler.expand.TestNearestPolicyTableEnumerator` --------- Co-authored-by: Cursor <[email protected]> --- .../gravitino/policy/ObjectPolicyResolver.java | 132 ++++++-- .../apache/gravitino/tag/EffectiveTagResolver.java | 89 ++++- .../gravitino/policy/TestObjectPolicyResolver.java | 51 ++- .../gravitino/tag/TestEffectiveTagResolver.java | 29 ++ .../maintenance/scheduler/ExpandOccupancyGate.java | 38 +++ .../scheduler/expand/ExpandTableCandidate.java | 88 +++++ .../expand/NearestMaintenancePolicyResolver.java | 184 +++++++++++ .../expand/NearestPolicyTableEnumerator.java | 311 ++++++++++++++++++ .../scheduler/expand/PolicyTableEnumerator.java | 48 +++ .../TestNearestMaintenancePolicyResolver.java | 163 ++++++++++ .../expand/TestNearestPolicyTableEnumerator.java | 358 +++++++++++++++++++++ 11 files changed, 1463 insertions(+), 28 deletions(-) diff --git a/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java b/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java index eb1676b7da..1a244fd068 100644 --- a/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java +++ b/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java @@ -19,10 +19,15 @@ package org.apache.gravitino.policy; import java.io.IOException; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; import org.apache.gravitino.Entity; import org.apache.gravitino.EntityStore; @@ -34,6 +39,7 @@ import org.apache.gravitino.json.PolicyAssociationSelectorSerde; import org.apache.gravitino.meta.PolicyEntity; import org.apache.gravitino.meta.TagEntity; import org.apache.gravitino.tag.EffectiveTagResolver; +import org.apache.gravitino.tag.EffectiveTagResolver.ResolvedTag; import org.apache.gravitino.tag.TagAssignment; import org.apache.gravitino.utils.NameIdentifierUtil; import org.slf4j.Logger; @@ -43,7 +49,7 @@ import org.slf4j.LoggerFactory; * Resolves policies for a metadata object from its effective tag assignments. * * <p>The resolver evaluates relation selectors, filters disabled policies, and deduplicates - * repeated matches by policy entity ID. + * repeated matches by policy entity ID while preserving contributing metadata-object sources. */ public class ObjectPolicyResolver { @@ -74,16 +80,38 @@ public class ObjectPolicyResolver { * @return Enabled policies selected by the object's effective tags. */ public PolicyEntity[] resolve(String metalake, MetadataObject metadataObject) { - TagEntity[] effectiveTags = effectiveTagResolver.resolve(metalake, metadataObject); + Match[] matches = resolveMatches(metalake, metadataObject); + PolicyEntity[] policies = new PolicyEntity[matches.length]; + for (int i = 0; i < matches.length; i++) { + policies[i] = matches[i].policy(); + } + return policies; + } + + /** + * Resolves enabled policies together with the metadata-object sources that caused each match. + * + * <p>{@link #resolve(String, MetadataObject)} stays the PolicyManager list shape; callers that + * need contributing sources should use this method. + * + * @param metalake The metalake name. + * @param metadataObject The metadata object. + * @return Matched enabled policies with contributing sources. + */ + public Match[] resolveMatches(String metalake, MetadataObject metadataObject) { + ResolvedTag[] effectiveTags = effectiveTagResolver.resolveWithSources(metalake, metadataObject); if (effectiveTags.length == 0) { - return new PolicyEntity[0]; + return new Match[0]; } - Map<String, TagEntity> tagsByName = + Map<String, ResolvedTag> tagsByName = Arrays.stream(effectiveTags) .collect( Collectors.toMap( - TagEntity::name, tag -> tag, (left, right) -> left, LinkedHashMap::new)); + resolvedTag -> resolvedTag.tag().name(), + resolvedTag -> resolvedTag, + (left, right) -> left, + LinkedHashMap::new)); List<NameIdentifier> tagIdentifiers = tagsByName.keySet().stream() .map(tagName -> NameIdentifierUtil.ofTag(metalake, tagName)) @@ -105,10 +133,11 @@ public class ObjectPolicyResolver { Map<Long, MatchState> matchStates = new LinkedHashMap<>(); for (RelationalEntity<?> relation : relations) { - TagEntity tag = tagsByName.get(relation.source().name()); - if (tag == null) { + ResolvedTag resolvedTag = tagsByName.get(relation.source().name()); + if (resolvedTag == null) { continue; } + TagEntity tag = resolvedTag.tag(); PolicyEntity policy = (PolicyEntity) relation.targetEntity(); PolicyAssociationSelector selector = relation @@ -119,14 +148,14 @@ public class ObjectPolicyResolver { boolean matches = matches(selector, assignment); MatchState state = matchStates.computeIfAbsent(policy.id(), ignored -> new MatchState(policy)); - state.record(matches, tag.inherited().orElse(false)); + state.record(matches, tag.inherited().orElse(false), resolvedTag.source()); } return matchStates.values().stream() .filter(MatchState::matched) - .map(MatchState::policy) - .filter(PolicyEntity::enabled) - .toArray(PolicyEntity[]::new); + .map(MatchState::toMatch) + .filter(match -> match.policy().enabled()) + .toArray(Match[]::new); } private static boolean matches(PolicyAssociationSelector selector, TagAssignment assignment) { @@ -141,27 +170,94 @@ public class ObjectPolicyResolver { "Unsupported policy association selector type: " + selector.type()); } + /** + * An enabled policy matched for a metadata object, together with the metadata-object sources that + * contributed matching effective tags. + */ + public static final class Match { + + private final PolicyEntity policy; + private final List<MetadataObject> sources; + + /** + * Creates a policy match. + * + * @param policy The matched enabled policy (inherited flag already applied). + * @param sources The contributing metadata objects. + */ + public Match(PolicyEntity policy, List<MetadataObject> sources) { + this.policy = Objects.requireNonNull(policy, "policy"); + this.sources = + Collections.unmodifiableList(new ArrayList<>(Objects.requireNonNull(sources, "sources"))); + } + + /** + * Returns the matched policy. + * + * @return The matched policy. + */ + public PolicyEntity policy() { + return policy; + } + + /** + * Returns metadata objects whose effective tag assignments caused this policy to match. + * + * <p>The list order follows policy-tag relation enumeration. Callers should test membership. + * + * @return The contributing metadata objects. + */ + public List<MetadataObject> sources() { + return sources; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Match)) { + return false; + } + Match that = (Match) o; + return Objects.equals(policy, that.policy) && Objects.equals(sources, that.sources); + } + + @Override + public int hashCode() { + return Objects.hash(policy, sources); + } + + @Override + public String toString() { + return "Match{policy=" + policy + ", sources=" + sources + '}'; + } + } + private static final class MatchState { private final PolicyEntity policy; - private boolean matched; + private final Set<MetadataObject> sources = new LinkedHashSet<>(); private boolean directMatch; private MatchState(PolicyEntity policy) { this.policy = policy; } - private void record(boolean matches, boolean inherited) { - matched |= matches; - directMatch |= matches && !inherited; + private void record(boolean matches, boolean inherited, MetadataObject source) { + if (!matches) { + return; + } + sources.add(source); + directMatch |= !inherited; } private boolean matched() { - return matched; + return !sources.isEmpty(); } - private PolicyEntity policy() { - return policy.copyWithInherited(!directMatch); + private Match toMatch() { + return new Match(policy.copyWithInherited(!directMatch), new ArrayList<>(sources)); } } } diff --git a/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java b/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java index b446177add..65b4dbe0db 100644 --- a/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java +++ b/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import org.apache.gravitino.Entity; import org.apache.gravitino.EntityStore; import org.apache.gravitino.MetadataObject; @@ -65,12 +66,28 @@ public class EffectiveTagResolver { * @return Effective tags in deterministic nearest-assignment order. */ public TagEntity[] resolve(String metalake, MetadataObject metadataObject) { + ResolvedTag[] resolved = resolveWithSources(metalake, metadataObject); + TagEntity[] tags = new TagEntity[resolved.length]; + for (int i = 0; i < resolved.length; i++) { + tags[i] = resolved[i].tag(); + } + return tags; + } + + /** + * Resolves effective tags and the metadata object that contributed each assignment. + * + * @param metalake The metalake name. + * @param metadataObject The metadata object. + * @return Effective tags with sources in deterministic nearest-assignment order. + */ + public ResolvedTag[] resolveWithSources(String metalake, MetadataObject metadataObject) { MetadataObjectUtil.checkMetadataObject(metalake, metadataObject); List<MetadataObject> resolutionOrder = new ArrayList<>(); resolutionOrder.add(metadataObject); resolutionOrder.addAll(MetadataObjectUtil.getParentMetadataObjects(metadataObject)); - Map<String, TagEntity> effectiveTags = new LinkedHashMap<>(); + Map<String, ResolvedTag> effectiveTags = new LinkedHashMap<>(); for (int index = 0; index < resolutionOrder.size(); index++) { MetadataObject object = resolutionOrder.get(index); boolean inherited = index > 0; @@ -84,8 +101,10 @@ public class EffectiveTagResolver { SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL, identifier, entityType); - tags.forEach( - tag -> effectiveTags.putIfAbsent(tag.name(), tag.copyWithInherited(inherited))); + for (TagEntity tag : tags) { + effectiveTags.putIfAbsent( + tag.name(), new ResolvedTag(tag.copyWithInherited(inherited), object)); + } } catch (NoSuchEntityException e) { throw new NoSuchMetadataObjectException( e, "Failed to resolve effective tags for metadata object %s due to not found", object); @@ -94,6 +113,68 @@ public class EffectiveTagResolver { throw new RuntimeException(e); } } - return effectiveTags.values().toArray(new TagEntity[0]); + return effectiveTags.values().toArray(new ResolvedTag[0]); + } + + /** + * An effective tag assignment together with the metadata object that contributed it. + * + * <p>The source is the nearest object in the inheritance chain that assigned the tag name. + */ + public static final class ResolvedTag { + + private final TagEntity tag; + private final MetadataObject source; + + /** + * Creates a resolved tag. + * + * @param tag The effective tag entity, including assignment and inherited flag. + * @param source The metadata object that contributed this assignment. + */ + public ResolvedTag(TagEntity tag, MetadataObject source) { + this.tag = Objects.requireNonNull(tag, "tag"); + this.source = Objects.requireNonNull(source, "source"); + } + + /** + * Returns the effective tag entity. + * + * @return The effective tag entity. + */ + public TagEntity tag() { + return tag; + } + + /** + * Returns the metadata object that contributed this assignment. + * + * @return The contributing metadata object. + */ + public MetadataObject source() { + return source; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ResolvedTag)) { + return false; + } + ResolvedTag that = (ResolvedTag) o; + return Objects.equals(tag, that.tag) && Objects.equals(source, that.source); + } + + @Override + public int hashCode() { + return Objects.hash(tag, source); + } + + @Override + public String toString() { + return "ResolvedTag{tag=" + tag + ", source=" + source + '}'; + } } } diff --git a/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java b/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java index d448b925d4..a787ca0532 100644 --- a/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java +++ b/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java @@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableSet; import java.time.Instant; import java.util.Arrays; import java.util.Collections; +import java.util.List; import org.apache.gravitino.Entity; import org.apache.gravitino.EntityStore; import org.apache.gravitino.MetadataObject; @@ -37,6 +38,7 @@ import org.apache.gravitino.meta.AuditInfo; import org.apache.gravitino.meta.PolicyEntity; import org.apache.gravitino.meta.TagEntity; import org.apache.gravitino.tag.EffectiveTagResolver; +import org.apache.gravitino.tag.EffectiveTagResolver.ResolvedTag; import org.apache.gravitino.tag.TagAssignment; import org.apache.gravitino.utils.NameIdentifierUtil; import org.apache.gravitino.utils.NamespaceUtil; @@ -49,6 +51,8 @@ public class TestObjectPolicyResolver { private static final String METALAKE = "metalake"; private static final MetadataObject OBJECT = MetadataObjects.of(Arrays.asList("catalog", "schema", "table"), MetadataObject.Type.TABLE); + private static final MetadataObject CATALOG = + MetadataObjects.of(Collections.singletonList("catalog"), MetadataObject.Type.CATALOG); private EntityStore entityStore; private SupportsRelationOperations relationOperations; @@ -71,8 +75,11 @@ public class TestObjectPolicyResolver { TagEntity classified = tag(2L, "classified", TagAssignment.noValue()).copyWithInherited(true); PolicyEntity selected = policy(10L, "selected", true); PolicyEntity disabled = policy(11L, "disabled", false); - when(effectiveTagResolver.resolve(METALAKE, OBJECT)) - .thenReturn(new TagEntity[] {domain, classified}); + when(effectiveTagResolver.resolveWithSources(METALAKE, OBJECT)) + .thenReturn( + new ResolvedTag[] { + new ResolvedTag(domain, OBJECT), new ResolvedTag(classified, CATALOG) + }); when(relationOperations.batchListEntitiesByRelation( SupportsRelationOperations.Type.POLICY_TAG_REL, Arrays.asList( @@ -89,13 +96,18 @@ public class TestObjectPolicyResolver { Assertions.assertArrayEquals(new PolicyEntity[] {selected}, policies); Assertions.assertFalse(policies[0].inherited().orElseThrow()); + + ObjectPolicyResolver.Match[] matches = resolver.resolveMatches(METALAKE, OBJECT); + Assertions.assertEquals(1, matches.length); + Assertions.assertEquals(List.of(OBJECT, CATALOG), matches[0].sources()); } @Test public void testResolveMissingSelectorAsAllValues() throws Exception { TagEntity domain = tag(1L, "domain", TagAssignment.ofValues("finance")).copyWithInherited(true); PolicyEntity policy = policy(10L, "policy", true); - when(effectiveTagResolver.resolve(METALAKE, OBJECT)).thenReturn(new TagEntity[] {domain}); + when(effectiveTagResolver.resolveWithSources(METALAKE, OBJECT)) + .thenReturn(new ResolvedTag[] {new ResolvedTag(domain, CATALOG)}); when(relationOperations.batchListEntitiesByRelation( SupportsRelationOperations.Type.POLICY_TAG_REL, Collections.singletonList(NameIdentifierUtil.ofTag(METALAKE, "domain")), @@ -112,13 +124,16 @@ public class TestObjectPolicyResolver { PolicyEntity[] policies = resolver.resolve(METALAKE, OBJECT); Assertions.assertArrayEquals(new PolicyEntity[] {policy}, policies); Assertions.assertTrue(policies[0].inherited().orElseThrow()); + Assertions.assertEquals( + List.of(CATALOG), resolver.resolveMatches(METALAKE, OBJECT)[0].sources()); } @Test public void testDropNonMatchingSelector() throws Exception { TagEntity domain = tag(1L, "domain", TagAssignment.ofValues("engineering")); PolicyEntity policy = policy(10L, "policy", true); - when(effectiveTagResolver.resolve(METALAKE, OBJECT)).thenReturn(new TagEntity[] {domain}); + when(effectiveTagResolver.resolveWithSources(METALAKE, OBJECT)) + .thenReturn(new ResolvedTag[] {new ResolvedTag(domain, OBJECT)}); when(relationOperations.batchListEntitiesByRelation( SupportsRelationOperations.Type.POLICY_TAG_REL, Collections.singletonList(NameIdentifierUtil.ofTag(METALAKE, "domain")), @@ -127,6 +142,7 @@ public class TestObjectPolicyResolver { Collections.singletonList(relation(domain, policy, TagValueSelector.of("finance")))); Assertions.assertEquals(0, resolver.resolve(METALAKE, OBJECT).length); + Assertions.assertEquals(0, resolver.resolveMatches(METALAKE, OBJECT).length); } @Test @@ -134,8 +150,11 @@ public class TestObjectPolicyResolver { TagEntity domain = tag(1L, "domain", TagAssignment.ofValues("finance")); TagEntity classified = tag(2L, "classified", TagAssignment.ofValues("public")); PolicyEntity policy = policy(10L, "policy", true); - when(effectiveTagResolver.resolve(METALAKE, OBJECT)) - .thenReturn(new TagEntity[] {domain, classified}); + when(effectiveTagResolver.resolveWithSources(METALAKE, OBJECT)) + .thenReturn( + new ResolvedTag[] { + new ResolvedTag(domain, OBJECT), new ResolvedTag(classified, OBJECT) + }); when(relationOperations.batchListEntitiesByRelation( SupportsRelationOperations.Type.POLICY_TAG_REL, Arrays.asList( @@ -150,6 +169,26 @@ public class TestObjectPolicyResolver { Assertions.assertArrayEquals(new PolicyEntity[] {policy}, resolver.resolve(METALAKE, OBJECT)); } + @Test + public void testChildTagValueOverrideDropsAncestorTagValuePolicy() throws Exception { + // Effective tags after override: table data_domain=risk (catalog finance is shadowed). + TagEntity domain = + tag(1L, "data_domain", TagAssignment.ofValues("risk")).copyWithInherited(false); + PolicyEntity financePolicy = policy(30L, "finance-policy", true); + when(effectiveTagResolver.resolveWithSources(METALAKE, OBJECT)) + .thenReturn(new ResolvedTag[] {new ResolvedTag(domain, OBJECT)}); + when(relationOperations.batchListEntitiesByRelation( + SupportsRelationOperations.Type.POLICY_TAG_REL, + Collections.singletonList(NameIdentifierUtil.ofTag(METALAKE, "data_domain")), + Entity.EntityType.TAG)) + .thenReturn( + Collections.singletonList( + relation(domain, financePolicy, TagValueSelector.of("finance")))); + + Assertions.assertEquals(0, resolver.resolve(METALAKE, OBJECT).length); + Assertions.assertEquals(0, resolver.resolveMatches(METALAKE, OBJECT).length); + } + private static RelationalEntity<PolicyEntity> relation( TagEntity tag, PolicyEntity policy, PolicyAssociationSelector selector) { return new RelationalEntity<>( diff --git a/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java b/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java index 7f633218ed..0736256ace 100644 --- a/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java +++ b/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java @@ -107,6 +107,35 @@ public class TestEffectiveTagResolver { Assertions.assertTrue(effectiveTags[2].inherited().orElseThrow()); } + @Test + void testResolveWithSourcesPreservesContributingObject() throws Exception { + TagEntity directDomain = tag(1L, "domain", TagAssignment.ofValues("risk")); + TagEntity catalogDomain = tag(2L, "domain", TagAssignment.ofValues("finance")); + when(relationOperations.listEntitiesByRelation( + SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL, + MetadataObjectUtil.toEntityIdent(METALAKE, TABLE), + Entity.EntityType.TABLE)) + .thenReturn(Collections.singletonList(directDomain)); + when(relationOperations.listEntitiesByRelation( + SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL, + MetadataObjectUtil.toEntityIdent(METALAKE, SCHEMA), + Entity.EntityType.SCHEMA)) + .thenReturn(Collections.emptyList()); + when(relationOperations.listEntitiesByRelation( + SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL, + MetadataObjectUtil.toEntityIdent(METALAKE, CATALOG), + Entity.EntityType.CATALOG)) + .thenReturn(Collections.singletonList(catalogDomain)); + + EffectiveTagResolver.ResolvedTag[] effectiveTags = resolver.resolveWithSources(METALAKE, TABLE); + + Assertions.assertEquals(1, effectiveTags.length); + Assertions.assertEquals(TABLE, effectiveTags[0].source()); + Assertions.assertArrayEquals( + new String[] {"risk"}, effectiveTags[0].tag().assignment().orElseThrow().values()); + Assertions.assertFalse(effectiveTags[0].tag().inherited().orElseThrow()); + } + @Test void testMissingObjectIsTranslated() throws Exception { NoSuchEntityException exception = new NoSuchEntityException("missing"); diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.java new file mode 100644 index 0000000000..6d36baf0f5 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.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.gravitino.maintenance.scheduler; + +/** + * Gate used by expand before inserting a {@code cron-job} row. + * + * <p>Production expand wires {@link org.apache.gravitino.maintenance.occupancy.OccupancyGate} and + * may also skip when a pending {@code cron-job} row already covers the same table/policy pair. + */ +@FunctionalInterface +public interface ExpandOccupancyGate { + + /** + * @param tableId table entity id + * @param policyId policy entity id + * @param minIntervalMs resolved cooldown for the policy type + * @param nowMs current epoch millis + * @return true when expand must not enqueue a cron-job for this pair + */ + boolean shouldSkip(long tableId, long policyId, long minIntervalMs, long nowMs); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/ExpandTableCandidate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/ExpandTableCandidate.java new file mode 100644 index 0000000000..3311aafb8c --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/ExpandTableCandidate.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.maintenance.scheduler.expand; + +import com.google.common.base.Preconditions; +import java.util.Objects; +import org.apache.gravitino.MetadataObject; + +/** + * One enqueueable table after nearest-policy and occupancy gating, with the metadata level where + * the winning policy was selected. + */ +public final class ExpandTableCandidate { + + private final long tableId; + private final MetadataObject.Type winnerLevel; + + /** + * @param tableId table entity id + * @param winnerLevel {@link MetadataObject.Type#TABLE}, {@link MetadataObject.Type#SCHEMA}, or + * {@link MetadataObject.Type#CATALOG} + */ + public ExpandTableCandidate(long tableId, MetadataObject.Type winnerLevel) { + Preconditions.checkArgument(tableId > 0, "tableId must be positive"); + Preconditions.checkNotNull(winnerLevel, "winnerLevel must not be null"); + this.tableId = tableId; + this.winnerLevel = winnerLevel; + } + + /** + * @return table entity id + */ + public long tableId() { + return tableId; + } + + /** + * @return level where the nearest winning policy was attached + */ + public MetadataObject.Type winnerLevel() { + return winnerLevel; + } + + /** + * @return true when snapshot-expiry may pack this table into a {@code batch:...} cron-job + */ + public boolean batchableForSnapshotExpiry() { + return winnerLevel == MetadataObject.Type.SCHEMA || winnerLevel == MetadataObject.Type.CATALOG; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ExpandTableCandidate)) { + return false; + } + ExpandTableCandidate that = (ExpandTableCandidate) o; + return tableId == that.tableId && winnerLevel == that.winnerLevel; + } + + @Override + public int hashCode() { + return Objects.hash(tableId, winnerLevel); + } + + @Override + public String toString() { + return "ExpandTableCandidate{tableId=" + tableId + ", winnerLevel=" + winnerLevel + '}'; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestMaintenancePolicyResolver.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestMaintenancePolicyResolver.java new file mode 100644 index 0000000000..34fa73a2c1 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestMaintenancePolicyResolver.java @@ -0,0 +1,184 @@ +/* + * 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.maintenance.scheduler.expand; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.TreeSet; +import javax.annotation.Nullable; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.policy.ObjectPolicyResolver; +import org.apache.gravitino.policy.ObjectPolicyResolver.Match; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.utils.MetadataObjectUtil; + +/** + * Picks the nearest maintenance policy of a given {@link Policy.BuiltInType} for a table. + * + * <p>Matching and effective-tag inheritance come from {@link ObjectPolicyResolver}. This class only + * applies maintenance-specific nearest-level and same-level conflict rules over the returned + * metadata-object sources. + */ +public class NearestMaintenancePolicyResolver { + + private final ObjectPolicyResolver objectPolicyResolver; + + /** + * @param entityStore relational entity store + */ + public NearestMaintenancePolicyResolver(EntityStore entityStore) { + this(new ObjectPolicyResolver(entityStore)); + } + + /** + * @param objectPolicyResolver shared object-policy resolver + */ + NearestMaintenancePolicyResolver(ObjectPolicyResolver objectPolicyResolver) { + this.objectPolicyResolver = objectPolicyResolver; + } + + /** + * @param metalake metalake name + * @param table table metadata object + * @param policyType maintenance built-in type + * @return winner, none, or same-level conflict + */ + public Decision resolve(String metalake, MetadataObject table, Policy.BuiltInType policyType) { + Match[] matches = objectPolicyResolver.resolveMatches(metalake, table); + List<MetadataObject> chain = new ArrayList<>(); + chain.add(table); + chain.addAll(MetadataObjectUtil.getParentMetadataObjects(table)); + for (MetadataObject level : chain) { + TreeSet<Long> atLevel = new TreeSet<>(); + for (Match match : matches) { + if (match.policy().policyType() != policyType) { + continue; + } + if (match.sources().contains(level)) { + atLevel.add(match.policy().id()); + } + } + if (atLevel.isEmpty()) { + continue; + } + if (atLevel.size() > 1) { + return Decision.conflict(level, new ArrayList<>(atLevel)); + } + return Decision.winner(atLevel.first(), level); + } + return Decision.none(); + } + + /** Result of nearest maintenance-policy resolution for one table and built-in type. */ + public static final class Decision { + + /** Outcome of nearest-policy resolution. */ + public enum Kind { + /** Exactly one matching policy on the nearest non-empty level. */ + WINNER, + /** No matching policy on the inheritance chain. */ + NONE, + /** Two or more same-type policies on the nearest non-empty level. */ + CONFLICT + } + + private final Kind kind; + @Nullable private final Long winnerPolicyId; + @Nullable private final MetadataObject winnerLevel; + @Nullable private final MetadataObject conflictLevel; + private final List<Long> conflictingPolicyIds; + + private Decision( + Kind kind, + @Nullable Long winnerPolicyId, + @Nullable MetadataObject winnerLevel, + @Nullable MetadataObject conflictLevel, + List<Long> conflictingPolicyIds) { + this.kind = kind; + this.winnerPolicyId = winnerPolicyId; + this.winnerLevel = winnerLevel; + this.conflictLevel = conflictLevel; + this.conflictingPolicyIds = conflictingPolicyIds; + } + + /** + * @param policyId winning policy id + * @param level metadata level that produced the winner + * @return winner decision + */ + public static Decision winner(long policyId, MetadataObject level) { + return new Decision(Kind.WINNER, policyId, level, null, List.of()); + } + + /** + * @return none decision + */ + public static Decision none() { + return new Decision(Kind.NONE, null, null, null, List.of()); + } + + /** + * @param level metadata level that produced the conflict + * @param policyIds conflicting policy ids + * @return conflict decision + */ + public static Decision conflict(MetadataObject level, List<Long> policyIds) { + return new Decision(Kind.CONFLICT, null, null, level, List.copyOf(policyIds)); + } + + /** + * @return decision kind + */ + public Kind kind() { + return kind; + } + + /** + * @return winning policy id when {@link Kind#WINNER} + */ + public Optional<Long> winnerPolicyId() { + return Optional.ofNullable(winnerPolicyId); + } + + /** + * @return metadata level that produced the winner, or {@code null} + */ + @Nullable + public MetadataObject winnerLevel() { + return winnerLevel; + } + + /** + * @return metadata level that produced the conflict, or {@code null} + */ + @Nullable + public MetadataObject conflictLevel() { + return conflictLevel; + } + + /** + * @return conflicting policy ids when {@link Kind#CONFLICT} + */ + public List<Long> conflictingPolicyIds() { + return conflictingPolicyIds; + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestPolicyTableEnumerator.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestPolicyTableEnumerator.java new file mode 100644 index 0000000000..d182f3c56c --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/NearestPolicyTableEnumerator.java @@ -0,0 +1,311 @@ +/* + * 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.maintenance.scheduler.expand; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; +import org.apache.gravitino.Entity; +import org.apache.gravitino.GravitinoEnv; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.maintenance.scheduler.ExpandOccupancyGate; +import org.apache.gravitino.meta.GenericEntity; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.SchemaEntity; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.storage.relational.service.PolicyTagRelService; +import org.apache.gravitino.storage.relational.service.SchemaMetaService; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.apache.gravitino.storage.relational.service.TagMetaService; +import org.apache.gravitino.utils.HierarchicalSchemaUtil; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.apache.gravitino.utils.NamespaceUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Lists tables for which {@code policy} is the nearest maintenance policy of its type. + * + * <p>Walks {@code policy → tags → metadata objects}, expands catalog/schema attachments to tables + * (including tables in hierarchical descendant schemas), then keeps a table only when {@link + * NearestMaintenancePolicyResolver} selects this policy id and {@link ExpandOccupancyGate} allows + * enqueue. Scanning continues until {@code limit} enqueueable tables are found (cooldown / + * in-flight winners do not consume the batch budget). + */ +public final class NearestPolicyTableEnumerator implements PolicyTableEnumerator { + + private static final Logger LOG = LoggerFactory.getLogger(NearestPolicyTableEnumerator.class); + + private final PolicyTagRelService policyTagRelService; + private final TagMetaService tagMetaService; + private final SchemaMetaService schemaMetaService; + private final TableMetaService tableMetaService; + private final Supplier<NearestMaintenancePolicyResolver> nearestResolverSupplier; + + /** Uses relational singletons and {@link GravitinoEnv#entityStore()}. */ + public NearestPolicyTableEnumerator() { + this( + PolicyTagRelService.getInstance(), + TagMetaService.getInstance(), + SchemaMetaService.getInstance(), + TableMetaService.getInstance(), + () -> new NearestMaintenancePolicyResolver(GravitinoEnv.getInstance().entityStore())); + } + + /** + * @param policyTagRelService policy-tag relations + * @param tagMetaService tag metadata + * @param schemaMetaService schema listing + * @param tableMetaService table listing + * @param nearestResolverSupplier nearest maintenance-policy resolver + */ + public NearestPolicyTableEnumerator( + PolicyTagRelService policyTagRelService, + TagMetaService tagMetaService, + SchemaMetaService schemaMetaService, + TableMetaService tableMetaService, + Supplier<NearestMaintenancePolicyResolver> nearestResolverSupplier) { + this.policyTagRelService = policyTagRelService; + this.tagMetaService = tagMetaService; + this.schemaMetaService = schemaMetaService; + this.tableMetaService = tableMetaService; + this.nearestResolverSupplier = nearestResolverSupplier; + } + + @Override + public List<ExpandTableCandidate> listTableIds( + PolicyEntity policy, int limit, ExpandOccupancyGate occupancyGate, long minIntervalMs) { + if (limit <= 0) { + return List.of(); + } + Policy.BuiltInType policyType = policy.policyType(); + String metalake = policy.namespace().level(0); + NameIdentifier policyIdent = NameIdentifierUtil.ofPolicy(metalake, policy.name()); + List<RelationalEntity<?>> tagEdges = + policyTagRelService.listRelations(List.of(policyIdent), Entity.EntityType.POLICY); + + NearestMaintenancePolicyResolver nearest = nearestResolverSupplier.get(); + Set<Long> seenTables = new LinkedHashSet<>(); + Map<Long, ExpandTableCandidate> enqueueable = new LinkedHashMap<>(); + long nowMs = System.currentTimeMillis(); + ScanContext ctx = + new ScanContext( + metalake, + policy.id(), + policyType, + nearest, + occupancyGate, + minIntervalMs, + nowMs, + seenTables, + enqueueable, + limit); + + for (RelationalEntity<?> edge : tagEdges) { + Object target = edge.targetEntity(); + if (!(target instanceof TagEntity)) { + continue; + } + TagEntity tag = (TagEntity) target; + List<GenericEntity> objects; + try { + objects = tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier()); + } catch (IOException e) { + throw new UncheckedIOException("Failed to list objects for tag " + tag.nameIdentifier(), e); + } + for (GenericEntity object : objects) { + if (enqueueable.size() >= limit) { + return new ArrayList<>(enqueueable.values()); + } + collectCandidates(object, ctx); + } + } + return new ArrayList<>(enqueueable.values()); + } + + private void collectCandidates(GenericEntity object, ScanContext ctx) { + MetadataObject.Type objectType; + try { + objectType = MetadataObject.Type.valueOf(object.type().name()); + } catch (IllegalArgumentException e) { + return; + } + if (object.name() == null || object.name().isBlank()) { + return; + } + MetadataObject metadataObject; + try { + metadataObject = MetadataObjects.parse(object.name(), objectType); + } catch (RuntimeException e) { + return; + } + + switch (objectType) { + case TABLE: + if (object.id() == null) { + return; + } + considerTable(object.id(), metadataObject, ctx); + break; + case SCHEMA: + // A tag on a parent schema also covers tables in child schemas such as {@code a:b}. + expandSchema(metadataObject, ctx, true); + break; + case CATALOG: + expandCatalog(metadataObject, ctx); + break; + default: + break; + } + } + + private void expandCatalog(MetadataObject catalog, ScanContext ctx) { + List<SchemaEntity> schemas = + schemaMetaService.listSchemasByNamespace( + NamespaceUtil.ofSchema(ctx.metalake, catalog.name())); + for (SchemaEntity schema : schemas) { + if (ctx.enqueueable.size() >= ctx.limit) { + return; + } + MetadataObject schemaObject = + MetadataObjects.of(catalog.name(), schema.name(), MetadataObject.Type.SCHEMA); + // Catalog listing already includes every hierarchical schema name. + expandSchema(schemaObject, ctx, false); + } + } + + private void expandSchema(MetadataObject schema, ScanContext ctx, boolean includeDescendants) { + String catalogName = schema.parent(); + if (catalogName == null || schema.name() == null || schema.name().isBlank()) { + return; + } + expandTables(catalogName, schema.name(), ctx); + if (!includeDescendants || ctx.enqueueable.size() >= ctx.limit) { + return; + } + String separator = HierarchicalSchemaUtil.schemaSeparator(); + String descendantPrefix = schema.name() + separator; + List<SchemaEntity> schemas = + schemaMetaService.listSchemasByNamespace(NamespaceUtil.ofSchema(ctx.metalake, catalogName)); + for (SchemaEntity descendant : schemas) { + if (ctx.enqueueable.size() >= ctx.limit) { + return; + } + if (descendant.name() == null || !descendant.name().startsWith(descendantPrefix)) { + continue; + } + expandTables(catalogName, descendant.name(), ctx); + } + } + + private void expandTables(String catalogName, String schemaName, ScanContext ctx) { + List<TableEntity> tables = + tableMetaService.listTablesByNamespace( + NamespaceUtil.ofTable(ctx.metalake, catalogName, schemaName)); + for (TableEntity table : tables) { + if (ctx.enqueueable.size() >= ctx.limit) { + return; + } + MetadataObject tableObject = + NameIdentifierUtil.toMetadataObject(table.nameIdentifier(), Entity.EntityType.TABLE); + considerTable(table.id(), tableObject, ctx); + } + } + + private static void considerTable(long tableId, MetadataObject tableObject, ScanContext ctx) { + if (tableId <= 0 || !ctx.seenTables.add(tableId)) { + return; + } + NearestMaintenancePolicyResolver.Decision decision = + ctx.nearest.resolve(ctx.metalake, tableObject, ctx.policyType); + if (decision.kind() == NearestMaintenancePolicyResolver.Decision.Kind.CONFLICT) { + LOG.error( + "Skip table {} for maintenance type {}: ambiguous policies {} on nearest level {} ({})", + tableObject.fullName(), + ctx.policyType, + decision.conflictingPolicyIds(), + decision.conflictLevel() == null ? "?" : decision.conflictLevel().type(), + decision.conflictLevel() == null ? "?" : decision.conflictLevel().fullName()); + return; + } + if (decision.winnerPolicyId().orElse(-1L) != ctx.policyId) { + return; + } + MetadataObject winnerLevel = decision.winnerLevel(); + if (winnerLevel == null) { + return; + } + if (ctx.occupancyGate.shouldSkip(tableId, ctx.policyId, ctx.minIntervalMs, ctx.nowMs)) { + LOG.debug( + "Skip cron-job enqueue for table {} policy {} (in-flight or within minIntervalMs={})", + tableId, + ctx.policyId, + ctx.minIntervalMs); + return; + } + ctx.enqueueable.put(tableId, new ExpandTableCandidate(tableId, winnerLevel.type())); + } + + private static final class ScanContext { + private final String metalake; + private final long policyId; + private final Policy.BuiltInType policyType; + private final NearestMaintenancePolicyResolver nearest; + private final ExpandOccupancyGate occupancyGate; + private final long minIntervalMs; + private final long nowMs; + private final Set<Long> seenTables; + private final Map<Long, ExpandTableCandidate> enqueueable; + private final int limit; + + private ScanContext( + String metalake, + long policyId, + Policy.BuiltInType policyType, + NearestMaintenancePolicyResolver nearest, + ExpandOccupancyGate occupancyGate, + long minIntervalMs, + long nowMs, + Set<Long> seenTables, + Map<Long, ExpandTableCandidate> enqueueable, + int limit) { + this.metalake = metalake; + this.policyId = policyId; + this.policyType = policyType; + this.nearest = nearest; + this.occupancyGate = occupancyGate; + this.minIntervalMs = minIntervalMs; + this.nowMs = nowMs; + this.seenTables = seenTables; + this.enqueueable = enqueueable; + this.limit = limit; + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/PolicyTableEnumerator.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/PolicyTableEnumerator.java new file mode 100644 index 0000000000..8d65ee667a --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/expand/PolicyTableEnumerator.java @@ -0,0 +1,48 @@ +/* + * 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.maintenance.scheduler.expand; + +import java.util.List; +import org.apache.gravitino.maintenance.scheduler.ExpandOccupancyGate; +import org.apache.gravitino.meta.PolicyEntity; + +/** + * Lists enqueueable tables that a maintenance policy should expand into cron jobs. + * + * <p>Implementations expand tag-bound catalog/schema/table attachments and keep tables for which + * the policy is the nearest winner of its {@link org.apache.gravitino.policy.Policy.BuiltInType}. + * {@code limit} is the max number of <em>enqueueable</em> tables after occupancy / cooldown gating. + * Each candidate carries the winner attachment level for snapshot-expiry batching. + */ +@FunctionalInterface +public interface PolicyTableEnumerator { + + /** + * Lists up to {@code limit} tables that are nearest-policy winners and pass {@code + * occupancyGate}. + * + * @param policy loaded policy + * @param limit maximum enqueueable candidates to return + * @param occupancyGate skip in-flight / cooldown / already-queued tables while scanning + * @param minIntervalMs resolved cooldown for {@link ExpandOccupancyGate#shouldSkip} + * @return distinct enqueueable candidates, size at most {@code limit} + */ + List<ExpandTableCandidate> listTableIds( + PolicyEntity policy, int limit, ExpandOccupancyGate occupancyGate, long minIntervalMs); +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestMaintenancePolicyResolver.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestMaintenancePolicyResolver.java new file mode 100644 index 0000000000..09229b47c6 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestMaintenancePolicyResolver.java @@ -0,0 +1,163 @@ +/* + * 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.maintenance.scheduler.expand; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import java.util.List; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.ObjectPolicyResolver; +import org.apache.gravitino.policy.ObjectPolicyResolver.Match; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestNearestMaintenancePolicyResolver { + + private static final String METALAKE = "metalake"; + private static final MetadataObject TABLE = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "table"), MetadataObject.Type.TABLE); + private static final MetadataObject CATALOG = + MetadataObjects.of( + java.util.Collections.singletonList("catalog"), MetadataObject.Type.CATALOG); + + private ObjectPolicyResolver objectPolicyResolver; + private NearestMaintenancePolicyResolver resolver; + + @BeforeEach + void setUp() { + objectPolicyResolver = mock(ObjectPolicyResolver.class); + resolver = new NearestMaintenancePolicyResolver(objectPolicyResolver); + } + + @Test + void testTableLevelBeatsCatalogLevel() { + PolicyEntity tablePolicy = policy(10L, "table-compaction"); + PolicyEntity catalogPolicy = policy(20L, "catalog-compaction"); + when(objectPolicyResolver.resolveMatches(METALAKE, TABLE)) + .thenReturn( + new Match[] { + new Match(tablePolicy, List.of(TABLE)), new Match(catalogPolicy, List.of(CATALOG)) + }); + + NearestMaintenancePolicyResolver.Decision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestMaintenancePolicyResolver.Decision.Kind.WINNER, decision.kind()); + Assertions.assertEquals(java.util.Optional.of(10L), decision.winnerPolicyId()); + Assertions.assertEquals(TABLE, decision.winnerLevel()); + } + + @Test + void testFallsBackToAncestorWhenTableHasNoMatch() { + PolicyEntity catalogPolicy = policy(20L, "catalog-compaction"); + when(objectPolicyResolver.resolveMatches(METALAKE, TABLE)) + .thenReturn(new Match[] {new Match(catalogPolicy, List.of(CATALOG))}); + + NearestMaintenancePolicyResolver.Decision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestMaintenancePolicyResolver.Decision.Kind.WINNER, decision.kind()); + Assertions.assertEquals(java.util.Optional.of(20L), decision.winnerPolicyId()); + Assertions.assertEquals(CATALOG, decision.winnerLevel()); + } + + @Test + void testEmptyWhenNoMatchingType() { + when(objectPolicyResolver.resolveMatches(METALAKE, TABLE)).thenReturn(new Match[0]); + + NearestMaintenancePolicyResolver.Decision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestMaintenancePolicyResolver.Decision.Kind.NONE, decision.kind()); + } + + @Test + void testSameLevelConflict() { + PolicyEntity p1 = policy(10L, "compaction-a"); + PolicyEntity p2 = policy(11L, "compaction-b"); + when(objectPolicyResolver.resolveMatches(METALAKE, TABLE)) + .thenReturn(new Match[] {new Match(p1, List.of(TABLE)), new Match(p2, List.of(TABLE))}); + + NearestMaintenancePolicyResolver.Decision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals( + NearestMaintenancePolicyResolver.Decision.Kind.CONFLICT, decision.kind()); + Assertions.assertEquals(TABLE, decision.conflictLevel()); + Assertions.assertEquals(List.of(10L, 11L), decision.conflictingPolicyIds()); + Assertions.assertTrue(decision.winnerPolicyId().isEmpty()); + } + + @Test + void testIgnoresNonMatchingBuiltInType() { + PolicyEntity custom = + PolicyEntity.builder() + .withId(99L) + .withName("custom") + .withNamespace(NamespaceUtil.ofPolicy(METALAKE)) + .withPolicyType(Policy.BuiltInType.CUSTOM) + .withEnabled(true) + .withContent( + PolicyContents.custom( + java.util.Map.of("k", "v"), java.util.Set.of(MetadataObject.Type.TABLE), null)) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .build(); + when(objectPolicyResolver.resolveMatches(METALAKE, TABLE)) + .thenReturn(new Match[] {new Match(custom, List.of(TABLE))}); + + NearestMaintenancePolicyResolver.Decision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestMaintenancePolicyResolver.Decision.Kind.NONE, decision.kind()); + } + + private static PolicyEntity policy(long id, String name) { + return PolicyEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofPolicy(METALAKE)) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + "0 0 2 * * ?", + null, + null, + null, + null)) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestPolicyTableEnumerator.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestPolicyTableEnumerator.java new file mode 100644 index 0000000000..4a02e6faae --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/expand/TestNearestPolicyTableEnumerator.java @@ -0,0 +1,358 @@ +/* + * 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.maintenance.scheduler.expand; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import java.util.List; +import org.apache.gravitino.Entity; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.SupportsRelationOperations; +import org.apache.gravitino.maintenance.scheduler.ExpandOccupancyGate; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.GenericEntity; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.SchemaEntity; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.storage.relational.service.PolicyTagRelService; +import org.apache.gravitino.storage.relational.service.SchemaMetaService; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.apache.gravitino.storage.relational.service.TagMetaService; +import org.apache.gravitino.tag.TagAssignment; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestNearestPolicyTableEnumerator { + + private static final String METALAKE = "metalake"; + + private PolicyTagRelService policyTagRelService; + private TagMetaService tagMetaService; + private SchemaMetaService schemaMetaService; + private TableMetaService tableMetaService; + private NearestMaintenancePolicyResolver nearest; + private NearestPolicyTableEnumerator enumerator; + + @BeforeEach + void setUp() { + policyTagRelService = mock(PolicyTagRelService.class); + tagMetaService = mock(TagMetaService.class); + schemaMetaService = mock(SchemaMetaService.class); + tableMetaService = mock(TableMetaService.class); + nearest = mock(NearestMaintenancePolicyResolver.class); + enumerator = + new NearestPolicyTableEnumerator( + policyTagRelService, + tagMetaService, + schemaMetaService, + tableMetaService, + () -> nearest); + when(schemaMetaService.listSchemasByNamespace(any())).thenReturn(List.of()); + } + + @Test + void testDirectTableIncludedWhenNearest() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tag = tag(1L, "maint"); + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(100L) + .withName("catalog.schema.t1") + .withEntityType(Entity.EntityType.TABLE) + .build())); + + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + when(nearest.resolve(METALAKE, table, Policy.BuiltInType.ICEBERG_COMPACTION)) + .thenReturn(NearestMaintenancePolicyResolver.Decision.winner(10L, table)); + + Assertions.assertEquals( + List.of(100L), tableIds(enumerator.listTableIds(policy, 10, allowAll(), 0L))); + } + + @Test + void testSchemaExpansionSkipsOverriddenTable() throws Exception { + PolicyEntity catalogPolicy = policy(10L, "catalog-compaction"); + TagEntity catalogTag = tag(1L, "maint"); + + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "catalog-compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(catalogPolicy, catalogTag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(catalogTag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(1L) + .withName("catalog.schema") + .withEntityType(Entity.EntityType.SCHEMA) + .build())); + + TableEntity t1 = + TableEntity.builder() + .withId(100L) + .withName("t1") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t2 = + TableEntity.builder() + .withId(101L) + .withName("t2") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")))) + .thenReturn(List.of(t1, t2)); + + MetadataObject table1 = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + MetadataObject table2 = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t2"), MetadataObject.Type.TABLE); + MetadataObject schema = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema"), MetadataObject.Type.SCHEMA); + + when(nearest.resolve(METALAKE, table1, Policy.BuiltInType.ICEBERG_COMPACTION)) + .thenReturn(NearestMaintenancePolicyResolver.Decision.winner(11L, table1)); + when(nearest.resolve(METALAKE, table2, Policy.BuiltInType.ICEBERG_COMPACTION)) + .thenReturn(NearestMaintenancePolicyResolver.Decision.winner(10L, schema)); + + Assertions.assertEquals( + List.of(101L), tableIds(enumerator.listTableIds(catalogPolicy, 10, allowAll(), 0L))); + } + + @Test + void testSchemaExpansionIncludesHierarchicalDescendantTables() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tag = tag(1L, "maint"); + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(1L) + .withName("catalog.a") + .withEntityType(Entity.EntityType.SCHEMA) + .build())); + + when(schemaMetaService.listSchemasByNamespace(eq(NamespaceUtil.ofSchema(METALAKE, "catalog")))) + .thenReturn( + List.of(schema(2L, "a"), schema(3L, "a:b"), schema(4L, "ab"), schema(5L, "other"))); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "a")))) + .thenReturn(List.of()); + TableEntity childTable = + TableEntity.builder() + .withId(200L) + .withName("t1") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "a:b")) + .withAuditInfo(audit()) + .build(); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "a:b")))) + .thenReturn(List.of(childTable)); + + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "a:b", "t1"), MetadataObject.Type.TABLE); + MetadataObject parentSchema = + MetadataObjects.of(java.util.Arrays.asList("catalog", "a"), MetadataObject.Type.SCHEMA); + when(nearest.resolve(METALAKE, table, Policy.BuiltInType.ICEBERG_COMPACTION)) + .thenReturn(NearestMaintenancePolicyResolver.Decision.winner(10L, parentSchema)); + + Assertions.assertEquals( + List.of(200L), tableIds(enumerator.listTableIds(policy, 10, allowAll(), 0L))); + } + + @Test + void testFillsBatchSkippingCooldownWinners() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tag = tag(1L, "maint"); + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(1L) + .withName("catalog.schema") + .withEntityType(Entity.EntityType.SCHEMA) + .build())); + + TableEntity t1 = + TableEntity.builder() + .withId(100L) + .withName("t1") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t2 = + TableEntity.builder() + .withId(101L) + .withName("t2") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t3 = + TableEntity.builder() + .withId(102L) + .withName("t3") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")))) + .thenReturn(List.of(t1, t2, t3)); + + when(nearest.resolve( + eq(METALAKE), any(MetadataObject.class), eq(Policy.BuiltInType.ICEBERG_COMPACTION))) + .thenAnswer( + invocation -> { + MetadataObject table = invocation.getArgument(1); + return NearestMaintenancePolicyResolver.Decision.winner(10L, table); + }); + + ExpandOccupancyGate gate = (tableId, policyId, minIntervalMs, nowMs) -> tableId != 102L; + Assertions.assertEquals( + List.of(102L), tableIds(enumerator.listTableIds(policy, 1, gate, 60_000L))); + } + + @Test + void testSameLevelConflictSkipsTable() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tagA = tag(1L, "t-a"); + + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tagA))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tagA.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(100L) + .withName("catalog.schema.t1") + .withEntityType(Entity.EntityType.TABLE) + .build())); + + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + when(nearest.resolve(METALAKE, table, Policy.BuiltInType.ICEBERG_COMPACTION)) + .thenReturn(NearestMaintenancePolicyResolver.Decision.conflict(table, List.of(10L, 11L))); + + Assertions.assertEquals( + List.of(), tableIds(enumerator.listTableIds(policy, 10, allowAll(), 0L))); + } + + private static List<Long> tableIds(List<ExpandTableCandidate> candidates) { + return candidates.stream().map(ExpandTableCandidate::tableId).toList(); + } + + private static ExpandOccupancyGate allowAll() { + return (tableId, policyId, minIntervalMs, nowMs) -> false; + } + + private static RelationalEntity<TagEntity> policyTagEdge(PolicyEntity policy, TagEntity tag) { + return new RelationalEntity<>( + SupportsRelationOperations.Type.POLICY_TAG_REL, + policy.nameIdentifier(), + Entity.EntityType.POLICY, + tag, + null); + } + + private static TagEntity tag(long id, String name) { + return TagEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofTag(METALAKE)) + .withAssignment(TagAssignment.noValue()) + .withAuditInfo(audit()) + .build(); + } + + private static PolicyEntity policy(long id, String name) { + return PolicyEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofPolicy(METALAKE)) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + "0 0 2 * * ?", + null, + null, + null, + null)) + .withAuditInfo(audit()) + .build(); + } + + private static SchemaEntity schema(long id, String name) { + return SchemaEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofSchema(METALAKE, "catalog")) + .withAuditInfo(audit()) + .build(); + } + + private static AuditInfo audit() { + return AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build(); + } +}
