From 719f3acbff52de604a7d18f270097d11686c71de Mon Sep 17 00:00:00 2001
From: Hannu Krosing <hannuk@google.com>
Date: Sun, 11 Jan 2026 11:54:55 +0100
Subject: [PATCH v5] rebased, changed --no-functions flag id from 18 to 19

---
 doc/src/sgml/ref/pgbench.sgml |  74 ++++++++++++--
 src/bin/pgbench/pgbench.c     | 187 +++++++++++++++++++++++++++++++++-
 2 files changed, 247 insertions(+), 14 deletions(-)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 2e401d1ceb8..d7b21ccd1a6 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -101,12 +101,12 @@ pgbench -i <optional> <replaceable>other-options</replaceable> </optional> <repl
 
   <caution>
    <para>
-    <literal>pgbench -i</literal> creates four tables <structname>pgbench_accounts</structname>,
-    <structname>pgbench_branches</structname>, <structname>pgbench_history</structname>, and
-    <structname>pgbench_tellers</structname>,
-    destroying any existing tables of these names.
-    Be very careful to use another database if you have tables having these
-    names!
+    <literal>pgbench -i</literal> creates four tables (<structname>pgbench_accounts</structname>,
+    <structname>pgbench_branches</structname>, <structname>pgbench_history</structname>,
+    and <structname>pgbench_tellers</structname>) and six functions with names
+    begining with <structname>pgbench_</structname>. This operation will drop
+    any existing tables or functions with these names, including all dependent
+    objects.
    </para>
   </caution>
 
@@ -192,18 +192,26 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
         <replaceable>init_steps</replaceable> specifies the
         initialization steps to be performed, using one character per step.
         Each step is invoked in the specified order.
-        The default is <literal>dtgvp</literal>.
+        The default is <literal>dYtgvpy</literal>.
         The available steps are:
 
         <variablelist>
          <varlistentry id="pgbench-option-init-steps-d">
-          <term><literal>d</literal> (Drop)</term>
+          <term><literal>d</literal> (Drop Tables)</term>
           <listitem>
            <para>
             Drop any existing <application>pgbench</application> tables.
            </para>
           </listitem>
          </varlistentry>
+         <varlistentry id="pgbench-option-init-steps-Y">
+          <term><literal>Y</literal> (Drop Functions)</term>
+          <listitem>
+           <para>
+            Drop any existing <application>pgbench</application> functions.
+           </para>
+          </listitem>
+         </varlistentry>
          <varlistentry id="pgbench-option-init-steps-t">
           <term><literal>t</literal> (create Tables)</term>
           <listitem>
@@ -268,7 +276,15 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
            </para>
           </listitem>
          </varlistentry>
-         <varlistentry id="pgbench-option-init-steps-f">
+         <varlistentry id="pgbench-option-init-steps-y">
+          <term><literal>y</literal> (create Functions)</term>
+          <listitem>
+           <para>
+            Create any neccessary <application>pgbench</application> functions.
+           </para>
+          </listitem>
+         </varlistentry>
+          <varlistentry id="pgbench-option-init-steps-f">
          <term><literal>f</literal> (create Foreign keys)</term>
           <listitem>
            <para>
@@ -360,6 +376,17 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
       </listitem>
      </varlistentry>
 
+     <varlistentry id="pgbench-option-no-functions">
+      <term><option>--no-functions</option></term>
+      <listitem>
+       <para>
+        Do not create pl/pgsql or SQL functions for internal scripts.
+        (This option suppresses the <literal>y</literal> initialization step,
+        even if it was specified in <option>-I</option>.)
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry id="pgbench-option-partition-method">
       <term><option>--partition-method=<replaceable>NAME</replaceable></option></term>
       <listitem>
@@ -426,8 +453,35 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
         Available built-in scripts are: <literal>tpcb-like</literal>,
         <literal>simple-update</literal> and <literal>select-only</literal>.
         Unambiguous prefixes of built-in names are accepted.
+       </para>
+       <para>
+        Unless disabled with the <literal>--no-functions</literal> option at database
+        init, the <literal>tpcb-like</literal> and <literal>simple-update</literal>
+        scripts are also implemented as User-Defined functions in the database which
+        can be tested using scripts named <literal>plpgsql-tpcb-like</literal>,
+        <literal>sqlfunc-tpcb-like</literal>, <literal>oldsqlf-tpcb-like</literal>,
+        <literal>plpgsql-simple-update</literal>, <literal>sqlfunc-simple-update</literal>
+        and <literal>oldsqlf-simple-update</literal>.
+        The <literal>sqlfunc-*</literal> versions use the new SQL-standard SQL functions and
+        the <literal>oldsqlf-*</literal> use the SQL functions defined using <literal>LANGUAGE SQL</literal>.
+        Use <literal>--show-script=scriptname</literal> to see what is actually run.
+       </para>
+       <para>
         With the special name <literal>list</literal>, show the list of built-in scripts
-        and exit immediately.
+        and exit immediately :
+<programlisting>
+$ pgbench -b list
+Available builtin scripts:
+              tpcb-like: &lt;builtin: TPC-B (sort of)&gt;
+      plpgsql-tpcb-like: &lt;builtin: TPC-B (sort of) - pl/pgsql UDF&gt;
+      sqlfunc-tpcb-like: &lt;builtin: TPC-B (sort of) - 'BEGIN ATOMIC' SQL UDF&gt;
+      oldsqlf-tpcb-like: &lt;builtin: TPC-B (sort of) - LANGUAGE SQL UDF&gt;
+          simple-update: &lt;builtin: simple update&gt;
+  plpgsql-simple-update: &lt;builtin: simple update - pl/pgsql UDF&gt;
+  sqlfunc-simple-update: &lt;builtin: simple update - 'BEGIN ATOMIC' SQL UDF&gt;
+  oldsqlf-simple-update: &lt;builtin: simple update - LANGUAGE SQL UDF&gt;
+            select-only: &lt;builtin: select only&gt;
+</programlisting>
        </para>
        <para>
         Optionally, write an integer weight after <literal>@</literal> to
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 58735871c17..64bfe9fefea 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -160,8 +160,8 @@ typedef struct socket_set
 /********************************************************************
  * some configurable parameters */
 
-#define DEFAULT_INIT_STEPS "dtgvp"	/* default -I setting */
-#define ALL_INIT_STEPS "dtgGvpf"	/* all possible steps */
+#define DEFAULT_INIT_STEPS "dYtgvpy"	/* default -I setting */
+#define ALL_INIT_STEPS "dYtgGvpfy"	/* all possible steps */
 
 #define LOG_STEP_SECONDS	5	/* seconds between log messages */
 #define DEFAULT_NXACTS	10		/* default nxacts */
@@ -804,6 +804,33 @@ static const BuiltinScript builtin_script[] =
 		"INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n"
 		"END;\n"
 	},
+	{
+		"plpgsql-tpcb-like",
+		"<builtin: TPC-B (sort of) - pl/pgsql UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_tpcb_like(:aid, :bid, :tid, :delta);\n"
+	},
+	{
+		"sqlfunc-tpcb-like",
+		"<builtin: TPC-B (sort of) - 'BEGIN ATOMIC' SQL UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_tpcb_like_sqlfunc(:aid, :bid, :tid, :delta);\n"
+	},
+	{
+		"oldsqlf-tpcb-like",
+		"<builtin: TPC-B (sort of) - LANGUAGE SQL UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_tpcb_like_oldsqlfunc(:aid, :bid, :tid, :delta);\n"
+	},
 	{
 		"simple-update",
 		"<builtin: simple update>",
@@ -817,6 +844,33 @@ static const BuiltinScript builtin_script[] =
 		"INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n"
 		"END;\n"
 	},
+	{
+		"plpgsql-simple-update",
+		"<builtin: simple update - pl/pgsql UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_simple_update(:aid, :bid, :tid, :delta);\n"
+	},
+	{
+		"sqlfunc-simple-update",
+		"<builtin: simple update - 'BEGIN ATOMIC' SQL UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_simple_update_sqlfunc(:aid, :bid, :tid, :delta);\n"
+	},
+	{
+		"oldsqlf-simple-update",
+		"<builtin: simple update - LANGUAGE SQL UDF>",
+		"\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n"
+		"\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n"
+		"\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n"
+		"\\set delta random(-5000, 5000)\n"
+		"SELECT 1 FROM pgbench_simple_update_oldsqlfunc(:aid, :bid, :tid, :delta);\n"
+	},
 	{
 		"select-only",
 		"<builtin: select only>",
@@ -925,6 +979,7 @@ usage(void)
 		   "  --foreign-keys           create foreign key constraints between tables\n"
 		   "  --index-tablespace=TABLESPACE\n"
 		   "                           create indexes in the specified tablespace\n"
+		   "  --no-functions           do not create pl/pgsql or SQL functions for internal scripts\n"
 		   "  --partition-method=(range|hash)\n"
 		   "                           partition pgbench_accounts with this method (default: range)\n"
 		   "  --partitions=NUM         partition pgbench_accounts into NUM parts (default: 0)\n"
@@ -4814,7 +4869,7 @@ initDropTables(PGconn *con)
 					 "pgbench_accounts, "
 					 "pgbench_branches, "
 					 "pgbench_history, "
-					 "pgbench_tellers");
+					 "pgbench_tellers cascade");
 }
 
 /*
@@ -4889,6 +4944,107 @@ createPartitions(PGconn *con)
 	termPQExpBuffer(&query);
 }
 
+/*
+ * Create the functions needed for plpgsql-* builtin scripts
+ */
+static void
+initCreateFunctions(PGconn *con)
+{
+	fprintf(stderr, "creating functions...\n");
+
+	executeStatement(con,
+		"CREATE FUNCTION pgbench_tpcb_like(_aid int, _bid int, _tid int, _delta int)\n"
+		"RETURNS void\n"
+		"LANGUAGE plpgsql\n"
+		"AS $plpgsql$\n"
+		"BEGIN\n"
+		"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+		"    PERFORM abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+		"    UPDATE pgbench_tellers SET tbalance = tbalance + _delta WHERE tid = _tid;\n"
+		"    UPDATE pgbench_branches SET bbalance = bbalance + _delta WHERE bid = _bid;\n"
+		"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+		"END;\n"
+		"$plpgsql$;\n");
+	executeStatement(con,
+		"CREATE FUNCTION pgbench_simple_update(_aid int, _bid int, _tid int, _delta int)\n"
+		"RETURNS void\n"
+		"LANGUAGE plpgsql\n"
+		"AS $plpgsql$\n"
+		"BEGIN\n"
+		"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+		"    PERFORM abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+		"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+		"END;\n"
+		"$plpgsql$;\n");
+	if ((PQserverVersion(con) >= 140000))
+	{
+		executeStatement(con,
+			"CREATE FUNCTION pgbench_tpcb_like_sqlfunc(_aid int, _bid int, _tid int, _delta int)\n"
+			"RETURNS void\n"
+			"BEGIN ATOMIC\n"
+			"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+			"    SELECT abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+			"    UPDATE pgbench_tellers SET tbalance = tbalance + _delta WHERE tid = _tid;\n"
+			"    UPDATE pgbench_branches SET bbalance = bbalance + _delta WHERE bid = _bid;\n"
+			"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+			"END;\n");
+		executeStatement(con,
+			"CREATE FUNCTION pgbench_simple_update_sqlfunc(_aid int, _bid int, _tid int, _delta int)\n"
+			"RETURNS void\n"
+			"BEGIN ATOMIC\n"
+			"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+			"    SELECT abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+			"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+			"END;\n");
+	}
+	executeStatement(con,
+		"CREATE FUNCTION pgbench_tpcb_like_oldsqlfunc(_aid int, _bid int, _tid int, _delta int)\n"
+		"RETURNS void\n"
+		"LANGUAGE sql\n"
+		"AS $sql$\n"
+		"-- BEGIN\n"
+		"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+		"    SELECT abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+		"    UPDATE pgbench_tellers SET tbalance = tbalance + _delta WHERE tid = _tid;\n"
+		"    UPDATE pgbench_branches SET bbalance = bbalance + _delta WHERE bid = _bid;\n"
+		"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+		"-- END;\n"
+		"$sql$;\n");
+	executeStatement(con,
+		"CREATE FUNCTION pgbench_simple_update_oldsqlfunc(_aid int, _bid int, _tid int, _delta int)\n"
+		"RETURNS void\n"
+		"LANGUAGE sql\n"
+		"AS $sql$\n"
+		"-- BEGIN\n"
+		"    UPDATE pgbench_accounts SET abalance = abalance + _delta WHERE aid = _aid;\n"
+		"    SELECT abalance FROM pgbench_accounts WHERE aid = _aid;\n"
+		"    INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (_tid, _bid, _aid, _delta, CURRENT_TIMESTAMP);\n"
+		"-- END;\n"
+		"$sql$;\n");
+}
+
+/*
+ * Remove old pgbench functions, if any exist
+ */
+static void
+initDropFunctions(PGconn *con)
+{
+	fprintf(stderr, "dropping old functions...\n");
+
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_tpcb_like(_aid int, _bid int, _tid int, _delta int);\n");
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_simple_update(_aid int, _bid int, _tid int, _delta int);\n");
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_tpcb_like_sqlfunc(_aid int, _bid int, _tid int, _delta int);\n");
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_simple_update_sqlfunc(_aid int, _bid int, _tid int, _delta int);\n");
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_tpcb_like_oldsqlfunc(_aid int, _bid int, _tid int, _delta int);\n");
+	executeStatement(con,
+		"DROP FUNCTION IF EXISTS pgbench_simple_update_oldsqlfunc(_aid int, _bid int, _tid int, _delta int);\n");
+}
+
 /*
  * Create pgbench's standard tables
  */
@@ -5375,6 +5531,14 @@ runInitSteps(const char *initialize_steps)
 				op = "foreign keys";
 				initCreateFKeys(con);
 				break;
+			case 'Y':
+				op = "drop functions";
+				initDropFunctions(con);
+				break;
+			case 'y':
+				op = "create functions";
+				initCreateFunctions(con);
+				break;
 			case ' ':
 				break;			/* ignore */
 			default:
@@ -6210,7 +6374,7 @@ listAvailableScripts(void)
 
 	fprintf(stderr, "Available builtin scripts:\n");
 	for (i = 0; i < lengthof(builtin_script); i++)
-		fprintf(stderr, "  %13s: %s\n", builtin_script[i].name, builtin_script[i].desc);
+		fprintf(stderr, "  %21s: %s\n", builtin_script[i].name, builtin_script[i].desc);
 	fprintf(stderr, "\n");
 }
 
@@ -6779,6 +6943,7 @@ main(int argc, char **argv)
 		{"exit-on-abort", no_argument, NULL, 16},
 		{"debug", no_argument, NULL, 17},
 		{"continue-on-error", no_argument, NULL, 18},
+		{"no-functions", no_argument, NULL, 19},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -6786,6 +6951,7 @@ main(int argc, char **argv)
 	bool		is_init_mode = false;	/* initialize mode? */
 	char	   *initialize_steps = NULL;
 	bool		foreign_keys = false;
+	bool		no_functions = false;
 	bool		is_no_vacuum = false;
 	bool		do_vacuum_accounts = false; /* vacuum accounts table? */
 	int			optindex;
@@ -7136,6 +7302,10 @@ main(int argc, char **argv)
 				benchmarking_option_set = true;
 				continue_on_error = true;
 				break;
+			case 19:				/* no-functions */
+				initialization_option_set = true;
+				no_functions = true;
+				break;
 			default:
 				/* getopt_long already emitted a complaint */
 				pg_log_error_hint("Try \"%s --help\" for more information.", progname);
@@ -7233,6 +7403,15 @@ main(int argc, char **argv)
 				*p = ' ';
 		}
 
+		if (no_functions)
+		{
+			/* Remove create function step in initialize_steps */
+			char	   *p;
+
+			while ((p = strchr(initialize_steps, 'y')) != NULL)
+				*p = ' ';
+		}
+
 		if (foreign_keys)
 		{
 			/* Add 'f' to end of initialize_steps, if not already there */
-- 
2.43.0

