pgrx 0.18.0

//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.
//! Helper functions for creating Postgres UDFs
//!
//! Typically these functions are not necessary to call directly as they're used behind
//! the scenes by the code generated by the `#[pg_extern]` macro.

use crate::{FromDatum, PgBox, PgMemoryContexts, pg_sys, void_mut_ptr};
use core::{ptr, slice};
use pgrx_pg_sys::ffi::pg_guard_ffi_boundary;

/// A macro for specifying default argument values so they get properly translated to SQL in
/// `CREATE FUNCTION` statements
///
/// ## Examples
///
/// This example will create a SQL function like so:
///
/// ```sql
/// CREATE FUNCTION fun_with_default_arg_value(a integer, b integer DEFAULT 99) RETURNS integer ...;
/// ```
///
/// ```rust
/// use pgrx::prelude::*;
///
/// #[pg_extern]
/// fn fun_with_default_arg_value(a: i32, b: default!(i32, 99)) -> i32 {
///    a + b
/// }
/// ```
///
/// This allows users of this function, from within Postgres, to elide the `b` argument, and
/// Postgres will automatically use `99`.
#[macro_export]
macro_rules! default {
    ($ty:ty, $val:tt) => {
        $ty
    };
    ($ty:ty, $val:path) => {
        $ty
    };
    ($ty:ty, $val:expr) => {
        $ty
    };
}

/// The equivalent of a PostgreSQL `NULL`
///
/// This is used primarily in `default!()` macros.
pub struct NULL;

/// A macro for providing SQL names for the returned fields for functions that return a Rust tuple,
/// especially those that return a [`TableIterator`][crate::iter::TableIterator].
///
/// ## Examples
///
/// This example will create a SQL function like so:
///
/// ```sql
/// CREATE FUNCTION get_a_set() RETURNS TABLE (id integer, title text) ...;
/// ```
///
/// ```rust
/// use pgrx::prelude::*;
///
/// #[pg_extern]
/// fn get_a_set() -> TableIterator<'static, (name!(id, i32), name!(title, &'static str))> {
///     TableIterator::new(vec![1, 2, 3].into_iter().zip(vec!["A", "B", "C"].into_iter()))
/// }
/// ```
#[macro_export]
macro_rules! name {
    ($name:tt, $ty:ty) => {
        $ty
    };
}

/// Get a numbered argument for a `PG_FUNCTION_INFO_V1` function as the specified Rust type.
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
///
/// We also cannot ensure that the specified Rust type `T` is compatible with whatever the
/// underlying datum is at the argument `num` position.  This too, is your responsibility
// TODO: should relate back to lifetimes and use a bound like T: UnboxDatum<'dat>?
#[inline]
pub unsafe fn pg_getarg<T: FromDatum>(fcinfo: pg_sys::FunctionCallInfo, num: usize) -> Option<T> {
    let datum = pg_get_nullable_datum(fcinfo, num);
    unsafe {
        if T::GET_TYPOID {
            T::from_polymorphic_datum(datum.value, datum.isnull, super::pg_getarg_type(fcinfo, num))
        } else {
            T::from_datum(datum.value, datum.isnull)
        }
    }
}

/// Is the specified argument for a `PG_FUNCTION_INFO_V1` function NULL?
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
#[inline]
pub unsafe fn pg_arg_is_null(fcinfo: pg_sys::FunctionCallInfo, num: usize) -> bool {
    pg_get_nullable_datum(fcinfo, num).isnull
}

/// Get a numbered argument for a `PG_FUNCTION_INFO_V1` function as an Option containing a
/// [`pg_sys::Datum`].
///
/// If the specified argument Datum is NULL, returns [`Option::None`].
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
#[inline]
pub unsafe fn pg_getarg_datum(
    fcinfo: pg_sys::FunctionCallInfo,
    num: usize,
) -> Option<pg_sys::Datum> {
    if pg_arg_is_null(fcinfo, num) { None } else { Some(pg_get_nullable_datum(fcinfo, num).value) }
}

/// Get a numbered argument for a `PG_FUNCTION_INFO_V1` function as a raw [`pg_sys::Datum`].
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
#[inline]
pub unsafe fn pg_getarg_datum_raw(fcinfo: pg_sys::FunctionCallInfo, num: usize) -> pg_sys::Datum {
    pg_get_nullable_datum(fcinfo, num).value
}

/// Returns the [`pg_sys::NullableDatum`] for a given arg.
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
#[inline]
pub unsafe fn pg_get_nullable_datum(
    fcinfo: pg_sys::FunctionCallInfo,
    num: usize,
) -> pg_sys::NullableDatum {
    let _nullptr_check = ptr::NonNull::new(fcinfo).expect("fcinfo pointer must be non-null");
    unsafe {
        let nargs = (*fcinfo).nargs;
        let args_ptr: *const pg_sys::NullableDatum = ptr::addr_of!((*fcinfo).args).cast();
        let args = slice::from_raw_parts(args_ptr, nargs as _);
        args[num]
    }
}

/// Modifies the specified `fcinfo` struct to indicate that its return value is null.
///
/// # Examples
///
/// ```rust,no_run
/// use pgrx::pg_return_null;
/// use pgrx::prelude::*;
/// fn foo(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
///     return unsafe { pg_return_null(fcinfo) };
/// }
/// ```
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
#[inline]
pub unsafe fn pg_return_null(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
    let fcinfo = unsafe { fcinfo.as_mut() }.unwrap();
    fcinfo.isnull = true;
    pg_sys::Datum::from(0)
}

/// Get the collation the function call should use
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid
/// [`pg_sys::FunctionCallInfo`] pointer.  This is your responsibility.
pub unsafe fn pg_get_collation(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Oid {
    let fcinfo = unsafe { fcinfo.as_mut() }.unwrap();
    fcinfo.fncollation
}

/// # Safety
///
/// The provided `fcinfo` must be valid otherwise this function results in undefined behavior due
/// to an out of bounds read.
#[inline]
pub unsafe fn pg_getarg_type(fcinfo: pg_sys::FunctionCallInfo, num: usize) -> pg_sys::Oid {
    pg_sys::get_fn_expr_argtype(fcinfo.as_ref().unwrap().flinfo, num as std::os::raw::c_int)
}

/// Indicates that a `PG_FUNCTION_INFO_V1` function is returning a SQL "void".
///
/// # Example
///
/// ```rust,no_run
/// use pgrx::pg_return_void;
/// use pgrx::prelude::*;
///
/// fn foo(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
///     pg_return_void()
/// }
///```
#[inline]
pub fn pg_return_void() -> pg_sys::Datum {
    pg_sys::Datum::from(0)
}

/// Retrieve the `.flinfo.fn_extra` pointer (as a PgBox'd type) from [`pg_sys::FunctionCallInfo`].
///
/// This function is unsafe as we cannot guarantee the provided [`pg_sys::FunctionCallInfo`] pointer is valid
pub unsafe fn pg_func_extra<ReturnType, DefaultValue: FnOnce() -> ReturnType>(
    fcinfo: pg_sys::FunctionCallInfo,
    default: DefaultValue,
) -> PgBox<ReturnType> {
    let fcinfo = PgBox::from_pg(fcinfo);
    let mut flinfo = PgBox::from_pg(fcinfo.flinfo);
    if flinfo.fn_extra.is_null() {
        flinfo.fn_extra = PgMemoryContexts::For(flinfo.fn_mcxt).leak_and_drop_on_delete(default())
            as void_mut_ptr;
    }

    PgBox::from_pg(flinfo.fn_extra as *mut ReturnType)
}

/// This mimics the functionality of Postgres' `DirectFunctionCall` macros, allowing you to call
/// internal Postgres functions using its "V1" calling convention.  Unlike the Postgres C macros,
/// the function is allowed to return a NULL datum.
///
/// ## Safety
///
/// This function is unsafe as the underlying function being called is likely unsafe
///
/// ## Examples
/// ```rust
/// use std::ffi::CString;
///
/// use pgrx::prelude::*;
/// use pgrx::direct_function_call;
///
/// fn some_func() {
///     let result = unsafe {
///         direct_function_call::<pg_sys::Oid>(
///             pg_sys::regclassin,
///             &[CString::new("pg_class").unwrap().as_c_str().into_datum()]
///         )
///     };
///     let oid = result.expect("failed to lookup oid for pg_class");
///     assert_eq!(oid, pg_sys::RelationRelationId); // your value could be different, maybe
/// }
/// ```
pub unsafe fn direct_function_call<R: FromDatum>(
    func: unsafe fn(pg_sys::FunctionCallInfo) -> pg_sys::Datum,
    // TODO: this could take an iterator, but it would break turbofish :(
    args: &[Option<pg_sys::Datum>],
) -> Option<R> {
    direct_function_call_as_datum(func, args).and_then(|d| R::from_datum(d, false))
}

/// Akin to [direct_function_call], but specifically for calling those functions declared with the
/// `#[pg_extern]` attribute.
///
/// When using this, you'll want to suffix the function you want to call with `_wrapper`.
///
/// ## Example
///
/// ```rust,no_run
/// use pgrx::prelude::*;
/// use pgrx::direct_pg_extern_function_call;
///
/// #[pg_extern]
/// fn add_numbers(a: i32, b: i32) -> i32 {
///     a + b
/// }
///
/// /* NOTE:  _wrapper suffix! */
/// let result = unsafe {
///     direct_pg_extern_function_call::<i32>(add_numbers_wrapper, &[1_i32.into_datum(), 2_i32.into_datum()])
/// };
/// let sum = result.expect("add_numbers_wrapper returned NULL");
/// assert_eq!(3, sum);
/// ```
///
/// ## Safety
///
/// This function is unsafe as the function you're calling is also unsafe
pub unsafe fn direct_pg_extern_function_call<R: FromDatum>(
    func: unsafe extern "C-unwind" fn(pg_sys::FunctionCallInfo) -> pg_sys::Datum,
    args: &[Option<pg_sys::Datum>],
) -> Option<R> {
    direct_pg_extern_function_call_as_datum(func, args).and_then(|d| R::from_datum(d, false))
}

/// Same as [direct_function_call] but instead returns the direct `Option<pg_sys::Datum>` instead
/// of converting it to a value
///
/// ## Safety
///
/// This function is unsafe as the function you're calling is also unsafe
pub unsafe fn direct_function_call_as_datum(
    func: unsafe fn(pg_sys::FunctionCallInfo) -> pg_sys::Datum,
    args: &[Option<pg_sys::Datum>],
) -> Option<pg_sys::Datum> {
    direct_function_call_as_datum_internal(|fcinfo| func(fcinfo), args)
}

unsafe fn direct_function_call_as_datum_internal(
    func: impl FnOnce(pg_sys::FunctionCallInfo) -> pg_sys::Datum,
    args: &[Option<pg_sys::Datum>],
) -> Option<pg_sys::Datum> {
    let nargs: i16 = args.len().try_into().expect("too many args passed to function");
    let fcinfo = pg_sys::palloc0(
        std::mem::size_of::<pg_sys::FunctionCallInfoBaseData>()
            + std::mem::size_of::<pg_sys::NullableDatum>() * args.len(),
    )
    .cast::<pg_sys::FunctionCallInfoBaseData>();

    (*fcinfo).flinfo = std::ptr::null_mut();
    (*fcinfo).context = std::ptr::null_mut();
    (*fcinfo).resultinfo = std::ptr::null_mut();
    (*fcinfo).fncollation = pg_sys::InvalidOid;
    (*fcinfo).isnull = false;
    (*fcinfo).nargs = nargs;

    let args_ptr: *mut pg_sys::NullableDatum = ptr::addr_of_mut!((*fcinfo).args).cast();
    // This block is necessary for soundness. This way, we confine the slice's lifetime
    // to the bounds of mutating the arguments. We later will call a function on the fcinfo,
    // so we want all `&mut T` to be out-of-scope by the time we do that.
    {
        let arg_slice = slice::from_raw_parts_mut(args_ptr, args.len());
        for (n_datum, arg) in arg_slice.iter_mut().zip(args) {
            n_datum.isnull = arg.is_none();
            n_datum.value = arg.unwrap_or(pg_sys::Datum::from(0));
        }
    }

    let result = func(fcinfo);
    let result = if (*fcinfo).isnull { None } else { Some(result) };

    pg_sys::pfree(fcinfo.cast());
    result
}

/// Same as [direct_pg_extern_function_call] but instead returns the direct `Option<pg_sys::Datum>` instead
/// of converting it to a value
///
/// ## Safety
///
/// This function is unsafe as the function you're calling is also unsafe
pub unsafe fn direct_pg_extern_function_call_as_datum(
    func: unsafe extern "C-unwind" fn(pg_sys::FunctionCallInfo) -> pg_sys::Datum,
    args: &[Option<pg_sys::Datum>],
) -> Option<pg_sys::Datum> {
    direct_function_call_as_datum_internal(|fcinfo| pg_guard_ffi_boundary(|| func(fcinfo)), args)
}

#[inline]
pub unsafe fn srf_is_first_call(fcinfo: pg_sys::FunctionCallInfo) -> bool {
    (*(*fcinfo).flinfo).fn_extra.is_null()
}

#[inline]
#[deprecated(since = "0.12.0", note = "you want pg_sys::init_MultiFuncCall")]
pub unsafe fn srf_first_call_init(
    fcinfo: pg_sys::FunctionCallInfo,
) -> *mut pg_sys::FuncCallContext {
    pg_sys::init_MultiFuncCall(fcinfo)
}

#[inline]
#[deprecated(since = "0.12.0", note = "you want pg_sys::per_MultiFuncCall")]
pub unsafe fn srf_per_call_setup(fcinfo: pg_sys::FunctionCallInfo) -> *mut pg_sys::FuncCallContext {
    pg_sys::per_MultiFuncCall(fcinfo)
}

#[inline]
pub unsafe fn srf_return_next(
    fcinfo: pg_sys::FunctionCallInfo,
    funcctx: *mut pg_sys::FuncCallContext,
) {
    (*funcctx).call_cntr += 1;
    (*((*fcinfo).resultinfo as *mut pg_sys::ReturnSetInfo)).isDone =
        pg_sys::ExprDoneCond::ExprMultipleResult;
}

#[inline]
pub unsafe fn srf_return_done(
    fcinfo: pg_sys::FunctionCallInfo,
    funcctx: *mut pg_sys::FuncCallContext,
) {
    pg_sys::end_MultiFuncCall(fcinfo, funcctx);
    (*((*fcinfo).resultinfo as *mut pg_sys::ReturnSetInfo)).isDone =
        pg_sys::ExprDoneCond::ExprEndResult;
}