Refactor code

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>

Author: nathaliellenaa

grpc/src/main/java/org/opensearch/ml/grpc/GrpcStatusMapper.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.grpc;
 
 import java.io.IOException;
+import java.util.Map;
 
 import org.opensearch.OpenSearchException;
 import org.opensearch.OpenSearchSecurityException;
@@ -20,6 +21,21 @@
  * Maps OpenSearch and ML Commons exceptions to gRPC Status codes.
  */
 public class GrpcStatusMapper {
+    private static final Map<RestStatus, Status> STATUS_MAP = Map
+        .of(
+            RestStatus.FORBIDDEN,
+            Status.PERMISSION_DENIED,
+            RestStatus.NOT_FOUND,
+            Status.NOT_FOUND,
+            RestStatus.TOO_MANY_REQUESTS,
+            Status.RESOURCE_EXHAUSTED,
+            RestStatus.BAD_REQUEST,
+            Status.INVALID_ARGUMENT,
+            RestStatus.SERVICE_UNAVAILABLE,
+            Status.UNAVAILABLE,
+            RestStatus.UNAUTHORIZED,
+            Status.UNAUTHENTICATED
+        );
 
     /**
      * Converts an exception to a gRPC Status.
@@ -48,17 +64,8 @@ public static Status toGrpcStatus(Exception exception) {
 
         // Handle OpenSearch exceptions with status codes
         if (exception instanceof OpenSearchException osException) {
-            RestStatus status = osException.status();
-
-            return switch (status) {
-                case FORBIDDEN -> Status.PERMISSION_DENIED.withDescription(exception.getMessage()).withCause(exception);
-                case NOT_FOUND -> Status.NOT_FOUND.withDescription(exception.getMessage()).withCause(exception);
-                case TOO_MANY_REQUESTS -> Status.RESOURCE_EXHAUSTED.withDescription(exception.getMessage()).withCause(exception);
-                case BAD_REQUEST -> Status.INVALID_ARGUMENT.withDescription(exception.getMessage()).withCause(exception);
-                case SERVICE_UNAVAILABLE -> Status.UNAVAILABLE.withDescription(exception.getMessage()).withCause(exception);
-                case UNAUTHORIZED -> Status.UNAUTHENTICATED.withDescription(exception.getMessage()).withCause(exception);
-                default -> Status.INTERNAL.withDescription(exception.getMessage()).withCause(exception);
-            };
+            Status grpcStatus = STATUS_MAP.getOrDefault(osException.status(), Status.INTERNAL);
+            return grpcStatus.withDescription(exception.getMessage()).withCause(exception);
         }
 
         // Handle I/O exceptions

grpc/src/main/java/org/opensearch/ml/grpc/MLGrpcServiceFactory.java

@@ -9,6 +9,12 @@
 import java.util.List;
 
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
+import org.opensearch.ml.grpc.interfaces.MLClient;
+import org.opensearch.ml.grpc.interfaces.MLModelAccessControlHelper;
+import org.opensearch.ml.grpc.interfaces.MLModelManager;
+import org.opensearch.ml.grpc.interfaces.MLSdkClient;
+import org.opensearch.ml.grpc.interfaces.MLTaskRunner;
+import org.opensearch.ml.grpc.interfaces.MLUserContextProvider;
 import org.opensearch.transport.grpc.spi.GrpcServiceFactory;
 
 import io.grpc.BindableService;
@@ -22,13 +28,14 @@
 @Log4j2
 public class MLGrpcServiceFactory implements GrpcServiceFactory {
 
-    private static volatile Object modelManager;
-    private static volatile Object predictTaskRunner;
-    private static volatile Object executeTaskRunner;
+    private static volatile MLModelManager modelManager;
+    private static volatile MLTaskRunner predictTaskRunner;
+    private static volatile MLTaskRunner executeTaskRunner;
     private static volatile MLFeatureEnabledSetting mlFeatureEnabledSetting;
-    private static volatile Object modelAccessControlHelper;
-    private static volatile Object client;
-    private static volatile Object sdkClient;
+    private static volatile MLModelAccessControlHelper modelAccessControlHelper;
+    private static volatile MLClient client;
+    private static volatile MLSdkClient sdkClient;
+    private static volatile MLUserContextProvider userContextProvider;
 
     /**
      * No-arg constructor required by SPI.
@@ -48,23 +55,30 @@ public MLGrpcServiceFactory() {
      * @param modelAccessControlHelper helper for validating model access control
      * @param client OpenSearch client for validation
      * @param sdkClient SDK client for multi-tenant operations
+     * @param userContextProvider provider for extracting user from security context
      */
     public static void initialize(
-        Object modelManager,
-        Object predictTaskRunner,
-        Object executeTaskRunner,
+        MLModelManager modelManager,
+        MLTaskRunner predictTaskRunner,
+        MLTaskRunner executeTaskRunner,
         MLFeatureEnabledSetting mlFeatureEnabledSetting,
-        Object modelAccessControlHelper,
-        Object client,
-        Object sdkClient
+        MLModelAccessControlHelper modelAccessControlHelper,
+        MLClient client,
+        MLSdkClient sdkClient,
+        MLUserContextProvider userContextProvider
     ) {
+        if (modelManager == null || predictTaskRunner == null || executeTaskRunner == null) {
+            throw new IllegalArgumentException("Required dependencies cannot be null");
+        }
+
         MLGrpcServiceFactory.modelManager = modelManager;
         MLGrpcServiceFactory.predictTaskRunner = predictTaskRunner;
         MLGrpcServiceFactory.executeTaskRunner = executeTaskRunner;
         MLGrpcServiceFactory.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
         MLGrpcServiceFactory.modelAccessControlHelper = modelAccessControlHelper;
         MLGrpcServiceFactory.client = client;
         MLGrpcServiceFactory.sdkClient = sdkClient;
+        MLGrpcServiceFactory.userContextProvider = userContextProvider;
     }
 
     @Override
@@ -77,6 +91,11 @@ public List<BindableService> build() {
         List<BindableService> services = new ArrayList<>();
 
         try {
+            // Validate initialization
+            if (modelManager == null || predictTaskRunner == null || executeTaskRunner == null) {
+                throw new IllegalStateException("MLGrpcServiceFactory not initialized. Call initialize() first.");
+            }
+
             // Create ML streaming service
             MLStreamingService streamingService = new MLStreamingService(
                 modelManager,
@@ -85,7 +104,8 @@ public List<BindableService> build() {
                 mlFeatureEnabledSetting,
                 modelAccessControlHelper,
                 client,
-                sdkClient
+                sdkClient,
+                userContextProvider
             );
 
             // Wrap service with tenant ID interceptor to extract metadata