Extract query ID from all kill_query procedure variations #425

willmostly · 2024-07-29T19:07:37Z

Description

Trino Gateway attempts to extract the query id from the body of queries using the runtime.kill_query procedure from the system catalog. This adds a test for this functionality.

Additional context and related issues

Release notes

( x) This is not user-visible or is docs only, and no release notes are required.
( ) Release notes are required. Please propose a release note for me.
( ) Release notes are required, with the following suggested text:

mosabua

Looks fine .. I wish we would not have to mock around for stuff like that but a full integration test seems worse..

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

ebyhr

The regular expression is incorrect in some cases.

mosabua · 2024-09-05T16:58:41Z

Needs a rebase now..

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

.

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

willmostly · 2024-09-23T19:10:25Z

@ebyhr I switched to using the full query parser, which should be resilient to whitespace, comments, the string literals containing the procedure name, etc

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

ebyhr · 2024-09-23T22:51:18Z

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

+            procedure = Optional.of((Call) node);
+            procedureArgs = ((Call) node).getArguments();


The method name is getNames. Setting these fields changes doesn't match the naming now. I would recommend moving this logic to ProxyUtils and reverting changes in this class.

i renamed the method to be more generic. The initial implementation was only concerned with extracting qualified names from SQL, but there is no reason IMO not to extract other information as well. In addition to this specific case of the kill_query call which contains a query id, the Call and CallArguments could be useful for routing logic. For example, perhaps the sync_partition_metadata procedure should always be run on a specific set of clusters.

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/main/java/io/trino/gateway/ha/config/RequestAnalyzerConfig.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

ebyhr · 2024-09-26T13:29:49Z

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

-        catch (Exception e) {
-            log.error(e, "Error extracting query payload from request");
+        catch (IOException e) {
+            log.error(e, "Error reading request body");


Why do we suppress IOException here and throw an exception for non string literal argument of kill_query procedure?

The request.getInputStream() may contain kill_query, but it can go to extractQueryIdIfPresent in case of IOException, right?

If the request body contains the kill_query SQL, then I already know at this point that extractQueryIdIfPresent will not find a query id because Trino doesn't have any endpoints with query id embedded in the path or queryParams that accept a POST. So falling back won't help. If I don't throw on a non-string literal argument, then this request will most likely get routed to a cluster that the target query is not running on, and the user will get a potentially confusing error message back about the query not existing. So throwing on a non-string literal allows us to fail fast with a more meaningful error.

Do you think we should throw on IOException as well? I was just trying to match the existing behavior, but i'm not sure that it is correct.

Just noting that this has been changed such that the IOException is not suppressed, and is treated as an error. If the gateway can't read the request body, Trino will not be able to either, so there is no reason to forward the request.

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java

ebyhr · 2024-10-01T22:09:27Z

Could you rebase on main to resolve conflicts?

…inoQueryProperties

cla-bot bot added the cla-signed label Jul 29, 2024

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 0f19b85 to 1ad0e45 Compare July 29, 2024 19:15

mosabua approved these changes Aug 21, 2024

View reviewed changes

ebyhr reviewed Aug 24, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 1ad0e45 to a8355b6 Compare August 30, 2024 21:42

willmostly requested a review from ebyhr August 30, 2024 21:45

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from a8355b6 to 83760bd Compare August 30, 2024 21:53

ebyhr reviewed Aug 30, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 9a00f3e to 95e782a Compare September 3, 2024 19:40

willmostly changed the title ~~Test extraction of query ID from kill_query procedure~~ Extract query ID from all kill_query procedure variations Sep 3, 2024

willmostly requested a review from ebyhr September 3, 2024 19:43

ebyhr reviewed Sep 3, 2024

View reviewed changes

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java Outdated Show resolved Hide resolved

ebyhr previously requested changes Sep 4, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 95e782a to c311199 Compare September 4, 2024 20:10

willmostly requested a review from ebyhr September 4, 2024 20:13

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from c311199 to e9226b6 Compare September 9, 2024 19:29

ebyhr reviewed Sep 13, 2024

View reviewed changes

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java Show resolved Hide resolved

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java Outdated Show resolved Hide resolved

ebyhr reviewed Sep 16, 2024

View reviewed changes

gateway-ha/src/test/java/io/trino/gateway/ha/handler/TestQueryIdCachingProxyHandler.java Outdated Show resolved Hide resolved

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch 2 times, most recently from 805fd81 to e805660 Compare September 23, 2024 19:05

willmostly requested a review from ebyhr September 23, 2024 19:10

ebyhr reviewed Sep 23, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from e805660 to 6e9842a Compare September 25, 2024 23:56

ebyhr requested changes Sep 26, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 6e9842a to 76f2c68 Compare September 26, 2024 13:05

ebyhr reviewed Sep 26, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from 76f2c68 to f4ee990 Compare September 26, 2024 20:57

ebyhr reviewed Sep 26, 2024

View reviewed changes

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch 4 times, most recently from 9e35c95 to c8fe87d Compare October 1, 2024 23:39

willmostly requested a review from ebyhr October 2, 2024 00:11

willmostly added 2 commits October 2, 2024 10:07

Extract query ID from all kill_query procedure variations

02b2b03

Update maxBodySize type in RequestAnalyzerConfig to match usage in Tr…

131228b

…inoQueryProperties

willmostly force-pushed the will/test_extract_query_id_from_system_runtime_kill_query branch from c8fe87d to 131228b Compare October 2, 2024 14:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Extract query ID from all kill_query procedure variations #425

Extract query ID from all kill_query procedure variations #425

willmostly commented Jul 29, 2024

mosabua left a comment

ebyhr left a comment

mosabua commented Sep 5, 2024

willmostly commented Sep 23, 2024

ebyhr Sep 23, 2024

willmostly Sep 25, 2024

ebyhr Sep 26, 2024

willmostly Sep 26, 2024

willmostly Oct 1, 2024

ebyhr commented Oct 1, 2024

		procedure = Optional.of((Call) node);
		procedureArgs = ((Call) node).getArguments();

Extract query ID from all kill_query procedure variations #425

Are you sure you want to change the base?

Extract query ID from all kill_query procedure variations #425

Conversation

willmostly commented Jul 29, 2024

Description

Additional context and related issues

Release notes

mosabua left a comment

Choose a reason for hiding this comment

ebyhr left a comment

Choose a reason for hiding this comment

mosabua commented Sep 5, 2024

willmostly commented Sep 23, 2024

ebyhr Sep 23, 2024

Choose a reason for hiding this comment

willmostly Sep 25, 2024

Choose a reason for hiding this comment

ebyhr Sep 26, 2024

Choose a reason for hiding this comment

willmostly Sep 26, 2024

Choose a reason for hiding this comment

willmostly Oct 1, 2024

Choose a reason for hiding this comment

ebyhr commented Oct 1, 2024