pgbench: Use PQExpBuffer to simplify code that constructs SQL.

hlinnaka · hlinnaka · commit 9796f455c38e · 2020-09-30T10:58:09.000+03:00
Author: Fabien Coelho Reviewed-by: Jeevan Ladhe Discussion: https://www.postgresql.org/message-id/alpine.DEB.2.21.1910220826570.15559%40lancre
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
@@ -603,7 +603,6 @@ static void doLog(TState *thread, CState *st,
 				  StatsData *agg, bool skipped, double latency, double lag);
 static void processXactStats(TState *thread, CState *st, instr_time *now,
 							 bool skipped, StatsData *agg);
-static void append_fillfactor(char *opts, int len);
 static void addScript(ParsedScript script);
 static void *threadRun(void *arg);
 static void finishCon(CState *st);
@@ -3630,30 +3629,26 @@ initDropTables(PGconn *con)
 static void
 createPartitions(PGconn *con)
 {
-	char		ff[64];
-
-	ff[0] = '\0';
-
-	/*
-	 * Per ddlinfo in initCreateTables, fillfactor is needed on table
-	 * pgbench_accounts.
-	 */
-	append_fillfactor(ff, sizeof(ff));
+	PQExpBufferData query;
 
 	/* we must have to create some partitions */
 	Assert(partitions > 0);
 
 	fprintf(stderr, "creating %d partitions...\n", partitions);
 
+	initPQExpBuffer(&query);
+
 	for (int p = 1; p <= partitions; p++)
 	{
-		char		query[256];
-
 		if (partition_method == PART_RANGE)
 		{
 			int64		part_size = (naccounts * (int64) scale + partitions - 1) / partitions;
-			char		minvalue[32],
-						maxvalue[32];
+
+			printfPQExpBuffer(&query,
+							  "create%s table pgbench_accounts_%d\n"
+							  "  partition of pgbench_accounts\n"
+							  "  for values from (",
+							  unlogged_tables ? " unlogged" : "", p);
 
 			/*
 			 * For RANGE, we use open-ended partitions at the beginning and
@@ -3662,34 +3657,39 @@ createPartitions(PGconn *con)
 			 * scale, it is more generic and the performance is better.
 			 */
 			if (p == 1)
-				sprintf(minvalue, "minvalue");
+				appendPQExpBufferStr(&query, "minvalue");
 			else
-				sprintf(minvalue, INT64_FORMAT, (p - 1) * part_size + 1);
+				appendPQExpBuffer(&query, INT64_FORMAT, (p - 1) * part_size + 1);
+
+			appendPQExpBufferStr(&query, ") to (");
 
 			if (p < partitions)
-				sprintf(maxvalue, INT64_FORMAT, p * part_size + 1);
+				appendPQExpBuffer(&query, INT64_FORMAT, p * part_size + 1);
 			else
-				sprintf(maxvalue, "maxvalue");
-
-			snprintf(query, sizeof(query),
-					 "create%s table pgbench_accounts_%d\n"
-					 "  partition of pgbench_accounts\n"
-					 "  for values from (%s) to (%s)%s\n",
-					 unlogged_tables ? " unlogged" : "", p,
-					 minvalue, maxvalue, ff);
+				appendPQExpBufferStr(&query, "maxvalue");
+
+			appendPQExpBufferChar(&query, ')');
 		}
 		else if (partition_method == PART_HASH)
-			snprintf(query, sizeof(query),
-					 "create%s table pgbench_accounts_%d\n"
-					 "  partition of pgbench_accounts\n"
-					 "  for values with (modulus %d, remainder %d)%s\n",
-					 unlogged_tables ? " unlogged" : "", p,
-					 partitions, p - 1, ff);
+			printfPQExpBuffer(&query,
+							  "create%s table pgbench_accounts_%d\n"
+							  "  partition of pgbench_accounts\n"
+							  "  for values with (modulus %d, remainder %d)",
+							  unlogged_tables ? " unlogged" : "", p,
+							  partitions, p - 1);
 		else					/* cannot get there */
 			Assert(0);
 
-		executeStatement(con, query);
+		/*
+		 * Per ddlinfo in initCreateTables, fillfactor is needed on table
+		 * pgbench_accounts.
+		 */
+		appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor);
+
+		executeStatement(con, query.data);
 	}
+
+	termPQExpBuffer(&query);
 }
 
 /*
@@ -3743,63 +3743,50 @@ initCreateTables(PGconn *con)
 		}
 	};
 	int			i;
+	PQExpBufferData query;
 
 	fprintf(stderr, "creating tables...\n");
 
+	initPQExpBuffer(&query);
+
 	for (i = 0; i < lengthof(DDLs); i++)
 	{
-		char		opts[256];
-		char		buffer[256];
 		const struct ddlinfo *ddl = &DDLs[i];
-		const char *cols;
 
 		/* Construct new create table statement. */
-		opts[0] = '\0';
+		printfPQExpBuffer(&query, "create%s table %s(%s)",
+						  unlogged_tables ? " unlogged" : "",
+						  ddl->table,
+						  (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols);
 
 		/* Partition pgbench_accounts table */
 		if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0)
-			snprintf(opts + strlen(opts), sizeof(opts) - strlen(opts),
-					 " partition by %s (aid)", PARTITION_METHOD[partition_method]);
+			appendPQExpBuffer(&query,
+							  " partition by %s (aid)", PARTITION_METHOD[partition_method]);
 		else if (ddl->declare_fillfactor)
+		{
 			/* fillfactor is only expected on actual tables */
-			append_fillfactor(opts, sizeof(opts));
+			appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor);
+		}
 
 		if (tablespace != NULL)
 		{
 			char	   *escape_tablespace;
 
-			escape_tablespace = PQescapeIdentifier(con, tablespace,
-												   strlen(tablespace));
-			snprintf(opts + strlen(opts), sizeof(opts) - strlen(opts),
-					 " tablespace %s", escape_tablespace);
+			escape_tablespace = PQescapeIdentifier(con, tablespace, strlen(tablespace));
+			appendPQExpBuffer(&query, " tablespace %s", escape_tablespace);
 			PQfreemem(escape_tablespace);
 		}
 
-		cols = (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols;
-
-		snprintf(buffer, sizeof(buffer), "create%s table %s(%s)%s",
-				 unlogged_tables ? " unlogged" : "",
-				 ddl->table, cols, opts);
-
-		executeStatement(con, buffer);
+		executeStatement(con, query.data);
 	}
 
+	termPQExpBuffer(&query);
+
 	if (partition_method != PART_NONE)
 		createPartitions(con);
 }
 
-/*
- * add fillfactor percent option.
- *
- * XXX - As default is 100, it could be removed in this case.
- */
-static void
-append_fillfactor(char *opts, int len)
-{
-	snprintf(opts + strlen(opts), len - strlen(opts),
-			 " with (fillfactor=%d)", fillfactor);
-}
-
 /*
  * Truncate away any old data, in one command in case there are foreign keys
  */
@@ -3819,7 +3806,7 @@ initTruncateTables(PGconn *con)
 static void
 initGenerateDataClientSide(PGconn *con)
 {
-	char		sql[256];
+	PQExpBufferData sql;
 	PGresult   *res;
 	int			i;
 	int64		k;
@@ -3845,26 +3832,28 @@ initGenerateDataClientSide(PGconn *con)
 	/* truncate away any old data */
 	initTruncateTables(con);
 
+	initPQExpBuffer(&sql);
+
 	/*
 	 * fill branches, tellers, accounts in that order in case foreign keys
 	 * already exist
 	 */
 	for (i = 0; i < nbranches * scale; i++)
 	{
 		/* "filler" column defaults to NULL */
-		snprintf(sql, sizeof(sql),
-				 "insert into pgbench_branches(bid,bbalance) values(%d,0)",
-				 i + 1);
-		executeStatement(con, sql);
+		printfPQExpBuffer(&sql,
+						  "insert into pgbench_branches(bid,bbalance) values(%d,0)",
+						  i + 1);
+		executeStatement(con, sql.data);
 	}
 
 	for (i = 0; i < ntellers * scale; i++)
 	{
 		/* "filler" column defaults to NULL */
-		snprintf(sql, sizeof(sql),
-				 "insert into pgbench_tellers(tid,bid,tbalance) values (%d,%d,0)",
-				 i + 1, i / ntellers + 1);
-		executeStatement(con, sql);
+		printfPQExpBuffer(&sql,
+						  "insert into pgbench_tellers(tid,bid,tbalance) values (%d,%d,0)",
+						  i + 1, i / ntellers + 1);
+		executeStatement(con, sql.data);
 	}
 
 	/*
@@ -3885,10 +3874,10 @@ initGenerateDataClientSide(PGconn *con)
 		int64		j = k + 1;
 
 		/* "filler" column defaults to blank padded empty string */
-		snprintf(sql, sizeof(sql),
-				 INT64_FORMAT "\t" INT64_FORMAT "\t%d\t\n",
-				 j, k / naccounts + 1, 0);
-		if (PQputline(con, sql))
+		printfPQExpBuffer(&sql,
+						  INT64_FORMAT "\t" INT64_FORMAT "\t%d\t\n",
+						  j, k / naccounts + 1, 0);
+		if (PQputline(con, sql.data))
 		{
 			pg_log_fatal("PQputline failed");
 			exit(1);
@@ -3950,6 +3939,8 @@ initGenerateDataClientSide(PGconn *con)
 		exit(1);
 	}
 
+	termPQExpBuffer(&sql);
+
 	executeStatement(con, "commit");
 }
 
@@ -3963,7 +3954,7 @@ initGenerateDataClientSide(PGconn *con)
 static void
 initGenerateDataServerSide(PGconn *con)
 {
-	char		sql[256];
+	PQExpBufferData sql;
 
 	fprintf(stderr, "generating data (server-side)...\n");
 
@@ -3976,24 +3967,28 @@ initGenerateDataServerSide(PGconn *con)
 	/* truncate away any old data */
 	initTruncateTables(con);
 
-	snprintf(sql, sizeof(sql),
-			 "insert into pgbench_branches(bid,bbalance) "
-			 "select bid, 0 "
-			 "from generate_series(1, %d) as bid", nbranches * scale);
-	executeStatement(con, sql);
-
-	snprintf(sql, sizeof(sql),
-			 "insert into pgbench_tellers(tid,bid,tbalance) "
-			 "select tid, (tid - 1) / %d + 1, 0 "
-			 "from generate_series(1, %d) as tid", ntellers, ntellers * scale);
-	executeStatement(con, sql);
-
-	snprintf(sql, sizeof(sql),
-			 "insert into pgbench_accounts(aid,bid,abalance,filler) "
-			 "select aid, (aid - 1) / %d + 1, 0, '' "
-			 "from generate_series(1, " INT64_FORMAT ") as aid",
-			 naccounts, (int64) naccounts * scale);
-	executeStatement(con, sql);
+	initPQExpBuffer(&sql);
+
+	printfPQExpBuffer(&sql,
+					  "insert into pgbench_branches(bid,bbalance) "
+					  "select bid, 0 "
+					  "from generate_series(1, %d) as bid", nbranches * scale);
+	executeStatement(con, sql.data);
+
+	printfPQExpBuffer(&sql,
+					  "insert into pgbench_tellers(tid,bid,tbalance) "
+					  "select tid, (tid - 1) / %d + 1, 0 "
+					  "from generate_series(1, %d) as tid", ntellers, ntellers * scale);
+	executeStatement(con, sql.data);
+
+	printfPQExpBuffer(&sql,
+					  "insert into pgbench_accounts(aid,bid,abalance,filler) "
+					  "select aid, (aid - 1) / %d + 1, 0, '' "
+					  "from generate_series(1, " INT64_FORMAT ") as aid",
+					  naccounts, (int64) naccounts * scale);
+	executeStatement(con, sql.data);
+
+	termPQExpBuffer(&sql);
 
 	executeStatement(con, "commit");
 }
@@ -4023,27 +4018,30 @@ initCreatePKeys(PGconn *con)
 		"alter table pgbench_accounts add primary key (aid)"
 	};
 	int			i;
+	PQExpBufferData query;
 
 	fprintf(stderr, "creating primary keys...\n");
+	initPQExpBuffer(&query);
+
 	for (i = 0; i < lengthof(DDLINDEXes); i++)
 	{
-		char		buffer[256];
-
-		strlcpy(buffer, DDLINDEXes[i], sizeof(buffer));
+		resetPQExpBuffer(&query);
+		appendPQExpBufferStr(&query, DDLINDEXes[i]);
 
 		if (index_tablespace != NULL)
 		{
 			char	   *escape_tablespace;
 
 			escape_tablespace = PQescapeIdentifier(con, index_tablespace,
 												   strlen(index_tablespace));
-			snprintf(buffer + strlen(buffer), sizeof(buffer) - strlen(buffer),
-					 " using index tablespace %s", escape_tablespace);
+			appendPQExpBuffer(&query, " using index tablespace %s", escape_tablespace);
 			PQfreemem(escape_tablespace);
 		}
 
-		executeStatement(con, buffer);
+		executeStatement(con, query.data);
 	}
+
+	termPQExpBuffer(&query);
 }
 
 /*
