This is an automated email from the ASF dual-hosted git repository.
jmclean 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 bfb8568066 [#5960] fix(CLI): Add register and link commands to CLI for
model (#6066)
bfb8568066 is described below
commit bfb85680668b91d68f8190a8247156477f326039
Author: Lord of Abyss <[email protected]>
AuthorDate: Sat Jan 4 06:25:18 2025 +0800
[#5960] fix(CLI): Add register and link commands to CLI for model (#6066)
### What changes were proposed in this pull request?
Add register and link commands to CLI for model
- register a model：`model create`
- link a model：`model update <—uri uri> [--alias aliaA aliaB]`
meantime, add two options
- `—uri` ：The URI of the model version artifact.
- `—alias` ：The aliases of the model version.
The documentation will be updated after #6047 merge and I will create a
new issue.
### Why are the changes needed?
Fix: #5960
### Does this PR introduce _any_ user-facing change?
NO
### How was this patch tested?
#### register test
```bash
# register a model
gcli model create -m demo_metalake --name Hive_catalog.default.model
# register a model with comment
gcli model create -m demo_metalake --name Hive_catalog.default.model
--comment comment
# register a model with properties
gcli model create -m demo_metalake --name Hive_catalog.default.model
--properties key1=val1 key2=val2
# register a model with properties and comment
gcli model create -m demo_metalake --name Hive_catalog.default.model
--properties key1=val1 klinkey2=val2 --comment comment
```
#### link test
```bash
# link a model
gcli model update -m demo_metalake --name Hive_catalog.default.model --uri
file:///tmp/file
# link a model with alias
gcli model update -m demo_metalake --name Hive_catalog.default.model --uri
file:///tmp/file --alias aliasA aliasB
# link a model with all component
gcli model update -m demo_metalake --name Hive_catalog.default.model --uri
file:///tmp/file --alias aliasA aliasB --comment comment --properties
key1=val1 key2=val2
# link a model without uri
gcli model update -m demo_metalake --name Hive_catalog.default.model
```
---
.../org/apache/gravitino/cli/ErrorMessages.java | 3 +-
.../apache/gravitino/cli/GravitinoCommandLine.java | 34 +++
.../org/apache/gravitino/cli/GravitinoOptions.java | 6 +
.../apache/gravitino/cli/TestableCommandLine.java | 29 +++
.../apache/gravitino/cli/commands/LinkModel.java | 106 ++++++++
.../gravitino/cli/commands/RegisterModel.java | 103 ++++++++
.../apache/gravitino/cli/TestModelCommands.java | 284 +++++++++++++++++++++
7 files changed, 564 insertions(+), 1 deletion(-)
diff --git
a/clients/cli/src/main/java/org/apache/gravitino/cli/ErrorMessages.java
b/clients/cli/src/main/java/org/apache/gravitino/cli/ErrorMessages.java
index 084b5c34c8..e90c525963 100644
--- a/clients/cli/src/main/java/org/apache/gravitino/cli/ErrorMessages.java
+++ b/clients/cli/src/main/java/org/apache/gravitino/cli/ErrorMessages.java
@@ -36,6 +36,7 @@ public class ErrorMessages {
public static final String MISSING_USER = "Missing --user option.";
public static final String MISSING_ROLE = "Missing --role option.";
public static final String MISSING_TAG = "Missing --tag option.";
+ public static final String MISSING_URI = "Missing --uri option.";
public static final String METALAKE_EXISTS = "Metalake already exists.";
public static final String CATALOG_EXISTS = "Catalog already exists.";
public static final String SCHEMA_EXISTS = "Schema already exists.";
@@ -51,13 +52,13 @@ public class ErrorMessages {
public static final String COLUMN_EXISTS = "Column already exists.";
public static final String UNKNOWN_TOPIC = "Unknown topic.";
public static final String TOPIC_EXISTS = "Topic already exists.";
+ public static final String MODEL_EXISTS = "Model already exists.";
public static final String UNKNOWN_FILESET = "Unknown fileset.";
public static final String FILESET_EXISTS = "Fileset already exists.";
public static final String TAG_EMPTY = "Error: Must configure --tag option.";
public static final String UNKNOWN_ROLE = "Unknown role.";
public static final String ROLE_EXISTS = "Role already exists.";
public static final String TABLE_EXISTS = "Table already exists.";
- public static final String MODEL_EXISTS = "Model already exists.";
public static final String INVALID_SET_COMMAND =
"Unsupported combination of options either use --name, --user, --group
or --property and --value.";
public static final String INVALID_REMOVE_COMMAND =
diff --git
a/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoCommandLine.java
b/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoCommandLine.java
index c23fb8b7cd..3a9322d010 100644
---
a/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoCommandLine.java
+++
b/clients/cli/src/main/java/org/apache/gravitino/cli/GravitinoCommandLine.java
@@ -1192,6 +1192,40 @@ public class GravitinoCommandLine extends
TestableCommandLine {
}
break;
+ case CommandActions.CREATE:
+ String createComment = line.getOptionValue(GravitinoOptions.COMMENT);
+ String[] createProperties =
line.getOptionValues(GravitinoOptions.PROPERTIES);
+ Map<String, String> createPropertyMap = new
Properties().parse(createProperties);
+ newCreateModel(
+ url, ignore, metalake, catalog, schema, model, createComment,
createPropertyMap)
+ .handle();
+ break;
+
+ case CommandActions.UPDATE:
+ String[] alias = line.getOptionValues(GravitinoOptions.ALIAS);
+ String uri = line.getOptionValue(GravitinoOptions.URI);
+ if (uri == null) {
+ System.err.println(ErrorMessages.MISSING_URI);
+ Main.exit(-1);
+ }
+
+ String linkComment = line.getOptionValue(GravitinoOptions.COMMENT);
+ String[] linkProperties =
line.getOptionValues(CommandActions.PROPERTIES);
+ Map<String, String> linkPropertityMap = new
Properties().parse(linkProperties);
+ newLinkModel(
+ url,
+ ignore,
+ metalake,
+ catalog,
+ schema,
+ model,
+ uri,
+ alias,
+ linkComment,
+ linkPropertityMap)
+ .handle();
+ break;
+
default:
System.err.println(ErrorMessages.UNSUPPORTED_ACTION);
break;
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 657566036d..aaeb8f0184 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
@@ -62,6 +62,8 @@ public class GravitinoOptions {
public static final String ALL = "all";
public static final String ENABLE = "enable";
public static final String DISABLE = "disable";
+ public static final String ALIAS = "alias";
+ public static final String URI = "uri";
/**
* Builds and returns the CLI options for Gravitino.
@@ -109,6 +111,10 @@ public class GravitinoOptions {
options.addOption(createArgOption(COLUMNFILE, "CSV file describing
columns"));
options.addOption(createSimpleOption(null, ALL, "all operation for
--enable"));
+ // model options
+ options.addOption(createArgOption(null, URI, "model version artifact"));
+ options.addOption(createArgsOption(null, ALIAS, "model aliases"));
+
// Options that support multiple values
options.addOption(createArgsOption("p", PROPERTIES, "property name/value
pairs"));
options.addOption(createArgsOption("t", TAG, "tag name"));
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 6a46874917..8df9498d97 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
@@ -55,6 +55,7 @@ import org.apache.gravitino.cli.commands.FilesetDetails;
import org.apache.gravitino.cli.commands.GrantPrivilegesToRole;
import org.apache.gravitino.cli.commands.GroupAudit;
import org.apache.gravitino.cli.commands.GroupDetails;
+import org.apache.gravitino.cli.commands.LinkModel;
import org.apache.gravitino.cli.commands.ListAllTags;
import org.apache.gravitino.cli.commands.ListCatalogProperties;
import org.apache.gravitino.cli.commands.ListCatalogs;
@@ -83,6 +84,7 @@ import org.apache.gravitino.cli.commands.MetalakeEnable;
import org.apache.gravitino.cli.commands.ModelAudit;
import org.apache.gravitino.cli.commands.ModelDetails;
import org.apache.gravitino.cli.commands.OwnerDetails;
+import org.apache.gravitino.cli.commands.RegisterModel;
import org.apache.gravitino.cli.commands.RemoveAllTags;
import org.apache.gravitino.cli.commands.RemoveCatalogProperty;
import org.apache.gravitino.cli.commands.RemoveFilesetProperty;
@@ -925,4 +927,31 @@ public class TestableCommandLine {
String url, boolean ignore, String metalake, String catalog, String
schema, String model) {
return new ModelDetails(url, ignore, metalake, catalog, schema, model);
}
+
+ protected RegisterModel newCreateModel(
+ String url,
+ boolean ignore,
+ String metalake,
+ String catalog,
+ String schema,
+ String model,
+ String comment,
+ Map<String, String> properties) {
+ return new RegisterModel(url, ignore, metalake, catalog, schema, model,
comment, properties);
+ }
+
+ protected LinkModel newLinkModel(
+ String url,
+ boolean ignore,
+ String metalake,
+ String catalog,
+ String schema,
+ String model,
+ String uri,
+ String[] alias,
+ String comment,
+ Map<String, String> properties) {
+ return new LinkModel(
+ url, ignore, metalake, catalog, schema, model, uri, alias, comment,
properties);
+ }
}
diff --git
a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/LinkModel.java
b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/LinkModel.java
new file mode 100644
index 0000000000..6e8a4ffb76
--- /dev/null
+++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/LinkModel.java
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.cli.commands;
+
+/** Link a new model version to the registered model. */
+import java.util.Arrays;
+import java.util.Map;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.cli.ErrorMessages;
+import org.apache.gravitino.client.GravitinoClient;
+import
org.apache.gravitino.exceptions.ModelVersionAliasesAlreadyExistException;
+import org.apache.gravitino.exceptions.NoSuchCatalogException;
+import org.apache.gravitino.exceptions.NoSuchMetalakeException;
+import org.apache.gravitino.exceptions.NoSuchModelException;
+import org.apache.gravitino.exceptions.NoSuchSchemaException;
+import org.apache.gravitino.model.ModelCatalog;
+
+public class LinkModel extends Command {
+ protected final String metalake;
+ protected final String catalog;
+ protected final String schema;
+ protected final String model;
+ protected final String uri;
+ protected final String[] alias;
+ protected final String comment;
+ protected final Map<String, String> properties;
+
+ /**
+ * Link a new model version to the registered model.
+ *
+ * @param url The URL of the Gravitino server.
+ * @param ignoreVersions If true don't check the client/server versions
match.
+ * @param metalake The name of the metalake.
+ * @param catalog The name of the catalog.
+ * @param schema The name of schema.
+ * @param model The name of model.
+ * @param uri The URI of the model version artifact.
+ * @param alias The aliases of the model version.
+ * @param comment The comment of the model version.
+ * @param properties The properties of the model version.
+ */
+ public LinkModel(
+ String url,
+ boolean ignoreVersions,
+ String metalake,
+ String catalog,
+ String schema,
+ String model,
+ String uri,
+ String[] alias,
+ String comment,
+ Map<String, String> properties) {
+ super(url, ignoreVersions);
+ this.metalake = metalake;
+ this.catalog = catalog;
+ this.schema = schema;
+ this.model = model;
+ this.uri = uri;
+ this.alias = alias;
+ this.comment = comment;
+ this.properties = properties;
+ }
+
+ /** Link a new model version to the registered model. */
+ @Override
+ public void handle() {
+ NameIdentifier name = NameIdentifier.of(schema, model);
+
+ try {
+ GravitinoClient client = buildClient(metalake);
+ ModelCatalog modelCatalog = client.loadCatalog(catalog).asModelCatalog();
+ modelCatalog.linkModelVersion(name, uri, alias, comment, properties);
+ } catch (NoSuchMetalakeException err) {
+ exitWithError(ErrorMessages.UNKNOWN_METALAKE);
+ } catch (NoSuchCatalogException err) {
+ exitWithError(ErrorMessages.UNKNOWN_CATALOG);
+ } catch (NoSuchSchemaException err) {
+ exitWithError(ErrorMessages.UNKNOWN_SCHEMA);
+ } catch (NoSuchModelException err) {
+ exitWithError(ErrorMessages.UNKNOWN_MODEL);
+ } catch (ModelVersionAliasesAlreadyExistException err) {
+ exitWithError(Arrays.toString(alias) + " already exist.");
+ } catch (Exception err) {
+ exitWithError(err.getMessage());
+ }
+
+ System.out.println(
+ "Linked model " + model + " to " + uri + " with aliases " +
Arrays.toString(alias));
+ }
+}
diff --git
a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RegisterModel.java
b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RegisterModel.java
new file mode 100644
index 0000000000..d50dbed50e
--- /dev/null
+++
b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RegisterModel.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.cli.commands;
+
+import java.util.Map;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.cli.ErrorMessages;
+import org.apache.gravitino.cli.Main;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.exceptions.ModelAlreadyExistsException;
+import org.apache.gravitino.exceptions.NoSuchCatalogException;
+import org.apache.gravitino.exceptions.NoSuchMetalakeException;
+import org.apache.gravitino.exceptions.NoSuchSchemaException;
+import org.apache.gravitino.model.Model;
+import org.apache.gravitino.model.ModelCatalog;
+
+/** Register a model in the catalog */
+public class RegisterModel extends Command {
+
+ protected final String metalake;
+ protected final String catalog;
+ protected final String schema;
+ protected final String model;
+ protected final String comment;
+ protected final Map<String, String> properties;
+
+ /**
+ * Register a model in the catalog
+ *
+ * @param url The URL of the Gravitino server.
+ * @param ignoreVersions If true don't check the client/server versions
match.
+ * @param metalake The name of the metalake.
+ * @param catalog The name of the catalog.
+ * @param schema The name of schema.
+ * @param model The name of model.
+ * @param comment The comment of the model version.
+ * @param properties The properties of the model version.
+ */
+ public RegisterModel(
+ String url,
+ boolean ignoreVersions,
+ String metalake,
+ String catalog,
+ String schema,
+ String model,
+ String comment,
+ Map<String, String> properties) {
+ super(url, ignoreVersions);
+ this.metalake = metalake;
+ this.catalog = catalog;
+ this.schema = schema;
+ this.model = model;
+ this.comment = comment;
+ this.properties = properties;
+ }
+
+ /** Register a model in the catalog */
+ @Override
+ public void handle() {
+ NameIdentifier name = NameIdentifier.of(schema, model);
+ Model registeredModel = null;
+
+ try {
+ GravitinoClient client = buildClient(metalake);
+ ModelCatalog modelCatalog = client.loadCatalog(catalog).asModelCatalog();
+ registeredModel = modelCatalog.registerModel(name, comment, properties);
+ } catch (NoSuchMetalakeException err) {
+ exitWithError(ErrorMessages.UNKNOWN_METALAKE);
+ } catch (NoSuchCatalogException err) {
+ exitWithError(ErrorMessages.UNKNOWN_CATALOG);
+ } catch (NoSuchSchemaException err) {
+ exitWithError(ErrorMessages.UNKNOWN_SCHEMA);
+ } catch (ModelAlreadyExistsException err) {
+ exitWithError(ErrorMessages.MODEL_EXISTS);
+ } catch (Exception err) {
+ exitWithError(err.getMessage());
+ }
+
+ if (registeredModel != null) {
+ System.out.println("Successful register " + registeredModel.name() +
".");
+ } else {
+ System.err.println("Failed to register model: " + model + ".");
+ Main.exit(-1);
+ }
+ }
+}
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 e486c41a9d..391201f292 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
@@ -21,6 +21,7 @@ package org.apache.gravitino.cli;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.isNull;
import static org.mockito.Mockito.doReturn;
@@ -35,11 +36,14 @@ import java.io.PrintStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collections;
+import java.util.Map;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Options;
+import org.apache.gravitino.cli.commands.LinkModel;
import org.apache.gravitino.cli.commands.ListModel;
import org.apache.gravitino.cli.commands.ModelAudit;
import org.apache.gravitino.cli.commands.ModelDetails;
+import org.apache.gravitino.cli.commands.RegisterModel;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -289,4 +293,284 @@ public class TestModelCommands {
commandLine.handleCommandLine();
verify(mockAudit).handle();
}
+
+ @Test
+ void testRegisterModelCommand() {
+ RegisterModel mockCreate = mock(RegisterModel.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.PROPERTIES)).thenReturn(false);
+
when(mockCommandLine.hasOption(GravitinoOptions.COMMENT)).thenReturn(false);
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.CREATE));
+
+ doReturn(mockCreate)
+ .when(commandLine)
+ .newCreateModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ isNull(),
+ argThat(Map::isEmpty));
+ commandLine.handleCommandLine();
+ verify(mockCreate).handle();
+ }
+
+ @Test
+ void testRegisterModelCommandWithComment() {
+ RegisterModel mockCreate = mock(RegisterModel.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.PROPERTIES)).thenReturn(false);
+ when(mockCommandLine.hasOption(GravitinoOptions.COMMENT)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.COMMENT)).thenReturn("comment");
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.CREATE));
+
+ doReturn(mockCreate)
+ .when(commandLine)
+ .newCreateModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ eq("comment"),
+ argThat(Map::isEmpty));
+ commandLine.handleCommandLine();
+ verify(mockCreate).handle();
+ }
+
+ @Test
+ void testRegisterModelCommandWithProperties() {
+ RegisterModel mockCreate = mock(RegisterModel.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.PROPERTIES)).thenReturn(true);
+ when(mockCommandLine.getOptionValues(GravitinoOptions.PROPERTIES))
+ .thenReturn(new String[] {"key1=val1", "key2" + "=val2"});
+
when(mockCommandLine.hasOption(GravitinoOptions.COMMENT)).thenReturn(false);
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.CREATE));
+
+ doReturn(mockCreate)
+ .when(commandLine)
+ .newCreateModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ isNull(),
+ argThat(
+ argument ->
+ argument.size() == 2
+ && argument.containsKey("key1")
+ && argument.get("key1").equals("val1")));
+ commandLine.handleCommandLine();
+ verify(mockCreate).handle();
+ }
+
+ @Test
+ void testRegisterModelCommandWithCommentAndProperties() {
+ RegisterModel mockCreate = mock(RegisterModel.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.PROPERTIES)).thenReturn(true);
+ when(mockCommandLine.getOptionValues(GravitinoOptions.PROPERTIES))
+ .thenReturn(new String[] {"key1=val1", "key2" + "=val2"});
+ when(mockCommandLine.hasOption(GravitinoOptions.COMMENT)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.COMMENT)).thenReturn("comment");
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.CREATE));
+
+ doReturn(mockCreate)
+ .when(commandLine)
+ .newCreateModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ eq("comment"),
+ argThat(
+ argument ->
+ argument.size() == 2
+ && argument.containsKey("key1")
+ && argument.get("key1").equals("val1")));
+ commandLine.handleCommandLine();
+ verify(mockCreate).handle();
+ }
+
+ @Test
+ void testLinkModelCommandWithoutAlias() {
+ LinkModel linkModelMock = mock(LinkModel.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.URI)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.URI)).thenReturn("file:///tmp/file");
+ when(mockCommandLine.hasOption(GravitinoOptions.ALIAS)).thenReturn(false);
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.UPDATE));
+
+ doReturn(linkModelMock)
+ .when(commandLine)
+ .newLinkModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ eq("file:///tmp/file"),
+ isNull(),
+ isNull(),
+ argThat(Map::isEmpty));
+ commandLine.handleCommandLine();
+ verify(linkModelMock).handle();
+ }
+
+ @Test
+ void testLinkModelCommandWithAlias() {
+ LinkModel linkModelMock = mock(LinkModel.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.URI)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.URI)).thenReturn("file:///tmp/file");
+ when(mockCommandLine.hasOption(GravitinoOptions.ALIAS)).thenReturn(true);
+ when(mockCommandLine.getOptionValues(GravitinoOptions.ALIAS))
+ .thenReturn(new String[] {"aliasA", "aliasB"});
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.UPDATE));
+
+ doReturn(linkModelMock)
+ .when(commandLine)
+ .newLinkModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ eq("file:///tmp/file"),
+ argThat(
+ argument ->
+ argument.length == 2
+ && "aliasA".equals(argument[0])
+ && "aliasB".equals(argument[1])),
+ isNull(),
+ argThat(Map::isEmpty));
+ commandLine.handleCommandLine();
+ verify(linkModelMock).handle();
+ }
+
+ @Test
+ void testLinkModelCommandWithoutURI() {
+ Main.useExit = false;
+
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.URI)).thenReturn(false);
+ when(mockCommandLine.hasOption(GravitinoOptions.ALIAS)).thenReturn(false);
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.UPDATE));
+
+ assertThrows(RuntimeException.class, commandLine::handleCommandLine);
+ verify(commandLine, never())
+ .newLinkModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ isNull(),
+ isNull(),
+ isNull(),
+ argThat(Map::isEmpty));
+ String output = new String(errContent.toByteArray(),
StandardCharsets.UTF_8).trim();
+ assertEquals(ErrorMessages.MISSING_URI, output);
+ }
+
+ @Test
+ void testLinkModelCommandWithAllComponent() {
+ LinkModel linkModelMock = mock(LinkModel.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.URI)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.URI)).thenReturn("file:///tmp/file");
+ when(mockCommandLine.hasOption(GravitinoOptions.ALIAS)).thenReturn(true);
+ when(mockCommandLine.getOptionValues(GravitinoOptions.ALIAS))
+ .thenReturn(new String[] {"aliasA", "aliasB"});
+ when(mockCommandLine.hasOption(GravitinoOptions.COMMENT)).thenReturn(true);
+
when(mockCommandLine.getOptionValue(GravitinoOptions.COMMENT)).thenReturn("comment");
+
when(mockCommandLine.hasOption(GravitinoOptions.PROPERTIES)).thenReturn(true);
+ when(mockCommandLine.getOptionValues(GravitinoOptions.PROPERTIES))
+ .thenReturn(new String[] {"key1=val1", "key2" + "=val2"});
+ GravitinoCommandLine commandLine =
+ spy(
+ new GravitinoCommandLine(
+ mockCommandLine, mockOptions, CommandEntities.MODEL,
CommandActions.UPDATE));
+
+ doReturn(linkModelMock)
+ .when(commandLine)
+ .newLinkModel(
+ eq(GravitinoCommandLine.DEFAULT_URL),
+ eq(false),
+ eq("metalake_demo"),
+ eq("catalog"),
+ eq("schema"),
+ eq("model"),
+ eq("file:///tmp/file"),
+ argThat(
+ argument ->
+ argument.length == 2
+ && "aliasA".equals(argument[0])
+ && "aliasB".equals(argument[1])),
+ eq("comment"),
+ argThat(
+ argument ->
+ argument.size() == 2
+ && argument.containsKey("key1")
+ && argument.containsKey("key2")
+ && "val1".equals(argument.get("key1"))
+ && "val2".equals(argument.get("key2"))));
+ commandLine.handleCommandLine();
+ verify(linkModelMock).handle();
+ }
}
