use anyhow::{Context, Result};
use deadpool_postgres::Pool;
use fraiseql_db::postgres::PostgresTlsConfig;
#[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,
pub language: String,
pub source: String,
}
#[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>,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChangeLogRlsStatus {
pub rls_enabled: bool,
pub can_bypass: bool,
pub is_owner: bool,
pub role_name: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublicGrant {
pub relname: String,
pub privileges: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChangeLogActorStats {
pub null_rows: i64,
pub unknown_values: Vec<String>,
pub constraint_installed: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CaptureFnSecurity {
pub security_definer: bool,
pub search_path: Option<String>,
}
#[derive(Debug)]
pub struct SecurityInvokerAudit {
pub rls_in_use: bool,
pub views_without_invoker: Vec<String>,
}
pub struct PgCatalog {
pool: Pool,
}
impl PgCatalog {
#[must_use]
pub(crate) const fn from_pool(pool: Pool) -> Self {
Self { pool }
}
pub async fn connect(db_url: &str, tls: &PostgresTlsConfig) -> Result<Self> {
let pool = crate::connection::postgres_pool(db_url, "catalogue introspection", tls).await?;
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, \
l.lanname, COALESCE(p.prosrc, '') AS prosrc, \
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 \
JOIN pg_language l ON l.oid = p.prolang ";
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"),
language: row.get("lanname"),
source: row.get("prosrc"),
});
}
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())
}
pub async fn change_log_rls_status(&self) -> Result<Option<ChangeLogRlsStatus>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let rows = client
.query(
"SELECT c.relrowsecurity AS rls_enabled, \
(r.rolbypassrls OR r.rolsuper) AS can_bypass, \
(pg_get_userbyid(c.relowner) = current_user) AS is_owner, \
current_user::text AS role_name \
FROM pg_class c \
JOIN pg_namespace n ON n.oid = c.relnamespace \
JOIN pg_roles r ON r.rolname = current_user \
WHERE n.nspname = 'core' AND c.relname = 'tb_entity_change_log'",
&[],
)
.await
.context("failed to query change-log RLS status from pg_class/pg_roles")?;
Ok(rows.first().map(|row| ChangeLogRlsStatus {
rls_enabled: row.get("rls_enabled"),
can_bypass: row.get("can_bypass"),
is_owner: row.get("is_owner"),
role_name: row.get("role_name"),
}))
}
pub async fn security_invoker_audit(&self, views: &[String]) -> Result<SecurityInvokerAudit> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let rls_in_use: bool = client
.query_one("SELECT EXISTS (SELECT 1 FROM pg_policies) AS in_use", &[])
.await
.context("failed to query pg_policies for RLS usage")?
.get("in_use");
let rows = client
.query(
"SELECT c.relname \
FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace \
WHERE c.relkind = 'v' AND c.relname = ANY($1) \
AND NOT EXISTS ( \
SELECT 1 FROM pg_options_to_table(c.reloptions) o \
WHERE o.option_name = 'security_invoker' \
AND lower(o.option_value) IN ('true', 'on', '1'))",
&[&views],
)
.await
.context("failed to audit view security_invoker options from pg_class")?;
Ok(SecurityInvokerAudit {
rls_in_use,
views_without_invoker: rows.iter().map(|r| r.get::<_, String>("relname")).collect(),
})
}
pub async fn change_log_public_grants(&self) -> Result<Vec<PublicGrant>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let rows = client
.query(
"SELECT c.relname, \
COALESCE( \
array_agg(DISTINCT a.privilege_type) FILTER (WHERE a.grantee = 0), \
ARRAY[]::text[]) AS public_privs \
FROM pg_class c \
JOIN pg_namespace n ON n.oid = c.relnamespace \
LEFT JOIN LATERAL aclexplode(c.relacl) a ON true \
WHERE n.nspname = 'core' \
AND c.relname IN ('tb_entity_change_log', 'v_entity_change_log', \
'v_entity_change_log_debezium') \
GROUP BY c.relname \
ORDER BY c.relname",
&[],
)
.await
.context("failed to query change-log PUBLIC grants from pg_class.relacl")?;
Ok(rows
.iter()
.map(|row| PublicGrant {
relname: row.get("relname"),
privileges: row.get("public_privs"),
})
.collect())
}
pub async fn change_log_actor_stats(
&self,
valid_tokens: &[String],
) -> Result<Option<ChangeLogActorStats>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let table_exists: bool = client
.query_one(
"SELECT EXISTS (SELECT 1 FROM pg_class c \
JOIN pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = 'core' AND c.relname = 'tb_entity_change_log') AS present",
&[],
)
.await
.context("failed to check for core.tb_entity_change_log in pg_class")?
.get("present");
if !table_exists {
return Ok(None);
}
let row = client
.query_one(
"SELECT \
(SELECT count(*) FROM core.tb_entity_change_log \
WHERE actor_type IS NULL) AS null_rows, \
(SELECT COALESCE(array_agg(DISTINCT actor_type ORDER BY actor_type), \
ARRAY[]::text[]) \
FROM core.tb_entity_change_log \
WHERE actor_type IS NOT NULL AND actor_type <> ALL($1)) AS unknown_values, \
EXISTS (SELECT 1 FROM pg_constraint con \
JOIN pg_class c ON c.oid = con.conrelid \
JOIN pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = 'core' AND c.relname = 'tb_entity_change_log' \
AND con.conname = 'chk_entity_change_log_actor_type') \
AS constraint_installed",
&[&valid_tokens],
)
.await
.context("failed to query change-log actor attribution stats")?;
Ok(Some(ChangeLogActorStats {
null_rows: row.get("null_rows"),
unknown_values: row.get("unknown_values"),
constraint_installed: row.get("constraint_installed"),
}))
}
pub async fn capture_fn_security(&self) -> Result<Option<CaptureFnSecurity>> {
let client = self.pool.get().await.context("failed to acquire DB connection")?;
let rows = client
.query(
"SELECT p.prosecdef AS security_definer, \
(SELECT cfg FROM unnest(p.proconfig) AS cfg \
WHERE cfg LIKE 'search_path=%' LIMIT 1) AS search_path_cfg \
FROM pg_proc p \
JOIN pg_namespace n ON n.oid = p.pronamespace \
WHERE n.nspname = 'core' AND p.proname = 'fn_entity_change_log_capture' \
LIMIT 1",
&[],
)
.await
.context("failed to query capture-function security from pg_proc")?;
Ok(rows.first().map(|row| {
let raw: Option<String> = row.get("search_path_cfg");
CaptureFnSecurity {
security_definer: row.get("security_definer"),
search_path: raw
.map(|s| s.strip_prefix("search_path=").unwrap_or(&s).to_string()),
}
}))
}
}
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()),
}
}