add support for postgrest prepared statements

Author: olirice

README.md

@@ -9,6 +9,38 @@
 
 ---
 
+```sql
+select
+    *
+from
+    index_advisor('
+        select
+            book.id,
+            book.title,
+            publisher.name as publisher_name,
+            author.name as author_name,
+            review.body review_body
+        from
+            book
+            join publisher
+                on book.publisher_id = publisher.id
+            join author
+                on book.author_id = author.id
+            join review
+                on book.id = review.book_id
+        where
+            author.id = $1
+            and publisher.id = $2
+    ');
+
+ startup_cost_before | startup_cost_after | total_cost_before | total_cost_after |                  index_statements
+---------------------+--------------------+-------------------+------------------+----------------------------------------------------------
+ 27.26               | 12.77              | 68.48             | 42.37            | {"CREATE INDEX ON public.book USING btree (author_id)",
+                                                                                    "CREATE INDEX ON public.book USING btree (publisher_id)",
+                                                                                    "CREATE INDEX ON public.review USING btree (book_id)"}
+(1 row)
+```
+
 
 A PostgreSQL extension for recommending indexes to improve query performance.
 

index_advisor--0.1.0.sql

@@ -1,7 +1,21 @@
+create type index_advisor_output as (
+    index_statements text[],
+    startup_cost_before jsonb,
+    startup_cost_after jsonb,
+    total_cost_before jsonb,
+    total_cost_after jsonb
+);
+
 create function index_advisor(
     query text
 )
-    returns table( index_statement text )
+    returns table  (
+        startup_cost_before jsonb,
+        startup_cost_after jsonb,
+        total_cost_before jsonb,
+        total_cost_after jsonb,
+        index_statements text[]
+    )
     volatile
     language plpgsql
     as $$
@@ -21,6 +35,9 @@ begin
         raise exception 'query must not contain a semicolon';
     end if;
 
+    -- Hack to support PostgREST because the prepared statement for args incorrectly defaults to text
+    query := replace(query, 'WITH pgrst_payload AS (SELECT $1 AS json_data)', 'WITH pgrst_payload AS (SELECT $1::json AS json_data)');
+
     -- Create a prepared statement for the given query
     deallocate all;
     execute format('prepare %I as %s', prepared_statement_name, query);
@@ -137,7 +154,13 @@ begin
     -- Reset prepared statements
     deallocate all;
 
-    return query select * from unnest(statements);
+    return query values (
+        (plan_initial -> 0 -> 'Plan' -> 'Startup Cost'),
+        (plan_final -> 0 -> 'Plan' -> 'Startup Cost'),
+        (plan_initial -> 0 -> 'Plan' -> 'Total Cost'),
+        (plan_final -> 0 -> 'Plan' -> 'Total Cost'),
+        statements::text[]
+    );
 
 end;
 $$;

test/expected/integration.out

@@ -1,16 +1,16 @@
 begin;
     create extension index_advisor cascade;
 NOTICE:  installing required extension "hypopg"
-create table public.book(
-    id int,
-    name text
-);
-select index_advisor($$
-    select * from book where id = $1
-$$);
-                index_advisor                 
-----------------------------------------------
- CREATE INDEX ON public.book USING btree (id)
+    create table public.book(
+        id int,
+        name text
+    );
+    select index_advisor($$
+        select * from book where id = $1
+    $$);
+                                index_advisor                                 
+------------------------------------------------------------------------------
+ (0.00,4.07,25.88,13.54,"{""CREATE INDEX ON public.book USING btree (id)""}")
 (1 row)
 
 rollback;

test/expected/multi_index.out

@@ -22,6 +22,8 @@ NOTICE:  installing required extension "hypopg"
         body text not null
     );
     select
+        *
+    from
         index_advisor('
             select
                 book.id,
@@ -41,11 +43,9 @@ NOTICE:  installing required extension "hypopg"
                 author.id = $1
                 and publisher.id = $2
         ');
-                     index_advisor                      
---------------------------------------------------------
- CREATE INDEX ON public.book USING btree (author_id)
- CREATE INDEX ON public.book USING btree (publisher_id)
- CREATE INDEX ON public.review USING btree (book_id)
-(3 rows)
+ startup_cost_before | startup_cost_after | total_cost_before | total_cost_after |                                                                            index_statements                                                                            
+---------------------+--------------------+-------------------+------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+ 27.26               | 12.77              | 68.48             | 42.37            | {"CREATE INDEX ON public.book USING btree (author_id)","CREATE INDEX ON public.book USING btree (publisher_id)","CREATE INDEX ON public.review USING btree (book_id)"}
+(1 row)
 
 rollback;

test/expected/postgrest_query.out

@@ -0,0 +1,13 @@
+begin;
+    create extension index_advisor cascade;
+NOTICE:  installing required extension "hypopg"
+    create function get_info(x int) returns text language sql as $$ select 'foo' $$;
+    select index_advisor($$
+        WITH pgrst_source AS (WITH pgrst_payload AS (SELECT $1 AS json_data), pgrst_body AS ( SELECT CASE WHEN json_typeof(json_data) = $4 THEN json_data ELSE json_build_array(json_data) END AS val FROM pgrst_payload), pgrst_args AS ( SELECT * FROM json_to_recordset((SELECT val FROM pgrst_body)) AS _("x" integer) )SELECT "public"."get_info"("x" := (SELECT "x" FROM pgrst_args LIMIT $5)) AS pgrst_scalar) SELECT $6::bigint AS total_result_set, pg_catalog.count(_postgrest_t) AS page_total, coalesce((json_agg(_postgrest_t.pgrst_scalar)->$7)::text, $8) AS body, nullif(current_setting($9, $10), $11) AS response_headers, nullif(current_setting($12, $13), $14) AS response_status FROM (SELECT "get_info".* FROM "pgrst_source" AS "get_info"    LIMIT $2 OFFSET $3) _postgrest_t
+    $$);
+      index_advisor       
+--------------------------
+ (0.07,0.07,0.10,0.10,{})
+(1 row)
+
+rollback;

test/sql/integration.sql

@@ -2,13 +2,13 @@ begin;
 
     create extension index_advisor cascade;
 
-create table public.book(
-    id int,
-    name text
-);
-
-select index_advisor($$
-    select * from book where id = $1
-$$);
+    create table public.book(
+        id int,
+        name text
+    );
+
+    select index_advisor($$
+        select * from book where id = $1
+    $$);
 
 rollback;

test/sql/multi_index.sql

@@ -27,6 +27,8 @@ begin;
     );
 
     select
+        *
+    from
         index_advisor('
             select
                 book.id,

test/sql/postgrest_query.sql

@@ -0,0 +1,11 @@
+begin;
+
+    create extension index_advisor cascade;
+
+    create function get_info(x int) returns text language sql as $$ select 'foo' $$;
+
+    select index_advisor($$
+        WITH pgrst_source AS (WITH pgrst_payload AS (SELECT $1 AS json_data), pgrst_body AS ( SELECT CASE WHEN json_typeof(json_data) = $4 THEN json_data ELSE json_build_array(json_data) END AS val FROM pgrst_payload), pgrst_args AS ( SELECT * FROM json_to_recordset((SELECT val FROM pgrst_body)) AS _("x" integer) )SELECT "public"."get_info"("x" := (SELECT "x" FROM pgrst_args LIMIT $5)) AS pgrst_scalar) SELECT $6::bigint AS total_result_set, pg_catalog.count(_postgrest_t) AS page_total, coalesce((json_agg(_postgrest_t.pgrst_scalar)->$7)::text, $8) AS body, nullif(current_setting($9, $10), $11) AS response_headers, nullif(current_setting($12, $13), $14) AS response_status FROM (SELECT "get_info".* FROM "pgrst_source" AS "get_info"    LIMIT $2 OFFSET $3) _postgrest_t
+    $$);
+
+rollback;