Drop certain tables in a function managing partitions

Question

I needed a function to automatically manage partitions in a DB. I found a function that created 2 months at a time and adapted it to work with quarters and multiple tables with different primary keys and creation date keys.
Now I need to change it from dropping the last quarter´s partition. Instead I only want it to drop a quarter´s partitions if they are from the previous year. How do I change the code for previous quarters to do that?

Relevant code:

-- check if the partition for the previous quarter exists
v_date_from := date_trunc('quarter', v_current_date - '3 month'::interval);
v_partition_name := master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);
v_rule_name := 'rule_' || master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);

SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables WHERE schemaname = 'public' AND tablename = v_partition_name;

IF (v_exists) THEN

    EXECUTE 'DROP RULE ' || v_rule_name || ' ON ' || master_table;
    EXECUTE 'DROP TABLE ' || v_partition_name;

END IF;

Complete function:

CREATE OR REPLACE FUNCTION manage_partitions(timestamp without time zone, master_table character varying, prime_key character varying, prime_date character varying) RETURNS void AS
$BODY$
DECLARE

    -- name of the next partition and rule (and interval boundaries)
    v_partition_name    VARCHAR(32);
    v_rule_name         VARCHAR(32);

    v_date_from         TIMESTAMP;
    v_date_to           TIMESTAMP;

    -- current date (if NULL, a current timestamp is used)
    v_date              ALIAS FOR $1;
    v_current_date      TIMESTAMP;

    -- used just for checking existence of the partitions
    v_exists            BOOLEAN;

BEGIN

    IF (v_date IS NULL) THEN
        v_current_date := current_timestamp;
    ELSE
        v_current_date := v_date;
    END IF;

    -- check if the partition for the previous quarter exists
    v_date_from := date_trunc('quarter', v_current_date - '3 month'::interval);
    v_partition_name := master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);
    v_rule_name := 'rule_' || master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);

    SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables WHERE schemaname = 'public' AND tablename = v_partition_name;

    IF (v_exists) THEN

        EXECUTE 'DROP RULE ' || v_rule_name || ' ON ' || master_table;
        EXECUTE 'DROP TABLE ' || v_partition_name;

    END IF;

    -- create a partition for this quarter
    v_date_from := date_trunc('quarter', v_current_date);
    v_date_to := v_date_from + '3 month';
    v_partition_name := master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);
    v_rule_name := 'rule_' || master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);

    SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables WHERE schemaname = 'public' AND tablename = v_partition_name;

    IF (NOT v_exists) THEN

        EXECUTE 'CREATE TABLE ' || v_partition_name || ' (PRIMARY KEY (' || prime_key || '), CHECK (' || prime_date || ' >= ''' || v_date_from || ''' AND ' || prime_date || ' < ''' || v_date_to || ''')) INHERITS (' || master_table || ')';
        EXECUTE 'CREATE RULE ' || v_rule_name || ' AS ON INSERT TO ' || master_table || ' DO INSTEAD INSERT INTO ' || v_partition_name || ' VALUES (NEW.*)';
        -- if you need to create indexes / foreign keys / whatever on the partition, you may do it here

    END IF;

    -- create a partition for next quarter
    v_date_from := date_trunc('quarter', v_current_date + '3 month'::interval);
    v_date_to := v_date_from + '3 month';
    v_partition_name := master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);
    v_rule_name := 'rule_' || master_table || '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from);

    SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables WHERE schemaname = 'public' AND tablename = v_partition_name;

    IF (NOT v_exists) THEN

        EXECUTE 'CREATE TABLE ' || v_partition_name || ' (PRIMARY KEY (' || prime_key || '), CHECK (' || prime_date || ' >= ''' || v_date_from || ''' AND ' || prime_date || ' < ''' || v_date_to || ''')) INHERITS (' || master_table || ')';
        EXECUTE 'CREATE RULE ' || v_rule_name || ' AS ON INSERT TO ' || master_table || ' DO INSTEAD INSERT INTO ' || v_partition_name || ' VALUES (NEW.*)';

        -- if you need to create indexes / foreign keys / whatever on the partition, you may do it here

    END IF;

END;
$BODY$
LANGUAGE plpgsql VOLATILE
COST 100;

You should provide the source where you `found a function that created 2 months at a time` and (as always!) your version of Postgres. — Erwin Brandstetter, Feb 03 '14 at 23:20

Erwin Brandstetter · Answer 1 · 2020-08-08T02:49:51.030

The code you are working off is outdated. It is also rather inefficient for any version of Postgres. (I'd be interested where you got that from?)

Completely rewrite the whole function with features of modern-day SQL and PL/pgSQL:

CREATE OR REPLACE FUNCTION manage_partitions(
     v_date       timestamp
   , master_table regclass
   , prime_key    text
   , prime_date   text)
  RETURNS void
  LANGUAGE plpgsql AS
$func$
DECLARE
   v_current_date   timestamp := COALESCE(v_date, now()); -- fallback for NULL
   v_date_from      timestamp;
   v_partition_name text;

BEGIN
   -- drop partition for previous quarter of previous year if exists --
   v_date_from := date_trunc('quarter', v_current_date - interval '3 month');

   IF v_date_from < date_trunc('year', now()) THEN -- your (odd?) condition
      v_partition_name := master_table || to_char(v_date_from, '"_q"Q_YYYY');

      EXECUTE format(
         'DROP RULE  IF EXISTS %I ON %s;
          DROP TABLE IF EXISTS %I'
        , 'rule_' || master_table || to_char(v_date_from, '"_q"Q_YYYY')
        , master_table
        , v_partition_name);
   END IF;

   -- create partition for this quarter --
   v_date_from := date_trunc('quarter', v_current_date);
   v_partition_name := master_table || to_char(v_date_from, '"_q"Q_YYYY');

   IF NOT EXISTS (
      SELECT FROM pg_tables t
      WHERE  t.schemaname = 'public'
      AND    t.tablename  = v_partition_name) THEN

      EXECUTE format(
         'CREATE TABLE %$1I (
             PRIMARY KEY (%$2I), CHECK (%$3L >= %$4L AND %$3L < %$5L))
             INHERITS (%$6s);
          CREATE RULE %$7I AS ON INSERT TO %$6s DO INSTEAD
          INSERT INTO %$1I VALUES (NEW.*)'
        , v_partition_name
        , prime_key
        , prime_date
        , v_date_from
        , v_date_from + interval '3 month'
        , master_table
        , 'rule_' || master_table || to_char(v_date_from, '"_q"Q_YYYY')
      );
   END IF;

   -- create partition for next quarter --
   v_date_from := date_trunc('quarter', v_current_date + interval '3 month');
   v_partition_name := master_table || to_char(v_date_from, '"_q"Q_YYYY');

   IF NOT EXISTS (
      SELECT FROM pg_tables t
      WHERE  t.schemaname = 'public'
      AND    t.tablename = v_partition_name) THEN

      EXECUTE format(
         'CREATE TABLE %$1I (
             PRIMARY KEY (%$2I), CHECK (%$3L >= %$4L AND %$3L < %$5L))
             INHERITS (%$6s);
          CREATE RULE %$7I AS ON INSERT TO %$6s DO INSTEAD
          INSERT INTO %$1I VALUES (NEW.*)'
        , v_partition_name
        , prime_key
        , prime_date
        , v_date_from
        , v_date_from + interval '3 month'
        , master_table
        , 'rule_' || master_table || to_char(v_date_from, '"_q"Q_YYYY')
      );
   END IF;

END
$func$;

You can assign variables at declaration time to simplify code.

Replace:

  '_Q' || EXTRACT(QUARTER FROM v_date_from) || '_' || EXTRACT(YEAR FROM v_date_from)

with the simpler and faster:

  to_char(v_date_from, '"Q_"Q_YYYY')

to_char() in the manual.

use the superior for IF EXISTS (...) THEN .... We can then drop the useless variable v_exists. See:
- PL/pgSQL checking if a row exists
Have a look at the manual page on format() (Postgres 9.1+). If you are going to work with dynamic SQL you need to know about it.
Use the object identifier type regclass for master_table to verify that it exists and is visible with the current setting for search_path and prevent SQL injection at the same time. See:
- Table name as a PostgreSQL function parameter
Don't use the outdated and discouraged ALIAS FOR clause. Use parameter names instead, like you already do for all other function params.
Note how I replaced capital letters in identifiers (which are a bad idea) with lower case letters ('Q' -> 'q'). More in the last paragraph of this related answer:
- Define table and column names as arguments in a plpgsql function?
I also dropped the variables v_date_to and v_rule_name and replaced them with expressions since those are used once only in my code.

score 0 · Accepted Answer · answered Jan 31 '14 at 07:16

I think you need to change the following line:

SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables 
WHERE schemaname = 'public' AND tablename = v_partition_name;

Into something like:

SELECT COUNT(*) = 1 INTO v_exists FROM pg_tables 
WHERE schemaname = 'public' AND tablename = v_partition_name 
and v_date_from < date_trunc('year', current_date);

This will return 1 only in the case when partition is from previous year.

Drop certain tables in a function managing partitions

2 Answers2