Skip to content

[Enhancement] Support to specify warehouse for query_plan rest query #57930

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zhangbutao wants to merge 3 commits into from
Closed

[Enhancement] Support to specify warehouse for query_plan rest query #57930

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7ce923a
[Enhancement] Support to specify warehouse for rest query
zhangbutao Apr 11, 2025
0432095
Added test
zhangbutao Apr 17, 2025
ecc40f2
check and set warehouse before execute query_plan
zhangbutao Apr 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/http/rest/RestBaseAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,9 @@ public void execute(BaseRequest request, BaseResponse response) throws DdlExcept
ctx.setCurrentUserIdentity(currentUser);
ctx.setCurrentRoleIds(currentUser);
ctx.setThreadLocalInfo();
if (request.getRequest().headers().contains(WAREHOUSE_KEY)) {

@kevincai kevincai Apr 16, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add some unit tests about the new parameters, and what's the response looks like if the specified warehouse doesn't exist.

@zhangbutao zhangbutao Apr 17, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a UT to post a query_plan request, just like sr flink source connector:
https://github.com/StarRocks/starrocks-connector-for-apache-flink/blob/3e17b7b0531f128fb722b4892fdf608931854202/src/main/java/com/starrocks/connector/flink/manager/StarRocksQueryPlanVisitor.java#L95-#L102

ctx.setCurrentWarehouse(request.getRequest().headers().get(WAREHOUSE_KEY));
}
executeWithoutPassword(request, response);
}

Expand Down
23 changes: 23 additions & 0 deletions fe/fe-core/src/test/java/com/starrocks/http/TableQueryPlanActionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(500, jsonObject.getInt("status"));

@kevincai kevincai Apr 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is not expected, FE should respond HTTP 400 BAD request.

@zhangbutao zhangbutao Apr 17, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The exception comes from here, which all exception wiil be 500

starrocks/fe/fe-core/src/main/java/com/starrocks/http/rest/TableQueryPlanAction.java

Lines 221 to 225 in 6996187

} catch (Exception e) {
LOG.error("Get query plan for sql[{}] error, queryId: {}", sql, context.getQueryId(), e);
throw new StarRocksHttpException(
HttpResponseStatus.INTERNAL_SERVER_ERROR,
"Invalid SQL: " + sql);

@kevincai kevincai Apr 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should do input parameter check before execute the parameter, or do fine grain exception handling. It is the request's parameter invalidation, should not be considered as a HTTP 5XX error. Usually a HTTP 5XX error means the http client is not doing anything wrong, and can do endless retry without modification of the request as wish.

@zhangbutao zhangbutao Apr 17, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kevincai I changed the code to check and set warehouse before execute the query_plan, so that we can get the real error status and exception.
So this PR is for TableQueryPlanAction, which can solve the sr flink source connector warehouse issue.

For other http Action, i am not sure if need the warehouse. But we can also do similar check like this PR if other http Actions have any need.

@kevincai kevincai Apr 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is good enough to fix only in query_plan API in this PR.

String exception = jsonObject.getString("exception");
Assert.assertNotNull(exception);
Assert.assertTrue(
exception.startsWith("Invalid SQL: select k1 from testDb.testTbl "));
}
}
}