Limit pgbench COPY FREEZE to ordinary relations

pgbench client-side data generation uses COPY FREEZE to load data for most
tables. COPY FREEZE isn't supported for partitioned tables and since pgbench
only supports partitioning pgbench_accounts, pgbench used a hard-coded check to
skip COPY FREEZE and use plain COPY for a partitioned pgbench_accounts.

If the user has manually partitioned one of the other pgbench tables, this
causes client-side data generation to error out with:

ERROR:  cannot perform COPY FREEZE on a partitioned table

Fix this by limiting COPY FREEZE to ordinary tables (RELKIND_RELATION).

Author: Sergey Tatarintsev <[email protected]>
Reviewed-by: Álvaro Herrera <[email protected]>
Discussion: https://postgr.es/m/flat/97f55fca-8a7b-4da8-b413-7d1c57010676%40postgrespro.ru

2 files changed, 35 insertions, 15 deletions

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 46240e3f725..3f488b44d7b 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -231,11 +231,10 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
             sent to the server. This uses the client/server bandwidth
             extensively through a <command>COPY</command>.
             <command>pgbench</command> uses the <option>FREEZE</option> option
-            with version 14 or later
-            of <productname>PostgreSQL</productname> to speed up
-            subsequent <command>VACUUM</command>, except on the
-            <literal>pgbench_accounts</literal> table if partitions are
-            enabled. Using <literal>g</literal> causes logging to
+            to load data into ordinary (non-partition) tables with version 14
+            or later of <productname>PostgreSQL</productname> to speed up
+            subsequent <command>VACUUM</command>.
+            Using <literal>g</literal> causes logging to
             print one message every 100,000 rows while generating data for all
             tables.
            </para>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index f303bdeec8d..5e1fcf59c61 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -53,6 +53,7 @@
 #include <sys/select.h>
 #endif
 
+#include "catalog/pg_class_d.h"
 #include "common/int.h"
 #include "common/logging.h"
 #include "common/pg_prng.h"
@@ -848,6 +849,31 @@ static const PsqlScanCallbacks pgbench_callbacks = {
 	NULL,						/* don't need get_variable functionality */
 };
 
+static char
+get_table_relkind(PGconn *con, const char *table)
+{
+	PGresult   *res;
+	char	   *val;
+	char		relkind;
+	const char *params[1] = {table};
+	const char *sql =
+		"SELECT relkind FROM pg_catalog.pg_class WHERE oid=$1::pg_catalog.regclass";
+
+	res = PQexecParams(con, sql, 1, NULL, params, NULL, NULL, 0);
+	if (PQresultStatus(res) != PGRES_TUPLES_OK)
+	{
+		pg_log_error("query failed: %s", PQerrorMessage(con));
+		pg_log_error_detail("Query was: %s", sql);
+		exit(1);
+	}
+	val = PQgetvalue(res, 0, 0);
+	Assert(strlen(val) == 1);
+	relkind = val[0];
+	PQclear(res);
+
+	return relkind;
+}
+
 static inline pg_time_usec_t
 pg_time_now(void)
 {
@@ -4962,16 +4988,11 @@ initPopulateTable(PGconn *con, const char *table, int64 base,
 
 	initPQExpBuffer(&sql);
 
-	/*
-	 * Use COPY with FREEZE on v14 and later for all the tables except
-	 * pgbench_accounts when it is partitioned.
-	 */
-	if (PQserverVersion(con) >= 140000)
-	{
-		if (strcmp(table, "pgbench_accounts") != 0 ||
-			partitions == 0)
-			copy_statement_fmt = "copy %s from stdin with (freeze on)";
-	}
+	/* Use COPY with FREEZE on v14 and later for all ordinary tables */
+	if ((PQserverVersion(con) >= 140000) &&
+		get_table_relkind(con, table) == RELKIND_RELATION)
+		copy_statement_fmt = "copy %s from stdin with (freeze on)";
+
 
 	n = pg_snprintf(copy_statement, sizeof(copy_statement), copy_statement_fmt, table);
 	if (n >= sizeof(copy_statement))