use crate::error::{TViewError, TViewResult};
use pgrx::datum::DatumWithOid;
use pgrx::pg_sys::{self, Oid};
use pgrx::prelude::*;
use std::ffi::CStr;
pub struct AsOwner {
saved_user: Oid,
saved_context: i32,
guc_level: i32,
}
impl AsOwner {
pub fn of_entity(entity: &str) -> TViewResult<Self> {
let meta = crate::catalog::TviewMeta::load_by_entity(entity)?.ok_or_else(|| {
TViewError::MetadataNotFound {
entity: entity.to_string(),
}
})?;
Self::of_table(meta.tview_oid)
}
pub fn of_table(table: Oid) -> TViewResult<Self> {
Ok(Self::role(relation_owner(table)?))
}
pub fn of_extension() -> TViewResult<Self> {
let owner = Spi::connect(|client| {
client
.select(
"SELECT extowner FROM pg_catalog.pg_extension WHERE extname = 'pg_tviews'",
None,
&[],
)?
.first()
.get_one::<Oid>()
})
.map_err(|e| TViewError::CatalogError {
operation: "Look up the owner of pg_tviews".to_string(),
pg_error: e.to_string(),
})?
.ok_or_else(|| TViewError::CatalogError {
operation: "Look up the owner of pg_tviews".to_string(),
pg_error: "extension pg_tviews is not installed".to_string(),
})?;
Ok(Self::role(owner))
}
fn role(owner: Oid) -> Self {
let mut saved_user = pg_sys::InvalidOid;
let mut saved_context = 0;
let guc_level = unsafe {
pg_sys::GetUserIdAndSecContext(&raw mut saved_user, &raw mut saved_context);
pg_sys::SetUserIdAndSecContext(
owner,
saved_context
| (pg_sys::SECURITY_LOCAL_USERID_CHANGE
| pg_sys::SECURITY_RESTRICTED_OPERATION)
.cast_signed(),
);
let level = pg_sys::NewGUCNestLevel();
set_local(c"search_path", c"pg_catalog, pg_temp");
pin_settings();
level
};
Self {
saved_user,
saved_context,
guc_level,
}
}
}
impl Drop for AsOwner {
fn drop(&mut self) {
if std::thread::panicking() {
return;
}
unsafe {
pg_sys::AtEOXact_GUC(false, self.guc_level);
pg_sys::SetUserIdAndSecContext(self.saved_user, self.saved_context);
}
}
}
pub const RENDER_SETTINGS: [(&CStr, &CStr); 5] = [
(c"TimeZone", c"UTC"),
(c"DateStyle", c"ISO, YMD"),
(c"IntervalStyle", c"postgres"),
(c"extra_float_digits", c"1"),
(c"bytea_output", c"hex"),
];
pub struct RenderPin {
guc_level: i32,
}
impl RenderPin {
#[must_use]
pub fn new() -> Self {
let guc_level = unsafe {
let level = pg_sys::NewGUCNestLevel();
pin_settings();
level
};
Self { guc_level }
}
}
impl Default for RenderPin {
fn default() -> Self {
Self::new()
}
}
impl Drop for RenderPin {
fn drop(&mut self) {
if std::thread::panicking() {
return;
}
unsafe { pg_sys::AtEOXact_GUC(false, self.guc_level) };
}
}
unsafe fn set_local(name: &CStr, value: &CStr) {
unsafe {
pg_sys::set_config_option(
name.as_ptr(),
value.as_ptr(),
pg_sys::GucContext::PGC_USERSET,
pg_sys::GucSource::PGC_S_SESSION,
pg_sys::GucAction::GUC_ACTION_SAVE,
true,
0,
false,
);
}
}
unsafe fn pin_settings() {
for (name, value) in RENDER_SETTINGS {
unsafe { set_local(name, value) };
}
}
pub fn require_owner(table: Oid, tview: &str) -> TViewResult<()> {
let allowed = Spi::connect(|client| {
let args =
[unsafe { DatumWithOid::new(table, PgOid::BuiltIn(PgBuiltInOids::OIDOID).value()) }];
client
.select(
"SELECT COALESCE((SELECT pg_catalog.pg_has_role(c.relowner, 'USAGE') \
FROM pg_catalog.pg_class c WHERE c.oid = $1), false) \
OR pg_catalog.pg_has_role(e.extowner, 'USAGE') \
FROM pg_catalog.pg_extension e WHERE e.extname = 'pg_tviews'",
None,
&args,
)?
.first()
.get_one::<bool>()
})
.map_err(|e| TViewError::CatalogError {
operation: format!("Check ownership of {tview}"),
pg_error: e.to_string(),
})?;
if allowed != Some(true) {
pg_sys::panic::ErrorReport::new(
PgSqlErrorCode::ERRCODE_INSUFFICIENT_PRIVILEGE,
format!("must be owner of TVIEW {tview}"),
function_name!(),
)
.report(PgLogLevel::ERROR);
}
Ok(())
}
fn relation_owner(table: Oid) -> TViewResult<Oid> {
let owner = unsafe {
let tuple = pg_sys::SearchSysCache1(
pg_sys::SysCacheIdentifier::RELOID.cast_signed(),
pg_sys::Datum::from(table),
);
if tuple.is_null() {
None
} else {
#[allow(clippy::cast_ptr_alignment)]
let form = pg_sys::GETSTRUCT(tuple).cast::<pg_sys::FormData_pg_class>();
let owner = (*form).relowner;
pg_sys::ReleaseSysCache(tuple);
Some(owner)
}
};
owner.ok_or_else(|| TViewError::CatalogError {
operation: format!("Look up the owner of relation {table:?}"),
pg_error: "relation does not exist".to_string(),
})
}