fix: Correctly proxy HEAD requests

Author: sbernauer

CHANGELOG.md

@@ -26,6 +26,8 @@ All notable changes to this project will be documented in this file.
 
 - Set connection and response timeout for Redis connections ([#85]).
 - Only remove queries from the persistence in case they don't send a `nextUri` and are in state `FINISHED` ([#98]).
+- Correctly proxy HEAD requests to `/v1/statement/executing/{queryId}/{slug}/{token}`.
+  Previously, we would GET (instead of HEAD) the URL at the Trino cluster, which resulted in trino-lb dropping the HTTP body, causing problems ([#100]).
 
 [#68]: https://github.com/stackabletech/trino-lb/pull/68
 [#85]: https://github.com/stackabletech/trino-lb/pull/85

trino-lb/src/cluster_group_manager.rs

@@ -35,6 +35,9 @@ pub enum Error {
     #[snafu(display("Failed to contact Trino API to post query"))]
     ContactTrinoPostQuery { source: reqwest::Error },
 
+    #[snafu(display("Failed to call Trino HEAD URL {url:?}"))]
+    CallTrinoHeadUrl { source: reqwest::Error, url: Url },
+
     #[snafu(display("Failed to decode Trino API response"))]
     DecodeTrinoResponse { source: reqwest::Error },
 
@@ -232,7 +235,7 @@ impl ClusterGroupManager {
             .get(next_uri)
             .headers(headers)
             .send()
-            .instrument(info_span!("Send HTTP get to Trino"))
+            .instrument(info_span!("Send HTTP GET to Trino"))
             .await
             .context(ContactTrinoPostQuerySnafu)?;
         let headers = response.headers();
@@ -249,6 +252,33 @@ impl ClusterGroupManager {
         Ok((trino_query_api_response, headers))
     }
 
+    /// Sometimes the trino-client HEADs a /executing/xxx endpoint instead of GETing it.
+    /// We need to proxy this as a HEAD request as well.
+    #[instrument(
+        skip(self),
+        fields(head_uri = %head_uri, headers = ?headers.sanitize())
+    )]
+    pub async fn send_head_to_trino(
+        &self,
+        head_uri: Url,
+        mut headers: HeaderMap,
+    ) -> Result<HeaderMap, Error> {
+        add_current_context_to_client_request(tracing::Span::current().context(), &mut headers);
+
+        let response = self
+            .http_client
+            .head(head_uri.clone())
+            .headers(headers)
+            .send()
+            .instrument(info_span!("Send HTTP HEAD to Trino"))
+            .await
+            .with_context(|_| CallTrinoHeadUrlSnafu { url: head_uri })?;
+        let headers = response.headers();
+        let headers = filter_to_trino_headers(headers);
+
+        Ok(headers)
+    }
+
     #[instrument(
         skip(self),
         fields(request_headers = ?request_headers.sanitize())

trino-lb/src/http_server/v1/statement.rs

@@ -8,6 +8,7 @@ use std::{
 
 use axum::{
     Json,
+    body::Body,
     extract::{Path, State},
     response::{IntoResponse, Response},
 };
@@ -93,6 +94,11 @@ pub enum Error {
         source: cluster_group_manager::Error,
     },
 
+    #[snafu(display("Failed to send HEAD request to trino"))]
+    SendHeadToTrino {
+        source: cluster_group_manager::Error,
+    },
+
     #[snafu(display(
         "Failed to decrement the query counter query trino cluster {trino_cluster:?}"
     ))]
@@ -216,22 +222,45 @@ pub async fn get_trino_queued_statement(
 /// Trino cluster and currently running.
 ///
 /// In case the nextUri is null, the query will be stopped and removed from trino-lb.
+///
+/// Please note that sometimes the Trino clients also HEAD this endpoint, in which case we need some
+/// special handling.
 #[instrument(
-    name = "GET /v1/statement/executing/{queryId}/{slug}/{token}",
+    name = "GET (or HEAD) /v1/statement/executing/{queryId}/{slug}/{token}",
     skip(state, headers)
 )]
 pub async fn get_trino_executing_statement(
+    method: http::Method,
     headers: HeaderMap,
     State(state): State<Arc<AppState>>,
     Path((query_id, _, _)): Path<(TrinoQueryId, String, u64)>,
     uri: Uri,
-) -> Result<(HeaderMap, Json<TrinoQueryApiResponse>), Error> {
+) -> Result<Response, Error> {
+    if method == http::Method::HEAD {
+        state.metrics.http_counter.add(
+            1,
+            &[KeyValue::new("resource", "head_trino_executing_statement")],
+        );
+
+        let headers = handle_head_request_to_trino(&state, headers, query_id, uri.path()).await?;
+
+        // For a HEAD request we don't need (nor can) return a body.
+        let mut response = Response::new(Body::empty());
+        *response.status_mut() = StatusCode::OK;
+        *response.headers_mut() = headers;
+        return Ok(response);
+    }
     state.metrics.http_counter.add(
         1,
         &[KeyValue::new("resource", "get_trino_executing_statement")],
     );
 
-    handle_query_running_on_trino(&state, headers, query_id, uri.path()).await
+    let (headers, body) =
+        handle_query_running_on_trino(&state, headers, query_id, uri.path()).await?;
+
+    let mut response = body.into_response();
+    *response.headers_mut() = headers;
+    Ok(response)
 }
 
 #[instrument(skip(
@@ -495,6 +524,38 @@ async fn handle_query_running_on_trino(
     Ok((trino_headers, Json(trino_query_api_response)))
 }
 
+async fn handle_head_request_to_trino(
+    state: &Arc<AppState>,
+    headers: HeaderMap,
+    query_id: TrinoQueryId,
+    requested_path: &str,
+) -> Result<HeaderMap, Error> {
+    let query =
+        state
+            .persistence
+            .load_query(&query_id)
+            .await
+            .context(LoadQueryFromPersistenceSnafu {
+                query_id: query_id.clone(),
+            })?;
+
+    let headers = state
+        .cluster_group_manager
+        .send_head_to_trino(
+            query.trino_endpoint.join(requested_path).context(
+                JoinRequestPathToTrinoEndpointSnafu {
+                    requested_path,
+                    trino_endpoint: query.trino_endpoint.clone(),
+                },
+            )?,
+            headers,
+        )
+        .await
+        .context(SendHeadToTrinoSnafu)?;
+
+    Ok(headers)
+}
+
 /// This function get's asked to delete the queued query.
 /// IMPORTANT: It does not check that the user is authorized to delete the queued query. Instead we assume that the
 /// random part of the queryId trino-lb generates provides sufficient protection, as other clients can not extract