use anyhow::{Context, Result};
use deadpool_postgres::{Config, ManagerConfig, Pool, RecyclingMethod, Runtime};
use tokio_postgres::NoTls;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutColumn {
pub name: String,
pub type_name: String,
pub is_enum: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PgFunction {
pub schema: String,
pub name: String,
pub in_types: Vec<String>,
pub in_names: Vec<Option<String>>,
pub out_columns: Vec<OutColumn>,
pub returns_set: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BodyError {
pub caller: String,
pub lineno: Option<i32>,
pub message: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LiveColumn {
pub name: String,
pub udt_name: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PlpgsqlCheckOutcome {
Unavailable,
Ran {
errors: Vec<BodyError>,
},
}
pub struct PgCatalog {
pool: Pool,
}
impl PgCatalog {
pub fn connect(db_url: &str) -> Result<Self> {
if !db_url.starts_with("postgres") {
anyhow::bail!(
"--against-db requires a PostgreSQL connection URL (postgres://…); got: {db_url}"
);
}
let mut cfg = Config::new();
cfg.url = Some(db_url.to_string());
cfg.manager = Some(ManagerConfig {
recycling_method: RecyclingMethod::Fast,
});
cfg.pool = Some(deadpool_postgres::PoolConfig::new(2));
let pool = cfg
.create_pool(Some(Runtime::Tokio1), NoTls)
.context("failed to create PostgreSQL connection pool for --against-db")?;
Ok(Self { pool })
}
pub async fn resolve_functions(&self, sql_source: &str) -> Result<Vec<PgFunction>> {
const BASE: &str = "\
SELECT p.oid, n.nspname, p.proname, p.proretset, \
COALESCE((SELECT array_agg(format_type(t, NULL) ORDER BY ord) \
FROM unnest(COALESCE(p.proallargtypes, p.proargtypes::oid[])) \
WITH ORDINALITY AS a(t, ord) \
LEFT JOIN unnest(p.proargmodes) WITH ORDINALITY AS m(mode, mord) ON mord = ord \
WHERE COALESCE(m.mode, 'i') IN ('i','b','v')), ARRAY[]::text[]) AS in_types, \
COALESCE((SELECT array_agg(COALESCE(nm, '') ORDER BY ord) \
FROM unnest(COALESCE(p.proallargtypes, p.proargtypes::oid[])) \
WITH ORDINALITY AS a(t, ord) \
LEFT JOIN unnest(p.proargmodes) WITH ORDINALITY AS m(mode, mord) ON mord = ord \
LEFT JOIN unnest(p.proargnames) WITH ORDINALITY AS nmt(nm, nord) ON nord = ord \
WHERE COALESCE(m.mode, 'i') IN ('i','b','v')), ARRAY[]::text[]) AS in_names \
FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace ";
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let (schema, name) = split_qualified(sql_source);
let rows = if let Some(schema) = schema {
client
.query(
&format!("{BASE} WHERE n.nspname = $1 AND p.proname = $2"),
&[&schema, &name],
)
.await
} else {
client
.query(
&format!(
"{BASE} WHERE p.proname = $1 AND n.nspname = ANY(current_schemas(false))"
),
&[&name],
)
.await
}
.context("failed to query pg_proc for function resolution")?;
let mut functions = Vec::with_capacity(rows.len());
for row in rows {
let oid: u32 = row.get("oid");
let in_names: Vec<String> = row.get("in_names");
functions.push(PgFunction {
schema: row.get("nspname"),
name: row.get("proname"),
in_types: row.get("in_types"),
in_names: in_names
.into_iter()
.map(|n| if n.is_empty() { None } else { Some(n) })
.collect(),
out_columns: self.out_columns(oid).await?,
returns_set: row.get("proretset"),
});
}
Ok(functions)
}
async fn out_columns(&self, oid: u32) -> Result<Vec<OutColumn>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let table_rows = client
.query(
"SELECT COALESCE(nm, '') AS name, format_type(t, NULL) AS type_name, \
(tt.typtype = 'e') AS is_enum \
FROM pg_proc p, \
unnest(p.proallargtypes) WITH ORDINALITY AS a(t, ord) \
JOIN unnest(p.proargmodes) WITH ORDINALITY AS m(mode, mord) ON mord = ord \
LEFT JOIN unnest(p.proargnames) WITH ORDINALITY AS nmt(nm, nord) ON nord = ord \
JOIN pg_type tt ON tt.oid = t \
WHERE p.oid = $1 AND m.mode IN ('o','b','t') ORDER BY ord",
&[&oid],
)
.await
.context("failed to query OUT/TABLE result columns")?;
if !table_rows.is_empty() {
return Ok(table_rows.iter().map(row_to_out_column).collect());
}
let comp_rows = client
.query(
"SELECT att.attname AS name, \
format_type(att.atttypid, att.atttypmod) AS type_name, \
(at.typtype = 'e') AS is_enum \
FROM pg_proc p \
JOIN pg_type rt ON rt.oid = p.prorettype \
JOIN pg_attribute att ON att.attrelid = rt.typrelid \
AND att.attnum > 0 AND NOT att.attisdropped \
JOIN pg_type at ON at.oid = att.atttypid \
WHERE p.oid = $1 ORDER BY att.attnum",
&[&oid],
)
.await
.context("failed to query composite result columns")?;
Ok(comp_rows.iter().map(row_to_out_column).collect())
}
pub async fn plpgsql_check_available(&self) -> Result<bool> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let row = client
.query_one(
"SELECT EXISTS(SELECT 1 FROM pg_available_extensions WHERE name = 'plpgsql_check')",
&[],
)
.await
.context("failed to probe for the plpgsql_check extension")?;
Ok(row.get(0))
}
pub async fn plpgsql_check_unresolved_calls(
&self,
schemas: &[String],
) -> Result<PlpgsqlCheckOutcome> {
if !self.plpgsql_check_available().await? {
return Ok(PlpgsqlCheckOutcome::Unavailable);
}
let client = self.pool.get().await.context("failed to acquire DB connection")?;
client
.batch_execute("CREATE EXTENSION IF NOT EXISTS plpgsql_check")
.await
.context("failed to CREATE EXTENSION plpgsql_check")?;
let rows = client
.query(
"SELECT (p.oid::regprocedure)::text AS caller, t.lineno, t.message \
FROM pg_proc p \
JOIN pg_namespace n ON n.oid = p.pronamespace \
JOIN pg_language l ON l.oid = p.prolang AND l.lanname = 'plpgsql' \
CROSS JOIN LATERAL plpgsql_check_function_tb(p.oid, fatal_errors := false) t \
WHERE n.nspname = ANY($1) \
AND p.prorettype <> 'trigger'::regtype \
AND t.level = 'error' \
AND t.message ~* 'function .* does not exist' \
ORDER BY 1, 2",
&[&schemas],
)
.await
.context("failed to run the plpgsql_check body-resolution pass")?;
let errors = rows
.iter()
.map(|row| BodyError {
caller: row.get("caller"),
lineno: row.get("lineno"),
message: row.get("message"),
})
.collect();
Ok(PlpgsqlCheckOutcome::Ran { errors })
}
pub async fn table_columns(&self, schema: &str, table: &str) -> Result<Vec<LiveColumn>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let rows = client
.query(
"SELECT column_name, udt_name FROM information_schema.columns \
WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position",
&[&schema, &table],
)
.await
.context("failed to query information_schema.columns")?;
Ok(rows
.iter()
.map(|row| LiveColumn {
name: row.get("column_name"),
udt_name: row.get("udt_name"),
})
.collect())
}
}
fn row_to_out_column(row: &tokio_postgres::Row) -> OutColumn {
OutColumn {
name: row.get("name"),
type_name: row.get("type_name"),
is_enum: row.get("is_enum"),
}
}
fn split_qualified(sql_source: &str) -> (Option<String>, String) {
match sql_source.split_once('.') {
Some((schema, name)) => (Some(schema.to_string()), name.to_string()),
None => (None, sql_source.to_string()),
}
}