use std::error::Error as StdError;
use std::sync::{Arc, Mutex};
use tokio::runtime::Handle;
use tokio_postgres::Row as PgRow;
use tokio_postgres::error::DbError;
use tokio_postgres::types::Type;
use super::queries::query_for;
use super::{CatalogError, CatalogQuerier, CatalogQuery, PgVersion, Row, Value};
pub struct PgCatalogQuerier {
client: Arc<tokio_postgres::Client>,
runtime: Handle,
version: Mutex<Option<PgVersion>>,
}
impl std::fmt::Debug for PgCatalogQuerier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PgCatalogQuerier").finish_non_exhaustive()
}
}
#[derive(Debug, thiserror::Error)]
#[error("PgCatalogQuerier must be constructed from within a Tokio runtime: {0}")]
pub struct NoRuntimeError(#[from] tokio::runtime::TryCurrentError);
impl PgCatalogQuerier {
pub fn new(client: tokio_postgres::Client) -> Result<Self, NoRuntimeError> {
Ok(Self {
client: Arc::new(client),
runtime: Handle::try_current()?,
version: Mutex::new(None),
})
}
pub fn from_arc(client: Arc<tokio_postgres::Client>) -> Result<Self, NoRuntimeError> {
Ok(Self {
client,
runtime: Handle::try_current()?,
version: Mutex::new(None),
})
}
pub fn version(&self) -> Result<PgVersion, CatalogError> {
let cached = *self.version.lock().expect("poisoned");
if let Some(v) = cached {
return Ok(v);
}
let v = PgVersion::detect(self)?;
*self.version.lock().expect("poisoned") = Some(v);
Ok(v)
}
}
impl CatalogQuerier for PgCatalogQuerier {
fn fetch(
&self,
query: CatalogQuery,
text_array_param: &[&str],
) -> Result<Vec<Row>, CatalogError> {
let sql = if matches!(query, CatalogQuery::PgVersion) {
query_for(PgVersion::Pg16, query)
} else {
query_for(self.version()?, query)
};
let client = Arc::clone(&self.client);
let owned: Vec<String> = text_array_param.iter().map(ToString::to_string).collect();
let pg_rows: Vec<PgRow> = tokio::task::block_in_place(|| {
self.runtime.block_on(async move {
if query.takes_text_array_param() {
client.query(sql, &[&owned]).await
} else {
client.query(sql, &[]).await
}
})
})
.map_err(|e| {
let message = StdError::source(&e)
.and_then(|s| s.downcast_ref::<DbError>())
.map_or_else(
|| e.to_string(),
|db_err| {
format!(
"{} (code={}, detail={})",
db_err.message(),
db_err.code().code(),
db_err.detail().unwrap_or("")
)
},
);
CatalogError::QueryFailed { query, message }
})?;
pg_rows
.into_iter()
.map(|r| pg_row_to_row(&r, query))
.collect()
}
}
fn pg_row_to_row(row: &PgRow, query: CatalogQuery) -> Result<Row, CatalogError> {
let mut out = Row::new();
for (i, col) in row.columns().iter().enumerate() {
let name = col.name();
let value = pg_value(row, i, col.type_(), name, query)?;
out.insert(name, value);
}
Ok(out)
}
#[allow(clippy::too_many_lines)] fn pg_value(
row: &PgRow,
idx: usize,
ty: &Type,
column_name: &str,
query: CatalogQuery,
) -> Result<Value, CatalogError> {
let bad = |msg: String| CatalogError::BadColumnType {
query,
column: column_name.to_string(),
message: msg,
};
macro_rules! get_opt {
($t:ty) => {
row.try_get::<_, Option<$t>>(idx)
.map_err(|e| bad(format!("decode {} as {}: {e}", ty.name(), stringify!($t))))?
};
}
match *ty {
Type::BOOL => Ok(get_opt!(bool).map_or(Value::Null, Value::Bool)),
Type::INT2 => Ok(get_opt!(i16).map_or(Value::Null, Value::SmallInt)),
Type::INT4 => Ok(get_opt!(i32).map_or(Value::Null, |v| Value::Integer(i64::from(v)))),
Type::INT8 => Ok(get_opt!(i64).map_or(Value::Null, Value::Integer)),
Type::OID => Ok(get_opt!(u32).map_or(Value::Null, |v| Value::Integer(i64::from(v)))),
Type::TEXT | Type::VARCHAR | Type::NAME | Type::BPCHAR => {
Ok(get_opt!(String).map_or(Value::Null, Value::Text))
}
Type::CHAR => {
let v = get_opt!(i8);
#[allow(clippy::cast_sign_loss)]
Ok(v.map_or(Value::Null, |b| Value::Char(b as u8 as char)))
}
Type::INT2_ARRAY => Ok(get_opt!(Vec<i16>).map_or(Value::Null, |v| {
Value::IntegerArray(v.into_iter().map(i64::from).collect())
})),
Type::INT4_ARRAY => Ok(get_opt!(Vec<i32>).map_or(Value::Null, |v| {
Value::IntegerArray(v.into_iter().map(i64::from).collect())
})),
Type::INT8_ARRAY => Ok(get_opt!(Vec<i64>).map_or(Value::Null, Value::IntegerArray)),
Type::TEXT_ARRAY | Type::NAME_ARRAY | Type::VARCHAR_ARRAY => {
Ok(get_opt!(Vec<String>).map_or(Value::Null, Value::TextArray))
}
Type::BYTEA => Ok(get_opt!(Vec<u8>).map_or(Value::Null, Value::Bytes)),
_ => {
row.try_get::<_, Option<String>>(idx)
.map(|v| v.map_or(Value::Null, Value::Text))
.map_err(|e| bad(format!("unsupported type {} ({e})", ty.name())))
}
}
}