Add pg_upgrade support functions

age--1.7.0--y.y.y.sql

@@ -30,3 +30,341 @@
 --* Please add all additions, deletions, and modifications to the end of this
 --* file. We need to keep the order of these changes.
 --* REMOVE ALL LINES ABOVE, and this one, that start with --*
+
+--
+-- pg_upgrade support functions
+--
+-- These functions help users upgrade PostgreSQL major versions using pg_upgrade
+-- while preserving Apache AGE graph data.
+--
+
+CREATE FUNCTION ag_catalog.age_prepare_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = ag_catalog, pg_catalog
+    AS $function$
+DECLARE
+    has_graphs boolean;
+BEGIN
+    -- Check if there are any graphs to process
+    SELECT EXISTS(SELECT 1 FROM ag_catalog.ag_graph) INTO has_graphs;
+
+    IF NOT has_graphs THEN
+        RAISE NOTICE 'No graphs found. Nothing to prepare for pg_upgrade.';
+        RETURN;
+    END IF;
+
+    -- Check if namespace column is already oid type (already prepared)
+    IF EXISTS (
+        SELECT 1 FROM information_schema.columns
+        WHERE table_schema = 'ag_catalog'
+          AND table_name = 'ag_graph'
+          AND column_name = 'namespace'
+          AND data_type = 'oid'
+    ) THEN
+        RAISE NOTICE 'Database already prepared for pg_upgrade (namespace is oid type).';
+        RETURN;
+    END IF;
+
+    -- Drop existing backup table if it exists (from a previous failed attempt)
+    DROP TABLE IF EXISTS public._age_pg_upgrade_backup;
+
+    -- Create backup table with graph names mapped to namespace names
+    -- We store names (not OIDs) because names survive pg_upgrade while OIDs don't
+    CREATE TABLE public._age_pg_upgrade_backup AS
+    SELECT
+        g.graphid AS old_graphid,
+        g.name AS graph_name,
+        g.namespace::regnamespace::text AS namespace_name
+    FROM ag_catalog.ag_graph g;
+
+    RAISE NOTICE 'Created backup table public._age_pg_upgrade_backup with % graph(s)',
+        (SELECT count(*) FROM public._age_pg_upgrade_backup);
+
+    -- Drop the existing regnamespace-based index
+    DROP INDEX IF EXISTS ag_catalog.ag_graph_namespace_index;
+
+    -- Convert namespace column from regnamespace to oid
+    ALTER TABLE ag_catalog.ag_graph
+        ALTER COLUMN namespace TYPE oid USING namespace::oid;
+
+    -- Recreate the index with oid type
+    CREATE UNIQUE INDEX ag_graph_namespace_index
+        ON ag_catalog.ag_graph USING btree (namespace);
+
+    -- Create a view for backward-compatible display of namespace as schema name
+    CREATE OR REPLACE VIEW ag_catalog.ag_graph_view AS
+    SELECT graphid, name, namespace::regnamespace AS namespace
+    FROM ag_catalog.ag_graph;
+
+    RAISE NOTICE 'Successfully prepared database for pg_upgrade.';
+    RAISE NOTICE 'The ag_graph.namespace column has been converted from regnamespace to oid.';
+    RAISE NOTICE 'You can now run pg_upgrade.';
+    RAISE NOTICE 'After pg_upgrade completes, run: SELECT age_finish_pg_upgrade();';
+END;
+$function$;
+
+COMMENT ON FUNCTION ag_catalog.age_prepare_pg_upgrade() IS
+'Prepares an AGE database for pg_upgrade by converting ag_graph.namespace from regnamespace to oid type. Run this before pg_upgrade.';
+
+CREATE FUNCTION ag_catalog.age_finish_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = ag_catalog, pg_catalog
+    AS $function$
+DECLARE
+    mapping_count integer;
+    updated_labels integer;
+    updated_graphs integer;
+BEGIN
+    -- Check if backup table exists
+    IF NOT EXISTS (
+        SELECT 1 FROM information_schema.tables
+        WHERE table_schema = 'public'
+          AND table_name = '_age_pg_upgrade_backup'
+    ) THEN
+        RAISE EXCEPTION 'Backup table public._age_pg_upgrade_backup not found. '
+            'Did you run age_prepare_pg_upgrade() before pg_upgrade?';
+    END IF;
+
+    -- Check if namespace column is oid type (was properly prepared)
+    IF NOT EXISTS (
+        SELECT 1 FROM information_schema.columns
+        WHERE table_schema = 'ag_catalog'
+          AND table_name = 'ag_graph'
+          AND column_name = 'namespace'
+          AND data_type = 'oid'
+    ) THEN
+        RAISE EXCEPTION 'ag_graph.namespace is not oid type. '
+            'Did you run age_prepare_pg_upgrade() before pg_upgrade?';
+    END IF;
+
+    -- Create temporary mapping table with old and new OIDs
+    CREATE TEMP TABLE _graphid_mapping AS
+    SELECT
+        b.old_graphid,
+        b.graph_name,
+        n.oid AS new_graphid
+    FROM public._age_pg_upgrade_backup b
+    JOIN pg_namespace n ON n.nspname = b.namespace_name;
+
+    GET DIAGNOSTICS mapping_count = ROW_COUNT;
+
+    IF mapping_count = 0 THEN
+        RAISE EXCEPTION 'No OID mappings found. Schema names may have changed.';
+    END IF;
+
+    RAISE NOTICE 'Found % graph(s) to remap', mapping_count;
+
+    -- Temporarily drop foreign key constraint
+    ALTER TABLE ag_catalog.ag_label DROP CONSTRAINT IF EXISTS fk_graph_oid;
+
+    -- Update ag_label.graph references to use new OIDs
+    UPDATE ag_catalog.ag_label l
+    SET graph = m.new_graphid
+    FROM _graphid_mapping m
+    WHERE l.graph = m.old_graphid;
+
+    GET DIAGNOSTICS updated_labels = ROW_COUNT;
+    RAISE NOTICE 'Updated % label record(s)', updated_labels;
+
+    -- Update ag_graph.graphid and ag_graph.namespace to new OIDs
+    UPDATE ag_catalog.ag_graph g
+    SET graphid = m.new_graphid,
+        namespace = m.new_graphid
+    FROM _graphid_mapping m
+    WHERE g.graphid = m.old_graphid;
+
+    RAISE NOTICE 'Updated % graph record(s)', updated_graphs;
+
+    -- Restore foreign key constraint
+    ALTER TABLE ag_catalog.ag_label
+    ADD CONSTRAINT fk_graph_oid
+        FOREIGN KEY(graph) REFERENCES ag_catalog.ag_graph(graphid);
+
+    -- Clean up temporary mapping table
+    DROP TABLE _graphid_mapping;
+    DROP TABLE public._age_pg_upgrade_backup;
+
+    RAISE NOTICE 'Successfully completed pg_upgrade OID remapping.';
+
+    --
+    -- Restore original schema (revert namespace to regnamespace)
+    --
+    RAISE NOTICE 'Restoring original schema...';
+
+    -- Drop the view (no longer needed with regnamespace)
+    DROP VIEW IF EXISTS ag_catalog.ag_graph_view;
+
+    -- Drop the existing oid-based index
+    DROP INDEX IF EXISTS ag_catalog.ag_graph_namespace_index;
+
+    -- Convert namespace column back to regnamespace
+    ALTER TABLE ag_catalog.ag_graph
+        ALTER COLUMN namespace TYPE regnamespace USING namespace::regnamespace;
+
+    -- Recreate the index with regnamespace type
+    CREATE UNIQUE INDEX ag_graph_namespace_index
+        ON ag_catalog.ag_graph USING btree (namespace);
+
+    RAISE NOTICE 'Successfully restored ag_graph.namespace to regnamespace type.';
+
+    --
+    -- Invalidate AGE's internal caches by touching each graph's namespace
+    -- AGE registers a syscache callback on NAMESPACEOID, so altering a schema
+    -- triggers cache invalidation. This ensures cypher queries work immediately
+    -- without requiring a session reconnect.
+    --
+    RAISE NOTICE 'Invalidating AGE caches...';
+    PERFORM pg_catalog.pg_advisory_lock(hashtext('age_finish_pg_upgrade'));
+    BEGIN
+        -- Touch each graph's namespace to invalidate caches
+        DECLARE
+            graph_rec RECORD;
+        BEGIN
+            FOR graph_rec IN SELECT namespace::text AS ns_name FROM ag_catalog.ag_graph
+            LOOP
+                EXECUTE format('ALTER SCHEMA %I OWNER TO CURRENT_USER', graph_rec.ns_name);
+            END LOOP;
+        END;
+    END;
+    PERFORM pg_catalog.pg_advisory_unlock(hashtext('age_finish_pg_upgrade'));
+
+    RAISE NOTICE '';
+    RAISE NOTICE 'pg_upgrade complete. All graph data has been preserved.';
+END;
+$function$;
+
+COMMENT ON FUNCTION ag_catalog.age_finish_pg_upgrade() IS
+'Completes pg_upgrade by remapping stale OIDs and restoring the original schema. Run this after pg_upgrade.';
+
+CREATE FUNCTION ag_catalog.age_revert_pg_upgrade_changes()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = ag_catalog, pg_catalog
+    AS $function$
+BEGIN
+    -- Check if namespace column is oid type (needs reverting)
+    IF NOT EXISTS (
+        SELECT 1 FROM information_schema.columns
+        WHERE table_schema = 'ag_catalog'
+          AND table_name = 'ag_graph'
+          AND column_name = 'namespace'
+          AND data_type = 'oid'
+    ) THEN
+        RAISE NOTICE 'ag_graph.namespace is already regnamespace type. Nothing to revert.';
+        RETURN;
+    END IF;
+
+    -- Drop the view (no longer needed with regnamespace)
+    DROP VIEW IF EXISTS ag_catalog.ag_graph_view;
+
+    -- Drop the existing oid-based index
+    DROP INDEX IF EXISTS ag_catalog.ag_graph_namespace_index;
+
+    -- Convert namespace column back to regnamespace
+    ALTER TABLE ag_catalog.ag_graph
+        ALTER COLUMN namespace TYPE regnamespace USING namespace::regnamespace;
+
+    -- Recreate the index with regnamespace type
+    CREATE UNIQUE INDEX ag_graph_namespace_index
+        ON ag_catalog.ag_graph USING btree (namespace);
+
+    --
+    -- Invalidate AGE's internal caches by touching each graph's namespace
+    --
+    PERFORM pg_catalog.pg_advisory_lock(hashtext('age_revert_pg_upgrade'));
+    BEGIN
+        DECLARE
+            graph_rec RECORD;
+        BEGIN
+            FOR graph_rec IN SELECT namespace::text AS ns_name FROM ag_catalog.ag_graph
+            LOOP
+                EXECUTE format('ALTER SCHEMA %I OWNER TO CURRENT_USER', graph_rec.ns_name);
+            END LOOP;
+        END;
+    END;
+    PERFORM pg_catalog.pg_advisory_unlock(hashtext('age_revert_pg_upgrade'));
+
+    RAISE NOTICE 'Successfully reverted ag_graph.namespace to regnamespace type.';
+    RAISE NOTICE '';
+    RAISE NOTICE 'Upgrade preparation has been cancelled.';
+    RAISE NOTICE 'You may want to drop the backup table: DROP TABLE IF EXISTS public._age_pg_upgrade_backup;';
+END;
+$function$;
+
+COMMENT ON FUNCTION ag_catalog.age_revert_pg_upgrade_changes() IS
+'Reverts schema changes if you need to cancel after age_prepare_pg_upgrade() but before pg_upgrade. Not needed after age_finish_pg_upgrade().';
+
+CREATE FUNCTION ag_catalog.age_pg_upgrade_status()
+    RETURNS TABLE (
+        status text,
+        namespace_type text,
+        graph_count bigint,
+        backup_exists boolean,
+        message text
+    )
+    LANGUAGE plpgsql
+    SET search_path = ag_catalog, pg_catalog
+    AS $function$
+DECLARE
+    ns_type text;
+    g_count bigint;
+    backup_exists boolean;
+BEGIN
+    -- Get namespace column type
+    SELECT data_type INTO ns_type
+    FROM information_schema.columns
+    WHERE table_schema = 'ag_catalog'
+      AND table_name = 'ag_graph'
+      AND column_name = 'namespace';
+
+    -- Get graph count
+    SELECT count(*) INTO g_count FROM ag_catalog.ag_graph;
+
+    -- Check for backup table
+    SELECT EXISTS(
+        SELECT 1 FROM information_schema.tables
+        WHERE table_schema = 'public'
+          AND table_name = '_age_pg_upgrade_backup'
+    ) INTO backup_exists;
+
+    -- Determine status and message
+    IF ns_type = 'regnamespace' AND NOT backup_exists THEN
+        -- Normal state - ready for use, needs prep before pg_upgrade
+        RETURN QUERY SELECT
+            'NORMAL'::text,
+            ns_type,
+            g_count,
+            backup_exists,
+            'Run SELECT age_prepare_pg_upgrade(); before pg_upgrade'::text;
+    ELSIF ns_type = 'regnamespace' AND backup_exists THEN
+        -- Unusual state - backup exists but schema wasn't converted
+        RETURN QUERY SELECT
+            'WARNING'::text,
+            ns_type,
+            g_count,
+            backup_exists,
+            'Backup table exists but schema not converted. Run age_prepare_pg_upgrade() again.'::text;
+    ELSIF ns_type = 'oid' AND backup_exists THEN
+        -- Prepared and ready for pg_upgrade, or awaiting finish after pg_upgrade
+        RETURN QUERY SELECT
+            'PREPARED - AWAITING FINISH'::text,
+            ns_type,
+            g_count,
+            backup_exists,
+            'After pg_upgrade, run SELECT age_finish_pg_upgrade();'::text;
+    ELSE
+        -- oid type without backup - manually converted or partial state
+        RETURN QUERY SELECT
+            'CONVERTED'::text,
+            ns_type,
+            g_count,
+            backup_exists,
+            'Namespace is oid type. If upgrading, ensure backup table exists.'::text;
+    END IF;
+END;
+$function$;
+
+COMMENT ON FUNCTION ag_catalog.age_pg_upgrade_status() IS
+'Returns the current pg_upgrade readiness status of the AGE installation.';