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>, } impl PgClass { pub async fn fetch_namespace( &self, client: &mut WorkspaceClient, ) -> Result { 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>(kinds: I) -> WithKindInQuery { WithKindInQuery { kinds: kinds.into_iter().collect(), } } pub fn belonging_to_namespace< T: Clone + Copy + Display + sqlx::Type + for<'a> Encode<'a, Postgres>, >( namespace: T, ) -> BelongingToNamespaceQuery { BelongingToNamespaceQuery { namespace } } } #[derive(Clone, Copy, Debug)] pub struct BelongingToNamespaceQuery< T: Clone + Copy + Display + sqlx::Type + for<'a> Encode<'a, Postgres>, > { namespace: T, } impl + for<'a> Encode<'a, Postgres>> BelongingToNamespaceQuery { pub async fn fetch_all(self, client: &mut WorkspaceClient) -> sqlx::Result> { // `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" 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" from pg_class where oid = $1 "#, $value, ) }; } impl WithOidQuery { pub async fn fetch_one(self, client: &mut WorkspaceClient) -> Result { with_oid_sqlx_query!(self.oid) .fetch_one(client.get_conn()) .await } pub async fn fetch_optional( self, client: &mut WorkspaceClient, ) -> Result, sqlx::Error> { with_oid_sqlx_query!(self.oid) .fetch_optional(client.get_conn()) .await } } pub struct WithKindInQuery { kinds: Vec, } impl WithKindInQuery { pub async fn fetch_all( self, client: &mut WorkspaceClient, ) -> Result, 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" 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 } }