Skip to content

Instantly share code, notes, and snippets.

@Harkishen-Singh
Last active August 2, 2023 09:14
Show Gist options
  • Save Harkishen-Singh/ff06a40b8c90353704a97120d44b16ff to your computer and use it in GitHub Desktop.
Save Harkishen-Singh/ff06a40b8c90353704a97120d44b16ff to your computer and use it in GitHub Desktop.
A PLPGSQL function to increase database size until a `minimum_db_size` (in GBs) is reached. It does this by cloning `source_schema` into new schemas with prefix `dest_schema`. Useful for creating large datasets.
-- This function increases the DB size upto `minimum_db_size` by cloning the `source_schema` into a new name `dest_schema_<dynamic_count>`
CREATE OR REPLACE FUNCTION increase_db_size_upto(source_schema text, dest_schema text, minimum_db_size int, multiplier_factor int)
RETURNS void AS
$BODY$
DECLARE
    objeto record;
    buffer text;
    db_size int;
    schema_count int;
    schema_exists boolean;
    new_schema_name text;
BEGIN
    db_size := (SELECT pg_database_size(current_database()) / 1024 / 1024 / 1024);
    schema_count := 0;
    WHILE db_size < minimum_db_size LOOP
        schema_count := schema_count + 1;
        new_schema_name := dest_schema || '_' || multiplier_factor * schema_count;
        RAISE NOTICE 'Current database size: % GB. Minimum database size: % GB. Cloning schema to %.', db_size, minimum_db_size, new_schema_name;
        schema_exists := EXISTS (SELECT schema_name FROM information_schema.schemata WHERE schema_name = new_schema_name);
        IF NOT schema_exists THEN
            EXECUTE 'CREATE SCHEMA ' || new_schema_name ;
            FOR objeto IN SELECT c.relname::text, c.relkind FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = source_schema LOOP
                buffer := new_schema_name || '.' || objeto.relname;
                IF objeto.relkind = 'r' THEN
                    EXECUTE 'CREATE TABLE ' || buffer || ' (LIKE ' || source_schema || '.' || objeto.relname || ' INCLUDING DEFAULTS)';
                    EXECUTE 'INSERT INTO ' || buffer || ' (SELECT * FROM ' || source_schema || '.' || objeto.relname || ')';
                ELSIF objeto.relkind = 'v' THEN
                    EXECUTE (SELECT 'CREATE OR REPLACE VIEW ' || buffer || ' AS ' || definition FROM pg_views WHERE viewname = objeto.relname AND schemaname = source_schema);
                ELSIF objeto.relkind = 'm' THEN
                    EXECUTE (SELECT 'CREATE MATERIALIZED VIEW ' || buffer || ' AS ' || definition FROM pg_matviews WHERE matviewname = objeto.relname AND schemaname = source_schema);
                END IF;
            END LOOP;
        END IF;
        db_size := (SELECT pg_database_size(current_database()) / 1024 / 1024 / 1024);
    END LOOP;
END;
$BODY$
LANGUAGE plpgsql VOLATILE;

-- Example: SELECT clone_schema_upto('iot_1', 'iot_prefix_', 1500, 3); -- Increase database size till at least 1.5TBs.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment