Examples with User org.opensearch.commons.authuser.User used on opensource projects

Example 1 with User

use of org.opensearch.commons.authuser.User in project ml-commons by opensearch-project.

the class MLPredictTaskRunner method predict.

private void predict(MLTask mlTask, DataFrame inputDataFrame, MLPredictionTaskRequest request, ActionListener<MLTaskResponse> listener) {
    ActionListener<MLTaskResponse> internalListener = wrappedCleanupListener(listener, mlTask.getTaskId());
    // track ML task count and add ML task into cache
    mlStats.getStat(ML_EXECUTING_TASK_COUNT).increment();
    mlStats.getStat(ML_TOTAL_REQUEST_COUNT).increment();
    mlStats.createCounterStatIfAbsent(requestCountStat(mlTask.getFunctionName(), ActionName.PREDICT)).increment();
    mlTaskManager.add(mlTask);
    // run predict
    if (request.getModelId() != null) {
        // search model by model id.
        try (ThreadContext.StoredContext context = threadPool.getThreadContext().stashContext()) {
            MLInput mlInput = request.getMlInput();
            ActionListener<GetResponse> getResponseListener = ActionListener.wrap(r -> {
                if (r == null || !r.isExists()) {
                    internalListener.onFailure(new ResourceNotFoundException("No model found, please check the modelId."));
                    return;
                }
                Map<String, Object> source = r.getSourceAsMap();
                User requestUser = getUserContext(client);
                User resourceUser = User.parse((String) source.get(USER));
                if (!checkUserPermissions(requestUser, resourceUser, request.getModelId())) {
                    // The backend roles of request user and resource user doesn't have intersection
                    OpenSearchException e = new OpenSearchException("User: " + requestUser.getName() + " does not have permissions to run predict by model: " + request.getModelId());
                    handlePredictFailure(mlTask, internalListener, e, false);
                    return;
                }
                Model model = new Model();
                model.setName((String) source.get(MLModel.MODEL_NAME));
                model.setVersion((Integer) source.get(MLModel.MODEL_VERSION));
                byte[] decoded = Base64.getDecoder().decode((String) source.get(MLModel.MODEL_CONTENT));
                model.setContent(decoded);
                // run predict
                mlTaskManager.updateTaskState(mlTask.getTaskId(), MLTaskState.RUNNING, mlTask.isAsync());
                MLOutput output = MLEngine.predict(mlInput.toBuilder().inputDataset(new DataFrameInputDataset(inputDataFrame)).build(), model);
                if (output instanceof MLPredictionOutput) {
                    ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
                }
                // Once prediction complete, reduce ML_EXECUTING_TASK_COUNT and update task state
                handleAsyncMLTaskComplete(mlTask);
                MLTaskResponse response = MLTaskResponse.builder().output(output).build();
                internalListener.onResponse(response);
            }, e -> {
                log.error("Failed to predict " + mlInput.getAlgorithm() + ", modelId: " + mlTask.getModelId(), e);
                handlePredictFailure(mlTask, internalListener, e, true);
            });
            GetRequest getRequest = new GetRequest(ML_MODEL_INDEX, mlTask.getModelId());
            client.get(getRequest, ActionListener.runBefore(getResponseListener, () -> context.restore()));
        } catch (Exception e) {
            log.error("Failed to get model " + mlTask.getModelId(), e);
            handlePredictFailure(mlTask, internalListener, e, true);
        }
    } else {
        IllegalArgumentException e = new IllegalArgumentException("ModelId is invalid");
        log.error("ModelId is invalid", e);
        handlePredictFailure(mlTask, internalListener, e, false);
    }
}

Example 2 with User

use of org.opensearch.commons.authuser.User in project ml-commons by opensearch-project.

the class AccessControllerTests method testCheckUserPermissionsWithNullBackendRoles.

public void testCheckUserPermissionsWithNullBackendRoles() {
    User requestUser = User.parse("requestuser||role1,role2|myTenant");
    User resourceUser = User.parse("resourceuser||role1,role2|myTenant");
    boolean hasPermission = AccessController.checkUserPermissions(requestUser, resourceUser, "");
    assertFalse(hasPermission);
}

Example 3 with User

use of org.opensearch.commons.authuser.User in project ml-commons by opensearch-project.

the class AccessControllerTests method testCheckUserPermissionsWithoutMatch.

public void testCheckUserPermissionsWithoutMatch() {
    User requestUser = User.parse("requestuser|bckrole1,bckrole2|role1,role2|myTenant");
    User resourceUser = User.parse("resourceuser|bckrole3,bckrole4|role1,role2|myTenant");
    boolean hasPermission = AccessController.checkUserPermissions(requestUser, resourceUser, "");
    assertFalse(hasPermission);
}

Example 4 with User

use of org.opensearch.commons.authuser.User in project ml-commons by opensearch-project.

the class AccessControllerTests method testCheckUserPermissionsWithMatch.

public void testCheckUserPermissionsWithMatch() {
    User requestUser = User.parse("requestuser|bckrole1,bckrole2|role1,role2|myTenant");
    User resourceUser = User.parse("resourceuser|bckrole2,bckrole3|role1,role2|myTenant");
    boolean hasPermission = AccessController.checkUserPermissions(requestUser, resourceUser, "");
    assertTrue(hasPermission);
}

Example 5 with User

use of org.opensearch.commons.authuser.User in project ml-commons by opensearch-project.

the class MLModel method parse.

public static MLModel parse(XContentParser parser) throws IOException {
    String name = null;
    FunctionName algorithm = null;
    Integer version = null;
    String content = null;
    User user = null;
    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
    while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
        String fieldName = parser.currentName();
        parser.nextToken();
        switch(fieldName) {
            case MODEL_NAME:
                name = parser.text();
                break;
            case MODEL_CONTENT:
                content = parser.text();
                break;
            case MODEL_VERSION:
                version = parser.intValue(false);
                break;
            case USER:
                user = User.parse(parser);
                break;
            case ALGORITHM:
                algorithm = FunctionName.from(parser.text());
                break;
            default:
                parser.skipChildren();
                break;
        }
    }
    return MLModel.builder().name(name).algorithm(algorithm).version(version).content(content).user(user).build();
}