diff --git a/Cargo.lock b/Cargo.lock index c148fb3d0b..dfd0f9d9ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,6 +130,12 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "array-init" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d62b7694a562cdf5a74227903507c56ab2cc8bdd1f781ed5cb4cf9c9f810bfc" + [[package]] name = "asn1-rs" version = "0.5.2" @@ -2885,6 +2891,7 @@ name = "postgres-types" version = "0.2.4" source = "git+https://github.com/neondatabase/rust-postgres.git?rev=43e6db254a97fdecbce33d8bc0890accfd74495e#43e6db254a97fdecbce33d8bc0890accfd74495e" dependencies = [ + "array-init", "bytes", "chrono", "fallible-iterator", diff --git a/Cargo.toml b/Cargo.toml index 520e48baad..ddb07dea6e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -124,10 +124,10 @@ env_logger = "0.10" log = "0.4" ## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed -postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e", features = ["with-chrono-0_4"] } +postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e", features = ["with-chrono-0_4", "array-impls"] } postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" } -postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" } -tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" } +postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e", features = ["array-impls"] } +tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e", features = ["array-impls"] } tokio-tar = { git = "https://github.com/neondatabase/tokio-tar.git", rev="404df61437de0feef49ba2ccdbdd94eb8ad6e142" } ## Other git libraries diff --git a/proxy/src/http/websocket.rs b/proxy/src/http/websocket.rs index 3d9850fed9..072a88dc46 100644 --- a/proxy/src/http/websocket.rs +++ b/proxy/src/http/websocket.rs @@ -6,6 +6,7 @@ use crate::{ use bytes::{Buf, Bytes}; use futures::{Sink, Stream, StreamExt, TryStreamExt}; use tokio_postgres::Row; +use tokio_postgres::error::DbError; use std::collections::HashMap; use hyper::{ server::{accept, conn::AddrIncoming}, @@ -16,7 +17,7 @@ use hyper_tungstenite::{tungstenite::Message, HyperWebsocket, WebSocketStream}; use percent_encoding::percent_decode; use pin_project_lite::pin_project; use pq_proto::StartupMessageParams; -use serde_json::Value; +use serde_json::{Value, json}; use tokio::sync::Mutex; use tokio_postgres::types::{ToSql}; @@ -196,16 +197,32 @@ async fn ws_handler( Ok(response) } else if request.uri().path() == "/sql" && request.method() == Method::POST { - match handle_sql(config, request).await { - Ok(resp) => json_response(StatusCode::OK, resp).map(|mut r| { - r.headers_mut().insert( - "Access-Control-Allow-Origin", - hyper::http::HeaderValue::from_static("*"), - ); - r - }), - Err(e) => json_response(StatusCode::BAD_REQUEST, format!("error: {e:?}")), - } + let result = handle_sql(config, request).await; + let status_code = match result { + Ok(_) => StatusCode::OK, + Err(_) => StatusCode::BAD_REQUEST + }; + let json = match result { + Ok(r) => serde_json::to_value(r).unwrap(), + Err(e) => { + let message = format!("{:?}", e); + let code = match e.downcast_ref::() { + Some(e) => match e.code() { + Some(e) => serde_json::to_value(e.code()).unwrap(), + None => Value::Null + }, + None => Value::Null + }; + json!({ "message": message, "code": code }) + } + }; + json_response(status_code, json).map(|mut r| { + r.headers_mut().insert( + "Access-Control-Allow-Origin", + hyper::http::HeaderValue::from_static("*"), + ); + r + }) } else if request.uri().path() == "/sleep" { // sleep 15ms @@ -217,6 +234,21 @@ async fn ws_handler( } } +fn boxed_from_json(json: Vec) -> Result>, anyhow::Error> { + json.iter().map(|value| { + let boxed: Result, anyhow::Error> = match value { + Value::Bool(b) => Ok(Box::new(b.clone())), + Value::Number(n) => Ok(Box::new(n.as_f64().unwrap())), + Value::String(s) => Ok(Box::new(s.clone())), + // TODO: support null (not like this: `Value::Null => Ok(Box::new(None::)),`) + // TODO: support arrays + x => Err(anyhow::anyhow!("unsupported param {:?}", x)) + }; + boxed + }) + .collect::, anyhow::Error>>() +} + // XXX: return different error codes async fn handle_sql( config: &'static ProxyConfig, @@ -232,6 +264,11 @@ async fn handle_sql( let connection_url = Url::parse(connection_string)?; + let protocol = connection_url.scheme(); + if protocol != "postgres" && protocol != "postgresql" { + return Err(anyhow::anyhow!("connection string must start with postgres: or postgresql:")) + } + let mut url_path = connection_url .path_segments() .ok_or(anyhow::anyhow!("missing database name"))?; @@ -320,17 +357,9 @@ async fn handle_sql( }); let query = &query_data.query; - let query_params = query_data.params.iter().map(|value| { - let boxed: Box = match value { - Value::Null => Box::new(None::), - Value::Bool(b) => Box::new(b.clone()), - Value::Number(n) => Box::new(n.as_f64().unwrap()), - Value::String(s) => Box::new(s.clone()), - _ => panic!("wrong parameter type") - }; - boxed - }).collect::>(); + let query_params = boxed_from_json(query_data.params)?; + // TODO: find a way to catch the panic and return an error if the number of params is wrong let pg_rows: Vec = client .query_raw(query, query_params) .await?