From 30d53be7643c26c2cac6e0acd5ade0f9bc6d00eb Mon Sep 17 00:00:00 2001 From: Jerry Shao Date: Tue, 14 Jan 2025 17:53:07 +0800 Subject: [PATCH] Add tag support for model --- .../org/apache/gravitino/MetadataObjects.java | 1 + .../apache/gravitino/client/GenericModel.java | 42 ++++++- .../gravitino/client/GenericModelCatalog.java | 4 +- .../client/integration/test/TagIT.java | 113 ++++++++++++++++++ docs/manage-tags-in-gravitino.md | 8 +- docs/open-api/tags.yaml | 4 +- 6 files changed, 161 insertions(+), 11 deletions(-) diff --git a/api/src/main/java/org/apache/gravitino/MetadataObjects.java b/api/src/main/java/org/apache/gravitino/MetadataObjects.java index 557ccdefc49..e96b6e7e4a6 100644 --- a/api/src/main/java/org/apache/gravitino/MetadataObjects.java +++ b/api/src/main/java/org/apache/gravitino/MetadataObjects.java @@ -126,6 +126,7 @@ public static MetadataObject parent(MetadataObject object) { case TABLE: case FILESET: case TOPIC: + case MODEL: parentType = MetadataObject.Type.SCHEMA; break; case SCHEMA: diff --git a/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModel.java b/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModel.java index 2d356b712fe..d4c0c24f1a9 100644 --- a/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModel.java +++ b/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModel.java @@ -18,20 +18,35 @@ */ package org.apache.gravitino.client; +import com.google.common.collect.Lists; +import java.util.List; import java.util.Map; import org.apache.gravitino.Audit; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.Namespace; import org.apache.gravitino.authorization.SupportsRoles; import org.apache.gravitino.dto.model.ModelDTO; +import org.apache.gravitino.exceptions.NoSuchTagException; +import org.apache.gravitino.exceptions.TagAlreadyAssociatedException; import org.apache.gravitino.model.Model; import org.apache.gravitino.tag.SupportsTags; +import org.apache.gravitino.tag.Tag; /** Represents a generic model. */ -class GenericModel implements Model { +class GenericModel implements Model, SupportsTags { private final ModelDTO modelDTO; - GenericModel(ModelDTO modelDTO) { + private final MetadataObjectTagOperations objectTagOperations; + + GenericModel(ModelDTO modelDTO, RESTClient restClient, Namespace modelNs) { this.modelDTO = modelDTO; + List modelFullName = + Lists.newArrayList(modelNs.level(1), modelNs.level(2), modelDTO.name()); + MetadataObject modelObject = MetadataObjects.of(modelFullName, MetadataObject.Type.MODEL); + this.objectTagOperations = + new MetadataObjectTagOperations(modelNs.level(0), modelObject, restClient); } @Override @@ -61,7 +76,7 @@ public int latestVersion() { @Override public SupportsTags supportsTags() { - throw new UnsupportedOperationException("Not supported yet."); + return this; } @Override @@ -86,4 +101,25 @@ public boolean equals(Object o) { public int hashCode() { return modelDTO.hashCode(); } + + @Override + public String[] listTags() { + return objectTagOperations.listTags(); + } + + @Override + public Tag[] listTagsInfo() { + return objectTagOperations.listTagsInfo(); + } + + @Override + public Tag getTag(String name) throws NoSuchTagException { + return objectTagOperations.getTag(name); + } + + @Override + public String[] associateTags(String[] tagsToAdd, String[] tagsToRemove) + throws TagAlreadyAssociatedException { + return objectTagOperations.associateTags(tagsToAdd, tagsToRemove); + } } diff --git a/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModelCatalog.java b/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModelCatalog.java index 50e9eb246ac..8658a972596 100644 --- a/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModelCatalog.java +++ b/clients/client-java/src/main/java/org/apache/gravitino/client/GenericModelCatalog.java @@ -97,7 +97,7 @@ public Model getModel(NameIdentifier ident) throws NoSuchModelException { ErrorHandlers.modelErrorHandler()); resp.validate(); - return new GenericModel(resp.getModel()); + return new GenericModel(resp.getModel(), restClient, modelFullNs); } @Override @@ -118,7 +118,7 @@ public Model registerModel(NameIdentifier ident, String comment, Map c.name().equals("col1")).findFirst().get(); + + // Create model catalog + String modelCatalogName = GravitinoITUtils.genRandomName("tag_it_model_catalog"); + Assertions.assertFalse(metalake.catalogExists(modelCatalogName)); + modelCatalog = + metalake.createCatalog( + modelCatalogName, Catalog.Type.MODEL, "comment", Collections.emptyMap()); + + // Create model schema + String modelSchemaName = GravitinoITUtils.genRandomName("tag_it_model_schema"); + Assertions.assertFalse(modelCatalog.asSchemas().schemaExists(modelSchemaName)); + modelSchema = + modelCatalog.asSchemas().createSchema(modelSchemaName, "comment", Collections.emptyMap()); + + // Create model + String modelName = GravitinoITUtils.genRandomName("tag_it_model"); + Assertions.assertFalse( + modelCatalog.asModelCatalog().modelExists(NameIdentifier.of(modelSchemaName, modelName))); + model = + modelCatalog + .asModelCatalog() + .registerModel( + NameIdentifier.of(modelSchemaName, modelName), "comment", Collections.emptyMap()); } @AfterAll @@ -115,6 +143,11 @@ public void tearDown() { relationalCatalog.asTableCatalog().dropTable(NameIdentifier.of(schema.name(), table.name())); relationalCatalog.asSchemas().dropSchema(schema.name(), true); metalake.dropCatalog(relationalCatalog.name(), true); + + modelCatalog.asModelCatalog().deleteModel(NameIdentifier.of(modelSchema.name(), model.name())); + modelCatalog.asSchemas().dropSchema(modelSchema.name(), true); + metalake.dropCatalog(modelCatalog.name(), true); + client.dropMetalake(metalakeName, true); if (client != null) { @@ -653,4 +686,84 @@ public void testAssociateAndDeleteTags() { String[] associatedTags1 = relationalCatalog.supportsTags().listTags(); Assertions.assertArrayEquals(new String[] {tag2.name()}, associatedTags1); } + + @Test + public void testAssociateTagsToModel() { + Tag tag1 = + metalake.createTag( + GravitinoITUtils.genRandomName("tag_it_model_tag1"), + "comment1", + Collections.emptyMap()); + Tag tag2 = + metalake.createTag( + GravitinoITUtils.genRandomName("tag_it_model_tag2"), + "comment2", + Collections.emptyMap()); + Tag tag3 = + metalake.createTag( + GravitinoITUtils.genRandomName("tag_it_model_tag3"), + "comment3", + Collections.emptyMap()); + + // Associate tags to catalog + modelCatalog.supportsTags().associateTags(new String[] {tag1.name()}, null); + + // Associate tags to schema + modelSchema.supportsTags().associateTags(new String[] {tag2.name()}, null); + + // Associate tags to model + model.supportsTags().associateTags(new String[] {tag3.name()}, null); + + // Test list associated tags for model + String[] tags1 = model.supportsTags().listTags(); + Assertions.assertEquals(3, tags1.length); + Set tagNames = Sets.newHashSet(tags1); + Assertions.assertTrue(tagNames.contains(tag1.name())); + Assertions.assertTrue(tagNames.contains(tag2.name())); + Assertions.assertTrue(tagNames.contains(tag3.name())); + + // Test list associated tags with details for model + Tag[] tags2 = model.supportsTags().listTagsInfo(); + Assertions.assertEquals(3, tags2.length); + + Set nonInheritedTags = + Arrays.stream(tags2).filter(tag -> !tag.inherited().get()).collect(Collectors.toSet()); + Set inheritedTags = + Arrays.stream(tags2).filter(tag -> tag.inherited().get()).collect(Collectors.toSet()); + + Assertions.assertEquals(1, nonInheritedTags.size()); + Assertions.assertEquals(2, inheritedTags.size()); + Assertions.assertTrue(nonInheritedTags.contains(tag3)); + Assertions.assertTrue(inheritedTags.contains(tag1)); + Assertions.assertTrue(inheritedTags.contains(tag2)); + + // Test get associated tag for model + Tag resultTag1 = model.supportsTags().getTag(tag1.name()); + Assertions.assertEquals(tag1, resultTag1); + Assertions.assertTrue(resultTag1.inherited().get()); + + Tag resultTag2 = model.supportsTags().getTag(tag2.name()); + Assertions.assertEquals(tag2, resultTag2); + Assertions.assertTrue(resultTag2.inherited().get()); + + Tag resultTag3 = model.supportsTags().getTag(tag3.name()); + Assertions.assertEquals(tag3, resultTag3); + Assertions.assertFalse(resultTag3.inherited().get()); + + // Test get objects associated with tag + Assertions.assertEquals(1, tag1.associatedObjects().count()); + Assertions.assertEquals(modelCatalog.name(), tag1.associatedObjects().objects()[0].name()); + Assertions.assertEquals( + MetadataObject.Type.CATALOG, tag1.associatedObjects().objects()[0].type()); + + Assertions.assertEquals(1, tag2.associatedObjects().count()); + Assertions.assertEquals(modelSchema.name(), tag2.associatedObjects().objects()[0].name()); + Assertions.assertEquals( + MetadataObject.Type.SCHEMA, tag2.associatedObjects().objects()[0].type()); + + Assertions.assertEquals(1, tag3.associatedObjects().count()); + Assertions.assertEquals(model.name(), tag3.associatedObjects().objects()[0].name()); + Assertions.assertEquals( + MetadataObject.Type.MODEL, tag3.associatedObjects().objects()[0].type()); + } } diff --git a/docs/manage-tags-in-gravitino.md b/docs/manage-tags-in-gravitino.md index 4163ca89d2c..04e5b79ec87 100644 --- a/docs/manage-tags-in-gravitino.md +++ b/docs/manage-tags-in-gravitino.md @@ -23,12 +23,12 @@ the future versions. :::info 1. Metadata objects are objects that are managed in Gravitino, such as `CATALOG`, `SCHEMA`, `TABLE`, - `COLUMN`, `FILESET`, `TOPIC`, `COLUMN`, etc. A metadata object is combined by a `type` and a + `COLUMN`, `FILESET`, `TOPIC`, `COLUMN`, `MODEL`, etc. A metadata object is combined by a `type` and a comma-separated `name`. For example, a `CATAGLOG` object has a name "catalog1" with type "CATALOG", a `SCHEMA` object has a name "catalog1.schema1" with type "SCHEMA", a `TABLE` object has a name "catalog1.schema1.table1" with type "TABLE", a `COLUMN` object has a name "catalog1.schema1.table1.column1" with type "COLUMN". -2. Currently, `CATALOG`, `SCHEMA`, `TABLE`, `FILESET`, `TOPIC`, and `COLUMN` objects can be tagged. +2. Currently, `CATALOG`, `SCHEMA`, `TABLE`, `FILESET`, `TOPIC`, `MODEL`, and `COLUMN` objects can be tagged. 3. Tags in Gravitino is inheritable, so listing tags of a metadata object will also list the tags of its parent metadata objects. For example, listing tags of a `Table` will also list the tags of its parent `Schema` and `Catalog`. @@ -204,8 +204,8 @@ client.deleteTag("tag2"); ## Tag associations -Gravitino allows you to associate and disassociate tags with metadata objects. Currently, only -`CATALOG`, `SCHEMA`, `TABLE`, `FILESET`, `TOPIC` objects can be tagged. +Gravitino allows you to associate and disassociate tags with metadata objects. Currently, +`CATALOG`, `SCHEMA`, `TABLE`, `FILESET`, `TOPIC`, `MODEL`, and `COLUMN` objects can be tagged. ### Associate and disassociate tags with a metadata object diff --git a/docs/open-api/tags.yaml b/docs/open-api/tags.yaml index a3be5230b94..a8b3ac053c2 100644 --- a/docs/open-api/tags.yaml +++ b/docs/open-api/tags.yaml @@ -394,8 +394,8 @@ components: - "TABLE" - "FILESET" - "TOPIC" - - "ROLE" - - "METALAKE" + - "MODEL" + - "COLUMN" requests: