fix: resolve PR comments

epam · Sep 19, 2024 · 705c4ce · 705c4ce
1 parent 3e8cc6b
commit 705c4ce
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 24 deletions.
diff --git a/src/main/java/com/epam/aidial/core/controller/LimitController.java b/src/main/java/com/epam/aidial/core/controller/LimitController.java
@@ -22,35 +22,34 @@ public LimitController(Proxy proxy, ProxyContext context) {
     }
 
     public Future<?> getLimits(String deploymentId) {
+
         Deployment deployment = context.getConfig().selectDeployment(deploymentId);
 
         Future<Deployment> deploymentFuture;
         if (deployment != null) {
-            if (!DeploymentController.hasAccess(context, deployment)) {
-                log.error("LimitController. Forbidden deployment {}. Key: {}. User sub: {}", deploymentId, context.getProject(), context.getUserSub());
-                return context.respond(HttpStatus.FORBIDDEN, "Forbidden deployment: " + deploymentId);
+            if (DeploymentController.hasAccess(context, deployment)) {
+                deploymentFuture = Future.succeededFuture(deployment);
+            } else {
+                deploymentFuture = Future.failedFuture(new PermissionDeniedException("Forbidden deployment: " + deploymentId));
             }
-            deploymentFuture = Future.succeededFuture(deployment);
         } else {
             deploymentFuture = proxy.getVertx().executeBlocking(() ->
                     proxy.getCustomApplicationService().getCustomApplication(deploymentId, context), false);
         }
+
         deploymentFuture.compose(dep -> {
             if (dep == null) {
-                String error = String.format("LimitController. Deployment not found %s", deploymentId);
-                log.error(error);
-                context.respond(HttpStatus.NOT_FOUND, error);
-                return Future.succeededFuture();
+                throw new ResourceNotFoundException("Deployment " + deploymentId + " not found");
+            }
+            return proxy.getRateLimiter().getLimitStats(dep, context);
+        }).onSuccess(limitStats -> {
+            if (limitStats == null) {
+                context.respond(HttpStatus.NOT_FOUND);
+            } else {
+                context.respond(HttpStatus.OK, limitStats);
             }
-
-            return proxy.getRateLimiter().getLimitStats(dep, context).onSuccess(limitStats -> {
-                if (limitStats == null) {
-                    context.respond(HttpStatus.NOT_FOUND);
-                } else {
-                    context.respond(HttpStatus.OK, limitStats);
-                }
-            });
         }).onFailure(error -> handleRequestError(deploymentId, error));
+
         return Future.succeededFuture();
     }
 

diff --git a/src/main/java/com/epam/aidial/core/limiter/RateLimiter.java b/src/main/java/com/epam/aidial/core/limiter/RateLimiter.java
@@ -2,7 +2,6 @@
 
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.config.Deployment;
-import com.epam.aidial.core.config.Key;
 import com.epam.aidial.core.config.Limit;
 import com.epam.aidial.core.config.Role;
 import com.epam.aidial.core.data.ItemLimitStats;
@@ -85,14 +84,7 @@ public Future<LimitStats> getLimitStats(Deployment deployment, ProxyContext cont
             if (resourceService == null) {
                 return Future.succeededFuture();
             }
-            Key key = context.getKey();
             Limit limit = getLimitByUser(context, deployment);
-            if (limit == null) {
-                log.warn("Limit is not found. Trace: {}. Span: {}. Key: {}. User sub: {}. Deployment: {}",
-                        context.getTraceId(), context.getSpanId(), key == null ? null : key.getProject(),
-                        context.getUserSub(), deployment.getName());
-                return Future.succeededFuture();
-            }
             return vertx.executeBlocking(() -> getLimitStats(context, limit, deployment.getName()), false);
         } catch (Throwable e) {
             return Future.failedFuture(e);