Refactor code structure for improved readability and maintainability; optimize performance in key functions.

deploy/postgres-partitioning/001_partition_infrastructure.sql

@@ -21,7 +21,25 @@ COMMENT ON SCHEMA partition_mgmt IS
     'Partition management utilities for time-series tables';
 
 -- ============================================================================
--- Step 2: Partition creation function
+-- Step 2: Managed table registration
+-- ============================================================================
+
+CREATE TABLE IF NOT EXISTS partition_mgmt.managed_tables (
+    schema_name TEXT NOT NULL,
+    table_name TEXT NOT NULL,
+    partition_key TEXT NOT NULL,
+    partition_type TEXT NOT NULL,
+    retention_months INT NOT NULL DEFAULT 0,
+    months_ahead INT NOT NULL DEFAULT 3,
+    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    PRIMARY KEY (schema_name, table_name)
+);
+
+COMMENT ON TABLE partition_mgmt.managed_tables IS
+    'Tracks partitioned tables with retention and creation settings';
+
+-- ============================================================================
+-- Step 3: Partition creation function
 -- ============================================================================
 
 -- Creates a new partition for a given table and date range
@@ -78,7 +96,7 @@ END;
 $$;
 
 -- ============================================================================
--- Step 3: Monthly partition creation helper
+-- Step 4: Monthly partition creation helper
 -- ============================================================================
 
 CREATE OR REPLACE FUNCTION partition_mgmt.create_monthly_partitions(
@@ -114,7 +132,7 @@ END;
 $$;
 
 -- ============================================================================
--- Step 4: Quarterly partition creation helper
+-- Step 5: Quarterly partition creation helper
 -- ============================================================================
 
 CREATE OR REPLACE FUNCTION partition_mgmt.create_quarterly_partitions(
@@ -156,7 +174,155 @@ END;
 $$;
 
 -- ============================================================================
--- Step 5: Partition detach and archive function
+-- Step 6: Ensure future partitions exist
+-- ============================================================================
+
+CREATE OR REPLACE FUNCTION partition_mgmt.ensure_future_partitions(
+    p_schema_name TEXT,
+    p_table_name TEXT,
+    p_months_ahead INT
+)
+RETURNS INT
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    v_partition_key TEXT;
+    v_partition_type TEXT;
+    v_months_ahead INT;
+    v_created INT := 0;
+    v_current DATE;
+    v_end DATE;
+    v_suffix TEXT;
+    v_partition_name TEXT;
+BEGIN
+    SELECT partition_key, partition_type, months_ahead
+      INTO v_partition_key, v_partition_type, v_months_ahead
+      FROM partition_mgmt.managed_tables
+     WHERE schema_name = p_schema_name
+       AND table_name = p_table_name;
+
+    IF v_partition_key IS NULL THEN
+        RETURN 0;
+    END IF;
+
+    IF p_months_ahead IS NOT NULL AND p_months_ahead > 0 THEN
+        v_months_ahead := p_months_ahead;
+    END IF;
+
+    IF v_months_ahead IS NULL OR v_months_ahead <= 0 THEN
+        RETURN 0;
+    END IF;
+
+    v_partition_type := lower(coalesce(v_partition_type, 'monthly'));
+
+    IF v_partition_type = 'monthly' THEN
+        v_current := date_trunc('month', NOW())::DATE;
+        v_end := date_trunc('month', NOW() + (v_months_ahead || ' months')::INTERVAL)::DATE;
+
+        WHILE v_current <= v_end LOOP
+            v_partition_name := format('%s_%s', p_table_name, to_char(v_current, 'YYYY_MM'));
+            IF NOT EXISTS (
+                SELECT 1 FROM pg_class c
+                JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE n.nspname = p_schema_name AND c.relname = v_partition_name
+            ) THEN
+                PERFORM partition_mgmt.create_partition(
+                    p_schema_name,
+                    p_table_name,
+                    v_partition_key,
+                    v_current,
+                    (v_current + INTERVAL '1 month')::DATE
+                );
+                v_created := v_created + 1;
+            END IF;
+
+            v_current := (v_current + INTERVAL '1 month')::DATE;
+        END LOOP;
+    ELSIF v_partition_type = 'quarterly' THEN
+        v_current := date_trunc('quarter', NOW())::DATE;
+        v_end := date_trunc('quarter', NOW() + (v_months_ahead || ' months')::INTERVAL)::DATE;
+
+        WHILE v_current <= v_end LOOP
+            v_suffix := to_char(v_current, 'YYYY') || '_Q' ||
+                        EXTRACT(QUARTER FROM v_current)::TEXT;
+            v_partition_name := format('%s_%s', p_table_name, v_suffix);
+
+            IF NOT EXISTS (
+                SELECT 1 FROM pg_class c
+                JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE n.nspname = p_schema_name AND c.relname = v_partition_name
+            ) THEN
+                PERFORM partition_mgmt.create_partition(
+                    p_schema_name,
+                    p_table_name,
+                    v_partition_key,
+                    v_current,
+                    (v_current + INTERVAL '3 months')::DATE,
+                    v_suffix
+                );
+                v_created := v_created + 1;
+            END IF;
+
+            v_current := (v_current + INTERVAL '3 months')::DATE;
+        END LOOP;
+    END IF;
+
+    RETURN v_created;
+END;
+$$;
+
+-- ============================================================================
+-- Step 7: Retention enforcement function
+-- ============================================================================
+
+CREATE OR REPLACE FUNCTION partition_mgmt.enforce_retention(
+    p_schema_name TEXT,
+    p_table_name TEXT,
+    p_retention_months INT
+)
+RETURNS INT
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    v_retention_months INT;
+    v_cutoff_date DATE;
+    v_partition RECORD;
+    v_dropped INT := 0;
+BEGIN
+    SELECT retention_months
+      INTO v_retention_months
+      FROM partition_mgmt.managed_tables
+     WHERE schema_name = p_schema_name
+       AND table_name = p_table_name;
+
+    IF p_retention_months IS NOT NULL AND p_retention_months > 0 THEN
+        v_retention_months := p_retention_months;
+    END IF;
+
+    IF v_retention_months IS NULL OR v_retention_months <= 0 THEN
+        RETURN 0;
+    END IF;
+
+    v_cutoff_date := (NOW() - (v_retention_months || ' months')::INTERVAL)::DATE;
+
+    FOR v_partition IN
+        SELECT partition_name, partition_end
+          FROM partition_mgmt.partition_stats
+         WHERE schema_name = p_schema_name
+           AND table_name = p_table_name
+    LOOP
+        IF v_partition.partition_end IS NOT NULL AND v_partition.partition_end < v_cutoff_date THEN
+            EXECUTE format('DROP TABLE IF EXISTS %I.%I', p_schema_name, v_partition.partition_name);
+            v_dropped := v_dropped + 1;
+        END IF;
+    END LOOP;
+
+    RETURN v_dropped;
+END;
+$$;
+
+-- ============================================================================
+-- Step 8: Partition detach and archive function
 -- ============================================================================
 
 CREATE OR REPLACE FUNCTION partition_mgmt.detach_partition(
@@ -204,7 +370,7 @@ END;
 $$;
 
 -- ============================================================================
--- Step 6: Partition retention cleanup function
+-- Step 9: Partition retention cleanup function
 -- ============================================================================
 
 CREATE OR REPLACE FUNCTION partition_mgmt.cleanup_old_partitions(
@@ -262,7 +428,7 @@ END;
 $$;
 
 -- ============================================================================
--- Step 7: Partition statistics view
+-- Step 10: Partition statistics view
 -- ============================================================================
 
 CREATE OR REPLACE VIEW partition_mgmt.partition_stats AS
@@ -271,6 +437,8 @@ SELECT
     parent.relname AS table_name,
     c.relname AS partition_name,
     pg_get_expr(c.relpartbound, c.oid) AS partition_range,
+    (regexp_match(pg_get_expr(c.relpartbound, c.oid), 'FROM \(''([^'']+)''\)'))[1]::DATE AS partition_start,
+    (regexp_match(pg_get_expr(c.relpartbound, c.oid), 'TO \(''([^'']+)''\)'))[1]::DATE AS partition_end,
     pg_size_pretty(pg_relation_size(c.oid)) AS size,
     pg_relation_size(c.oid) AS size_bytes,
     COALESCE(s.n_live_tup, 0) AS estimated_rows,
@@ -291,7 +459,7 @@ COMMENT ON VIEW partition_mgmt.partition_stats IS
     'Statistics for all partitioned tables in the database';
 
 -- ============================================================================
--- Step 8: BRIN index optimization helper
+-- Step 11: BRIN index optimization helper
 -- ============================================================================
 
 CREATE OR REPLACE FUNCTION partition_mgmt.create_brin_index_if_not_exists(
@@ -336,7 +504,7 @@ END;
 $$;
 
 -- ============================================================================
--- Step 9: Maintenance job tracking table
+-- Step 12: Maintenance job tracking table
 -- ============================================================================
 
 CREATE TABLE IF NOT EXISTS partition_mgmt.maintenance_log (
@@ -356,7 +524,7 @@ CREATE INDEX idx_maintenance_log_table ON partition_mgmt.maintenance_log(schema_
 CREATE INDEX idx_maintenance_log_status ON partition_mgmt.maintenance_log(status, started_at);
 
 -- ============================================================================
--- Step 10: Archive schema for detached partitions
+-- Step 13: Archive schema for detached partitions
 -- ============================================================================
 
 CREATE SCHEMA IF NOT EXISTS archive;