Support TRUNCATE for foreign tables

src/backend/distributed/commands/truncate.c

@@ -267,13 +267,17 @@ ErrorIfUnsupportedTruncateStmt(TruncateStmt *truncateStatement)
 
 		ErrorIfIllegallyChangingKnownShard(relationId);
 
-		if (IsCitusTable(relationId) && IsForeignTable(relationId))
+		/*
+		 * We allow truncating foreign tables that are added to metadata
+		 * only on the coordinator, as user mappings are not propagated.
+		 */
+		if (IsForeignTable(relationId) &&
+			IsCitusTableType(relationId, CITUS_LOCAL_TABLE) &&
+			!IsCoordinator())
 		{
 			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							errmsg("truncating distributed foreign tables is "
-								   "currently unsupported"),
-							errhint("Consider undistributing table before TRUNCATE, "
-									"and then distribute or add to metadata again")));
+							errmsg("truncating foreign tables that are added to metadata "
+								   "can only be excuted on the coordinator")));
 		}
 	}
 }

src/backend/distributed/utils/resource_lock.c

@@ -1012,8 +1012,8 @@ CitusRangeVarCallbackForLockTable(const RangeVar *rangeVar, Oid relationId,
 		return;
 	}
 
-	/* we only allow tables and views to be locked */
-	if (!RegularTable(relationId))
+	/* we only allow tables, views and foreign tables to be locked */
+	if (!RegularTable(relationId) && !IsForeignTable(relationId))
 	{
 		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
 						errmsg("\"%s\" is not a table", rangeVar->relname)));

src/test/regress/expected/pg14.out

@@ -1272,3 +1272,72 @@ SELECT count(*) FROM
 
 set client_min_messages to error;
 drop schema pg14 cascade;
+create schema pg14;
+set search_path to pg14;
+select 1 from citus_add_node('localhost',:master_port,groupid=>0);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+-- test adding foreign table to metadata with the guc
+-- will test truncating foreign tables later
+CREATE TABLE foreign_table_test (id integer NOT NULL, data text, a bigserial);
+INSERT INTO foreign_table_test VALUES (1, 'text_test');
+SELECT citus_add_local_table_to_metadata('foreign_table_test');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE EXTENSION postgres_fdw;
+CREATE SERVER foreign_server
+        FOREIGN DATA WRAPPER postgres_fdw
+        OPTIONS (host 'localhost', port :'master_port', dbname 'regression');
+CREATE USER MAPPING FOR CURRENT_USER
+        SERVER foreign_server
+        OPTIONS (user 'postgres');
+CREATE FOREIGN TABLE foreign_table (
+        id integer NOT NULL,
+        data text,
+        a bigserial
+)
+        SERVER foreign_server
+        OPTIONS (schema_name 'pg14', table_name 'foreign_table_test');
+SELECT citus_add_local_table_to_metadata('foreign_table');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT count(*) FROM foreign_table;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+TRUNCATE foreign_table;
+\c - - - :worker_1_port
+set search_path to pg14;
+-- verify the foreign table is truncated
+SELECT count(*) FROM pg14.foreign_table;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+-- should error out
+TRUNCATE foreign_table;
+ERROR:  truncating foreign tables that are added to metadata can only be excuted on the coordinator
+\c - - - :master_port
+-- cleanup
+set client_min_messages to error;
+drop extension postgres_fdw cascade;
+drop schema pg14 cascade;
+reset client_min_messages;
+select 1 from citus_remove_node('localhost',:master_port);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+

src/test/regress/sql/pg14.sql

@@ -642,3 +642,46 @@ SELECT count(*) FROM
 
 set client_min_messages to error;
 drop schema pg14 cascade;
+
+create schema pg14;
+set search_path to pg14;
+
+select 1 from citus_add_node('localhost',:master_port,groupid=>0);
+
+-- test adding foreign table to metadata with the guc
+-- will test truncating foreign tables later
+CREATE TABLE foreign_table_test (id integer NOT NULL, data text, a bigserial);
+INSERT INTO foreign_table_test VALUES (1, 'text_test');
+SELECT citus_add_local_table_to_metadata('foreign_table_test');
+CREATE EXTENSION postgres_fdw;
+CREATE SERVER foreign_server
+        FOREIGN DATA WRAPPER postgres_fdw
+        OPTIONS (host 'localhost', port :'master_port', dbname 'regression');
+CREATE USER MAPPING FOR CURRENT_USER
+        SERVER foreign_server
+        OPTIONS (user 'postgres');
+CREATE FOREIGN TABLE foreign_table (
+        id integer NOT NULL,
+        data text,
+        a bigserial
+)
+        SERVER foreign_server
+        OPTIONS (schema_name 'pg14', table_name 'foreign_table_test');
+SELECT citus_add_local_table_to_metadata('foreign_table');
+
+SELECT count(*) FROM foreign_table;
+TRUNCATE foreign_table;
+\c - - - :worker_1_port
+set search_path to pg14;
+-- verify the foreign table is truncated
+SELECT count(*) FROM pg14.foreign_table;
+
+-- should error out
+TRUNCATE foreign_table;
+\c - - - :master_port
+-- cleanup
+set client_min_messages to error;
+drop extension postgres_fdw cascade;
+drop schema pg14 cascade;
+reset client_min_messages;
+select 1 from citus_remove_node('localhost',:master_port);