From 93bb8e2dfeab17d5291273e6343f61e40e501633 Mon Sep 17 00:00:00 2001
From: Corey Huinker <corey.huinker@gmail.com>
Date: Wed, 24 Jul 2024 23:45:26 -0400
Subject: [PATCH v27j 1/2] Create function pg_set_relation_stats.

This function is used to tweak statistics on any relation that the user
owns.

The first parameter, relation, is used to identify the the relation to be
modified.

The remaining parameters correspond to the statistics attributes in
pg_class: relpages, reltuples, and relallisvible.

This function allows the user to tweak pg_class statistics values,
allowing the user to inflate rowcounts, table size, and visibility to see
what effect those changes will have on the the query planner.

The function has no return value.

Author: Corey Huinker
Discussion: https://postgr.es/m/CADkLM=e0jM7m0GFV44nNtvL22CtnFUu+pekppCVE=DOBe58RTQ@mail.gmail.com
---
 doc/src/sgml/func.sgml                     |  65 +++++++
 src/backend/statistics/Makefile            |   1 +
 src/backend/statistics/import_stats.c      | 207 +++++++++++++++++++++
 src/backend/statistics/meson.build         |   1 +
 src/include/catalog/pg_proc.dat            |   9 +
 src/test/regress/expected/stats_import.out |  99 ++++++++++
 src/test/regress/parallel_schedule         |   2 +-
 src/test/regress/sql/stats_import.sql      |  79 ++++++++
 8 files changed, 462 insertions(+), 1 deletion(-)
 create mode 100644 src/backend/statistics/import_stats.c
 create mode 100644 src/test/regress/expected/stats_import.out
 create mode 100644 src/test/regress/sql/stats_import.sql

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 5dd95d73a1a..02d93f6c925 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -29938,6 +29938,71 @@ DETAIL:  Make sure pg_wal_replay_wait() isn't called within a transaction with a
     in identifying the specific disk files associated with database objects.
    </para>
 
+   <table id="functions-admin-statsimport">
+    <title>Database Object Statistics Import Functions</title>
+    <tgroup cols="1">
+     <thead>
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        Function
+       </para>
+       <para>
+        Description
+       </para></entry>
+      </row>
+     </thead>
+
+     <tbody>
+      <row>
+       <entry role="func_table_entry">
+        <para role="func_signature">
+         <indexterm>
+          <primary>pg_set_relation_stats</primary>
+         </indexterm>
+         <function>pg_set_relation_stats</function> ( <parameter>relation</parameter> <type>regclass</type>, <parameter>relpages</parameter> <type>integer</type>, <parameter>reltuples</parameter> <type>real</type>, <parameter>relallvisible</parameter> <type>integer</type> )
+         <returnvalue>void</returnvalue>
+        </para>
+        <para>
+         Updates table-level statistics for the given relation to the
+         specified values. The parameters correspond to columns in <link
+         linkend="catalog-pg-class"><structname>pg_class</structname></link>.
+        </para>
+        <para>
+         Ordinarily, these statistics are collected automatically or updated
+         as a part of <xref linkend="sql-vacuum"/> or <xref
+         linkend="sql-analyze"/>, so it's not necessary to call this
+         function. However, it may be useful when testing the effects of
+         statistics on the planner to understand or anticipate plan changes.
+        </para>
+        <para>
+         The caller must have the <literal>MAINTAIN</literal> privilege on
+         the table or be the owner of the database.
+        </para>
+        <para>
+         The value of <structfield>relpages</structfield> must be greater than
+         or equal to <literal>0</literal>,
+         <structfield>reltuples</structfield> must be greater than or equal to
+         <literal>-1.0</literal>, and <structfield>relallvisible</structfield>
+         must be greater than or equal to <literal>0</literal>.
+        </para>
+        <warning>
+         <para>
+          Changes made by this function are likely to be overwritten by <link
+          linkend="autovacuum">autovacuum</link> (or manual
+          <command>VACUUM</command> or <command>ANALYZE</command>) and should
+          be considered temporary.
+         </para>
+         <para>
+          The signature of this function may change in new major releases if
+          there are changes to the statistics being tracked.
+         </para>
+        </warning>
+       </entry>
+      </row>
+     </tbody>
+    </tgroup>
+   </table>
+
    <table id="functions-admin-dblocation">
     <title>Database Object Location Functions</title>
     <tgroup cols="1">
diff --git a/src/backend/statistics/Makefile b/src/backend/statistics/Makefile
index 89cf8c27973..5e776c02218 100644
--- a/src/backend/statistics/Makefile
+++ b/src/backend/statistics/Makefile
@@ -15,6 +15,7 @@ include $(top_builddir)/src/Makefile.global
 OBJS = \
 	dependencies.o \
 	extended_stats.o \
+	import_stats.o \
 	mcv.o \
 	mvdistinct.o
 
diff --git a/src/backend/statistics/import_stats.c b/src/backend/statistics/import_stats.c
new file mode 100644
index 00000000000..fede53f6634
--- /dev/null
+++ b/src/backend/statistics/import_stats.c
@@ -0,0 +1,207 @@
+/*-------------------------------------------------------------------------
+ * import_stats.c
+ *
+ *	  PostgreSQL statistics import
+ *
+ * Code supporting the direct importation 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/import_stats.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/heapam.h"
+#include "catalog/indexing.h"
+#include "catalog/pg_database.h"
+#include "miscadmin.h"
+#include "utils/acl.h"
+#include "utils/builtins.h"
+#include "utils/syscache.h"
+
+static bool relation_statistics_update(Oid reloid, int version,
+									   int32 relpages, float4 reltuples,
+									   int32 relallvisible, int elevel);
+static void check_privileges(Relation rel);
+
+/*
+ * Internal function for modifying statistics for a relation.
+ */
+static bool
+relation_statistics_update(Oid reloid, int version, int32 relpages,
+						   float4 reltuples, int32 relallvisible, int elevel)
+{
+	Relation		relation;
+	Relation		crel;
+	HeapTuple		ctup;
+	Form_pg_class	pgcform;
+
+	if (relpages < -1)
+	{
+		ereport(elevel,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("relpages cannot be < -1")));
+		return false;
+	}
+
+	if (reltuples < -1.0)
+	{
+		ereport(elevel,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("reltuples cannot be < -1.0")));
+		return false;
+	}
+
+	if (relallvisible < -1)
+	{
+		ereport(elevel,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("relallvisible cannot be < -1")));
+		return false;
+	}
+
+	/*
+	 * Open relation with ShareUpdateExclusiveLock, consistent with
+	 * ANALYZE. Only needed for permission check and then we close it (but
+	 * retain the lock).
+	 */
+	relation = table_open(reloid, ShareUpdateExclusiveLock);
+
+	check_privileges(relation);
+
+	table_close(relation, NoLock);
+
+	/*
+	 * Take RowExclusiveLock on pg_class, consistent with
+	 * vac_update_relstats().
+	 */
+	crel = table_open(RelationRelationId, RowExclusiveLock);
+
+	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)));
+		return false;
+	}
+
+	pgcform = (Form_pg_class) GETSTRUCT(ctup);
+
+	/* only update pg_class if there is a meaningful change */
+	if (pgcform->relpages != relpages ||
+		pgcform->reltuples != reltuples ||
+		pgcform->relallvisible != relallvisible)
+	{
+		int			cols[3] = {
+			Anum_pg_class_relpages,
+			Anum_pg_class_reltuples,
+			Anum_pg_class_relallvisible,
+		};
+		Datum		values[3] = {
+			Int32GetDatum(relpages),
+			Float4GetDatum(reltuples),
+			Int32GetDatum(relallvisible),
+		};
+		bool		nulls[3] = {false, false, false};
+
+		TupleDesc	tupdesc = RelationGetDescr(crel);
+		HeapTuple	newtup;
+
+		CatalogIndexState indstate = CatalogOpenIndexes(crel);
+
+		newtup = heap_modify_tuple_by_cols(ctup, tupdesc, 3, cols, values,
+										   nulls);
+
+		CatalogTupleUpdateWithInfo(crel, &newtup->t_self, newtup, indstate);
+		heap_freetuple(newtup);
+		CatalogCloseIndexes(indstate);
+	}
+
+	/* release the lock, consistent with vac_update_relstats() */
+	table_close(crel, RowExclusiveLock);
+
+	PG_RETURN_BOOL(true);
+}
+
+/*
+ * 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
+ */
+static void
+check_privileges(Relation rel)
+{
+	AclResult               aclresult;
+
+	if (object_ownercheck(DatabaseRelationId, MyDatabaseId, GetUserId()) &&
+		!rel->rd_rel->relisshared)
+		return;
+
+	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));
+}
+
+/*
+ * Set statistics for a given pg_class entry.
+ *
+ * Use a transactional update, and assume statistics come from the current
+ * server version.
+ *
+ * Not intended for bulk import of statistics from older versions.
+ */
+Datum
+pg_set_relation_stats(PG_FUNCTION_ARGS)
+{
+	Oid			reloid;
+	int			version = PG_VERSION_NUM;
+	int			elevel = ERROR;
+	int32		relpages;
+	float		reltuples;
+	int32		relallvisible;
+
+	if (PG_ARGISNULL(0))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("relation cannot be NULL")));
+	else
+		reloid = PG_GETARG_OID(0);
+
+	if (PG_ARGISNULL(1))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("relpages cannot be NULL")));
+	else
+		relpages = PG_GETARG_INT32(1);
+
+	if (PG_ARGISNULL(2))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("reltuples cannot be NULL")));
+	else
+		reltuples = PG_GETARG_FLOAT4(2);
+
+	if (PG_ARGISNULL(3))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("relallvisible cannot be NULL")));
+	else
+		relallvisible = PG_GETARG_INT32(3);
+
+
+	relation_statistics_update(reloid, version, relpages, reltuples,
+							   relallvisible, elevel);
+
+	PG_RETURN_VOID();
+}
diff --git a/src/backend/statistics/meson.build b/src/backend/statistics/meson.build
index 73b29a3d50a..849df3bf323 100644
--- a/src/backend/statistics/meson.build
+++ b/src/backend/statistics/meson.build
@@ -3,6 +3,7 @@
 backend_sources += files(
   'dependencies.c',
   'extended_stats.c',
+  'import_stats.c',
   'mcv.c',
   'mvdistinct.c',
 )
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 4abc6d95262..d700dd50f7b 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -12255,4 +12255,13 @@
   proargnames => '{summarized_tli,summarized_lsn,pending_lsn,summarizer_pid}',
   prosrc => 'pg_get_wal_summarizer_state' },
 
+# Statistics Import
+{ oid => '8048',
+  descr => 'set statistics on relation',
+  proname => 'pg_set_relation_stats', provolatile => 'v', proisstrict => 'f',
+  proparallel => 'u', prorettype => 'void',
+  proargtypes => 'oid int4 float4 int4',
+  proargnames => '{relation,relpages,reltuples,relallvisible}',
+  prosrc => 'pg_set_relation_stats' },
+
 ]
diff --git a/src/test/regress/expected/stats_import.out b/src/test/regress/expected/stats_import.out
new file mode 100644
index 00000000000..f727cce9765
--- /dev/null
+++ b/src/test/regress/expected/stats_import.out
@@ -0,0 +1,99 @@
+CREATE SCHEMA stats_import;
+CREATE TYPE stats_import.complex_type AS (
+    a integer,
+    b real,
+    c text,
+    d date,
+    e jsonb);
+CREATE TABLE stats_import.test(
+    id INTEGER PRIMARY KEY,
+    name text,
+    comp stats_import.complex_type,
+    arange int4range,
+    tags text[]
+);
+-- starting stats
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+ relpages | reltuples | relallvisible 
+----------+-----------+---------------
+        0 |        -1 |             0
+(1 row)
+
+-- error: regclass not found
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 0::Oid,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+ERROR:  could not open relation with OID 0
+-- error: relpages NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => NULL::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+ERROR:  relpages cannot be NULL
+-- error: reltuples NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => NULL::real,
+        relallvisible => 4::integer);
+ERROR:  reltuples cannot be NULL
+-- error: relallvisible NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => NULL::integer);
+ERROR:  relallvisible cannot be NULL
+-- named arguments
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+ relpages | reltuples | relallvisible 
+----------+-----------+---------------
+       17 |       400 |             4
+(1 row)
+
+-- positional arguments
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        'stats_import.test'::regclass,
+        18::integer,
+        401.0::real,
+        5::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+ relpages | reltuples | relallvisible 
+----------+-----------+---------------
+       18 |       401 |             5
+(1 row)
+
+DROP SCHEMA stats_import CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to type stats_import.complex_type
+drop cascades to table stats_import.test
diff --git a/src/test/regress/parallel_schedule b/src/test/regress/parallel_schedule
index 2429ec2bbaa..85fc85bfa03 100644
--- a/src/test/regress/parallel_schedule
+++ b/src/test/regress/parallel_schedule
@@ -28,7 +28,7 @@ test: strings md5 numerology point lseg line box path polygon circle date time t
 # geometry depends on point, lseg, line, box, path, polygon, circle
 # horology depends on date, time, timetz, timestamp, timestamptz, interval
 # ----------
-test: geometry horology tstypes regex type_sanity opr_sanity misc_sanity comments expressions unicode xid mvcc database
+test: geometry horology tstypes regex type_sanity opr_sanity misc_sanity comments expressions unicode xid mvcc database stats_import
 
 # ----------
 # Load huge amounts of data
diff --git a/src/test/regress/sql/stats_import.sql b/src/test/regress/sql/stats_import.sql
new file mode 100644
index 00000000000..effd5b892bf
--- /dev/null
+++ b/src/test/regress/sql/stats_import.sql
@@ -0,0 +1,79 @@
+CREATE SCHEMA stats_import;
+
+CREATE TYPE stats_import.complex_type AS (
+    a integer,
+    b real,
+    c text,
+    d date,
+    e jsonb);
+
+CREATE TABLE stats_import.test(
+    id INTEGER PRIMARY KEY,
+    name text,
+    comp stats_import.complex_type,
+    arange int4range,
+    tags text[]
+);
+
+-- starting stats
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+
+-- error: regclass not found
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 0::Oid,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+
+-- error: relpages NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => NULL::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+
+-- error: reltuples NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => NULL::real,
+        relallvisible => 4::integer);
+
+-- error: relallvisible NULL
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => NULL::integer);
+
+-- named arguments
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.test'::regclass,
+        relpages => 17::integer,
+        reltuples => 400.0::real,
+        relallvisible => 4::integer);
+
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+
+-- positional arguments
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        'stats_import.test'::regclass,
+        18::integer,
+        401.0::real,
+        5::integer);
+
+SELECT relpages, reltuples, relallvisible
+FROM pg_class
+WHERE oid = 'stats_import.test'::regclass;
+
+DROP SCHEMA stats_import CASCADE;
-- 
2.34.1

