Add TRUNCATE arbitrary config tests (#5848)

Adds TRUNCATE arbitrary config tests.
Also adds the ability to skip tests from particular configs.

src/test/regress/citus_tests/arbitrary_configs/README.md

@@ -69,6 +69,10 @@ So the infrastructure tests:
 When you want to add a new test, you can add the create statements to `create_schedule` and add the sql queries to `sql_schedule`.
 If you are adding Citus UDFs that should be a NO-OP for Postgres, make sure to override the UDFs in `postgres.sql`.
 
+If the test needs to be skipped in some configs, you can do that by adding the test names in the `skip_tests` array for
+each config. The test files associated with the skipped test will be set to empty so the test will pass without the actual test
+being run.
+
 ## Adding a new config
 
 You can add your new config to `config.py`. Make sure to extend either `CitusDefaultClusterConfig` or `CitusMXBaseClusterConfig`.

src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py

@@ -128,13 +128,24 @@ def copy_test_files(config):
                 colon_index = line.index(":")
                 line = line[colon_index + 1 :].strip()
                 test_names = line.split(" ")
-                copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path)
+                copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path, config)
 
 
-def copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path):
+def copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path, config):
     for test_name in test_names:
+        # make empty files for the skipped tests
+        if test_name in config.skip_tests:
+            expected_sql_file = os.path.join(sql_dir_path, test_name + ".sql")
+            open(expected_sql_file, 'x').close()
+
+            expected_out_file = os.path.join(expected_dir_path, test_name + ".out")
+            open(expected_out_file, 'x').close()
+
+            continue
+
         sql_name = os.path.join("./sql", test_name + ".sql")
         output_name = os.path.join("./expected", test_name + ".out")
+
         shutil.copy(sql_name, sql_dir_path)
         if os.path.isfile(output_name):
             # it might be the first time we run this test and the expected file

src/test/regress/citus_tests/config.py

@@ -109,6 +109,7 @@ class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
         self.new_settings = {}
         self.add_coordinator_to_metadata = False
         self.env_variables = {}
+        self.skip_tests = []
 
     def post_init(self):
         self._init_node_name_ports()
@@ -297,6 +298,13 @@ class CitusUnusualQuerySettingsConfig(CitusDefaultClusterConfig):
             "citus.values_materialization_threshold": "0",
         }
 
+        self.skip_tests = [
+            # Creating a reference table from a table referred to by a fk
+            # requires the table with the fk to be converted to a citus_local_table.
+            # As of c11, there is no way to do that through remote execution so this test
+            # will fail
+            "arbitrary_configs_truncate_cascade_create", "arbitrary_configs_truncate_cascade"]
+
 
 class CitusSingleNodeSingleShardClusterConfig(CitusDefaultClusterConfig):
     def __init__(self, arguments):
@@ -312,6 +320,13 @@ class CitusShardReplicationFactorClusterConfig(CitusDefaultClusterConfig):
     def __init__(self, arguments):
         super().__init__(arguments)
         self.new_settings = {"citus.shard_replication_factor": 2}
+        self.skip_tests = [
+            # citus does not support foreign keys in distributed tables
+            # when citus.shard_replication_factor > 2
+            "arbitrary_configs_truncate_partition_create", "arbitrary_configs_truncate_partition",
+            # citus does not support modifying a partition when
+            # citus.shard_replication_factor > 2
+            "arbitrary_configs_truncate_cascade_create", "arbitrary_configs_truncate_cascade"]
 
 
 class CitusSingleShardClusterConfig(CitusDefaultClusterConfig):

src/test/regress/create_schedule

@@ -7,3 +7,6 @@ test: connectivity_checks
 test: views_create
 test: sequences_create
 test: index_create
+test: arbitrary_configs_truncate_create
+test: arbitrary_configs_truncate_cascade_create
+test: arbitrary_configs_truncate_partition_create

src/test/regress/expected/arbitrary_configs_truncate.out

@@ -0,0 +1,85 @@
+SET search_path TO truncate_tests_schema;
+-- Test truncate rollback on a basic table
+SELECT COUNT(*) FROM basic_table;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+BEGIN;
+TRUNCATE basic_table;
+SELECT COUNT(*) FROM basic_table;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+ROLLBACK;
+SELECT COUNT(*) FROM basic_table;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+-- Test truncate on a basic table
+SELECT COUNT(*) FROM basic_table;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+TRUNCATE basic_table;
+SELECT COUNT(*) FROM basic_table;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+-- Test trucate rollback on partitioned table
+SELECT COUNT(*) FROM partitioned_table_0;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+BEGIN;
+TRUNCATE partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_0;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+ROLLBACK;
+SELECT COUNT(*) FROM partitioned_table_0;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+-- Test truncate a partioned table
+SELECT COUNT(*) FROM partitioned_table;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_1;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+TRUNCATE partitioned_table;
+SELECT COUNT(*) FROM partitioned_table;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_1;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+

src/test/regress/expected/arbitrary_configs_truncate_cascade.out

@@ -0,0 +1,84 @@
+SET search_path TO truncate_cascade_tests_schema;
+-- Test truncate error on table with dependencies
+TRUNCATE table_with_pk;
+ERROR:  cannot truncate a table referenced in a foreign key constraint
+DETAIL:  Table "table_with_fk_1" references "table_with_pk".
+HINT:  Truncate table "table_with_fk_1" at the same time, or use TRUNCATE ... CASCADE.
+-- Test truncate rollback on table with dependencies
+SELECT COUNT(*) FROM table_with_fk_1;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_2;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+BEGIN;
+TRUNCATE table_with_pk CASCADE;
+SELECT COUNT(*) FROM table_with_fk_1;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_2;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+ROLLBACK;
+SELECT COUNT(*) FROM table_with_fk_1;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_2;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+-- Test truncate on table with dependencies
+SELECT COUNT(*) FROM table_with_pk;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_1;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_2;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+TRUNCATE table_with_pk CASCADE;
+SELECT COUNT(*) FROM table_with_pk;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_1;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT COUNT(*) FROM table_with_fk_2;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+

src/test/regress/expected/arbitrary_configs_truncate_cascade_create.out

@@ -0,0 +1,29 @@
+CREATE SCHEMA truncate_cascade_tests_schema;
+SET search_path TO truncate_cascade_tests_schema;
+-- tables connected with foreign keys
+CREATE TABLE table_with_pk(a bigint PRIMARY KEY);
+CREATE TABLE table_with_fk_1(a bigint, b bigint, FOREIGN KEY (b) REFERENCES table_with_pk(a));
+CREATE TABLE table_with_fk_2(a bigint, b bigint, FOREIGN KEY (b) REFERENCES table_with_pk(a));
+-- distribute tables
+SELECT create_reference_table('table_with_pk');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('table_with_fk_1', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_reference_table('table_with_fk_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- fill tables with data
+INSERT INTO table_with_pk(a) SELECT n FROM generate_series(1, 10) n;
+INSERT INTO table_with_fk_1(a, b) SELECT n, n FROM generate_series(1, 10) n;
+INSERT INTO table_with_fk_2(a, b) SELECT n, n FROM generate_series(1, 10) n;

src/test/regress/expected/arbitrary_configs_truncate_create.out

@@ -0,0 +1,26 @@
+CREATE SCHEMA truncate_tests_schema;
+SET search_path TO truncate_tests_schema;
+-- simple table
+CREATE TABLE basic_table(a int);
+-- partioned table
+CREATE TABLE partitioned_table(a int) PARTITION BY RANGE(a);
+CREATE TABLE partitioned_table_0 PARTITION OF partitioned_table
+FOR VALUES FROM (1) TO (6);
+CREATE TABLE partitioned_table_1 PARTITION OF partitioned_table
+FOR VALUES FROM (6) TO (11);
+-- distribute tables
+SELECT create_distributed_table('basic_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('partitioned_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- fill tables with data
+INSERT INTO basic_table(a) SELECT n FROM generate_series(1, 10) n;
+INSERT INTO partitioned_table(a) SELECT n FROM generate_series(1, 10) n;

src/test/regress/expected/arbitrary_configs_truncate_partition.out

@@ -0,0 +1,39 @@
+SET search_path TO truncate_partition_tests_schema;
+-- Test truncate on a partition
+SELECT COUNT(*) FROM partitioned_table;
+ count
+---------------------------------------------------------------------
+    10
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_0;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_1;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+TRUNCATE partitioned_table_0;
+SELECT COUNT(*) FROM partitioned_table;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_0;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT COUNT(*) FROM partitioned_table_1;
+ count
+---------------------------------------------------------------------
+     5
+(1 row)
+

src/test/regress/expected/arbitrary_configs_truncate_partition_create.out

@@ -0,0 +1,17 @@
+CREATE SCHEMA truncate_partition_tests_schema;
+SET search_path TO truncate_partition_tests_schema;
+-- partioned table
+CREATE TABLE partitioned_table(a int) PARTITION BY RANGE(a);
+CREATE TABLE partitioned_table_0 PARTITION OF partitioned_table
+FOR VALUES FROM (1) TO (6);
+CREATE TABLE partitioned_table_1 PARTITION OF partitioned_table
+FOR VALUES FROM (6) TO (11);
+-- distribute tables
+SELECT create_distributed_table('partitioned_table', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- fill tables with data
+INSERT INTO partitioned_table(a) SELECT n FROM generate_series(1, 10) n;

src/test/regress/sql/arbitrary_configs_truncate.sql

@@ -0,0 +1,37 @@
+SET search_path TO truncate_tests_schema;
+
+-- Test truncate rollback on a basic table
+SELECT COUNT(*) FROM basic_table;
+
+BEGIN;
+TRUNCATE basic_table;
+SELECT COUNT(*) FROM basic_table;
+ROLLBACK;
+
+SELECT COUNT(*) FROM basic_table;
+
+-- Test truncate on a basic table
+SELECT COUNT(*) FROM basic_table;
+
+TRUNCATE basic_table;
+
+SELECT COUNT(*) FROM basic_table;
+
+-- Test trucate rollback on partitioned table
+SELECT COUNT(*) FROM partitioned_table_0;
+
+BEGIN;
+TRUNCATE partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_0;
+ROLLBACK;
+
+SELECT COUNT(*) FROM partitioned_table_0;
+
+-- Test truncate a partioned table
+SELECT COUNT(*) FROM partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_1;
+
+TRUNCATE partitioned_table;
+
+SELECT COUNT(*) FROM partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_1;

src/test/regress/sql/arbitrary_configs_truncate_cascade.sql

@@ -0,0 +1,28 @@
+SET search_path TO truncate_cascade_tests_schema;
+
+-- Test truncate error on table with dependencies
+TRUNCATE table_with_pk;
+
+-- Test truncate rollback on table with dependencies
+SELECT COUNT(*) FROM table_with_fk_1;
+SELECT COUNT(*) FROM table_with_fk_2;
+
+BEGIN;
+TRUNCATE table_with_pk CASCADE;
+SELECT COUNT(*) FROM table_with_fk_1;
+SELECT COUNT(*) FROM table_with_fk_2;
+ROLLBACK;
+
+SELECT COUNT(*) FROM table_with_fk_1;
+SELECT COUNT(*) FROM table_with_fk_2;
+
+-- Test truncate on table with dependencies
+SELECT COUNT(*) FROM table_with_pk;
+SELECT COUNT(*) FROM table_with_fk_1;
+SELECT COUNT(*) FROM table_with_fk_2;
+
+TRUNCATE table_with_pk CASCADE;
+
+SELECT COUNT(*) FROM table_with_pk;
+SELECT COUNT(*) FROM table_with_fk_1;
+SELECT COUNT(*) FROM table_with_fk_2;

src/test/regress/sql/arbitrary_configs_truncate_cascade_create.sql

@@ -0,0 +1,17 @@
+CREATE SCHEMA truncate_cascade_tests_schema;
+SET search_path TO truncate_cascade_tests_schema;
+
+-- tables connected with foreign keys
+CREATE TABLE table_with_pk(a bigint PRIMARY KEY);
+CREATE TABLE table_with_fk_1(a bigint, b bigint, FOREIGN KEY (b) REFERENCES table_with_pk(a));
+CREATE TABLE table_with_fk_2(a bigint, b bigint, FOREIGN KEY (b) REFERENCES table_with_pk(a));
+
+-- distribute tables
+SELECT create_reference_table('table_with_pk');
+SELECT create_distributed_table('table_with_fk_1', 'a');
+SELECT create_reference_table('table_with_fk_2');
+
+-- fill tables with data
+INSERT INTO table_with_pk(a) SELECT n FROM generate_series(1, 10) n;
+INSERT INTO table_with_fk_1(a, b) SELECT n, n FROM generate_series(1, 10) n;
+INSERT INTO table_with_fk_2(a, b) SELECT n, n FROM generate_series(1, 10) n;

src/test/regress/sql/arbitrary_configs_truncate_create.sql

@@ -0,0 +1,23 @@
+CREATE SCHEMA truncate_tests_schema;
+SET search_path TO truncate_tests_schema;
+
+-- simple table
+CREATE TABLE basic_table(a int);
+
+-- partioned table
+CREATE TABLE partitioned_table(a int) PARTITION BY RANGE(a);
+CREATE TABLE partitioned_table_0 PARTITION OF partitioned_table
+FOR VALUES FROM (1) TO (6);
+CREATE TABLE partitioned_table_1 PARTITION OF partitioned_table
+FOR VALUES FROM (6) TO (11);
+
+-- distribute tables
+SELECT create_distributed_table('basic_table', 'a');
+SELECT create_distributed_table('partitioned_table', 'a');
+
+-- fill tables with data
+INSERT INTO basic_table(a) SELECT n FROM generate_series(1, 10) n;
+INSERT INTO partitioned_table(a) SELECT n FROM generate_series(1, 10) n;
+
+
+

src/test/regress/sql/arbitrary_configs_truncate_partition.sql

@@ -0,0 +1,12 @@
+SET search_path TO truncate_partition_tests_schema;
+
+-- Test truncate on a partition
+SELECT COUNT(*) FROM partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_0;
+SELECT COUNT(*) FROM partitioned_table_1;
+
+TRUNCATE partitioned_table_0;
+
+SELECT COUNT(*) FROM partitioned_table;
+SELECT COUNT(*) FROM partitioned_table_0;
+SELECT COUNT(*) FROM partitioned_table_1;

src/test/regress/sql/arbitrary_configs_truncate_partition_create.sql

@@ -0,0 +1,15 @@
+CREATE SCHEMA truncate_partition_tests_schema;
+SET search_path TO truncate_partition_tests_schema;
+
+-- partioned table
+CREATE TABLE partitioned_table(a int) PARTITION BY RANGE(a);
+CREATE TABLE partitioned_table_0 PARTITION OF partitioned_table
+FOR VALUES FROM (1) TO (6);
+CREATE TABLE partitioned_table_1 PARTITION OF partitioned_table
+FOR VALUES FROM (6) TO (11);
+
+-- distribute tables
+SELECT create_distributed_table('partitioned_table', 'a');
+
+-- fill tables with data
+INSERT INTO partitioned_table(a) SELECT n FROM generate_series(1, 10) n;

src/test/regress/sql_schedule

@@ -7,3 +7,6 @@ test: dropped_columns_1 distributed_planning
 test: local_dist_join
 test: connectivity_checks citus_run_command
 test: sequences
+test: arbitrary_configs_truncate
+test: arbitrary_configs_truncate_cascade
+test: arbitrary_configs_truncate_partition