LCOV - code coverage report
Current view: top level - proxy/src/serverless - sql_over_http.rs (source / functions) Coverage Total Hit
Test: 09e7485004805bd42b53a0c369170b3228136512.info Lines: 6.0 % 762 46
Test Date: 2024-11-21 18:36:18 Functions: 6.5 % 138 9

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

Generated by: LCOV version 2.1-beta