Create functions pg_set_relation_stats, pg_clear_relation_stats.

These functions are used to tweak statistics on any relation, provided
that the user has MAINTAIN privilege on the relation, or is the database
owner.

Bump catalog version.

Author: Corey Huinker
Discussion: https://postgr.es/m/CADkLM=eErgzn7ECDpwFcptJKOk9SxZEk5Pot4d94eVTZsvj3gw@mail.gmail.com

5 files changed, 319 insertions, 1 deletions

diff --git a/src/backend/catalog/system_functions.sql b/src/backend/catalog/system_functions.sql
index fd6b606ae90..92b6aefe7aa 100644
--- a/src/backend/catalog/system_functions.sql
+++ b/src/backend/catalog/system_functions.sql
@@ -639,6 +639,16 @@ LANGUAGE INTERNAL
 CALLED ON NULL INPUT VOLATILE PARALLEL SAFE
 AS 'pg_stat_reset_slru';
 
+CREATE OR REPLACE FUNCTION
+  pg_set_relation_stats(relation regclass,
+                        relpages integer DEFAULT NULL,
+                        reltuples real DEFAULT NULL,
+                        relallvisible integer DEFAULT NULL)
+RETURNS bool
+LANGUAGE INTERNAL
+CALLED ON NULL INPUT VOLATILE
+AS 'pg_set_relation_stats';
+
 --
 -- The default permissions for functions mean that anyone can execute them.
 -- A number of functions shouldn't be executable by just anyone, but rather
diff --git a/src/backend/statistics/Makefile b/src/backend/statistics/Makefile
index 89cf8c27973..041f5f8a581 100644
--- a/src/backend/statistics/Makefile
+++ b/src/backend/statistics/Makefile
@@ -16,6 +16,8 @@ OBJS = \
 	dependencies.o \
 	extended_stats.o \
 	mcv.o \
-	mvdistinct.o
+	mvdistinct.o \
+	relation_stats.o \
+	stat_utils.o
 
 include $(top_srcdir)/src/backend/common.mk
diff --git a/src/backend/statistics/meson.build b/src/backend/statistics/meson.build
index 73b29a3d50a..23648b3775f 100644
--- a/src/backend/statistics/meson.build
+++ b/src/backend/statistics/meson.build
@@ -5,4 +5,6 @@ backend_sources += files(
   'extended_stats.c',
   'mcv.c',
   'mvdistinct.c',
+  'relation_stats.c',
+  'stat_utils.c'
 )
diff --git a/src/backend/statistics/relation_stats.c b/src/backend/statistics/relation_stats.c
new file mode 100644
index 00000000000..ffa3d83a878
--- /dev/null
+++ b/src/backend/statistics/relation_stats.c
@@ -0,0 +1,210 @@
+/*-------------------------------------------------------------------------
+ * relation_stats.c
+ *
+ *	  PostgreSQL relation statistics manipulation
+ *
+ * Code supporting the direct import of relation statistics, similar to
+ * what is done by the ANALYZE command.
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *       src/backend/statistics/relation_stats.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/heapam.h"
+#include "catalog/indexing.h"
+#include "statistics/stat_utils.h"
+#include "utils/fmgrprotos.h"
+#include "utils/syscache.h"
+
+#define DEFAULT_RELPAGES Int32GetDatum(0)
+#define DEFAULT_RELTUPLES Float4GetDatum(-1.0)
+#define DEFAULT_RELALLVISIBLE Int32GetDatum(0)
+
+/*
+ * Positional argument numbers, names, and types for
+ * relation_statistics_update().
+ */
+
+enum relation_stats_argnum
+{
+	RELATION_ARG = 0,
+	RELPAGES_ARG,
+	RELTUPLES_ARG,
+	RELALLVISIBLE_ARG,
+	NUM_RELATION_STATS_ARGS
+};
+
+static struct StatsArgInfo relarginfo[] =
+{
+	[RELATION_ARG] = {"relation", REGCLASSOID},
+	[RELPAGES_ARG] = {"relpages", INT4OID},
+	[RELTUPLES_ARG] = {"reltuples", FLOAT4OID},
+	[RELALLVISIBLE_ARG] = {"relallvisible", INT4OID},
+	[NUM_RELATION_STATS_ARGS] = {0}
+};
+
+static bool relation_statistics_update(FunctionCallInfo fcinfo, int elevel);
+
+/*
+ * Internal function for modifying statistics for a relation.
+ */
+static bool
+relation_statistics_update(FunctionCallInfo fcinfo, int elevel)
+{
+	Oid			reloid;
+	Relation	crel;
+	HeapTuple	ctup;
+	Form_pg_class pgcform;
+	int			replaces[3] = {0};
+	Datum		values[3] = {0};
+	bool		nulls[3] = {0};
+	int			ncols = 0;
+	TupleDesc	tupdesc;
+	HeapTuple	newtup;
+
+
+	stats_check_required_arg(fcinfo, relarginfo, RELATION_ARG);
+	reloid = PG_GETARG_OID(RELATION_ARG);
+
+	stats_lock_check_privileges(reloid);
+
+	/*
+	 * Take RowExclusiveLock on pg_class, consistent with
+	 * vac_update_relstats().
+	 */
+	crel = table_open(RelationRelationId, RowExclusiveLock);
+
+	tupdesc = RelationGetDescr(crel);
+	ctup = SearchSysCacheCopy1(RELOID, ObjectIdGetDatum(reloid));
+	if (!HeapTupleIsValid(ctup))
+	{
+		ereport(elevel,
+				(errcode(ERRCODE_OBJECT_IN_USE),
+				 errmsg("pg_class entry for relid %u not found", reloid)));
+		table_close(crel, RowExclusiveLock);
+		return false;
+	}
+
+	pgcform = (Form_pg_class) GETSTRUCT(ctup);
+
+	/* relpages */
+	if (!PG_ARGISNULL(RELPAGES_ARG))
+	{
+		int32		relpages = PG_GETARG_INT32(RELPAGES_ARG);
+
+		if (relpages < -1)
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("relpages cannot be < -1")));
+			table_close(crel, RowExclusiveLock);
+			return false;
+		}
+
+		if (relpages != pgcform->relpages)
+		{
+			replaces[ncols] = Anum_pg_class_relpages;
+			values[ncols] = Int32GetDatum(relpages);
+			ncols++;
+		}
+	}
+
+	if (!PG_ARGISNULL(RELTUPLES_ARG))
+	{
+		float		reltuples = PG_GETARG_FLOAT4(RELTUPLES_ARG);
+
+		if (reltuples < -1.0)
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("reltuples cannot be < -1.0")));
+			table_close(crel, RowExclusiveLock);
+			return false;
+		}
+
+		if (reltuples != pgcform->reltuples)
+		{
+			replaces[ncols] = Anum_pg_class_reltuples;
+			values[ncols] = Float4GetDatum(reltuples);
+			ncols++;
+		}
+	}
+
+	if (!PG_ARGISNULL(RELALLVISIBLE_ARG))
+	{
+		int32		relallvisible = PG_GETARG_INT32(RELALLVISIBLE_ARG);
+
+		if (relallvisible < 0)
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("relallvisible cannot be < 0")));
+			table_close(crel, RowExclusiveLock);
+			return false;
+		}
+
+		if (relallvisible != pgcform->relallvisible)
+		{
+			replaces[ncols] = Anum_pg_class_relallvisible;
+			values[ncols] = Int32GetDatum(relallvisible);
+			ncols++;
+		}
+	}
+
+	/* only update pg_class if there is a meaningful change */
+	if (ncols == 0)
+	{
+		table_close(crel, RowExclusiveLock);
+		return false;
+	}
+
+	newtup = heap_modify_tuple_by_cols(ctup, tupdesc, ncols, replaces, values,
+									   nulls);
+
+	CatalogTupleUpdate(crel, &newtup->t_self, newtup);
+	heap_freetuple(newtup);
+
+	/* release the lock, consistent with vac_update_relstats() */
+	table_close(crel, RowExclusiveLock);
+
+	return true;
+}
+
+/*
+ * Set statistics for a given pg_class entry.
+ */
+Datum
+pg_set_relation_stats(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(relation_statistics_update(fcinfo, ERROR));
+}
+
+/*
+ * Clear statistics for a given pg_class entry; that is, set back to initial
+ * stats for a newly-created table.
+ */
+Datum
+pg_clear_relation_stats(PG_FUNCTION_ARGS)
+{
+	LOCAL_FCINFO(newfcinfo, 4);
+
+	InitFunctionCallInfoData(*newfcinfo, NULL, 4, InvalidOid, NULL, NULL);
+
+	newfcinfo->args[0].value = PG_GETARG_OID(0);
+	newfcinfo->args[0].isnull = PG_ARGISNULL(0);
+	newfcinfo->args[1].value = DEFAULT_RELPAGES;
+	newfcinfo->args[1].isnull = false;
+	newfcinfo->args[2].value = DEFAULT_RELTUPLES;
+	newfcinfo->args[2].isnull = false;
+	newfcinfo->args[3].value = DEFAULT_RELALLVISIBLE;
+	newfcinfo->args[3].isnull = false;
+
+	PG_RETURN_BOOL(relation_statistics_update(newfcinfo, ERROR));
+}
diff --git a/src/backend/statistics/stat_utils.c b/src/backend/statistics/stat_utils.c
new file mode 100644
index 00000000000..4babed2e5d7
--- /dev/null
+++ b/src/backend/statistics/stat_utils.c
@@ -0,0 +1,94 @@
+/*-------------------------------------------------------------------------
+ * stat_utils.c
+ *
+ *	  PostgreSQL statistics manipulation utilities.
+ *
+ * Code supporting the direct manipulation of statistics.
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *       src/backend/statistics/stat_utils.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/relation.h"
+#include "catalog/pg_database.h"
+#include "miscadmin.h"
+#include "statistics/stat_utils.h"
+#include "utils/array.h"
+#include "utils/builtins.h"
+#include "utils/acl.h"
+#include "utils/rel.h"
+
+/*
+ * Ensure that a given argument is not null.
+ */
+void
+stats_check_required_arg(FunctionCallInfo fcinfo,
+						 struct StatsArgInfo *arginfo,
+						 int argnum)
+{
+	if (PG_ARGISNULL(argnum))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("\"%s\" cannot be NULL",
+						arginfo[argnum].argname)));
+}
+
+/*
+ * Lock relation in ShareUpdateExclusive mode, check privileges, and close the
+ * relation (but retain the lock).
+ *
+ * A role has privileges to set statistics on the relation if any of the
+ * following are true:
+ *   - the role owns the current database and the relation is not shared
+ *   - the role has the MAINTAIN privilege on the relation
+ */
+void
+stats_lock_check_privileges(Oid reloid)
+{
+	Relation	rel = relation_open(reloid, ShareUpdateExclusiveLock);
+	const char	relkind = rel->rd_rel->relkind;
+
+	/* All of the types that can be used with ANALYZE, plus indexes */
+	switch (relkind)
+	{
+		case RELKIND_RELATION:
+		case RELKIND_INDEX:
+		case RELKIND_MATVIEW:
+		case RELKIND_FOREIGN_TABLE:
+		case RELKIND_PARTITIONED_TABLE:
+		case RELKIND_PARTITIONED_INDEX:
+			break;
+		default:
+			ereport(ERROR,
+					(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+					 errmsg("cannot modify statistics for relation \"%s\"",
+							RelationGetRelationName(rel)),
+					 errdetail_relkind_not_supported(rel->rd_rel->relkind)));
+	}
+
+	if (rel->rd_rel->relisshared)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot modify statistics for shared relation")));
+
+	if (!object_ownercheck(DatabaseRelationId, MyDatabaseId, GetUserId()))
+	{
+		AclResult	aclresult = pg_class_aclcheck(RelationGetRelid(rel),
+												  GetUserId(),
+												  ACL_MAINTAIN);
+
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult,
+						   get_relkind_objtype(rel->rd_rel->relkind),
+						   NameStr(rel->rd_rel->relname));
+	}
+
+	relation_close(rel, NoLock);
+}