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
The following commit(s) were added to refs/heads/main by this push:
new d0567913ee [#14060] feat(core): Batch policy resolution for metadata
objects (#14194)
d0567913ee is described below
commit d0567913eeb0c8dc608e5d159c46d187162ef3ad
Author: roryqi <[email protected]>
AuthorDate: Wed Oct 28 18:26:14 2026 +0800
[#14060] feat(core): Batch policy resolution for metadata objects (#14194)
### What changes were proposed in this pull request?
- Add a core API to resolve effective policies for multiple metadata
objects in one call, sharing ancestor tag reads and policy association
reads.
- Batch column ID and `TAG_METADATA_OBJECT_REL` lookups in the
relational store.
- Resolve an Iceberg table and its columns together when loading read
restrictions.
### Why are the changes needed?
Wide Iceberg tables currently repeat tag and policy lookups for every
column. This makes storage queries grow with the column count.
Fix: #14060. Depends on #13778.
### Does this PR introduce _any_ user-facing change?
Adds `PolicyDispatcher.listPolicyInfosForMetadataObjects` to the core
API. No REST API or configuration changes.
### How was this patch tested?
- Added unit tests for shared ancestor resolution, selector matching,
policy deduplication, and policy events.
- Extended the H2 policy manager test to cover batched table and column
reads.
- Updated the experimental Iceberg loader test to verify one batch call.
- Ran the targeted core and Iceberg tests and Spotless checks.
---
.../gravitino/hook/PolicyHookDispatcher.java | 8 ++
.../gravitino/listener/PolicyEventDispatcher.java | 34 +++++
.../gravitino/policy/ObjectPolicyResolver.java | 99 +++++++++++---
.../apache/gravitino/policy/PolicyDispatcher.java | 18 +++
.../org/apache/gravitino/policy/PolicyManager.java | 13 ++
.../gravitino/storage/relational/JDBCBackend.java | 2 +
.../relational/mapper/TableColumnMapper.java | 13 ++
.../mapper/TableColumnSQLProviderFactory.java | 6 +
.../mapper/TagMetadataObjectRelMapper.java | 15 ++
.../TagMetadataObjectRelSQLProviderFactory.java | 8 ++
.../provider/base/TableColumnBaseSQLProvider.java | 16 +++
.../base/TagMetadataObjectRelBaseSQLProvider.java | 21 +++
.../storage/relational/po/TagRelationTagPO.java | 38 ++++++
.../relational/service/TableColumnMetaService.java | 24 ++++
.../storage/relational/service/TagMetaService.java | 71 ++++++++++
.../apache/gravitino/tag/EffectiveTagResolver.java | 103 ++++++++++++++
.../apache/gravitino/utils/MetadataObjectUtil.java | 38 ++++++
.../listener/api/event/TestPolicyEvent.java | 34 +++++
.../gravitino/policy/TestObjectPolicyResolver.java | 126 +++++++++++++++++
.../apache/gravitino/policy/TestPolicyManager.java | 21 +++
.../service/TestTableColumnMetaService.java | 109 +++++++++++++++
.../relational/service/TestTagMetaService.java | 152 +++++++++++++++++++++
.../gravitino/tag/TestEffectiveTagResolver.java | 77 +++++++++++
.../gravitino/utils/TestMetadataObjectUtil.java | 104 ++++++++++++++
.../ExperimentalIcebergReadRestrictionLoader.java | 17 ++-
...stExperimentalIcebergReadRestrictionLoader.java | 12 +-
26 files changed, 1152 insertions(+), 27 deletions(-)
diff --git
a/core/src/main/java/org/apache/gravitino/hook/PolicyHookDispatcher.java
b/core/src/main/java/org/apache/gravitino/hook/PolicyHookDispatcher.java
index 1d08d80f53..d853b842cf 100644
--- a/core/src/main/java/org/apache/gravitino/hook/PolicyHookDispatcher.java
+++ b/core/src/main/java/org/apache/gravitino/hook/PolicyHookDispatcher.java
@@ -18,6 +18,8 @@
package org.apache.gravitino.hook;
import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
import org.apache.gravitino.Entity;
import org.apache.gravitino.GravitinoEnv;
import org.apache.gravitino.MetadataObject;
@@ -123,4 +125,10 @@ public class PolicyHookDispatcher implements
PolicyDispatcher {
String metalake, MetadataObject metadataObject) {
return dispatcher.listPolicyInfosForMetadataObject(metalake,
metadataObject);
}
+
+ @Override
+ public Map<MetadataObject, PolicyEntity[]> listPolicyInfosForMetadataObjects(
+ String metalake, List<MetadataObject> metadataObjects) {
+ return dispatcher.listPolicyInfosForMetadataObjects(metalake,
metadataObjects);
+ }
}
diff --git
a/core/src/main/java/org/apache/gravitino/listener/PolicyEventDispatcher.java
b/core/src/main/java/org/apache/gravitino/listener/PolicyEventDispatcher.java
index 8869222622..f721e08105 100644
---
a/core/src/main/java/org/apache/gravitino/listener/PolicyEventDispatcher.java
+++
b/core/src/main/java/org/apache/gravitino/listener/PolicyEventDispatcher.java
@@ -18,6 +18,8 @@
*/
package org.apache.gravitino.listener;
+import java.util.List;
+import java.util.Map;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.RelationalEntity;
@@ -284,6 +286,38 @@ public class PolicyEventDispatcher implements
PolicyDispatcher {
}
}
+ @Override
+ public Map<MetadataObject, PolicyEntity[]> listPolicyInfosForMetadataObjects(
+ String metalake, List<MetadataObject> metadataObjects) {
+ metadataObjects.forEach(
+ object ->
+ eventBus.dispatchEvent(
+ new ListPolicyInfosForMetadataObjectPreEvent(
+ PrincipalUtils.getCurrentUserName(), metalake, object)));
+ try {
+ Map<MetadataObject, PolicyEntity[]> results =
+ dispatcher.listPolicyInfosForMetadataObjects(metalake,
metadataObjects);
+ metadataObjects.forEach(
+ object -> {
+ PolicyEntity[] policies = results.get(object);
+ eventBus.dispatchEvent(
+ new ListPolicyInfosForMetadataObjectEvent(
+ PrincipalUtils.getCurrentUserName(),
+ metalake,
+ object,
+ policies != null ? policies.length : -1));
+ });
+ return results;
+ } catch (Exception e) {
+ metadataObjects.forEach(
+ object ->
+ eventBus.dispatchEvent(
+ new ListPolicyInfosForMetadataObjectFailureEvent(
+ PrincipalUtils.getCurrentUserName(), metalake, object,
e)));
+ throw e;
+ }
+ }
+
private PolicyInfo toPolicyInfo(PolicyEntity policy) {
return new PolicyInfo(
policy.name(),
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 1a244fd068..73b22b013a 100644
--- a/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java
+++ b/core/src/main/java/org/apache/gravitino/policy/ObjectPolicyResolver.java
@@ -80,12 +80,7 @@ public class ObjectPolicyResolver {
* @return Enabled policies selected by the object's effective tags.
*/
public PolicyEntity[] resolve(String metalake, MetadataObject
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;
+ return policies(resolveMatches(metalake, metadataObject));
}
/**
@@ -104,18 +99,73 @@ public class ObjectPolicyResolver {
return new Match[0];
}
- Map<String, ResolvedTag> tagsByName =
- Arrays.stream(effectiveTags)
- .collect(
- Collectors.toMap(
- resolvedTag -> resolvedTag.tag().name(),
- resolvedTag -> resolvedTag,
- (left, right) -> left,
- LinkedHashMap::new));
List<NameIdentifier> tagIdentifiers =
- tagsByName.keySet().stream()
+ Arrays.stream(effectiveTags)
+ .map(resolvedTag -> resolvedTag.tag().name())
+ .distinct()
.map(tagName -> NameIdentifierUtil.ofTag(metalake, tagName))
.collect(Collectors.toList());
+ return selectMatches(effectiveTags, loadRelations(tagIdentifiers,
metadataObject));
+ }
+
+ /**
+ * Resolves effective policies for multiple objects using shared tag and
policy reads.
+ *
+ * @param metalake the metalake name
+ * @param metadataObjects the objects to resolve
+ * @return policies keyed by requested object
+ */
+ public Map<MetadataObject, PolicyEntity[]> resolve(
+ String metalake, List<MetadataObject> metadataObjects) {
+ if (metadataObjects.isEmpty()) {
+ return new LinkedHashMap<>();
+ }
+ Map<MetadataObject, ResolvedTag[]> tagsByObject =
+ effectiveTagResolver.resolveWithSources(metalake, metadataObjects);
+ Map<String, NameIdentifier> tagIdentifiers = new LinkedHashMap<>();
+ tagsByObject
+ .values()
+ .forEach(
+ tags -> {
+ for (ResolvedTag resolvedTag : tags) {
+ TagEntity tag = resolvedTag.tag();
+ tagIdentifiers.putIfAbsent(
+ tag.name(), NameIdentifierUtil.ofTag(metalake,
tag.name()));
+ }
+ });
+ List<RelationalEntity<?>> relations =
+ tagIdentifiers.isEmpty()
+ ? Collections.emptyList()
+ : loadRelations(new ArrayList<>(tagIdentifiers.values()),
metadataObjects);
+ Map<String, List<Integer>> relationPositionsByTag = new LinkedHashMap<>();
+ for (int i = 0; i < relations.size(); i++) {
+ relationPositionsByTag
+ .computeIfAbsent(relations.get(i).source().name(), ignored -> new
ArrayList<>())
+ .add(i);
+ }
+ Map<MetadataObject, PolicyEntity[]> result = new LinkedHashMap<>();
+ tagsByObject.forEach(
+ (object, tags) -> {
+ // Retain relation order while visiting only relations for this
object's effective tags.
+ List<RelationalEntity<?>> objectRelations =
+ Arrays.stream(tags)
+ .map(resolvedTag -> resolvedTag.tag().name())
+ .distinct()
+ .flatMap(
+ name ->
+ relationPositionsByTag
+ .getOrDefault(name, Collections.emptyList())
+ .stream())
+ .sorted()
+ .map(relations::get)
+ .collect(Collectors.toList());
+ result.put(object, policies(selectMatches(tags, objectRelations)));
+ });
+ return result;
+ }
+
+ private List<RelationalEntity<?>> loadRelations(
+ List<NameIdentifier> tagIdentifiers, Object request) {
List<RelationalEntity<?>> relations;
try {
@@ -127,9 +177,26 @@ public class ObjectPolicyResolver {
tagIdentifiers,
Entity.EntityType.TAG);
} catch (IOException e) {
- LOG.error("Failed to resolve policies for metadata object {}",
metadataObject, e);
+ LOG.error("Failed to resolve policies for metadata objects {}", request,
e);
throw new RuntimeException(e);
}
+ return relations;
+ }
+
+ private static PolicyEntity[] policies(Match[] matches) {
+ return
Arrays.stream(matches).map(Match::policy).toArray(PolicyEntity[]::new);
+ }
+
+ private static Match[] selectMatches(
+ ResolvedTag[] effectiveTags, List<RelationalEntity<?>> relations) {
+ Map<String, ResolvedTag> tagsByName =
+ Arrays.stream(effectiveTags)
+ .collect(
+ Collectors.toMap(
+ resolvedTag -> resolvedTag.tag().name(),
+ resolvedTag -> resolvedTag,
+ (left, right) -> left,
+ LinkedHashMap::new));
Map<Long, MatchState> matchStates = new LinkedHashMap<>();
for (RelationalEntity<?> relation : relations) {
diff --git
a/core/src/main/java/org/apache/gravitino/policy/PolicyDispatcher.java
b/core/src/main/java/org/apache/gravitino/policy/PolicyDispatcher.java
index 0f1ee9203a..199af5ed45 100644
--- a/core/src/main/java/org/apache/gravitino/policy/PolicyDispatcher.java
+++ b/core/src/main/java/org/apache/gravitino/policy/PolicyDispatcher.java
@@ -20,6 +20,9 @@
package org.apache.gravitino.policy;
import java.util.Arrays;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.RelationalEntity;
import org.apache.gravitino.annotation.Evolving;
@@ -163,4 +166,19 @@ public interface PolicyDispatcher {
* @return The array of policies associated with the specified metadata
object.
*/
PolicyEntity[] listPolicyInfosForMetadataObject(String metalake,
MetadataObject metadataObject);
+
+ /**
+ * Lists effective policies for several metadata objects in the same
metalake.
+ *
+ * @param metalake the metalake name
+ * @param metadataObjects the objects to resolve
+ * @return policies keyed by metadata object, including empty results
+ */
+ default Map<MetadataObject, PolicyEntity[]>
listPolicyInfosForMetadataObjects(
+ String metalake, List<MetadataObject> metadataObjects) {
+ Map<MetadataObject, PolicyEntity[]> result = new LinkedHashMap<>();
+ metadataObjects.forEach(
+ object -> result.put(object,
listPolicyInfosForMetadataObject(metalake, object)));
+ return result;
+ }
}
diff --git a/core/src/main/java/org/apache/gravitino/policy/PolicyManager.java
b/core/src/main/java/org/apache/gravitino/policy/PolicyManager.java
index 52d8d91c5c..55e623c9d7 100644
--- a/core/src/main/java/org/apache/gravitino/policy/PolicyManager.java
+++ b/core/src/main/java/org/apache/gravitino/policy/PolicyManager.java
@@ -27,6 +27,9 @@ import java.io.IOException;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
import org.apache.gravitino.Entity;
import org.apache.gravitino.EntityAlreadyExistsException;
import org.apache.gravitino.EntityStore;
@@ -233,6 +236,16 @@ public class PolicyManager implements PolicyDispatcher {
return objectPolicyResolver.resolve(metalake, metadataObject);
}
+ @Override
+ public Map<MetadataObject, PolicyEntity[]> listPolicyInfosForMetadataObjects(
+ String metalake, List<MetadataObject> metadataObjects) {
+ if (metadataObjects.isEmpty()) {
+ return new LinkedHashMap<>();
+ }
+ checkMetalake(NameIdentifier.of(metalake), entityStore);
+ return objectPolicyResolver.resolve(metalake, metadataObjects);
+ }
+
private PolicyEntity loadPolicyEntity(String metalake, String policyName) {
try {
return entityStore.get(
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/JDBCBackend.java
b/core/src/main/java/org/apache/gravitino/storage/relational/JDBCBackend.java
index 83b6c2409f..e25d804426 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/JDBCBackend.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/JDBCBackend.java
@@ -730,6 +730,8 @@ public class JDBCBackend implements RelationalBackend,
SupportsOrphanedRelationC
return OwnerMetaService.getInstance().batchGetOwner(nameIdentifiers,
identType);
case POLICY_TAG_REL:
return
PolicyTagRelService.getInstance().listRelations(nameIdentifiers, identType);
+ case TAG_METADATA_OBJECT_REL:
+ return
TagMetaService.getInstance().listTagsForMetadataObjects(nameIdentifiers,
identType);
default:
throw new IllegalArgumentException(
String.format("Doesn't support the relation type %s", relType));
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnMapper.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnMapper.java
index 906ad01c8a..3c4824c50c 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnMapper.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnMapper.java
@@ -73,6 +73,19 @@ public interface TableColumnMapper {
Long selectColumnIdByTableIdAndName(
@Param("tableId") Long tableId, @Param("columnName") String name);
+ /**
+ * Selects live column IDs for a table and a set of names.
+ *
+ * @param tableId the table ID
+ * @param columnNames the column names
+ * @return the matching column rows
+ */
+ @SelectProvider(
+ type = TableColumnSQLProviderFactory.class,
+ method = "selectColumnIdsByTableIdAndNames")
+ List<ColumnPO> selectColumnIdsByTableIdAndNames(
+ @Param("tableId") Long tableId, @Param("columnNames") List<String>
columnNames);
+
@SelectProvider(type = TableColumnSQLProviderFactory.class, method =
"selectColumnPOById")
ColumnPO selectColumnPOById(@Param("columnId") Long columnId);
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnSQLProviderFactory.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnSQLProviderFactory.java
index c00fe1554c..ed3fb3bf1a 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnSQLProviderFactory.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableColumnSQLProviderFactory.java
@@ -92,6 +92,12 @@ public class TableColumnSQLProviderFactory {
return getProvider().selectColumnIdByTableIdAndName(tableId, name);
}
+ /** Returns the SQL for a batched active-column lookup. */
+ public static String selectColumnIdsByTableIdAndNames(
+ @Param("tableId") Long tableId, @Param("columnNames") List<String>
columnNames) {
+ return getProvider().selectColumnIdsByTableIdAndNames(tableId,
columnNames);
+ }
+
public static String selectColumnPOById(@Param("columnId") Long columnId) {
return getProvider().selectColumnPOById(columnId);
}
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelMapper.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelMapper.java
index d375184121..50b4f68117 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelMapper.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelMapper.java
@@ -21,6 +21,7 @@ package org.apache.gravitino.storage.relational.mapper;
import java.util.List;
import org.apache.gravitino.storage.relational.po.TagMetadataObjectRelPO;
import org.apache.gravitino.storage.relational.po.TagPO;
+import org.apache.gravitino.storage.relational.po.TagRelationTagPO;
import org.apache.ibatis.annotations.DeleteProvider;
import org.apache.ibatis.annotations.InsertProvider;
import org.apache.ibatis.annotations.Param;
@@ -37,6 +38,20 @@ public interface TagMetadataObjectRelMapper {
@Param("metadataObjectId") Long metadataObjectId,
@Param("metadataObjectType") String metadataObjectType);
+ /**
+ * Selects tag assignment rows for multiple objects of one type.
+ *
+ * @param metadataObjectIds the source object IDs
+ * @param metadataObjectType the source object type
+ * @return the tag rows with source IDs
+ */
+ @SelectProvider(
+ type = TagMetadataObjectRelSQLProviderFactory.class,
+ method = "listTagPOsByMetadataObjectIdsAndType")
+ List<TagRelationTagPO> listTagPOsByMetadataObjectIdsAndType(
+ @Param("metadataObjectIds") List<Long> metadataObjectIds,
+ @Param("metadataObjectType") String metadataObjectType);
+
@SelectProvider(
type = TagMetadataObjectRelSQLProviderFactory.class,
method = "getTagPOsByMetadataObjectAndTagName")
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelSQLProviderFactory.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelSQLProviderFactory.java
index db6f12e938..d65bb86d66 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelSQLProviderFactory.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/TagMetadataObjectRelSQLProviderFactory.java
@@ -58,6 +58,14 @@ public class TagMetadataObjectRelSQLProviderFactory {
return getProvider().listTagPOsByMetadataObjectIdAndType(metadataObjectId,
metadataObjectType);
}
+ /** Returns the SQL for a batched tag assignment lookup. */
+ public static String listTagPOsByMetadataObjectIdsAndType(
+ @Param("metadataObjectIds") List<Long> metadataObjectIds,
+ @Param("metadataObjectType") String metadataObjectType) {
+ return getProvider()
+ .listTagPOsByMetadataObjectIdsAndType(metadataObjectIds,
metadataObjectType);
+ }
+
public static String getTagPOsByMetadataObjectAndTagName(
@Param("metadataObjectId") Long metadataObjectId,
@Param("metadataObjectType") String metadataObjectType,
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TableColumnBaseSQLProvider.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TableColumnBaseSQLProvider.java
index fa6e2fdd83..d1123f3081 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TableColumnBaseSQLProvider.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TableColumnBaseSQLProvider.java
@@ -141,6 +141,22 @@ public class TableColumnBaseSQLProvider {
+ " ORDER BY c.table_version DESC, c.id DESC LIMIT 1";
}
+ /** Returns the SQL for a batched active-column lookup. */
+ public String selectColumnIdsByTableIdAndNames(Long tableId, List<String>
columnNames) {
+ return "<script>SELECT c.column_id AS columnId, c.column_name AS
columnName FROM "
+ + TableColumnMapper.COLUMN_TABLE_NAME
+ + " c JOIN (SELECT column_id, MAX(table_version) AS max_version FROM "
+ + TableColumnMapper.COLUMN_TABLE_NAME
+ + " WHERE table_id = #{tableId} AND deleted_at = 0 GROUP BY column_id)
latest"
+ + " ON c.column_id = latest.column_id AND c.table_version =
latest.max_version"
+ + " WHERE c.table_id = #{tableId} AND c.column_name IN "
+ + "<foreach collection='columnNames' item='name' open='('
separator=',' close=')'>"
+ + "#{name}</foreach>"
+ + " AND c.deleted_at = 0 AND c.column_op_type <> "
+ + ColumnPO.ColumnOpType.DELETE.value()
+ + "</script>";
+ }
+
public String selectColumnPOById(@Param("columnId") Long columnId) {
return "SELECT column_id AS columnId, column_name AS columnName,"
+ " column_position AS columnPosition, metalake_id AS metalakeId,
catalog_id AS catalogId,"
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TagMetadataObjectRelBaseSQLProvider.java
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TagMetadataObjectRelBaseSQLProvider.java
index ea7d11a56c..c1893694d7 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TagMetadataObjectRelBaseSQLProvider.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/TagMetadataObjectRelBaseSQLProvider.java
@@ -58,6 +58,27 @@ public class TagMetadataObjectRelBaseSQLProvider {
+ " AND tm.deleted_at = 0";
}
+ /** Returns the SQL for a batched tag assignment lookup. */
+ public String listTagPOsByMetadataObjectIdsAndType(
+ List<Long> metadataObjectIds, String metadataObjectType) {
+ return "<script>SELECT te.metadata_object_id as metadataObjectId,"
+ + " tm.tag_id as tagId, tm.tag_name as tagName,"
+ + " tm.metalake_id as metalakeId, tm.tag_comment as comment,
tm.properties as properties,"
+ + " tm.allowed_values as allowedValues, NULLIF(te.tag_value, '') as
assignmentValue,"
+ + " tm.audit_info as auditInfo, tm.current_version as currentVersion,"
+ + " tm.last_version as lastVersion, tm.deleted_at as deletedAt"
+ + " FROM "
+ + TagMetaMapper.TAG_TABLE_NAME
+ + " tm JOIN "
+ + TagMetadataObjectRelMapper.TAG_METADATA_OBJECT_RELATION_TABLE_NAME
+ + " te ON tm.tag_id = te.tag_id"
+ + " WHERE te.metadata_object_id IN "
+ + "<foreach collection='metadataObjectIds' item='id' open='('
separator=',' close=')'>"
+ + "#{id}</foreach>"
+ + " AND te.metadata_object_type = #{metadataObjectType}"
+ + " AND te.deleted_at = 0 AND tm.deleted_at = 0</script>";
+ }
+
public String getTagPOsByMetadataObjectAndTagName(
@Param("metadataObjectId") Long metadataObjectId,
@Param("metadataObjectType") String metadataObjectType,
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/po/TagRelationTagPO.java
b/core/src/main/java/org/apache/gravitino/storage/relational/po/TagRelationTagPO.java
new file mode 100644
index 0000000000..b96cf77f20
--- /dev/null
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/po/TagRelationTagPO.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.storage.relational.po;
+
+/** A tag assignment row with the source metadata object ID. */
+public class TagRelationTagPO extends TagPO {
+ private Long metadataObjectId;
+
+ /**
+ * @return the source metadata object ID
+ */
+ public Long getMetadataObjectId() {
+ return metadataObjectId;
+ }
+
+ /**
+ * @param metadataObjectId the source metadata object ID
+ */
+ public void setMetadataObjectId(Long metadataObjectId) {
+ this.metadataObjectId = metadataObjectId;
+ }
+}
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/service/TableColumnMetaService.java
b/core/src/main/java/org/apache/gravitino/storage/relational/service/TableColumnMetaService.java
index fd64a42b80..467f537d0a 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/service/TableColumnMetaService.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/service/TableColumnMetaService.java
@@ -91,6 +91,30 @@ public class TableColumnMetaService {
return columnId;
}
+ /** Returns active column IDs by name using batched reads of one table. */
+ public Map<String, Long> getColumnIdsByTableIdAndNames(Long tableId,
List<String> columnNames) {
+ if (columnNames.isEmpty()) {
+ return new HashMap<>();
+ }
+ Map<String, Long> ids = new HashMap<>();
+ for (List<String> batch : Lists.partition(columnNames, 500)) {
+ List<ColumnPO> rows =
+ SessionUtils.getWithoutCommit(
+ TableColumnMapper.class,
+ mapper -> mapper.selectColumnIdsByTableIdAndNames(tableId,
batch));
+ rows.forEach(row -> ids.put(row.getColumnName(), row.getColumnId()));
+ }
+ for (String name : columnNames) {
+ if (!ids.containsKey(name)) {
+ throw new NoSuchEntityException(
+ NoSuchEntityException.NO_SUCH_ENTITY_MESSAGE,
+ Entity.EntityType.COLUMN.name().toLowerCase(Locale.ROOT),
+ name);
+ }
+ }
+ return ids;
+ }
+
@Monitored(
metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME,
baseMetricName = "getColumnPOById")
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/service/TagMetaService.java
b/core/src/main/java/org/apache/gravitino/storage/relational/service/TagMetaService.java
index fdfa3c4f20..1c84f7027e 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/service/TagMetaService.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/service/TagMetaService.java
@@ -46,6 +46,8 @@ import org.apache.gravitino.HasIdentifier;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.Namespace;
+import org.apache.gravitino.RelationalEntity;
+import org.apache.gravitino.SupportsRelationOperations;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.exceptions.NoSuchTagException;
import org.apache.gravitino.json.JsonUtils;
@@ -62,6 +64,7 @@ import
org.apache.gravitino.storage.relational.mapper.TagMetadataObjectRelMapper
import org.apache.gravitino.storage.relational.po.MetalakePO;
import org.apache.gravitino.storage.relational.po.TagMetadataObjectRelPO;
import org.apache.gravitino.storage.relational.po.TagPO;
+import org.apache.gravitino.storage.relational.po.TagRelationTagPO;
import org.apache.gravitino.storage.relational.utils.ExceptionUtils;
import org.apache.gravitino.storage.relational.utils.POConverters;
import org.apache.gravitino.storage.relational.utils.SessionUtils;
@@ -213,6 +216,74 @@ public class TagMetaService {
return tagPOsToTagEntities(tagPOs, NamespaceUtil.ofTag(metalake));
}
+ /**
+ * Lists direct tag assignments for metadata objects of one type in one
metalake.
+ *
+ * @param objectIdentifiers source identifiers
+ * @param objectType source entity type
+ * @return tag relation edges with their source identifiers
+ */
+ public List<RelationalEntity<?>> listTagsForMetadataObjects(
+ List<NameIdentifier> objectIdentifiers, Entity.EntityType objectType) {
+ if (objectIdentifiers.isEmpty()) {
+ return new ArrayList<>();
+ }
+ String metalake = objectIdentifiers.get(0).namespace().level(0);
+ Map<Long, NameIdentifier> identifiersById = new LinkedHashMap<>();
+ Map<NameIdentifier, List<NameIdentifier>> columnsByTable = new
LinkedHashMap<>();
+ for (NameIdentifier identifier : objectIdentifiers) {
+ Preconditions.checkArgument(
+ metalake.equals(identifier.namespace().level(0)),
+ "All metadata objects must belong to the same metalake");
+ if (objectType == Entity.EntityType.COLUMN) {
+ columnsByTable
+ .computeIfAbsent(
+ NameIdentifier.of(identifier.namespace().levels()), ignored ->
new ArrayList<>())
+ .add(identifier);
+ } else {
+ identifiersById.put(EntityIdService.getEntityId(identifier,
objectType), identifier);
+ }
+ }
+ columnsByTable.forEach(
+ (table, columns) -> {
+ Long tableId = EntityIdService.getEntityId(table,
Entity.EntityType.TABLE);
+ Map<String, Long> ids =
+ TableColumnMetaService.getInstance()
+ .getColumnIdsByTableIdAndNames(
+ tableId,
+
columns.stream().map(NameIdentifier::name).collect(Collectors.toList()));
+ columns.forEach(column ->
identifiersById.put(ids.get(column.name()), column));
+ });
+ List<TagRelationTagPO> rows = new ArrayList<>();
+ String metadataObjectType =
+ NameIdentifierUtil.toMetadataObject(objectIdentifiers.get(0),
objectType).type().toString();
+ for (List<Long> batch : Lists.partition(new
ArrayList<>(identifiersById.keySet()), 500)) {
+ rows.addAll(
+ SessionUtils.getWithoutCommit(
+ TagMetadataObjectRelMapper.class,
+ mapper -> mapper.listTagPOsByMetadataObjectIdsAndType(batch,
metadataObjectType)));
+ }
+ Map<Long, List<TagPO>> rowsByObject = new LinkedHashMap<>();
+ rows.forEach(
+ row ->
+ rowsByObject
+ .computeIfAbsent(row.getMetadataObjectId(), ignored -> new
ArrayList<>())
+ .add(row));
+ List<RelationalEntity<?>> result = new ArrayList<>();
+ rowsByObject.forEach(
+ (id, tagRows) ->
+ tagPOsToTagEntities(tagRows, NamespaceUtil.ofTag(metalake))
+ .forEach(
+ tag ->
+ result.add(
+ new RelationalEntity<>(
+
SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
+ identifiersById.get(id),
+ objectType,
+ tag))));
+ return result;
+ }
+
@Monitored(
metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME,
baseMetricName = "getTagForMetadataObject")
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 65b4dbe0db..83b5f4ca9d 100644
--- a/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java
+++ b/core/src/main/java/org/apache/gravitino/tag/EffectiveTagResolver.java
@@ -20,14 +20,20 @@ package org.apache.gravitino.tag;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.EnumMap;
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 org.apache.gravitino.Entity;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.RelationalEntity;
import org.apache.gravitino.SupportsRelationOperations;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.exceptions.NoSuchMetadataObjectException;
@@ -177,4 +183,101 @@ public class EffectiveTagResolver {
return "ResolvedTag{tag=" + tag + ", source=" + source + '}';
}
}
+
+ /**
+ * Resolves effective tags for multiple objects, reading each distinct
ancestor only once.
+ *
+ * @param metalake the metalake name
+ * @param metadataObjects the objects to resolve
+ * @return effective tags keyed by requested object
+ */
+ public Map<MetadataObject, TagEntity[]> resolve(
+ String metalake, List<MetadataObject> metadataObjects) {
+ Map<MetadataObject, TagEntity[]> result = new LinkedHashMap<>();
+ resolveWithSources(metalake, metadataObjects)
+ .forEach(
+ (object, tags) ->
+ result.put(
+ object,
Arrays.stream(tags).map(ResolvedTag::tag).toArray(TagEntity[]::new)));
+ return result;
+ }
+
+ /**
+ * Resolves effective tags with their sources, reading each distinct
ancestor only once.
+ *
+ * @param metalake the metalake name
+ * @param metadataObjects the objects to resolve
+ * @return effective tags with sources keyed by requested object
+ */
+ public Map<MetadataObject, ResolvedTag[]> resolveWithSources(
+ String metalake, List<MetadataObject> metadataObjects) {
+ if (metadataObjects.isEmpty()) {
+ return new LinkedHashMap<>();
+ }
+ MetadataObjectUtil.checkMetadataObjects(metalake, metadataObjects);
+ Map<MetadataObject, List<MetadataObject>> resolutionOrders = new
LinkedHashMap<>();
+ Map<Entity.EntityType, Set<MetadataObject>> objectsByType =
+ new EnumMap<>(Entity.EntityType.class);
+ for (MetadataObject requested : metadataObjects) {
+ List<MetadataObject> order = new ArrayList<>();
+ order.add(requested);
+ order.addAll(MetadataObjectUtil.getParentMetadataObjects(requested));
+ resolutionOrders.put(requested, order);
+ for (MetadataObject object : order) {
+ objectsByType
+ .computeIfAbsent(
+ MetadataObjectUtil.toEntityType(object), ignored -> new
LinkedHashSet<>())
+ .add(object);
+ }
+ }
+
+ Map<MetadataObject, List<TagEntity>> directTags = new LinkedHashMap<>();
+ for (Map.Entry<Entity.EntityType, Set<MetadataObject>> entry :
objectsByType.entrySet()) {
+ List<NameIdentifier> identifiers = new ArrayList<>();
+ Map<NameIdentifier, MetadataObject> objectsByIdentifier = new
LinkedHashMap<>();
+ for (MetadataObject object : entry.getValue()) {
+ NameIdentifier identifier = MetadataObjectUtil.toEntityIdent(metalake,
object);
+ identifiers.add(identifier);
+ objectsByIdentifier.put(identifier, object);
+ }
+ try {
+ for (RelationalEntity<?> relation :
+ entityStore
+ .relationOperations()
+ .batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
+ identifiers,
+ entry.getKey())) {
+ directTags
+ .computeIfAbsent(
+ objectsByIdentifier.get(relation.source()), ignored -> new
ArrayList<>())
+ .add((TagEntity) relation.targetEntity());
+ }
+ } catch (NoSuchEntityException e) {
+ throw new NoSuchMetadataObjectException(
+ e, "Failed to resolve effective tags for metadata objects due to
not found");
+ } catch (IOException e) {
+ LOG.error("Failed to resolve effective tags for metadata objects", e);
+ throw new RuntimeException(e);
+ }
+ }
+
+ Map<MetadataObject, ResolvedTag[]> result = new LinkedHashMap<>();
+ resolutionOrders.forEach(
+ (requested, order) -> {
+ Map<String, ResolvedTag> effectiveTags = new LinkedHashMap<>();
+ for (int index = 0; index < order.size(); index++) {
+ boolean inherited = index > 0;
+ MetadataObject source = order.get(index);
+ directTags
+ .getOrDefault(source, Collections.emptyList())
+ .forEach(
+ tag ->
+ effectiveTags.putIfAbsent(
+ tag.name(), new
ResolvedTag(tag.copyWithInherited(inherited), source)));
+ }
+ result.put(requested, effectiveTags.values().toArray(new
ResolvedTag[0]));
+ });
+ return result;
+ }
}
diff --git
a/core/src/main/java/org/apache/gravitino/utils/MetadataObjectUtil.java
b/core/src/main/java/org/apache/gravitino/utils/MetadataObjectUtil.java
index 47f056498e..70af3cf295 100644
--- a/core/src/main/java/org/apache/gravitino/utils/MetadataObjectUtil.java
+++ b/core/src/main/java/org/apache/gravitino/utils/MetadataObjectUtil.java
@@ -25,8 +25,10 @@ import com.google.common.base.Preconditions;
import com.google.common.collect.BiMap;
import com.google.common.collect.ImmutableBiMap;
import java.util.ArrayList;
+import java.util.HashSet;
import java.util.List;
import java.util.Optional;
+import java.util.Set;
import java.util.function.Supplier;
import org.apache.commons.lang3.StringUtils;
import org.apache.gravitino.Entity;
@@ -342,6 +344,42 @@ public class MetadataObjectUtil {
}
}
+ /**
+ * Checks metadata objects, checking each distinct table only once for
tables and their columns.
+ *
+ * <p>Every object's identifier is validated. Successful table checks are
shared only within this
+ * call. Like {@link #checkMetadataObject(String, MetadataObject)}, this
must be called outside
+ * the tree lock because checking an object may load it into the entity
store.
+ *
+ * @param metalake The metalake name.
+ * @param objects The metadata objects to check, in order.
+ * @throws NoSuchMetadataObjectException if a metadata object does not exist.
+ */
+ public static void checkMetadataObjects(String metalake,
List<MetadataObject> objects) {
+ Set<NameIdentifier> checkedTables = new HashSet<>();
+ for (MetadataObject object : objects) {
+ Preconditions.checkArgument(object != null, "metadataObject cannot be
null");
+ if (object.type() != MetadataObject.Type.TABLE
+ && object.type() != MetadataObject.Type.COLUMN) {
+ checkMetadataObject(metalake, object);
+ continue;
+ }
+
+ NameIdentifier identifier = toEntityIdent(metalake, object);
+ NameIdentifier tableIdentifier;
+ if (object.type() == MetadataObject.Type.COLUMN) {
+ NameIdentifierUtil.checkColumn(identifier);
+ tableIdentifier = NameIdentifier.of(identifier.namespace().levels());
+ } else {
+ NameIdentifierUtil.checkTable(identifier);
+ tableIdentifier = identifier;
+ }
+ if (checkedTables.add(tableIdentifier)) {
+ checkMetadataObject(metalake, object);
+ }
+ }
+ }
+
private static void check(
final boolean expression, Supplier<? extends RuntimeException>
exceptionToThrowSupplier) {
if (!expression) {
diff --git
a/core/src/test/java/org/apache/gravitino/listener/api/event/TestPolicyEvent.java
b/core/src/test/java/org/apache/gravitino/listener/api/event/TestPolicyEvent.java
index 49c9b05fe0..efc9248351 100644
---
a/core/src/test/java/org/apache/gravitino/listener/api/event/TestPolicyEvent.java
+++
b/core/src/test/java/org/apache/gravitino/listener/api/event/TestPolicyEvent.java
@@ -26,6 +26,8 @@ import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
import java.util.Objects;
import org.apache.gravitino.Entity;
import org.apache.gravitino.MetadataObject;
@@ -307,6 +309,38 @@ public class TestPolicyEvent {
Assertions.assertEquals(2, ((ListPolicyInfosForMetadataObjectEvent)
postevent).resultCount());
}
+ @Test
+ void testBatchListPolicyInfosEmitsEventsForEachObject() {
+ MetadataObject catalog =
+ NameIdentifierUtil.toMetadataObject(
+ NameIdentifierUtil.ofCatalog("metalake", "catalog_for_test"),
+ Entity.EntityType.CATALOG);
+ MetadataObject schema =
+ NameIdentifierUtil.toMetadataObject(
+ NameIdentifierUtil.ofSchema("metalake", "catalog_for_test",
"schema_for_test"),
+ Entity.EntityType.SCHEMA);
+ PolicyDispatcher delegate = mock(PolicyDispatcher.class);
+ when(delegate.listPolicyInfosForMetadataObjects("metalake",
List.of(catalog, schema)))
+ .thenReturn(Map.of(catalog, new PolicyEntity[] {policy}, schema, new
PolicyEntity[0]));
+ DummyEventListener listener = new DummyEventListener();
+ PolicyEventDispatcher batchDispatcher =
+ new PolicyEventDispatcher(new EventBus(List.of(listener)), delegate);
+
+ Map<MetadataObject, PolicyEntity[]> result =
+ batchDispatcher.listPolicyInfosForMetadataObjects("metalake",
List.of(catalog, schema));
+
+ Assertions.assertEquals(1, result.get(catalog).length);
+ Assertions.assertEquals(0, result.get(schema).length);
+ Assertions.assertEquals(
+ ListPolicyInfosForMetadataObjectPreEvent.class,
listener.popPreEvent().getClass());
+ Assertions.assertEquals(
+ ListPolicyInfosForMetadataObjectPreEvent.class,
listener.popPreEvent().getClass());
+ Assertions.assertEquals(
+ 0, ((ListPolicyInfosForMetadataObjectEvent)
listener.popPostEvent()).resultCount());
+ Assertions.assertEquals(
+ 1, ((ListPolicyInfosForMetadataObjectEvent)
listener.popPostEvent()).resultCount());
+ }
+
@Test
@SuppressWarnings("deprecation")
void
testListPolicyInfosForMetadataObjectEventDeprecatedConstructorReturnsNegativeCount()
{
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 a787ca0532..2520d7e40d 100644
---
a/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java
+++
b/core/src/test/java/org/apache/gravitino/policy/TestObjectPolicyResolver.java
@@ -19,18 +19,25 @@
package org.apache.gravitino.policy;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.time.Instant;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
+import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
import org.apache.gravitino.Entity;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.MetadataObjects;
+import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.RelationalEntity;
import org.apache.gravitino.SupportsRelationOperations;
import org.apache.gravitino.json.PolicyAssociationSelectorSerde;
@@ -189,6 +196,125 @@ public class TestObjectPolicyResolver {
Assertions.assertEquals(0, resolver.resolveMatches(METALAKE,
OBJECT).length);
}
+ @Test
+ public void testBatchSharesPolicyReadAndPreservesSelectors() throws
Exception {
+ MetadataObject column =
+ MetadataObjects.of(
+ Arrays.asList("catalog", "schema", "table", "column"),
MetadataObject.Type.COLUMN);
+ TagEntity tableTag = tag(1L, "domain", TagAssignment.ofValues("finance"));
+ TagEntity columnTag = tag(1L, "domain",
TagAssignment.ofValues("engineering"));
+ PolicyEntity policy = policy(10L, "selected", true);
+ Map<MetadataObject, ResolvedTag[]> tags = new LinkedHashMap<>();
+ tags.put(
+ OBJECT, new ResolvedTag[] {new
ResolvedTag(tableTag.copyWithInherited(false), OBJECT)});
+ tags.put(
+ column, new ResolvedTag[] {new
ResolvedTag(columnTag.copyWithInherited(true), CATALOG)});
+ when(effectiveTagResolver.resolveWithSources(METALAKE, List.of(OBJECT,
column)))
+ .thenReturn(tags);
+ when(relationOperations.batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.POLICY_TAG_REL,
+ List.of(NameIdentifierUtil.ofTag(METALAKE, "domain")),
+ Entity.EntityType.TAG))
+ .thenReturn(List.of(relation(tableTag, policy,
TagValueSelector.of("finance"))));
+
+ Map<MetadataObject, PolicyEntity[]> result =
+ resolver.resolve(METALAKE, List.of(OBJECT, column));
+
+ Assertions.assertEquals(1, result.get(OBJECT).length);
+ Assertions.assertFalse(result.get(OBJECT)[0].inherited().orElseThrow());
+ Assertions.assertEquals(0, result.get(column).length);
+ verify(relationOperations, times(1))
+ .batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.POLICY_TAG_REL,
+ List.of(NameIdentifierUtil.ofTag(METALAKE, "domain")),
+ Entity.EntityType.TAG);
+ }
+
+ /** Verifies distinct column tags do not cause every column to scan all
policy relations. */
+ @Test
+ public void testBatchVisitsOnlyRelationsForEffectiveTags() throws Exception {
+ List<MetadataObject> objects = new ArrayList<>();
+ Map<MetadataObject, ResolvedTag[]> tags = new LinkedHashMap<>();
+ List<NameIdentifier> tagIdentifiers = new ArrayList<>();
+ List<RelationalEntity<?>> relations = new ArrayList<>();
+ List<PolicyEntity> expectedPolicies = new ArrayList<>();
+ for (int i = 0; i < 1000; i++) {
+ MetadataObject column =
+ MetadataObjects.of(
+ List.of("catalog", "schema", "table", "column_" + i),
MetadataObject.Type.COLUMN);
+ TagEntity tag = tag(i, "tag_" + i,
TagAssignment.noValue()).copyWithInherited(false);
+ PolicyEntity policy = policy(1000L + i, "policy_" + i, true);
+ objects.add(column);
+ tags.put(column, new ResolvedTag[] {new ResolvedTag(tag, column)});
+ tagIdentifiers.add(tag.nameIdentifier());
+ relations.add(spy(relation(tag, policy, AllValuesSelector.get())));
+ expectedPolicies.add(policy);
+ }
+ MetadataObject untagged =
+ MetadataObjects.of(
+ List.of("catalog", "schema", "table", "untagged"),
MetadataObject.Type.COLUMN);
+ objects.add(untagged);
+ tags.put(untagged, new ResolvedTag[0]);
+ when(effectiveTagResolver.resolveWithSources(METALAKE,
objects)).thenReturn(tags);
+ when(relationOperations.batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.POLICY_TAG_REL, tagIdentifiers,
Entity.EntityType.TAG))
+ .thenReturn(relations);
+
+ Map<MetadataObject, PolicyEntity[]> result = resolver.resolve(METALAKE,
objects);
+
+ Assertions.assertEquals(objects.size(), result.size());
+ for (int i = 0; i < expectedPolicies.size(); i++) {
+ Assertions.assertArrayEquals(
+ new PolicyEntity[] {expectedPolicies.get(i)},
result.get(objects.get(i)));
+ // One visit to build the index and one visit for the column that owns
this tag.
+ verify(relations.get(i), times(2)).source();
+ }
+ Assertions.assertEquals(0, result.get(untagged).length);
+ verify(relationOperations, times(1))
+ .batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.POLICY_TAG_REL, tagIdentifiers,
Entity.EntityType.TAG);
+ }
+
+ /**
+ * Verifies indexed matching retains relation order, deduplication and
direct-match precedence.
+ */
+ @Test
+ public void testBatchIndexPreservesPolicyOrderAndInheritance() throws
Exception {
+ TagEntity direct = tag(1L, "direct", TagAssignment.ofValues("finance"));
+ TagEntity inherited = tag(2L, "inherited",
TagAssignment.noValue()).copyWithInherited(true);
+ TagEntity unrelated = tag(3L, "unrelated", TagAssignment.noValue());
+ PolicyEntity first = policy(10L, "first", true);
+ PolicyEntity second = policy(11L, "second", true);
+ PolicyEntity disabled = policy(12L, "disabled", false);
+ Map<MetadataObject, ResolvedTag[]> tags = new LinkedHashMap<>();
+ tags.put(
+ OBJECT,
+ new ResolvedTag[] {new ResolvedTag(direct, OBJECT), new
ResolvedTag(inherited, CATALOG)});
+ tags.put(CATALOG, new ResolvedTag[] {new ResolvedTag(unrelated, CATALOG)});
+ when(effectiveTagResolver.resolveWithSources(METALAKE, List.of(OBJECT,
CATALOG)))
+ .thenReturn(tags);
+ when(relationOperations.batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.POLICY_TAG_REL,
+ List.of(
+ direct.nameIdentifier(), inherited.nameIdentifier(),
unrelated.nameIdentifier()),
+ Entity.EntityType.TAG))
+ .thenReturn(
+ List.of(
+ relation(inherited, first, AllValuesSelector.get()),
+ relation(unrelated, second, AllValuesSelector.get()),
+ relation(direct, second, TagValueSelector.of("finance")),
+ relation(direct, first, TagValueSelector.of("finance")),
+ relation(direct, disabled, AllValuesSelector.get())));
+
+ Map<MetadataObject, PolicyEntity[]> result =
+ resolver.resolve(METALAKE, List.of(OBJECT, CATALOG));
+
+ Assertions.assertArrayEquals(new PolicyEntity[] {first, second},
result.get(OBJECT));
+ Assertions.assertFalse(result.get(OBJECT)[0].inherited().orElseThrow());
+ Assertions.assertFalse(result.get(OBJECT)[1].inherited().orElseThrow());
+ Assertions.assertArrayEquals(new PolicyEntity[] {second},
result.get(CATALOG));
+ }
+
private static RelationalEntity<PolicyEntity> relation(
TagEntity tag, PolicyEntity policy, PolicyAssociationSelector selector) {
return new RelationalEntity<>(
diff --git
a/core/src/test/java/org/apache/gravitino/policy/TestPolicyManager.java
b/core/src/test/java/org/apache/gravitino/policy/TestPolicyManager.java
index 51887d8b05..cf74141a2b 100644
--- a/core/src/test/java/org/apache/gravitino/policy/TestPolicyManager.java
+++ b/core/src/test/java/org/apache/gravitino/policy/TestPolicyManager.java
@@ -750,6 +750,27 @@ public class TestPolicyManager {
Assertions.assertEquals(
ImmutableSet.of(policy1, policy2, policy3),
ImmutableSet.copyOf(policiesInfo2));
+ MetadataObject columnObject =
+ NameIdentifierUtil.toMetadataObject(
+ NameIdentifierUtil.ofColumn(METALAKE, CATALOG, SCHEMA, TABLE,
COLUMN),
+ Entity.EntityType.COLUMN);
+ associatePolicyThroughTag(policy1, columnObject);
+ Map<MetadataObject, PolicyEntity[]> batched =
+ policyManager.listPolicyInfosForMetadataObjects(
+ METALAKE, List.of(tableObject, columnObject));
+ Assertions.assertEquals(3, batched.get(tableObject).length);
+ Assertions.assertEquals(3, batched.get(columnObject).length);
+ Assertions.assertFalse(
+ Arrays.stream(batched.get(columnObject))
+ .filter(policy -> policy.name().equals(policy1.name()))
+ .findFirst()
+ .orElseThrow()
+ .inherited()
+ .orElseThrow());
+ Assertions.assertEquals(
+
ImmutableSet.copyOf(policyManager.listPolicyInfosForMetadataObject(METALAKE,
columnObject)),
+ ImmutableSet.copyOf(batched.get(columnObject)));
+
// List policies for non-existent metadata object
MetadataObject nonExistentObject =
NameIdentifierUtil.toMetadataObject(
diff --git
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTableColumnMetaService.java
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTableColumnMetaService.java
index f53defa175..681aec9566 100644
---
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTableColumnMetaService.java
+++
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTableColumnMetaService.java
@@ -190,6 +190,115 @@ public class TestTableColumnMetaService extends
TestJDBCBackend {
compareTwoColumns(createdTable.columns(), retrievedTable.columns());
}
+ /** Verifies batched column ID reads across query boundaries on every
backend. */
+ @TestTemplate
+ public void testBatchGetColumnIdsForWideTable() throws IOException {
+ createParentEntities(METALAKE_NAME, "catalog1", "schema1", AUDIT_INFO);
+ List<ColumnEntity> columns = new ArrayList<>();
+ for (int i = 0; i < WIDE_TABLE_COLUMN_COUNT; i++) {
+ columns.add(newIntColumn("column_" + i, i, "comment"));
+ }
+ TableEntity table =
+ TableEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("wide_table")
+ .withNamespace(Namespace.of(METALAKE_NAME, "catalog1", "schema1"))
+ .withColumns(columns)
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ TableMetaService.getInstance().insertTable(table, false);
+ TableColumnMetaService service = TableColumnMetaService.getInstance();
+ for (int count : new int[] {0, 1, 499, 500, 501, 1001,
WIDE_TABLE_COLUMN_COUNT}) {
+ List<ColumnEntity> requested = columns.subList(0, count);
+ Map<String, Long> expected =
+ requested.stream().collect(Collectors.toMap(ColumnEntity::name,
ColumnEntity::id));
+ Assertions.assertEquals(
+ expected,
+ service.getColumnIdsByTableIdAndNames(
+ table.id(),
requested.stream().map(ColumnEntity::name).collect(Collectors.toList())));
+ }
+ List<String> names =
columns.stream().map(ColumnEntity::name).collect(Collectors.toList());
+ names.add(columns.get(0).name());
+ Assertions.assertEquals(
+ WIDE_TABLE_COLUMN_COUNT,
service.getColumnIdsByTableIdAndNames(table.id(), names).size());
+ names.add("missing_column");
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () -> service.getColumnIdsByTableIdAndNames(table.id(), names));
+ }
+
+ /** Verifies batch lookups follow renames, deletions and replacement IDs
within one table. */
+ @TestTemplate
+ public void testBatchGetColumnIdsFollowsColumnLifecycle() throws IOException
{
+ createParentEntities(METALAKE_NAME, "catalog1", "schema1", AUDIT_INFO);
+ ColumnEntity dropped = newIntColumn("dropped", 0, "comment");
+ ColumnEntity renamed = newIntColumn("renamed", 1, "comment");
+ ColumnEntity kept = newIntColumn("kept", 2, "comment");
+ TableEntity table =
+ TableEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("table1")
+ .withNamespace(Namespace.of(METALAKE_NAME, "catalog1", "schema1"))
+ .withColumns(Lists.newArrayList(dropped, renamed, kept))
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ TableMetaService.getInstance().insertTable(table, false);
+ TableColumnMetaService service = TableColumnMetaService.getInstance();
+ Assertions.assertEquals(
+ Map.of("dropped", dropped.id(), "renamed", renamed.id(), "kept",
kept.id()),
+ service.getColumnIdsByTableIdAndNames(table.id(), List.of("dropped",
"renamed", "kept")));
+
+ ColumnEntity renamedColumn = newIntColumn("new_name", 1, "comment",
renamed.id());
+ TableEntity updated =
+ TableEntity.builder()
+ .withId(table.id())
+ .withName(table.name())
+ .withNamespace(table.namespace())
+ .withColumns(Lists.newArrayList(renamedColumn, kept))
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ TableMetaService.getInstance().updateTable(table.nameIdentifier(), ignored
-> updated);
+ Assertions.assertEquals(
+ Map.of("new_name", renamed.id(), "kept", kept.id()),
+ service.getColumnIdsByTableIdAndNames(table.id(), List.of("new_name",
"kept")));
+ for (String missing : List.of("dropped", "renamed", "missing")) {
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () -> service.getColumnIdsByTableIdAndNames(table.id(),
List.of("kept", missing)));
+ }
+
+ ColumnEntity replacement = newIntColumn("dropped", 0, "replacement");
+ TableEntity replaced =
+ TableEntity.builder()
+ .withId(table.id())
+ .withName(table.name())
+ .withNamespace(table.namespace())
+ .withColumns(Lists.newArrayList(replacement, renamedColumn, kept))
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ TableMetaService.getInstance().updateTable(table.nameIdentifier(), ignored
-> replaced);
+ Assertions.assertEquals(
+ Map.of("dropped", replacement.id(), "new_name", renamed.id(), "kept",
kept.id()),
+ service.getColumnIdsByTableIdAndNames(table.id(), List.of("dropped",
"new_name", "kept")));
+
+ ColumnEntity otherColumn = newIntColumn("kept", 0, "other table");
+ TableEntity otherTable =
+ TableEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("table2")
+ .withNamespace(table.namespace())
+ .withColumns(Lists.newArrayList(otherColumn))
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ TableMetaService.getInstance().insertTable(otherTable, false);
+ Assertions.assertEquals(
+ Map.of("kept", kept.id()),
+ service.getColumnIdsByTableIdAndNames(table.id(), List.of("kept")));
+ Assertions.assertEquals(
+ Map.of("kept", otherColumn.id()),
+ service.getColumnIdsByTableIdAndNames(otherTable.id(),
List.of("kept")));
+ }
+
@TestTemplate
public void testDropColumnsRemovesTheirRelationsInBatches() throws Exception
{
String catalogName = "catalog1";
diff --git
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTagMetaService.java
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTagMetaService.java
index 7781afbd51..7194e5dae3 100644
---
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTagMetaService.java
+++
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestTagMetaService.java
@@ -33,9 +33,11 @@ import java.sql.Statement;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -51,6 +53,7 @@ import org.apache.gravitino.Namespace;
import org.apache.gravitino.RelationEdgeTarget;
import org.apache.gravitino.RelationQuery;
import org.apache.gravitino.RelationUpdate;
+import org.apache.gravitino.RelationalEntity;
import org.apache.gravitino.SupportsRelationOperations;
import org.apache.gravitino.authorization.AuthorizationUtils;
import org.apache.gravitino.authorization.Privileges;
@@ -1476,6 +1479,103 @@ public class TestTagMetaService extends TestJDBCBackend
{
Assertions.assertThrows(SQLException.class, () ->
insertDuplicateActiveTagRel(tagEntity.id()));
}
+ /**
+ * Verifies batch boundaries and per-source tag values, including columns
from different tables.
+ */
+ @TestTemplate
+ public void testBatchListTagsForMetadataObjects() throws IOException {
+ createAndInsertMakeLake(METALAKE_NAME);
+ createAndInsertCatalog(METALAKE_NAME, "catalog1");
+ createAndInsertSchema(METALAKE_NAME, "catalog1", "schema1");
+ List<ColumnEntity> columns = new ArrayList<>();
+ for (int i = 0; i < 502; i++) {
+ columns.add(newColumn("column_" + i, i));
+ }
+ TableEntity table =
+ TableEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("table1")
+ .withNamespace(Namespace.of(METALAKE_NAME, "catalog1", "schema1"))
+ .withColumns(columns)
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ backend.insert(table, false);
+ TableEntity otherTable =
+ TableEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("table2")
+ .withNamespace(table.namespace())
+ .withColumns(Lists.newArrayList(newColumn("column_0", 0)))
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ backend.insert(otherTable, false);
+ TagMetaService service = TagMetaService.getInstance();
+ for (String name : List.of("stage", "flag")) {
+ service.insertTag(
+ TagEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName(name)
+ .withNamespace(NamespaceUtil.ofTag(METALAKE_NAME))
+ .withAuditInfo(AUDIT_INFO)
+ .build(),
+ false);
+ }
+ List<NameIdentifier> identifiers = new ArrayList<>();
+ Map<NameIdentifier, Map<String, Set<String>>> expected = new
LinkedHashMap<>();
+ for (int i = 0; i < 501; i++) {
+ NameIdentifier identifier = columnIdent(table, columns.get(i).name());
+ identifiers.add(identifier);
+ service.associateTagValuesWithMetadataObject(
+ identifier,
+ Entity.EntityType.COLUMN,
+ new TagValue[] {TagValue.of("stage", "common"), TagValue.of("stage",
"value_" + i)},
+ new TagValue[0]);
+ expected.put(identifier, Map.of("stage", Set.of("common", "value_" +
i)));
+ }
+ service.associateTagValuesWithMetadataObject(
+ identifiers.get(0),
+ Entity.EntityType.COLUMN,
+ new TagValue[] {TagValue.noValue("flag")},
+ new TagValue[0]);
+ expected.put(
+ identifiers.get(0), Map.of("stage", Set.of("common", "value_0"),
"flag", Set.of()));
+ for (int count : new int[] {0, 1, 499, 500, 501}) {
+ List<NameIdentifier> requested = identifiers.subList(0, count);
+ assertBatchTagAssignments(
+ service,
+ requested,
+ Entity.EntityType.COLUMN,
+ requested.stream().collect(Collectors.toMap(identifier ->
identifier, expected::get)));
+ }
+
+ NameIdentifier otherColumn = columnIdent(otherTable, "column_0");
+ service.associateTagValuesWithMetadataObject(
+ otherColumn,
+ Entity.EntityType.COLUMN,
+ new TagValue[] {TagValue.of("stage", "other_table")},
+ new TagValue[0]);
+ expected.put(otherColumn, Map.of("stage", Set.of("other_table")));
+ identifiers.add(otherColumn);
+ identifiers.add(columnIdent(table, "column_501")); // An existing column
without tags.
+ identifiers.add(identifiers.get(0)); // Duplicate requests must not
duplicate relation edges.
+ assertBatchTagAssignments(service, identifiers, Entity.EntityType.COLUMN,
expected);
+
+ service.associateTagValuesWithMetadataObject(
+ table.nameIdentifier(),
+ Entity.EntityType.TABLE,
+ new TagValue[] {TagValue.of("stage", "table_value")},
+ new TagValue[0]);
+ assertBatchTagAssignments(
+ service,
+ List.of(table.nameIdentifier(), otherTable.nameIdentifier()),
+ Entity.EntityType.TABLE,
+ Map.of(table.nameIdentifier(), Map.of("stage",
Set.of("table_value"))));
+ identifiers.add(columnIdent(table, "missing"));
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () -> service.listTagsForMetadataObjects(identifiers,
Entity.EntityType.COLUMN));
+ }
+
@TestTemplate
public void testListTagsForMetadataObject() throws IOException {
testAssociateAndDisassociateTagsWithMetadataObject();
@@ -2098,6 +2198,22 @@ public class TestTagMetaService extends TestJDBCBackend {
// The dropped column's relation row is removed, not just hidden.
Assertions.assertEquals(2, countActiveTagRel(tag.id()));
+ assertBatchTagAssignments(
+ tagMetaService,
+ List.of(columnIdent(table, "c2_new"), columnIdent(table, "c3")),
+ Entity.EntityType.COLUMN,
+ Map.of(
+ columnIdent(table, "c2_new"), Map.of(tag.name(), Set.of()),
+ columnIdent(table, "c3"), Map.of(tag.name(), Set.of())));
+ for (String missing : List.of("c1", "c2", "missing")) {
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () ->
+ tagMetaService.listTagsForMetadataObjects(
+ List.of(columnIdent(table, "c3"), columnIdent(table,
missing)),
+ Entity.EntityType.COLUMN));
+ }
+
Assertions.assertEquals(
1,
tagMetaService
@@ -2127,6 +2243,13 @@ public class TestTagMetaService extends TestJDBCBackend {
.isEmpty());
Assertions.assertEquals(
2,
tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier()).size());
+ assertBatchTagAssignments(
+ tagMetaService,
+ List.of(columnIdent(table, "c1"), columnIdent(table, "c2_new"),
columnIdent(table, "c3")),
+ Entity.EntityType.COLUMN,
+ Map.of(
+ columnIdent(table, "c2_new"), Map.of(tag.name(), Set.of()),
+ columnIdent(table, "c3"), Map.of(tag.name(), Set.of())));
}
private static ColumnEntity newColumn(String name, int position) {
@@ -2179,6 +2302,35 @@ public class TestTagMetaService extends TestJDBCBackend {
() -> tagMetaService.getTagIdByTagName(metalakeId, "missing_tag"));
}
+ private void assertBatchTagAssignments(
+ TagMetaService service,
+ List<NameIdentifier> identifiers,
+ Entity.EntityType objectType,
+ Map<NameIdentifier, Map<String, Set<String>>> expected) {
+ List<RelationalEntity<?>> relations =
+ service.listTagsForMetadataObjects(identifiers, objectType);
+ relations.forEach(
+ relation -> {
+
assertEquals(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
relation.type());
+ assertEquals(objectType, relation.sourceType());
+ assertEquals(NamespaceUtil.ofTag(METALAKE_NAME),
relation.targetEntity().namespace());
+ });
+ Map<NameIdentifier, Map<String, Set<String>>> actual =
+ relations.stream()
+ .collect(
+ Collectors.groupingBy(
+ RelationalEntity::source,
+ Collectors.toMap(
+ relation -> relation.targetEntity().name(),
+ relation -> {
+ TagEntity tag = (TagEntity) relation.targetEntity();
+ assertTrue(tag.assignment().isPresent());
+ return new LinkedHashSet<>(
+ Arrays.asList(tag.assignment().get().values()));
+ })));
+ assertEquals(expected, actual);
+ }
+
private TagEntity copyTagWithComment(TagEntity tag, String comment) {
return copyTag(tag, tag.name(), comment);
}
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 0736256ace..7b66d309f0 100644
--- a/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java
+++ b/core/src/test/java/org/apache/gravitino/tag/TestEffectiveTagResolver.java
@@ -18,18 +18,26 @@
*/
package org.apache.gravitino.tag;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.CALLS_REAL_METHODS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
+import java.util.List;
+import java.util.Map;
import org.apache.gravitino.Entity;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.MetadataObject;
import org.apache.gravitino.MetadataObjects;
+import org.apache.gravitino.RelationalEntity;
import org.apache.gravitino.SupportsRelationOperations;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.exceptions.NoSuchMetadataObjectException;
@@ -66,6 +74,9 @@ public class TestEffectiveTagResolver {
metadataObjectUtil
.when(() -> MetadataObjectUtil.checkMetadataObject(METALAKE, TABLE))
.thenAnswer(invocation -> null);
+ metadataObjectUtil
+ .when(() -> MetadataObjectUtil.checkMetadataObjects(eq(METALAKE),
anyList()))
+ .thenAnswer(invocation -> null);
}
@AfterEach
@@ -151,6 +162,72 @@ public class TestEffectiveTagResolver {
Assertions.assertSame(exception, actual.getCause());
}
+ @Test
+ void testBatchSharesAncestorsAndPreservesNearestAssignment() throws
Exception {
+ MetadataObject first =
+ MetadataObjects.of(
+ Arrays.asList("catalog", "schema", "table", "first"),
MetadataObject.Type.COLUMN);
+ MetadataObject second =
+ MetadataObjects.of(
+ Arrays.asList("catalog", "schema", "table", "second"),
MetadataObject.Type.COLUMN);
+ TagEntity inherited = tag(1L, "domain", TagAssignment.ofValues("finance"));
+ TagEntity direct = tag(1L, "domain",
TagAssignment.ofValues("engineering"));
+ when(relationOperations.batchListEntitiesByRelation(any(), any(), any()))
+ .thenAnswer(
+ invocation -> {
+ Entity.EntityType type = invocation.getArgument(2);
+ if (type == Entity.EntityType.TABLE) {
+ return Collections.singletonList(
+ new RelationalEntity<>(
+
SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
+ MetadataObjectUtil.toEntityIdent(METALAKE, TABLE),
+ type,
+ inherited));
+ }
+ if (type == Entity.EntityType.COLUMN) {
+ return Collections.singletonList(
+ new RelationalEntity<>(
+
SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
+ MetadataObjectUtil.toEntityIdent(METALAKE, first),
+ type,
+ direct));
+ }
+ return Collections.emptyList();
+ });
+
+ Map<MetadataObject, TagEntity[]> result = resolver.resolve(METALAKE,
List.of(first, second));
+
+ metadataObjectUtil.verify(
+ () -> MetadataObjectUtil.checkMetadataObjects(METALAKE, List.of(first,
second)));
+
+ Assertions.assertArrayEquals(
+ new String[] {"engineering"},
result.get(first)[0].assignment().orElseThrow().values());
+ Assertions.assertFalse(result.get(first)[0].inherited().orElseThrow());
+ Assertions.assertArrayEquals(
+ new String[] {"finance"},
result.get(second)[0].assignment().orElseThrow().values());
+ Assertions.assertTrue(result.get(second)[0].inherited().orElseThrow());
+ verify(relationOperations, times(1))
+ .batchListEntitiesByRelation(
+ SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL,
+ List.of(MetadataObjectUtil.toEntityIdent(METALAKE, TABLE)),
+ Entity.EntityType.TABLE);
+
+ Map<MetadataObject, EffectiveTagResolver.ResolvedTag[]> withSources =
+ resolver.resolveWithSources(METALAKE, List.of(first, second));
+ Assertions.assertEquals(first, withSources.get(first)[0].source());
+ Assertions.assertEquals(TABLE, withSources.get(second)[0].source());
+ Assertions.assertArrayEquals(
+ result.get(first),
+ Arrays.stream(withSources.get(first))
+ .map(EffectiveTagResolver.ResolvedTag::tag)
+ .toArray(TagEntity[]::new));
+ Assertions.assertArrayEquals(
+ result.get(second),
+ Arrays.stream(withSources.get(second))
+ .map(EffectiveTagResolver.ResolvedTag::tag)
+ .toArray(TagEntity[]::new));
+ }
+
private static TagEntity tag(long id, String name, TagAssignment assignment)
{
return TagEntity.builder()
.withId(id)
diff --git
a/core/src/test/java/org/apache/gravitino/utils/TestMetadataObjectUtil.java
b/core/src/test/java/org/apache/gravitino/utils/TestMetadataObjectUtil.java
index 4cd8916e04..7ecd04fa24 100644
--- a/core/src/test/java/org/apache/gravitino/utils/TestMetadataObjectUtil.java
+++ b/core/src/test/java/org/apache/gravitino/utils/TestMetadataObjectUtil.java
@@ -18,6 +18,7 @@
*/
package org.apache.gravitino.utils;
+import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.never;
@@ -25,6 +26,7 @@ import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
+import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.gravitino.Entity;
@@ -42,6 +44,8 @@ import org.apache.gravitino.catalog.SemanticModelDispatcher;
import org.apache.gravitino.catalog.TableDispatcher;
import org.apache.gravitino.catalog.TopicDispatcher;
import org.apache.gravitino.catalog.ViewDispatcher;
+import org.apache.gravitino.exceptions.IllegalNamespaceException;
+import org.apache.gravitino.exceptions.NoSuchMetadataObjectException;
import org.apache.gravitino.job.JobOperationDispatcher;
import org.apache.gravitino.metalake.MetalakeDispatcher;
import org.apache.gravitino.policy.PolicyDispatcher;
@@ -361,6 +365,106 @@ public class TestMetadataObjectUtil {
verify(jobDispatcher).getJobTemplate("metalake", "template");
}
+ @Test
+ void testBatchChecksEachTableOnceAndDoesNotReuseChecksAcrossCalls() {
+ GravitinoEnv env = mock(GravitinoEnv.class);
+ TableDispatcher dispatcher = mock(TableDispatcher.class);
+ when(env.internalTableDispatcher()).thenReturn(dispatcher);
+ NameIdentifier tableIdentifier = NameIdentifier.of("metalake", "catalog",
"schema", "table");
+ when(dispatcher.tableExists(tableIdentifier)).thenReturn(true);
+ MetadataObject table = MetadataObjects.of("catalog.schema", "table",
MetadataObject.Type.TABLE);
+ MetadataObject first =
+ MetadataObjects.of("catalog.schema.table", "first",
MetadataObject.Type.COLUMN);
+ MetadataObject second =
+ MetadataObjects.of("catalog.schema.table", "second",
MetadataObject.Type.COLUMN);
+ try (MockedStatic<GravitinoEnv> mockedEnv =
mockStatic(GravitinoEnv.class)) {
+ mockedEnv.when(GravitinoEnv::getInstance).thenReturn(env);
+ MetadataObjectUtil.checkMetadataObjects("metalake", List.of(table,
first, second, table));
+ verify(dispatcher, times(1)).tableExists(tableIdentifier);
+ MetadataObjectUtil.checkMetadataObjects("metalake", List.of(first,
second, table));
+ verify(dispatcher, times(2)).tableExists(tableIdentifier);
+
+ when(dispatcher.tableExists(tableIdentifier)).thenReturn(false);
+ NoSuchMetadataObjectException exception =
+ Assertions.assertThrows(
+ NoSuchMetadataObjectException.class,
+ () -> MetadataObjectUtil.checkMetadataObjects("metalake",
List.of(first, table)));
+ Assertions.assertEquals(
+ "Metadata object catalog.schema.table.first type COLUMN doesn't
exist",
+ exception.getMessage());
+ verify(dispatcher, times(3)).tableExists(tableIdentifier);
+ }
+ }
+
+ @Test
+ void testBatchChecksTablesWithTheSameNameInDifferentNamespaces() {
+ GravitinoEnv env = mock(GravitinoEnv.class);
+ TableDispatcher dispatcher = mock(TableDispatcher.class);
+ when(env.internalTableDispatcher()).thenReturn(dispatcher);
+ when(dispatcher.tableExists(any())).thenReturn(true);
+ try (MockedStatic<GravitinoEnv> mockedEnv =
mockStatic(GravitinoEnv.class)) {
+ mockedEnv.when(GravitinoEnv::getInstance).thenReturn(env);
+ MetadataObjectUtil.checkMetadataObjects(
+ "metalake",
+ List.of(
+ MetadataObjects.of("catalog.schema.table", "first",
MetadataObject.Type.COLUMN),
+ MetadataObjects.of("catalog.other.table", "second",
MetadataObject.Type.COLUMN),
+ MetadataObjects.of("other.schema", "table",
MetadataObject.Type.TABLE)));
+ verify(dispatcher).tableExists(NameIdentifier.of("metalake", "catalog",
"schema", "table"));
+ verify(dispatcher).tableExists(NameIdentifier.of("metalake", "catalog",
"other", "table"));
+ verify(dispatcher).tableExists(NameIdentifier.of("metalake", "other",
"schema", "table"));
+ }
+ }
+
+ @Test
+ void testBatchStillChecksNonTableObjects() {
+ GravitinoEnv env = mock(GravitinoEnv.class);
+ SchemaDispatcher dispatcher = mock(SchemaDispatcher.class);
+ when(env.internalSchemaDispatcher()).thenReturn(dispatcher);
+ NameIdentifier schemaIdentifier = NameIdentifier.of("metalake", "catalog",
"schema");
+ when(dispatcher.schemaExists(schemaIdentifier)).thenReturn(false);
+ MetadataObject schema = MetadataObjects.of("catalog", "schema",
MetadataObject.Type.SCHEMA);
+ try (MockedStatic<GravitinoEnv> mockedEnv =
mockStatic(GravitinoEnv.class)) {
+ mockedEnv.when(GravitinoEnv::getInstance).thenReturn(env);
+ Assertions.assertThrows(
+ NoSuchMetadataObjectException.class,
+ () -> MetadataObjectUtil.checkMetadataObjects("metalake",
List.of(schema)));
+ verify(dispatcher).schemaExists(schemaIdentifier);
+ }
+ }
+
+ @Test
+ void testBatchRejectsNullObject() {
+ IllegalArgumentException exception =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ MetadataObjectUtil.checkMetadataObjects(
+ "metalake", Collections.singletonList(null)));
+ Assertions.assertEquals("metadataObject cannot be null",
exception.getMessage());
+ }
+
+ @Test
+ void testBatchRejectsMalformedColumnIdentifier() {
+ GravitinoEnv env = mock(GravitinoEnv.class);
+ TableDispatcher dispatcher = mock(TableDispatcher.class);
+ when(env.internalTableDispatcher()).thenReturn(dispatcher);
+ NameIdentifier tableIdentifier = NameIdentifier.of("metalake", "catalog",
"schema", "table");
+ when(dispatcher.tableExists(tableIdentifier)).thenReturn(true);
+ MetadataObject table = MetadataObjects.of("catalog.schema", "table",
MetadataObject.Type.TABLE);
+ MetadataObject malformedColumn = mock(MetadataObject.class);
+ when(malformedColumn.type()).thenReturn(MetadataObject.Type.COLUMN);
+
when(malformedColumn.fullName()).thenReturn("catalog.schema.table.column.extra");
+ try (MockedStatic<GravitinoEnv> mockedEnv =
mockStatic(GravitinoEnv.class)) {
+ mockedEnv.when(GravitinoEnv::getInstance).thenReturn(env);
+ Assertions.assertThrows(
+ IllegalNamespaceException.class,
+ () ->
+ MetadataObjectUtil.checkMetadataObjects("metalake",
List.of(table, malformedColumn)));
+ verify(dispatcher, times(1)).tableExists(tableIdentifier);
+ }
+ }
+
private static List<String> describe(List<MetadataObject> objects) {
return objects.stream().map(o -> o.type() + ":" +
o.fullName()).collect(Collectors.toList());
}
diff --git
a/iceberg/iceberg-rest-server-experiment/src/main/java/org/apache/gravitino/iceberg/service/restriction/ExperimentalIcebergReadRestrictionLoader.java
b/iceberg/iceberg-rest-server-experiment/src/main/java/org/apache/gravitino/iceberg/service/restriction/ExperimentalIcebergReadRestrictionLoader.java
index b67adf17aa..19b13ac185 100644
---
a/iceberg/iceberg-rest-server-experiment/src/main/java/org/apache/gravitino/iceberg/service/restriction/ExperimentalIcebergReadRestrictionLoader.java
+++
b/iceberg/iceberg-rest-server-experiment/src/main/java/org/apache/gravitino/iceberg/service/restriction/ExperimentalIcebergReadRestrictionLoader.java
@@ -104,10 +104,10 @@ public final class
ExperimentalIcebergReadRestrictionLoader
metalake, context.catalogName(), tableName.namespace().level(2)),
Entity.EntityType.SCHEMA);
Schema schema = response.tableMetadata().schema();
- PolicyEntity[] rowPolicies =
- policyDispatcher.listPolicyInfosForMetadataObject(
- metalake, NameIdentifierUtil.toMetadataObject(tableName,
Entity.EntityType.TABLE));
- Map<Integer, PolicyEntity[]> maskPolicies = new LinkedHashMap<>();
+ MetadataObject table = NameIdentifierUtil.toMetadataObject(tableName,
Entity.EntityType.TABLE);
+ List<MetadataObject> objects = new ArrayList<>();
+ objects.add(table);
+ Map<Integer, MetadataObject> columns = new LinkedHashMap<>();
for (Types.NestedField field : schema.columns()) {
NameIdentifier columnName =
NameIdentifierUtil.ofColumn(
@@ -118,9 +118,14 @@ public final class ExperimentalIcebergReadRestrictionLoader
field.name());
MetadataObject column =
NameIdentifierUtil.toMetadataObject(columnName,
Entity.EntityType.COLUMN);
- maskPolicies.put(
- field.fieldId(),
policyDispatcher.listPolicyInfosForMetadataObject(metalake, column));
+ objects.add(column);
+ columns.put(field.fieldId(), column);
}
+ Map<MetadataObject, PolicyEntity[]> policies =
+ policyDispatcher.listPolicyInfosForMetadataObjects(metalake, objects);
+ PolicyEntity[] rowPolicies = policies.get(table);
+ Map<Integer, PolicyEntity[]> maskPolicies = new LinkedHashMap<>();
+ columns.forEach((fieldId, column) -> maskPolicies.put(fieldId,
policies.get(column)));
IdentityBindingContext bindingContext =
new IdentityBindingContext(context.userName(), metalake,
accessControlDispatcher);
diff --git
a/iceberg/iceberg-rest-server-experiment/src/test/java/org/apache/gravitino/iceberg/service/restriction/TestExperimentalIcebergReadRestrictionLoader.java
b/iceberg/iceberg-rest-server-experiment/src/test/java/org/apache/gravitino/iceberg/service/restriction/TestExperimentalIcebergReadRestrictionLoader.java
index 4376efa863..2b4b5893cf 100644
---
a/iceberg/iceberg-rest-server-experiment/src/test/java/org/apache/gravitino/iceberg/service/restriction/TestExperimentalIcebergReadRestrictionLoader.java
+++
b/iceberg/iceberg-rest-server-experiment/src/test/java/org/apache/gravitino/iceberg/service/restriction/TestExperimentalIcebergReadRestrictionLoader.java
@@ -28,6 +28,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -177,7 +178,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
@Test
void testLoadWithoutApplicablePoliciesQueriesTableAndColumns() throws
Exception {
- PolicyDispatcher dispatcher = mock(PolicyDispatcher.class);
+ PolicyDispatcher dispatcher = mock(PolicyDispatcher.class,
CALLS_REAL_METHODS);
when(dispatcher.listPolicyInfosForMetadataObject(eq(METALAKE),
any(MetadataObject.class)))
.thenReturn(new PolicyEntity[0]);
LoadTableResponse response = response(SCHEMA);
@@ -189,6 +190,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
assertFalse(result.etagSeed().isPresent());
ArgumentCaptor<MetadataObject> objects =
ArgumentCaptor.forClass(MetadataObject.class);
verify(dispatcher,
times(4)).listPolicyInfosForMetadataObject(eq(METALAKE), objects.capture());
+ verify(dispatcher,
times(1)).listPolicyInfosForMetadataObjects(eq(METALAKE), any());
assertEquals(MetadataObject.Type.TABLE,
objects.getAllValues().get(0).type());
for (int index = 1; index < 4; index++) {
assertEquals(MetadataObject.Type.COLUMN,
objects.getAllValues().get(index).type());
@@ -197,7 +199,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
@Test
void testLoadBuildsRestrictedResponseAndSeed() throws Exception {
- PolicyDispatcher dispatcher = mock(PolicyDispatcher.class);
+ PolicyDispatcher dispatcher = mock(PolicyDispatcher.class,
CALLS_REAL_METHODS);
when(dispatcher.listPolicyInfosForMetadataObject(eq(METALAKE),
any(MetadataObject.class)))
.thenAnswer(
invocation -> {
@@ -239,7 +241,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
@Test
void testLoadRejectsDottedTableNameWithPolicyDispatcher() throws Exception {
- PolicyDispatcher dispatcher = mock(PolicyDispatcher.class);
+ PolicyDispatcher dispatcher = mock(PolicyDispatcher.class,
CALLS_REAL_METHODS);
IllegalNameIdentifierException exception =
assertThrows(
@@ -254,7 +256,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
@Test
void testLoadRejectsDottedColumnNameWithPolicyDispatcher() throws Exception {
Schema schema = new Schema(Types.NestedField.optional(1, "a.b",
Types.StringType.get()));
- PolicyDispatcher dispatcher = mock(PolicyDispatcher.class);
+ PolicyDispatcher dispatcher = mock(PolicyDispatcher.class,
CALLS_REAL_METHODS);
when(dispatcher.listPolicyInfosForMetadataObject(eq(METALAKE),
any(MetadataObject.class)))
.thenReturn(new PolicyEntity[0]);
@@ -268,7 +270,7 @@ public class TestExperimentalIcebergReadRestrictionLoader {
@Test
void testLoadRejectsDottedAncestorNamesWithPolicyDispatcher() throws
Exception {
- PolicyDispatcher dispatcher = mock(PolicyDispatcher.class);
+ PolicyDispatcher dispatcher = mock(PolicyDispatcher.class,
CALLS_REAL_METHODS);
TableIdentifier table = TableIdentifier.of("schema", "table");
IllegalNameIdentifierException metalakeError =
