/*-------------------------------------------------------------------------
 *
 * partcache.c
 *  Support routines for manipulating partition information cached in
 *  relcache
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * IDENTIFICATION
 *    src/backend/utils/cache/partcache.c
 *
 *-------------------------------------------------------------------------
*/
#include "postgres.h"

#include "access/hash.h"
#include "access/htup_details.h"
#include "access/nbtree.h"
#include "access/relation.h"
#include "catalog/partition.h"
#include "catalog/pg_opclass.h"
#include "catalog/pg_partitioned_table.h"
#include "miscadmin.h"
#include "nodes/makefuncs.h"
#include "nodes/nodeFuncs.h"
#include "optimizer/optimizer.h"
#include "partitioning/partbounds.h"
#include "utils/builtins.h"
#include "utils/lsyscache.h"
#include "utils/memutils.h"
#include "utils/partcache.h"
#include "utils/rel.h"
#include "utils/syscache.h"


static void RelationBuildPartitionKey(Relation relation);
static List *generate_partition_qual(Relation rel);

/*
 * RelationGetPartitionKey -- get partition key, if relation is partitioned
 *
 * Note: partition keys are not allowed to change after the partitioned rel
 * is created.  RelationClearRelation knows this and preserves rd_partkey
 * across relcache rebuilds, as long as the relation is open.  Therefore,
 * even though we hand back a direct pointer into the relcache entry, it's
 * safe for callers to continue to use that pointer as long as they hold
 * the relation open.
 */
PartitionKey
RelationGetPartitionKey(Relation rel)
{
 if (rel->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
  return NULL;

 if (unlikely(rel->rd_partkey == NULL))
  RelationBuildPartitionKey(rel);

 return rel->rd_partkey;
}

/*
 * RelationBuildPartitionKey
 *  Build partition key data of relation, and attach to relcache
 *
 * Partitioning key data is a complex structure; to avoid complicated logic to
 * free individual elements whenever the relcache entry is flushed, we give it
 * its own memory context, a child of CacheMemoryContext, which can easily be
 * deleted on its own.  To avoid leaking memory in that context in case of an
 * error partway through this function, the context is initially created as a
 * child of CurTransactionContext and only re-parented to CacheMemoryContext
 * at the end, when no further errors are possible.  Also, we don't make this
 * context the current context except in very brief code sections, out of fear
 * that some of our callees allocate memory on their own which would be leaked
 * permanently.
 */
static void
RelationBuildPartitionKey(Relation relation)
{
 Form_pg_partitioned_table form;
 HeapTuple tuple;
 bool  isnull;
 int   i;
 PartitionKey key;
 AttrNumber *attrs;
 oidvector  *opclass;
 oidvector  *collation;
 ListCell   *partexprs_item;
 Datum  datum;
 MemoryContext partkeycxt,
    oldcxt;
 int16  procnum;

 tuple = SearchSysCache1(PARTRELID,
       ObjectIdGetDatum(RelationGetRelid(relation)));

 if (!HeapTupleIsValid(tuple))
  elog(ERROR, "cache lookup failed for partition key of relation %u",
    RelationGetRelid(relation));

 partkeycxt = AllocSetContextCreate(CurTransactionContext,
            "partition key",
            ALLOCSET_SMALL_SIZES);
 MemoryContextCopyAndSetIdentifier(partkeycxt,
           RelationGetRelationName(relation));

 key = (PartitionKey) MemoryContextAllocZero(partkeycxt,
            sizeof(PartitionKeyData));

 /* Fixed-length attributes */
 form = (Form_pg_partitioned_table) GETSTRUCT(tuple);
 key->strategy = form->partstrat;
 key->partnatts = form->partnatts;

 /* Validate partition strategy code */
 if (key->strategy != PARTITION_STRATEGY_LIST &&
  key->strategy != PARTITION_STRATEGY_RANGE &&
  key->strategy != PARTITION_STRATEGY_HASH)
  elog(ERROR, "invalid partition strategy \"%c\"", key->strategy);

 /*
  * We can rely on the first variable-length attribute being mapped to the
  * relevant field of the catalog's C struct, because all previous
  * attributes are non-nullable and fixed-length.
 */
 attrs = form->partattrs.values;

 /* But use the hard way to retrieve further variable-length attributes */
 /* Operator class */
 datum = SysCacheGetAttrNotNull(PARTRELID, tuple,
           Anum_pg_partitioned_table_partclass);
 opclass = (oidvector *) DatumGetPointer(datum);

 /* Collation */
 datum = SysCacheGetAttrNotNull(PARTRELID, tuple,
           Anum_pg_partitioned_table_partcollation);
 collation = (oidvector *) DatumGetPointer(datum);

 /* Expressions */
 datum = SysCacheGetAttr(PARTRELID, tuple,
       Anum_pg_partitioned_table_partexprs, &isnull);
 if (!isnull)
 {
  char    *exprString;
  Node    *expr;

  exprString = TextDatumGetCString(datum);
  expr = stringToNode(exprString);
  pfree(exprString);

  /*
   * Run the expressions through const-simplification since the planner
   * will be comparing them to similarly-processed qual clause operands,
   * and may fail to detect valid matches without this step; fix
   * opfuncids while at it.  We don't need to bother with
   * canonicalize_qual() though, because partition expressions should be
   * in canonical form already (ie, no need for OR-merging or constant
   * elimination).
 */
  expr = eval_const_expressions(NULL, expr);
  fix_opfuncids(expr);

  oldcxt = MemoryContextSwitchTo(partkeycxt);
  key->partexprs = (List *) copyObject(expr);
  MemoryContextSwitchTo(oldcxt);
 }

 /* Allocate assorted arrays in the partkeycxt, which we'll fill below */
 oldcxt = MemoryContextSwitchTo(partkeycxt);
 key->partattrs = (AttrNumber *) palloc0(key->partnatts * sizeof(AttrNumber));
 key->partopfamily = (Oid *) palloc0(key->partnatts * sizeof(Oid));
 key->partopcintype = (Oid *) palloc0(key->partnatts * sizeof(Oid));
 key->partsupfunc = (FmgrInfo *) palloc0(key->partnatts * sizeof(FmgrInfo));

 key->partcollation = (Oid *) palloc0(key->partnatts * sizeof(Oid));
 key->parttypid = (Oid *) palloc0(key->partnatts * sizeof(Oid));
 key->parttypmod = (int32 *) palloc0(key->partnatts * sizeof(int32));
 key->parttyplen = (int16 *) palloc0(key->partnatts * sizeof(int16));
 key->parttypbyval = (bool *) palloc0(key->partnatts * sizeof(bool));
 key->parttypalign = (char *) palloc0(key->partnatts * sizeof(char));
 key->parttypcoll = (Oid *) palloc0(key->partnatts * sizeof(Oid));
 MemoryContextSwitchTo(oldcxt);

 /* determine support function number to search for */
 procnum = (key->strategy == PARTITION_STRATEGY_HASH) ?
  HASHEXTENDED_PROC : BTORDER_PROC;

 /* Copy partattrs and fill other per-attribute info */
 memcpy(key->partattrs, attrs, key->partnatts * sizeof(int16));
 partexprs_item = list_head(key->partexprs);
 for (i = 0; i < key->partnatts; i++)
 {
  AttrNumber attno = key->partattrs[i];
  HeapTuple opclasstup;
  Form_pg_opclass opclassform;
  Oid   funcid;

  /* Collect opfamily information */
  opclasstup = SearchSysCache1(CLAOID,
          ObjectIdGetDatum(opclass->values[i]));
  if (!HeapTupleIsValid(opclasstup))
   elog(ERROR, "cache lookup failed for opclass %u", opclass->values[i]);

  opclassform = (Form_pg_opclass) GETSTRUCT(opclasstup);
  key->partopfamily[i] = opclassform->opcfamily;
  key->partopcintype[i] = opclassform->opcintype;

  /* Get a support function for the specified opfamily and datatypes */
  funcid = get_opfamily_proc(opclassform->opcfamily,
           opclassform->opcintype,
           opclassform->opcintype,
           procnum);
  if (!OidIsValid(funcid))
   ereport(ERROR,
     (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
      errmsg("operator class \"%s\" of access method %s is missing support function %d for type %s",
       NameStr(opclassform->opcname),
       (key->strategy == PARTITION_STRATEGY_HASH) ?
       "hash" : "btree",
       procnum,
       format_type_be(opclassform->opcintype))));

  fmgr_info_cxt(funcid, &key->partsupfunc[i], partkeycxt);

  /* Collation */
  key->partcollation[i] = collation->values[i];

  /* Collect type information */
  if (attno != 0)
  {
   Form_pg_attribute att = TupleDescAttr(relation->rd_att, attno - 1);

   key->parttypid[i] = att->atttypid;
   key->parttypmod[i] = att->atttypmod;
   key->parttypcoll[i] = att->attcollation;
  }
  else
  {
   if (partexprs_item == NULL)
    elog(ERROR, "wrong number of partition key expressions");

   key->parttypid[i] = exprType(lfirst(partexprs_item));
   key->parttypmod[i] = exprTypmod(lfirst(partexprs_item));
   key->parttypcoll[i] = exprCollation(lfirst(partexprs_item));

   partexprs_item = lnext(key->partexprs, partexprs_item);
  }
  get_typlenbyvalalign(key->parttypid[i],
        &key->parttyplen[i],
        &key->parttypbyval[i],
        &key->parttypalign[i]);

  ReleaseSysCache(opclasstup);
 }

 ReleaseSysCache(tuple);

 /* Assert that we're not leaking any old data during assignments below */
 Assert(relation->rd_partkeycxt == NULL);
 Assert(relation->rd_partkey == NULL);

 /*
  * Success --- reparent our context and make the relcache point to the
  * newly constructed key
 */
 MemoryContextSetParent(partkeycxt, CacheMemoryContext);
 relation->rd_partkeycxt = partkeycxt;
 relation->rd_partkey = key;
}

/*
 * RelationGetPartitionQual
 *
 * Returns a list of partition quals
 */
List *
RelationGetPartitionQual(Relation rel)
{
 /* Quick exit */
 if (!rel->rd_rel->relispartition)
  return NIL;

 return generate_partition_qual(rel);
}

/*
 * get_partition_qual_relid
 *
 * Returns an expression tree describing the passed-in relation's partition
 * constraint.
 *
 * If the relation is not found, or is not a partition, or there is no
 * partition constraint, return NULL.  We must guard against the first two
 * cases because this supports a SQL function that could be passed any OID.
 * The last case can happen even if relispartition is true, when a default
 * partition is the only partition.
 */
Expr *
get_partition_qual_relid(Oid relid)
{
 Expr    *result = NULL;

 /* Do the work only if this relation exists and is a partition. */
 if (get_rel_relispartition(relid))
 {
  Relation rel = relation_open(relid, AccessShareLock);
  List    *and_args;

  and_args = generate_partition_qual(rel);

  /* Convert implicit-AND list format to boolean expression */
  if (and_args == NIL)
   result = NULL;
  else if (list_length(and_args) > 1)
   result = makeBoolExpr(AND_EXPR, and_args, -1);
  else
   result = linitial(and_args);

  /* Keep the lock, to allow safe deparsing against the rel by caller. */
  relation_close(rel, NoLock);
 }

 return result;
}

/*
 * generate_partition_qual
 *
 * Generate partition predicate from rel's partition bound expression. The
 * function returns a NIL list if there is no predicate.
 *
 * We cache a copy of the result in the relcache entry, after constructing
 * it using the caller's context.  This approach avoids leaking any data
 * into long-lived cache contexts, especially if we fail partway through.
 */
static List *
generate_partition_qual(Relation rel)
{
 HeapTuple tuple;
 MemoryContext oldcxt;
 Datum  boundDatum;
 bool  isnull;
 List    *my_qual = NIL,
      *result = NIL;
 Oid   parentrelid;
 Relation parent;

 /* Guard against stack overflow due to overly deep partition tree */
 check_stack_depth();

 /* If we already cached the result, just return a copy */
 if (rel->rd_partcheckvalid)
  return copyObject(rel->rd_partcheck);

 /*
  * Grab at least an AccessShareLock on the parent table.  Must do this
  * even if the partition has been partially detached, because transactions
  * concurrent with the detach might still be trying to use a partition
  * descriptor that includes it.
 */
 parentrelid = get_partition_parent(RelationGetRelid(rel), true);
 parent = relation_open(parentrelid, AccessShareLock);

 /* Get pg_class.relpartbound */
 tuple = SearchSysCache1(RELOID,
       ObjectIdGetDatum(RelationGetRelid(rel)));
 if (!HeapTupleIsValid(tuple))
  elog(ERROR, "cache lookup failed for relation %u",
    RelationGetRelid(rel));

 boundDatum = SysCacheGetAttr(RELOID, tuple,
         Anum_pg_class_relpartbound,
         &isnull);
 if (!isnull)
 {
  PartitionBoundSpec *bound;

  bound = castNode(PartitionBoundSpec,
       stringToNode(TextDatumGetCString(boundDatum)));

  my_qual = get_qual_from_partbound(parent, bound);
 }

 ReleaseSysCache(tuple);

 /* Add the parent's quals to the list (if any) */
 if (parent->rd_rel->relispartition)
  result = list_concat(generate_partition_qual(parent), my_qual);
 else
  result = my_qual;

 /*
  * Change Vars to have partition's attnos instead of the parent's. We do
  * this after we concatenate the parent's quals, because we want every Var
  * in it to bear this relation's attnos. It's safe to assume varno = 1
  * here.
 */
 result = map_partition_varattnos(result, 1, rel, parent);

 /* Assert that we're not leaking any old data during assignments below */
 Assert(rel->rd_partcheckcxt == NULL);
 Assert(rel->rd_partcheck == NIL);

 /*
  * Save a copy in the relcache.  The order of these operations is fairly
  * critical to avoid memory leaks and ensure that we don't leave a corrupt
  * relcache entry if we fail partway through copyObject.
  *
  * If, as is definitely possible, the partcheck list is NIL, then we do
  * not need to make a context to hold it.
 */
 if (result != NIL)
 {
  rel->rd_partcheckcxt = AllocSetContextCreate(CacheMemoryContext,
              "partition constraint",
              ALLOCSET_SMALL_SIZES);
  MemoryContextCopyAndSetIdentifier(rel->rd_partcheckcxt,
            RelationGetRelationName(rel));
  oldcxt = MemoryContextSwitchTo(rel->rd_partcheckcxt);
  rel->rd_partcheck = copyObject(result);
  MemoryContextSwitchTo(oldcxt);
 }
 else
  rel->rd_partcheck = NIL;
 rel->rd_partcheckvalid = true;

 /* Keep the parent locked until commit */
 relation_close(parent, NoLock);

 /* Return the working copy to the caller */
 return result;
}