This is an automated email from the ASF dual-hosted git repository. jshao pushed a commit to branch branch-0.9 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit ea1788196c56f7e9d4ddf2cc22d5e96a69ac8aa5 Author: Lord of Abyss <[email protected]> AuthorDate: Mon May 12 17:49:46 2025 +0800 [#6814] feat(core): Support update aliases for model version (#7133) ### What changes were proposed in this pull request? - [X] PR1: Add ModelVersionChange API interface, Implement the update alias logic in model catalog and JDBC backend logic, update related event. - [X] PR2: Add REST endpoint to support model version change, add Java client and Python client for model version alias update. Meanwhile, fixed some typos and added reload logic to each model version change test. ### Why are the changes needed? Fix: #6814 ### Does this PR introduce _any_ user-facing change? user now can update aliases of a model version. ### How was this patch tested? local test(ut + it). original model version  `bin/gcli.sh model update -m demo_metalake --name model_catalog.schema.model2 --newalias 'alias3' --alias test`  `bin/gcli.sh model remove -m demo_metalake --name model_catalog.schema.model2 --removealias "alias1" "alias2" "alias3" --alias "alias1"`  --- .../apache/gravitino/model/ModelVersionChange.java | 30 ++-- .../gravitino/model/TestModelVersionChange.java | 40 ++--- .../catalog/model/ModelCatalogOperations.java | 2 +- .../integration/test/ModelCatalogOperationsIT.java | 192 +++++++++++++++++++++ .../org/apache/gravitino/cli/GravitinoOptions.java | 8 + .../apache/gravitino/cli/ModelCommandHandler.java | 38 +++- .../apache/gravitino/cli/TestableCommandLine.java | 15 ++ ...Comment.java => UpdateModelVersionAliases.java} | 39 +++-- .../cli/commands/UpdateModelVersionComment.java | 9 +- .../cli/commands/UpdateModelVersionUri.java | 9 +- .../apache/gravitino/cli/TestModelCommands.java | 93 ++++++++++ .../org/apache/gravitino/client/DTOConverters.java | 6 + .../gravitino/api/model_version_change.py | 68 ++++++++ .../gravitino/client/generic_model_catalog.py | 5 + .../dto/requests/model_version_update_request.py | 50 +++++- .../gravitino/filesystem/gvfs_base_operations.py | 8 +- .../tests/integration/test_model_catalog.py | 45 +++++ .../tests/unittests/test_gravitino_version.py | 16 +- .../dto/requests/ModelVersionUpdateRequest.java | 37 +++- .../gravitino/connector/TestCatalogOperations.java | 2 +- docs/manage-model-metadata-using-gravitino.md | 54 ++++-- docs/open-api/models.yaml | 31 ++++ 22 files changed, 713 insertions(+), 84 deletions(-) diff --git a/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java b/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java index 4e28071e05..55fa015453 100644 --- a/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java +++ b/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java @@ -84,16 +84,16 @@ public interface ModelVersionChange { * Create a ModelVersionChange for updating the aliases of a model version. * * @param aliasesToAdd The new aliases to be added for the model version. - * @param aliasesToDelete The aliases to be removed from the model version. + * @param aliasesToRemove The aliases to be removed from the model version. * @return A new ModelVersionChange instance for updating the aliases of a model version. */ - static ModelVersionChange updateAliases(String[] aliasesToAdd, String[] aliasesToDelete) { + static ModelVersionChange updateAliases(String[] aliasesToAdd, String[] aliasesToRemove) { String[] toAdd = aliasesToAdd == null ? new String[0] : aliasesToAdd; - String[] toDelete = aliasesToDelete == null ? new String[0] : aliasesToDelete; + String[] toRemove = aliasesToRemove == null ? new String[0] : aliasesToRemove; return new UpdateAliases( Arrays.stream(toAdd).collect(Collectors.toList()), - Arrays.stream(toDelete).collect(Collectors.toList())); + Arrays.stream(toRemove).collect(Collectors.toList())); } /** A ModelVersionChange to update the model version comment. */ @@ -361,20 +361,20 @@ public interface ModelVersionChange { */ final class UpdateAliases implements ModelVersionChange { private final ImmutableSortedSet<String> aliasesToAdd; - private final ImmutableSortedSet<String> aliasesToDelete; + private final ImmutableSortedSet<String> aliasesToRemove; /** * Constructs a new aliases-update operation, specifying the aliases to add and remove. * * @param aliasesToAdd the aliases to add, or null for none - * @param aliasesToDelete the aliases to remove, or null for none + * @param aliasesToRemove the aliases to remove, or null for none */ - public UpdateAliases(List<String> aliasesToAdd, List<String> aliasesToDelete) { + public UpdateAliases(List<String> aliasesToAdd, List<String> aliasesToRemove) { this.aliasesToAdd = ImmutableSortedSet.copyOf(aliasesToAdd != null ? aliasesToAdd : Lists.newArrayList()); - this.aliasesToDelete = + this.aliasesToRemove = ImmutableSortedSet.copyOf( - aliasesToDelete != null ? aliasesToDelete : Lists.newArrayList()); + aliasesToRemove != null ? aliasesToRemove : Lists.newArrayList()); } /** @@ -391,8 +391,8 @@ public interface ModelVersionChange { * * @return an immutable, sorted set of aliases to remove */ - public Set<String> aliasesToDelete() { - return aliasesToDelete; + public Set<String> aliasesToRemove() { + return aliasesToRemove; } /** @@ -408,7 +408,7 @@ public interface ModelVersionChange { if (this == o) return true; if (!(o instanceof UpdateAliases)) return false; UpdateAliases that = (UpdateAliases) o; - return aliasesToAdd.equals(that.aliasesToAdd) && aliasesToDelete.equals(that.aliasesToDelete); + return aliasesToAdd.equals(that.aliasesToAdd) && aliasesToRemove.equals(that.aliasesToRemove); } /** @@ -419,7 +419,7 @@ public interface ModelVersionChange { */ @Override public int hashCode() { - return Objects.hash(aliasesToAdd, aliasesToDelete); + return Objects.hash(aliasesToAdd, aliasesToRemove); } /** @@ -435,8 +435,8 @@ public interface ModelVersionChange { + COMMA_JOINER.join(aliasesToAdd) + ")" + " " - + "AliasToDelete: (" - + COMMA_JOINER.join(aliasesToDelete) + + "AliasToRemove: (" + + COMMA_JOINER.join(aliasesToRemove) + ")"; } } diff --git a/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java b/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java index d2b6526841..cd64776390 100644 --- a/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java +++ b/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java @@ -206,10 +206,10 @@ public class TestModelVersionChange { @Test void testCreateUpdateVersionAliasUseStaticMethod() { String[] aliasesToAdd = {"alias add 1", "alias add 2"}; - String[] aliasesToDelete = {"alias delete 1", "alias delete 2"}; + String[] aliasesToRemove = {"alias remove 1", "alias remove 2"}; ModelVersionChange modelVersionChange = - ModelVersionChange.updateAliases(aliasesToAdd, aliasesToDelete); + ModelVersionChange.updateAliases(aliasesToAdd, aliasesToRemove); Assertions.assertEquals(ModelVersionChange.UpdateAliases.class, modelVersionChange.getClass()); @@ -218,13 +218,13 @@ public class TestModelVersionChange { Assertions.assertEquals( ImmutableSet.of("alias add 1", "alias add 2"), updateAliasesChange.aliasesToAdd()); Assertions.assertEquals( - ImmutableSet.of("alias delete 1", "alias delete 2"), updateAliasesChange.aliasesToDelete()); + ImmutableSet.of("alias remove 1", "alias remove 2"), updateAliasesChange.aliasesToRemove()); Assertions.assertEquals( "UpdateAlias " + "AliasToAdd: (alias add 1,alias add 2)" + " " - + "AliasToDelete: (alias " - + "delete 1,alias delete 2)", + + "AliasToRemove: (alias " + + "remove 1,alias remove 2)", updateAliasesChange.toString()); } @@ -236,18 +236,18 @@ public class TestModelVersionChange { ModelVersionChange.UpdateAliases updateAliasesChange = (ModelVersionChange.UpdateAliases) modelVersionChange; Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToAdd()); - Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToDelete()); + Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToRemove()); Assertions.assertEquals( - "UpdateAlias AliasToAdd: () AliasToDelete: ()", updateAliasesChange.toString()); + "UpdateAlias AliasToAdd: () AliasToRemove: ()", updateAliasesChange.toString()); } @Test void testCreateUpdateVersionAliasUseConstructor() { List<String> aliasesToAdd = Lists.newArrayList("alias add 1", "alias add 2"); - List<String> aliasesToDelete = Lists.newArrayList("alias delete 1", "alias delete 2"); + List<String> aliasesToRemove = Lists.newArrayList("alias remove 1", "alias remove 2"); ModelVersionChange modelVersionChange = - new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToDelete); + new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToRemove); Assertions.assertEquals(ModelVersionChange.UpdateAliases.class, modelVersionChange.getClass()); @@ -256,13 +256,13 @@ public class TestModelVersionChange { Assertions.assertEquals( ImmutableSet.of("alias add 1", "alias add 2"), updateAliasesChange.aliasesToAdd()); Assertions.assertEquals( - ImmutableSet.of("alias delete 1", "alias delete 2"), updateAliasesChange.aliasesToDelete()); + ImmutableSet.of("alias remove 1", "alias remove 2"), updateAliasesChange.aliasesToRemove()); Assertions.assertEquals( "UpdateAlias " + "AliasToAdd: (alias add 1,alias add 2)" + " " - + "AliasToDelete: (alias " - + "delete 1,alias delete 2)", + + "AliasToRemove: (alias " + + "remove 1,alias remove 2)", updateAliasesChange.toString()); } @@ -274,27 +274,27 @@ public class TestModelVersionChange { ModelVersionChange.UpdateAliases updateAliasesChange = (ModelVersionChange.UpdateAliases) modelVersionChange; Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToAdd()); - Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToDelete()); + Assertions.assertEquals(ImmutableSet.of(), updateAliasesChange.aliasesToRemove()); Assertions.assertEquals( - "UpdateAlias AliasToAdd: () AliasToDelete: ()", updateAliasesChange.toString()); + "UpdateAlias AliasToAdd: () AliasToRemove: ()", updateAliasesChange.toString()); } @Test void testUpdateVersionAliasChangeEquals() { List<String> aliasesToAdd = Lists.newArrayList("alias add 1", "alias add 2"); - List<String> aliasesToDelete = Lists.newArrayList("alias delete 1", "alias delete 2"); + List<String> aliasesToRemove = Lists.newArrayList("alias remove 1", "alias remove 2"); List<String> differentAliasesToAdd = Lists.newArrayList("alias add 1", "alias add 3"); - List<String> differentAliasesToDelete = Lists.newArrayList("alias delete 1", "alias delete 3"); + List<String> differentAliasesToRemove = Lists.newArrayList("alias remove 1", "alias remove 3"); ModelVersionChange modelVersionChange1 = - new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToDelete); + new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToRemove); ModelVersionChange modelVersionChange2 = - new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToDelete); + new ModelVersionChange.UpdateAliases(aliasesToAdd, aliasesToRemove); ModelVersionChange modelVersionChange3 = - new ModelVersionChange.UpdateAliases(differentAliasesToAdd, aliasesToDelete); + new ModelVersionChange.UpdateAliases(differentAliasesToAdd, aliasesToRemove); ModelVersionChange modelVersionChange4 = - new ModelVersionChange.UpdateAliases(aliasesToAdd, differentAliasesToDelete); + new ModelVersionChange.UpdateAliases(aliasesToAdd, differentAliasesToRemove); Assertions.assertEquals(modelVersionChange1, modelVersionChange2); Assertions.assertNotEquals(modelVersionChange1, modelVersionChange3); diff --git a/catalogs/catalog-model/src/main/java/org/apache/gravitino/catalog/model/ModelCatalogOperations.java b/catalogs/catalog-model/src/main/java/org/apache/gravitino/catalog/model/ModelCatalogOperations.java index 11c81ed9b0..924057e1ca 100644 --- a/catalogs/catalog-model/src/main/java/org/apache/gravitino/catalog/model/ModelCatalogOperations.java +++ b/catalogs/catalog-model/src/main/java/org/apache/gravitino/catalog/model/ModelCatalogOperations.java @@ -442,7 +442,7 @@ public class ModelCatalogOperations extends ManagedSchemaOperations ModelVersionChange.UpdateAliases updateAliasesChange = (ModelVersionChange.UpdateAliases) change; Set<String> addTmpSet = updateAliasesChange.aliasesToAdd(); - Set<String> deleteTmpSet = updateAliasesChange.aliasesToDelete(); + Set<String> deleteTmpSet = updateAliasesChange.aliasesToRemove(); Set<String> aliasToAdd = Sets.difference(addTmpSet, deleteTmpSet).immutableCopy(); Set<String> aliasToDelete = Sets.difference(deleteTmpSet, addTmpSet).immutableCopy(); diff --git a/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/integration/test/ModelCatalogOperationsIT.java b/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/integration/test/ModelCatalogOperationsIT.java index ecda97cbad..56fed62a3f 100644 --- a/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/integration/test/ModelCatalogOperationsIT.java +++ b/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/integration/test/ModelCatalogOperationsIT.java @@ -423,6 +423,15 @@ public class ModelCatalogOperationsIT extends BaseIT { gravitinoCatalog .asModelCatalog() .alterModel(NameIdentifier.of(schemaName, null), updateName)); + + // reload model and check name + Model reloadedModel = + gravitinoCatalog + .asModelCatalog() + .getModel(NameIdentifier.of(modelIdent.namespace(), newName)); + Assertions.assertEquals(newName, reloadedModel.name()); + Assertions.assertEquals(createdModel.properties(), reloadedModel.properties()); + Assertions.assertEquals(createdModel.comment(), reloadedModel.comment()); } @Test @@ -444,6 +453,13 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertNotEquals(createdModel.properties(), alteredModel.properties()); Assertions.assertEquals(newProperties, alteredModel.properties()); Assertions.assertEquals(createdModel.comment(), alteredModel.comment()); + + // reload model and check properties + Model reloadedModel = gravitinoCatalog.asModelCatalog().getModel(modelIdent); + Assertions.assertEquals(modelName, reloadedModel.name()); + Assertions.assertNotEquals(createdModel.properties(), reloadedModel.properties()); + Assertions.assertEquals(newProperties, reloadedModel.properties()); + Assertions.assertEquals(createdModel.comment(), reloadedModel.comment()); } @Test @@ -462,6 +478,12 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertEquals(modelName, alteredModel.name()); Assertions.assertEquals(newProperties, alteredModel.properties()); Assertions.assertEquals(createdModel.comment(), alteredModel.comment()); + + // reload model and check properties + Model reloadedModel = gravitinoCatalog.asModelCatalog().getModel(modelIdent); + Assertions.assertEquals(modelName, reloadedModel.name()); + Assertions.assertEquals(newProperties, reloadedModel.properties()); + Assertions.assertEquals(createdModel.comment(), reloadedModel.comment()); } @Test @@ -480,6 +502,12 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertEquals(modelName, alteredModel.name()); Assertions.assertEquals(newProperties, alteredModel.properties()); Assertions.assertEquals(createdModel.comment(), alteredModel.comment()); + + // reload model and check properties + Model reloadedModel = gravitinoCatalog.asModelCatalog().getModel(modelIdent); + Assertions.assertEquals(modelName, reloadedModel.name()); + Assertions.assertEquals(newProperties, reloadedModel.properties()); + Assertions.assertEquals(createdModel.comment(), reloadedModel.comment()); } @Test @@ -498,6 +526,12 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertEquals(createdModel.name(), alteredModel.name()); Assertions.assertEquals(createdModel.properties(), alteredModel.properties()); Assertions.assertEquals(newComment, alteredModel.comment()); + + // reload model and check comment + Model reloadedModel = gravitinoCatalog.asModelCatalog().getModel(modelIdent); + Assertions.assertEquals(createdModel.name(), reloadedModel.name()); + Assertions.assertEquals(createdModel.properties(), reloadedModel.properties()); + Assertions.assertEquals(newComment, reloadedModel.comment()); } @Test @@ -535,6 +569,16 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(newProperties, updatedModelVersion.properties()); + + // reload model version and check properties + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, 0); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(newProperties, reloadedModelVersion.properties()); } @Test @@ -573,6 +617,16 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(newProperties, updatedModelVersion.properties()); + + // reload model version and check properties + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, aliases[0]); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(newProperties, reloadedModelVersion.properties()); } @Test @@ -609,6 +663,16 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(modelVersion.properties(), updatedModelVersion.properties()); + + // reload model version and check uri + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, 0); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(newUri, reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), reloadedModelVersion.properties()); } @Test @@ -648,6 +712,16 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(modelVersion.properties(), updatedModelVersion.properties()); + + // reload model version and check uri + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, aliases[0]); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(newUri, reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), reloadedModelVersion.properties()); } @Test @@ -681,6 +755,16 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(newProperties, updatedModelVersion.properties()); + + // reload model version and check properties + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, 0); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(newProperties, reloadedModelVersion.properties()); } @Test @@ -715,6 +799,114 @@ public class ModelCatalogOperationsIT extends BaseIT { Assertions.assertArrayEquals(modelVersion.aliases(), updatedModelVersion.aliases()); Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); Assertions.assertEquals(newProperties, updatedModelVersion.properties()); + + // reload model version and check properties + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, aliases[0]); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(modelVersion.aliases(), reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(newProperties, reloadedModelVersion.properties()); + } + + @Test + void testLinkAndUpdateModelVersionAliases() { + String modelName = RandomNameUtils.genRandomName("model1"); + String[] aliases = {"alias1", "alias2"}; + Map<String, String> properties = ImmutableMap.of("key1", "val1", "key2", "val2"); + NameIdentifier modelIdent = NameIdentifier.of(schemaName, modelName); + + gravitinoCatalog.asModelCatalog().registerModel(modelIdent, null, null); + + gravitinoCatalog + .asModelCatalog() + .linkModelVersion(modelIdent, "uri", aliases, "comment", properties); + + ModelVersion modelVersion = gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, 0); + + Assertions.assertEquals(0, modelVersion.version()); + Assertions.assertEquals("uri", modelVersion.uri()); + Assertions.assertArrayEquals(aliases, modelVersion.aliases()); + Assertions.assertEquals("comment", modelVersion.comment()); + Assertions.assertEquals(properties, modelVersion.properties()); + + ModelVersionChange change = + ModelVersionChange.updateAliases( + new String[] {"alias1", "alias3"}, new String[] {"alias1", "alias2"}); + ModelVersion updatedModelVersion = + gravitinoCatalog.asModelCatalog().alterModelVersion(modelIdent, 0, change); + String[] newAliases = {"alias1", "alias3"}; + + Assertions.assertEquals(modelVersion.version(), updatedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), updatedModelVersion.uri()); + Assertions.assertArrayEquals(newAliases, updatedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), updatedModelVersion.properties()); + + // reload model version and check aliases + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, 0); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(newAliases, reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), reloadedModelVersion.properties()); + } + + @Test + void testLinkAndUpdateModelVersionAliasesByAlias() { + String modelName = RandomNameUtils.genRandomName("model1"); + String[] aliases = {"alias1", "alias2"}; + Map<String, String> properties = ImmutableMap.of("key1", "val1", "key2", "val2"); + NameIdentifier modelIdent = NameIdentifier.of(schemaName, modelName); + + gravitinoCatalog.asModelCatalog().registerModel(modelIdent, null, null); + + gravitinoCatalog + .asModelCatalog() + .linkModelVersion(modelIdent, "uri", aliases, "comment", properties); + + ModelVersion modelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, "alias1"); + + Assertions.assertEquals(0, modelVersion.version()); + Assertions.assertEquals("uri", modelVersion.uri()); + Assertions.assertArrayEquals(aliases, modelVersion.aliases()); + Assertions.assertEquals("comment", modelVersion.comment()); + Assertions.assertEquals(properties, modelVersion.properties()); + + ModelVersionChange change = + ModelVersionChange.updateAliases( + new String[] {"alias1", "alias3"}, new String[] {"alias1", "alias2"}); + ModelVersion updatedModelVersion = + gravitinoCatalog.asModelCatalog().alterModelVersion(modelIdent, "alias1", change); + String[] newAliases = {"alias1", "alias3"}; + + Assertions.assertEquals(modelVersion.version(), updatedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), updatedModelVersion.uri()); + Assertions.assertArrayEquals(newAliases, updatedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), updatedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), updatedModelVersion.properties()); + + // reload model version and check aliases + Assertions.assertThrows( + NoSuchModelVersionException.class, + () -> gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, "alias2")); + Assertions.assertDoesNotThrow( + () -> gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, "alias1")); + Assertions.assertDoesNotThrow( + () -> gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, "alias3")); + ModelVersion reloadedModelVersion = + gravitinoCatalog.asModelCatalog().getModelVersion(modelIdent, "alias3"); + + Assertions.assertEquals(modelVersion.version(), reloadedModelVersion.version()); + Assertions.assertEquals(modelVersion.uri(), reloadedModelVersion.uri()); + Assertions.assertArrayEquals(newAliases, reloadedModelVersion.aliases()); + Assertions.assertEquals(modelVersion.comment(), reloadedModelVersion.comment()); + Assertions.assertEquals(modelVersion.properties(), reloadedModelVersion.properties()); } private void createMetalake() { diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoOptions.java b/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoOptions.java index 5fd861f078..ca017eb312 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoOptions.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoOptions.java @@ -69,6 +69,12 @@ public class GravitinoOptions { // TODO: temporary option for model version update, it will be refactored in the future, just // prove the E2E flow. public static final String NEW_URI = "newuri"; + // TODO: temporary option for model version update, it will be refactored in the future, just + // prove the E2E flow. + public static final String NEW_ALIAS = "newalias"; + // TODO: temporary option for model version update, it will be refactored in the future, just + // prove the E2E flow. + public static final String REMOVE_ALIAS = "removealias"; /** * Builds and returns the CLI options for Gravitino. @@ -122,6 +128,8 @@ public class GravitinoOptions { options.addOption(createArgsOption(null, ALIAS, "model aliases")); options.addOption(createArgOption(null, VERSION, "Gravitino client version")); options.addOption(createArgOption(null, NEW_URI, "New uri of a model version")); + options.addOption(createArgsOption(null, NEW_ALIAS, "New alias of a model version")); + options.addOption(createArgsOption(null, REMOVE_ALIAS, "Remove alias of a model version")); // Options that support multiple values options.addOption(createArgsOption("p", PROPERTIES, "property name/value pairs")); diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java index 7dec8bfb05..222e1d56e6 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java @@ -215,6 +215,19 @@ public class ModelCommandHandler extends CommandHandler { .validate() .handle(); } + + if (line.hasOption(GravitinoOptions.NEW_ALIAS) + && (line.hasOption(GravitinoOptions.ALIAS) || line.hasOption(GravitinoOptions.VERSION))) { + String[] newAliases = line.getOptionValues(GravitinoOptions.NEW_ALIAS); + Integer version = getVersionFromLine(line); + String alias = getAliasFromLine(line); + + gravitinoCommandLine + .newUpdateModelVersionAliases( + context, metalake, catalog, schema, model, version, alias, newAliases, new String[0]) + .validate() + .handle(); + } } /** Handles the "LIST" command. */ @@ -243,10 +256,31 @@ public class ModelCommandHandler extends CommandHandler { } } + /** Handles the "REMOVE" command. */ private void handleRemoveCommand() { String property = line.getOptionValue(GravitinoOptions.PROPERTY); - if (line.hasOption(GravitinoOptions.ALIAS) || line.hasOption(GravitinoOptions.VERSION)) { + if (line.hasOption(GravitinoOptions.REMOVE_ALIAS) + && (line.hasOption(GravitinoOptions.ALIAS) || line.hasOption(GravitinoOptions.VERSION))) { + String[] removeAliases = line.getOptionValues(GravitinoOptions.REMOVE_ALIAS); + Integer version = getVersionFromLine(line); + String alias = getAliasFromLine(line); + + gravitinoCommandLine + .newUpdateModelVersionAliases( + context, + metalake, + catalog, + schema, + model, + version, + alias, + new String[0], + removeAliases) + .validate() + .handle(); + + } else if (line.hasOption(GravitinoOptions.ALIAS) || line.hasOption(GravitinoOptions.VERSION)) { Integer version = getVersionFromLine(line); String alias = getAliasFromLine(line); @@ -255,8 +289,8 @@ public class ModelCommandHandler extends CommandHandler { context, metalake, catalog, schema, model, version, alias, property) .validate() .handle(); - } else { + } else { gravitinoCommandLine .newRemoveModelProperty(context, metalake, catalog, schema, model, property) .validate() diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java b/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java index 81453fc305..7b4c809d40 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java @@ -138,6 +138,7 @@ import org.apache.gravitino.cli.commands.UpdateMetalakeComment; import org.apache.gravitino.cli.commands.UpdateMetalakeName; import org.apache.gravitino.cli.commands.UpdateModelComment; import org.apache.gravitino.cli.commands.UpdateModelName; +import org.apache.gravitino.cli.commands.UpdateModelVersionAliases; import org.apache.gravitino.cli.commands.UpdateModelVersionComment; import org.apache.gravitino.cli.commands.UpdateModelVersionUri; import org.apache.gravitino.cli.commands.UpdateTableComment; @@ -926,6 +927,20 @@ public class TestableCommandLine { context, metalake, catalog, schema, model, version, alias, uri); } + protected UpdateModelVersionAliases newUpdateModelVersionAliases( + CommandContext context, + String metalake, + String catalog, + String schema, + String model, + Integer version, + String alias, + String[] aliasesToAdd, + String[] aliasesToRemove) { + return new UpdateModelVersionAliases( + context, metalake, catalog, schema, model, version, alias, aliasesToAdd, aliasesToRemove); + } + protected SetModelProperty newSetModelProperty( CommandContext context, String metalake, diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionAliases.java similarity index 72% copy from clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java copy to clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionAliases.java index 283848d322..9de73d0a8d 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionAliases.java @@ -25,22 +25,24 @@ import org.apache.gravitino.cli.ErrorMessages; import org.apache.gravitino.client.GravitinoClient; import org.apache.gravitino.exceptions.NoSuchCatalogException; import org.apache.gravitino.exceptions.NoSuchMetalakeException; +import org.apache.gravitino.exceptions.NoSuchModelException; +import org.apache.gravitino.exceptions.NoSuchModelVersionException; import org.apache.gravitino.exceptions.NoSuchSchemaException; -import org.apache.gravitino.exceptions.NoSuchTableException; import org.apache.gravitino.model.ModelVersionChange; -/** Update the comment of a model version. */ -public class UpdateModelVersionComment extends Command { +/** Update the aliases of a model version. */ +public class UpdateModelVersionAliases extends Command { protected final String metalake; protected final String catalog; protected final String schema; protected final String model; protected final Integer version; private final String alias; - private final String comment; + private final String[] aliasesToAdd; + private final String[] aliasesToRemove; /** - * Constructs a new {@link UpdateModelVersionComment} instance. + * Constructs a new {@link UpdateModelVersionAliases} instance with the specified parameters. * * @param context The command context. * @param metalake The name of the metalake. @@ -48,10 +50,11 @@ public class UpdateModelVersionComment extends Command { * @param schema The name of the schema. * @param model The name of the model. * @param version The version of the model. - * @param alias The alias of the model version. - * @param comment The new comment for the model version. + * @param alias The alias to update. + * @param aliasesToAdd The aliases to add. + * @param aliasesToRemove The aliases to remove. */ - public UpdateModelVersionComment( + public UpdateModelVersionAliases( CommandContext context, String metalake, String catalog, @@ -59,7 +62,8 @@ public class UpdateModelVersionComment extends Command { String model, Integer version, String alias, - String comment) { + String[] aliasesToAdd, + String[] aliasesToRemove) { super(context); this.metalake = metalake; this.catalog = catalog; @@ -67,16 +71,17 @@ public class UpdateModelVersionComment extends Command { this.model = model; this.version = version; this.alias = alias; - this.comment = comment; + this.aliasesToAdd = aliasesToAdd; + this.aliasesToRemove = aliasesToRemove; } - /** Update the comment of a model version. */ + /** Update the aliases of a model version. */ @Override public void handle() { try { NameIdentifier modelIdent = NameIdentifier.of(schema, model); GravitinoClient client = buildClient(metalake); - ModelVersionChange change = ModelVersionChange.updateComment(comment); + ModelVersionChange change = ModelVersionChange.updateAliases(aliasesToAdd, aliasesToRemove); if (alias != null) { client.loadCatalog(catalog).asModelCatalog().alterModelVersion(modelIdent, alias, change); } else { @@ -88,16 +93,18 @@ public class UpdateModelVersionComment extends Command { exitWithError(ErrorMessages.UNKNOWN_CATALOG); } catch (NoSuchSchemaException err) { exitWithError(ErrorMessages.UNKNOWN_SCHEMA); - } catch (NoSuchTableException err) { - exitWithError(ErrorMessages.UNKNOWN_TABLE); + } catch (NoSuchModelException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL); + } catch (NoSuchModelVersionException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL_VERSION); } catch (Exception exp) { exitWithError(exp.getMessage()); } if (alias != null) { - printInformation(model + " alias " + alias + " comment changed."); + printInformation(model + " alias " + alias + " aliases changed."); } else { - printInformation(model + " version " + version + " comment changed."); + printInformation(model + " version " + version + " aliases changed."); } } diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java index 283848d322..5aed4b2f17 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java @@ -25,8 +25,9 @@ import org.apache.gravitino.cli.ErrorMessages; import org.apache.gravitino.client.GravitinoClient; import org.apache.gravitino.exceptions.NoSuchCatalogException; import org.apache.gravitino.exceptions.NoSuchMetalakeException; +import org.apache.gravitino.exceptions.NoSuchModelException; +import org.apache.gravitino.exceptions.NoSuchModelVersionException; import org.apache.gravitino.exceptions.NoSuchSchemaException; -import org.apache.gravitino.exceptions.NoSuchTableException; import org.apache.gravitino.model.ModelVersionChange; /** Update the comment of a model version. */ @@ -88,8 +89,10 @@ public class UpdateModelVersionComment extends Command { exitWithError(ErrorMessages.UNKNOWN_CATALOG); } catch (NoSuchSchemaException err) { exitWithError(ErrorMessages.UNKNOWN_SCHEMA); - } catch (NoSuchTableException err) { - exitWithError(ErrorMessages.UNKNOWN_TABLE); + } catch (NoSuchModelException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL); + } catch (NoSuchModelVersionException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL_VERSION); } catch (Exception exp) { exitWithError(exp.getMessage()); } diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java index 0b5634bcb1..d4543159b1 100644 --- a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java @@ -25,8 +25,9 @@ import org.apache.gravitino.cli.ErrorMessages; import org.apache.gravitino.client.GravitinoClient; import org.apache.gravitino.exceptions.NoSuchCatalogException; import org.apache.gravitino.exceptions.NoSuchMetalakeException; +import org.apache.gravitino.exceptions.NoSuchModelException; +import org.apache.gravitino.exceptions.NoSuchModelVersionException; import org.apache.gravitino.exceptions.NoSuchSchemaException; -import org.apache.gravitino.exceptions.NoSuchTableException; import org.apache.gravitino.model.ModelVersionChange; /** Update the uri of a model version. */ @@ -88,8 +89,10 @@ public class UpdateModelVersionUri extends Command { exitWithError(ErrorMessages.UNKNOWN_CATALOG); } catch (NoSuchSchemaException err) { exitWithError(ErrorMessages.UNKNOWN_SCHEMA); - } catch (NoSuchTableException err) { - exitWithError(ErrorMessages.UNKNOWN_TABLE); + } catch (NoSuchModelException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL); + } catch (NoSuchModelVersionException err) { + exitWithError(ErrorMessages.UNKNOWN_MODEL_VERSION); } catch (Exception exp) { exitWithError(exp.getMessage()); } diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/TestModelCommands.java b/clients/cli/src/test/java/org/apache/gravitino/cli/TestModelCommands.java index 7919e85200..e0f1797e69 100644 --- a/clients/cli/src/test/java/org/apache/gravitino/cli/TestModelCommands.java +++ b/clients/cli/src/test/java/org/apache/gravitino/cli/TestModelCommands.java @@ -52,6 +52,7 @@ import org.apache.gravitino.cli.commands.SetModelProperty; import org.apache.gravitino.cli.commands.SetModelVersionProperty; import org.apache.gravitino.cli.commands.UpdateModelComment; import org.apache.gravitino.cli.commands.UpdateModelName; +import org.apache.gravitino.cli.commands.UpdateModelVersionAliases; import org.apache.gravitino.cli.commands.UpdateModelVersionComment; import org.apache.gravitino.cli.commands.UpdateModelVersionUri; import org.junit.jupiter.api.AfterEach; @@ -1053,4 +1054,96 @@ public class TestModelCommands { Assertions.assertThrows(RuntimeException.class, commandLine::handleCommandLine); } + + @Test + void testAddModelVersionAliases() { + UpdateModelVersionAliases mock = mock(UpdateModelVersionAliases.class); + when(mockCommandLine.hasOption(GravitinoOptions.METALAKE)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.METALAKE)).thenReturn("metalake_demo"); + when(mockCommandLine.hasOption(GravitinoOptions.NAME)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.NAME)).thenReturn("catalog.schema.model"); + when(mockCommandLine.hasOption(GravitinoOptions.VERSION)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.VERSION)).thenReturn("1"); + when(mockCommandLine.hasOption(GravitinoOptions.NEW_ALIAS)).thenReturn(true); + when(mockCommandLine.getOptionValues(GravitinoOptions.NEW_ALIAS)) + .thenReturn(new String[] {"aliasA", "aliasB"}); + + GravitinoCommandLine commandLine = + spy( + new GravitinoCommandLine( + mockCommandLine, mockOptions, CommandEntities.MODEL, CommandActions.UPDATE)); + + doReturn(mock) + .when(commandLine) + .newUpdateModelVersionAliases( + any(CommandContext.class), + eq("metalake_demo"), + eq("catalog"), + eq("schema"), + eq("model"), + any(), + any(), + any(), + any()); + doReturn(mock).when(mock).validate(); + commandLine.handleCommandLine(); + verify(mock).handle(); + } + + @Test + void testRemoveModelVersionAliases() { + UpdateModelVersionAliases mock = mock(UpdateModelVersionAliases.class); + when(mockCommandLine.hasOption(GravitinoOptions.METALAKE)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.METALAKE)).thenReturn("metalake_demo"); + when(mockCommandLine.hasOption(GravitinoOptions.NAME)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.NAME)).thenReturn("catalog.schema.model"); + when(mockCommandLine.hasOption(GravitinoOptions.VERSION)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.VERSION)).thenReturn("1"); + when(mockCommandLine.hasOption(GravitinoOptions.REMOVE_ALIAS)).thenReturn(true); + when(mockCommandLine.getOptionValues(GravitinoOptions.REMOVE_ALIAS)) + .thenReturn(new String[] {"aliasA", "aliasB"}); + + GravitinoCommandLine commandLine = + spy( + new GravitinoCommandLine( + mockCommandLine, mockOptions, CommandEntities.MODEL, CommandActions.REMOVE)); + + doReturn(mock) + .when(commandLine) + .newUpdateModelVersionAliases( + any(CommandContext.class), + eq("metalake_demo"), + eq("catalog"), + eq("schema"), + eq("model"), + any(), + any(), + any(), + any()); + doReturn(mock).when(mock).validate(); + commandLine.handleCommandLine(); + verify(mock).handle(); + } + + @Test + void testUpdateModelVersionAliasesByAliasAndVersion() { + Main.useExit = false; + when(mockCommandLine.hasOption(GravitinoOptions.METALAKE)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.METALAKE)).thenReturn("metalake_demo"); + when(mockCommandLine.hasOption(GravitinoOptions.ALIAS)).thenReturn(true); + when(mockCommandLine.getOptionValues(GravitinoOptions.ALIAS)) + .thenReturn(new String[] {"aliasA"}); + when(mockCommandLine.hasOption(GravitinoOptions.VERSION)).thenReturn(true); + when(mockCommandLine.getOptionValue(GravitinoOptions.VERSION)).thenReturn("1"); + when(mockCommandLine.hasOption(GravitinoOptions.NEW_ALIAS)).thenReturn(true); + when(mockCommandLine.getOptionValues(GravitinoOptions.NEW_ALIAS)) + .thenReturn(new String[] {"aliasA", "aliasB"}); + + GravitinoCommandLine commandLine = + spy( + new GravitinoCommandLine( + mockCommandLine, mockOptions, CommandEntities.MODEL, CommandActions.UPDATE)); + + Assertions.assertThrows(RuntimeException.class, commandLine::handleCommandLine); + } } diff --git a/clients/client-java/src/main/java/org/apache/gravitino/client/DTOConverters.java b/clients/client-java/src/main/java/org/apache/gravitino/client/DTOConverters.java index f911d4874e..b33c27a9e6 100644 --- a/clients/client-java/src/main/java/org/apache/gravitino/client/DTOConverters.java +++ b/clients/client-java/src/main/java/org/apache/gravitino/client/DTOConverters.java @@ -408,6 +408,12 @@ class DTOConverters { return new ModelVersionUpdateRequest.UpdateModelVersionUriRequest( ((ModelVersionChange.UpdateUri) change).newUri()); + } else if (change instanceof ModelVersionChange.UpdateAliases) { + ModelVersionChange.UpdateAliases updateAliases = (ModelVersionChange.UpdateAliases) change; + return new ModelVersionUpdateRequest.UpdateModelVersionAliasesRequest( + (updateAliases.aliasesToAdd().toArray(new String[0])), + updateAliases.aliasesToRemove().toArray(new String[0])); + } else { throw new IllegalArgumentException( "Unknown model version change type: " + change.getClass().getSimpleName()); diff --git a/clients/client-python/gravitino/api/model_version_change.py b/clients/client-python/gravitino/api/model_version_change.py index 4418b51f7f..6ca2aff18b 100644 --- a/clients/client-python/gravitino/api/model_version_change.py +++ b/clients/client-python/gravitino/api/model_version_change.py @@ -65,6 +65,17 @@ class ModelVersionChange(ABC): """ return ModelVersionChange.UpdateUri(uri) + @staticmethod + def update_aliases(aliases_to_add, aliases_to_remove): + """Creates a new model version change to update the aliases of the model version. + Args: + aliases_to_add: The new aliases to add to the model version. + aliases_to_remove: The aliases to remove from the model version. + Returns: + The model version change. + """ + return ModelVersionChange.UpdateAliases(aliases_to_add, aliases_to_remove) + class UpdateComment: """A model version change to update the comment of the model version.""" @@ -240,3 +251,60 @@ class ModelVersionChange(ABC): A string summary of this URI update operation. """ return f"UpdateUri {self._new_uri}" + + class UpdateAliases: + """A model version change to update the aliases of the model version.""" + + def __init__(self, aliases_to_add, aliases_to_remove): + self._aliases_to_add = set(aliases_to_add) if aliases_to_add else set() + self._aliases_to_remove = ( + set(aliases_to_remove) if aliases_to_remove else set() + ) + + def aliases_to_add(self): + """Retrieves the aliases to add. + Returns: + The aliases to add. + """ + return self._aliases_to_add + + def aliases_to_remove(self): + """Retrieves the aliases to remove. + Returns: + The aliases to remove. + """ + return self._aliases_to_remove + + def __eq__(self, other) -> bool: + """Compares this UpdateAliases instance with another object for equality. Two instances are + considered equal if they designate the same aliases to add and remove from the model version. + Args: + other: The object to compare with this instance. + Returns: + true if the given object represents an identical model version alias update operation; + false otherwise. + """ + if not isinstance(other, ModelVersionChange.UpdateAliases): + return False + return ( + self.aliases_to_add() == other.aliases_to_add() + and self.aliases_to_remove() == other.aliases_to_remove() + ) + + def __hash__(self) -> int: + """Generates a hash code for this UpdateAliases instance. The hash code is primarily based on + the aliases to add and remove from the model version. + Returns: + A hash code value for this alias update operation. + """ + return hash((self._aliases_to_add, self._aliases_to_remove)) + + def __str__(self) -> str: + """Returns a string representation of the UpdateAliases instance. This string includes the + class name followed by the aliases to add and remove for the model version. + Returns: + A string summary of this alias update operation. + """ + add_str = ", ".join(sorted(self._aliases_to_add)) + remove_str = ", ".join(sorted(self._aliases_to_remove)) + return f"UpdateAlias AliasToAdd: ({add_str}) AliasToRemove: ({remove_str})" diff --git a/clients/client-python/gravitino/client/generic_model_catalog.py b/clients/client-python/gravitino/client/generic_model_catalog.py index 318748781a..97fc75d538 100644 --- a/clients/client-python/gravitino/client/generic_model_catalog.py +++ b/clients/client-python/gravitino/client/generic_model_catalog.py @@ -550,6 +550,11 @@ class GenericModelCatalog(BaseSchemaCatalog): change.new_uri() ) + if isinstance(change, ModelVersionChange.UpdateAliases): + return ModelVersionUpdateRequest.ModelVersionAliasesRequest( + change.aliases_to_add(), change.aliases_to_remove() + ) + raise ValueError(f"Unknown change type: {type(change).__name__}") def _check_model_namespace(self, namespace: Namespace): diff --git a/clients/client-python/gravitino/dto/requests/model_version_update_request.py b/clients/client-python/gravitino/dto/requests/model_version_update_request.py index 73fd43fb0f..962d0634b3 100644 --- a/clients/client-python/gravitino/dto/requests/model_version_update_request.py +++ b/clients/client-python/gravitino/dto/requests/model_version_update_request.py @@ -18,7 +18,7 @@ from abc import abstractmethod from dataclasses import dataclass, field -from typing import Optional +from typing import Optional, Set from dataclasses_json import config @@ -138,3 +138,51 @@ class ModelVersionUpdateRequest: ModelVersionChange: The ModelVersionChange object representing the update uri operation. """ return ModelVersionChange.update_uri(self._new_uri) + + @dataclass + class ModelVersionAliasesRequest(ModelVersionUpdateRequestBase): + """Request to update model version aliases""" + + _aliases_to_add: Optional[Set[str]] = field( + metadata=config(field_name="aliasesToAdd") + ) + _aliases_to_remove: Optional[Set[str]] = field( + metadata=config(field_name="aliasesToRemove") + ) + + def __init__(self, aliases_to_add: Set[str], aliases_to_remove: Set[str]): + super().__init__("updateAliases") + self._aliases_to_add = aliases_to_add + self._aliases_to_remove = aliases_to_remove + + def aliases_to_add(self): + """Retrieves the new aliases of the model version. + Returns: + The new aliases of the model version. + """ + return self._aliases_to_add + + def aliases_to_remove(self): + """Retrieves the new aliases of the model version. + Returns: + The new aliases of the model version. + """ + return self._aliases_to_remove + + def validate(self): + """Validates the fields of the request.""" + if self._aliases_to_add is None and self._aliases_to_remove is None: + raise ValueError( + "At least one of aliasesToAdd or aliasesToRemove must be non-null" + ) + + def model_version_change(self): + """ + Returns a ModelVersionChange object representing the update aliases operation. + Returns: + ModelVersionChange: The ModelVersionChange object representing the update aliases operation. + """ + return ModelVersionChange.UpdateAliases( + self._aliases_to_add, + self._aliases_to_remove, + ) diff --git a/clients/client-python/gravitino/filesystem/gvfs_base_operations.py b/clients/client-python/gravitino/filesystem/gvfs_base_operations.py index dc835efbca..02eb1cca0a 100644 --- a/clients/client-python/gravitino/filesystem/gvfs_base_operations.py +++ b/clients/client-python/gravitino/filesystem/gvfs_base_operations.py @@ -377,9 +377,11 @@ class BaseGVFSOperations(ABC): fileset = catalog.as_fileset_catalog().load_fileset( NameIdentifier.of(fileset_ident.namespace().level(2), fileset_ident.name()) ) - target_location_name = location_name or fileset.properties().get( - fileset.PROPERTY_DEFAULT_LOCATION_NAME - ) or fileset.LOCATION_NAME_UNKNOWN + target_location_name = ( + location_name + or fileset.properties().get(fileset.PROPERTY_DEFAULT_LOCATION_NAME) + or fileset.LOCATION_NAME_UNKNOWN + ) actual_location = fileset.storage_locations().get(target_location_name) if actual_location is None: raise NoSuchLocationNameException( diff --git a/clients/client-python/tests/integration/test_model_catalog.py b/clients/client-python/tests/integration/test_model_catalog.py index 5d0eb9218a..bd7d31d53c 100644 --- a/clients/client-python/tests/integration/test_model_catalog.py +++ b/clients/client-python/tests/integration/test_model_catalog.py @@ -440,6 +440,51 @@ class TestModelCatalog(IntegrationTestEnv): self.assertEqual("comment", updated_model_version.comment()) self.assertEqual({"k1": "v1", "k2": "v2"}, updated_model_version.properties()) + def test_link_update_model_version_aliases(self): + model_name = f"model_it_model{str(randint(0, 1000))}" + model_ident = NameIdentifier.of(self._schema_name, model_name) + aliases = ["alias1", "alias2"] + comment = "comment" + properties = {"k1": "v1", "k2": "v2"} + self._catalog.as_model_catalog().register_model( + model_ident, comment, properties + ) + self._catalog.as_model_catalog().link_model_version( + model_ident, + uri="uri", + aliases=aliases, + comment="comment", + properties={"k1": "v1", "k2": "v2"}, + ) + + original_model_version = self._catalog.as_model_catalog().get_model_version( + model_ident, 0 + ) + + self.assertEqual(0, original_model_version.version()) + self.assertEqual("uri", original_model_version.uri()) + self.assertEqual(["alias1", "alias2"], original_model_version.aliases()) + self.assertEqual("comment", original_model_version.comment()) + self.assertEqual({"k1": "v1", "k2": "v2"}, original_model_version.properties()) + + # todo + changes = [ + ModelVersionChange.update_aliases( + ["alias2", "alias3"], + ["alias1", "alias2"], + ) + ] + self._catalog.as_model_catalog().alter_model_version(model_ident, 0, *changes) + + updated_model_version = self._catalog.as_model_catalog().get_model_version( + model_ident, 0 + ) + self.assertEqual(0, updated_model_version.version()) + self.assertEqual("uri", updated_model_version.uri()) + self.assertEqual(["alias2", "alias3"], updated_model_version.aliases()) + self.assertEqual("comment", updated_model_version.comment()) + self.assertEqual({"k1": "v1", "k2": "v2"}, updated_model_version.properties()) + def test_link_get_model_version(self): model_name = "model_it_model" + str(randint(0, 1000)) model_ident = NameIdentifier.of(self._schema_name, model_name) diff --git a/clients/client-python/tests/unittests/test_gravitino_version.py b/clients/client-python/tests/unittests/test_gravitino_version.py index 5cc0f299bf..9e70b17bc1 100644 --- a/clients/client-python/tests/unittests/test_gravitino_version.py +++ b/clients/client-python/tests/unittests/test_gravitino_version.py @@ -57,7 +57,9 @@ class TestGravitinoVersion(unittest.TestCase): self.assertEqual(version.patch, 0) version = GravitinoVersion( - VersionDTO("0.9.0-incubating-SNAPSHOT-0.8.0-SNAPSHOT", "2023-01-01", "1234567") + VersionDTO( + "0.9.0-incubating-SNAPSHOT-0.8.0-SNAPSHOT", "2023-01-01", "1234567" + ) ) self.assertEqual(version.major, 0) @@ -72,7 +74,9 @@ class TestGravitinoVersion(unittest.TestCase): self.assertEqual(version.patch, 0) # Test a valid the version string with alpha and hyphen separator - version = GravitinoVersion(VersionDTO("0.6.0-alpha-0.6.0-alpha", "2023-01-01", "1234567")) + version = GravitinoVersion( + VersionDTO("0.6.0-alpha-0.6.0-alpha", "2023-01-01", "1234567") + ) self.assertEqual(version.major, 0) self.assertEqual(version.minor, 6) @@ -86,7 +90,9 @@ class TestGravitinoVersion(unittest.TestCase): self.assertEqual(version.patch, 0) # Test a valid the version string with pypi format and hyphen separator - version = GravitinoVersion(VersionDTO("0.6.0.dev21-0.6.0.dev21", "2023-01-01", "1234567")) + version = GravitinoVersion( + VersionDTO("0.6.0.dev21-0.6.0.dev21", "2023-01-01", "1234567") + ) self.assertEqual(version.major, 0) self.assertEqual(version.minor, 6) @@ -166,7 +172,9 @@ class TestGravitinoVersion(unittest.TestCase): version1 = GravitinoVersion( VersionDTO("0.6.0-SNAPSHOT-0.6.0", "2023-01-01", "1234567") ) - version2 = GravitinoVersion(VersionDTO("0.6.0-0.6.0-SNAPSHOT", "2023-01-01", "1234567")) + version2 = GravitinoVersion( + VersionDTO("0.6.0-0.6.0-SNAPSHOT", "2023-01-01", "1234567") + ) self.assertEqual(version1, version2) diff --git a/common/src/main/java/org/apache/gravitino/dto/requests/ModelVersionUpdateRequest.java b/common/src/main/java/org/apache/gravitino/dto/requests/ModelVersionUpdateRequest.java index 96f7ebc70b..395dfd08c6 100644 --- a/common/src/main/java/org/apache/gravitino/dto/requests/ModelVersionUpdateRequest.java +++ b/common/src/main/java/org/apache/gravitino/dto/requests/ModelVersionUpdateRequest.java @@ -48,7 +48,10 @@ import org.apache.gravitino.rest.RESTRequest; name = "removeProperty"), @JsonSubTypes.Type( value = ModelVersionUpdateRequest.UpdateModelVersionUriRequest.class, - name = "updateUri") + name = "updateUri"), + @JsonSubTypes.Type( + value = ModelVersionUpdateRequest.UpdateModelVersionAliasesRequest.class, + name = "updateAliases") }) public interface ModelVersionUpdateRequest extends RESTRequest { @@ -167,4 +170,36 @@ public interface ModelVersionUpdateRequest extends RESTRequest { StringUtils.isNotBlank(newUri), "\"newUri\" field is required and cannot be empty"); } } + + /** Request to update the aliases of a model version. */ + @EqualsAndHashCode + @AllArgsConstructor + @NoArgsConstructor(force = true) + @ToString + @Getter + class UpdateModelVersionAliasesRequest implements ModelVersionUpdateRequest { + @JsonProperty("aliasesToAdd") + private final String[] aliasesToAdd; + + @JsonProperty("aliasesToRemove") + private final String[] aliasesToRemove; + + /** {@inheritDoc} */ + @Override + public ModelVersionChange modelVersionChange() { + return ModelVersionChange.updateAliases(aliasesToAdd, aliasesToRemove); + } + + /** + * Validates the request + * + * @throws IllegalArgumentException If the request is invalid, this exception is thrown. + */ + @Override + public void validate() throws IllegalArgumentException { + Preconditions.checkArgument( + aliasesToAdd != null || aliasesToRemove != null, + "At least one of aliasesToAdd or aliasesToRemove must be non-null"); + } + } } diff --git a/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java b/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java index 8370f6a1ae..f0408223c9 100644 --- a/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java +++ b/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java @@ -1073,7 +1073,7 @@ public class TestCatalogOperations (ModelVersionChange.UpdateAliases) change; Set<String> addTmpSet = updateAliasesChange.aliasesToAdd(); - Set<String> deleteTmpSet = updateAliasesChange.aliasesToDelete(); + Set<String> deleteTmpSet = updateAliasesChange.aliasesToRemove(); Set<String> aliasToAdd = Sets.difference(addTmpSet, deleteTmpSet).immutableCopy(); Set<String> aliasToDelete = Sets.difference(deleteTmpSet, addTmpSet).immutableCopy(); diff --git a/docs/manage-model-metadata-using-gravitino.md b/docs/manage-model-metadata-using-gravitino.md index 46893fa443..8b7d40cd1b 100644 --- a/docs/manage-model-metadata-using-gravitino.md +++ b/docs/manage-model-metadata-using-gravitino.md @@ -661,6 +661,16 @@ cat <<EOF >model.json { "@type": "removeProperty", "property": "key" + }, + { + "@type": "updateAliases", + "aliasesToAdd": [ + "alias1", + "alias2" + ], + "aliasesToRemove": [ + "alias3" + ] } ] } @@ -690,7 +700,8 @@ ModelVersionChange[] changes = { ModelVersionChange.updateComment("Updated comment of model version"), ModelVersionChange.updateUri("new_uri"), ModelVersionChange.setProperty("key", "value"), - ModelVersionChange.removeProperty("key") + ModelVersionChange.removeProperty("key"), + ModelVersionChange.updateAliases(new String[] {"alias1", "alias2"}, new String[] {"alias3"}) }; // Apply changes @@ -718,6 +729,7 @@ changes = ( ModelVersionChange.update_uri("new_uri"), ModelVersionChange.set_property("k2", "v2"), ModelVersionChange.remove_property("k1"), + ModelVersionChange.update_aliases(["alias1", "alias2"], ["alias3"]) ) # Apply changes @@ -731,12 +743,13 @@ updated_model = model_catalog.alter_model_version( #### Supported modifications -| Operation | JSON Example | Java Method | Python Method | -|---------------------|------------------------------------------------------------|---------------------------------------------------|----------------------------------------------------| -| **Update uri** | `{"@type":"updateUri","newName":"new_uri"}` | `ModelVersionChange.updateUri("new_uri")` | `ModelVersionChange.update_uri("new_uri")` | -| **Update comment** | `{"@type":"updateComment","newComment":"new_comment"}` | `ModelVersionChange.updateComment("new_comment")` | `ModelVersionChange.update_comment("new_comment")` | -| **Set property** | `{"@type":"setProperty","property":"key","value":"value"}` | `ModelVersionChange.setProperty("key", "value")` | `ModelVersionChange.set_property("key", "value")` | -| **Remove property** | `{"@type":"removeProperty","property":"key"}` | `ModelVersionChange.removeProperty("key")` | `ModelVersionChange.remove_property("key")` | +| Operation | JSON Example | Java Method | Python Method | +|---------------------|------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------| +| **Update uri** | `{"@type":"updateUri","newName":"new_uri"}` | `ModelVersionChange.updateUri("new_uri")` | `ModelVersionChange.update_uri("new_uri")` | +| **Update comment** | `{"@type":"updateComment","newComment":"new_comment"}` | `ModelVersionChange.updateComment("new_comment")` | `ModelVersionChange.update_comment("new_comment")` | +| **Set property** | `{"@type":"setProperty","property":"key","value":"value"}` | `ModelVersionChange.setProperty("key", "value")` | `ModelVersionChange.set_property("key", "value")` | +| **Remove property** | `{"@type":"removeProperty","property":"key"}` | `ModelVersionChange.removeProperty("key")` | `ModelVersionChange.remove_property("key")` | +| **Update Aliases** | `{"@type": "updateAliases","aliasesToAdd": ["alias1","alias2"],"aliasesToRemove": ["alias3"]}` | `ModelVersionChange.updateAliases(new String[] {"alias1", "alias2"}, new String[] {"alias3"})` | `ModelVersionChange.update_aliases(["alias1", "alias2"], ["alias3"])` | :::note - Multiple modifications can be applied in a single request. @@ -774,6 +787,16 @@ cat <<EOF >model.json { "@type": "removeProperty", "property": "key" + }, + { + "@type": "updateAliases", + "aliasesToAdd": [ + "alias1", + "alias2" + ], + "aliasesToRemove": [ + "alias3" + ] } ] } @@ -803,7 +826,8 @@ ModelVersionChange[] changes = { ModelVersionChange.updateComment("Updated comment of model version"), ModelVersionChange.updateUri("new_uri"), ModelVersionChange.setProperty("key", "value"), - ModelVersionChange.removeProperty("key") + ModelVersionChange.removeProperty("key"), + ModelVersionChange.updateAliases(new String[] {"alias1", "alias2"}, new String[] {"alias3"}) }; // Apply changes @@ -831,6 +855,7 @@ changes = ( ModelVersionChange.update_uri("new_uri"), ModelVersionChange.set_property("k2", "v2"), ModelVersionChange.remove_property("k1"), + ModelVersionChange.update_aliases(["alias1", "alias2"], ["alias3"]) ) # Apply changes @@ -844,12 +869,13 @@ updated_model = model_catalog.alter_model_version_by_alias( #### Supported modifications -| Operation | JSON Example | Java Method | Python Method | -|---------------------|------------------------------------------------------------|---------------------------------------------------|----------------------------------------------------| -| **Update uri** | `{"@type":"updateUri","newName":"new_uri"}` | `ModelVersionChange.updateUri("new_uri")` | `ModelVersionChange.update_uri("new_uri")` | -| **Update comment** | `{"@type":"updateComment","newComment":"new_comment"}` | `ModelVersionChange.updateComment("new_comment")` | `ModelVersionChange.update_comment("new_comment")` | -| **Set property** | `{"@type":"setProperty","property":"key","value":"value"}` | `ModelVersionChange.setProperty("key", "value")` | `ModelVersionChange.set_property("key", "value")` | -| **Remove property** | `{"@type":"removeProperty","property":"key"}` | `ModelVersionChange.removeProperty("key")` | `ModelVersionChange.remove_property("key")` | +| Operation | JSON Example | Java Method | Python Method | +|---------------------|------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------| +| **Update uri** | `{"@type":"updateUri","newName":"new_uri"}` | `ModelVersionChange.updateUri("new_uri")` | `ModelVersionChange.update_uri("new_uri")` | +| **Update comment** | `{"@type":"updateComment","newComment":"new_comment"}` | `ModelVersionChange.updateComment("new_comment")` | `ModelVersionChange.update_comment("new_comment")` | +| **Set property** | `{"@type":"setProperty","property":"key","value":"value"}` | `ModelVersionChange.setProperty("key", "value")` | `ModelVersionChange.set_property("key", "value")` | +| **Remove property** | `{"@type":"removeProperty","property":"key"}` | `ModelVersionChange.removeProperty("key")` | `ModelVersionChange.remove_property("key")` | +| **Update Aliases** | `{"@type": "updateAliases","aliasesToAdd": ["alias1","alias2"],"aliasesToRemove": ["alias3"]}` | `ModelVersionChange.updateAliases(new String[] {"alias1", "alias2"}, new String[] {"alias3"})` | `ModelVersionChange.update_aliases(["alias1", "alias2"], ["alias3"])` | :::note diff --git a/docs/open-api/models.yaml b/docs/open-api/models.yaml index e64ad7bfb6..d16c739e4e 100644 --- a/docs/open-api/models.yaml +++ b/docs/open-api/models.yaml @@ -570,6 +570,7 @@ components: - $ref: "#/components/schemas/SetModelVersionPropertyRequest" - $ref: "#/components/schemas/RemoveModelVersionPropertyRequest" - $ref: "#/components/schemas/UpdateModelVersionUriRequest" + - $ref: "#/components/schemas/UpdateModelVersionAliasesRequest" discriminator: propertyName: "@type" mapping: @@ -577,6 +578,7 @@ components: setProperty: "#/components/schemas/SetModelVersionPropertyRequest" removeProperty: "#/components/schemas/RemoveModelVersionPropertyRequest" updateUri: "#/components/schemas/UpdateModelVersionUriRequest" + updateAliases: "#/components/schemas/UpdateModelVersionAliasesRequest" UpdateModelVersionCommentRequest: type: object @@ -657,6 +659,35 @@ components: "newUri": "s3://path/to/model" } + UpdateModelVersionAliasesRequest: + type: object + required: + - "@type" + - aliasesToAdd + - aliasesToRemove + properties: + "@type": + type: string + description: The type of the update + enum: + - updateAliases + aliasesToAdd: + type: array + description: The aliases to add + nullable: true + items: + type: string + aliasesToRemove: + type: array + description: The aliases to remove + nullable: true + items: + type: string + example: { + "@type": "updateAliases", + "aliasesToAdd": ["alias1", "alias2"], + "aliasesToRemove": ["alias3", "alias4"] + } responses: ModelResponse:
