Files
HugoCasa b753e4d06c feat: add support for bytea in pg (#3926)
* feat: add support for bytea in pg

* fix: editor nits
2024-06-18 17:38:46 +02:00

685 lines
24 KiB
Rust

use std::net::IpAddr;
use std::sync::atomic::{AtomicBool, AtomicU64};
use std::sync::Arc;
use std::time::Duration;
use anyhow::Context;
use base64::{engine, Engine as _};
use chrono::Utc;
use futures::TryStreamExt;
use native_tls::{Certificate, TlsConnector};
use postgres_native_tls::MakeTlsConnector;
use rust_decimal::{prelude::FromPrimitive, Decimal};
use serde::Deserialize;
use serde_json::value::RawValue;
use serde_json::Map;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio_postgres::types::IsNull;
use tokio_postgres::{
types::{to_sql_checked, ToSql},
NoTls, Row,
};
use tokio_postgres::{
types::{FromSql, Type},
Column,
};
use uuid::Uuid;
use windmill_common::error::{self, Error};
use windmill_common::worker::{to_raw_value, CLOUD_HOSTED};
use windmill_common::{error::to_anyhow, jobs::QueuedJob};
use windmill_parser::Typ;
use windmill_parser_sql::{parse_db_resource, parse_pgsql_sig};
use windmill_queue::CanceledBy;
use crate::common::{build_args_values, run_future_with_polling_update_job_poller, sizeof_val};
use crate::{AuthedClientBackgroundTask, MAX_RESULT_SIZE};
use bytes::{Buf, BytesMut};
use lazy_static::lazy_static;
use urlencoding::encode;
#[derive(Deserialize)]
struct PgDatabase {
host: String,
user: Option<String>,
password: Option<String>,
port: Option<u16>,
sslmode: Option<String>,
dbname: String,
root_certificate_pem: Option<String>,
}
lazy_static! {
pub static ref CONNECTION_CACHE: Arc<Mutex<Option<(String, tokio_postgres::Client)>>> =
Arc::new(Mutex::new(None));
pub static ref LAST_QUERY: AtomicU64 = AtomicU64::new(0);
pub static ref RUNNING: AtomicBool = AtomicBool::new(false);
}
pub async fn do_postgresql(
job: &QueuedJob,
client: &AuthedClientBackgroundTask,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
worker_name: &str,
column_order: &mut Option<Vec<String>>,
) -> error::Result<Box<RawValue>> {
let pg_args = build_args_values(job, client, db).await?;
let inline_db_res_path = parse_db_resource(&query);
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
)
.await?,
)
} else {
pg_args.get("database").cloned()
};
let database = if let Some(db) = db_arg {
serde_json::from_value::<PgDatabase>(db.clone())
.map_err(|e| Error::ExecutionErr(e.to_string()))?
} else {
return Err(Error::BadRequest("Missing database argument".to_string()));
};
let sslmode = match database.sslmode.as_deref() {
Some("allow") => "prefer".to_string(),
Some("verify-ca") | Some("verify-full") => "require".to_string(),
Some(s) => s.to_string(),
None => "prefer".to_string(),
};
let database_string = format!(
"postgres://{user}:{password}@{host}:{port}/{dbname}?sslmode={sslmode}",
user = encode(&database.user.unwrap_or("postgres".to_string())),
password = encode(&database.password.unwrap_or("".to_string())),
host = encode(&database.host),
port = database.port.unwrap_or(5432),
dbname = database.dbname,
sslmode = sslmode
);
let database_string_clone = database_string.clone();
RUNNING.store(true, std::sync::atomic::Ordering::Relaxed);
LAST_QUERY.store(
chrono::Utc::now().timestamp().try_into().unwrap_or(0),
std::sync::atomic::Ordering::Relaxed,
);
let mtex;
if !*CLOUD_HOSTED {
mtex = Some(CONNECTION_CACHE.lock().await);
} else {
mtex = None;
}
let has_cached_con = mtex
.as_ref()
.is_some_and(|x| x.as_ref().is_some_and(|y| y.0 == database_string));
let new_client = if has_cached_con {
tracing::info!("Using cached connection");
None
} else if sslmode == "require" {
tracing::info!("Creating new connection");
let mut connector = TlsConnector::builder();
if let Some(root_certificate_pem) = database.root_certificate_pem {
if !root_certificate_pem.is_empty() {
connector.add_root_certificate(
Certificate::from_pem(root_certificate_pem.as_bytes())
.map_err(|e| error::Error::BadConfig(format!("Invalid Certs: {e:#}")))?,
);
} else {
connector.danger_accept_invalid_certs(true);
connector.danger_accept_invalid_hostnames(true);
}
} else {
connector
.danger_accept_invalid_certs(true)
.danger_accept_invalid_hostnames(true);
}
let (client, connection) = tokio_postgres::connect(
&database_string,
MakeTlsConnector::new(connector.build().map_err(to_anyhow)?),
)
.await
.map_err(to_anyhow)?;
let handle = tokio::spawn(async move {
if let Err(e) = connection.await {
let mut mtex = CONNECTION_CACHE.lock().await;
*mtex = None;
tracing::error!("connection error: {}", e);
}
});
Some((client, handle))
} else {
tracing::info!("Creating new connection");
let (client, connection) = tokio_postgres::connect(&database_string, NoTls)
.await
.map_err(to_anyhow)?;
let handle = tokio::spawn(async move {
if let Err(e) = connection.await {
let mut mtex = CONNECTION_CACHE.lock().await;
*mtex = None;
tracing::error!("connection error: {}", e);
}
});
Some((client, handle))
};
let mut statement_values: Vec<serde_json::Value> = vec![];
let sig = parse_pgsql_sig(&query)
.map_err(|x| Error::ExecutionErr(x.to_string()))?
.args;
for arg in &sig {
statement_values.push(
pg_args
.get(&arg.name)
.map(|x| x.to_owned())
.unwrap_or_else(|| serde_json::Value::Null),
);
}
let query_params = statement_values
.iter()
.enumerate()
.map(|(i, value)| {
let arg_t = &sig[i]
.otyp
.as_ref()
.ok_or_else(|| anyhow::anyhow!("Missing otyp for pg arg"))?
.to_owned();
let typ = &sig[i].typ;
convert_val(value, arg_t, typ)
})
.collect::<windmill_common::error::Result<Vec<_>>>()?;
let (client, handle) = if let Some((client, handle)) = new_client.as_ref() {
(client, Some(handle))
} else {
let (_, client) = mtex.as_ref().unwrap().as_ref().unwrap();
(client, None)
};
let result_f = async {
// Now we can execute a simple statement that just returns its parameter.
let rows = client
.query_raw(query, query_params)
.await
.map_err(to_anyhow)?;
let rows = rows.try_collect::<Vec<Row>>().await.map_err(to_anyhow)?;
*column_order = Some(
rows.first()
.map(|x| {
x.columns()
.iter()
.map(|x| x.name().to_string())
.collect::<Vec<String>>()
})
.unwrap_or_default(),
);
let mut siz = 0;
let mut res: Vec<serde_json::Value> = vec![];
for row in rows.into_iter() {
let r = postgres_row_to_json_value(row);
if let Ok(v) = r.as_ref() {
let size = sizeof_val(v);
siz += size;
}
if *CLOUD_HOSTED && siz > MAX_RESULT_SIZE * 4 {
return Err(anyhow::anyhow!(
"Query result too large for cloud (size = {} > {})",
siz,
MAX_RESULT_SIZE & 4
));
}
if let Ok(v) = r {
res.push(v);
} else {
return Err(to_anyhow(r.err().unwrap()));
}
}
Ok((res, siz))
};
let (result, size) = run_future_with_polling_update_job_poller(
job.id,
job.timeout,
db,
mem_peak,
canceled_by,
result_f,
worker_name,
&job.workspace_id,
)
.await?;
*mem_peak = size as i32;
RUNNING.store(false, std::sync::atomic::Ordering::Relaxed);
if let Some(handle) = handle {
if let Some(mut mtex) = mtex {
let abort_handler = handle.abort_handle();
if let Some(new_client) = new_client {
*mtex = Some((database_string, new_client.0));
}
drop(mtex);
LAST_QUERY.store(
chrono::Utc::now().timestamp().try_into().unwrap_or(0),
std::sync::atomic::Ordering::Relaxed,
);
tokio::spawn(async move {
loop {
tokio::time::sleep(Duration::from_secs(5)).await;
let last_query = LAST_QUERY.load(std::sync::atomic::Ordering::Relaxed);
let now = chrono::Utc::now().timestamp().try_into().unwrap_or(0);
//we cache connection for 5 minutes at most
if last_query + 60 * 5 < now
&& !RUNNING.load(std::sync::atomic::Ordering::Relaxed)
{
tracing::info!("Closing cache connection due to inactivity");
break;
}
let mtex = CONNECTION_CACHE.lock().await;
if mtex.is_none() {
// connection is not in the mutex anymore
break;
} else if let Some(mtex) = mtex.as_ref() {
if mtex.0.as_str() != &database_string_clone {
// connection is not the latest one
break;
}
}
tracing::debug!("Keeping cached connection alive due to activity")
}
let mut mtex = CONNECTION_CACHE.lock().await;
*mtex = None;
abort_handler.abort();
});
} else {
handle.abort();
}
}
let raw_result = to_raw_value(&result);
*mem_peak = (raw_result.get().len() / 1000) as i32;
// And then check that we got back the same string we sent over.
return Ok(raw_result);
}
#[derive(Debug)]
enum PgType {
String(String),
Bool(bool),
I8(i8),
I16(i16),
I32(i32),
I64(i64),
U32(u32),
F32(f32),
F64(f64),
Uuid(Uuid),
Decimal(Decimal),
Date(chrono::NaiveDate),
Time(chrono::NaiveTime),
Timestamp(chrono::NaiveDateTime),
None(Option<bool>),
Array(Vec<PgType>),
Json(serde_json::Value),
Bytea(Vec<u8>),
}
impl ToSql for PgType {
fn to_sql(
&self,
ty: &Type,
out: &mut BytesMut,
) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
match *self {
PgType::String(ref val) => val.to_sql(ty, out),
PgType::Bool(ref val) => val.to_sql(ty, out),
PgType::I8(ref val) => val.to_sql(ty, out),
PgType::I16(ref val) => val.to_sql(ty, out),
PgType::I32(ref val) => val.to_sql(ty, out),
PgType::I64(ref val) => val.to_sql(ty, out),
PgType::U32(ref val) => val.to_sql(ty, out),
PgType::F32(ref val) => val.to_sql(ty, out),
PgType::F64(ref val) => val.to_sql(ty, out),
PgType::Uuid(ref val) => val.to_sql(ty, out),
PgType::Decimal(ref val) => val.to_sql(ty, out),
PgType::Date(ref val) => val.to_sql(ty, out),
PgType::Time(ref val) => val.to_sql(ty, out),
PgType::Timestamp(ref val) => val.to_sql(ty, out),
PgType::None(ref val) => val.to_sql(ty, out),
PgType::Array(ref val) => val.to_sql(ty, out),
PgType::Json(ref val) => val.to_sql(ty, out),
PgType::Bytea(ref val) => val.to_sql(ty, out),
}
}
fn accepts(_: &Type) -> bool {
true
}
to_sql_checked!();
}
fn convert_val(value: &Value, arg_t: &String, typ: &Typ) -> windmill_common::error::Result<PgType> {
match value {
Value::Array(vec) if arg_t.ends_with("[]") => {
let arg_t = arg_t.trim_end_matches("[]").to_string();
let mut result = vec![];
for val in vec {
result.push(convert_val(val, &arg_t, typ)?);
}
Ok(PgType::Array(result))
}
Value::Null => Ok(PgType::None(None::<bool>)),
Value::Bool(b) => Ok(PgType::Bool(b.clone())),
Value::Number(n) if matches!(typ, Typ::Str(_)) => Ok(PgType::String(n.to_string())),
Value::Number(n) if n.is_i64() && arg_t == "char" => {
Ok(PgType::I8(n.as_i64().unwrap() as i8))
}
Value::Number(n) if n.is_i64() && (arg_t == "smallint" || arg_t == "smallserial") => {
Ok(PgType::I16(n.as_i64().unwrap() as i16))
}
Value::Number(n)
if n.is_i64()
&& (arg_t == "int"
|| arg_t == "integer"
|| arg_t == "int4"
|| arg_t == "serial") =>
{
Ok(PgType::I32(n.as_i64().unwrap() as i32))
}
Value::Number(n) if n.is_i64() && (arg_t == "numeric" || arg_t == "decimal") => Ok(
PgType::Decimal(Decimal::from_i64(n.as_i64().unwrap()).unwrap()),
),
Value::Number(n) if n.is_i64() => Ok(PgType::I64(n.as_i64().unwrap())),
Value::Number(n) if n.is_u64() && arg_t == "oid" => {
Ok(PgType::U32(n.as_u64().unwrap() as u32))
}
Value::Number(n) if n.is_u64() && (arg_t == "bigint" || arg_t == "bigserial") => {
Ok(PgType::I64(n.as_u64().unwrap() as i64))
}
Value::Number(n) if n.is_f64() && arg_t == "real" => {
Ok(PgType::F32(n.as_f64().unwrap() as f32))
}
Value::Number(n) if n.is_f64() && arg_t == "double" => Ok(PgType::F64(n.as_f64().unwrap())),
Value::Number(n) if n.is_f64() && (arg_t == "numeric" || arg_t == "decimal") => Ok(
PgType::Decimal(Decimal::from_f64(n.as_f64().unwrap()).unwrap()),
),
Value::Number(n) => Ok(PgType::F64(n.as_f64().unwrap())),
Value::String(s) if arg_t == "uuid" => Ok(PgType::Uuid(Uuid::parse_str(s)?)),
Value::String(s) if arg_t == "date" => {
let date =
chrono::NaiveDate::parse_from_str(s, "%Y-%m-%dT%H:%M:%S.%3fZ").unwrap_or_default();
Ok(PgType::Date(date))
}
Value::String(s) if arg_t == "time" || arg_t == "timetz" => {
let time =
chrono::NaiveTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S.%3fZ").unwrap_or_default();
Ok(PgType::Time(time))
}
Value::String(s) if arg_t == "timestamp" || arg_t == "timestamptz" => {
let datetime = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S.%3fZ")
.unwrap_or_default();
Ok(PgType::Timestamp(datetime))
}
Value::String(s) if arg_t == "bytea" => {
let bytes = engine::general_purpose::STANDARD
.decode(s)
.unwrap_or(vec![]);
Ok(PgType::Bytea(bytes))
}
Value::Object(_) => Ok(PgType::Json(value.clone())),
Value::String(s) => Ok(PgType::String(s.clone())),
_ => Err(Error::ExecutionErr(format!(
"Unsupported type in query: {:?} and signature {arg_t:?}",
value
))),
}
}
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.into())
.ok_or(anyhow::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::BIT => get_basic(row, column, column_i, |a: bit_vec::BitVec| match a.len() {
1 => Ok(JSONValue::Bool(a.get(0).unwrap())),
_ => Ok(JSONValue::String(
a.iter()
.map(|x| if x { "1" } else { "0" })
.collect::<String>(),
)),
})?,
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::TIMESTAMP => get_basic(row, column, column_i, |a: chrono::NaiveDateTime| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::DATE => get_basic(row, column, column_i, |a: chrono::NaiveDate| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::TIME => get_basic(row, column, column_i, |a: chrono::NaiveTime| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::TIMESTAMPTZ => get_basic(row, column, column_i, |a: chrono::DateTime<Utc>| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::UUID => get_basic(row, column, column_i, |a: uuid::Uuid| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::INET => get_basic(row, column, column_i, |a: IpAddr| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::INTERVAL => get_basic(row, column, column_i, |a: IntervalStr| {
Ok(JSONValue::String(a.0))
})?,
Type::JSON | Type::JSONB => get_basic(row, column, column_i, |a: JSONValue| Ok(a))?,
Type::FLOAT4 => get_basic(row, column, column_i, |a: f32| {
Ok(f64_to_json_number(a.into())?)
})?,
Type::NUMERIC => get_basic(row, column, column_i, |a: Decimal| {
Ok(serde_json::to_value(a)
.map_err(|_| anyhow::anyhow!("Cannot convert decimal to json"))?)
})?,
Type::FLOAT8 => get_basic(row, column, column_i, |a: f64| f64_to_json_number(a))?,
Type::BYTEA => get_basic(row, column, column_i, |a: Vec<u8>| {
Ok(JSONValue::String(format!("\\x{}", hex::encode(a))))
})?,
// 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::BIT_ARRAY => get_array(row, column, column_i, |a: bit_vec::BitVec| match a.len() {
1 => Ok(JSONValue::Bool(a.get(0).unwrap())),
_ => Ok(JSONValue::String(
a.iter()
.map(|x| if x { "1" } else { "0" })
.collect::<String>(),
)),
})?,
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, |a: f32| {
Ok(f64_to_json_number(a.into())?)
})?,
Type::FLOAT8_ARRAY => {
get_array(row, column, column_i, |a: f64| Ok(f64_to_json_number(a)?))?
}
// 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))
})?,
Type::TIMESTAMP_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveDateTime| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::DATE_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveDate| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::TIME_ARRAY => get_array(row, column, column_i, |a: chrono::NaiveTime| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::TIMESTAMPTZ_ARRAY => get_array(row, column, column_i, |a: chrono::DateTime<Utc>| {
Ok(JSONValue::String(a.to_string()))
})?,
Type::BYTEA_ARRAY => get_array(row, column, column_i, |a: Vec<u8>| {
Ok(JSONValue::String(format!("\\x{}", hex::encode(a))))
})?,
_ => get_basic(row, column, column_i, |a: String| Ok(JSONValue::String(a)))?,
})
}
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))
}
// some type-aliases I use in my project
pub type JSONValue = serde_json::Value;
pub type RowData = Map<String, JSONValue>;
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)
}
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!(
"conversion issue for value at column_name `{}` with type {:?}",
column.name(),
column.type_()
)
})?;
raw_val.map_or(Ok(JSONValue::Null), val_to_json_val)
}
struct IntervalStr(String);
impl<'a> FromSql<'a> for IntervalStr {
fn from_sql(
_: &Type,
mut raw: &'a [u8],
) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
let microseconds = raw.get_i64();
let days = raw.get_i32();
let months = raw.get_i32();
Ok(IntervalStr(format!(
"{:?} months {:?} days {:?} ms",
months, days, microseconds
)))
}
fn accepts(ty: &Type) -> bool {
matches!(ty, &Type::INTERVAL)
}
}
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!(
"conversion issue for array at 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,
})
}
// you can remove this section if not using TS_VECTOR (or other types requiring an intermediary `FromSQL` struct)
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
}
}