This is an automated email from the ASF dual-hosted git repository.
jshao 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 c37e02c21e [#6813] feat(server): Support update URI for model version
(#7010)
c37e02c21e is described below
commit c37e02c21e700b669a29a1d9bd6d13b59f0c67e9
Author: Lord of Abyss <[email protected]>
AuthorDate: Mon Apr 21 14:01:28 2025 +0800
[#6813] feat(server): Support update URI for model version (#7010)
### What changes were proposed in this pull request?
Support update URI for model version
- [X] PR1: Add ModelVersionChange API interface, Implement the update
uri logic in model catalog and JDBC backend logic, update related event.
- [ ] PR2: Add REST endpoint to support model version change, add Java
client and Python client for model version comment update.
### Why are the changes needed?
Fix: #6813
### Does this PR introduce _any_ user-facing change?
User can update the uri of a model version.
### How was this patch tested?
local ut test
---
.../apache/gravitino/model/ModelVersionChange.java | 71 ++++++++++++++++
.../gravitino/model/TestModelVersionChange.java | 44 ++++++++++
.../catalog/model/ModelCatalogOperations.java | 4 +
.../catalog/model/TestModelCatalogOperations.java | 99 ++++++++++++++++++++++
.../catalog/TestModelOperationDispatcher.java | 73 ++++++++++++++++
.../gravitino/connector/TestCatalogOperations.java | 4 +
.../service/TestModelVersionMetaService.java | 83 ++++++++++++++++++
7 files changed, 378 insertions(+)
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 9a35c28f85..ced922cffc 100644
--- a/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java
+++ b/api/src/main/java/org/apache/gravitino/model/ModelVersionChange.java
@@ -61,6 +61,16 @@ public interface ModelVersionChange {
return new ModelVersionChange.RemoveProperty(property);
}
+ /**
+ * Create a ModelVersionChange for updating the uri of a model version.
+ *
+ * @param newUri The new uri to be set for the model version.
+ * @return A new ModelVersionChange instance for updating the uri of a model
version.
+ */
+ static ModelVersionChange updateUri(String newUri) {
+ return new ModelVersionChange.UpdateUri(newUri);
+ }
+
/** A ModelVersionChange to update the model version comment. */
final class UpdateComment implements ModelVersionChange {
@@ -256,4 +266,65 @@ public interface ModelVersionChange {
return "REMOVEPROPERTY " + property;
}
}
+
+ /** A ModelVersionChange to update the uri of a model version. */
+ final class UpdateUri implements ModelVersionChange {
+ private final String newUri;
+
+ /**
+ * Creates a new {@link UpdateUri} instance with the specified new uri.
+ *
+ * @param newUri The new uri to be set for the model version.
+ */
+ public UpdateUri(String newUri) {
+ this.newUri = newUri;
+ }
+
+ /**
+ * Returns the new uri to be set for the model version.
+ *
+ * @return The new uri to be set for the model version.
+ */
+ public String newUri() {
+ return newUri;
+ }
+
+ /**
+ * Compares this UpdateUri instance with another object for equality. The
comparison is based on
+ * the new uri of the model version.
+ *
+ * @param obj The object to compare with this instance.
+ * @return {@code true} if the given object represents the same model
update operation; {@code
+ * false} otherwise.
+ */
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == this) return true;
+ if (!(obj instanceof UpdateUri)) return false;
+ UpdateUri other = (UpdateUri) obj;
+ return Objects.equals(newUri, other.newUri);
+ }
+
+ /**
+ * Generates a hash code for this UpdateUri instance. The hash code is
based on the new uri of
+ * the model.
+ *
+ * @return A hash code value for this model renaming operation.
+ */
+ @Override
+ public int hashCode() {
+ return Objects.hash(newUri);
+ }
+
+ /**
+ * Provides a string representation of the UpdateUri instance. This string
format includes the
+ * class name followed by the new uri to be set.
+ *
+ * @return A string summary of the UpdateUri instance.
+ */
+ @Override
+ public String toString() {
+ return "UpdateUri " + newUri;
+ }
+ }
}
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 f978437f27..c4e24110e3 100644
--- a/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java
+++ b/api/src/test/java/org/apache/gravitino/model/TestModelVersionChange.java
@@ -155,4 +155,48 @@ public class TestModelVersionChange {
Assertions.assertNotEquals(modelVersionChange1.hashCode(),
modelVersionChange3.hashCode());
Assertions.assertNotEquals(modelVersionChange2.hashCode(),
modelVersionChange3.hashCode());
}
+
+ @Test
+ void testUpdateUriChangeUseStaticMethod() {
+ String newUri = "S3://bucket/key";
+ ModelVersionChange modelVersionChange =
ModelVersionChange.updateUri(newUri);
+
+ Assertions.assertEquals(ModelVersionChange.UpdateUri.class,
modelVersionChange.getClass());
+
+ ModelVersionChange.UpdateUri updateUriChange =
+ (ModelVersionChange.UpdateUri) modelVersionChange;
+ Assertions.assertEquals(newUri, updateUriChange.newUri());
+ Assertions.assertEquals("UpdateUri " + newUri, updateUriChange.toString());
+ }
+
+ @Test
+ void testUpdateUriChangeUseConstructor() {
+ String newUri = "S3://bucket/key";
+ ModelVersionChange modelVersionChange = new
ModelVersionChange.UpdateUri(newUri);
+
+ Assertions.assertEquals(ModelVersionChange.UpdateUri.class,
modelVersionChange.getClass());
+
+ ModelVersionChange.UpdateUri updateUriChange =
+ (ModelVersionChange.UpdateUri) modelVersionChange;
+ Assertions.assertEquals(newUri, updateUriChange.newUri());
+ Assertions.assertEquals("UpdateUri " + newUri, updateUriChange.toString());
+ }
+
+ @Test
+ void testUpdateUriChangeEquals() {
+ String uri1 = "S3://bucket/key1";
+ String uri2 = "S3://bucket/key2";
+
+ ModelVersionChange modelVersionChange1 =
ModelVersionChange.updateUri(uri1);
+ ModelVersionChange modelVersionChange2 =
ModelVersionChange.updateUri(uri1);
+ ModelVersionChange modelVersionChange3 =
ModelVersionChange.updateUri(uri2);
+
+ Assertions.assertEquals(modelVersionChange1, modelVersionChange2);
+ Assertions.assertNotEquals(modelVersionChange1, modelVersionChange3);
+ Assertions.assertNotEquals(modelVersionChange2, modelVersionChange3);
+
+ Assertions.assertEquals(modelVersionChange1.hashCode(),
modelVersionChange2.hashCode());
+ Assertions.assertNotEquals(modelVersionChange1.hashCode(),
modelVersionChange3.hashCode());
+ Assertions.assertNotEquals(modelVersionChange2.hashCode(),
modelVersionChange3.hashCode());
+ }
}
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 d0c9b89cb6..87b0d12d45 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
@@ -425,6 +425,10 @@ public class ModelCatalogOperations extends
ManagedSchemaOperations
(ModelVersionChange.RemoveProperty) change;
doRemoveProperty(entityProperties, removePropertyChange);
+ } else if (change instanceof ModelVersionChange.UpdateUri) {
+ ModelVersionChange.UpdateUri updateUriChange =
(ModelVersionChange.UpdateUri) change;
+ entityUri = updateUriChange.newUri();
+
} else {
throw new IllegalArgumentException(
"Unsupported model version change: " +
change.getClass().getSimpleName());
diff --git
a/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/TestModelCatalogOperations.java
b/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/TestModelCatalogOperations.java
index c84e2fc9cf..77331cf673 100644
---
a/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/TestModelCatalogOperations.java
+++
b/catalogs/catalog-model/src/test/java/org/apache/gravtitino/catalog/model/TestModelCatalogOperations.java
@@ -1110,6 +1110,105 @@ public class TestModelCatalogOperations {
Assertions.assertEquals(newVersionProperties,
updatedModelVersion.properties());
}
+ @Test
+ void testUpdateModelVersionUri() {
+ String schemaName = randomSchemaName();
+ createSchema(schemaName);
+
+ String modelName = "model1";
+ String modelComment = "model1 comment";
+
+ String versionComment = "version1 comment";
+ String versionUri = "model_version_path";
+ String[] versionAliases = new String[] {"alias1", "alias2"};
+ String newVersionUri = "s3://bucket/path/to/new/version";
+
+ NameIdentifier modelIdent =
+ NameIdentifierUtil.ofModel(METALAKE_NAME, CATALOG_NAME, schemaName,
modelName);
+ StringIdentifier stringId = StringIdentifier.fromId(idGenerator.nextId());
+ Map<String, String> properties =
StringIdentifier.newPropertiesWithId(stringId, null);
+
+ ops.registerModel(modelIdent, modelComment, properties);
+ StringIdentifier versionId = StringIdentifier.fromId(idGenerator.nextId());
+ Map<String, String> versionProperties =
+ StringIdentifier.newPropertiesWithId(
+ versionId, ImmutableMap.of("key1", "value1", "key2", "value2"));
+
+ ops.linkModelVersion(modelIdent, versionUri, versionAliases,
versionComment, versionProperties);
+
+ // validate loaded model
+ Model loadedModel = ops.getModel(modelIdent);
+ Assertions.assertEquals(1, loadedModel.latestVersion());
+
+ // validate loaded version
+ ModelVersion loadedVersion = ops.getModelVersion(modelIdent, 0);
+ Assertions.assertEquals(0, loadedVersion.version());
+ Assertions.assertArrayEquals(versionAliases, loadedVersion.aliases());
+ Assertions.assertEquals(versionComment, loadedVersion.comment());
+ Assertions.assertEquals(versionUri, loadedVersion.uri());
+ Assertions.assertEquals(versionProperties, loadedVersion.properties());
+
+ // validate update version uri
+ ModelVersionChange updateUriChange =
ModelVersionChange.updateUri(newVersionUri);
+ ModelVersion updatedModelVersion = ops.alterModelVersion(modelIdent, 0,
updateUriChange);
+
+ Assertions.assertEquals(0, updatedModelVersion.version());
+ Assertions.assertEquals(newVersionUri, updatedModelVersion.uri());
+ Assertions.assertEquals(versionComment, updatedModelVersion.comment());
+ Assertions.assertArrayEquals(versionAliases,
updatedModelVersion.aliases());
+ Assertions.assertEquals(versionProperties,
updatedModelVersion.properties());
+ }
+
+ @Test
+ void testUpdateModelVersionUriByAlias() {
+ String schemaName = randomSchemaName();
+ createSchema(schemaName);
+
+ String modelName = "model1";
+ String modelComment = "model1 comment";
+
+ String versionComment = "version1 comment";
+ String versionUri = "model_version_path";
+ String[] versionAliases = new String[] {"alias1", "alias2"};
+ String newVersionUri = "s3://bucket/path/to/new/version";
+
+ NameIdentifier modelIdent =
+ NameIdentifierUtil.ofModel(METALAKE_NAME, CATALOG_NAME, schemaName,
modelName);
+ StringIdentifier stringId = StringIdentifier.fromId(idGenerator.nextId());
+ Map<String, String> properties =
StringIdentifier.newPropertiesWithId(stringId, null);
+
+ ops.registerModel(modelIdent, modelComment, properties);
+ StringIdentifier versionId = StringIdentifier.fromId(idGenerator.nextId());
+ Map<String, String> versionProperties =
+ StringIdentifier.newPropertiesWithId(
+ versionId, ImmutableMap.of("key1", "value1", "key2", "value2"));
+
+ ops.linkModelVersion(modelIdent, versionUri, versionAliases,
versionComment, versionProperties);
+
+ // validate loaded model
+ Model loadedModel = ops.getModel(modelIdent);
+ Assertions.assertEquals(1, loadedModel.latestVersion());
+
+ // validate loaded version
+ ModelVersion loadedVersion = ops.getModelVersion(modelIdent,
versionAliases[0]);
+ Assertions.assertEquals(0, loadedVersion.version());
+ Assertions.assertArrayEquals(versionAliases, loadedVersion.aliases());
+ Assertions.assertEquals(versionComment, loadedVersion.comment());
+ Assertions.assertEquals(versionUri, loadedVersion.uri());
+ Assertions.assertEquals(versionProperties, loadedVersion.properties());
+
+ // validate update version uri
+ ModelVersionChange updateUriChange =
ModelVersionChange.updateUri(newVersionUri);
+ ModelVersion updatedModelVersion =
+ ops.alterModelVersion(modelIdent, versionAliases[0], updateUriChange);
+
+ Assertions.assertEquals(0, updatedModelVersion.version());
+ Assertions.assertEquals(newVersionUri, updatedModelVersion.uri());
+ Assertions.assertEquals(versionComment, updatedModelVersion.comment());
+ Assertions.assertArrayEquals(versionAliases,
updatedModelVersion.aliases());
+ Assertions.assertEquals(versionProperties,
updatedModelVersion.properties());
+ }
+
private String randomSchemaName() {
return "schema_" + UUID.randomUUID().toString().replace("-", "");
}
diff --git
a/core/src/test/java/org/apache/gravitino/catalog/TestModelOperationDispatcher.java
b/core/src/test/java/org/apache/gravitino/catalog/TestModelOperationDispatcher.java
index 5e27a77ef1..f059b857c1 100644
---
a/core/src/test/java/org/apache/gravitino/catalog/TestModelOperationDispatcher.java
+++
b/core/src/test/java/org/apache/gravitino/catalog/TestModelOperationDispatcher.java
@@ -554,6 +554,79 @@ public class TestModelOperationDispatcher extends
TestOperationDispatcher {
Assertions.assertEquals(newProps, alteredModelVersion.properties());
}
+ @Test
+ void testUpdateModelVersionUri() {
+ String schemaName = randomSchemaName();
+ String schemaComment = "schema which tests update";
+
+ String modelName = randomModelName();
+ String modelComment = "model which tests update";
+ Map<String, String> props = ImmutableMap.of("k1", "v1", "k2", "v2");
+ String newUri = "s3://new-bucket/new-path/new-model.json";
+
+ String versionUri = "s3://test-bucket/test-path/model.json";
+ String[] versionAliases = {"alias1", "alias2"};
+ String versionComment = "version which tests update";
+
+ NameIdentifier schemaIdent = NameIdentifier.of(metalake, catalog,
schemaName);
+ schemaOperationDispatcher.createSchema(schemaIdent, schemaComment, props);
+
+ NameIdentifier modelIdent =
+ NameIdentifierUtil.ofModel(metalake, catalog, schemaName, modelName);
+ modelOperationDispatcher.registerModel(modelIdent, modelComment, props);
+
+ modelOperationDispatcher.linkModelVersion(
+ modelIdent, versionUri, versionAliases, versionComment, props);
+
+ ModelVersionChange change = ModelVersionChange.updateUri(newUri);
+ ModelVersion modelVersion =
modelOperationDispatcher.getModelVersion(modelIdent, 0);
+ ModelVersion alteredModelVersion =
+ modelOperationDispatcher.alterModelVersion(modelIdent, 0, change);
+
+ Assertions.assertEquals(newUri, alteredModelVersion.uri());
+ Assertions.assertEquals(modelVersion.version(),
alteredModelVersion.version());
+ Assertions.assertEquals(modelVersion.aliases(),
alteredModelVersion.aliases());
+ Assertions.assertEquals(modelVersion.comment(),
alteredModelVersion.comment());
+ Assertions.assertEquals(modelVersion.properties(),
alteredModelVersion.properties());
+ }
+
+ @Test
+ void testUpdateModelVersionUriByAlias() {
+ String schemaName = randomSchemaName();
+ String schemaComment = "schema which tests update";
+
+ String modelName = randomModelName();
+ String modelComment = "model which tests update";
+ Map<String, String> props = ImmutableMap.of("k1", "v1", "k2", "v2");
+ String newUri = "s3://new-bucket/new-path/new-model.json";
+
+ String versionUri = "s3://test-bucket/test-path/model.json";
+ String[] versionAliases = {"alias1", "alias2"};
+ String versionComment = "version which tests update";
+
+ NameIdentifier schemaIdent = NameIdentifier.of(metalake, catalog,
schemaName);
+ schemaOperationDispatcher.createSchema(schemaIdent, schemaComment, props);
+
+ NameIdentifier modelIdent =
+ NameIdentifierUtil.ofModel(metalake, catalog, schemaName, modelName);
+ modelOperationDispatcher.registerModel(modelIdent, modelComment, props);
+
+ modelOperationDispatcher.linkModelVersion(
+ modelIdent, versionUri, versionAliases, versionComment, props);
+
+ ModelVersionChange change = ModelVersionChange.updateUri(newUri);
+ ModelVersion modelVersion =
+ modelOperationDispatcher.getModelVersion(modelIdent,
versionAliases[0]);
+ ModelVersion alteredModelVersion =
+ modelOperationDispatcher.alterModelVersion(modelIdent,
versionAliases[0], change);
+
+ Assertions.assertEquals(newUri, alteredModelVersion.uri());
+ Assertions.assertEquals(modelVersion.version(),
alteredModelVersion.version());
+ Assertions.assertEquals(modelVersion.aliases(),
alteredModelVersion.aliases());
+ Assertions.assertEquals(modelVersion.comment(),
alteredModelVersion.comment());
+ Assertions.assertEquals(modelVersion.properties(),
alteredModelVersion.properties());
+ }
+
private String randomSchemaName() {
return "schema_" + UUID.randomUUID().toString().replace("-", "");
}
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 94dcd420f0..a0105e0878 100644
---
a/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java
+++
b/core/src/test/java/org/apache/gravitino/connector/TestCatalogOperations.java
@@ -1055,6 +1055,10 @@ public class TestCatalogOperations
ModelVersionChange.SetProperty setProperty =
(ModelVersionChange.SetProperty) change;
newProps.put(setProperty.property(), setProperty.value());
+ } else if (change instanceof ModelVersionChange.UpdateUri) {
+ ModelVersionChange.UpdateUri updateUriChange =
(ModelVersionChange.UpdateUri) change;
+ newUri = updateUriChange.newUri();
+
} else {
throw new IllegalArgumentException("Unsupported model change: " +
change);
}
diff --git
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestModelVersionMetaService.java
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestModelVersionMetaService.java
index 4284537d76..38bfc1e94e 100644
---
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestModelVersionMetaService.java
+++
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestModelVersionMetaService.java
@@ -632,6 +632,89 @@ public class TestModelVersionMetaService extends
TestJDBCBackend {
updatePropertiesUpdater));
}
+ @Test
+ void testUpdateModelVersionUri() throws IOException {
+ createParentEntities(METALAKE_NAME, CATALOG_NAME, SCHEMA_NAME, auditInfo);
+
+ Map<String, String> properties = ImmutableMap.of("k1", "v1", "k2", "v2");
+ String modelName = randomModelName();
+ String modelComment = "model1 comment";
+ String modelVersionUri = "S3://test/path/to/model/version";
+ List<String> modelVersionAliases = ImmutableList.of("alias1", "alias2");
+ String modelVersionComment = "test comment";
+ String updatedUri = "S3://test/path/to/new/model/version";
+ int version = 0;
+
+ ModelEntity modelEntity =
+ createModelEntity(
+ RandomIdGenerator.INSTANCE.nextId(),
+ MODEL_NS,
+ modelName,
+ modelComment,
+ 0,
+ properties,
+ auditInfo);
+
+ ModelVersionEntity modelVersionEntity =
+ createModelVersionEntity(
+ modelEntity.nameIdentifier(),
+ version,
+ modelVersionUri,
+ modelVersionAliases,
+ modelVersionComment,
+ properties,
+ auditInfo);
+
+ ModelVersionEntity updatedModelVersionEntity =
+ createModelVersionEntity(
+ modelVersionEntity.modelIdentifier(),
+ modelVersionEntity.version(),
+ updatedUri,
+ modelVersionEntity.aliases(),
+ modelVersionEntity.comment(),
+ modelVersionEntity.properties(),
+ modelVersionEntity.auditInfo());
+
+ Assertions.assertDoesNotThrow(
+ () -> ModelMetaService.getInstance().insertModel(modelEntity, false));
+
+ Assertions.assertDoesNotThrow(
+ () ->
ModelVersionMetaService.getInstance().insertModelVersion(modelVersionEntity));
+
+ Function<ModelVersionEntity, ModelVersionEntity> updatePropertiesUpdater =
+ oldModelVersionEntity -> updatedModelVersionEntity;
+
+ ModelVersionEntity alteredModelVersionEntity =
+ ModelVersionMetaService.getInstance()
+ .updateModelVersion(modelVersionEntity.nameIdentifier(),
updatePropertiesUpdater);
+
+ Assertions.assertEquals(updatedModelVersionEntity,
alteredModelVersionEntity);
+
+ // Test update a non-exist model
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () ->
+ ModelVersionMetaService.getInstance()
+ .updateModelVersion(
+ NameIdentifierUtil.ofModelVersion(
+ METALAKE_NAME,
+ CATALOG_NAME,
+ SCHEMA_NAME,
+ "non_exist_model",
+ "non_exist_version"),
+ updatePropertiesUpdater));
+
+ // Test update a non-exist model version
+ Assertions.assertThrows(
+ NoSuchEntityException.class,
+ () ->
+ ModelVersionMetaService.getInstance()
+ .updateModelVersion(
+ NameIdentifierUtil.ofModelVersion(
+ METALAKE_NAME, CATALOG_NAME, SCHEMA_NAME, modelName,
"non_exist_version"),
+ updatePropertiesUpdater));
+ }
+
private NameIdentifier getModelVersionIdent(NameIdentifier modelIdent, int
version) {
List<String> parts = Lists.newArrayList(modelIdent.namespace().levels());
parts.add(modelIdent.name());
