From 72ede8812f61ed4879f5c005a18131c32ba2768b Mon Sep 17 00:00:00 2001
From: Hannu Krosing <hannuk@google.com>
Date: Sun, 6 Jul 2025 23:41:46 +0200
Subject: [PATCH v3 1/2] added new and old style SQL functions and
 documentation

---
 doc/src/sgml/ref/pgbench.sgml |  43 +++++++-
 src/bin/pgbench/pgbench.c     | 187 +++++++++++++++++++++++++++++++++-
 2 files changed, 223 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index ab252d9fc74..6d733cff1af 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -105,8 +105,9 @@ pgbench -i <optional> <replaceable>other-options</replaceable> </optional> <repl
     <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
+    destroying any existing tables of these names and their dependencies.
+    It also creates 6 functions starting with <structname>pgbench_</structname>.
+    Be very careful to use another database if you have tables or functions having these
     names!
    </para>
   </caution>
@@ -361,6 +362,15 @@ 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 and SQL functions for internal scripts.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry id="pgbench-option-partition-method">
       <term><option>--partition-method=<replaceable>NAME</replaceable></option></term>
       <listitem>
@@ -427,8 +437,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 <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: <builtin: TPC-B (sort of)>
+      plpgsql-tpcb-like: <builtin: TPC-B (sort of) - pl/pgsql UDF>
+      sqlfunc-tpcb-like: <builtin: TPC-B (sort of) - 'BEGIN ATOMIC' SQL UDF>
+      oldsqlf-tpcb-like: <builtin: TPC-B (sort of) - LANGUAGE SQL UDF>
+          simple-update: <builtin: simple update>
+  plpgsql-simple-update: <builtin: simple update - pl/pgsql UDF>
+  sqlfunc-simple-update: <builtin: simple update - 'BEGIN ATOMIC' SQL UDF>
+  oldsqlf-simple-update: <builtin: simple update - LANGUAGE SQL UDF>
+            select-only: <builtin: select only>
+</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 497a936c141..21d2fd64548 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 */
@@ -796,6 +796,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>",
@@ -809,6 +836,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>",
@@ -915,6 +969,7 @@ usage(void)
 		   "  -q, --quiet              quiet logging (one message each 5 seconds)\n"
 		   "  -s, --scale=NUM          scaling factor\n"
 		   "  --foreign-keys           create foreign key constraints between tables\n"
+		   "  --no-functions           do not create pl/pgsql and SQL functions for internal scripts\n"
 		   "  --index-tablespace=TABLESPACE\n"
 		   "                           create indexes in the specified tablespace\n"
 		   "  --partition-method=(range|hash)\n"
@@ -4750,7 +4805,7 @@ initDropTables(PGconn *con)
 					 "pgbench_accounts, "
 					 "pgbench_branches, "
 					 "pgbench_history, "
-					 "pgbench_tellers");
+					 "pgbench_tellers cascade");
 }
 
 /*
@@ -4825,6 +4880,107 @@ createPartitions(PGconn *con)
 	termPQExpBuffer(&query);
 }
 
+/*
+ * Create the functions needed for plpgsql-* builting scripts
+ */
+static void
+initCreateFuntions(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
  */
@@ -5311,6 +5467,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";
+				initCreateFuntions(con);
+				break;
 			case ' ':
 				break;			/* ignore */
 			default:
@@ -6146,7 +6310,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");
 }
 
@@ -6705,6 +6869,7 @@ main(int argc, char **argv)
 		{"verbose-errors", no_argument, NULL, 15},
 		{"exit-on-abort", no_argument, NULL, 16},
 		{"debug", no_argument, NULL, 17},
+		{"no-functions", no_argument, NULL, 18},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -6712,6 +6877,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;
@@ -7058,6 +7224,10 @@ main(int argc, char **argv)
 			case 17:			/* debug */
 				pg_logging_increase_verbosity();
 				break;
+			case 18:				/* 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);
@@ -7155,6 +7325,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.50.0.727.gbf7dc18ff4-goog

