Skip to content

Commit

Permalink
[apache#5960] fix(CLI): Add register and link commands to CLI for model
Browse files Browse the repository at this point in the history
1. Add register and link commands to CLI for model.
2. Add two option for model --uri and --alias.
  • Loading branch information
Abyss-lord committed Jan 3, 2025
1 parent e5bc05f commit 5eeec9a
Show file tree
Hide file tree
Showing 6 changed files with 278 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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.";
Expand All @@ -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.";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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<String, String> 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<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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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"));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<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);
}
}
Original file line number Diff line number Diff line change
@@ -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(
"Link model " + model + " to " + uri + " with alias " + Arrays.toString(alias));
}
}
Original file line number Diff line number Diff line change
@@ -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<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 + ".");
}
}
}

0 comments on commit 5eeec9a

Please sign in to comment.