phonograph/phono-backends/src/pg_class.rs

use std::fmt::Display;

use sqlx::{Encode, Postgres, postgres::types::Oid, query_as, query_as_unchecked};

use crate::{
    client::WorkspaceClient, escape_identifier, pg_acl::PgAclItem, pg_namespace::PgNamespace,
};

#[derive(Clone, Debug)]
pub struct PgClass {
    /// Row identifier
    pub oid: Oid,
    /// Name of the table, index, view, etc.
    pub relname: String,
    /// The OID of the namespace that contains this relation
    pub relnamespace: Oid,
    /// SYNTHESIZED: name of this relation's namespace as text
    pub regnamespace: String,
    /// The OID of the data type that corresponds to this table's row type, if any; zero for indexes, sequences, and toast tables, which have no pg_type entry
    pub reltype: Oid,
    /// For typed tables, the OID of the underlying composite type; zero for all other relations
    pub reloftype: Oid,
    /// Owner of the relation
    pub relowner: Oid,
    /// SYNTHESIZED: name of this relation's owner role as text
    pub regowner: String,
    /// r = ordinary table, i = index, S = sequence, t = TOAST table, v = view, m = materialized view, c = composite type, f = foreign table, p = partitioned table, I = partitioned index
    pub relkind: i8,
    /// Number of user columns in the relation (system columns not counted). There must be this many corresponding entries in pg_attribute. See also pg_attribute.attnum.
    pub relnatts: i16,
    /// Number of CHECK constraints on the table; see pg_constraint catalog
    pub relchecks: i16,
    /// True if table has (or once had) rules; see pg_rewrite catalog
    pub relhasrules: bool,
    /// True if table has (or once had) triggers; see pg_trigger catalog
    pub relhastriggers: bool,
    /// True if table or index has (or once had) any inheritance children or partitions
    pub relhassubclass: bool,
    /// True if table has row-level security enabled; see pg_policy catalog
    pub relrowsecurity: bool,
    /// True if row-level security (when enabled) will also apply to table owner; see pg_policy catalog
    pub relforcerowsecurity: bool,
    /// True if relation is populated (this is true for all relations other than some materialized views)
    pub relispopulated: bool,
    /// True if table or index is a partition
    pub relispartition: bool,
    pub relacl: Option<Vec<PgAclItem>>,
}

impl PgClass {
    pub async fn fetch_namespace(
        &self,
        client: &mut WorkspaceClient,
    ) -> Result<PgNamespace, sqlx::Error> {
        PgNamespace::fetch_by_oid(self.relnamespace, client.get_conn())
            .await?
            // If client has access to the class, it would expect to have access
            // to the namespace that contains it. If not, that's an error.
            .ok_or(sqlx::Error::RowNotFound)
    }

    /// Get ecaped identifier, including namespace.
    pub fn get_identifier(&self) -> String {
        format!(
            "{0}.{1}",
            escape_identifier(&self.regnamespace),
            escape_identifier(&self.relname)
        )
    }

    pub fn with_oid(oid: Oid) -> WithOidQuery {
        WithOidQuery { oid }
    }

    pub fn with_kind_in<I: IntoIterator<Item = PgRelKind>>(kinds: I) -> WithKindInQuery {
        WithKindInQuery {
            kinds: kinds.into_iter().collect(),
        }
    }

    pub fn belonging_to_namespace<
        T: Clone + Copy + Display + sqlx::Type<Postgres> + for<'a> Encode<'a, Postgres>,
    >(
        namespace: T,
    ) -> BelongingToNamespaceQuery<T> {
        BelongingToNamespaceQuery { namespace }
    }
}
#[derive(Clone, Copy, Debug)]
pub struct BelongingToNamespaceQuery<
    T: Clone + Copy + Display + sqlx::Type<Postgres> + for<'a> Encode<'a, Postgres>,
> {
    namespace: T,
}

impl<T: Clone + Copy + Display + sqlx::Type<Postgres> + for<'a> Encode<'a, Postgres>>
    BelongingToNamespaceQuery<T>
{
    pub async fn fetch_all(self, client: &mut WorkspaceClient) -> sqlx::Result<Vec<PgClass>> {
        // `query_as!()` rightly complains that there may not be a built-in type
        // mapping for `T` to `regnamespace`.
        // TODO: Figure out whether it's possible to add a trait bound that
        // ensures the mapping exists.
        query_as_unchecked!(
            PgClass,
            r#"
select
    oid,
    relname,
    relnamespace,
    relnamespace::regnamespace::text as "regnamespace!",
    reltype,
    reloftype,
    relowner,
    relowner::regrole::text as "regowner!",
    relkind,
    relnatts,
    relchecks,
    relhasrules,
    relhastriggers,
    relhassubclass,
    relrowsecurity,
    relforcerowsecurity,
    relispopulated,
    relispartition,
    relacl::text[] as "relacl: Vec<PgAclItem>"
from pg_class
where
    relnamespace = $1::regnamespace::oid
"#,
            self.namespace,
        )
        .fetch_all(client.get_conn())
        .await
    }
}

pub struct WithOidQuery {
    oid: Oid,
}

// Extracted as macro so that fetch_one() and fetch_optional() methods can
// reuse the same code.
macro_rules! with_oid_sqlx_query {
    ($value:expr) => {
        query_as!(
            PgClass,
            r#"
select
    oid,
    relname,
    relnamespace,
    relnamespace::regnamespace::text as "regnamespace!",
    reltype,
    reloftype,
    relowner,
    relowner::regrole::text as "regowner!",
    relkind,
    relnatts,
    relchecks,
    relhasrules,
    relhastriggers,
    relhassubclass,
    relrowsecurity,
    relforcerowsecurity,
    relispopulated,
    relispartition,
    relacl::text[] as "relacl: Vec<PgAclItem>"
from pg_class
where
    oid = $1
"#,
            $value,
        )
    };
}

impl WithOidQuery {
    pub async fn fetch_one(self, client: &mut WorkspaceClient) -> Result<PgClass, sqlx::Error> {
        with_oid_sqlx_query!(self.oid)
            .fetch_one(client.get_conn())
            .await
    }

    pub async fn fetch_optional(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<PgClass>, sqlx::Error> {
        with_oid_sqlx_query!(self.oid)
            .fetch_optional(client.get_conn())
            .await
    }
}

pub struct WithKindInQuery {
    kinds: Vec<PgRelKind>,
}

impl WithKindInQuery {
    pub async fn fetch_all(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Vec<PgClass>, sqlx::Error> {
        let kinds_i8: Vec<_> = self
            .kinds
            .into_iter()
            .map(|kind| kind.to_u8() as i8)
            .collect();
        query_as!(
            PgClass,
            r#"
select
    oid,
    relname,
    relnamespace,
    relnamespace::regnamespace::text as "regnamespace!",
    reltype,
    reloftype,
    relowner,
    relowner::regrole::text as "regowner!",
    relkind,
    relnatts,
    relchecks,
    relhasrules,
    relhastriggers,
    relhassubclass,
    relrowsecurity,
    relforcerowsecurity,
    relispopulated,
    relispartition,
    relacl::text[] as "relacl: Vec<PgAclItem>"
from pg_class
where
    relkind = any($1)
"#,
            kinds_i8.as_slice(),
        )
        .fetch_all(client.get_conn())
        .await
    }
}

pub enum PgRelKind {
    OrdinaryTable,
    Index,
    Sequence,
    ToastTable,
    View,
    MaterializedView,
    CompositeType,
    ForeignTable,
    PartitionedTable,
    PartitionedIndex,
}

impl PgRelKind {
    pub fn to_u8(&self) -> u8 {
        let ch = match self {
            Self::OrdinaryTable => 'r',
            Self::Index => 'i',
            Self::Sequence => 'S',
            Self::ToastTable => 't',
            Self::View => 'v',
            Self::MaterializedView => 'm',
            Self::CompositeType => 'c',
            Self::ForeignTable => 'f',
            Self::PartitionedTable => 'p',
            Self::PartitionedIndex => 'I',
        };
        ch as u8
    }
}