improve query param parsing

- add parse_query_param() - use Cow<> where possible - move param parsing code to utils::http::request This was originally PR https://github.com/neondatabase/neon/pull/3502 which targeted a different branch. closes #3510
2026-06-04 22:10:39 +00:00 · 2023-02-01 14:11:12 +01:00
parent 9bb6a6c77c
commit 590695e845
6 changed files with 78 additions and 57 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4337,6 +4337,7 @@ dependencies = [
 "tokio-rustls",
 "tracing",
 "tracing-subscriber",
+ "url",
 "workspace_hack",
 ]

--- a/libs/utils/Cargo.toml
+++ b/libs/utils/Cargo.toml
@@ -37,6 +37,7 @@ metrics.workspace = true
 pq_proto.workspace = true

 workspace_hack.workspace = true
+url.workspace = true

 [dev-dependencies]
 byteorder.workspace = true
--- a/libs/utils/src/http/request.rs
+++ b/libs/utils/src/http/request.rs
@@ -1,4 +1,5 @@
-use std::str::FromStr;
+use core::fmt;
+use std::{borrow::Cow, str::FromStr};

 use super::error::ApiError;
 use anyhow::anyhow;
@@ -29,6 +30,50 @@ pub fn parse_request_param<T: FromStr>(
    }
 }

+fn get_query_param<'a>(
+    request: &'a Request<Body>,
+    param_name: &str,
+) -> Result<Option<Cow<'a, str>>, ApiError> {
+    let query = match request.uri().query() {
+        Some(q) => q,
+        None => return Ok(None),
+    };
+    let mut values = url::form_urlencoded::parse(query.as_bytes())
+        .filter_map(|(k, v)| if k == param_name { Some(v) } else { None })
+        // we call .next() twice below. If it's None the first time, .fuse() ensures it's None afterwards
+        .fuse();
+
+    let value1 = values.next();
+    if values.next().is_some() {
+        return Err(ApiError::BadRequest(anyhow!(
+            "param {param_name} specified more than once"
+        )));
+    }
+    Ok(value1)
+}
+
+pub fn must_get_query_param<'a>(
+    request: &'a Request<Body>,
+    param_name: &str,
+) -> Result<Cow<'a, str>, ApiError> {
+    get_query_param(request, param_name)?.ok_or_else(|| {
+        ApiError::BadRequest(anyhow!("no {param_name} specified in query parameters"))
+    })
+}
+
+pub fn parse_query_param<E: fmt::Display, T: FromStr<Err = E>>(
+    request: &Request<Body>,
+    param_name: &str,
+) -> Result<Option<T>, ApiError> {
+    get_query_param(request, param_name)?
+        .map(|v| {
+            v.parse().map_err(|e| {
+                ApiError::BadRequest(anyhow!("cannot parse query param {param_name}: {e}"))
+            })
+        })
+        .transpose()
+}
+
 pub async fn ensure_no_body(request: &mut Request<Body>) -> Result<(), ApiError> {
    match request.body_mut().data().await {
        Some(_) => Err(ApiError::BadRequest(anyhow!("Unexpected request body"))),
--- a/pageserver/src/http/routes.rs
+++ b/pageserver/src/http/routes.rs
@@ -7,6 +7,7 @@ use pageserver_api::models::DownloadRemoteLayersTaskSpawnRequest;
 use remote_storage::GenericRemoteStorage;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
+use utils::http::request::{must_get_query_param, parse_query_param};

 use super::models::{
    StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
@@ -235,8 +236,8 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<

 async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
    let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
-    let include_non_incremental_logical_size =
-        query_param_present(&request, "include-non-incremental-logical-size");
+    let include_non_incremental_logical_size: Option<bool> =
+        parse_query_param(&request, "include-non-incremental-logical-size")?;
    check_permission(&request, Some(tenant_id))?;

    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
@@ -249,13 +250,14 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,

        let mut response_data = Vec::with_capacity(timelines.len());
        for timeline in timelines {
-            let timeline_info =
-                build_timeline_info(&timeline, include_non_incremental_logical_size, &ctx)
-                    .await
-                    .context(
-                        "Failed to convert tenant timeline {timeline_id} into the local one: {e:?}",
-                    )
-                    .map_err(ApiError::InternalServerError)?;
+            let timeline_info = build_timeline_info(
+                &timeline,
+                include_non_incremental_logical_size.unwrap_or(false),
+                &ctx,
+            )
+            .await
+            .context("Failed to convert tenant timeline {timeline_id} into the local one: {e:?}")
+            .map_err(ApiError::InternalServerError)?;

            response_data.push(timeline_info);
        }
@@ -267,36 +269,11 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
    json_response(StatusCode::OK, response_data)
 }

-/// Checks if a query param is present in the request's URL
-fn query_param_present(request: &Request<Body>, param: &str) -> bool {
-    request
-        .uri()
-        .query()
-        .map(|v| url::form_urlencoded::parse(v.as_bytes()).any(|(p, _)| p == param))
-        .unwrap_or(false)
-}
-
-fn get_query_param(request: &Request<Body>, param_name: &str) -> Result<String, ApiError> {
-    request.uri().query().map_or(
-        Err(ApiError::BadRequest(anyhow!("empty query in request"))),
-        |v| {
-            url::form_urlencoded::parse(v.as_bytes())
-                .find(|(k, _)| k == param_name)
-                .map_or(
-                    Err(ApiError::BadRequest(anyhow!(
-                        "no {param_name} specified in query parameters"
-                    ))),
-                    |(_, v)| Ok(v.into_owned()),
-                )
-        },
-    )
-}
-
 async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
    let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
-    let include_non_incremental_logical_size =
-        query_param_present(&request, "include-non-incremental-logical-size");
+    let include_non_incremental_logical_size: Option<bool> =
+        parse_query_param(&request, "include-non-incremental-logical-size")?;
    check_permission(&request, Some(tenant_id))?;

    // Logical size calculation needs downloading.
@@ -311,11 +288,14 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
            .get_timeline(timeline_id, false)
            .map_err(ApiError::NotFound)?;

-        let timeline_info =
-            build_timeline_info(&timeline, include_non_incremental_logical_size, &ctx)
-                .await
-                .context("get local timeline info")
-                .map_err(ApiError::InternalServerError)?;
+        let timeline_info = build_timeline_info(
+            &timeline,
+            include_non_incremental_logical_size.unwrap_or(false),
+            &ctx,
+        )
+        .await
+        .context("get local timeline info")
+        .map_err(ApiError::InternalServerError)?;

        Ok::<_, ApiError>(timeline_info)
    }
@@ -330,8 +310,8 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
    check_permission(&request, Some(tenant_id))?;

    let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
-    let timestamp_raw = get_query_param(&request, "timestamp")?;
-    let timestamp = humantime::parse_rfc3339(timestamp_raw.as_str())
+    let timestamp_raw = must_get_query_param(&request, "timestamp")?;
+    let timestamp = humantime::parse_rfc3339(&timestamp_raw)
        .with_context(|| format!("Invalid time: {:?}", timestamp_raw))
        .map_err(ApiError::BadRequest)?;
    let timestamp_pg = postgres_ffi::to_pg_timestamp(timestamp);
@@ -503,13 +483,7 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
    let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
    check_permission(&request, Some(tenant_id))?;

-    let inputs_only = if query_param_present(&request, "inputs_only") {
-        get_query_param(&request, "inputs_only")?
-            .parse()
-            .map_err(|_| ApiError::BadRequest(anyhow!("failed to parse inputs_only")))?
-    } else {
-        false
-    };
+    let inputs_only: Option<bool> = parse_query_param(&request, "inputs_only")?;

    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
    let tenant = mgr::get_tenant(tenant_id, true)
@@ -522,7 +496,7 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
        .await
        .map_err(ApiError::InternalServerError)?;

-    let size = if !inputs_only {
+    let size = if !inputs_only.unwrap_or(false) {
        Some(inputs.calculate().map_err(ApiError::InternalServerError)?)
    } else {
        None
--- a/scripts/export_import_between_pageservers.py
+++ b/scripts/export_import_between_pageservers.py
@@ -293,7 +293,7 @@ class NeonPageserverHttpClient(requests.Session):

    def timeline_detail(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID) -> Dict[Any, Any]:
        res = self.get(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id.hex}/timeline/{timeline_id.hex}?include-non-incremental-logical-size=1"
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id.hex}/timeline/{timeline_id.hex}?include-non-incremental-logical-size=true"
        )
        self.verbose_error(res)
        res_json = res.json()
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -1232,9 +1232,9 @@ class PageserverHttpClient(requests.Session):

        params = {}
        if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "yes"
+            params["include-non-incremental-logical-size"] = "true"
        if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "yes"
+            params["include-timeline-dir-layer-file-size-sum"] = "true"

        res = self.get(
            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline", params=params
@@ -1276,9 +1276,9 @@ class PageserverHttpClient(requests.Session):
    ) -> Dict[Any, Any]:
        params = {}
        if include_non_incremental_logical_size:
-            params["include-non-incremental-logical-size"] = "yes"
+            params["include-non-incremental-logical-size"] = "true"
        if include_timeline_dir_layer_file_size_sum:
-            params["include-timeline-dir-layer-file-size-sum"] = "yes"
+            params["include-timeline-dir-layer-file-size-sum"] = "true"

        res = self.get(
            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}",