Honour enable_ddl_propagation in truncate trigger

src/backend/distributed/Makefile

@@ -10,7 +10,7 @@ EXTVERSIONS = 5.0 5.0-1 5.0-2  \
 	5.2-1 5.2-2 5.2-3 5.2-4 \
 	6.0-1 6.0-2 6.0-3 6.0-4 6.0-5 6.0-6 6.0-7 6.0-8 6.0-9 6.0-10 6.0-11 6.0-12 6.0-13 6.0-14 6.0-15 6.0-16 6.0-17 6.0-18 \
 	6.1-1 6.1-2 6.1-3 6.1-4 6.1-5 6.1-6 6.1-7 6.1-8 6.1-9 6.1-10 6.1-11 6.1-12 6.1-13 6.1-14 6.1-15 6.1-16 6.1-17 \
-	6.2-1 6.2-2 6.2-3
+	6.2-1 6.2-2 6.2-3 6.2-4
 
 # All citus--*.sql files in the source directory
 DATA = $(patsubst $(citus_abs_srcdir)/%.sql,%.sql,$(wildcard $(citus_abs_srcdir)/$(EXTENSION)--*--*.sql))
@@ -136,6 +136,8 @@ $(EXTENSION)--6.2-2.sql: $(EXTENSION)--6.2-1.sql $(EXTENSION)--6.2-1--6.2-2.sql
 	cat $^ > $@
 $(EXTENSION)--6.2-3.sql: $(EXTENSION)--6.2-2.sql $(EXTENSION)--6.2-2--6.2-3.sql
 	cat $^ > $@
+$(EXTENSION)--6.2-4.sql: $(EXTENSION)--6.2-3.sql $(EXTENSION)--6.2-3--6.2-4.sql
+	cat $^ > $@
 
 NO_PGXS = 1
 

src/backend/distributed/citus--6.2-3--6.2-4.sql

@@ -0,0 +1,8 @@
+/* citus--6.2-3--6.2-4.sql */
+
+CREATE OR REPLACE FUNCTION pg_catalog.citus_truncate_trigger()
+    RETURNS trigger
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$citus_truncate_trigger$$;
+COMMENT ON FUNCTION pg_catalog.citus_truncate_trigger()
+    IS 'trigger function called when truncating the distributed table';

src/backend/distributed/citus.control

@@ -1,6 +1,6 @@
 # Citus extension
 comment = 'Citus distributed database'
-default_version = '6.2-3'
+default_version = '6.2-4'
 module_pathname = '$libdir/citus'
 relocatable = false
 schema = pg_catalog

src/backend/distributed/master/master_truncate.c

@@ -0,0 +1,85 @@
+/*-------------------------------------------------------------------------
+ *
+ * master_truncate.c
+ *
+ * Routine for truncating local data after a table has been distributed.
+ *
+ * Copyright (c) 2014-2017, Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include <stddef.h>
+
+#include "commands/tablecmds.h"
+#include "commands/trigger.h"
+#include "distributed/master_metadata_utility.h"
+#include "distributed/multi_join_order.h"
+#include "distributed/multi_utility.h"
+#include "distributed/pg_dist_partition.h"
+#include "utils/builtins.h"
+#include "utils/lsyscache.h"
+#include "utils/rel.h"
+
+
+/* exports for SQL callable functions */
+PG_FUNCTION_INFO_V1(citus_truncate_trigger);
+
+
+/*
+ * citus_truncate_trigger is called as a trigger when a distributed
+ * table is truncated.
+ */
+Datum
+citus_truncate_trigger(PG_FUNCTION_ARGS)
+{
+	TriggerData *triggerData = NULL;
+	Relation truncatedRelation = NULL;
+	Oid relationId = InvalidOid;
+	char *relationName = NULL;
+	Oid schemaId = InvalidOid;
+	char *schemaName = NULL;
+	char partitionMethod = 0;
+
+	if (!CALLED_AS_TRIGGER(fcinfo))
+	{
+		ereport(ERROR, (errcode(ERRCODE_E_R_I_E_TRIGGER_PROTOCOL_VIOLATED),
+						errmsg("must be called as trigger")));
+	}
+
+	triggerData = (TriggerData *) fcinfo->context;
+	truncatedRelation = triggerData->tg_relation;
+	relationId = RelationGetRelid(truncatedRelation);
+	relationName = get_rel_name(relationId);
+	schemaId = get_rel_namespace(relationId);
+	schemaName = get_namespace_name(schemaId);
+	partitionMethod = PartitionMethod(relationId);
+
+	if (!EnableDDLPropagation)
+	{
+		PG_RETURN_DATUM(PointerGetDatum(NULL));
+	}
+
+	if (partitionMethod == DISTRIBUTE_BY_APPEND)
+	{
+		DirectFunctionCall3(master_drop_all_shards,
+							ObjectIdGetDatum(relationId),
+							CStringGetTextDatum(relationName),
+							CStringGetTextDatum(schemaName));
+	}
+	else
+	{
+		StringInfo truncateStatement = makeStringInfo();
+		char *qualifiedTableName = quote_qualified_identifier(schemaName, relationName);
+
+		appendStringInfo(truncateStatement, "TRUNCATE TABLE %s CASCADE",
+						 qualifiedTableName);
+
+		DirectFunctionCall1(master_modify_multiple_shards,
+							CStringGetTextDatum(truncateStatement->data));
+	}
+
+	PG_RETURN_DATUM(PointerGetDatum(NULL));
+}

src/include/distributed/multi_utility.h

@@ -36,5 +36,8 @@ extern List * PlanGrantStmt(GrantStmt *grantStmt);
 extern void ErrorIfUnsupportedConstraint(Relation relation, char distributionMethod,
 										 Var *distributionColumn, uint32 colocationId);
 
+extern Datum master_drop_all_shards(PG_FUNCTION_ARGS);
+extern Datum master_modify_multiple_shards(PG_FUNCTION_ARGS);
+
 
 #endif /* MULTI_UTILITY_H */

src/test/regress/expected/multi_colocation_utils.out

@@ -1,4 +1,5 @@
 ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1300000;
+ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 4;
 -- ===================================================================
 -- create test utility function
 -- ===================================================================

src/test/regress/expected/multi_extension.out

@@ -79,6 +79,7 @@ ALTER EXTENSION citus UPDATE TO '6.1-17';
 ALTER EXTENSION citus UPDATE TO '6.2-1';
 ALTER EXTENSION citus UPDATE TO '6.2-2';
 ALTER EXTENSION citus UPDATE TO '6.2-3';
+ALTER EXTENSION citus UPDATE TO '6.2-4';
 -- show running version
 SHOW citus.version;
  citus.version 

src/test/regress/expected/multi_truncate.out

@@ -295,3 +295,76 @@ SELECT shardid FROM pg_dist_shard where logicalrelid = '"a b append"'::regclass;
 (0 rows)
 
 DROP TABLE "a b append";
+-- Truncate local data only
+CREATE TABLE test_local_truncate (x int, y int);
+INSERT INTO test_local_truncate VALUES (1,2);
+SELECT create_distributed_table('test_local_truncate', 'x', colocate_with => 'none');
+NOTICE:  Copying data from local table...
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+BEGIN;
+SET LOCAL citus.enable_ddl_propagation TO off;
+TRUNCATE test_local_truncate;
+COMMIT;
+-- Ensure distributed data is not truncated
+SELECT * FROM test_local_truncate;
+ x | y 
+---+---
+ 1 | 2
+(1 row)
+
+-- Undistribute table
+SELECT master_drop_all_shards('test_local_truncate', 'pubic', 'test_local_truncate');
+ master_drop_all_shards 
+------------------------
+                     32
+(1 row)
+
+DELETE FROM pg_dist_partition WHERE logicalrelid = 'test_local_truncate'::regclass;
+-- Ensure local data is truncated
+SELECT * FROM test_local_truncate;
+ x | y 
+---+---
+(0 rows)
+
+DROP TABLE test_local_truncate;
+-- Truncate local data, but roll back
+CREATE TABLE test_local_truncate (x int, y int);
+INSERT INTO test_local_truncate VALUES (1,2);
+SELECT create_distributed_table('test_local_truncate', 'x', colocate_with => 'none');
+NOTICE:  Copying data from local table...
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+BEGIN;
+SET LOCAL citus.enable_ddl_propagation TO off;
+TRUNCATE test_local_truncate;
+ROLLBACK;
+-- Ensure distributed data is not truncated
+SELECT * FROM test_local_truncate;
+ x | y 
+---+---
+ 1 | 2
+(1 row)
+
+-- Undistribute table
+SELECT master_drop_all_shards('test_local_truncate', 'pubic', 'test_local_truncate');
+ master_drop_all_shards 
+------------------------
+                     32
+(1 row)
+
+DELETE FROM pg_dist_partition WHERE logicalrelid = 'test_local_truncate'::regclass;
+-- Ensure local data is not truncated
+SELECT * FROM test_local_truncate;
+ x | y 
+---+---
+ 1 | 2
+(1 row)
+
+DROP TABLE test_local_truncate;

src/test/regress/expected/multi_unsupported_worker_operations.out

@@ -260,8 +260,6 @@ SELECT master_remove_node('localhost', 5432);
 TRUNCATE mx_table;
 ERROR:  operation is not allowed on this node
 HINT:  Connect to the coordinator and run it again.
-CONTEXT:  SQL statement "SELECT master_modify_multiple_shards(commandText)"
-PL/pgSQL function citus_truncate_trigger() line 17 at PERFORM
 SELECT count(*) FROM mx_table;
  count 
 -------

src/test/regress/expected/multi_utilities.out

@@ -346,6 +346,8 @@ SELECT worker_hash('(1, 2)'::test_composite_type);
  -1895345704
 (1 row)
 
+SELECT citus_truncate_trigger();
+ERROR:  must be called as trigger
 -- TODO: support VERBOSE
 -- VACUUM VERBOSE dustbunnies;
 -- VACUUM (FULL, VERBOSE) dustbunnies;

src/test/regress/sql/multi_colocation_utils.sql

@@ -1,5 +1,6 @@
 
 ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1300000;
+ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 4;
 
 -- ===================================================================
 -- create test utility function

src/test/regress/sql/multi_extension.sql

@@ -80,6 +80,7 @@ ALTER EXTENSION citus UPDATE TO '6.1-17';
 ALTER EXTENSION citus UPDATE TO '6.2-1';
 ALTER EXTENSION citus UPDATE TO '6.2-2';
 ALTER EXTENSION citus UPDATE TO '6.2-3';
+ALTER EXTENSION citus UPDATE TO '6.2-4';
 
 -- show running version
 SHOW citus.version;

src/test/regress/sql/multi_truncate.sql

@@ -177,3 +177,47 @@ TRUNCATE TABLE "a b append";
 SELECT shardid FROM pg_dist_shard where logicalrelid = '"a b append"'::regclass;
 
 DROP TABLE "a b append";
+
+-- Truncate local data only
+CREATE TABLE test_local_truncate (x int, y int);
+INSERT INTO test_local_truncate VALUES (1,2);
+SELECT create_distributed_table('test_local_truncate', 'x', colocate_with => 'none');
+
+BEGIN;
+SET LOCAL citus.enable_ddl_propagation TO off;
+TRUNCATE test_local_truncate;
+COMMIT;
+
+-- Ensure distributed data is not truncated
+SELECT * FROM test_local_truncate;
+
+-- Undistribute table
+SELECT master_drop_all_shards('test_local_truncate', 'pubic', 'test_local_truncate');
+DELETE FROM pg_dist_partition WHERE logicalrelid = 'test_local_truncate'::regclass;
+
+-- Ensure local data is truncated
+SELECT * FROM test_local_truncate;
+
+DROP TABLE test_local_truncate;
+
+-- Truncate local data, but roll back
+CREATE TABLE test_local_truncate (x int, y int);
+INSERT INTO test_local_truncate VALUES (1,2);
+SELECT create_distributed_table('test_local_truncate', 'x', colocate_with => 'none');
+
+BEGIN;
+SET LOCAL citus.enable_ddl_propagation TO off;
+TRUNCATE test_local_truncate;
+ROLLBACK;
+
+-- Ensure distributed data is not truncated
+SELECT * FROM test_local_truncate;
+
+-- Undistribute table
+SELECT master_drop_all_shards('test_local_truncate', 'pubic', 'test_local_truncate');
+DELETE FROM pg_dist_partition WHERE logicalrelid = 'test_local_truncate'::regclass;
+
+-- Ensure local data is not truncated
+SELECT * FROM test_local_truncate;
+
+DROP TABLE test_local_truncate;

src/test/regress/sql/multi_utilities.sql

@@ -222,6 +222,8 @@ SELECT worker_hash('1997-08-08'::date);
 SELECT worker_hash('(1, 2)');
 SELECT worker_hash('(1, 2)'::test_composite_type);
 
+SELECT citus_truncate_trigger();
+
 -- TODO: support VERBOSE
 -- VACUUM VERBOSE dustbunnies;
 -- VACUUM (FULL, VERBOSE) dustbunnies;