diff --git a/fe/fe-core/src/main/java/com/starrocks/http/rest/TableQueryPlanAction.java b/fe/fe-core/src/main/java/com/starrocks/http/rest/TableQueryPlanAction.java index f68eb72be6c196..41de55c7dbc819 100644 --- a/fe/fe-core/src/main/java/com/starrocks/http/rest/TableQueryPlanAction.java +++ b/fe/fe-core/src/main/java/com/starrocks/http/rest/TableQueryPlanAction.java @@ -75,6 +75,7 @@ import com.starrocks.thrift.TScanRangeLocations; import com.starrocks.thrift.TTabletVersionInfo; import com.starrocks.thrift.TUniqueId; +import com.starrocks.warehouse.Warehouse; import io.netty.handler.codec.http.HttpMethod; import io.netty.handler.codec.http.HttpResponseStatus; import org.apache.logging.log4j.LogManager; @@ -126,6 +127,16 @@ protected void executeWithoutPassword(BaseRequest request, BaseResponse response || Strings.isNullOrEmpty(tableName)) { throw new StarRocksHttpException(HttpResponseStatus.BAD_REQUEST, "{database}/{table} must be selected"); } + if (request.getRequest().headers().contains(WAREHOUSE_KEY)) { + String warehouseName = request.getRequest().headers().get(WAREHOUSE_KEY); + Warehouse wh = globalStateMgr.getWarehouseMgr().getWarehouseAllowNull(warehouseName); + if (wh == null) { + throw new StarRocksHttpException(HttpResponseStatus.NOT_FOUND, + "Warehouse [" + warehouseName + "] " + "does not exist"); + } else { + ConnectContext.get().setCurrentWarehouse(warehouseName); + } + } if (Strings.isNullOrEmpty(postContent)) { throw new StarRocksHttpException(HttpResponseStatus.BAD_REQUEST, "POST body must contains [sql] root object"); diff --git a/fe/fe-core/src/test/java/com/starrocks/http/TableQueryPlanActionTest.java b/fe/fe-core/src/test/java/com/starrocks/http/TableQueryPlanActionTest.java index 79a4e64da3e57f..c73f9fc9d94254 100644 --- a/fe/fe-core/src/test/java/com/starrocks/http/TableQueryPlanActionTest.java +++ b/fe/fe-core/src/test/java/com/starrocks/http/TableQueryPlanActionTest.java @@ -182,4 +182,27 @@ public void testQueryPlanActionPruneEmpty() throws Exception { Assert.assertEquals("{\"partitions\":{},\"opaqued_query_plan\":\"\",\"status\":200}", respStr); } } + + @Test + public void testQueryPlanActionWithInvalidWarehouse() throws Exception { + super.setUpWithCatalog(); + RequestBody body = + RequestBody.create(JSON, "{ \"sql\" : \" select k1 from " + DB_NAME + "." + TABLE_NAME + " \" }"); + Request request = new Request.Builder() + .post(body) + .addHeader("Authorization", rootAuth) + .addHeader("warehouse", "invalid_warehouse") + .url(URI + PATH_URI) + .build(); + try (Response response = networkClient.newCall(request).execute()) { + String respStr = Objects.requireNonNull(response.body()).string(); + JSONObject jsonObject = new JSONObject(respStr); + System.out.println(respStr); + Assert.assertEquals(404, jsonObject.getInt("status")); + String exception = jsonObject.getString("exception"); + Assert.assertNotNull(exception); + Assert.assertTrue( + exception.startsWith("Warehouse [invalid_warehouse] does not exist")); + } + } }