From 5eeec9af221daef9d5c5f1b608d318a31fce5c94 Mon Sep 17 00:00:00 2001 From: pancx Date: Thu, 2 Jan 2025 17:15:25 +0800 Subject: [PATCH] [#5960] fix(CLI): Add register and link commands to CLI for model 1. Add register and link commands to CLI for model. 2. Add two option for model --uri and --alias. --- .../apache/gravitino/cli/ErrorMessages.java | 2 + .../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 | 101 +++++++++++++++++ 6 files changed, 278 insertions(+) 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 4bd523ec280..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,6 +52,7 @@ 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."; 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 8cd335bebbe..67c5c7c2563 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 @@ -1175,6 +1175,8 @@ private void handleModelCommand() { } String model = name.getModelName(); + String[] alias = line.getOptionValues(GravitinoOptions.ALIAS); + String uri = line.getOptionValue(GravitinoOptions.URI); if (model == null) missingEntities.add(CommandEntities.MODEL); checkEntities(missingEntities); @@ -1183,6 +1185,38 @@ private void handleModelCommand() { newModelDetails(url, ignore, metalake, catalog, schema, model).handle(); 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: + 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 3cfd84ad83c..ef775a4c9b1 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; @@ -82,6 +83,7 @@ import org.apache.gravitino.cli.commands.MetalakeEnable; 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; @@ -919,4 +921,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..2660382d7cd --- /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( + "Link model " + model + " to " + uri + " with alias " + 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..5a099266c31 --- /dev/null +++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RegisterModel.java @@ -0,0 +1,101 @@ +/* + * 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.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 + "."); + } + } +}