dimitri · dimitri · Apr 17, 2024 · Apr 17, 2024
diff --git a/src/bin/pgcopydb/copydb.h b/src/bin/pgcopydb/copydb.h
@@ -446,6 +446,7 @@ bool copydb_copy_table(CopyDataSpec *specs, PGSQL *src, PGSQL *dst,
 
 bool copydb_table_create_lockfile(CopyDataSpec *specs,
 								  CopyTableDataSpec *tableSpecs,
+								  PGSQL *dst,
 								  bool *isDone);
 
 bool copydb_mark_table_as_done(CopyDataSpec *specs,

diff --git a/src/bin/pgcopydb/pgsql.c b/src/bin/pgcopydb/pgsql.c
@@ -1329,6 +1329,44 @@ pgsql_has_sequence_privilege(PGSQL *pgsql,
 }
 
 
+/*
+ * pgsql_has_table_privilege calls has_table_privilege() and copies the result
+ * in the granted boolean pointer given.
+ */
+bool
+pgsql_has_table_privilege(PGSQL *pgsql,
+						  const char *tablename,
+						  const char *privilege,
+						  bool *granted)
+{
+	SingleValueResultContext parseContext = { { 0 }, PGSQL_RESULT_BOOL, false };
+
+	char *sql = "select has_table_privilege($1, $2);";
+
+	int paramCount = 2;
+	Oid paramTypes[2] = { TEXTOID, TEXTOID };
+	const char *paramValues[2] = { tablename, privilege };
+
+	if (!pgsql_execute_with_params(pgsql, sql,
+								   paramCount, paramTypes, paramValues,
+								   &parseContext, &parseSingleValueResult))
+	{
+		log_error("Failed to query privileges for table \"%s\"", tablename);
+		return false;
+	}
+
+	if (!parseContext.parsedOk)
+	{
+		log_error("Failed to query privileges for table \"%s\"", tablename);
+		return false;
+	}
+
+	*granted = parseContext.boolVal;
+
+	return true;
+}
+
+
 /*
  * pgsql_get_search_path runs the query "show search_path" and copies the
  * result in the given pre-allocated string buffer.
@@ -2614,7 +2652,8 @@ pgsql_truncate(PGSQL *pgsql, const char *qname)
 
 	sformat(sql, sizeof(sql), "TRUNCATE ONLY %s", qname);
 
-	log_sql("%s", sql);
+	/* this being more like a DDL operation, proper log level is NOTICE */
+	log_notice("%s", sql);
 
 	return pgsql_execute(pgsql, sql);
 }
@@ -2687,6 +2726,9 @@ pg_copy_data(PGSQL *src, PGSQL *dst, CopyArgs *args)
 		}
 	}
 
+	/* make sure to log TRUNCATE before we log COPY, avoid confusion */
+	log_notice("%s", args->logCommand);
+
 	/* SRC: COPY schema.table TO STDOUT */
 	if (!pg_copy_send_query(src, args, PGRES_COPY_OUT))
 	{

diff --git a/src/bin/pgcopydb/pgsql.h b/src/bin/pgcopydb/pgsql.h
@@ -285,6 +285,11 @@ bool pgsql_has_sequence_privilege(PGSQL *pgsql,
 								  const char *privilege,
 								  bool *granted);
 
+bool pgsql_has_table_privilege(PGSQL *pgsql,
+							   const char *tablename,
+							   const char *privilege,
+							   bool *granted);
+
 bool pgsql_get_search_path(PGSQL *pgsql, char *search_path, size_t size);
 bool pgsql_set_search_path(PGSQL *pgsql, char *search_path, bool local);
 bool pgsql_prepend_search_path(PGSQL *pgsql, const char *namespace);
@@ -326,6 +331,7 @@ typedef struct CopyArgs
 	char *srcWhereClause;
 	char *dstQname;
 	char *dstAttrList;
+	char *logCommand;
 	bool truncate;
 	bool freeze;
 	uint64_t bytesTransmitted;

diff --git a/src/bin/pgcopydb/table-data.c b/src/bin/pgcopydb/table-data.c
@@ -574,12 +574,23 @@ copydb_copy_supervisor_add_table_hook(void *ctx, SourceTable *table)
 		 * Before adding the table to be processed by workers, truncate it on
 		 * the target database now, avoiding concurrency issues.
 		 */
-		if (!pgsql_truncate(dst, table->qname))
+		bool granted = false;
+
+		if (!pgsql_has_table_privilege(dst, table->qname, "TRUNCATE", &granted))
 		{
 			/* errors have already been logged */
 			return false;
 		}
 
+		if (granted)
+		{
+			if (!pgsql_truncate(dst, table->qname))
+			{
+				/* errors have already been logged */
+				return false;
+			}
+		}
+
 		for (int i = 0; i < table->partition.partCount; i++)
 		{
 			int partNumber = i + 1;
@@ -899,7 +910,7 @@ copydb_copy_data_by_oid(CopyDataSpec *specs, PGSQL *src, PGSQL *dst,
 	 */
 	bool isDone = false;
 
-	if (!copydb_table_create_lockfile(specs, tableSpecs, &isDone))
+	if (!copydb_table_create_lockfile(specs, tableSpecs, dst, &isDone))
 	{
 		/* errors have already been logged */
 		return false;
@@ -1026,6 +1037,7 @@ copydb_copy_data_by_oid(CopyDataSpec *specs, PGSQL *src, PGSQL *dst,
 bool
 copydb_table_create_lockfile(CopyDataSpec *specs,
 							 CopyTableDataSpec *tableSpecs,
+							 PGSQL *dst,
 							 bool *isDone)
 {
 	DatabaseCatalog *sourceDB = &(specs->catalogs.source);
@@ -1101,10 +1113,37 @@ copydb_table_create_lockfile(CopyDataSpec *specs,
 	args->srcWhereClause = NULL;
 	args->dstQname = tableSpecs->sourceTable->qname;
 	args->dstAttrList = tableSpecs->sourceTable->attrList;
-	args->truncate = tableSpecs->sourceTable->partition.partCount <= 1;
+	args->truncate = false;     /* default value, see below */
 	args->freeze = tableSpecs->sourceTable->partition.partCount <= 1;
 	args->bytesTransmitted = 0;
 
+	/*
+	 * Check to see if we want to TRUNCATE the table and benefit from the COPY
+	 * FREEZE optimisation.
+	 *
+	 * First, if the table COPY is partitionned then we truncate at the
+	 * top-level rather than for each partition, disabling the COPY FREEZE
+	 * optimisation.
+	 *
+	 * Second, we need the permission to run the TRUNCATE command on the target
+	 * table on the target database.
+	 */
+	if (tableSpecs->sourceTable->partition.partCount <= 1)
+	{
+		bool granted = false;
+
+		if (!pgsql_has_table_privilege(dst,
+									   tableSpecs->sourceTable->qname,
+									   "TRUNCATE",
+									   &granted))
+		{
+			/* errors have already been logged */
+			return false;
+		}
+
+		args->truncate = granted;
+	}
+
 	if (!copydb_prepare_copy_query(tableSpecs, args))
 	{
 		/* errors have already been logged */
@@ -1253,7 +1292,6 @@ copydb_copy_table(CopyDataSpec *specs, PGSQL *src, PGSQL *dst,
 
 	/* Now copy the data from source to target */
 	CopyTableSummary *summary = &(tableSpecs->summary);
-	log_notice("%s", summary->command);
 
 	int attempts = 0;
 	int maxAttempts = 5;        /* allow 5 attempts total, 4 retries */
@@ -1406,6 +1444,9 @@ copydb_prepare_summary_command(CopyTableDataSpec *tableSpecs)
 	tableSummary->command =
 		command->data != NULL ? strdup(command->data) : NULL;
 
+	/* also keep a pointer around in the copyArgs structure */
+	tableSpecs->copyArgs.logCommand = tableSummary->command;
+
 	if (PQExpBufferBroken(command) || tableSummary->command == NULL)
 	{
 		log_error("Failed to create summary command for %s: out of memory",