pgroles-inspect 0.3.0

//! Query object privileges from PostgreSQL catalog tables.
//!
//! Uses `aclexplode()` to decompose ACL arrays from `pg_class`, `pg_namespace`,
//! `pg_proc`, `pg_type`, and `pg_database`. NULL ACLs (meaning system defaults)
//! are handled via `COALESCE(acl_col, acldefault(type_char, owner_oid))`.
//!
//! The privilege character mapping:
//!   r = SELECT, a = INSERT, w = UPDATE, d = DELETE, D = TRUNCATE,
//!   x = REFERENCES, t = TRIGGER, X = EXECUTE, U = USAGE, C = CREATE,
//!   c = CONNECT, T = TEMPORARY

use std::collections::{BTreeMap, BTreeSet};

use sqlx::PgPool;

use crate::WildcardGrantPattern;
use pgroles_core::manifest::{ObjectType, Privilege};
use pgroles_core::model::{GrantKey, GrantState};

/// A raw ACL row returned by our `aclexplode()` queries.
#[derive(Debug, sqlx::FromRow)]
struct AclRow {
    /// The grantee role name. NULL means PUBLIC — we skip those.
    grantee: Option<String>,
    /// The privilege type as a single character (e.g. 'r' for SELECT).
    privilege_type: String,
    /// The schema name (NULL for database-level grants).
    schema_name: Option<String>,
    /// The object name (the schema name itself for schema-level grants).
    object_name: String,
    /// The object type discriminator we embed in the query.
    obj_type: String,
}

/// Map a PostgreSQL ACL privilege character to our `Privilege` enum.
fn acl_char_to_privilege(character: &str) -> Option<Privilege> {
    match character {
        "r" | "SELECT" => Some(Privilege::Select),
        "a" | "INSERT" => Some(Privilege::Insert),
        "w" | "UPDATE" => Some(Privilege::Update),
        "d" | "DELETE" => Some(Privilege::Delete),
        "D" | "TRUNCATE" => Some(Privilege::Truncate),
        "x" | "REFERENCES" => Some(Privilege::References),
        "t" | "TRIGGER" => Some(Privilege::Trigger),
        "X" | "EXECUTE" => Some(Privilege::Execute),
        "U" | "USAGE" => Some(Privilege::Usage),
        "C" | "CREATE" => Some(Privilege::Create),
        "c" | "CONNECT" => Some(Privilege::Connect),
        "T" | "TEMPORARY" => Some(Privilege::Temporary),
        _ => None,
    }
}

/// Map our query's `obj_type` discriminator string to an `ObjectType`.
fn obj_type_str_to_object_type(obj_type: &str) -> Option<ObjectType> {
    match obj_type {
        "table" => Some(ObjectType::Table),
        "view" => Some(ObjectType::View),
        "materialized_view" => Some(ObjectType::MaterializedView),
        "sequence" => Some(ObjectType::Sequence),
        "function" => Some(ObjectType::Function),
        "schema" => Some(ObjectType::Schema),
        "database" => Some(ObjectType::Database),
        "type" => Some(ObjectType::Type),
        _ => None,
    }
}

/// Fetch all object privileges from the database for the given schemas and roles.
///
/// Queries tables/views/sequences via `pg_class`, schemas via `pg_namespace`,
/// functions via `pg_proc`, types via `pg_type`, and (optionally) databases via
/// `pg_database`.
///
/// Returns a map of `GrantKey → GrantState` ready for insertion into a `RoleGraph`.
pub async fn fetch_privileges(
    pool: &PgPool,
    managed_schemas: &[&str],
    managed_roles: &[&str],
) -> Result<BTreeMap<GrantKey, GrantState>, sqlx::Error> {
    fetch_privileges_with_wildcards(pool, managed_schemas, managed_roles, &[]).await
}

/// Fetch schema-scoped relation names grouped by object type.
pub async fn fetch_relation_inventory(
    pool: &PgPool,
    managed_schemas: &[&str],
) -> Result<BTreeMap<(ObjectType, String), Vec<String>>, sqlx::Error> {
    let rows = sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            NULL::text AS grantee,
            '' AS privilege_type,
            n.nspname AS schema_name,
            c.relname AS object_name,
            CASE c.relkind
                WHEN 'r' THEN 'table'
                WHEN 'p' THEN 'table'
                WHEN 'v' THEN 'view'
                WHEN 'm' THEN 'materialized_view'
            END AS obj_type
        FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = ANY($1)
          AND c.relkind IN ('r', 'p', 'v', 'm')
        ORDER BY n.nspname, c.relkind, c.relname
        "#,
    )
    .bind(managed_schemas)
    .fetch_all(pool)
    .await?;

    let mut inventory = BTreeMap::new();
    for row in rows {
        let Some(object_type) = obj_type_str_to_object_type(&row.obj_type) else {
            continue;
        };
        inventory
            .entry((
                object_type,
                row.schema_name
                    .expect("relation inventory rows always include schema"),
            ))
            .or_insert_with(Vec::new)
            .push(row.object_name);
    }
    Ok(inventory)
}

pub(crate) async fn fetch_privileges_with_wildcards(
    pool: &PgPool,
    managed_schemas: &[&str],
    managed_roles: &[&str],
    wildcard_grants: &[WildcardGrantPattern],
) -> Result<BTreeMap<GrantKey, GrantState>, sqlx::Error> {
    let mut grants: BTreeMap<GrantKey, GrantState> = BTreeMap::new();
    let mut inventory: BTreeMap<(ObjectType, String), BTreeSet<String>> = BTreeMap::new();

    // Run all the independent queries and collect results.
    // We use separate queries per object type rather than one giant UNION
    // because the NULL-ACL handling (acldefault) differs per type.

    let relation_rows = fetch_relation_privileges(pool, managed_schemas).await?;
    let schema_rows = fetch_schema_privileges(pool, managed_schemas).await?;
    let function_rows = fetch_function_privileges(pool, managed_schemas).await?;
    let type_rows = fetch_type_privileges(pool, managed_schemas).await?;

    let all_rows: Vec<AclRow> = relation_rows
        .into_iter()
        .chain(schema_rows)
        .chain(function_rows)
        .chain(type_rows)
        .collect();

    for row in &all_rows {
        if let Some(object_type) = obj_type_str_to_object_type(&row.obj_type)
            && !matches!(object_type, ObjectType::Schema | ObjectType::Database)
            && let Some(schema_name) = &row.schema_name
        {
            inventory
                .entry((object_type, schema_name.clone()))
                .or_default()
                .insert(row.object_name.clone());
        }
    }

    for row in all_rows {
        // Skip PUBLIC grantee (NULL)
        let grantee = match row.grantee {
            Some(ref name) => name,
            None => continue,
        };

        // Skip if the grantee isn't in the managed set
        if !managed_roles.contains(&grantee.as_str()) {
            continue;
        }

        let privilege = match acl_char_to_privilege(&row.privilege_type) {
            Some(privilege) => privilege,
            None => continue,
        };

        let object_type = match obj_type_str_to_object_type(&row.obj_type) {
            Some(object_type) => object_type,
            None => continue,
        };

        // Build the GrantKey.
        // Schema-level grants: object_type=Schema, schema=None, name=Some(schema_name)
        // Database-level grants: object_type=Database, schema=None, name=Some(db_name)
        // Other: object_type, schema=Some(schema_name), name=Some(object_name)
        let (schema, name) = match object_type {
            ObjectType::Schema => (None, Some(row.object_name.clone())),
            ObjectType::Database => (None, Some(row.object_name.clone())),
            _ => (row.schema_name.clone(), Some(row.object_name.clone())),
        };

        let key = GrantKey {
            role: grantee.clone(),
            object_type,
            schema,
            name,
        };

        let entry = grants.entry(key).or_insert_with(|| GrantState {
            privileges: BTreeSet::new(),
        });
        entry.privileges.insert(privilege);
    }

    Ok(normalize_wildcard_grants(
        grants,
        &inventory,
        wildcard_grants,
    ))
}

fn normalize_wildcard_grants(
    mut grants: BTreeMap<GrantKey, GrantState>,
    inventory: &BTreeMap<(ObjectType, String), BTreeSet<String>>,
    wildcard_grants: &[WildcardGrantPattern],
) -> BTreeMap<GrantKey, GrantState> {
    for wildcard in wildcard_grants {
        let Some(object_names) = inventory.get(&(wildcard.object_type, wildcard.schema.clone()))
        else {
            continue;
        };

        if object_names.is_empty() {
            continue;
        }

        let mut shared_privileges = all_privileges();

        for object_name in object_names {
            let key = GrantKey {
                role: wildcard.role.clone(),
                object_type: wildcard.object_type,
                schema: Some(wildcard.schema.clone()),
                name: Some(object_name.clone()),
            };

            if let Some(state) = grants.get(&key) {
                shared_privileges.retain(|privilege| state.privileges.contains(privilege));
            } else {
                shared_privileges.clear();
                break;
            }
        }

        if shared_privileges.is_empty() {
            continue;
        }

        let wildcard_key = GrantKey {
            role: wildcard.role.clone(),
            object_type: wildcard.object_type,
            schema: Some(wildcard.schema.clone()),
            name: Some("*".to_string()),
        };

        grants.insert(
            wildcard_key,
            GrantState {
                privileges: shared_privileges.clone(),
            },
        );

        for object_name in object_names {
            let key = GrantKey {
                role: wildcard.role.clone(),
                object_type: wildcard.object_type,
                schema: Some(wildcard.schema.clone()),
                name: Some(object_name.clone()),
            };

            let remove_key = match grants.get_mut(&key) {
                Some(state) => {
                    state
                        .privileges
                        .retain(|privilege| !shared_privileges.contains(privilege));
                    state.privileges.is_empty()
                }
                None => false,
            };

            if remove_key {
                grants.remove(&key);
            }
        }
    }

    grants
}

fn all_privileges() -> BTreeSet<Privilege> {
    [
        Privilege::Select,
        Privilege::Insert,
        Privilege::Update,
        Privilege::Delete,
        Privilege::Truncate,
        Privilege::References,
        Privilege::Trigger,
        Privilege::Execute,
        Privilege::Usage,
        Privilege::Create,
        Privilege::Connect,
        Privilege::Temporary,
    ]
    .into_iter()
    .collect()
}

/// Fetch privileges on tables, views, materialized views, and sequences.
///
/// Uses `pg_class` joined with `pg_namespace`. The `relkind` column determines
/// the object type:
///   'r' = table, 'v' = view, 'm' = materialized view, 'S' = sequence, 'p' = partitioned table
///
/// NULL ACLs are handled with `acldefault('r', c.relowner)` for relations and
/// `acldefault('S', c.relowner)` for sequences.
async fn fetch_relation_privileges(
    pool: &PgPool,
    managed_schemas: &[&str],
) -> Result<Vec<AclRow>, sqlx::Error> {
    sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            grantee.rolname AS grantee,
            acl.privilege_type,
            n.nspname AS schema_name,
            c.relname AS object_name,
            CASE c.relkind
                WHEN 'r' THEN 'table'
                WHEN 'p' THEN 'table'
                WHEN 'v' THEN 'view'
                WHEN 'm' THEN 'materialized_view'
                WHEN 'S' THEN 'sequence'
            END AS obj_type
        FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        CROSS JOIN LATERAL aclexplode(
            COALESCE(
                c.relacl,
                acldefault(
                    CASE WHEN c.relkind = 'S' THEN 'S'::"char" ELSE 'r'::"char" END,
                    c.relowner
                )
            )
        ) AS acl
        LEFT JOIN pg_roles grantee ON grantee.oid = acl.grantee
        WHERE n.nspname = ANY($1)
          AND c.relkind IN ('r', 'p', 'v', 'm', 'S')
        ORDER BY n.nspname, c.relname
        "#,
    )
    .bind(managed_schemas)
    .fetch_all(pool)
    .await
}

/// Fetch privileges on schemas.
///
/// Uses `pg_namespace`. The ACL default type is 'n' (namespace).
/// For schema grants, the object_name is the schema name itself.
async fn fetch_schema_privileges(
    pool: &PgPool,
    managed_schemas: &[&str],
) -> Result<Vec<AclRow>, sqlx::Error> {
    sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            grantee.rolname AS grantee,
            acl.privilege_type,
            NULL::text AS schema_name,
            n.nspname AS object_name,
            'schema' AS obj_type
        FROM pg_namespace n
        CROSS JOIN LATERAL aclexplode(
            COALESCE(
                n.nspacl,
                acldefault('n'::"char", n.nspowner)
            )
        ) AS acl
        LEFT JOIN pg_roles grantee ON grantee.oid = acl.grantee
        WHERE n.nspname = ANY($1)
        ORDER BY n.nspname
        "#,
    )
    .bind(managed_schemas)
    .fetch_all(pool)
    .await
}

/// Fetch privileges on functions/procedures.
///
/// Uses `pg_proc` joined with `pg_namespace`. The ACL default type is 'f'.
/// Function names can be overloaded, so we include the OID-derived
/// identity signature via `pg_catalog.pg_get_function_identity_arguments()`.
async fn fetch_function_privileges(
    pool: &PgPool,
    managed_schemas: &[&str],
) -> Result<Vec<AclRow>, sqlx::Error> {
    sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            grantee.rolname AS grantee,
            acl.privilege_type,
            n.nspname AS schema_name,
            p.proname || '(' || pg_catalog.pg_get_function_identity_arguments(p.oid) || ')' AS object_name,
            'function' AS obj_type
        FROM pg_proc p
        JOIN pg_namespace n ON n.oid = p.pronamespace
        CROSS JOIN LATERAL aclexplode(
            COALESCE(
                p.proacl,
                acldefault('f'::"char", p.proowner)
            )
        ) AS acl
        LEFT JOIN pg_roles grantee ON grantee.oid = acl.grantee
        WHERE n.nspname = ANY($1)
        ORDER BY n.nspname, p.proname
        "#,
    )
    .bind(managed_schemas)
    .fetch_all(pool)
    .await
}

/// Fetch privileges on types/domains.
///
/// Uses `pg_type` joined with `pg_namespace`. The ACL default type is 'T'.
/// We filter out internal/array types (typname not starting with '_',
/// typtype not 'p' for pseudo-types).
async fn fetch_type_privileges(
    pool: &PgPool,
    managed_schemas: &[&str],
) -> Result<Vec<AclRow>, sqlx::Error> {
    sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            grantee.rolname AS grantee,
            acl.privilege_type,
            n.nspname AS schema_name,
            t.typname AS object_name,
            'type' AS obj_type
        FROM pg_type t
        JOIN pg_namespace n ON n.oid = t.typnamespace
        CROSS JOIN LATERAL aclexplode(
            COALESCE(
                t.typacl,
                acldefault('T'::"char", t.typowner)
            )
        ) AS acl
        LEFT JOIN pg_roles grantee ON grantee.oid = acl.grantee
        WHERE n.nspname = ANY($1)
          AND t.typname NOT LIKE '\_%'
          AND t.typtype <> 'p'
        ORDER BY n.nspname, t.typname
        "#,
    )
    .bind(managed_schemas)
    .fetch_all(pool)
    .await
}

/// Fetch database-level privileges on the current database.
///
/// Uses `pg_database`. The ACL default type is 'd'.
/// This is separate because it's not schema-scoped; we always query the
/// current database.
pub async fn fetch_database_privileges(
    pool: &PgPool,
    managed_roles: &[&str],
) -> Result<BTreeMap<GrantKey, GrantState>, sqlx::Error> {
    let rows = sqlx::query_as::<_, AclRow>(
        r#"
        SELECT
            grantee.rolname AS grantee,
            acl.privilege_type,
            NULL::text AS schema_name,
            db.datname AS object_name,
            'database' AS obj_type
        FROM pg_database db
        CROSS JOIN LATERAL aclexplode(
            COALESCE(
                db.datacl,
                acldefault('d'::"char", db.datdba)
            )
        ) AS acl
        LEFT JOIN pg_roles grantee ON grantee.oid = acl.grantee
        WHERE db.datname = current_database()
        ORDER BY db.datname
        "#,
    )
    .fetch_all(pool)
    .await?;

    let mut grants: BTreeMap<GrantKey, GrantState> = BTreeMap::new();

    for row in rows {
        let grantee = match row.grantee {
            Some(ref name) => name,
            None => continue,
        };

        if !managed_roles.contains(&grantee.as_str()) {
            continue;
        }

        let privilege = match acl_char_to_privilege(&row.privilege_type) {
            Some(privilege) => privilege,
            None => continue,
        };

        let key = GrantKey {
            role: grantee.clone(),
            object_type: ObjectType::Database,
            schema: None,
            name: Some(row.object_name.clone()),
        };

        let entry = grants.entry(key).or_insert_with(|| GrantState {
            privileges: std::collections::BTreeSet::new(),
        });
        entry.privileges.insert(privilege);
    }

    Ok(grants)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::WildcardGrantPattern;

    #[test]
    fn acl_char_mapping_covers_all_privileges() {
        // Standard PostgreSQL ACL characters
        let cases = vec![
            ("r", Privilege::Select),
            ("a", Privilege::Insert),
            ("w", Privilege::Update),
            ("d", Privilege::Delete),
            ("D", Privilege::Truncate),
            ("x", Privilege::References),
            ("t", Privilege::Trigger),
            ("X", Privilege::Execute),
            ("U", Privilege::Usage),
            ("C", Privilege::Create),
            ("c", Privilege::Connect),
            ("T", Privilege::Temporary),
        ];
        for (char, expected) in cases {
            assert_eq!(
                acl_char_to_privilege(char),
                Some(expected),
                "failed for char '{char}'"
            );
        }
        assert_eq!(acl_char_to_privilege("Z"), None);
    }

    #[test]
    fn obj_type_str_mapping_covers_all_types() {
        let cases = vec![
            ("table", ObjectType::Table),
            ("view", ObjectType::View),
            ("materialized_view", ObjectType::MaterializedView),
            ("sequence", ObjectType::Sequence),
            ("function", ObjectType::Function),
            ("schema", ObjectType::Schema),
            ("database", ObjectType::Database),
            ("type", ObjectType::Type),
        ];
        for (type_str, expected) in cases {
            assert_eq!(
                obj_type_str_to_object_type(type_str),
                Some(expected),
                "failed for type_str '{type_str}'"
            );
        }
        assert_eq!(obj_type_str_to_object_type("unknown"), None);
    }

    #[test]
    fn wildcard_normalization_promotes_shared_table_privileges() {
        let mut grants = BTreeMap::new();
        grants.insert(
            GrantKey {
                role: "inventory-editor".to_string(),
                object_type: ObjectType::Table,
                schema: Some("inventory".to_string()),
                name: Some("widgets".to_string()),
            },
            GrantState {
                privileges: [Privilege::Select, Privilege::Insert].into_iter().collect(),
            },
        );
        grants.insert(
            GrantKey {
                role: "inventory-editor".to_string(),
                object_type: ObjectType::Table,
                schema: Some("inventory".to_string()),
                name: Some("orders".to_string()),
            },
            GrantState {
                privileges: [Privilege::Select].into_iter().collect(),
            },
        );

        let inventory = BTreeMap::from([(
            (ObjectType::Table, "inventory".to_string()),
            BTreeSet::from(["orders".to_string(), "widgets".to_string()]),
        )]);
        let selectors = vec![WildcardGrantPattern {
            role: "inventory-editor".to_string(),
            object_type: ObjectType::Table,
            schema: "inventory".to_string(),
        }];

        let normalized = normalize_wildcard_grants(grants, &inventory, &selectors);

        let wildcard = normalized
            .get(&GrantKey {
                role: "inventory-editor".to_string(),
                object_type: ObjectType::Table,
                schema: Some("inventory".to_string()),
                name: Some("*".to_string()),
            })
            .expect("wildcard grant should be synthesized");
        assert_eq!(wildcard.privileges, BTreeSet::from([Privilege::Select]));

        let specific = normalized
            .get(&GrantKey {
                role: "inventory-editor".to_string(),
                object_type: ObjectType::Table,
                schema: Some("inventory".to_string()),
                name: Some("widgets".to_string()),
            })
            .expect("extra object-specific privileges should remain");
        assert_eq!(specific.privileges, BTreeSet::from([Privilege::Insert]));
    }
}