wip: now compiles

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;
 

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<String, JSONValue>;
+pub type Error = anyhow::Error; // from: https://github.com/dtolnay/anyhow
+
+pub fn postgres_row_to_json_value(row: Row) -> Result<JSONValue, Error> {
+    let row_data = postgres_row_to_row_data(row)?;
+    Ok(JSONValue::Object(row_data))
+}
+
+pub fn postgres_row_to_row_data(row: Row) -> Result<RowData, Error> {
+    let mut result: Map<String, JSONValue> = 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<JSONValue, Error> {
+    let f64_to_json_number = |raw_val: f64| -> Result<JSONValue, Error> {
+        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<JSONValue, Error>,
+) -> Result<JSONValue, Error> {
+    let raw_val = row
+        .try_get::<_, Option<T>>(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<JSONValue, Error>,
+) -> Result<JSONValue, Error> {
+    let raw_val_array = row
+        .try_get::<_, Option<Vec<T>>>(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<StringCollector, Box<dyn std::error::Error + Sync + Send>> {
+        let result = std::str::from_utf8(raw)?;
+        Ok(StringCollector(result.to_owned()))
+    }
+    fn accepts(_ty: &Type) -> bool {
+        true
+    }
+}

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<Body>,
+) -> anyhow::Result<String> {
+    let get_params = request
+        .uri()
+        .query()
+        .ok_or(anyhow::anyhow!("missing query string"))?;
+
+    let parsed_params: HashMap<String, String> = 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(&params, 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<Vec<serde_json::Value>, 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,