--
-- This file defines functions to (hopefully) reasonably safely enable and
-- disable enforcement of a foreign key constraint, written by Craig Ringer.
-- They're free for any use your care to make of them.
--
-- These functions work on the system that they're used on, but you
-- should still evaluate them for correctness and sanity before
-- adopting them yourself.
--
-- I make no guarantees that they won't destroy your data or steal your
-- lunch.
--


CREATE OR REPLACE FUNCTION disable_triggers_for_fkey_constraint(constraint_name VARCHAR) RETURNS void AS $$
DECLARE
    tgrec RECORD;
    relname VARCHAR;
    constraint_type CHAR;
BEGIN
    SELECT contype
    INTO constraint_type
    FROM pg_catalog.pg_constraint
    WHERE pg_catalog.pg_constraint.conname = constraint_name;

    IF constraint_type <> 'f' THEN
        RAISE EXCEPTION 'Can only disable triggers for foreign key constraints';
    END IF;

    FOR tgrec IN SELECT tgname FROM pg_catalog.pg_trigger WHERE tgconstrname = constraint_name
    LOOP
        -- Obtain the name of the table this trigger affects. Foreign key
        -- constraint triggers may affect the fkey or pkey tables and we have
        -- to find out which in order to disable the constraint.
        SELECT pg_catalog.pg_class.relname
        INTO STRICT relname
        FROM pg_catalog.pg_class INNER JOIN pg_catalog.pg_trigger
          ON pg_catalog.pg_trigger.tgrelid = pg_catalog.pg_class.oid
        WHERE pg_catalog.pg_trigger.tgname=tgrec.tgname;

        EXECUTE 'ALTER TABLE "'||relname||'" DISABLE TRIGGER "'||tgrec.tgname||'";';
    END LOOP;
END;
$$ LANGUAGE 'plpgsql' VOLATILE;

COMMENT ON FUNCTION disable_triggers_for_fkey_constraint(VARCHAR) IS 'Disable enforcement of foreign key constraint $1';


--
-- This stored procedure does a rapid check of the referential integrity protected by `constraint_name'
-- (MUCH faster than the incredibly slow one postgresql does during ALTER TABLE ... ADD CONSTRAINT ... FOREIGN KEY ...)
-- then re-enables the triggers that enforce the constraint.
--
-- It only works on foreign keys with only one column involved.
--
CREATE OR REPLACE FUNCTION enable_triggers_for_fkey_constraint(constraint_name VARCHAR) RETURNS void AS $$
DECLARE
    tgrec RECORD;
    relname VARCHAR;
    foreign_key_misses RECORD;
    constraint_info RECORD;
    fkey_table_name VARCHAR;
    pkey_table_name VARCHAR;
    fkey_col_list VARCHAR;
    fkey_col_not_null_clause VARCHAR;
    pkey_col_list VARCHAR;
    colname VARCHAR; -- temporary variable
    -- Used to control comma insertion in loops
    first BOOLEAN;
    -- Loop variables
    i INTEGER;
    -- Query text
    q VARCHAR;
BEGIN

    -- Look up the tables and columns that the foreign key involves
    SELECT
        contype,
        conrelid, -- oid of referencing relation
        confrelid, -- oid of referenced relation
        (SELECT pg_catalog.pg_type.typname FROM pg_catalog.pg_type WHERE pg_catalog.pg_type.typrelid = conrelid) AS conrelid_name, -- name of referencing relation
        (SELECT pg_catalog.pg_type.typname FROM pg_catalog.pg_type WHERE pg_catalog.pg_type.typrelid = confrelid) AS confrelid_name, -- name of referenced relation
        pg_catalog.pg_constraint.conkey, -- Position of referencing column, eg {14}
        pg_catalog.pg_constraint.confkey -- Position of referenced column, eg {1}
    INTO STRICT constraint_info
    FROM pg_catalog.pg_constraint
    WHERE pg_catalog.pg_constraint.conname = constraint_name;

    IF constraint_info.contype <> 'f' THEN
        RAISE EXCEPTION 'Can only enable triggers for foreign key constraints';
    END IF;

    fkey_table_name := constraint_info.conrelid_name;
    pkey_table_name := constraint_info.confrelid_name;

    -- Now we need to build SQL snippets for:
    --   the fkey table column list
    --   a WHERE clause snippet to exclude foreign key values where all fields are NULL from the check
    --   the pkey table column list
    first := 't';
    fkey_col_list := '';
    fkey_col_not_null_clause := '';
    FOR i IN array_lower(constraint_info.conkey,1)..array_upper(constraint_info.conkey,1) LOOP
        IF first THEN
            first := 'f';
            fkey_col_not_null_clause := '(';
        ELSE
            fkey_col_list := fkey_col_list||', ';
            fkey_col_not_null_clause := fkey_col_not_null_clause||' AND ';
        END IF;

        SELECT pg_catalog.pg_attribute.attname
        INTO colname
        FROM pg_catalog.pg_attribute
        WHERE pg_catalog.pg_attribute.attnum = constraint_info.conkey[i] AND pg_catalog.pg_attribute.attrelid = constraint_info.conrelid;

        fkey_col_list := fkey_col_list||'"'||colname||'"';
        fkey_col_not_null_clause := fkey_col_not_null_clause||colname||' IS NOT NULL';
    END LOOP;
    fkey_col_not_null_clause := fkey_col_not_null_clause||')';

    first := 't';
    pkey_col_list := '';
    FOR i IN array_lower(constraint_info.conkey,1)..array_upper(constraint_info.confkey,1) LOOP
        IF first THEN
            first := 'f';
        ELSE
            pkey_col_list := pkey_col_list||', ';
        END IF;

        SELECT pg_catalog.pg_attribute.attname
        INTO colname
        FROM pg_catalog.pg_attribute
        WHERE pg_catalog.pg_attribute.attnum = constraint_info.confkey[i] AND pg_catalog.pg_attribute.attrelid = constraint_info.confrelid;

        pkey_col_list := pkey_col_list||'"'||colname||'"';
    END LOOP;

    -- An optimised foreign key check, found at
    -- http://archives.postgresql.org/pgsql-performance/2003-03/msg00157.php
    -- and adapted for correct handling of NULLs, improved efficiency with many
    -- similar foreign key entries, handling of multiple key columns.
    --
    -- It's highly efficient if the foreign key table is a similar or larger
    -- size than the primary key table, but doesn't run as fast as an indexed
    -- SELECT when the primary key table is much larger then the list of
    -- foreign keys being checked. In that case, though, you won't be using
    -- these functions...
    --
    -- If any results are returned by this query, they're foreign key entries without
    -- matching primary key entries, so if the loop body is executed there's a referential
    -- integrity error.
    --
    q := '
        SELECT '||fkey_col_list||'
        FROM (
            SELECT DISTINCT '||fkey_col_list||', 0 AS pri FROM "'||fkey_table_name||'" WHERE '||fkey_col_not_null_clause||'
            UNION ALL
            SELECT '||pkey_col_list||', 1 AS pri FROM "'||pkey_table_name||'"
        ) AS key_info
        GROUP BY '||fkey_col_list||'
        HAVING sum(pri) = 0';

    --RAISE NOTICE 'About to execute: %',q;

    FOR foreign_key_misses IN EXECUTE q
    LOOP
        RAISE EXCEPTION 'Foreign key constraint check failed on %(%)',fkey_table_name,fkey_col_list;
    END LOOP;
    RAISE NOTICE 'Referential integrity check on %(%) REFERENCES %(%) passed - activating triggers',fkey_table_name,fkey_col_list,pkey_table_name,pkey_col_list;

    FOR tgrec IN SELECT tgname FROM pg_catalog.pg_trigger WHERE tgconstrname = constraint_name
    LOOP
        -- Obtain the name of the table this trigger affects. Foreign key
        -- constraint triggers may affect the fkey or pkey tables and we have
        -- to find out which in order to enable the constraint.
        SELECT pg_catalog.pg_class.relname
        INTO STRICT relname
        FROM pg_catalog.pg_class INNER JOIN pg_catalog.pg_trigger
          ON pg_catalog.pg_trigger.tgrelid = pg_catalog.pg_class.oid
        WHERE pg_catalog.pg_trigger.tgname=tgrec.tgname;

        EXECUTE 'ALTER TABLE "'||relname||'" ENABLE TRIGGER "'||tgrec.tgname||'";';
    END LOOP;
END;
$$ LANGUAGE 'plpgsql' VOLATILE;

COMMENT ON FUNCTION enable_triggers_for_fkey_constraint(VARCHAR) IS 'Quickly re-check and enable enforcement of foreign key constraint $1';