From 9dd7129318804ede9e412d886f3124e6d452a2d9 Mon Sep 17 00:00:00 2001
From: Corey Huinker <corey.huinker@gmail.com>
Date: Wed, 26 Feb 2025 00:39:17 -0500
Subject: [PATCH vViaDellaAttnums 1/2] Add ability to reference columns by
 attnum to pg_restore_attribute_stats().

We need this because pg_dump needs to address index expression column
statistics by attnum because the attribute name is not stable across
upgrades.
---
 src/backend/statistics/attribute_stats.c   | 84 +++++++++++++++++-----
 src/test/regress/expected/stats_import.out |  2 +-
 2 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/src/backend/statistics/attribute_stats.c b/src/backend/statistics/attribute_stats.c
index 66a5676c810..71a7a175d1f 100644
--- a/src/backend/statistics/attribute_stats.c
+++ b/src/backend/statistics/attribute_stats.c
@@ -38,6 +38,7 @@ enum attribute_stats_argnum
 {
 	ATTRELATION_ARG = 0,
 	ATTNAME_ARG,
+	ATTNUM_ARG,
 	INHERITED_ARG,
 	NULL_FRAC_ARG,
 	AVG_WIDTH_ARG,
@@ -59,6 +60,7 @@ static struct StatsArgInfo attarginfo[] =
 {
 	[ATTRELATION_ARG] = {"relation", REGCLASSOID},
 	[ATTNAME_ARG] = {"attname", NAMEOID},
+	[ATTNUM_ARG] = {"attnum", INT2OID},
 	[INHERITED_ARG] = {"inherited", BOOLOID},
 	[NULL_FRAC_ARG] = {"null_frac", FLOAT4OID},
 	[AVG_WIDTH_ARG] = {"avg_width", INT4OID},
@@ -76,6 +78,22 @@ static struct StatsArgInfo attarginfo[] =
 	[NUM_ATTRIBUTE_STATS_ARGS] = {0}
 };
 
+enum clear_attribute_stats_argnum
+{
+	C_ATTRELATION_ARG = 0,
+	C_ATTNAME_ARG,
+	C_INHERITED_ARG,
+	C_NUM_ATTRIBUTE_STATS_ARGS
+};
+
+static struct StatsArgInfo cleararginfo[] =
+{
+	[C_ATTRELATION_ARG] = {"relation", REGCLASSOID},
+	[C_ATTNAME_ARG] = {"attname", NAMEOID},
+	[C_INHERITED_ARG] = {"inherited", BOOLOID},
+	[C_NUM_ATTRIBUTE_STATS_ARGS] = {0}
+};
+
 static bool attribute_statistics_update(FunctionCallInfo fcinfo);
 static Node *get_attr_expr(Relation rel, int attnum);
 static void get_attr_stat_type(Oid reloid, AttrNumber attnum,
@@ -116,9 +134,9 @@ static bool
 attribute_statistics_update(FunctionCallInfo fcinfo)
 {
 	Oid			reloid;
-	Name		attname;
-	bool		inherited;
+	char	   *attname;
 	AttrNumber	attnum;
+	bool		inherited;
 
 	Relation	starel;
 	HeapTuple	statup;
@@ -164,21 +182,51 @@ attribute_statistics_update(FunctionCallInfo fcinfo)
 	/* lock before looking up attribute */
 	stats_lock_check_privileges(reloid);
 
-	stats_check_required_arg(fcinfo, attarginfo, ATTNAME_ARG);
-	attname = PG_GETARG_NAME(ATTNAME_ARG);
-	attnum = get_attnum(reloid, NameStr(*attname));
+	/* user can specify either attname or attnum, but not both */
+	if (!PG_ARGISNULL(ATTNAME_ARG))
+	{
+		Name		attnamename;
+
+		if (!PG_ARGISNULL(ATTNUM_ARG))
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("must specify one of attname and attnum")));
+		attnamename = PG_GETARG_NAME(ATTNAME_ARG);
+		attname = NameStr(*attnamename);
+		attnum = get_attnum(reloid, attname);
+		/* note that this test covers attisdropped cases too: */
+		if (attnum == InvalidAttrNumber)
+			ereport(ERROR,
+					(errcode(ERRCODE_UNDEFINED_COLUMN),
+					 errmsg("column \"%s\" of relation \"%s\" does not exist",
+							attname, get_rel_name(reloid))));
+	}
+	else if (!PG_ARGISNULL(ATTNUM_ARG))
+	{
+		attnum = PG_GETARG_INT16(ATTNUM_ARG);
+		attname = get_attname(reloid, attnum, true);
+		/* Annoyingly, get_attname doesn't check attisdropped */
+		if (attname == NULL ||
+			!SearchSysCacheExistsAttName(reloid, attname))
+			ereport(ERROR,
+					(errcode(ERRCODE_UNDEFINED_COLUMN),
+					 errmsg("column %d of relation \"%s\" does not exist",
+							attnum, get_rel_name(reloid))));
+	}
+	else
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("must specify one of attname and attnum")));
+		attname = NULL;			/* keep compiler quiet */
+		attnum = 0;
+	}
 
 	if (attnum < 0)
 		ereport(ERROR,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("cannot modify statistics on system column \"%s\"",
-						NameStr(*attname))));
-
-	if (attnum == InvalidAttrNumber)
-		ereport(ERROR,
-				(errcode(ERRCODE_UNDEFINED_COLUMN),
-				 errmsg("column \"%s\" of relation \"%s\" does not exist",
-						NameStr(*attname), get_rel_name(reloid))));
+						attname)));
 
 	stats_check_required_arg(fcinfo, attarginfo, INHERITED_ARG);
 	inherited = PG_GETARG_BOOL(INHERITED_ARG);
@@ -241,7 +289,7 @@ attribute_statistics_update(FunctionCallInfo fcinfo)
 								&elemtypid, &elem_eq_opr))
 		{
 			ereport(WARNING,
-					(errmsg("unable to determine element type of attribute \"%s\"", NameStr(*attname)),
+					(errmsg("unable to determine element type of attribute \"%s\"", attname),
 					 errdetail("Cannot set STATISTIC_KIND_MCELEM or STATISTIC_KIND_DECHIST.")));
 			elemtypid = InvalidOid;
 			elem_eq_opr = InvalidOid;
@@ -257,7 +305,7 @@ attribute_statistics_update(FunctionCallInfo fcinfo)
 	{
 		ereport(WARNING,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("could not determine less-than operator for attribute \"%s\"", NameStr(*attname)),
+				 errmsg("could not determine less-than operator for attribute \"%s\"", attname),
 				 errdetail("Cannot set STATISTIC_KIND_HISTOGRAM or STATISTIC_KIND_CORRELATION.")));
 
 		do_histogram = false;
@@ -271,7 +319,7 @@ attribute_statistics_update(FunctionCallInfo fcinfo)
 	{
 		ereport(WARNING,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("attribute \"%s\" is not a range type", NameStr(*attname)),
+				 errmsg("attribute \"%s\" is not a range type", attname),
 				 errdetail("Cannot set STATISTIC_KIND_RANGE_LENGTH_HISTOGRAM or STATISTIC_KIND_BOUNDS_HISTOGRAM.")));
 
 		do_bounds_histogram = false;
@@ -857,7 +905,7 @@ pg_clear_attribute_stats(PG_FUNCTION_ARGS)
 	AttrNumber	attnum;
 	bool		inherited;
 
-	stats_check_required_arg(fcinfo, attarginfo, ATTRELATION_ARG);
+	stats_check_required_arg(fcinfo, cleararginfo, C_ATTRELATION_ARG);
 	reloid = PG_GETARG_OID(ATTRELATION_ARG);
 
 	if (RecoveryInProgress())
@@ -868,7 +916,7 @@ pg_clear_attribute_stats(PG_FUNCTION_ARGS)
 
 	stats_lock_check_privileges(reloid);
 
-	stats_check_required_arg(fcinfo, attarginfo, ATTNAME_ARG);
+	stats_check_required_arg(fcinfo, cleararginfo, C_ATTNAME_ARG);
 	attname = PG_GETARG_NAME(ATTNAME_ARG);
 	attnum = get_attnum(reloid, NameStr(*attname));
 
@@ -884,7 +932,7 @@ pg_clear_attribute_stats(PG_FUNCTION_ARGS)
 				 errmsg("column \"%s\" of relation \"%s\" does not exist",
 						NameStr(*attname), get_rel_name(reloid))));
 
-	stats_check_required_arg(fcinfo, attarginfo, INHERITED_ARG);
+	stats_check_required_arg(fcinfo, attarginfo, C_INHERITED_ARG);
 	inherited = PG_GETARG_BOOL(INHERITED_ARG);
 
 	delete_pg_statistic(reloid, attnum, inherited);
diff --git a/src/test/regress/expected/stats_import.out b/src/test/regress/expected/stats_import.out
index 7e8b7f429c9..b8fc2fcce46 100644
--- a/src/test/regress/expected/stats_import.out
+++ b/src/test/regress/expected/stats_import.out
@@ -1250,7 +1250,7 @@ SELECT pg_catalog.pg_restore_attribute_stats(
     'null_frac', 0.1::real,
     'avg_width', 2::integer,
     'n_distinct', 0.3::real);
-ERROR:  "attname" cannot be NULL
+ERROR:  must specify one of attname and attnum
 -- error: attname doesn't exist
 SELECT pg_catalog.pg_restore_attribute_stats(
     'relation', 'stats_import.test'::regclass,

base-commit: 0e42d31b0b2273c376ce9de946b59d155fac589c
-- 
2.48.1

