diff --git a/proxy/src/http.rs b/proxy/src/http.rs index a544157800..4a2f7c1152 100644 --- a/proxy/src/http.rs +++ b/proxy/src/http.rs @@ -2,6 +2,7 @@ //! Other modules should use stuff from this module instead of //! directly relying on deps like `reqwest` (think loose coupling). +pub mod pg_to_json; pub mod server; pub mod websocket; diff --git a/proxy/src/http/pg_to_json.rs b/proxy/src/http/pg_to_json.rs new file mode 100644 index 0000000000..01a106f7b9 --- /dev/null +++ b/proxy/src/http/pg_to_json.rs @@ -0,0 +1,140 @@ +// use std::iter::Map; + +// use std::Map; +use anyhow::{anyhow, Context}; +use serde_json::Map; +use tokio_postgres::{ + types::{FromSql, Type}, + Column, Row, +}; + +// some type-aliases I use in my project +pub type JSONValue = serde_json::Value; +pub type RowData = Map; +pub type Error = anyhow::Error; // from: https://github.com/dtolnay/anyhow + +pub fn postgres_row_to_json_value(row: Row) -> Result { + let row_data = postgres_row_to_row_data(row)?; + Ok(JSONValue::Object(row_data)) +} + +pub fn postgres_row_to_row_data(row: Row) -> Result { + let mut result: Map = Map::new(); + for (i, column) in row.columns().iter().enumerate() { + let name = column.name(); + let json_value = pg_cell_to_json_value(&row, column, i)?; + result.insert(name.to_string(), json_value); + } + Ok(result) +} + +pub fn pg_cell_to_json_value( + row: &Row, + column: &Column, + column_i: usize, +) -> Result { + let f64_to_json_number = |raw_val: f64| -> Result { + let temp = serde_json::Number::from_f64(raw_val).ok_or(anyhow!("invalid json-float"))?; + Ok(JSONValue::Number(temp)) + }; + Ok(match *column.type_() { + // for rust-postgres <> postgres type-mappings: https://docs.rs/postgres/latest/postgres/types/trait.FromSql.html#types + // for postgres types: https://www.postgresql.org/docs/7.4/datatype.html#DATATYPE-TABLE + + // single types + Type::BOOL => get_basic(row, column, column_i, |a: bool| Ok(JSONValue::Bool(a)))?, + Type::INT2 => get_basic(row, column, column_i, |a: i16| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::INT4 => get_basic(row, column, column_i, |a: i32| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::INT8 => get_basic(row, column, column_i, |a: i64| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::TEXT | Type::VARCHAR => { + get_basic(row, column, column_i, |a: String| Ok(JSONValue::String(a)))? + } + // Type::JSON | Type::JSONB => get_basic(row, column, column_i, |a: JSONValue| Ok(a))?, + Type::FLOAT4 => get_basic(row, column, column_i, |a: f32| f64_to_json_number(a.into()))?, + Type::FLOAT8 => get_basic(row, column, column_i, f64_to_json_number)?, + // these types require a custom StringCollector struct as an intermediary (see struct at bottom) + Type::TS_VECTOR => get_basic(row, column, column_i, |a: StringCollector| { + Ok(JSONValue::String(a.0)) + })?, + + // array types + Type::BOOL_ARRAY => get_array(row, column, column_i, |a: bool| Ok(JSONValue::Bool(a)))?, + Type::INT2_ARRAY => get_array(row, column, column_i, |a: i16| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::INT4_ARRAY => get_array(row, column, column_i, |a: i32| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::INT8_ARRAY => get_array(row, column, column_i, |a: i64| { + Ok(JSONValue::Number(serde_json::Number::from(a))) + })?, + Type::TEXT_ARRAY | Type::VARCHAR_ARRAY => { + get_array(row, column, column_i, |a: String| Ok(JSONValue::String(a)))? + } + // Type::JSON_ARRAY | Type::JSONB_ARRAY => get_array(row, column, column_i, |a: JSONValue| Ok(a))?, + Type::FLOAT4_ARRAY => get_array(row, column, column_i, f64_to_json_number)?, + Type::FLOAT8_ARRAY => get_array(row, column, column_i, f64_to_json_number)?, + // these types require a custom StringCollector struct as an intermediary (see struct at bottom) + Type::TS_VECTOR_ARRAY => get_array(row, column, column_i, |a: StringCollector| { + Ok(JSONValue::String(a.0)) + })?, + + _ => anyhow::bail!( + "Cannot convert pg-cell \"{}\" of type \"{}\" to a JSONValue.", + column.name(), + column.type_().name() + ), + }) +} + +fn get_basic<'a, T: FromSql<'a>>( + row: &'a Row, + column: &Column, + column_i: usize, + val_to_json_val: impl Fn(T) -> Result, +) -> Result { + let raw_val = row + .try_get::<_, Option>(column_i) + .with_context(|| format!("column_name:{}", column.name()))?; + raw_val.map_or(Ok(JSONValue::Null), val_to_json_val) +} +fn get_array<'a, T: FromSql<'a>>( + row: &'a Row, + column: &Column, + column_i: usize, + val_to_json_val: impl Fn(T) -> Result, +) -> Result { + let raw_val_array = row + .try_get::<_, Option>>(column_i) + .with_context(|| format!("column_name:{}", column.name()))?; + Ok(match raw_val_array { + Some(val_array) => { + let mut result = vec![]; + for val in val_array { + result.push(val_to_json_val(val)?); + } + JSONValue::Array(result) + } + None => JSONValue::Null, + }) +} + +struct StringCollector(String); +impl FromSql<'_> for StringCollector { + fn from_sql( + _: &Type, + raw: &[u8], + ) -> Result> { + let result = std::str::from_utf8(raw)?; + Ok(StringCollector(result.to_owned())) + } + fn accepts(_ty: &Type) -> bool { + true + } +} diff --git a/proxy/src/http/websocket.rs b/proxy/src/http/websocket.rs index c7676e8e14..bd97f18209 100644 --- a/proxy/src/http/websocket.rs +++ b/proxy/src/http/websocket.rs @@ -1,8 +1,11 @@ +use crate::http::pg_to_json::postgres_row_to_json_value; use crate::{ - cancellation::CancelMap, config::ProxyConfig, error::io_error, proxy::handle_ws_client, + auth, cancellation::CancelMap, config::ProxyConfig, console, error::io_error, + proxy::handle_ws_client, }; use bytes::{Buf, Bytes}; use futures::{Sink, Stream, StreamExt}; +use hashbrown::HashMap; use hyper::{ server::{accept, conn::AddrIncoming}, upgrade::Upgraded, @@ -10,6 +13,8 @@ use hyper::{ }; use hyper_tungstenite::{tungstenite::Message, HyperWebsocket, WebSocketStream}; use pin_project_lite::pin_project; +use pq_proto::StartupMessageParams; + use std::{ convert::Infallible, future::ready, @@ -24,6 +29,7 @@ use tokio::{ }; use tokio_util::sync::CancellationToken; use tracing::{error, info, info_span, warn, Instrument}; +use url::form_urlencoded; use utils::http::{error::ApiError, json::json_response}; // TODO: use `std::sync::Exclusive` once it's stabilized. @@ -181,11 +187,104 @@ async fn ws_handler( // Return the response so the spawned future can continue. Ok(response) + } else if request.uri().path() == "/sql" { + match handle_sql(config, request).await { + Ok(resp) => json_response(StatusCode::OK, resp), + Err(e) => json_response(StatusCode::BAD_REQUEST, format!("error: {e:?}")), + } } else { - json_response(StatusCode::OK, "Connect with a websocket client") + json_response(StatusCode::BAD_REQUEST, "query is not supported") } } +// XXX: return different error codes +async fn handle_sql( + config: &'static ProxyConfig, + request: Request, +) -> anyhow::Result { + let get_params = request + .uri() + .query() + .ok_or(anyhow::anyhow!("missing query string"))?; + + let parsed_params: HashMap = form_urlencoded::parse(get_params.as_bytes()) + .into_owned() + .collect(); + + let sql = parsed_params + .get("query") + .ok_or(anyhow::anyhow!("missing query"))?; + let dbname = parsed_params + .get("dbname") + .ok_or(anyhow::anyhow!("missing dbname"))?; + let username = parsed_params + .get("username") + .ok_or(anyhow::anyhow!("missing username"))?; + let password = parsed_params + .get("password") + .ok_or(anyhow::anyhow!("missing password"))?; + // XXX: does URI includes host too? then Url::parse() should work for both host_str and params + let hostname = request + .headers() + .get("host") + .and_then(|h| h.to_str().ok()) + .and_then(|h| h.split(':').next()) + .map(|s| s.to_string()) + .ok_or(anyhow::anyhow!("missing host header"))?; + + let params = StartupMessageParams::new([]); + let tls = config.tls_config.as_ref(); + let common_names = tls.and_then(|tls| tls.common_names.clone()); + let creds = config + .auth_backend + .as_ref() + .map(|_| auth::ClientCredentials::parse(¶ms, Some(hostname.as_str()), common_names)) + .transpose()?; + + let extra = console::ConsoleReqExtra { + session_id: uuid::Uuid::new_v4(), + application_name: Some("proxy_http_sql"), + }; + + let node = creds.wake_compute(&extra).await?.expect("msg"); + let conf = node.value.config; + + let host = match conf.get_hosts().first().expect("no host") { + tokio_postgres::config::Host::Tcp(host) => host, + tokio_postgres::config::Host::Unix(_) => { + return Err(anyhow::anyhow!("unix socket is not supported")); + } + }; + + let (client, connection) = tokio_postgres::connect( + &format!( + "host={} port={} user={} password={} dbname={}", + host, + conf.get_ports().first().expect("no port"), + username, + password, + dbname + ), + tokio_postgres::NoTls, + ) + .await?; + + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {}", e); + } + }); + + let rows: Result, anyhow::Error> = client + .query(sql, &[]) + .await? + .into_iter() + .map(postgres_row_to_json_value) + .collect(); + + Ok(serde_json::to_string(&rows?)?) +} + pub async fn task_main( config: &'static ProxyConfig, ws_listener: TcpListener,