Introduce expire_old_time_partitions to drop partitions older than given threshold

src/backend/distributed/sql/citus--10.1-1--10.2-1.sql

@@ -16,3 +16,4 @@ ALTER TABLE pg_catalog.pg_dist_placement ADD CONSTRAINT placement_shardid_groupi
 #include "udfs/citus_internal_update_placement_metadata/10.2-1.sql";
 #include "udfs/citus_internal_delete_shard_metadata/10.2-1.sql";
 #include "udfs/citus_internal_update_relation_colocation/10.2-1.sql";
+#include "udfs/expire_old_time_partitions/10.2-1.sql";

src/backend/distributed/sql/downgrades/citus--10.2-1--10.1-1.sql

@@ -17,6 +17,7 @@ DROP FUNCTION pg_catalog.citus_internal_add_placement_metadata(bigint, integer,
 DROP FUNCTION pg_catalog.citus_internal_update_placement_metadata(bigint, integer, integer);
 DROP FUNCTION pg_catalog.citus_internal_delete_shard_metadata(bigint);
 DROP FUNCTION pg_catalog.citus_internal_update_relation_colocation(oid, integer);
+DROP PROCEDURE pg_catalog.expire_old_time_partitions(oid, timestamptz);
 
 REVOKE ALL ON FUNCTION pg_catalog.worker_record_sequence_dependency(regclass,regclass,name) FROM PUBLIC;
 ALTER TABLE pg_catalog.pg_dist_placement DROP CONSTRAINT placement_shardid_groupid_unique_index;

src/backend/distributed/sql/udfs/expire_old_time_partitions/10.2-1.sql

@@ -0,0 +1,39 @@
+-- Heavily inspired by the procedure alter_old_partitions_set_access_method
+CREATE OR REPLACE PROCEDURE pg_catalog.expire_old_time_partitions(
+		parent_table_name regclass,
+		older_than timestamptz)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    r record;
+BEGIN
+	/* first check whether we can convert all the to_value's to timestamptz */
+	BEGIN
+		PERFORM
+		FROM pg_catalog.time_partitions
+		WHERE parent_table = parent_table_name
+		AND to_value IS NOT NULL
+		AND to_value::timestamptz <= older_than;
+	EXCEPTION WHEN invalid_datetime_format THEN
+		RAISE 'partition column of % cannot be cast to a timestamptz', parent_table_name;
+	END;
+
+	/* now drop the partitions in separate transactions */
+    FOR r IN
+		SELECT partition, from_value, to_value
+		FROM pg_catalog.time_partitions
+		WHERE parent_table = parent_table_name
+		AND to_value IS NOT NULL
+		AND to_value::timestamptz <= older_than
+		ORDER BY to_value::timestamptz
+    LOOP
+        RAISE NOTICE 'dropping % with start time % and end time %', r.partition, r.from_value, r.to_value;
+        EXECUTE format('DROP TABLE %I', r.partition);
+        COMMIT;
+    END LOOP;
+END;
+$$;
+COMMENT ON PROCEDURE pg_catalog.expire_old_time_partitions(
+		parent_table_name regclass,
+		older_than timestamptz)
+IS 'drop old partitions of a time-partitioned table';

src/backend/distributed/sql/udfs/expire_old_time_partitions/latest.sql

@@ -0,0 +1,39 @@
+-- Heavily inspired by the procedure alter_old_partitions_set_access_method
+CREATE OR REPLACE PROCEDURE pg_catalog.expire_old_time_partitions(
+		parent_table_name regclass,
+		older_than timestamptz)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    r record;
+BEGIN
+	/* first check whether we can convert all the to_value's to timestamptz */
+	BEGIN
+		PERFORM
+		FROM pg_catalog.time_partitions
+		WHERE parent_table = parent_table_name
+		AND to_value IS NOT NULL
+		AND to_value::timestamptz <= older_than;
+	EXCEPTION WHEN invalid_datetime_format THEN
+		RAISE 'partition column of % cannot be cast to a timestamptz', parent_table_name;
+	END;
+
+	/* now drop the partitions in separate transactions */
+    FOR r IN
+		SELECT partition, from_value, to_value
+		FROM pg_catalog.time_partitions
+		WHERE parent_table = parent_table_name
+		AND to_value IS NOT NULL
+		AND to_value::timestamptz <= older_than
+		ORDER BY to_value::timestamptz
+    LOOP
+        RAISE NOTICE 'dropping % with start time % and end time %', r.partition, r.from_value, r.to_value;
+        EXECUTE format('DROP TABLE %I', r.partition);
+        COMMIT;
+    END LOOP;
+END;
+$$;
+COMMENT ON PROCEDURE pg_catalog.expire_old_time_partitions(
+		parent_table_name regclass,
+		older_than timestamptz)
+IS 'drop old partitions of a time-partitioned table';

src/test/regress/expected/multi_partitioning.out

@@ -2210,6 +2210,81 @@ SELECT worker_partitioned_relation_total_size(oid) FROM pg_class WHERE relname L
 
  \c - - - :master_port
 DROP TABLE "events.Energy Added";
+-- test expire_old_time_partitions
+-- test with date partitioned table
+CREATE TABLE date_partitioned_table_to_exp (event_date date, event int) partition by range (event_date);
+SELECT create_distributed_table('date_partitioned_table_to_exp', 'event');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE date_partitioned_table_to_exp_d00 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2000-01-01') TO ('2009-12-31');
+CREATE TABLE date_partitioned_table_to_exp_d10 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2010-01-01') TO ('2019-12-31');
+CREATE TABLE date_partitioned_table_to_exp_d20 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2020-01-01') TO ('2029-12-31');
+INSERT INTO date_partitioned_table_to_exp VALUES ('2005-01-01', 1);
+INSERT INTO date_partitioned_table_to_exp VALUES ('2015-01-01', 2);
+INSERT INTO date_partitioned_table_to_exp VALUES ('2025-01-01', 3);
+\set VERBOSITY terse
+-- expire no partitions
+CALL expire_old_time_partitions('date_partitioned_table_to_exp', '1999-01-01');
+SELECT partition FROM time_partitions WHERE parent_table = 'date_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+             partition
+---------------------------------------------------------------------
+ date_partitioned_table_to_exp_d00
+ date_partitioned_table_to_exp_d10
+ date_partitioned_table_to_exp_d20
+(3 rows)
+
+-- expire 2 old partitions
+CALL expire_old_time_partitions('date_partitioned_table_to_exp', '2021-01-01');
+NOTICE:  dropping date_partitioned_table_to_exp_d00 with start time 01-01-2000 and end time 12-31-2009
+NOTICE:  dropping date_partitioned_table_to_exp_d10 with start time 01-01-2010 and end time 12-31-2019
+SELECT partition FROM time_partitions WHERE parent_table = 'date_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+             partition
+---------------------------------------------------------------------
+ date_partitioned_table_to_exp_d20
+(1 row)
+
+\set VERBOSITY default
+DROP TABLE date_partitioned_table_to_exp;
+-- test with timestamptz partitioned table
+CREATE TABLE tstz_partitioned_table_to_exp (event_time timestamptz, event int) partition by range (event_time);
+SELECT create_distributed_table('tstz_partitioned_table_to_exp', 'event');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE tstz_partitioned_table_to_exp_d0 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 02:00:00+00') TO ('2021-01-01 06:00:00+00');
+CREATE TABLE tstz_partitioned_table_to_exp_d1 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 06:00:00+00') TO ('2021-01-01 10:00:00+00');
+CREATE TABLE tstz_partitioned_table_to_exp_d2 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 10:00:00+00') TO ('2021-01-01 14:00:00+00');
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 03:00:00+00', 1);
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 09:00:00+00', 2);
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 13:00:00+00', 3);
+\set VERBOSITY terse
+-- expire no partitions
+CALL expire_old_time_partitions('tstz_partitioned_table_to_exp', '2021-01-01 01:00:00+00');
+SELECT partition FROM time_partitions WHERE parent_table = 'tstz_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+            partition
+---------------------------------------------------------------------
+ tstz_partitioned_table_to_exp_d0
+ tstz_partitioned_table_to_exp_d1
+ tstz_partitioned_table_to_exp_d2
+(3 rows)
+
+-- expire 2 old partitions
+CALL expire_old_time_partitions('tstz_partitioned_table_to_exp', '2021-01-01 12:00:00+00');
+NOTICE:  dropping tstz_partitioned_table_to_exp_d0 with start time Thu Dec 31 18:00:00 2020 PST and end time Thu Dec 31 22:00:00 2020 PST
+NOTICE:  dropping tstz_partitioned_table_to_exp_d1 with start time Thu Dec 31 22:00:00 2020 PST and end time Fri Jan 01 02:00:00 2021 PST
+SELECT partition FROM time_partitions WHERE parent_table = 'tstz_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+            partition
+---------------------------------------------------------------------
+ tstz_partitioned_table_to_exp_d2
+(1 row)
+
+\set VERBOSITY default
+DROP TABLE tstz_partitioned_table_to_exp;
 DROP SCHEMA partitioning_schema CASCADE;
 NOTICE:  drop cascades to table partitioning_schema."schema-test"
 DROP TABLE IF EXISTS

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -120,6 +120,7 @@ ORDER BY 1;
  function distributed_tables_colocated(regclass,regclass)
  function dump_global_wait_edges()
  function dump_local_wait_edges()
+ function expire_old_time_partitions(regclass, timestamp with time zone)
  function fetch_intermediate_results(text[],text,integer)
  function fix_pre_citus10_partitioned_table_constraint_names()
  function fix_pre_citus10_partitioned_table_constraint_names(regclass)
@@ -254,5 +255,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(238 rows)
+(239 rows)
 

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -109,6 +109,7 @@ ORDER BY 1;
  function distributed_tables_colocated(regclass,regclass)
  function dump_global_wait_edges()
  function dump_local_wait_edges()
+ function expire_old_time_partitions(regclass, timestamp with time zone)
  function fetch_intermediate_results(text[],text,integer)
  function fix_pre_citus10_partitioned_table_constraint_names()
  function fix_pre_citus10_partitioned_table_constraint_names(regclass)
@@ -240,5 +241,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(224 rows)
+(225 rows)
 

src/test/regress/sql/multi_partitioning.sql

@@ -1303,6 +1303,55 @@ SELECT worker_partitioned_relation_total_size(oid) FROM pg_class WHERE relname L
  \c - - - :master_port
 DROP TABLE "events.Energy Added";
 
+-- test expire_old_time_partitions
+-- test with date partitioned table
+CREATE TABLE date_partitioned_table_to_exp (event_date date, event int) partition by range (event_date);
+SELECT create_distributed_table('date_partitioned_table_to_exp', 'event');
+
+CREATE TABLE date_partitioned_table_to_exp_d00 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2000-01-01') TO ('2009-12-31');
+CREATE TABLE date_partitioned_table_to_exp_d10 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2010-01-01') TO ('2019-12-31');
+CREATE TABLE date_partitioned_table_to_exp_d20 PARTITION OF date_partitioned_table_to_exp FOR VALUES FROM ('2020-01-01') TO ('2029-12-31');
+INSERT INTO date_partitioned_table_to_exp VALUES ('2005-01-01', 1);
+INSERT INTO date_partitioned_table_to_exp VALUES ('2015-01-01', 2);
+INSERT INTO date_partitioned_table_to_exp VALUES ('2025-01-01', 3);
+
+\set VERBOSITY terse
+
+-- expire no partitions
+CALL expire_old_time_partitions('date_partitioned_table_to_exp', '1999-01-01');
+SELECT partition FROM time_partitions WHERE parent_table = 'date_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+
+-- expire 2 old partitions
+CALL expire_old_time_partitions('date_partitioned_table_to_exp', '2021-01-01');
+SELECT partition FROM time_partitions WHERE parent_table = 'date_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+
+\set VERBOSITY default
+DROP TABLE date_partitioned_table_to_exp;
+
+-- test with timestamptz partitioned table
+CREATE TABLE tstz_partitioned_table_to_exp (event_time timestamptz, event int) partition by range (event_time);
+SELECT create_distributed_table('tstz_partitioned_table_to_exp', 'event');
+
+CREATE TABLE tstz_partitioned_table_to_exp_d0 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 02:00:00+00') TO ('2021-01-01 06:00:00+00');
+CREATE TABLE tstz_partitioned_table_to_exp_d1 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 06:00:00+00') TO ('2021-01-01 10:00:00+00');
+CREATE TABLE tstz_partitioned_table_to_exp_d2 PARTITION OF tstz_partitioned_table_to_exp FOR VALUES FROM ('2021-01-01 10:00:00+00') TO ('2021-01-01 14:00:00+00');
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 03:00:00+00', 1);
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 09:00:00+00', 2);
+INSERT INTO tstz_partitioned_table_to_exp VALUES ('2021-01-01 13:00:00+00', 3);
+
+\set VERBOSITY terse
+
+-- expire no partitions
+CALL expire_old_time_partitions('tstz_partitioned_table_to_exp', '2021-01-01 01:00:00+00');
+SELECT partition FROM time_partitions WHERE parent_table = 'tstz_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+
+-- expire 2 old partitions
+CALL expire_old_time_partitions('tstz_partitioned_table_to_exp', '2021-01-01 12:00:00+00');
+SELECT partition FROM time_partitions WHERE parent_table = 'tstz_partitioned_table_to_exp'::regclass ORDER BY partition::text;
+
+\set VERBOSITY default
+DROP TABLE tstz_partitioned_table_to_exp;
+
 DROP SCHEMA partitioning_schema CASCADE;
 DROP TABLE IF EXISTS
 	partitioning_hash_test,