1//! A row as JSON that does not depend on the session: the same row
2//! gives the same bytes (and so the same cache key) whoever asks.
3//!
4//! [`RowJson`] walks records, arrays and domains itself and writes dates
5//! and times as RFC 3339 with ECMAScript's expanded years; every other
6//! value is `to_json`'s, under the fixed [`CanonicalOutput`] settings.
7//! `to_json` alone writes BC years as a `" BC"` suffix, years past 9999
8//! with five digits, and `timetz` in its own offset (PG 18).
9
10use std::collections::HashMap;
11use std::ffi::CString;
12use std::fmt::Write as _;
13
14use pgrx::pg_sys;
15
16use crate::Error;
17
18/// Output settings `to_json` depends on, fixed for the call: timestamps
19/// in UTC (RFC 3339 with +00:00), intervals in ISO 8601, floats at their
20/// shortest exact form, bytea as hex, money in the C locale.
21const CANONICAL_OUTPUT: &[(&str, &str)] = &[
22    ("TimeZone", "UTC"),
23    ("IntervalStyle", "iso_8601"),
24    ("extra_float_digits", "1"),
25    ("bytea_output", "hex"),
26    ("lc_monetary", "C"),
27];
28
29/// The settings above for as long as it lives, by the mechanism a
30/// function's `SET` clause uses (fmgr.c): a GUC nest level, saved values,
31/// and `AtEOXact_GUC` to restore them. If an ERROR unwinds instead, the
32/// (sub)transaction abort restores them, so Drop does nothing then.
33pub(crate) struct CanonicalOutput {
34    nest_level: i32,
35}
36
37impl CanonicalOutput {
38    pub(crate) fn enter() -> Result<Self, Error> {
39        // SAFETY: called in a transaction, as every SQL function is.
40        let nest_level = unsafe { pg_sys::NewGUCNestLevel() };
41        let guard = CanonicalOutput { nest_level };
42        for (name, value) in CANONICAL_OUTPUT {
43            let (n, v) = (CString::new(*name).unwrap(), CString::new(*value).unwrap());
44            // SAFETY: valid NUL-terminated strings; GUC_ACTION_SAVE ties
45            // the change to `nest_level`; elevel ERROR reports a refusal.
46            let applied = unsafe {
47                pg_sys::set_config_option(
48                    n.as_ptr(),
49                    v.as_ptr(),
50                    pg_sys::GucContext::PGC_USERSET,
51                    pg_sys::GucSource::PGC_S_SESSION,
52                    pg_sys::GucAction::GUC_ACTION_SAVE,
53                    true,
54                    pg_sys::PGERROR as i32,
55                    false,
56                )
57            };
58            if applied <= 0 {
59                return Err(Error::Config(format!("could not set {name} for canonical output")));
60            }
61        }
62        Ok(guard)
63    }
64}
65
66impl Drop for CanonicalOutput {
67    fn drop(&mut self) {
68        if !std::thread::panicking() {
69            // SAFETY: pairs with NewGUCNestLevel in `enter`; `true` keeps
70            // no change (a SET clause's changes are always undone).
71            unsafe { pg_sys::AtEOXact_GUC(true, self.nest_level) };
72        }
73    }
74}
75
76/// Writes values as JSON: records as objects in attribute order (dropped
77/// attributes skipped), arrays as nested arrays by dimension, domains as
78/// their base type, dates and times by [`date`], [`timestamp`] and
79/// [`timetz`], and anything else through `to_json`.
80pub(crate) struct RowJson {
81    /// Where the `to_json` calls keep their per-type state: the scan's.
82    memory: pg_sys::MemoryContext,
83    /// A `to_json` call per leaf type, planned for that argument type.
84    leaves: HashMap<pg_sys::Oid, Box<pg_sys::FmgrInfo>>,
85}
86
87impl RowJson {
88    /// Keeps its state in the current memory context, which must outlive
89    /// every [`RowJson::write`] (the executor's per-query context).
90    pub(crate) fn new() -> Self {
91        // SAFETY: reading a backend global.
92        RowJson { memory: unsafe { pg_sys::CurrentMemoryContext }, leaves: HashMap::new() }
93    }
94
95    /// Appends `datum`, a non-NULL value of `typid`/`typmod`, as JSON.
96    ///
97    /// # Safety
98    ///
99    /// `datum` is a valid value of that type, in a transaction, with
100    /// [`CanonicalOutput`] entered.
101    pub(crate) unsafe fn write(&mut self, out: &mut String, datum: pg_sys::Datum, typid: pg_sys::Oid, typmod: i32) {
102        unsafe {
103            let mut typmod = typmod;
104            let base = pg_sys::getBaseTypeAndTypmod(typid, &mut typmod);
105            match base {
106                pg_sys::TIMESTAMPTZOID => out.push_str(&quoted(&timestamp(datum.value() as i64, true))),
107                pg_sys::TIMESTAMPOID => out.push_str(&quoted(&timestamp(datum.value() as i64, false))),
108                pg_sys::DATEOID => out.push_str(&quoted(&date(datum.value() as i32))),
109                pg_sys::TIMETZOID => {
110                    let t = datum.cast_mut_ptr::<pg_sys::TimeTzADT>();
111                    out.push_str(&quoted(&timetz((*t).time, (*t).zone)));
112                }
113                _ if pg_sys::type_is_rowtype(base) => self.record(out, datum),
114                _ if pg_sys::get_element_type(base) != pg_sys::InvalidOid => self.array(out, datum),
115                _ => self.leaf(out, datum, base, typmod),
116            }
117        }
118    }
119
120    unsafe fn record(&mut self, out: &mut String, datum: pg_sys::Datum) {
121        unsafe {
122            let header = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()).cast::<pg_sys::HeapTupleHeaderData>();
123            // The tuple's own type: an anonymous record's typmod names its
124            // blessed descriptor, as composite_to_json reads it.
125            let fields = (*header).t_choice.t_datum;
126            let desc = pg_sys::lookup_rowtype_tupdesc(fields.datum_typeid, fields.datum_typmod);
127            let natts = (*desc).natts as usize;
128            let mut tuple = pg_sys::HeapTupleData {
129                t_len: varsize(header.cast()),
130                t_self: std::mem::zeroed(),
131                t_tableOid: pg_sys::InvalidOid,
132                t_data: header,
133            };
134            let (mut values, mut nulls) = (vec![pg_sys::Datum::from(0usize); natts], vec![false; natts]);
135            pg_sys::heap_deform_tuple(&mut tuple, desc, values.as_mut_ptr(), nulls.as_mut_ptr());
136            let view = pgrx::PgTupleDesc::from_pg_unchecked(desc);
137            let attrs: Vec<(String, pg_sys::Oid, i32)> = (0..natts)
138                .map(|i| view.get(i).expect("an attribute per natts"))
139                .map(|a| {
140                    let name = if a.attisdropped { String::new() } else { std::ffi::CStr::from_ptr(a.attname.data.as_ptr()).to_string_lossy().into_owned() };
141                    (name, if a.attisdropped { pg_sys::InvalidOid } else { a.atttypid }, a.atttypmod)
142                })
143                .collect();
144            drop(view);
145            if (*desc).tdrefcount >= 0 {
146                pg_sys::DecrTupleDescRefCount(desc);
147            }
148            out.push('{');
149            let mut first = true;
150            for (i, (name, atttypid, atttypmod)) in attrs.into_iter().enumerate() {
151                if atttypid == pg_sys::InvalidOid {
152                    continue;
153                }
154                if !std::mem::take(&mut first) {
155                    out.push(',');
156                }
157                out.push_str(&quoted(&name));
158                out.push(':');
159                if nulls[i] {
160                    out.push_str("null");
161                } else {
162                    self.write(out, values[i], atttypid, atttypmod);
163                }
164            }
165            out.push('}');
166        }
167    }
168
169    unsafe fn array(&mut self, out: &mut String, datum: pg_sys::Datum) {
170        unsafe {
171            let array = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()).cast::<pg_sys::ArrayType>();
172            let ndim = (*array).ndim as usize;
173            let dims = std::slice::from_raw_parts(array.add(1).cast::<i32>(), ndim).to_vec();
174            let element = (*array).elemtype;
175            let (mut len, mut byval, mut align) = (0i16, false, 0 as std::ffi::c_char);
176            pg_sys::get_typlenbyvalalign(element, &mut len, &mut byval, &mut align);
177            let (mut values, mut nulls, mut n) = (std::ptr::null_mut(), std::ptr::null_mut(), 0);
178            pg_sys::deconstruct_array(array, element, len.into(), byval, align, &mut values, &mut nulls, &mut n);
179            let (values, nulls) = match n {
180                0 => (&[][..], &[][..]),
181                n => (std::slice::from_raw_parts(values, n as usize), std::slice::from_raw_parts(nulls, n as usize)),
182            };
183            if ndim == 0 {
184                out.push_str("[]");
185                return;
186            }
187            let mut next = 0;
188            self.dimension(out, &dims, values, nulls, &mut next, element);
189        }
190    }
191
192    unsafe fn dimension(
193        &mut self,
194        out: &mut String,
195        dims: &[i32],
196        values: &[pg_sys::Datum],
197        nulls: &[bool],
198        next: &mut usize,
199        element: pg_sys::Oid,
200    ) {
201        out.push('[');
202        for i in 0..dims[0] {
203            if i > 0 {
204                out.push(',');
205            }
206            if dims.len() > 1 {
207                unsafe { self.dimension(out, &dims[1..], values, nulls, next, element) };
208            } else {
209                if nulls[*next] {
210                    out.push_str("null");
211                } else {
212                    // Array elements carry no typmod.
213                    unsafe { self.write(out, values[*next], element, -1) };
214                }
215                *next += 1;
216            }
217        }
218        out.push(']');
219    }
220
221    /// `to_json` on one value, through a call whose argument has the
222    /// value's type: `to_json` reads the type off `fn_expr`.
223    unsafe fn leaf(&mut self, out: &mut String, datum: pg_sys::Datum, typid: pg_sys::Oid, typmod: i32) {
224        unsafe {
225            let memory = self.memory;
226            let call = self.leaves.entry(typid).or_insert_with(|| {
227                let mut call: Box<pg_sys::FmgrInfo> = Box::new(std::mem::zeroed());
228                pg_sys::fmgr_info_cxt(pg_sys::Oid::from(pg_sys::F_TO_JSON), &mut *call, memory);
229                let previous = pg_sys::MemoryContextSwitchTo(memory);
230                let arg = pg_sys::makeNullConst(typid, typmod, pg_sys::InvalidOid);
231                let args = pg_sys::lappend(std::ptr::null_mut(), arg.cast());
232                call.fn_expr = pg_sys::makeFuncExpr(
233                    pg_sys::Oid::from(pg_sys::F_TO_JSON),
234                    pg_sys::JSONOID,
235                    args,
236                    pg_sys::InvalidOid,
237                    pg_sys::InvalidOid,
238                    pg_sys::CoercionForm::COERCE_EXPLICIT_CALL,
239                )
240                .cast();
241                pg_sys::MemoryContextSwitchTo(previous);
242                call
243            });
244            let json = pg_sys::FunctionCall1Coll(&mut **call, pg_sys::InvalidOid, datum);
245            let text = pg_sys::text_to_cstring(json.cast_mut_ptr());
246            out.push_str(&std::ffi::CStr::from_ptr(text).to_string_lossy());
247            pg_sys::pfree(text.cast());
248        }
249    }
250}
251
252/// A varlena's total length.
253unsafe fn varsize(p: *const pg_sys::varlena) -> u32 {
254    // SAFETY: a detoasted (4-byte header) varlena.
255    unsafe { pgrx::varlena::varsize_4b(p) as u32 }
256}
257
258/// Microseconds from 2000-01-01, Postgres's epoch, to 1970-01-01's.
259const EPOCH_DAYS: i64 = 10_957;
260const DAY_US: i64 = 86_400_000_000;
261
262/// A `timestamp[tz]` (microseconds since 2000-01-01, UTC when `tz`) as
263/// RFC 3339, `+00:00` when `tz`; the infinities as `to_json` writes them.
264pub(crate) fn timestamp(us: i64, tz: bool) -> String {
265    match us {
266        i64::MIN => "-infinity".into(),
267        i64::MAX => "infinity".into(),
268        _ => {
269            let (days, time) = (us.div_euclid(DAY_US), us.rem_euclid(DAY_US));
270            format!("{}T{}{}", civil(days), clock(time), if tz { "+00:00" } else { "" })
271        }
272    }
273}
274
275/// A `date` (days since 2000-01-01) as an RFC 3339 full-date.
276pub(crate) fn date(days: i32) -> String {
277    match days {
278        i32::MIN => "-infinity".into(),
279        i32::MAX => "infinity".into(),
280        _ => civil(days.into()),
281    }
282}
283
284/// A `timetz` in UTC: `zone` is Postgres's, seconds west of UTC.
285pub(crate) fn timetz(time: i64, zone: i32) -> String {
286    format!("{}+00:00", clock((time + i64::from(zone) * 1_000_000).rem_euclid(DAY_US)))
287}
288
289/// `YYYY-MM-DD` in the proleptic Gregorian calendar with astronomical
290/// years, as ECMAScript's `toISOString`: four digits for 0000–9999,
291/// otherwise a sign and at least six (ISO 8601 expanded years).
292fn civil(days: i64) -> String {
293    // Howard Hinnant's civil_from_days, from 1970-01-01.
294    let z = days + EPOCH_DAYS + 719_468;
295    let era = z.div_euclid(146_097);
296    let doe = z.rem_euclid(146_097);
297    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
298    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
299    let mp = (5 * doy + 2) / 153;
300    let day = doy - (153 * mp + 2) / 5 + 1;
301    let month = if mp < 10 { mp + 3 } else { mp - 9 };
302    let year = yoe + era * 400 + i64::from(month <= 2);
303    let year = match year {
304        0..=9999 => format!("{year:04}"),
305        _ => format!("{}{:06}", if year < 0 { '-' } else { '+' }, year.abs()),
306    };
307    format!("{year}-{month:02}-{day:02}")
308}
309
310/// `HH:MM:SS` and a fraction with trailing zeros dropped, as `to_json`
311/// writes one. `24:00:00` stays as Postgres allows it.
312fn clock(us: i64) -> String {
313    let (secs, frac) = (us / 1_000_000, us % 1_000_000);
314    let mut s = format!("{:02}:{:02}:{:02}", secs / 3600, secs / 60 % 60, secs % 60);
315    if frac != 0 {
316        let _ = write!(s, ".{frac:06}");
317        s.truncate(s.trim_end_matches('0').len());
318    }
319    s
320}
321
322/// `s` as a JSON string.
323fn quoted(s: &str) -> String {
324    let mut out = String::with_capacity(s.len() + 2);
325    out.push('"');
326    for c in s.chars() {
327        match c {
328            '"' => out.push_str("\\\""),
329            '\\' => out.push_str("\\\\"),
330            c if (c as u32) < 0x20 => {
331                let _ = write!(out, "\\u{:04x}", c as u32);
332            }
333            c => out.push(c),
334        }
335    }
336    out.push('"');
337    out
338}