partitioning/postgresql/procedures/02_maintenance.sql

-- ============================================================================
-- SCRIPT: 02_maintenance.sql
-- DESCRIPTION: Core functions for Zabbix partitioning (Create, Drop, Maintain).
-- ============================================================================

-- Function to check if a partition exists
CREATE OR REPLACE FUNCTION partitions.partition_exists(p_partition_name text) 
RETURNS boolean AS $$
BEGIN
    RETURN EXISTS (
        SELECT 1 FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE c.relname = p_partition_name
        AND n.nspname = 'public'
    );
END;
$$ LANGUAGE plpgsql;

-- Function to create a partition
CREATE OR REPLACE PROCEDURE partitions.create_partition(
    p_parent_table text,
    p_start_time timestamp with time zone,
    p_end_time timestamp with time zone,
    p_period text
) LANGUAGE plpgsql AS $$
DECLARE
    v_partition_name text;
    v_start_ts bigint;
    v_end_ts bigint;
    v_suffix text;
BEGIN
    -- (No changes needed for time here as passed params are already UTC-adjusted in caller)
    v_start_ts := extract(epoch from p_start_time)::bigint;
    v_end_ts := extract(epoch from p_end_time)::bigint;
    
    IF p_period = 'month' THEN
        v_suffix := to_char(p_start_time, 'YYYYMM');
    ELSE
        v_suffix := to_char(p_start_time, 'YYYYMMDD');
    END IF;
    
    v_partition_name := p_parent_table || '_p' || v_suffix;

    IF NOT partitions.partition_exists(v_partition_name) THEN
        EXECUTE format(
            'CREATE TABLE public.%I PARTITION OF public.%I FOR VALUES FROM (%s) TO (%s)',
            v_partition_name, p_parent_table, v_start_ts, v_end_ts
        );
    END IF;
END;
$$;

-- Function to drop old partitions
CREATE OR REPLACE PROCEDURE partitions.drop_old_partitions(
    p_parent_table text,
    p_retention interval,
    p_period text
) LANGUAGE plpgsql AS $$
DECLARE
    v_cutoff_ts bigint;
    v_partition record;
    v_partition_date timestamp with time zone;
    v_suffix text;
BEGIN
    -- Calculate cutoff timestamp
    v_cutoff_ts := extract(epoch from (now() - p_retention))::bigint;

    FOR v_partition IN 
        SELECT 
            child.relname AS partition_name
        FROM pg_inherits
        JOIN pg_class parent ON pg_inherits.inhparent = parent.oid
        JOIN pg_class child ON pg_inherits.inhrelid = child.oid
        WHERE parent.relname = p_parent_table
    LOOP
        -- Parse partition suffix to determine age
        -- Format: parent_pYYYYMM or parent_pYYYYMMDD
        v_suffix := substring(v_partition.partition_name from length(p_parent_table) + 3);
        
        BEGIN
            IF length(v_suffix) = 6 THEN -- YYYYMM
                v_partition_date := to_timestamp(v_suffix || '01', 'YYYYMMDD') AT TIME ZONE 'UTC';
                -- For monthly, we check if the END of the month is older than retention? 
                -- Or just strict retention. 
                -- To be safe, adding 1 month to check vs cutoff.
                IF extract(epoch from (v_partition_date + '1 month'::interval)) < v_cutoff_ts THEN
                    RAISE NOTICE 'Dropping old partition %', v_partition.partition_name;
                    EXECUTE format('DROP TABLE public.%I', v_partition.partition_name);
                    COMMIT; -- Release lock immediately
                END IF;
            ELSIF length(v_suffix) = 8 THEN -- YYYYMMDD
                v_partition_date := to_timestamp(v_suffix, 'YYYYMMDD') AT TIME ZONE 'UTC';
                IF extract(epoch from (v_partition_date + '1 day'::interval)) < v_cutoff_ts THEN
                    RAISE NOTICE 'Dropping old partition %', v_partition.partition_name;
                    EXECUTE format('DROP TABLE public.%I', v_partition.partition_name);
                    COMMIT; -- Release lock immediately
                END IF;
            END IF;
        EXCEPTION WHEN OTHERS THEN
            -- Ignore parsing errors for non-standard partitions
            NULL;
        END;
    END LOOP;
END;
$$;

-- MAIN Procedure to maintain a single table
CREATE OR REPLACE PROCEDURE partitions.maintain_table(
    p_table_name text,
    p_period text,
    p_keep_history interval,
    p_future_partitions integer DEFAULT 5
) LANGUAGE plpgsql AS $$
DECLARE
    v_start_time timestamp with time zone;
    v_period_interval interval;
    i integer;
    v_past_iterations integer;
BEGIN
    IF p_period = 'day' THEN
        v_period_interval := '1 day'::interval;
        v_start_time := date_trunc('day', now() AT TIME ZONE 'UTC');
        -- Calculate how many past days cover the retention period (86400 seconds = 1 day)
        v_past_iterations := ceil(extract(epoch from p_keep_history) / 86400)::integer;
        
    ELSIF p_period = 'week' THEN
        v_period_interval := '1 week'::interval;
        v_start_time := date_trunc('week', now() AT TIME ZONE 'UTC');
        -- 604800 seconds = 1 week
        v_past_iterations := ceil(extract(epoch from p_keep_history) / 604800)::integer;

    ELSIF p_period = 'month' THEN
        v_period_interval := '1 month'::interval;
        v_start_time := date_trunc('month', now() AT TIME ZONE 'UTC');
        -- Approximate 30 days per month (2592000 seconds)
        v_past_iterations := ceil(extract(epoch from p_keep_history) / 2592000)::integer;
    ELSE
        RETURN;
    END IF;

    -- 1. Create Future Partitions (Current + Buffer)
    FOR i IN 0..p_future_partitions LOOP
        CALL partitions.create_partition(
            p_table_name,
            v_start_time + (i * v_period_interval),
            v_start_time + ((i + 1) * v_period_interval),
            p_period
        );
        COMMIT; -- Release lock immediately
    END LOOP;
    
    -- 2. Create Past Partitions (Covering retention period)
    IF v_past_iterations > 0 THEN
        FOR i IN 1..v_past_iterations LOOP
            CALL partitions.create_partition(
                p_table_name,
                v_start_time - (i * v_period_interval),
                v_start_time - ((i - 1) * v_period_interval),
                p_period
            );
            COMMIT; -- Release lock immediately
        END LOOP;
    END IF;

    -- 3. Drop Old Partitions
    CALL partitions.drop_old_partitions(p_table_name, p_keep_history, p_period);
    
    -- 4. Update Metadata
    UPDATE partitions.config SET last_updated = now() WHERE table_name = p_table_name;
END;
$$;

-- Global Maintenance Procedure
CREATE OR REPLACE PROCEDURE partitions.run_maintenance() 
LANGUAGE plpgsql AS $$
DECLARE
    v_row record;
BEGIN
    FOR v_row IN SELECT * FROM partitions.config LOOP
        CALL partitions.maintain_table(v_row.table_name, v_row.period, v_row.keep_history, v_row.future_partitions);
    END LOOP;
END;
$$;