postgres table partitioning auto generation script - including constraints and indexes

table_partitioning_auto_generation_script.sql

--
-- Update_partitions - Takes a begin time, schema name, primary (parent) table name,
--                     table owner, the name of the date column,
--                     and if we want 'week'ly or 'month'ly partitions,
--					   is_first_contains_prior - should first table contain all prior dates
--                     The number of created tables is returned.
--                     ex: SELECT public.update_partitions('2010-02-01','my_schema','my_data','postgres','create_date','week',FALSE)
--


-- Function: public.update_partitions(timestamp without time zone, text, text, text, text, text)

-- DROP FUNCTION public.update_partitions(timestamp without time zone, text, text, text, text, text);

CREATE OR REPLACE FUNCTION public.update_partitions(begin_time timestamp without time zone, schema_name text, primary_table_name text, table_owner text, date_column text, plan text, is_first_contains_prior boolean)
  RETURNS integer AS
$BODY$
declare startTime timestamp;
declare endTime timestamp;
declare intervalTime timestamp;
declare createStmts text;
declare executeText text;
declare createTrigger text;
declare fullTablename text;
declare unknownTablename text;
declare triggerName text;
declare createdTables integer;
declare dateFormat text;
declare planInterval interval;
declare index_record RECORD;
declare constraint_record RECORD;
declare isFirst BOOLEAN;
declare unkownTableCreated BOOLEAN;
declare needToCreateIndexForDateColumn BOOLEAN;


BEGIN
dateFormat:=CASE WHEN plan='month' THEN 'YYYYMM'
                 WHEN plan='week'  THEN 'IYYYIW'
                 WHEN plan='day'   THEN 'YYYYDDD'
		 WHEN plan='year'  THEN 'YYYY'
                 ELSE 'error'
            END;
IF dateFormat='error' THEN
  RAISE EXCEPTION 'Non valid plan --> %', plan;
END IF;
-- Store the incoming begin_time, and set the endTime to one month/week/day in the future 
--     (this allows use of a cronjob at any time during the month/week/day to generate next month/week/day's table)
startTime:=(date_trunc(plan,begin_time));
planInterval:=('1 '||plan)::interval;
endTime:=(date_trunc(plan,(current_timestamp + planInterval)));
createdTables:=0;


-- Begin creating the trigger function, we're going to generate it backwards.
createTrigger:=' 
        ELSE
            INSERT INTO '||schema_name||'.'||primary_table_name||'_unknowns VALUES (NEW.*);
        END IF;
        RETURN NULL;
    END;
    $$
    LANGUAGE plpgsql;';
isFirst:= TRUE;
unkownTableCreated:=FALSE;
WHILE (startTime <= endTime) LOOP

	IF (NOT unkownTableCreated) THEN
		fullTablename:= primary_table_name||'_unknowns';
	ELSE 
		fullTablename:=primary_table_name||'_'||to_char(startTime,dateFormat);
		intervalTime:= startTime + planInterval;
	END IF;
	
	-- The table creation sql statement
	IF NOT EXISTS(SELECT * FROM information_schema.tables WHERE table_schema = schema_name AND table_name = fullTablename) THEN
		RAISE INFO '**** table doesn''t exists creating it: %', fullTablename;
		
		IF (NOT unkownTableCreated) THEN
			createStmts :='CREATE TABLE '||schema_name||'.'||fullTablename||' (
				) INHERITS ('||schema_name||'.'||primary_table_name||');'; 
		ELSE 
			createStmts:='CREATE TABLE '||schema_name||'.'||fullTablename||' (
				  CHECK ('|| CASE WHEN isFirst AND is_first_contains_prior THEN 
								'' ELSE date_column||' >= '''||startTime||''' AND ' END
								||date_column||' < '''||intervalTime||''')
				  ) INHERITS ('||schema_name||'.'||primary_table_name||')';
		END IF;

		
		-- Run the table creation
		EXECUTE createStmts;

		-- Set the table owner
		createStmts :='ALTER TABLE '||schema_name||'.'||fullTablename||' OWNER TO '||table_owner||';';
		EXECUTE createStmts;
     
		-----------------------------------------------------------------------------
		-------------------------- AUTO GENERATE CONSTRAINTS ------------------------
		FOR constraint_record IN 
			SELECT nspname, 
				REPLACE(conname,primary_table_name,fullTablename) AS full_constraint_name, 
				pg_get_constraintdef(pg_constraint.oid) AS constraintdef 
			FROM pg_constraint 
			INNER JOIN pg_class ON conrelid=pg_class.oid
			INNER JOIN pg_namespace ON pg_namespace.oid=pg_class.relnamespace
			WHERE relname = primary_table_name
			ORDER BY CASE WHEN contype='f' THEN 0 ELSE 1 END DESC,contype DESC,relname DESC,conname DESC
		LOOP
			executeText := 'ALTER TABLE '||constraint_record.nspname||'.'||fullTablename||' ADD CONSTRAINT '||constraint_record.full_constraint_name||' '|| constraint_record.constraintdef||';';
			EXECUTE executeText;
			RAISE INFO 'Constraint was generated: %', executeText;
		END LOOP;	 
		-----------------------------------------------------------------------------

		-----------------------------------------------------------------------------
		--------------------------- AUTO GENERATE INDEXES ---------------------------
		needToCreateIndexForDateColumn:=TRUE;
		FOR index_record IN 
			SELECT indexname, 
				indexdef
			FROM pg_indexes 
			WHERE tablename = primary_table_name 
				AND REPLACE(indexname,primary_table_name,fullTablename) NOT IN (SELECT indexname FROM pg_indexes WHERE tablename = fullTablename) -- will skip the index created by the PK above 
		LOOP
			IF (POSITION('(' || date_column || ')' in index_record.indexdef) > 0) THEN
				needToCreateIndexForDateColumn:=FALSE;
			END IF;
			executeText := 
				REPLACE(
							REPLACE(index_record.indexdef,
								'INDEX ' || primary_table_name,
								'INDEX ' || fullTablename
							),
							'ON ' || primary_table_name,
							'ON ' || fullTablename
						) FROM pg_indexes WHERE indexname = index_record.indexname
				;
			EXECUTE executeText;
			RAISE INFO 'Index was generated: %', executeText;
		END LOOP;
		-----------------------------------------------------------------------------
	 

		-- Create an index on the timestamp
		IF (needToCreateIndexForDateColumn) THEN
			createStmts:='CREATE INDEX idx_'||fullTablename||'_'||date_column||' ON '||schema_name||'.'||fullTablename||' ('||date_column||');';
			EXECUTE createStmts;
			RAISE NOTICE 'INDEX WAS GENERATED SPECIALLY FOR DATE COLUMN. create statement: %', createStmts;
		END IF;
		

		-- Track how many tables we are creating (should likely be 1, except for initial run and backfilling efforts).
		createdTables:=createdTables+1;
	END IF;

	IF (unkownTableCreated) THEN
		-- Add case for this table to trigger creation sql statement.
		createTrigger:='(' || -- isFirst=' || isFirst || ',is_first_contains_prior=' || is_first_contains_prior || '
			--' ||
			CASE WHEN isFirst AND is_first_contains_prior THEN 
				'' ELSE ' NEW.'||date_column||' >= TIMESTAMP '''||startTime||''' AND ' END
					||'NEW.'||date_column||' < TIMESTAMP '''||intervalTime||''' ) THEN 
				  INSERT INTO '||schema_name||'.'||fullTablename||' VALUES (NEW.*); '||createTrigger;

		startTime:=intervalTime;

		IF (startTime <= endTime) 
		THEN
		  createTrigger:=' 
			ELSEIF '||createTrigger;
		END IF;
		
		isFirst:=FALSE;
	ELSE 
		unkownTableCreated:=TRUE;
	END IF;
		
END LOOP;

-- Finish creating the trigger function (at the beginning).
createTrigger:='CREATE OR REPLACE FUNCTION '||schema_name||'.'||primary_table_name||'_insert_trigger_function()
                RETURNS TRIGGER AS $$
    BEGIN
        IF '||createTrigger;

-- Run the trigger replacement;
EXECUTE createTrigger;

-- Create the trigger that uses the trigger function, if it isn't already created
triggerName:=primary_table_name||'_insert_trigger';
IF NOT EXISTS(SELECT * FROM information_schema.triggers WHERE trigger_name = triggerName) THEN
  createTrigger:='CREATE TRIGGER '||primary_table_name||'_insert_trigger
                  BEFORE INSERT ON '||schema_name||'.'||primary_table_name||' 
                  FOR EACH ROW EXECUTE PROCEDURE '||schema_name||'.'||primary_table_name||'_insert_trigger_function();';
  EXECUTE createTrigger;
END IF;
return createdTables;
END;
$BODY$
  LANGUAGE 'plpgsql' VOLATILE
  COST 100;
ALTER FUNCTION public.update_partitions(TIMESTAMP WITHOUT TIME ZONE, TEXT, TEXT, TEXT, TEXT, TEXT,BOOLEAN) OWNER TO postgres;