-- ===================================================================
-- create test functions
-- ===================================================================
CREATE FUNCTION generate_alter_table_detach_partition_command(regclass)
    RETURNS text
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION generate_alter_table_attach_partition_command(regclass)
    RETURNS text
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION generate_partition_information(regclass)
    RETURNS text
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION print_partitions(regclass)
    RETURNS text
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION table_inherits(regclass)
    RETURNS bool
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION table_inherited(regclass)
    RETURNS bool
    AS 'citus'
    LANGUAGE C STRICT;
CREATE OR REPLACE FUNCTION detach_and_attach_partition(partition_name regclass, parent_table_name regclass)
RETURNS void LANGUAGE plpgsql VOLATILE
AS $function$
DECLARE
  detach_partition_command text := '';
	attach_partition_command text := '';
	command_result text := '';

BEGIN
   -- first generate the command
   SELECT public.generate_alter_table_attach_partition_command(partition_name) INTO attach_partition_command;

   -- now genereate the detach command
   SELECT public.generate_alter_table_detach_partition_command(partition_name) INTO detach_partition_command;

   -- later detach the same partition
   EXECUTE detach_partition_command;

   -- not attach it again
   EXECUTE attach_partition_command;
END;
$function$;
CREATE OR REPLACE FUNCTION drop_and_recreate_partitioned_table(parent_table_name regclass)
RETURNS void LANGUAGE plpgsql VOLATILE
AS $function$
DECLARE
	command text := '';

BEGIN
   -- first generate the command
   CREATE TABLE partitioned_table_create_commands AS SELECT master_get_table_ddl_events(parent_table_name::text);

   -- later detach the same partition
   EXECUTE 'DROP TABLE ' || parent_table_name::text || ';';

   FOR command IN SELECT * FROM partitioned_table_create_commands
   LOOP
        -- can do some processing here
   		EXECUTE command;
    END LOOP;

    DROP TABLE partitioned_table_create_commands;

END;
$function$;
-- create a partitioned table
CREATE TABLE date_partitioned_table(id int, time date) PARTITION BY RANGE (time);
-- we should be able to get the partitioning information even if there are no partitions
SELECT generate_partition_information('date_partitioned_table');
 generate_partition_information
---------------------------------------------------------------------
 RANGE ("time")
(1 row)

-- we should be able to drop and re-create the partitioned table using the command that Citus generate
SELECT drop_and_recreate_partitioned_table('date_partitioned_table');
 drop_and_recreate_partitioned_table
---------------------------------------------------------------------

(1 row)

-- we should also be able to see the PARTITION BY ... for the parent table
SELECT master_get_table_ddl_events('date_partitioned_table');
                                    master_get_table_ddl_events
---------------------------------------------------------------------
 CREATE TABLE public.date_partitioned_table (id integer, "time" date) PARTITION BY RANGE ("time")
 ALTER TABLE public.date_partitioned_table OWNER TO postgres
(2 rows)

-- now create the partitions
CREATE TABLE date_partition_2006 PARTITION OF date_partitioned_table FOR VALUES FROM ('2006-01-01') TO ('2007-01-01');
CREATE TABLE date_partition_2007 PARTITION OF date_partitioned_table FOR VALUES FROM ('2007-01-01') TO ('2008-01-01');
-- we should be able to get the partitioning information after the partitions are created
SELECT generate_partition_information('date_partitioned_table');
 generate_partition_information
---------------------------------------------------------------------
 RANGE ("time")
(1 row)

-- lets get the attach partition commands
SELECT generate_alter_table_attach_partition_command('date_partition_2006');
                                              generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE public.date_partitioned_table ATTACH PARTITION public.date_partition_2006 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007');
(1 row)

SELECT generate_alter_table_attach_partition_command('date_partition_2007');
                                              generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE public.date_partitioned_table ATTACH PARTITION public.date_partition_2007 FOR VALUES FROM ('01-01-2007') TO ('01-01-2008');
(1 row)

-- detach and attach the partition by the command generated by us
\d+ date_partitioned_table
                    Partitioned table "public.date_partitioned_table"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           |          |         | plain   |              |
 time   | date    |           |          |         | plain   |              |
Partition key: RANGE ("time")
Partitions: date_partition_2006 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007'),
            date_partition_2007 FOR VALUES FROM ('01-01-2007') TO ('01-01-2008')

SELECT detach_and_attach_partition('date_partition_2007', 'date_partitioned_table');
 detach_and_attach_partition
---------------------------------------------------------------------

(1 row)

-- check that both partitions are visiable
\d+ date_partitioned_table
                    Partitioned table "public.date_partitioned_table"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           |          |         | plain   |              |
 time   | date    |           |          |         | plain   |              |
Partition key: RANGE ("time")
Partitions: date_partition_2006 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007'),
            date_partition_2007 FOR VALUES FROM ('01-01-2007') TO ('01-01-2008')

-- make sure that inter shard commands work as expected
-- assume that the shardId is 100
CREATE TABLE date_partitioned_table_100 (id int, time date) PARTITION BY RANGE (time);
CREATE TABLE date_partition_2007_100 (id int, time date );
-- now create the partitioning hierarcy
SELECT worker_apply_inter_shard_ddl_command(referencing_shard:=100, referencing_schema_name:='public',
											referenced_shard:=100, referenced_schema_name:='public',
											command:='ALTER TABLE date_partitioned_table ATTACH PARTITION date_partition_2007 FOR VALUES FROM (''2007-01-01'') TO (''2008-01-02'')' );
 worker_apply_inter_shard_ddl_command
---------------------------------------------------------------------

(1 row)

-- the hierarcy is successfully created
\d+ date_partitioned_table_100
                  Partitioned table "public.date_partitioned_table_100"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           |          |         | plain   |              |
 time   | date    |           |          |         | plain   |              |
Partition key: RANGE ("time")
Partitions: date_partition_2007_100 FOR VALUES FROM ('01-01-2007') TO ('01-02-2008')

-- Citus can also get the DDL events for the partitions as regular tables
SELECT master_get_table_ddl_events('date_partition_2007_100');
                           master_get_table_ddl_events
---------------------------------------------------------------------
 CREATE TABLE public.date_partition_2007_100 (id integer, "time" date) USING heap
 ALTER TABLE public.date_partition_2007_100 OWNER TO postgres
(2 rows)

-- now break the partitioning hierarcy
SELECT worker_apply_inter_shard_ddl_command(referencing_shard:=100, referencing_schema_name:='public',
                      referenced_shard:=100, referenced_schema_name:='public',
                      command:='ALTER TABLE date_partitioned_table DETACH PARTITION date_partition_2007' );
 worker_apply_inter_shard_ddl_command
---------------------------------------------------------------------

(1 row)

-- the hierarcy is successfully broken
\d+ date_partitioned_table_100
                  Partitioned table "public.date_partitioned_table_100"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           |          |         | plain   |              |
 time   | date    |           |          |         | plain   |              |
Partition key: RANGE ("time")
Number of partitions: 0

-- now lets have some more complex partitioning hierarcies with
-- tables on different schemas and constraints on the tables
CREATE SCHEMA partition_parent_schema;
CREATE TABLE partition_parent_schema.parent_table (id int NOT NULL, time date DEFAULT now()) PARTITION BY RANGE (time);
CREATE SCHEMA partition_child_1_schema;
CREATE TABLE partition_child_1_schema.child_1 (id int NOT NULL, time date );
CREATE SCHEMA partition_child_2_schema;
CREATE TABLE partition_child_2_schema.child_2 (id int NOT NULL, time date );
-- we should be able to get the partitioning information even if there are no partitions
SELECT generate_partition_information('partition_parent_schema.parent_table');
 generate_partition_information
---------------------------------------------------------------------
 RANGE ("time")
(1 row)

-- we should be able to drop and re-create the partitioned table using the command that Citus generate
SELECT drop_and_recreate_partitioned_table('partition_parent_schema.parent_table');
 drop_and_recreate_partitioned_table
---------------------------------------------------------------------

(1 row)

ALTER TABLE partition_parent_schema.parent_table ATTACH PARTITION partition_child_1_schema.child_1 FOR VALUES FROM ('2009-01-01') TO ('2010-01-02');
SET search_path = 'partition_parent_schema';
ALTER TABLE parent_table ATTACH PARTITION  partition_child_2_schema.child_2 FOR VALUES FROM  ('2006-01-01') TO ('2007-01-01');
SELECT public.generate_partition_information('parent_table');
 generate_partition_information
---------------------------------------------------------------------
 RANGE ("time")
(1 row)

-- lets get the attach partition commands
SELECT public.generate_alter_table_attach_partition_command('partition_child_1_schema.child_1');
                                                    generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE partition_parent_schema.parent_table ATTACH PARTITION partition_child_1_schema.child_1 FOR VALUES FROM ('01-01-2009') TO ('01-02-2010');
(1 row)

SET search_path = 'partition_child_2_schema';
SELECT public.generate_alter_table_attach_partition_command('child_2');
                                                    generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE partition_parent_schema.parent_table ATTACH PARTITION partition_child_2_schema.child_2 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007');
(1 row)

SET search_path = 'partition_parent_schema';
-- detach and attach the partition by the command generated by us
\d+ parent_table
                 Partitioned table "partition_parent_schema.parent_table"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           | not null |         | plain   |              |
 time   | date    |           |          | now()   | plain   |              |
Partition key: RANGE ("time")
Partitions: partition_child_1_schema.child_1 FOR VALUES FROM ('01-01-2009') TO ('01-02-2010'),
            partition_child_2_schema.child_2 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007')

SELECT public.detach_and_attach_partition('partition_child_1_schema.child_1', 'parent_table');
 detach_and_attach_partition
---------------------------------------------------------------------

(1 row)

-- check that both partitions are visiable
\d+ parent_table
                 Partitioned table "partition_parent_schema.parent_table"
 Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description
---------------------------------------------------------------------
 id     | integer |           | not null |         | plain   |              |
 time   | date    |           |          | now()   | plain   |              |
Partition key: RANGE ("time")
Partitions: partition_child_1_schema.child_1 FOR VALUES FROM ('01-01-2009') TO ('01-02-2010'),
            partition_child_2_schema.child_2 FOR VALUES FROM ('01-01-2006') TO ('01-01-2007')

-- some very simple checks that should error out
SELECT public.generate_alter_table_attach_partition_command('parent_table');
ERROR:  "parent_table" is not a partition
SELECT public.generate_partition_information('partition_child_1_schema.child_1');
ERROR:  "child_1" is not a parent table
SELECT public.print_partitions('partition_child_1_schema.child_1');
ERROR:  "child_1" is not a parent table
-- now print the partitions
SELECT public.print_partitions('parent_table');
 print_partitions
---------------------------------------------------------------------
 child_1,child_2
(1 row)

SET search_path = 'public';
-- test multi column / expression partitioning with UNBOUNDED ranges
CREATE OR REPLACE FUNCTION some_function(input_val text)
RETURNS text LANGUAGE plpgsql IMMUTABLE
AS $function$
BEGIN
		return reverse(input_val);
END;
$function$;
CREATE TABLE multi_column_partitioned (
   a int,
   b int,
   c text
 ) PARTITION BY RANGE (a, (a+b+1), some_function(upper(c)));
CREATE TABLE multi_column_partition_1(
   a int,
   b int,
   c text
);
CREATE TABLE multi_column_partition_2(
   a int,
   b int,
   c text
);
-- partitioning information
SELECT generate_partition_information('multi_column_partitioned');
           generate_partition_information
---------------------------------------------------------------------
 RANGE (a, (((a + b) + 1)), some_function(upper(c)))
(1 row)

SELECT master_get_table_ddl_events('multi_column_partitioned');
                                                             master_get_table_ddl_events
---------------------------------------------------------------------
 CREATE TABLE public.multi_column_partitioned (a integer, b integer, c text) PARTITION BY RANGE (a, (((a + b) + 1)), public.some_function(upper(c)))
 ALTER TABLE public.multi_column_partitioned OWNER TO postgres
(2 rows)

SELECT drop_and_recreate_partitioned_table('multi_column_partitioned');
 drop_and_recreate_partitioned_table
---------------------------------------------------------------------

(1 row)

-- partitions and their ranges
ALTER TABLE multi_column_partitioned ATTACH PARTITION multi_column_partition_1 FOR VALUES FROM (1, 10, '250') TO (1, 20, '250');
SELECT generate_alter_table_attach_partition_command('multi_column_partition_1');
                                                 generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE public.multi_column_partitioned ATTACH PARTITION public.multi_column_partition_1 FOR VALUES FROM (1, 10, '250') TO (1, 20, '250');
(1 row)

ALTER TABLE multi_column_partitioned ATTACH PARTITION multi_column_partition_2 FOR VALUES FROM (10, 1000, '2500') TO (MAXVALUE, MAXVALUE, MAXVALUE);
SELECT generate_alter_table_attach_partition_command('multi_column_partition_2');
                                                           generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE public.multi_column_partitioned ATTACH PARTITION public.multi_column_partition_2 FOR VALUES FROM (10, 1000, '2500') TO (MAXVALUE, MAXVALUE, MAXVALUE);
(1 row)

SELECT generate_alter_table_detach_partition_command('multi_column_partition_2');
                              generate_alter_table_detach_partition_command
---------------------------------------------------------------------
 ALTER TABLE IF EXISTS public.multi_column_partitioned DETACH PARTITION public.multi_column_partition_2;
(1 row)

-- finally a test with LIST partitioning
CREATE TABLE list_partitioned (col1 NUMERIC, col2 NUMERIC, col3 VARCHAR(10)) PARTITION BY LIST (col1) ;
SELECT generate_partition_information('list_partitioned');
 generate_partition_information
---------------------------------------------------------------------
 LIST (col1)
(1 row)

SELECT master_get_table_ddl_events('list_partitioned');
                                               master_get_table_ddl_events
---------------------------------------------------------------------
 CREATE TABLE public.list_partitioned (col1 numeric, col2 numeric, col3 character varying(10)) PARTITION BY LIST (col1)
 ALTER TABLE public.list_partitioned OWNER TO postgres
(2 rows)

SELECT drop_and_recreate_partitioned_table('list_partitioned');
 drop_and_recreate_partitioned_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE list_partitioned_1 PARTITION OF list_partitioned FOR VALUES IN (100, 101, 102, 103, 104);
SELECT generate_alter_table_attach_partition_command('list_partitioned_1');
                                           generate_alter_table_attach_partition_command
---------------------------------------------------------------------
 ALTER TABLE public.list_partitioned ATTACH PARTITION public.list_partitioned_1 FOR VALUES IN ('100', '101', '102', '103', '104');
(1 row)

-- also differentiate partitions and inhereted tables
CREATE TABLE cities (
    name            text,
    population      float,
    altitude        int     -- in feet
);
CREATE TABLE capitals (
    state           char(2)
) INHERITS (cities);
-- returns true since capitals inherits from cities
SELECT table_inherits('capitals');
 table_inherits
---------------------------------------------------------------------
 t
(1 row)

-- although date_partition_2006 inherits from its parent
-- returns false since the hierarcy is formed via partitioning
SELECT table_inherits('date_partition_2006');
 table_inherits
---------------------------------------------------------------------
 f
(1 row)

-- returns true since cities inherited by capitals
SELECT table_inherited('cities');
 table_inherited
---------------------------------------------------------------------
 t
(1 row)

-- although date_partitioned_table inherited by its partitions
-- returns false since the hierarcy is formed via partitioning
SELECT table_inherited('date_partitioned_table');
 table_inherited
---------------------------------------------------------------------
 f
(1 row)

-- also these are not supported
SELECT master_get_table_ddl_events('capitals');
ERROR:  capitals is not a regular, foreign or partitioned table
SELECT master_get_table_ddl_events('cities');
ERROR:  cities is not a regular, foreign or partitioned table
-- dropping parents frop the partitions
DROP TABLE date_partitioned_table, multi_column_partitioned, list_partitioned, partition_parent_schema.parent_table, cities, capitals;