count_aws_rows.plsql

CREATE OR REPLACE FUNCTION get_total_row_count_for_aws_tables()
RETURNS BIGINT AS $$
DECLARE
    -- Variable for dynamically constructed SQL to select counts
    dynamic_sql TEXT := 'SELECT SUM(row_count) FROM (';
    -- Variable for the final part of the SQL statement
    final_sql TEXT := ') AS counts(row_count);';
    -- Variable to hold each individual count query
    count_query TEXT;
    -- Variable to accumulate the full SQL query
    full_sql TEXT;
    -- Variable to hold the final total count
    total_count BIGINT;
BEGIN
    -- Loop through each matching table and construct a count query
    FOR count_query IN
        SELECT 'SELECT COUNT(*) AS row_count FROM ' || quote_ident(table_name)
        FROM information_schema.tables
        WHERE table_schema = 'public' -- Adjust your schema if necessary
          AND table_name LIKE 'aws\_%' ESCAPE '\' -- replace `aws` with `oracle` or `azure` for that plugin row count
          AND table_type = 'BASE TABLE'
    LOOP
        -- Append each count query with UNION ALL except the last one
        dynamic_sql := dynamic_sql || count_query || ' UNION ALL ';
    END LOOP;
    
    -- If no tables match, return 0 to avoid executing an invalid query
    IF dynamic_sql = 'SELECT SUM(row_count) FROM (' THEN
        RETURN 0;
    END IF;

    -- Remove the trailing 'UNION ALL ' and append the final SQL part
    dynamic_sql := rtrim(dynamic_sql, ' UNION ALL ');
    full_sql := dynamic_sql || final_sql;

    -- Execute the constructed query to calculate the total count
    EXECUTE full_sql INTO total_count;

    -- Return the total count
    RETURN COALESCE(total_count, 0);
END
$$ LANGUAGE plpgsql;
                                                  

SELECT get_total_row_count_for_aws_tables();