use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use tokio::sync::{Mutex, OwnedMutexGuard};
use tokio_postgres::types::ToSql;
use tokio_postgres::{NoTls, Statement};
use oxisql_core::{
ColumnInfo, Connection, ForeignKeyInfo, IndexInfo, OxiSqlError, PreparedStatement, Row,
TableInfo, TableType, ToSqlValue, Transaction,
};
use crate::builder::TlsMode;
use crate::copy;
use crate::error::PgError;
use crate::notify::{self, PgNotification};
use crate::pipeline::PgPipeline;
use crate::prepared::{PgPrepared, StmtCache};
use crate::tls::PgTls;
use crate::types::{pg_row_to_row, value_to_param, OwnedParam};
#[derive(Debug, Clone)]
pub struct ColumnDescription {
pub name: String,
pub type_name: String,
pub nullable: bool,
}
#[derive(Clone)]
pub struct PgConnection {
inner: Arc<Mutex<tokio_postgres::Client>>,
stmt_cache: StmtCache,
reconnect_uri: Option<String>,
reconnect_tls: TlsMode,
notif_tx: Option<tokio::sync::broadcast::Sender<PgNotification>>,
server_version: Option<String>,
}
impl PgConnection {
pub fn from_client(client: tokio_postgres::Client) -> Self {
Self {
inner: Arc::new(Mutex::new(client)),
stmt_cache: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
reconnect_uri: None,
reconnect_tls: TlsMode::Disabled,
notif_tx: None,
server_version: None,
}
}
pub async fn connect(conn_str: &str, tls: TlsMode) -> Result<Self, PgError> {
let tls_clone = tls.clone();
let (client, notif_tx, server_version) = match tls {
TlsMode::Disabled => {
let (client, connection) = tokio_postgres::connect(conn_str, NoTls).await?;
let server_version = connection.parameter("server_version").map(str::to_string);
let tx = notify::spawn_connection_driver(connection);
(client, tx, server_version)
}
TlsMode::Rustls(cfg) => {
let tls_connector = PgTls::new(cfg);
let (client, connection) = tokio_postgres::connect(conn_str, tls_connector).await?;
let server_version = connection.parameter("server_version").map(str::to_string);
let tx = notify::spawn_connection_driver(connection);
(client, tx, server_version)
}
};
Ok(Self {
inner: Arc::new(Mutex::new(client)),
stmt_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
reconnect_uri: Some(conn_str.to_string()),
reconnect_tls: tls_clone,
notif_tx: Some(notif_tx),
server_version,
})
}
pub async fn connect_with_timeout(
conn_str: &str,
tls: TlsMode,
timeout: Duration,
) -> Result<Self, PgError> {
tokio::time::timeout(timeout, Self::connect(conn_str, tls))
.await
.map_err(|_| {
PgError::Timeout(format!(
"connection timed out after {}ms",
timeout.as_millis()
))
})?
}
pub async fn connect_skip_verify(conn_str: &str) -> Result<Self, PgError> {
Self::connect(conn_str, TlsMode::skip_verify()?).await
}
pub async fn connect_with_ca(conn_str: &str, ca_pem: Vec<u8>) -> Result<Self, PgError> {
Self::connect(conn_str, TlsMode::with_ca_pem(ca_pem)?).await
}
#[must_use]
pub fn server_version(&self) -> Option<&str> {
self.server_version.as_deref()
}
}
impl PgConnection {
pub fn is_read_only_query(sql: &str) -> bool {
match oxisql_parse::parse_one(sql) {
Ok(stmt) => oxisql_parse::is_read_only(&stmt),
Err(_) => false,
}
}
pub fn normalize_query(sql: &str) -> String {
oxisql_parse::normalize(sql)
}
}
impl PgConnection {
pub async fn copy_in_text(
&self,
table: &str,
columns: &[&str],
rows: impl Iterator<Item = Vec<String>>,
) -> Result<u64, PgError> {
copy::copy_in_text(&self.inner, table, columns, rows).await
}
pub async fn copy_out_text(
&self,
table: &str,
columns: &[&str],
) -> Result<Vec<Vec<String>>, PgError> {
copy::copy_out_text(&self.inner, table, columns).await
}
pub async fn notify(&self, channel: &str, payload: &str) -> Result<(), PgError> {
notify::notify(&self.inner, channel, payload).await
}
pub async fn listen(&self, channel: &str) -> Result<notify::NotificationStream, PgError> {
notify::listen(&self.inner, &self.notif_tx, channel).await
}
pub fn pipeline(&self) -> PgPipeline {
PgPipeline::new(Arc::clone(&self.inner))
}
pub async fn cancel_token(&self) -> PostgresCancelToken {
let client = self.inner.lock().await;
PostgresCancelToken {
inner: client.cancel_token(),
}
}
pub async fn reconnect(&self) -> Result<PgConnection, PgError> {
match &self.reconnect_uri {
None => Err(PgError::Notify(
"cannot reconnect: connection was created via from_client (no URI stored)"
.to_string(),
)),
Some(uri) => PgConnection::connect(uri, self.reconnect_tls.clone()).await,
}
}
pub async fn query_binary(
&self,
sql: &str,
params: &[&dyn oxisql_core::ToSqlValue],
) -> Result<Vec<Row>, PgError> {
use tokio_postgres::types::Type;
let client = self.inner.lock().await;
let owned = build_params(params);
let typed_params: Vec<(&(dyn ToSql + Sync), Type)> = owned
.iter()
.map(|p| {
let ty = match p {
OwnedParam::Bool(_) => Type::BOOL,
OwnedParam::I64(_) => Type::INT8,
OwnedParam::F64(_) => Type::FLOAT8,
OwnedParam::Text(_) => Type::TEXT,
OwnedParam::Blob(_) => Type::BYTEA,
OwnedParam::Null => Type::UNKNOWN,
};
let dyn_ref: &(dyn ToSql + Sync) = p;
(dyn_ref, ty)
})
.collect();
let pg_rows = client
.query_typed(sql, typed_params.as_slice())
.await
.map_err(PgError::Postgres)?;
pg_rows
.into_iter()
.map(crate::types::pg_row_to_row)
.collect()
}
pub async fn describe(&self, sql: &str) -> Result<Vec<ColumnDescription>, PgError> {
let client = self.inner.lock().await;
let stmt = client.prepare(sql).await.map_err(PgError::Postgres)?;
Ok(stmt
.columns()
.iter()
.map(|col| ColumnDescription {
name: col.name().to_string(),
type_name: col.type_().name().to_string(),
nullable: true,
})
.collect())
}
}
fn build_params(params: &[&dyn ToSqlValue]) -> Vec<OwnedParam> {
params
.iter()
.map(|p| value_to_param(&p.to_value()))
.collect()
}
fn owned_refs(owned: &[OwnedParam]) -> Vec<&(dyn ToSql + Sync)> {
owned.iter().map(|p| p as &(dyn ToSql + Sync)).collect()
}
#[async_trait]
impl Connection for PgConnection {
async fn execute(&self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<u64, OxiSqlError> {
let client = self.inner.lock().await;
let owned = build_params(params);
let refs = owned_refs(&owned);
client
.execute(sql, refs.as_slice())
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn query(&self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<Vec<Row>, OxiSqlError> {
let client = self.inner.lock().await;
let owned = build_params(params);
let refs = owned_refs(&owned);
let pg_rows = client
.query(sql, refs.as_slice())
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))?;
pg_rows
.into_iter()
.map(|r| pg_row_to_row(r).map_err(OxiSqlError::from))
.collect()
}
async fn transaction(&self) -> Result<Box<dyn Transaction + '_>, OxiSqlError> {
let guard = Arc::clone(&self.inner).lock_owned().await;
guard
.batch_execute("BEGIN")
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))?;
Ok(Box::new(PgTransaction {
guard: Some(guard),
done: false,
}))
}
async fn execute_batch(&self, sql: &str) -> Result<u64, OxiSqlError> {
let client = self.inner.lock().await;
client
.batch_execute(sql)
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))?;
Ok(0)
}
async fn ping(&self) -> Result<(), OxiSqlError> {
let client = self.inner.lock().await;
client
.simple_query("")
.await
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
Ok(())
}
async fn tables(&self) -> Result<Vec<TableInfo>, OxiSqlError> {
let rows = self
.query(
"SELECT table_name, table_schema, table_type \
FROM information_schema.tables \
WHERE table_schema NOT IN ('pg_catalog', 'information_schema') \
ORDER BY table_schema, table_name",
&[],
)
.await?;
rows.into_iter()
.map(|row| {
let name = row
.try_get::<String>("table_name")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let schema: Option<String> = row.try_get::<String>("table_schema").ok();
let type_str: String = row
.try_get::<String>("table_type")
.unwrap_or_else(|_| "TABLE".to_string());
Ok(TableInfo {
name,
schema,
table_type: TableType::from(type_str.as_str()),
})
})
.collect()
}
async fn columns(&self, table: &str) -> Result<Vec<ColumnInfo>, OxiSqlError> {
let rows = self
.query(
"SELECT column_name, ordinal_position, data_type, is_nullable, column_default, \
character_maximum_length, numeric_precision, numeric_scale \
FROM information_schema.columns \
WHERE table_name = $1 \
ORDER BY ordinal_position",
&[&table],
)
.await?;
rows.into_iter()
.map(|row| {
let name = row
.try_get::<String>("column_name")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let ordinal: i64 = row
.try_get::<i64>("ordinal_position")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let data_type = row
.try_get::<String>("data_type")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let nullable_str = row
.try_get::<String>("is_nullable")
.unwrap_or_else(|_| "YES".to_string());
let default = row.try_get::<String>("column_default").ok();
let max_len: Option<i64> = row.try_get::<i64>("character_maximum_length").ok();
let num_prec: Option<i64> = row.try_get::<i64>("numeric_precision").ok();
let num_scale: Option<i64> = row.try_get::<i64>("numeric_scale").ok();
Ok(ColumnInfo {
name,
ordinal_position: u32::try_from(ordinal).unwrap_or(0),
data_type,
nullable: !nullable_str.eq_ignore_ascii_case("NO"),
default,
max_length: max_len.and_then(|n| u64::try_from(n).ok()),
numeric_precision: num_prec.and_then(|n| u32::try_from(n).ok()),
numeric_scale: num_scale.and_then(|n| u32::try_from(n).ok()),
})
})
.collect()
}
async fn indexes(&self, table: &str) -> Result<Vec<IndexInfo>, OxiSqlError> {
let rows = self
.query(
"SELECT indexname, indexdef \
FROM pg_indexes \
WHERE tablename = $1 \
ORDER BY indexname",
&[&table],
)
.await?;
rows.into_iter()
.map(|row| {
let name = row
.try_get::<String>("indexname")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let def = row.try_get::<String>("indexdef").unwrap_or_default();
let unique = def.contains("UNIQUE");
let primary = name.ends_with("_pkey");
let columns = if let (Some(start), Some(end)) = (def.rfind('('), def.rfind(')')) {
def[start + 1..end]
.split(',')
.map(|s| s.trim().to_string())
.collect()
} else {
Vec::new()
};
Ok(IndexInfo {
name,
columns,
unique,
primary,
})
})
.collect()
}
async fn foreign_keys(&self, table: &str) -> Result<Vec<ForeignKeyInfo>, OxiSqlError> {
let rows = self
.query(
"SELECT tc.constraint_name, kcu.column_name, \
ccu.table_name AS foreign_table, \
ccu.column_name AS foreign_column \
FROM information_schema.table_constraints tc \
JOIN information_schema.key_column_usage kcu \
ON tc.constraint_name = kcu.constraint_name \
AND tc.table_schema = kcu.table_schema \
JOIN information_schema.constraint_column_usage ccu \
ON ccu.constraint_name = tc.constraint_name \
AND ccu.table_schema = tc.table_schema \
WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name = $1",
&[&table],
)
.await?;
rows.into_iter()
.map(|row| {
let constraint_name = row
.try_get::<String>("constraint_name")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let column = row
.try_get::<String>("column_name")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let foreign_table = row
.try_get::<String>("foreign_table")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
let foreign_column = row
.try_get::<String>("foreign_column")
.map_err(|e| OxiSqlError::Other(e.to_string()))?;
Ok(ForeignKeyInfo {
constraint_name,
column,
foreign_table,
foreign_column,
..Default::default()
})
})
.collect()
}
fn query_stream<'a>(
&'a self,
sql: &'a str,
params: &'a [&'a dyn oxisql_core::ToSqlValue],
) -> std::pin::Pin<
Box<
dyn futures::Stream<Item = Result<oxisql_core::Row, oxisql_core::OxiSqlError>>
+ Send
+ 'a,
>,
> {
use futures::StreamExt;
let fut = self.query(sql, params);
Box::pin(futures::stream::once(fut).flat_map(|result| match result {
Ok(rows) => futures::stream::iter(rows.into_iter().map(Ok)).left_stream(),
Err(e) => futures::stream::once(async move { Err(e) }).right_stream(),
}))
}
async fn prepare(&self, sql: &str) -> Result<Box<dyn PreparedStatement + '_>, OxiSqlError> {
{
let cache = self
.stmt_cache
.lock()
.map_err(|_| OxiSqlError::Other("statement cache lock poisoned".into()))?;
if let Some(stmt) = cache.get(sql) {
return Ok(Box::new(PgPrepared {
client: Arc::clone(&self.inner),
stmt: stmt.clone(),
sql_text: sql.to_string(),
}));
}
}
let stmt: Statement = {
let client = self.inner.lock().await;
client
.prepare(sql)
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))?
};
{
let mut cache = self
.stmt_cache
.lock()
.map_err(|_| OxiSqlError::Other("statement cache lock poisoned".into()))?;
cache.insert(sql.to_string(), stmt.clone());
}
Ok(Box::new(PgPrepared {
client: Arc::clone(&self.inner),
stmt,
sql_text: sql.to_string(),
}))
}
}
pub struct PgTransaction {
guard: Option<OwnedMutexGuard<tokio_postgres::Client>>,
done: bool,
}
impl Drop for PgTransaction {
fn drop(&mut self) {
if let Some(guard) = self.guard.take() {
if !self.done {
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
handle.spawn(async move {
let _ = guard.batch_execute("ROLLBACK").await;
});
}
Err(_) => {
drop(guard);
}
}
}
}
}
}
fn guard_ref(
guard: &mut Option<OwnedMutexGuard<tokio_postgres::Client>>,
) -> Result<&mut OwnedMutexGuard<tokio_postgres::Client>, OxiSqlError> {
guard.as_mut().ok_or_else(|| {
OxiSqlError::Execution("transaction already consumed (committed or rolled back)".into())
})
}
#[async_trait]
impl Transaction for PgTransaction {
async fn execute(&mut self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<u64, OxiSqlError> {
let owned = build_params(params);
let refs = owned_refs(&owned);
guard_ref(&mut self.guard)?
.execute(sql, refs.as_slice())
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn query(
&mut self,
sql: &str,
params: &[&dyn ToSqlValue],
) -> Result<Vec<Row>, OxiSqlError> {
let owned = build_params(params);
let refs = owned_refs(&owned);
let pg_rows = guard_ref(&mut self.guard)?
.query(sql, refs.as_slice())
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))?;
pg_rows
.into_iter()
.map(|r| pg_row_to_row(r).map_err(OxiSqlError::from))
.collect()
}
async fn commit(mut self: Box<Self>) -> Result<(), OxiSqlError> {
self.done = true;
let guard = self
.guard
.take()
.ok_or_else(|| OxiSqlError::Execution("transaction already consumed".into()))?;
guard
.batch_execute("COMMIT")
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn rollback(mut self: Box<Self>) -> Result<(), OxiSqlError> {
self.done = true;
let guard = self
.guard
.take()
.ok_or_else(|| OxiSqlError::Execution("transaction already consumed".into()))?;
guard
.batch_execute("ROLLBACK")
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
validate_savepoint_name(name)?;
guard_ref(&mut self.guard)?
.batch_execute(&format!("SAVEPOINT {name}"))
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn release_savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
validate_savepoint_name(name)?;
guard_ref(&mut self.guard)?
.batch_execute(&format!("RELEASE SAVEPOINT {name}"))
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
async fn rollback_to_savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
validate_savepoint_name(name)?;
guard_ref(&mut self.guard)?
.batch_execute(&format!("ROLLBACK TO SAVEPOINT {name}"))
.await
.map_err(|e| OxiSqlError::Execution(e.to_string()))
}
}
impl PgTransaction {
pub async fn savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
if !is_valid_savepoint_name(name) {
return Err(PgError::Notify(format!(
"invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
)));
}
let guard = self
.guard
.as_ref()
.ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
guard
.batch_execute(&format!("SAVEPOINT {name}"))
.await
.map_err(PgError::Postgres)
}
pub async fn rollback_to_savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
if !is_valid_savepoint_name(name) {
return Err(PgError::Notify(format!(
"invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
)));
}
let guard = self
.guard
.as_ref()
.ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
guard
.batch_execute(&format!("ROLLBACK TO SAVEPOINT {name}"))
.await
.map_err(PgError::Postgres)
}
pub async fn release_savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
if !is_valid_savepoint_name(name) {
return Err(PgError::Notify(format!(
"invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
)));
}
let guard = self
.guard
.as_ref()
.ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
guard
.batch_execute(&format!("RELEASE SAVEPOINT {name}"))
.await
.map_err(PgError::Postgres)
}
}
fn is_valid_savepoint_name(name: &str) -> bool {
!name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}
fn validate_savepoint_name(name: &str) -> Result<(), OxiSqlError> {
if name.is_empty() {
return Err(OxiSqlError::Other(
"savepoint name must not be empty".into(),
));
}
if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
return Err(OxiSqlError::Other(format!(
"invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
)));
}
Ok(())
}
pub struct PostgresCancelToken {
inner: tokio_postgres::CancelToken,
}
impl PostgresCancelToken {
pub async fn cancel_query<T>(&self, tls: T) -> Result<(), PgError>
where
T: tokio_postgres::tls::MakeTlsConnect<tokio_postgres::Socket> + Clone,
{
self.inner
.cancel_query(tls)
.await
.map_err(|e| PgError::Connection(e.to_string()))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PgConnParts {
pub host: String,
pub port: u16,
pub dbname: Option<String>,
pub user: Option<String>,
}
pub fn parse_pg_conn_str(conn_str: &str) -> Result<PgConnParts, OxiSqlError> {
let trimmed = conn_str.trim();
if trimmed.starts_with("postgres://") || trimmed.starts_with("postgresql://") {
parse_pg_uri(trimmed)
} else {
parse_pg_kv(trimmed)
}
}
fn parse_pg_kv(conn_str: &str) -> Result<PgConnParts, OxiSqlError> {
let mut parts = PgConnParts {
host: "localhost".to_string(),
port: 5432,
dbname: None,
user: None,
};
for token in conn_str.split_whitespace() {
if let Some((k, v)) = token.split_once('=') {
match k {
"host" | "hostaddr" => parts.host = v.to_string(),
"port" if !v.is_empty() => {
parts.port = v.parse::<u16>().map_err(|_| {
OxiSqlError::Other(format!(
"invalid port value in connection string: {v:?}"
))
})?;
}
"dbname" => parts.dbname = Some(v.to_string()),
"user" => parts.user = Some(v.to_string()),
_ => {}
}
}
}
Ok(parts)
}
fn parse_pg_uri(uri: &str) -> Result<PgConnParts, OxiSqlError> {
let rest = uri
.strip_prefix("postgresql://")
.or_else(|| uri.strip_prefix("postgres://"))
.ok_or_else(|| {
OxiSqlError::Other(format!("unrecognised connection string scheme: {uri:?}"))
})?;
let (rest, _query) = rest.split_once('?').unwrap_or((rest, ""));
let (user_info, host_part) = if let Some(at_pos) = rest.rfind('@') {
(Some(&rest[..at_pos]), &rest[at_pos + 1..])
} else {
(None, rest)
};
let user = user_info.and_then(|ui| {
let cred = ui.split_once(':').map(|(u, _)| u).unwrap_or(ui);
if cred.is_empty() {
None
} else {
Some(cred.to_string())
}
});
let (authority, dbname) = if let Some(slash_pos) = host_part.find('/') {
let db = &host_part[slash_pos + 1..];
(
&host_part[..slash_pos],
if db.is_empty() {
None
} else {
Some(db.to_string())
},
)
} else {
(host_part, None)
};
let (host, port) = if authority.starts_with('[') {
if let Some(bracket_end) = authority.find(']') {
let h = &authority[1..bracket_end];
let port_str = authority[bracket_end + 1..].trim_start_matches(':');
let p = if port_str.is_empty() {
5432
} else {
port_str
.parse::<u16>()
.map_err(|_| OxiSqlError::Other(format!("invalid port in URI: {port_str:?}")))?
};
(h.to_string(), p)
} else {
(authority.to_string(), 5432)
}
} else if let Some(colon_pos) = authority.rfind(':') {
let h = &authority[..colon_pos];
let port_str = &authority[colon_pos + 1..];
let p = if port_str.is_empty() {
5432
} else {
port_str
.parse::<u16>()
.map_err(|_| OxiSqlError::Other(format!("invalid port in URI: {port_str:?}")))?
};
(h.to_string(), p)
} else {
(authority.to_string(), 5432)
};
let host = if host.is_empty() {
"localhost".to_string()
} else {
host
};
Ok(PgConnParts {
host,
port,
dbname,
user,
})
}
#[cfg(test)]
mod cancel_token_tests {
use crate::PostgresCancelToken;
fn _assert_send_sync<T: Send + Sync>() {}
#[test]
fn cancel_token_is_send_sync() {
_assert_send_sync::<PostgresCancelToken>();
}
}