LCOV - code coverage report
Current view: top level - proxy/src/serverless - sql_over_http.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 5.8 % 798 46
Test Date: 2025-02-20 13:11:02 Functions: 3.9 % 102 4

            Line data    Source code
       1              : use std::pin::pin;
       2              : use std::sync::Arc;
       3              : 
       4              : use bytes::Bytes;
       5              : use futures::future::{select, try_join, Either};
       6              : use futures::{StreamExt, TryFutureExt};
       7              : use http::header::AUTHORIZATION;
       8              : use http::Method;
       9              : use http_body_util::combinators::BoxBody;
      10              : use http_body_util::{BodyExt, Full};
      11              : use http_utils::error::ApiError;
      12              : use hyper::body::Incoming;
      13              : use hyper::http::{HeaderName, HeaderValue};
      14              : use hyper::{header, HeaderMap, Request, Response, StatusCode};
      15              : use indexmap::IndexMap;
      16              : use postgres_client::error::{DbError, ErrorPosition, SqlState};
      17              : use postgres_client::{GenericClient, IsolationLevel, NoTls, ReadyForQueryStatus, Transaction};
      18              : use pq_proto::StartupMessageParamsBuilder;
      19              : use serde::Serialize;
      20              : use serde_json::value::RawValue;
      21              : use serde_json::Value;
      22              : use tokio::time::{self, Instant};
      23              : use tokio_util::sync::CancellationToken;
      24              : use tracing::{debug, error, info};
      25              : use typed_json::json;
      26              : use url::Url;
      27              : use uuid::Uuid;
      28              : 
      29              : use super::backend::{LocalProxyConnError, PoolingBackend};
      30              : use super::conn_pool::{AuthData, ConnInfoWithAuth};
      31              : use super::conn_pool_lib::{self, ConnInfo};
      32              : use super::error::HttpCodeError;
      33              : use super::http_util::json_response;
      34              : use super::json::{json_to_pg_text, pg_text_row_to_json, JsonConversionError};
      35              : use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
      36              : use crate::auth::{endpoint_sni, ComputeUserInfoParseError};
      37              : use crate::config::{AuthenticationConfig, HttpConfig, ProxyConfig, TlsConfig};
      38              : use crate::context::RequestContext;
      39              : use crate::error::{ErrorKind, ReportableError, UserFacingError};
      40              : use crate::http::{read_body_with_limit, ReadBodyError};
      41              : use crate::metrics::{HttpDirection, Metrics};
      42              : use crate::proxy::{run_until_cancelled, NeonOptions};
      43              : use crate::serverless::backend::HttpConnError;
      44              : use crate::types::{DbName, RoleName};
      45              : use crate::usage_metrics::{MetricCounter, MetricCounterRecorder};
      46              : 
      47           11 : #[derive(serde::Deserialize)]
      48              : #[serde(rename_all = "camelCase")]
      49              : struct QueryData {
      50              :     query: String,
      51              :     #[serde(deserialize_with = "bytes_to_pg_text")]
      52              :     #[serde(default)]
      53              :     params: Vec<Option<String>>,
      54              :     #[serde(default)]
      55              :     array_mode: Option<bool>,
      56              : }
      57              : 
      58            1 : #[derive(serde::Deserialize)]
      59              : struct BatchQueryData {
      60              :     queries: Vec<QueryData>,
      61              : }
      62              : 
      63              : #[derive(serde::Deserialize)]
      64              : #[serde(untagged)]
      65              : enum Payload {
      66              :     Single(QueryData),
      67              :     Batch(BatchQueryData),
      68              : }
      69              : 
      70              : pub(super) static NEON_REQUEST_ID: HeaderName = HeaderName::from_static("neon-request-id");
      71              : 
      72              : static CONN_STRING: HeaderName = HeaderName::from_static("neon-connection-string");
      73              : static RAW_TEXT_OUTPUT: HeaderName = HeaderName::from_static("neon-raw-text-output");
      74              : static ARRAY_MODE: HeaderName = HeaderName::from_static("neon-array-mode");
      75              : static ALLOW_POOL: HeaderName = HeaderName::from_static("neon-pool-opt-in");
      76              : static TXN_ISOLATION_LEVEL: HeaderName = HeaderName::from_static("neon-batch-isolation-level");
      77              : static TXN_READ_ONLY: HeaderName = HeaderName::from_static("neon-batch-read-only");
      78              : static TXN_DEFERRABLE: HeaderName = HeaderName::from_static("neon-batch-deferrable");
      79              : 
      80              : static HEADER_VALUE_TRUE: HeaderValue = HeaderValue::from_static("true");
      81              : 
      82            3 : fn bytes_to_pg_text<'de, D>(deserializer: D) -> Result<Vec<Option<String>>, D::Error>
      83            3 : where
      84            3 :     D: serde::de::Deserializer<'de>,
      85            3 : {
      86              :     // TODO: consider avoiding the allocation here.
      87            3 :     let json: Vec<Value> = serde::de::Deserialize::deserialize(deserializer)?;
      88            3 :     Ok(json_to_pg_text(json))
      89            3 : }
      90              : 
      91              : #[derive(Debug, thiserror::Error)]
      92              : pub(crate) enum ConnInfoError {
      93              :     #[error("invalid header: {0}")]
      94              :     InvalidHeader(&'static HeaderName),
      95              :     #[error("invalid connection string: {0}")]
      96              :     UrlParseError(#[from] url::ParseError),
      97              :     #[error("incorrect scheme")]
      98              :     IncorrectScheme,
      99              :     #[error("missing database name")]
     100              :     MissingDbName,
     101              :     #[error("invalid database name")]
     102              :     InvalidDbName,
     103              :     #[error("missing username")]
     104              :     MissingUsername,
     105              :     #[error("invalid username: {0}")]
     106              :     InvalidUsername(#[from] std::string::FromUtf8Error),
     107              :     #[error("missing authentication credentials: {0}")]
     108              :     MissingCredentials(Credentials),
     109              :     #[error("missing hostname")]
     110              :     MissingHostname,
     111              :     #[error("invalid hostname: {0}")]
     112              :     InvalidEndpoint(#[from] ComputeUserInfoParseError),
     113              :     #[error("malformed endpoint")]
     114              :     MalformedEndpoint,
     115              : }
     116              : 
     117              : #[derive(Debug, thiserror::Error)]
     118              : pub(crate) enum Credentials {
     119              :     #[error("required password")]
     120              :     Password,
     121              :     #[error("required authorization bearer token in JWT format")]
     122              :     BearerJwt,
     123              : }
     124              : 
     125              : impl ReportableError for ConnInfoError {
     126            0 :     fn get_error_kind(&self) -> ErrorKind {
     127            0 :         ErrorKind::User
     128            0 :     }
     129              : }
     130              : 
     131              : impl UserFacingError for ConnInfoError {
     132            0 :     fn to_string_client(&self) -> String {
     133            0 :         self.to_string()
     134            0 :     }
     135              : }
     136              : 
     137            0 : fn get_conn_info(
     138            0 :     config: &'static AuthenticationConfig,
     139            0 :     ctx: &RequestContext,
     140            0 :     headers: &HeaderMap,
     141            0 :     tls: Option<&TlsConfig>,
     142            0 : ) -> Result<ConnInfoWithAuth, ConnInfoError> {
     143            0 :     let connection_string = headers
     144            0 :         .get(&CONN_STRING)
     145            0 :         .ok_or(ConnInfoError::InvalidHeader(&CONN_STRING))?
     146            0 :         .to_str()
     147            0 :         .map_err(|_| ConnInfoError::InvalidHeader(&CONN_STRING))?;
     148              : 
     149            0 :     let connection_url = Url::parse(connection_string)?;
     150              : 
     151            0 :     let protocol = connection_url.scheme();
     152            0 :     if protocol != "postgres" && protocol != "postgresql" {
     153            0 :         return Err(ConnInfoError::IncorrectScheme);
     154            0 :     }
     155              : 
     156            0 :     let mut url_path = connection_url
     157            0 :         .path_segments()
     158            0 :         .ok_or(ConnInfoError::MissingDbName)?;
     159              : 
     160            0 :     let dbname: DbName =
     161            0 :         urlencoding::decode(url_path.next().ok_or(ConnInfoError::InvalidDbName)?)?.into();
     162            0 :     ctx.set_dbname(dbname.clone());
     163              : 
     164            0 :     let username = RoleName::from(urlencoding::decode(connection_url.username())?);
     165            0 :     if username.is_empty() {
     166            0 :         return Err(ConnInfoError::MissingUsername);
     167            0 :     }
     168            0 :     ctx.set_user(username.clone());
     169              : 
     170            0 :     let auth = if let Some(auth) = headers.get(&AUTHORIZATION) {
     171            0 :         if !config.accept_jwts {
     172            0 :             return Err(ConnInfoError::MissingCredentials(Credentials::Password));
     173            0 :         }
     174              : 
     175            0 :         let auth = auth
     176            0 :             .to_str()
     177            0 :             .map_err(|_| ConnInfoError::InvalidHeader(&AUTHORIZATION))?;
     178              :         AuthData::Jwt(
     179            0 :             auth.strip_prefix("Bearer ")
     180            0 :                 .ok_or(ConnInfoError::MissingCredentials(Credentials::BearerJwt))?
     181            0 :                 .into(),
     182              :         )
     183            0 :     } else if let Some(pass) = connection_url.password() {
     184              :         // wrong credentials provided
     185            0 :         if config.accept_jwts {
     186            0 :             return Err(ConnInfoError::MissingCredentials(Credentials::BearerJwt));
     187            0 :         }
     188            0 : 
     189            0 :         AuthData::Password(match urlencoding::decode_binary(pass.as_bytes()) {
     190            0 :             std::borrow::Cow::Borrowed(b) => b.into(),
     191            0 :             std::borrow::Cow::Owned(b) => b.into(),
     192              :         })
     193            0 :     } else if config.accept_jwts {
     194            0 :         return Err(ConnInfoError::MissingCredentials(Credentials::BearerJwt));
     195              :     } else {
     196            0 :         return Err(ConnInfoError::MissingCredentials(Credentials::Password));
     197              :     };
     198              : 
     199            0 :     let endpoint = match connection_url.host() {
     200            0 :         Some(url::Host::Domain(hostname)) => {
     201            0 :             if let Some(tls) = tls {
     202            0 :                 endpoint_sni(hostname, &tls.common_names)?
     203            0 :                     .ok_or(ConnInfoError::MalformedEndpoint)?
     204              :             } else {
     205            0 :                 hostname
     206            0 :                     .split_once('.')
     207            0 :                     .map_or(hostname, |(prefix, _)| prefix)
     208            0 :                     .into()
     209              :             }
     210              :         }
     211              :         Some(url::Host::Ipv4(_) | url::Host::Ipv6(_)) | None => {
     212            0 :             return Err(ConnInfoError::MissingHostname)
     213              :         }
     214              :     };
     215            0 :     ctx.set_endpoint_id(endpoint.clone());
     216            0 : 
     217            0 :     let pairs = connection_url.query_pairs();
     218            0 : 
     219            0 :     let mut options = Option::None;
     220            0 : 
     221            0 :     let mut params = StartupMessageParamsBuilder::default();
     222            0 :     params.insert("user", &username);
     223            0 :     params.insert("database", &dbname);
     224            0 :     for (key, value) in pairs {
     225            0 :         params.insert(&key, &value);
     226            0 :         if key == "options" {
     227            0 :             options = Some(NeonOptions::parse_options_raw(&value));
     228            0 :         }
     229              :     }
     230              : 
     231            0 :     let user_info = ComputeUserInfo {
     232            0 :         endpoint,
     233            0 :         user: username,
     234            0 :         options: options.unwrap_or_default(),
     235            0 :     };
     236            0 : 
     237            0 :     let conn_info = ConnInfo { user_info, dbname };
     238            0 :     Ok(ConnInfoWithAuth { conn_info, auth })
     239            0 : }
     240              : 
     241            0 : pub(crate) async fn handle(
     242            0 :     config: &'static ProxyConfig,
     243            0 :     ctx: RequestContext,
     244            0 :     request: Request<Incoming>,
     245            0 :     backend: Arc<PoolingBackend>,
     246            0 :     cancel: CancellationToken,
     247            0 : ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     248            0 :     let result = handle_inner(cancel, config, &ctx, request, backend).await;
     249              : 
     250            0 :     let mut response = match result {
     251            0 :         Ok(r) => {
     252            0 :             ctx.set_success();
     253            0 : 
     254            0 :             // Handling the error response from local proxy here
     255            0 :             if config.authentication_config.is_auth_broker && r.status().is_server_error() {
     256            0 :                 let status = r.status();
     257              : 
     258            0 :                 let body_bytes = r
     259            0 :                     .collect()
     260            0 :                     .await
     261            0 :                     .map_err(|e| {
     262            0 :                         ApiError::InternalServerError(anyhow::Error::msg(format!(
     263            0 :                             "could not collect http body: {e}"
     264            0 :                         )))
     265            0 :                     })?
     266            0 :                     .to_bytes();
     267              : 
     268            0 :                 if let Ok(mut json_map) =
     269            0 :                     serde_json::from_slice::<IndexMap<&str, &RawValue>>(&body_bytes)
     270              :                 {
     271            0 :                     let message = json_map.get("message");
     272            0 :                     if let Some(message) = message {
     273            0 :                         let msg: String = match serde_json::from_str(message.get()) {
     274            0 :                             Ok(msg) => msg,
     275              :                             Err(_) => {
     276            0 :                                 "Unable to parse the response message from server".to_string()
     277              :                             }
     278              :                         };
     279              : 
     280            0 :                         error!("Error response from local_proxy: {status} {msg}");
     281              : 
     282            0 :                         json_map.retain(|key, _| !key.starts_with("neon:")); // remove all the neon-related keys
     283            0 : 
     284            0 :                         let resp_json = serde_json::to_string(&json_map)
     285            0 :                             .unwrap_or("failed to serialize the response message".to_string());
     286            0 : 
     287            0 :                         return json_response(status, resp_json);
     288            0 :                     }
     289            0 :                 }
     290              : 
     291            0 :                 error!("Unable to parse the response message from local_proxy");
     292            0 :                 return json_response(
     293            0 :                     status,
     294            0 :                     json!({ "message": "Unable to parse the response message from server".to_string() }),
     295            0 :                 );
     296            0 :             }
     297            0 :             r
     298              :         }
     299            0 :         Err(e @ SqlOverHttpError::Cancelled(_)) => {
     300            0 :             let error_kind = e.get_error_kind();
     301            0 :             ctx.set_error_kind(error_kind);
     302            0 : 
     303            0 :             let message = "Query cancelled, connection was terminated";
     304            0 : 
     305            0 :             tracing::info!(
     306            0 :                 kind=error_kind.to_metric_label(),
     307            0 :                 error=%e,
     308            0 :                 msg=message,
     309            0 :                 "forwarding error to user"
     310              :             );
     311              : 
     312            0 :             json_response(
     313            0 :                 StatusCode::BAD_REQUEST,
     314            0 :                 json!({ "message": message, "code": SqlState::PROTOCOL_VIOLATION.code() }),
     315            0 :             )?
     316              :         }
     317            0 :         Err(e) => {
     318            0 :             let error_kind = e.get_error_kind();
     319            0 :             ctx.set_error_kind(error_kind);
     320            0 : 
     321            0 :             let mut message = e.to_string_client();
     322            0 :             let db_error = match &e {
     323            0 :                 SqlOverHttpError::ConnectCompute(HttpConnError::PostgresConnectionError(e))
     324            0 :                 | SqlOverHttpError::Postgres(e) => e.as_db_error(),
     325            0 :                 _ => None,
     326              :             };
     327            0 :             fn get<'a, T: Default>(db: Option<&'a DbError>, x: impl FnOnce(&'a DbError) -> T) -> T {
     328            0 :                 db.map(x).unwrap_or_default()
     329            0 :             }
     330              : 
     331            0 :             if let Some(db_error) = db_error {
     332            0 :                 db_error.message().clone_into(&mut message);
     333            0 :             }
     334              : 
     335            0 :             let position = db_error.and_then(|db| db.position());
     336            0 :             let (position, internal_position, internal_query) = match position {
     337            0 :                 Some(ErrorPosition::Original(position)) => (Some(position.to_string()), None, None),
     338            0 :                 Some(ErrorPosition::Internal { position, query }) => {
     339            0 :                     (None, Some(position.to_string()), Some(query.clone()))
     340              :                 }
     341            0 :                 None => (None, None, None),
     342              :             };
     343              : 
     344            0 :             let code = get(db_error, |db| db.code().code());
     345            0 :             let severity = get(db_error, |db| db.severity());
     346            0 :             let detail = get(db_error, |db| db.detail());
     347            0 :             let hint = get(db_error, |db| db.hint());
     348            0 :             let where_ = get(db_error, |db| db.where_());
     349            0 :             let table = get(db_error, |db| db.table());
     350            0 :             let column = get(db_error, |db| db.column());
     351            0 :             let schema = get(db_error, |db| db.schema());
     352            0 :             let datatype = get(db_error, |db| db.datatype());
     353            0 :             let constraint = get(db_error, |db| db.constraint());
     354            0 :             let file = get(db_error, |db| db.file());
     355            0 :             let line = get(db_error, |db| db.line().map(|l| l.to_string()));
     356            0 :             let routine = get(db_error, |db| db.routine());
     357            0 : 
     358            0 :             tracing::info!(
     359            0 :                 kind=error_kind.to_metric_label(),
     360            0 :                 error=%e,
     361            0 :                 msg=message,
     362            0 :                 "forwarding error to user"
     363              :             );
     364              : 
     365            0 :             json_response(
     366            0 :                 e.get_http_status_code(),
     367            0 :                 json!({
     368            0 :                     "message": message,
     369            0 :                     "code": code,
     370            0 :                     "detail": detail,
     371            0 :                     "hint": hint,
     372            0 :                     "position": position,
     373            0 :                     "internalPosition": internal_position,
     374            0 :                     "internalQuery": internal_query,
     375            0 :                     "severity": severity,
     376            0 :                     "where": where_,
     377            0 :                     "table": table,
     378            0 :                     "column": column,
     379            0 :                     "schema": schema,
     380            0 :                     "dataType": datatype,
     381            0 :                     "constraint": constraint,
     382            0 :                     "file": file,
     383            0 :                     "line": line,
     384            0 :                     "routine": routine,
     385            0 :                 }),
     386            0 :             )?
     387              :         }
     388              :     };
     389              : 
     390            0 :     response
     391            0 :         .headers_mut()
     392            0 :         .insert("Access-Control-Allow-Origin", HeaderValue::from_static("*"));
     393            0 :     Ok(response)
     394            0 : }
     395              : 
     396              : #[derive(Debug, thiserror::Error)]
     397              : pub(crate) enum SqlOverHttpError {
     398              :     #[error("{0}")]
     399              :     ReadPayload(#[from] ReadPayloadError),
     400              :     #[error("{0}")]
     401              :     ConnectCompute(#[from] HttpConnError),
     402              :     #[error("{0}")]
     403              :     ConnInfo(#[from] ConnInfoError),
     404              :     #[error("response is too large (max is {0} bytes)")]
     405              :     ResponseTooLarge(usize),
     406              :     #[error("invalid isolation level")]
     407              :     InvalidIsolationLevel,
     408              :     #[error("{0}")]
     409              :     Postgres(#[from] postgres_client::Error),
     410              :     #[error("{0}")]
     411              :     JsonConversion(#[from] JsonConversionError),
     412              :     #[error("{0}")]
     413              :     Cancelled(SqlOverHttpCancel),
     414              : }
     415              : 
     416              : impl ReportableError for SqlOverHttpError {
     417            0 :     fn get_error_kind(&self) -> ErrorKind {
     418            0 :         match self {
     419            0 :             SqlOverHttpError::ReadPayload(e) => e.get_error_kind(),
     420            0 :             SqlOverHttpError::ConnectCompute(e) => e.get_error_kind(),
     421            0 :             SqlOverHttpError::ConnInfo(e) => e.get_error_kind(),
     422            0 :             SqlOverHttpError::ResponseTooLarge(_) => ErrorKind::User,
     423            0 :             SqlOverHttpError::InvalidIsolationLevel => ErrorKind::User,
     424            0 :             SqlOverHttpError::Postgres(p) => p.get_error_kind(),
     425            0 :             SqlOverHttpError::JsonConversion(_) => ErrorKind::Postgres,
     426            0 :             SqlOverHttpError::Cancelled(c) => c.get_error_kind(),
     427              :         }
     428            0 :     }
     429              : }
     430              : 
     431              : impl UserFacingError for SqlOverHttpError {
     432            0 :     fn to_string_client(&self) -> String {
     433            0 :         match self {
     434            0 :             SqlOverHttpError::ReadPayload(p) => p.to_string(),
     435            0 :             SqlOverHttpError::ConnectCompute(c) => c.to_string_client(),
     436            0 :             SqlOverHttpError::ConnInfo(c) => c.to_string_client(),
     437            0 :             SqlOverHttpError::ResponseTooLarge(_) => self.to_string(),
     438            0 :             SqlOverHttpError::InvalidIsolationLevel => self.to_string(),
     439            0 :             SqlOverHttpError::Postgres(p) => p.to_string(),
     440            0 :             SqlOverHttpError::JsonConversion(_) => "could not parse postgres response".to_string(),
     441            0 :             SqlOverHttpError::Cancelled(_) => self.to_string(),
     442              :         }
     443            0 :     }
     444              : }
     445              : 
     446              : impl HttpCodeError for SqlOverHttpError {
     447            0 :     fn get_http_status_code(&self) -> StatusCode {
     448            0 :         match self {
     449            0 :             SqlOverHttpError::ReadPayload(e) => e.get_http_status_code(),
     450            0 :             SqlOverHttpError::ConnectCompute(h) => match h.get_error_kind() {
     451            0 :                 ErrorKind::User => StatusCode::BAD_REQUEST,
     452            0 :                 _ => StatusCode::INTERNAL_SERVER_ERROR,
     453              :             },
     454            0 :             SqlOverHttpError::ConnInfo(_) => StatusCode::BAD_REQUEST,
     455            0 :             SqlOverHttpError::ResponseTooLarge(_) => StatusCode::INSUFFICIENT_STORAGE,
     456            0 :             SqlOverHttpError::InvalidIsolationLevel => StatusCode::BAD_REQUEST,
     457            0 :             SqlOverHttpError::Postgres(_) => StatusCode::BAD_REQUEST,
     458            0 :             SqlOverHttpError::JsonConversion(_) => StatusCode::INTERNAL_SERVER_ERROR,
     459            0 :             SqlOverHttpError::Cancelled(_) => StatusCode::INTERNAL_SERVER_ERROR,
     460              :         }
     461            0 :     }
     462              : }
     463              : 
     464              : #[derive(Debug, thiserror::Error)]
     465              : pub(crate) enum ReadPayloadError {
     466              :     #[error("could not read the HTTP request body: {0}")]
     467              :     Read(#[from] hyper::Error),
     468              :     #[error("request is too large (max is {limit} bytes)")]
     469              :     BodyTooLarge { limit: usize },
     470              :     #[error("could not parse the HTTP request body: {0}")]
     471              :     Parse(#[from] serde_json::Error),
     472              : }
     473              : 
     474              : impl From<ReadBodyError<hyper::Error>> for ReadPayloadError {
     475            0 :     fn from(value: ReadBodyError<hyper::Error>) -> Self {
     476            0 :         match value {
     477            0 :             ReadBodyError::BodyTooLarge { limit } => Self::BodyTooLarge { limit },
     478            0 :             ReadBodyError::Read(e) => Self::Read(e),
     479              :         }
     480            0 :     }
     481              : }
     482              : 
     483              : impl ReportableError for ReadPayloadError {
     484            0 :     fn get_error_kind(&self) -> ErrorKind {
     485            0 :         match self {
     486            0 :             ReadPayloadError::Read(_) => ErrorKind::ClientDisconnect,
     487            0 :             ReadPayloadError::BodyTooLarge { .. } => ErrorKind::User,
     488            0 :             ReadPayloadError::Parse(_) => ErrorKind::User,
     489              :         }
     490            0 :     }
     491              : }
     492              : 
     493              : impl HttpCodeError for ReadPayloadError {
     494            0 :     fn get_http_status_code(&self) -> StatusCode {
     495            0 :         match self {
     496            0 :             ReadPayloadError::Read(_) => StatusCode::BAD_REQUEST,
     497            0 :             ReadPayloadError::BodyTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE,
     498            0 :             ReadPayloadError::Parse(_) => StatusCode::BAD_REQUEST,
     499              :         }
     500            0 :     }
     501              : }
     502              : 
     503              : #[derive(Debug, thiserror::Error)]
     504              : pub(crate) enum SqlOverHttpCancel {
     505              :     #[error("query was cancelled")]
     506              :     Postgres,
     507              :     #[error("query was cancelled while stuck trying to connect to the database")]
     508              :     Connect,
     509              : }
     510              : 
     511              : impl ReportableError for SqlOverHttpCancel {
     512            0 :     fn get_error_kind(&self) -> ErrorKind {
     513            0 :         match self {
     514            0 :             SqlOverHttpCancel::Postgres => ErrorKind::ClientDisconnect,
     515            0 :             SqlOverHttpCancel::Connect => ErrorKind::ClientDisconnect,
     516              :         }
     517            0 :     }
     518              : }
     519              : 
     520              : #[derive(Clone, Copy, Debug)]
     521              : struct HttpHeaders {
     522              :     raw_output: bool,
     523              :     default_array_mode: bool,
     524              :     txn_isolation_level: Option<IsolationLevel>,
     525              :     txn_read_only: bool,
     526              :     txn_deferrable: bool,
     527              : }
     528              : 
     529              : impl HttpHeaders {
     530            0 :     fn try_parse(headers: &hyper::http::HeaderMap) -> Result<Self, SqlOverHttpError> {
     531            0 :         // Determine the output options. Default behaviour is 'false'. Anything that is not
     532            0 :         // strictly 'true' assumed to be false.
     533            0 :         let raw_output = headers.get(&RAW_TEXT_OUTPUT) == Some(&HEADER_VALUE_TRUE);
     534            0 :         let default_array_mode = headers.get(&ARRAY_MODE) == Some(&HEADER_VALUE_TRUE);
     535              : 
     536              :         // isolation level, read only and deferrable
     537            0 :         let txn_isolation_level = match headers.get(&TXN_ISOLATION_LEVEL) {
     538            0 :             Some(x) => Some(
     539            0 :                 map_header_to_isolation_level(x).ok_or(SqlOverHttpError::InvalidIsolationLevel)?,
     540              :             ),
     541            0 :             None => None,
     542              :         };
     543              : 
     544            0 :         let txn_read_only = headers.get(&TXN_READ_ONLY) == Some(&HEADER_VALUE_TRUE);
     545            0 :         let txn_deferrable = headers.get(&TXN_DEFERRABLE) == Some(&HEADER_VALUE_TRUE);
     546            0 : 
     547            0 :         Ok(Self {
     548            0 :             raw_output,
     549            0 :             default_array_mode,
     550            0 :             txn_isolation_level,
     551            0 :             txn_read_only,
     552            0 :             txn_deferrable,
     553            0 :         })
     554            0 :     }
     555              : }
     556              : 
     557            0 : fn map_header_to_isolation_level(level: &HeaderValue) -> Option<IsolationLevel> {
     558            0 :     match level.as_bytes() {
     559            0 :         b"Serializable" => Some(IsolationLevel::Serializable),
     560            0 :         b"ReadUncommitted" => Some(IsolationLevel::ReadUncommitted),
     561            0 :         b"ReadCommitted" => Some(IsolationLevel::ReadCommitted),
     562            0 :         b"RepeatableRead" => Some(IsolationLevel::RepeatableRead),
     563            0 :         _ => None,
     564              :     }
     565            0 : }
     566              : 
     567            0 : fn map_isolation_level_to_headers(level: IsolationLevel) -> Option<HeaderValue> {
     568            0 :     match level {
     569            0 :         IsolationLevel::ReadUncommitted => Some(HeaderValue::from_static("ReadUncommitted")),
     570            0 :         IsolationLevel::ReadCommitted => Some(HeaderValue::from_static("ReadCommitted")),
     571            0 :         IsolationLevel::RepeatableRead => Some(HeaderValue::from_static("RepeatableRead")),
     572            0 :         IsolationLevel::Serializable => Some(HeaderValue::from_static("Serializable")),
     573            0 :         _ => None,
     574              :     }
     575            0 : }
     576              : 
     577            0 : async fn handle_inner(
     578            0 :     cancel: CancellationToken,
     579            0 :     config: &'static ProxyConfig,
     580            0 :     ctx: &RequestContext,
     581            0 :     request: Request<Incoming>,
     582            0 :     backend: Arc<PoolingBackend>,
     583            0 : ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     584            0 :     let _requeset_gauge = Metrics::get()
     585            0 :         .proxy
     586            0 :         .connection_requests
     587            0 :         .guard(ctx.protocol());
     588            0 :     info!(
     589            0 :         protocol = %ctx.protocol(),
     590            0 :         "handling interactive connection from client"
     591              :     );
     592              : 
     593            0 :     let conn_info = get_conn_info(
     594            0 :         &config.authentication_config,
     595            0 :         ctx,
     596            0 :         request.headers(),
     597            0 :         config.tls_config.as_ref(),
     598            0 :     )?;
     599            0 :     info!(
     600            0 :         user = conn_info.conn_info.user_info.user.as_str(),
     601            0 :         "credentials"
     602              :     );
     603              : 
     604            0 :     match conn_info.auth {
     605            0 :         AuthData::Jwt(jwt) if config.authentication_config.is_auth_broker => {
     606            0 :             handle_auth_broker_inner(ctx, request, conn_info.conn_info, jwt, backend).await
     607              :         }
     608            0 :         auth => {
     609            0 :             handle_db_inner(
     610            0 :                 cancel,
     611            0 :                 config,
     612            0 :                 ctx,
     613            0 :                 request,
     614            0 :                 conn_info.conn_info,
     615            0 :                 auth,
     616            0 :                 backend,
     617            0 :             )
     618            0 :             .await
     619              :         }
     620              :     }
     621            0 : }
     622              : 
     623            0 : async fn handle_db_inner(
     624            0 :     cancel: CancellationToken,
     625            0 :     config: &'static ProxyConfig,
     626            0 :     ctx: &RequestContext,
     627            0 :     request: Request<Incoming>,
     628            0 :     conn_info: ConnInfo,
     629            0 :     auth: AuthData,
     630            0 :     backend: Arc<PoolingBackend>,
     631            0 : ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     632            0 :     //
     633            0 :     // Determine the destination and connection params
     634            0 :     //
     635            0 :     let headers = request.headers();
     636              : 
     637              :     // Allow connection pooling only if explicitly requested
     638              :     // or if we have decided that http pool is no longer opt-in
     639            0 :     let allow_pool = !config.http_config.pool_options.opt_in
     640            0 :         || headers.get(&ALLOW_POOL) == Some(&HEADER_VALUE_TRUE);
     641              : 
     642            0 :     let parsed_headers = HttpHeaders::try_parse(headers)?;
     643              : 
     644            0 :     let fetch_and_process_request = Box::pin(
     645            0 :         async {
     646            0 :             let body = read_body_with_limit(
     647            0 :                 request.into_body(),
     648            0 :                 config.http_config.max_request_size_bytes,
     649            0 :             )
     650            0 :             .await?;
     651              : 
     652            0 :             Metrics::get()
     653            0 :                 .proxy
     654            0 :                 .http_conn_content_length_bytes
     655            0 :                 .observe(HttpDirection::Request, body.len() as f64);
     656            0 : 
     657            0 :             debug!(length = body.len(), "request payload read");
     658            0 :             let payload: Payload = serde_json::from_slice(&body)?;
     659            0 :             Ok::<Payload, ReadPayloadError>(payload) // Adjust error type accordingly
     660            0 :         }
     661            0 :         .map_err(SqlOverHttpError::from),
     662            0 :     );
     663            0 : 
     664            0 :     let authenticate_and_connect = Box::pin(
     665            0 :         async {
     666            0 :             let keys = match auth {
     667            0 :                 AuthData::Password(pw) => {
     668            0 :                     backend
     669            0 :                         .authenticate_with_password(ctx, &conn_info.user_info, &pw)
     670            0 :                         .await?
     671              :                 }
     672            0 :                 AuthData::Jwt(jwt) => {
     673            0 :                     backend
     674            0 :                         .authenticate_with_jwt(ctx, &conn_info.user_info, jwt)
     675            0 :                         .await?
     676              :                 }
     677              :             };
     678              : 
     679            0 :             let client = match keys.keys {
     680            0 :                 ComputeCredentialKeys::JwtPayload(payload)
     681            0 :                     if backend.auth_backend.is_local_proxy() =>
     682              :                 {
     683            0 :                     let mut client = backend.connect_to_local_postgres(ctx, conn_info).await?;
     684            0 :                     let (cli_inner, _dsc) = client.client_inner();
     685            0 :                     cli_inner.set_jwt_session(&payload).await?;
     686            0 :                     Client::Local(client)
     687              :                 }
     688              :                 _ => {
     689            0 :                     let client = backend
     690            0 :                         .connect_to_compute(ctx, conn_info, keys, !allow_pool)
     691            0 :                         .await?;
     692            0 :                     Client::Remote(client)
     693              :                 }
     694              :             };
     695              : 
     696              :             // not strictly necessary to mark success here,
     697              :             // but it's just insurance for if we forget it somewhere else
     698            0 :             ctx.success();
     699            0 :             Ok::<_, HttpConnError>(client)
     700            0 :         }
     701            0 :         .map_err(SqlOverHttpError::from),
     702            0 :     );
     703              : 
     704            0 :     let (payload, mut client) = match run_until_cancelled(
     705            0 :         // Run both operations in parallel
     706            0 :         try_join(
     707            0 :             pin!(fetch_and_process_request),
     708            0 :             pin!(authenticate_and_connect),
     709            0 :         ),
     710            0 :         &cancel,
     711            0 :     )
     712            0 :     .await
     713              :     {
     714            0 :         Some(result) => result?,
     715            0 :         None => return Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Connect)),
     716              :     };
     717              : 
     718            0 :     let mut response = Response::builder()
     719            0 :         .status(StatusCode::OK)
     720            0 :         .header(header::CONTENT_TYPE, "application/json");
     721              : 
     722              :     // Now execute the query and return the result.
     723            0 :     let json_output = match payload {
     724            0 :         Payload::Single(stmt) => {
     725            0 :             stmt.process(&config.http_config, cancel, &mut client, parsed_headers)
     726            0 :                 .await?
     727              :         }
     728            0 :         Payload::Batch(statements) => {
     729            0 :             if parsed_headers.txn_read_only {
     730            0 :                 response = response.header(TXN_READ_ONLY.clone(), &HEADER_VALUE_TRUE);
     731            0 :             }
     732            0 :             if parsed_headers.txn_deferrable {
     733            0 :                 response = response.header(TXN_DEFERRABLE.clone(), &HEADER_VALUE_TRUE);
     734            0 :             }
     735            0 :             if let Some(txn_isolation_level) = parsed_headers
     736            0 :                 .txn_isolation_level
     737            0 :                 .and_then(map_isolation_level_to_headers)
     738            0 :             {
     739            0 :                 response = response.header(TXN_ISOLATION_LEVEL.clone(), txn_isolation_level);
     740            0 :             }
     741              : 
     742            0 :             statements
     743            0 :                 .process(&config.http_config, cancel, &mut client, parsed_headers)
     744            0 :                 .await?
     745              :         }
     746              :     };
     747              : 
     748            0 :     let metrics = client.metrics();
     749            0 : 
     750            0 :     let len = json_output.len();
     751            0 :     let response = response
     752            0 :         .body(
     753            0 :             Full::new(Bytes::from(json_output))
     754            0 :                 .map_err(|x| match x {})
     755            0 :                 .boxed(),
     756            0 :         )
     757            0 :         // only fails if invalid status code or invalid header/values are given.
     758            0 :         // these are not user configurable so it cannot fail dynamically
     759            0 :         .expect("building response payload should not fail");
     760            0 : 
     761            0 :     // count the egress bytes - we miss the TLS and header overhead but oh well...
     762            0 :     // moving this later in the stack is going to be a lot of effort and ehhhh
     763            0 :     metrics.record_egress(len as u64);
     764            0 :     Metrics::get()
     765            0 :         .proxy
     766            0 :         .http_conn_content_length_bytes
     767            0 :         .observe(HttpDirection::Response, len as f64);
     768            0 : 
     769            0 :     Ok(response)
     770            0 : }
     771              : 
     772              : static HEADERS_TO_FORWARD: &[&HeaderName] = &[
     773              :     &AUTHORIZATION,
     774              :     &CONN_STRING,
     775              :     &RAW_TEXT_OUTPUT,
     776              :     &ARRAY_MODE,
     777              :     &TXN_ISOLATION_LEVEL,
     778              :     &TXN_READ_ONLY,
     779              :     &TXN_DEFERRABLE,
     780              : ];
     781              : 
     782            0 : pub(crate) fn uuid_to_header_value(id: Uuid) -> HeaderValue {
     783            0 :     let mut uuid = [0; uuid::fmt::Hyphenated::LENGTH];
     784            0 :     HeaderValue::from_str(id.as_hyphenated().encode_lower(&mut uuid[..]))
     785            0 :         .expect("uuid hyphenated format should be all valid header characters")
     786            0 : }
     787              : 
     788            0 : async fn handle_auth_broker_inner(
     789            0 :     ctx: &RequestContext,
     790            0 :     request: Request<Incoming>,
     791            0 :     conn_info: ConnInfo,
     792            0 :     jwt: String,
     793            0 :     backend: Arc<PoolingBackend>,
     794            0 : ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, SqlOverHttpError> {
     795            0 :     backend
     796            0 :         .authenticate_with_jwt(ctx, &conn_info.user_info, jwt)
     797            0 :         .await
     798            0 :         .map_err(HttpConnError::from)?;
     799              : 
     800            0 :     let mut client = backend.connect_to_local_proxy(ctx, conn_info).await?;
     801              : 
     802            0 :     let local_proxy_uri = ::http::Uri::from_static("http://proxy.local/sql");
     803            0 : 
     804            0 :     let (mut parts, body) = request.into_parts();
     805            0 :     let mut req = Request::builder().method(Method::POST).uri(local_proxy_uri);
     806              : 
     807              :     // todo(conradludgate): maybe auth-broker should parse these and re-serialize
     808              :     // these instead just to ensure they remain normalised.
     809            0 :     for &h in HEADERS_TO_FORWARD {
     810            0 :         if let Some(hv) = parts.headers.remove(h) {
     811            0 :             req = req.header(h, hv);
     812            0 :         }
     813              :     }
     814            0 :     req = req.header(&NEON_REQUEST_ID, uuid_to_header_value(ctx.session_id()));
     815            0 : 
     816            0 :     let req = req
     817            0 :         .body(body)
     818            0 :         .expect("all headers and params received via hyper should be valid for request");
     819            0 : 
     820            0 :     // todo: map body to count egress
     821            0 :     let _metrics = client.metrics();
     822            0 : 
     823            0 :     Ok(client
     824            0 :         .inner
     825            0 :         .inner
     826            0 :         .send_request(req)
     827            0 :         .await
     828            0 :         .map_err(LocalProxyConnError::from)
     829            0 :         .map_err(HttpConnError::from)?
     830            0 :         .map(|b| b.boxed()))
     831            0 : }
     832              : 
     833              : impl QueryData {
     834            0 :     async fn process(
     835            0 :         self,
     836            0 :         config: &'static HttpConfig,
     837            0 :         cancel: CancellationToken,
     838            0 :         client: &mut Client,
     839            0 :         parsed_headers: HttpHeaders,
     840            0 :     ) -> Result<String, SqlOverHttpError> {
     841            0 :         let (inner, mut discard) = client.inner();
     842            0 :         let cancel_token = inner.cancel_token();
     843              : 
     844            0 :         let res = match select(
     845            0 :             pin!(query_to_json(config, &*inner, self, &mut 0, parsed_headers)),
     846            0 :             pin!(cancel.cancelled()),
     847            0 :         )
     848            0 :         .await
     849              :         {
     850              :             // The query successfully completed.
     851            0 :             Either::Left((Ok((status, results)), __not_yet_cancelled)) => {
     852            0 :                 discard.check_idle(status);
     853            0 : 
     854            0 :                 let json_output =
     855            0 :                     serde_json::to_string(&results).expect("json serialization should not fail");
     856            0 :                 Ok(json_output)
     857              :             }
     858              :             // The query failed with an error
     859            0 :             Either::Left((Err(e), __not_yet_cancelled)) => {
     860            0 :                 discard.discard();
     861            0 :                 return Err(e);
     862              :             }
     863              :             // The query was cancelled.
     864            0 :             Either::Right((_cancelled, query)) => {
     865            0 :                 tracing::info!("cancelling query");
     866            0 :                 if let Err(err) = cancel_token.cancel_query(NoTls).await {
     867            0 :                     tracing::warn!(?err, "could not cancel query");
     868            0 :                 }
     869              :                 // wait for the query cancellation
     870            0 :                 match time::timeout(time::Duration::from_millis(100), query).await {
     871              :                     // query successed before it was cancelled.
     872            0 :                     Ok(Ok((status, results))) => {
     873            0 :                         discard.check_idle(status);
     874            0 : 
     875            0 :                         let json_output = serde_json::to_string(&results)
     876            0 :                             .expect("json serialization should not fail");
     877            0 :                         Ok(json_output)
     878              :                     }
     879              :                     // query failed or was cancelled.
     880            0 :                     Ok(Err(error)) => {
     881            0 :                         let db_error = match &error {
     882              :                             SqlOverHttpError::ConnectCompute(
     883            0 :                                 HttpConnError::PostgresConnectionError(e),
     884              :                             )
     885            0 :                             | SqlOverHttpError::Postgres(e) => e.as_db_error(),
     886            0 :                             _ => None,
     887              :                         };
     888              : 
     889              :                         // if errored for some other reason, it might not be safe to return
     890            0 :                         if !db_error.is_some_and(|e| *e.code() == SqlState::QUERY_CANCELED) {
     891            0 :                             discard.discard();
     892            0 :                         }
     893              : 
     894            0 :                         Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Postgres))
     895              :                     }
     896            0 :                     Err(_timeout) => {
     897            0 :                         discard.discard();
     898            0 :                         Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Postgres))
     899              :                     }
     900              :                 }
     901              :             }
     902              :         };
     903            0 :         res
     904            0 :     }
     905              : }
     906              : 
     907              : impl BatchQueryData {
     908            0 :     async fn process(
     909            0 :         self,
     910            0 :         config: &'static HttpConfig,
     911            0 :         cancel: CancellationToken,
     912            0 :         client: &mut Client,
     913            0 :         parsed_headers: HttpHeaders,
     914            0 :     ) -> Result<String, SqlOverHttpError> {
     915            0 :         info!("starting transaction");
     916            0 :         let (inner, mut discard) = client.inner();
     917            0 :         let cancel_token = inner.cancel_token();
     918            0 :         let mut builder = inner.build_transaction();
     919            0 :         if let Some(isolation_level) = parsed_headers.txn_isolation_level {
     920            0 :             builder = builder.isolation_level(isolation_level);
     921            0 :         }
     922            0 :         if parsed_headers.txn_read_only {
     923            0 :             builder = builder.read_only(true);
     924            0 :         }
     925            0 :         if parsed_headers.txn_deferrable {
     926            0 :             builder = builder.deferrable(true);
     927            0 :         }
     928              : 
     929            0 :         let transaction = builder.start().await.inspect_err(|_| {
     930            0 :             // if we cannot start a transaction, we should return immediately
     931            0 :             // and not return to the pool. connection is clearly broken
     932            0 :             discard.discard();
     933            0 :         })?;
     934              : 
     935            0 :         let json_output = match query_batch(
     936            0 :             config,
     937            0 :             cancel.child_token(),
     938            0 :             &transaction,
     939            0 :             self,
     940            0 :             parsed_headers,
     941            0 :         )
     942            0 :         .await
     943              :         {
     944            0 :             Ok(json_output) => {
     945            0 :                 info!("commit");
     946            0 :                 let status = transaction.commit().await.inspect_err(|_| {
     947            0 :                     // if we cannot commit - for now don't return connection to pool
     948            0 :                     // TODO: get a query status from the error
     949            0 :                     discard.discard();
     950            0 :                 })?;
     951            0 :                 discard.check_idle(status);
     952            0 :                 json_output
     953              :             }
     954              :             Err(SqlOverHttpError::Cancelled(_)) => {
     955            0 :                 if let Err(err) = cancel_token.cancel_query(NoTls).await {
     956            0 :                     tracing::warn!(?err, "could not cancel query");
     957            0 :                 }
     958              :                 // TODO: after cancelling, wait to see if we can get a status. maybe the connection is still safe.
     959            0 :                 discard.discard();
     960            0 : 
     961            0 :                 return Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Postgres));
     962              :             }
     963            0 :             Err(err) => {
     964            0 :                 info!("rollback");
     965            0 :                 let status = transaction.rollback().await.inspect_err(|_| {
     966            0 :                     // if we cannot rollback - for now don't return connection to pool
     967            0 :                     // TODO: get a query status from the error
     968            0 :                     discard.discard();
     969            0 :                 })?;
     970            0 :                 discard.check_idle(status);
     971            0 :                 return Err(err);
     972              :             }
     973              :         };
     974              : 
     975            0 :         Ok(json_output)
     976            0 :     }
     977              : }
     978              : 
     979            0 : async fn query_batch(
     980            0 :     config: &'static HttpConfig,
     981            0 :     cancel: CancellationToken,
     982            0 :     transaction: &Transaction<'_>,
     983            0 :     queries: BatchQueryData,
     984            0 :     parsed_headers: HttpHeaders,
     985            0 : ) -> Result<String, SqlOverHttpError> {
     986            0 :     let mut results = Vec::with_capacity(queries.queries.len());
     987            0 :     let mut current_size = 0;
     988            0 :     for stmt in queries.queries {
     989            0 :         let query = pin!(query_to_json(
     990            0 :             config,
     991            0 :             transaction,
     992            0 :             stmt,
     993            0 :             &mut current_size,
     994            0 :             parsed_headers,
     995            0 :         ));
     996            0 :         let cancelled = pin!(cancel.cancelled());
     997            0 :         let res = select(query, cancelled).await;
     998            0 :         match res {
     999              :             // TODO: maybe we should check that the transaction bit is set here
    1000            0 :             Either::Left((Ok((_, values)), _cancelled)) => {
    1001            0 :                 results.push(values);
    1002            0 :             }
    1003            0 :             Either::Left((Err(e), _cancelled)) => {
    1004            0 :                 return Err(e);
    1005              :             }
    1006            0 :             Either::Right((_cancelled, _)) => {
    1007            0 :                 return Err(SqlOverHttpError::Cancelled(SqlOverHttpCancel::Postgres));
    1008              :             }
    1009              :         }
    1010              :     }
    1011              : 
    1012            0 :     let results = json!({ "results": results });
    1013            0 :     let json_output = serde_json::to_string(&results).expect("json serialization should not fail");
    1014            0 : 
    1015            0 :     Ok(json_output)
    1016            0 : }
    1017              : 
    1018            0 : async fn query_to_json<T: GenericClient>(
    1019            0 :     config: &'static HttpConfig,
    1020            0 :     client: &T,
    1021            0 :     data: QueryData,
    1022            0 :     current_size: &mut usize,
    1023            0 :     parsed_headers: HttpHeaders,
    1024            0 : ) -> Result<(ReadyForQueryStatus, impl Serialize), SqlOverHttpError> {
    1025            0 :     let query_start = Instant::now();
    1026            0 : 
    1027            0 :     let query_params = data.params;
    1028            0 :     let mut row_stream = std::pin::pin!(client.query_raw_txt(&data.query, query_params).await?);
    1029            0 :     let query_acknowledged = Instant::now();
    1030            0 : 
    1031            0 :     // Manually drain the stream into a vector to leave row_stream hanging
    1032            0 :     // around to get a command tag. Also check that the response is not too
    1033            0 :     // big.
    1034            0 :     let mut rows: Vec<postgres_client::Row> = Vec::new();
    1035            0 :     while let Some(row) = row_stream.next().await {
    1036            0 :         let row = row?;
    1037            0 :         *current_size += row.body_len();
    1038            0 :         rows.push(row);
    1039            0 :         // we don't have a streaming response support yet so this is to prevent OOM
    1040            0 :         // from a malicious query (eg a cross join)
    1041            0 :         if *current_size > config.max_response_size_bytes {
    1042            0 :             return Err(SqlOverHttpError::ResponseTooLarge(
    1043            0 :                 config.max_response_size_bytes,
    1044            0 :             ));
    1045            0 :         }
    1046              :     }
    1047              : 
    1048            0 :     let query_resp_end = Instant::now();
    1049            0 :     let ready = row_stream.ready_status();
    1050            0 : 
    1051            0 :     // grab the command tag and number of rows affected
    1052            0 :     let command_tag = row_stream.command_tag().unwrap_or_default();
    1053            0 :     let mut command_tag_split = command_tag.split(' ');
    1054            0 :     let command_tag_name = command_tag_split.next().unwrap_or_default();
    1055            0 :     let command_tag_count = if command_tag_name == "INSERT" {
    1056              :         // INSERT returns OID first and then number of rows
    1057            0 :         command_tag_split.nth(1)
    1058              :     } else {
    1059              :         // other commands return number of rows (if any)
    1060            0 :         command_tag_split.next()
    1061              :     }
    1062            0 :     .and_then(|s| s.parse::<i64>().ok());
    1063            0 : 
    1064            0 :     info!(
    1065            0 :         rows = rows.len(),
    1066            0 :         ?ready,
    1067            0 :         command_tag,
    1068            0 :         acknowledgement = ?(query_acknowledged - query_start),
    1069            0 :         response = ?(query_resp_end - query_start),
    1070            0 :         "finished executing query"
    1071              :     );
    1072              : 
    1073            0 :     let columns_len = row_stream.columns().len();
    1074            0 :     let mut fields = Vec::with_capacity(columns_len);
    1075            0 :     let mut columns = Vec::with_capacity(columns_len);
    1076              : 
    1077            0 :     for c in row_stream.columns() {
    1078            0 :         fields.push(json!({
    1079            0 :             "name": c.name().to_owned(),
    1080            0 :             "dataTypeID": c.type_().oid(),
    1081            0 :             "tableID": c.table_oid(),
    1082            0 :             "columnID": c.column_id(),
    1083            0 :             "dataTypeSize": c.type_size(),
    1084            0 :             "dataTypeModifier": c.type_modifier(),
    1085            0 :             "format": "text",
    1086            0 :         }));
    1087            0 :         columns.push(client.get_type(c.type_oid()).await?);
    1088              :     }
    1089              : 
    1090            0 :     let array_mode = data.array_mode.unwrap_or(parsed_headers.default_array_mode);
    1091              : 
    1092              :     // convert rows to JSON
    1093            0 :     let rows = rows
    1094            0 :         .iter()
    1095            0 :         .map(|row| pg_text_row_to_json(row, &columns, parsed_headers.raw_output, array_mode))
    1096            0 :         .collect::<Result<Vec<_>, _>>()?;
    1097              : 
    1098              :     // Resulting JSON format is based on the format of node-postgres result.
    1099            0 :     let results = json!({
    1100            0 :         "command": command_tag_name.to_string(),
    1101            0 :         "rowCount": command_tag_count,
    1102            0 :         "rows": rows,
    1103            0 :         "fields": fields,
    1104            0 :         "rowAsArray": array_mode,
    1105            0 :     });
    1106            0 : 
    1107            0 :     Ok((ready, results))
    1108            0 : }
    1109              : 
    1110              : enum Client {
    1111              :     Remote(conn_pool_lib::Client<postgres_client::Client>),
    1112              :     Local(conn_pool_lib::Client<postgres_client::Client>),
    1113              : }
    1114              : 
    1115              : enum Discard<'a> {
    1116              :     Remote(conn_pool_lib::Discard<'a, postgres_client::Client>),
    1117              :     Local(conn_pool_lib::Discard<'a, postgres_client::Client>),
    1118              : }
    1119              : 
    1120              : impl Client {
    1121            0 :     fn metrics(&self) -> Arc<MetricCounter> {
    1122            0 :         match self {
    1123            0 :             Client::Remote(client) => client.metrics(),
    1124            0 :             Client::Local(local_client) => local_client.metrics(),
    1125              :         }
    1126            0 :     }
    1127              : 
    1128            0 :     fn inner(&mut self) -> (&mut postgres_client::Client, Discard<'_>) {
    1129            0 :         match self {
    1130            0 :             Client::Remote(client) => {
    1131            0 :                 let (c, d) = client.inner();
    1132            0 :                 (c, Discard::Remote(d))
    1133              :             }
    1134            0 :             Client::Local(local_client) => {
    1135            0 :                 let (c, d) = local_client.inner();
    1136            0 :                 (c, Discard::Local(d))
    1137              :             }
    1138              :         }
    1139            0 :     }
    1140              : }
    1141              : 
    1142              : impl Discard<'_> {
    1143            0 :     fn check_idle(&mut self, status: ReadyForQueryStatus) {
    1144            0 :         match self {
    1145            0 :             Discard::Remote(discard) => discard.check_idle(status),
    1146            0 :             Discard::Local(discard) => discard.check_idle(status),
    1147              :         }
    1148            0 :     }
    1149            0 :     fn discard(&mut self) {
    1150            0 :         match self {
    1151            0 :             Discard::Remote(discard) => discard.discard(),
    1152            0 :             Discard::Local(discard) => discard.discard(),
    1153              :         }
    1154            0 :     }
    1155              : }
    1156              : 
    1157              : #[cfg(test)]
    1158              : #[expect(clippy::unwrap_used)]
    1159              : mod tests {
    1160              :     use super::*;
    1161              : 
    1162              :     #[test]
    1163            1 :     fn test_payload() {
    1164            1 :         let payload = "{\"query\":\"SELECT * FROM users WHERE name = ?\",\"params\":[\"test\"],\"arrayMode\":true}";
    1165            1 :         let deserialized_payload: Payload = serde_json::from_str(payload).unwrap();
    1166            1 : 
    1167            1 :         match deserialized_payload {
    1168              :             Payload::Single(QueryData {
    1169            1 :                 query,
    1170            1 :                 params,
    1171            1 :                 array_mode,
    1172            1 :             }) => {
    1173            1 :                 assert_eq!(query, "SELECT * FROM users WHERE name = ?");
    1174            1 :                 assert_eq!(params, vec![Some(String::from("test"))]);
    1175            1 :                 assert!(array_mode.unwrap());
    1176              :             }
    1177              :             Payload::Batch(_) => {
    1178            0 :                 panic!("deserialization failed: case with single query, one param, and array mode")
    1179              :             }
    1180              :         }
    1181              : 
    1182            1 :         let payload = "{\"queries\":[{\"query\":\"SELECT * FROM users0 WHERE name = ?\",\"params\":[\"test0\"], \"arrayMode\":false},{\"query\":\"SELECT * FROM users1 WHERE name = ?\",\"params\":[\"test1\"],\"arrayMode\":true}]}";
    1183            1 :         let deserialized_payload: Payload = serde_json::from_str(payload).unwrap();
    1184            1 : 
    1185            1 :         match deserialized_payload {
    1186            1 :             Payload::Batch(BatchQueryData { queries }) => {
    1187            1 :                 assert_eq!(queries.len(), 2);
    1188            2 :                 for (i, query) in queries.into_iter().enumerate() {
    1189            2 :                     assert_eq!(
    1190            2 :                         query.query,
    1191            2 :                         format!("SELECT * FROM users{i} WHERE name = ?")
    1192            2 :                     );
    1193            2 :                     assert_eq!(query.params, vec![Some(format!("test{i}"))]);
    1194            2 :                     assert_eq!(query.array_mode.unwrap(), i > 0);
    1195              :                 }
    1196              :             }
    1197            0 :             Payload::Single(_) => panic!("deserialization failed: case with multiple queries"),
    1198              :         }
    1199              : 
    1200            1 :         let payload = "{\"query\":\"SELECT 1\"}";
    1201            1 :         let deserialized_payload: Payload = serde_json::from_str(payload).unwrap();
    1202            1 : 
    1203            1 :         match deserialized_payload {
    1204              :             Payload::Single(QueryData {
    1205            1 :                 query,
    1206            1 :                 params,
    1207            1 :                 array_mode,
    1208            1 :             }) => {
    1209            1 :                 assert_eq!(query, "SELECT 1");
    1210            1 :                 assert_eq!(params, vec![]);
    1211            1 :                 assert!(array_mode.is_none());
    1212              :             }
    1213            0 :             Payload::Batch(_) => panic!("deserialization failed: case with only one query"),
    1214              :         }
    1215            1 :     }
    1216              : }
        

Generated by: LCOV version 2.1-beta