1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
//LICENSE
//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
//LICENSE
//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
//LICENSE
//LICENSE All rights reserved.
//LICENSE
//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
//! Provides a safe wrapper around Postgres' `pg_sys::RelationData` struct
use crate::{
    direct_function_call, name_data_to_str, pg_sys, FromDatum, IntoDatum, PgBox, PgTupleDesc,
};
use pgrx_sql_entity_graph::metadata::{
    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
};
use std::ops::Deref;
use std::os::raw::c_char;

pub struct PgRelation {
    boxed: PgBox<pg_sys::RelationData>,
    need_close: bool,
    lockmode: Option<pg_sys::LOCKMODE>,
}

impl PgRelation {
    /// Wrap a Postgres-provided `pg_sys::Relation`.
    ///
    /// It is assumed that Postgres will later `RelationClose()` the provided relation pointer.
    /// As such, it is not closed when this instance is dropped
    ///
    /// ## Safety
    ///
    /// This method is unsafe as we cannot ensure that this relation will later be closed by Postgres
    pub unsafe fn from_pg(ptr: pg_sys::Relation) -> Self {
        PgRelation { boxed: PgBox::from_pg(ptr), need_close: false, lockmode: None }
    }

    /// Wrap a Postgres-provided `pg_sys::Relation`.
    ///
    /// The provided `Relation` will be closed via `pg_sys::RelationClose` when this instance is dropped
    pub unsafe fn from_pg_owned(ptr: pg_sys::Relation) -> Self {
        PgRelation { boxed: PgBox::from_pg(ptr), need_close: true, lockmode: None }
    }

    /// Given a relation oid, use `pg_sys::RelationIdGetRelation()` to open the relation
    ///
    /// If the specified relation oid was recently deleted, this function will panic.
    ///
    /// Additionally, the relation is closed via `pg_sys::RelationClose()` when this instance is
    /// dropped.
    ///
    /// ## Safety
    ///
    /// The caller should already have at least AccessShareLock on the relation ID, else there are
    /// nasty race conditions.
    ///
    /// As such, this function is unsafe as we cannot guarantee that this requirement is true.
    pub unsafe fn open(oid: pg_sys::Oid) -> Self {
        let rel = pg_sys::RelationIdGetRelation(oid);
        if rel.is_null() {
            // relation was recently deleted
            panic!("Cannot open relation with oid={:?}", oid);
        }

        PgRelation { boxed: PgBox::from_pg(rel), need_close: true, lockmode: None }
    }

    /// relation_open - open any relation by relation OID
    ///
    /// If lockmode is not "NoLock", the specified kind of lock is
    /// obtained on the relation.
    ///
    /// An error is raised if the relation does not exist.
    ///
    /// # Safety
    ///
    /// Using an inappropriate lockmode, such as using NoLock on a relation
    /// that has not been previously locked, may result in undefined behavior.
    ///
    /// NB: a "relation" is anything with a pg_class entry.  The caller is
    /// expected to check whether the relkind is something it can handle.
    ///
    /// The opened relation is automatically closed via `pg_sys::relation_close()`
    /// when this instance is dropped
    #[deny(unsafe_op_in_unsafe_fn)]
    pub unsafe fn with_lock(oid: pg_sys::Oid, lockmode: pg_sys::LOCKMODE) -> Self {
        unsafe {
            PgRelation {
                boxed: PgBox::from_pg(pg_sys::relation_open(oid, lockmode)),
                need_close: true,
                lockmode: Some(lockmode),
            }
        }
    }

    /// Given a relation name, use `pg_sys::to_regclass` to look up its oid, and then
    /// `pg_sys::RelationIdGetRelation()` to open the relation.
    ///
    /// If the specified relation name is not found, we return an `Err(&str)`.
    ///
    /// If the specified relation was recently deleted, this function will panic.
    ///
    /// Additionally, the relation is closed via `pg_sys::RelationClose()` when this instance is
    /// dropped.
    ///
    /// ## Safety
    ///
    /// The caller should already have at least AccessShareLock on the relation ID, else there are
    /// nasty race conditions.
    ///
    /// As such, this function is unsafe as we cannot guarantee that this requirement is true.
    pub unsafe fn open_with_name(relname: &str) -> std::result::Result<Self, &'static str> {
        match direct_function_call::<pg_sys::Oid>(pg_sys::to_regclass, &[relname.into_datum()]) {
            Some(oid) => Ok(PgRelation::open(oid)),
            None => Err("no such relation"),
        }
    }

    /// Given a relation name, use `pg_sys::to_regclass` to look up its oid, and then
    /// open it with an AccessShareLock
    ///
    /// If the specified relation name is not found, we return an `Err(&str)`.
    ///
    /// If the specified relation was recently deleted, this function will panic.
    ///
    /// Additionally, the relation is closed via `pg_sys::RelationClose()` when this instance is
    /// dropped.
    pub fn open_with_name_and_share_lock(relname: &str) -> std::result::Result<Self, &'static str> {
        unsafe {
            match direct_function_call::<pg_sys::Oid>(pg_sys::to_regclass, &[relname.into_datum()])
            {
                Some(oid) => {
                    Ok(PgRelation::with_lock(oid, pg_sys::AccessShareLock as pg_sys::LOCKMODE))
                }
                None => Err("no such relation"),
            }
        }
    }

    /// RelationGetRelationName
    ///            Returns the rel's name.
    ///
    /// Note that the name is only unique within the containing namespace.
    pub fn name(&self) -> &str {
        let rd_rel = unsafe { self.boxed.rd_rel.as_ref() }.unwrap();
        name_data_to_str(&rd_rel.relname)
    }

    /// RelationGetRelid
    ///          Returns the OID of the relation
    #[inline]
    pub fn oid(&self) -> pg_sys::Oid {
        let rel = &self.boxed;
        rel.rd_id
    }

    /// RelationGetNamespace
    ///            Returns the rel's namespace OID.
    pub fn namespace_oid(&self) -> pg_sys::Oid {
        // SAFETY: we know self.boxed and its members are correct as we created it
        let rd_rel: PgBox<pg_sys::FormData_pg_class> = unsafe { PgBox::from_pg(self.boxed.rd_rel) };
        rd_rel.relnamespace
    }

    /// What is the name of the namespace in which this relation is located?
    pub fn namespace(&self) -> &str {
        unsafe { core::ffi::CStr::from_ptr(pg_sys::get_namespace_name(self.namespace_oid())) }
            .to_str()
            .expect("unable to convert namespace name to UTF8")
    }

    /// If this `PgRelation` represents an index, return the `PgRelation` for the heap
    /// relation to which it is attached
    pub fn heap_relation(&self) -> Option<PgRelation> {
        // SAFETY: we know self.boxed and its members are correct as we created it
        let rd_index: PgBox<pg_sys::FormData_pg_index> =
            unsafe { PgBox::from_pg(self.boxed.rd_index) };
        if rd_index.is_null() {
            None
        } else {
            Some(unsafe { PgRelation::open(rd_index.indrelid) })
        }
    }

    /// Return an iterator of indices, as `PgRelation`s, attached to this relation
    #[cfg(feature = "cshim")]
    pub fn indices(
        &self,
        lockmode: pg_sys::LOCKMODE,
    ) -> impl std::iter::Iterator<Item = PgRelation> {
        use crate::PgList;
        // SAFETY: we know self.boxed is a valid pointer as we created it
        let list = unsafe {
            PgList::<pg_sys::Oid>::from_pg(pg_sys::RelationGetIndexList(self.boxed.as_ptr()))
        };

        list.iter_oid()
            .filter(|oid| *oid != pg_sys::InvalidOid)
            .map(|oid| unsafe { PgRelation::with_lock(oid, lockmode) })
            .collect::<Vec<PgRelation>>()
            .into_iter()
    }

    /// Returned a wrapped `PgTupleDesc`
    ///
    /// The returned `PgTupleDesc` is tied to the lifetime of this `PgRelation` instance.
    ///
    /// ```rust,no_run
    /// use pgrx::{PgRelation, pg_sys};
    /// # let example_pg_class_oid = |i| { unsafe { pg_sys::Oid::from_u32_unchecked(i) } };
    /// let oid = example_pg_class_oid(42); // a valid pg_class "oid" value
    /// let relation = unsafe { PgRelation::from_pg(pg_sys::RelationIdGetRelation(oid) ) };
    /// let tupdesc = relation.tuple_desc();
    ///
    /// // assert that the tuple descriptor has 12 attributes
    /// assert_eq!(tupdesc.len(), 12);
    /// ```
    pub fn tuple_desc(&self) -> PgTupleDesc {
        PgTupleDesc::from_relation(&self)
    }

    /// Number of tuples in this relation (not always up-to-date)
    pub fn reltuples(&self) -> Option<f32> {
        let reltuples = unsafe { self.boxed.rd_rel.as_ref() }.expect("rd_rel is NULL").reltuples;

        if reltuples == 0f32 {
            None
        } else {
            Some(reltuples)
        }
    }

    pub fn is_table(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_RELATION as c_char
    }

    pub fn is_matview(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_MATVIEW as c_char
    }

    pub fn is_index(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_INDEX as c_char
    }

    pub fn is_view(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_VIEW as c_char
    }

    pub fn is_sequence(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_SEQUENCE as c_char
    }

    pub fn is_composite_type(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_COMPOSITE_TYPE as c_char
    }

    pub fn is_foreign_table(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_FOREIGN_TABLE as c_char
    }

    pub fn is_partitioned_table(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_PARTITIONED_TABLE as c_char
    }

    pub fn is_toast_value(&self) -> bool {
        let rd_rel: &pg_sys::FormData_pg_class =
            unsafe { self.boxed.rd_rel.as_ref().expect("rd_rel is NULL") };
        rd_rel.relkind == pg_sys::RELKIND_TOASTVALUE as c_char
    }

    /// ensures that the returned `PgRelation` is closed by Rust when it is dropped
    pub fn to_owned(mut self) -> Self {
        self.need_close = true;
        self
    }
}

impl Clone for PgRelation {
    /// Same as calling `PgRelation::with_lock(AccessShareLock)` on the underlying relation id
    fn clone(&self) -> Self {
        unsafe { PgRelation::with_lock(self.rd_id, pg_sys::AccessShareLock as pg_sys::LOCKMODE) }
    }
}

impl FromDatum for PgRelation {
    unsafe fn from_polymorphic_datum(
        datum: pg_sys::Datum,
        is_null: bool,
        _typoid: pg_sys::Oid,
    ) -> Option<PgRelation> {
        if is_null {
            None
        } else {
            Some(PgRelation::with_lock(
                unsafe { pg_sys::Oid::from_u32_unchecked(u32::try_from(datum.value()).ok()?) },
                pg_sys::AccessShareLock as pg_sys::LOCKMODE,
            ))
        }
    }
}

impl IntoDatum for PgRelation {
    fn into_datum(self) -> Option<pg_sys::Datum> {
        Some(pg_sys::Datum::from(self.oid()))
    }

    fn type_oid() -> pg_sys::Oid {
        pg_sys::REGCLASSOID
    }
}

impl Deref for PgRelation {
    type Target = PgBox<pg_sys::RelationData>;

    fn deref(&self) -> &Self::Target {
        &self.boxed
    }
}

impl Drop for PgRelation {
    fn drop(&mut self) {
        if !self.boxed.is_null() {
            if self.need_close {
                match self.lockmode {
                    None => unsafe { pg_sys::RelationClose(self.boxed.as_ptr()) },
                    Some(lockmode) => unsafe {
                        pg_sys::relation_close(self.boxed.as_ptr(), lockmode)
                    },
                }
            }
        }
    }
}

unsafe impl SqlTranslatable for PgRelation {
    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
        Ok(SqlMapping::literal("regclass"))
    }
    fn return_sql() -> Result<Returns, ReturnsError> {
        Ok(Returns::One(SqlMapping::literal("regclass")))
    }
}