From bfb85680668b91d68f8190a8247156477f326039 Mon Sep 17 00:00:00 2001 From: Lord of Abyss <103809695+Abyss-lord@users.noreply.github.com> Date: Sat, 4 Jan 2025 06:25:18 +0800 Subject: [PATCH] [#5960] fix(CLI): Add register and link commands to CLI for model (#6066) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### 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 ``` --- .../apache/gravitino/cli/ErrorMessages.java | 3 +- .../gravitino/cli/GravitinoCommandLine.java | 34 +++ .../gravitino/cli/GravitinoOptions.java | 6 + .../gravitino/cli/TestableCommandLine.java | 29 ++ .../gravitino/cli/commands/LinkModel.java | 106 +++++++ .../gravitino/cli/commands/RegisterModel.java | 103 +++++++ .../gravitino/cli/TestModelCommands.java | 284 ++++++++++++++++++ 7 files changed, 564 insertions(+), 1 deletion(-) create mode 100644 clients/cli/src/main/java/org/apache/gravitino/cli/commands/LinkModel.java create mode 100644 clients/cli/src/main/java/org/apache/gravitino/cli/commands/RegisterModel.java 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 084b5c34c85..e90c5259638 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 c23fb8b7cd0..3a9322d010e 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 @@ private void handleModelCommand() { } break; + case CommandActions.CREATE: + String createComment = line.getOptionValue(GravitinoOptions.COMMENT); + String[] createProperties = line.getOptionValues(GravitinoOptions.PROPERTIES); + Map 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 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 657566036dc..aaeb8f0184f 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 Options options() { 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 6a468749178..8df9498d97b 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.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.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 @@ protected ModelDetails newModelDetails( 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 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 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 00000000000..6e8a4ffb76d --- /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 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 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 00000000000..d50dbed50e2 --- /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 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 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 e486c41a9d1..391201f292f 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 @@ 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.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 @@ void testModelAuditCommand() { 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(); + } }