partitionfuncs.c

Go to the documentation of this file.

/*-------------------------------------------------------------------------
 *
 * partitionfuncs.c
 *    Functions for accessing partition-related metadata
 *
 * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 *
 * IDENTIFICATION
 *    src/backend/utils/adt/partitionfuncs.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/htup_details.h"
#include "catalog/partition.h"
#include "catalog/pg_class.h"
#include "catalog/pg_inherits.h"
#include "catalog/pg_type.h"
#include "funcapi.h"
#include "utils/fmgrprotos.h"
#include "utils/lsyscache.h"
#include "utils/syscache.h"

/*
 * Checks if a given relation can be part of a partition tree.  Returns
 * false if the relation cannot be processed, in which case it is up to
 * the caller to decide what to do, by either raising an error or doing
 * something else.
 */
static bool
check_rel_can_be_partition(Oid relid)
{
    char        relkind;
    bool        relispartition;

    /* Check if relation exists */
    if (!SearchSysCacheExists1(RELOID, ObjectIdGetDatum(relid)))
        return false;

    relkind = get_rel_relkind(relid);
    relispartition = get_rel_relispartition(relid);

    /* Only allow relation types that can appear in partition trees. */
    if (!relispartition &&
        relkind != RELKIND_PARTITIONED_TABLE &&
        relkind != RELKIND_PARTITIONED_INDEX)
        return false;

    return true;
}

/*
 * pg_partition_tree
 *
 * Produce a view with one row per member of a partition tree, beginning
 * from the top-most parent given by the caller.  This gives information
 * about each partition, its immediate partitioned parent, if it is
 * a leaf partition and its level in the hierarchy.
 */
Datum
pg_partition_tree(PG_FUNCTION_ARGS)
{
#define PG_PARTITION_TREE_COLS  4
    Oid         rootrelid = PG_GETARG_OID(0);
    FuncCallContext *funcctx;
    List       *partitions;

    /* stuff done only on the first call of the function */
    if (SRF_IS_FIRSTCALL())
    {
        MemoryContext oldcxt;
        TupleDesc   tupdesc;

        /* create a function context for cross-call persistence */
        funcctx = SRF_FIRSTCALL_INIT();

        if (!check_rel_can_be_partition(rootrelid))
            SRF_RETURN_DONE(funcctx);

        /* switch to memory context appropriate for multiple function calls */
        oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

        /*
         * Find all members of inheritance set.  We only need AccessShareLock
         * on the children for the partition information lookup.
         */
        partitions = find_all_inheritors(rootrelid, AccessShareLock, NULL);

        tupdesc = CreateTemplateTupleDesc(PG_PARTITION_TREE_COLS);
        TupleDescInitEntry(tupdesc, (AttrNumber) 1, "relid",
                           REGCLASSOID, -1, 0);
        TupleDescInitEntry(tupdesc, (AttrNumber) 2, "parentid",
                           REGCLASSOID, -1, 0);
        TupleDescInitEntry(tupdesc, (AttrNumber) 3, "isleaf",
                           BOOLOID, -1, 0);
        TupleDescInitEntry(tupdesc, (AttrNumber) 4, "level",
                           INT4OID, -1, 0);

        funcctx->tuple_desc = BlessTupleDesc(tupdesc);

        /* The only state we need is the partition list */
        funcctx->user_fctx = (void *) partitions;

        MemoryContextSwitchTo(oldcxt);
    }

    /* stuff done on every call of the function */
    funcctx = SRF_PERCALL_SETUP();
    partitions = (List *) funcctx->user_fctx;

    if (funcctx->call_cntr < list_length(partitions))
    {
        Datum       result;
        Datum       values[PG_PARTITION_TREE_COLS];
        bool        nulls[PG_PARTITION_TREE_COLS];
        HeapTuple   tuple;
        Oid         parentid = InvalidOid;
        Oid         relid = list_nth_oid(partitions, funcctx->call_cntr);
        char        relkind = get_rel_relkind(relid);
        int         level = 0;
        List       *ancestors = get_partition_ancestors(relid);
        ListCell   *lc;

        /*
         * Form tuple with appropriate data.
         */
        MemSet(nulls, 0, sizeof(nulls));
        MemSet(values, 0, sizeof(values));

        /* relid */
        values[0] = ObjectIdGetDatum(relid);

        /* parentid */
        if (ancestors != NIL)
            parentid = linitial_oid(ancestors);
        if (OidIsValid(parentid))
            values[1] = ObjectIdGetDatum(parentid);
        else
            nulls[1] = true;

        /* isleaf */
        values[2] = BoolGetDatum(relkind != RELKIND_PARTITIONED_TABLE &&
                                 relkind != RELKIND_PARTITIONED_INDEX);

        /* level */
        if (relid != rootrelid)
        {
            foreach(lc, ancestors)
            {
                level++;
                if (lfirst_oid(lc) == rootrelid)
                    break;
            }
        }
        values[3] = Int32GetDatum(level);

        tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
        result = HeapTupleGetDatum(tuple);
        SRF_RETURN_NEXT(funcctx, result);
    }

    /* done when there are no more elements left */
    SRF_RETURN_DONE(funcctx);
}

/*
 * pg_partition_root
 *
 * Returns the top-most parent of the partition tree to which a given
 * relation belongs, or NULL if it's not (or cannot be) part of any
 * partition tree.
 */
Datum
pg_partition_root(PG_FUNCTION_ARGS)
{
    Oid         relid = PG_GETARG_OID(0);
    Oid         rootrelid;
    List       *ancestors;

    if (!check_rel_can_be_partition(relid))
        PG_RETURN_NULL();

    /* fetch the list of ancestors */
    ancestors = get_partition_ancestors(relid);

    /*
     * If the input relation is already the top-most parent, just return
     * itself.
     */
    if (ancestors == NIL)
        PG_RETURN_OID(relid);

    rootrelid = llast_oid(ancestors);
    list_free(ancestors);

    /*
     * "rootrelid" must contain a valid OID, given that the input relation is
     * a valid partition tree member as checked above.
     */
    Assert(OidIsValid(rootrelid));
    PG_RETURN_OID(rootrelid);
}

/*
 * pg_partition_ancestors
 *
 * Produces a view with one row per ancestor of the given partition,
 * including the input relation itself.
 */
Datum
pg_partition_ancestors(PG_FUNCTION_ARGS)
{
    Oid         relid = PG_GETARG_OID(0);
    FuncCallContext *funcctx;
    List       *ancestors;

    if (SRF_IS_FIRSTCALL())
    {
        MemoryContext oldcxt;

        funcctx = SRF_FIRSTCALL_INIT();

        if (!check_rel_can_be_partition(relid))
            SRF_RETURN_DONE(funcctx);

        oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

        ancestors = get_partition_ancestors(relid);
        ancestors = lcons_oid(relid, ancestors);

        /* The only state we need is the ancestors list */
        funcctx->user_fctx = (void *) ancestors;

        MemoryContextSwitchTo(oldcxt);
    }

    funcctx = SRF_PERCALL_SETUP();
    ancestors = (List *) funcctx->user_fctx;

    if (funcctx->call_cntr < list_length(ancestors))
    {
        Oid         relid = list_nth_oid(ancestors, funcctx->call_cntr);

        SRF_RETURN_NEXT(funcctx, ObjectIdGetDatum(relid));
    }

    SRF_RETURN_DONE(funcctx);
}