PostgreSQL extension

GRANT pg_monitor TO dynatrace;

CREATE SCHEMA IF NOT EXISTS dynatrace;

GRANT USAGE ON SCHEMA dynatrace TO dynatrace;

CREATE OR REPLACE FUNCTION dynatrace.dynatrace_execution_plan(
   query text,
   OUT explain JSON
) RETURNS SETOF JSON
   LANGUAGE plpgsql
   VOLATILE
   RETURNS NULL ON NULL INPUT
   SECURITY DEFINER
   ROWS 1
   SET plan_cache_mode = force_generic_plan
   SET statement_timeout = '10s'
   SET lock_timeout = '200ms'
AS
$$DECLARE
   sanitized   text;
   nparams     integer;
   explain_cmd text;
   json_result json;
BEGIN
   /* remove literals, comments and quoted identifiers */
   sanitized := dynatrace_execution_plan.query;
   sanitized := pg_catalog.regexp_replace(sanitized, '\$([A-Za-z_]\w*)?\$.*?\$\1\$', '', 'gs');
   sanitized := pg_catalog.regexp_replace(sanitized, '''(''''|\\.|[^''\\])*''', '', 'g');
   sanitized := pg_catalog.regexp_replace(sanitized, '"(""|[^"])*"', '', 'g');
   sanitized := pg_catalog.regexp_replace(sanitized, '/\*.*?\*/', '', 'gs');
   sanitized := pg_catalog.regexp_replace(sanitized, '--[^\n]*', '', 'g');
   /* reject statements containing a semicolon in the middle */
   IF pg_catalog.strpos(
         pg_catalog.regexp_replace(sanitized, '[;[:space:]]+$', ''),
         ';'
      ) OPERATOR(pg_catalog.>) 0 THEN
      RAISE EXCEPTION 'query string must not contain a semicolon';
   END IF;
   /* drop a statement left behind by an earlier failure */
   IF EXISTS (
         SELECT 1
         FROM pg_catalog.pg_prepared_statements ps
         WHERE ps.name OPERATOR(pg_catalog.=) '_dt_plan_stmt_'
      ) THEN
      DEALLOCATE _dt_plan_stmt_;
   END IF;
   /* construct a prepared statement, letting the planner infer the parameter types */
   EXECUTE
      pg_catalog.concat(
         'PREPARE _dt_plan_stmt_ AS ',
         dynatrace_execution_plan.query
      );
   BEGIN
      /* get the parameter count */
      SELECT COALESCE(pg_catalog.array_length(ps.parameter_types, 1), 0)
      INTO nparams
      FROM pg_catalog.pg_prepared_statements ps
      WHERE ps.name OPERATOR(pg_catalog.=) '_dt_plan_stmt_';
      /* construct an EXPLAIN statement */
      explain_cmd :=
         pg_catalog.concat(
            'EXPLAIN (FORMAT JSON, ANALYZE FALSE) EXECUTE _dt_plan_stmt_',
            CASE
               WHEN nparams OPERATOR(pg_catalog.=) 0 THEN ''
               ELSE pg_catalog.concat(
                       '(',
                       pg_catalog.rtrim(
                          pg_catalog.repeat('NULL,', nparams),
                          ','
                       ),
                       ')'
                    )
            END
         );
      /* get the plan */
      EXECUTE explain_cmd INTO json_result;
   EXCEPTION
      WHEN OTHERS THEN
         /* delete the prepared statement, then report the original error */
         DEALLOCATE _dt_plan_stmt_;
         RAISE;
   END;
   /* delete the prepared statement */
   DEALLOCATE _dt_plan_stmt_;
   /* return the plan */
   RETURN QUERY SELECT json_result;
END;$$;

ALTER USER dynatrace SET search_path to dynatrace, public;

CREATE EXTENSION IF NOT EXISTS pg_stat_statements SCHEMA public;

CREATE EXTENSION IF NOT EXISTS hypopg SCHEMA public;

CREATE OR REPLACE FUNCTION dynatrace.dynatrace_index_advisor(query text, OUT result json)
 RETURNS SETOF json
 LANGUAGE plpgsql
 STRICT SECURITY DEFINER ROWS 1
 SET search_path TO 'pg_catalog', 'dynatrace', 'public'
AS $function$
DECLARE
    n_args                  int;
    prepared_statement_name text = 'dynatrace_ia_stmt';
    hypopg_schema_name      text;
    explain_plan_statement  text;
    error_message           text;
    rec                     record;
    plan_initial            jsonb;
    plan_final              jsonb;
    statements              text[] = '{}';
    referenced_oids         oid[];
    initial_cost            numeric;
    final_cost              numeric;
    initial_startup         numeric;
    final_startup           numeric;
    cost_reduction_pct      numeric;
    cost_reduction          numeric;
    min_cost_threshold      numeric := 200;
BEGIN
    hypopg_schema_name := (SELECT extnamespace::regnamespace::text FROM pg_extension WHERE extname = 'hypopg');

    IF hypopg_schema_name IS NULL THEN
        RETURN QUERY SELECT json_build_object(
            'database', current_database(),
            'query',    query,
            'errors',   array['hypopg extension is required but not installed']::text[]
        );
        RETURN;
    END IF;

    query := trim(
        regexp_replace(
            regexp_replace(
                regexp_replace(query, '\/\*[\s\S]+?\*\/', '', 'g'),
            '--[^\r\n]*', ' ', 'g'),
        '\s+', ' ', 'g')
    );
    query := regexp_replace(query, ';\s*$', '');

    BEGIN
        IF query ILIKE '%;%' THEN
            RAISE EXCEPTION 'Query must not contain a semicolon';
        END IF;

        query := replace(query, 'WITH pgrst_payload AS (SELECT $1 AS json_data)', 'WITH pgrst_payload AS (SELECT $1::json AS json_data)');

        query := regexp_replace(
            query,
            '\m(INTERVAL|TIMESTAMPTZ|TIMESTAMP|DATE|TIME)\s+(\$\d+)',
            '\2::\1',
            'gi'
        );

        query := regexp_replace(
            query,
            '\mEXTRACT\s*\(\s*(\$\d+)\s+FROM\y',
            'date_part(\1::text,',
            'gi'
        );

        BEGIN
            EXECUTE format('DEALLOCATE %I', prepared_statement_name);
        EXCEPTION WHEN OTHERS THEN
            NULL;
        END;

        EXECUTE format('PREPARE %I AS %s', prepared_statement_name, query);

        n_args := (
            SELECT COALESCE(array_length(parameter_types, 1), 0)
            FROM pg_prepared_statements
            WHERE name = prepared_statement_name
            LIMIT 1
        );

        explain_plan_statement := format(
            'SET LOCAL plan_cache_mode = force_generic_plan; EXPLAIN (FORMAT JSON, VERBOSE) EXECUTE %I%s',
            prepared_statement_name,
            CASE
                WHEN n_args = 0 THEN ''
                ELSE format('(%s)', array_to_string(array_fill('null'::text, array[n_args]), ','))
            END
        );

        EXECUTE explain_plan_statement INTO plan_initial;
        initial_cost    := (plan_initial -> 0 -> 'Plan' ->> 'Total Cost')::numeric;
        initial_startup := (plan_initial -> 0 -> 'Plan' ->> 'Startup Cost')::numeric;

        IF initial_cost IS NULL OR initial_cost < min_cost_threshold THEN
            EXECUTE format('DEALLOCATE %I', prepared_statement_name);
            RETURN;
        END IF;

        BEGIN
            WITH RECURSIVE nodes AS (
                SELECT (plan_initial -> 0 -> 'Plan') AS node
                UNION ALL
                SELECT child.value
                FROM nodes n
                CROSS JOIN LATERAL jsonb_array_elements(n.node -> 'Plans') AS child(value)
                WHERE n.node ? 'Plans'
            )
            SELECT array_agg(DISTINCT rel.oid)
            INTO referenced_oids
            FROM nodes n
            CROSS JOIN LATERAL (
                SELECT format('%I.%I', n.node ->> 'Schema', n.node ->> 'Relation Name')::regclass::oid AS oid
                WHERE (n.node ->> 'Relation Name') IS NOT NULL
                  AND (n.node ->> 'Schema') IS NOT NULL
            ) rel;
        EXCEPTION WHEN OTHERS THEN
            referenced_oids := NULL;
        END;

        IF referenced_oids IS NULL OR array_length(referenced_oids, 1) IS NULL THEN
            EXECUTE format('DEALLOCATE %I', prepared_statement_name);
            RETURN QUERY SELECT json_build_object(
                'database',            current_database(),
                'query',               query,
                'startup_cost_before', initial_startup,
                'startup_cost_after',  initial_startup,
                'total_cost_before',   initial_cost,
                'total_cost_after',    initial_cost,
                'index_statements',    array[]::text[]
            );
            RETURN;
        END IF;

        FOR rec IN (
            WITH extension_regclass AS (
                SELECT DISTINCT objid AS oid
                FROM pg_catalog.pg_depend
                WHERE deptype = 'e'
            )
            SELECT
                format(
                    'SELECT %I.hypopg_create_index(''CREATE INDEX ON %I.%I(%I)'')',
                    hypopg_schema_name,
                    pc.relnamespace::regnamespace::text,
                    pc.relname,
                    pa.attname
                ) AS hypopg_statement
            FROM pg_catalog.pg_class pc
            JOIN pg_catalog.pg_attribute pa ON pc.oid = pa.attrelid
            LEFT JOIN extension_regclass er ON pc.oid = er.oid
            LEFT JOIN pg_catalog.pg_index pi
                ON pc.oid = pi.indrelid
                AND pi.indkey[0] = pa.attnum
                AND pi.indisvalid
                AND pi.indexprs IS NULL
                AND pi.indpred IS NULL
                AND (SELECT am.amname FROM pg_catalog.pg_class ic
                     JOIN pg_catalog.pg_am am ON am.oid = ic.relam
                     WHERE ic.oid = pi.indexrelid) = 'btree'
            WHERE pc.relnamespace::regnamespace::text NOT IN ('pg_catalog', 'pg_toast', 'information_schema')
              AND er.oid IS NULL
              AND pc.relkind IN ('r', 'm')
              AND pc.relpersistence = 'p'
              AND pa.attnum > 0
              AND NOT pa.attisdropped
              AND pi.indrelid IS NULL
              AND pa.atttypid IN (20,16,1082,1184,1114,701,23,21,700,1083,2950,1700,25,18,1042,1043,774,790,829,869,1186,1266)
              AND pc.oid = ANY(referenced_oids)
        )
        LOOP
            EXECUTE rec.hypopg_statement;
        END LOOP;

        EXECUTE format('DEALLOCATE %I', prepared_statement_name);
        EXECUTE format('PREPARE %I AS %s', prepared_statement_name, query);

        EXECUTE explain_plan_statement INTO plan_final;
        final_cost    := (plan_final -> 0 -> 'Plan' ->> 'Total Cost')::numeric;
        final_startup := (plan_final -> 0 -> 'Plan' ->> 'Startup Cost')::numeric;

        EXECUTE format(
            'SELECT COALESCE(array_agg(%I.hypopg_get_indexdef(indexrelid) ORDER BY indrelid, indkey::text), ''{}''::text[])
            FROM %I.hypopg()
            WHERE position(indexname IN %s) > 0',
            hypopg_schema_name,
            hypopg_schema_name,
            quote_literal(plan_final)::text
        ) INTO statements;

        EXECUTE format('SELECT %I.hypopg_reset()', hypopg_schema_name);

        BEGIN
            EXECUTE format('DEALLOCATE %I', prepared_statement_name);
        EXCEPTION WHEN OTHERS THEN
            NULL;
        END;

        IF final_cost IS NULL OR initial_startup IS NULL OR final_startup IS NULL THEN
            RETURN;
        END IF;

        IF final_cost > initial_cost THEN
            RETURN;
        END IF;

        IF initial_cost > 0 THEN
            cost_reduction_pct := round(
                (initial_cost - final_cost) / initial_cost * 100,
                2
            );
        ELSE
            cost_reduction_pct := 0;
        END IF;

        cost_reduction := GREATEST(initial_cost - final_cost, 0);

        RETURN QUERY SELECT json_build_object(
            'database',            current_database(),
            'query',               query,
            'startup_cost_before', initial_startup,
            'startup_cost_after',  final_startup,
            'total_cost_before',   initial_cost,
            'total_cost_after',    final_cost,
            'cost_reduction_pct',  cost_reduction_pct,
            'cost_reduction',      cost_reduction,
            'index_statements',    statements
        );

    EXCEPTION WHEN OTHERS THEN
        GET STACKED DIAGNOSTICS error_message = MESSAGE_TEXT;
        BEGIN EXECUTE format('SELECT %I.hypopg_reset()', hypopg_schema_name); EXCEPTION WHEN OTHERS THEN NULL; END;
        BEGIN EXECUTE format('DEALLOCATE %I', prepared_statement_name); EXCEPTION WHEN OTHERS THEN NULL; END;

        RETURN QUERY SELECT json_build_object(
            'database',         current_database(),
            'query',            query,
            'index_statements', array[]::text[],
            'errors',           array[error_message]::text[]
        );
    END;
END;
$function$;