add cursor_tuple_fraction setting to search function (#143)

Author: bitner

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [v0.6.9]
+
+### Fixed
+- Set cursor_tuple_fraction to 1 in search function to let query planner know to expect the entire table result within the search function to be returned. The default cursor_tuple_fraction of .1 within that function was at times creating bad query plans leading to slow queries.
+
 ## [v0.6.8]
 
 ### Added

pypgstac/pypgstac/migrations/pgstac.0.6.8-0.6.9.sql

@@ -0,0 +1,194 @@
+SET SEARCH_PATH to pgstac, public;
+set check_function_bodies = off;
+
+CREATE OR REPLACE FUNCTION pgstac.search(_search jsonb DEFAULT '{}'::jsonb)
+ RETURNS jsonb
+ LANGUAGE plpgsql
+ SECURITY DEFINER
+ SET search_path TO 'pgstac', 'public'
+ SET cursor_tuple_fraction TO '1'
+AS $function$
+DECLARE
+    searches searches%ROWTYPE;
+    _where text;
+    token_where text;
+    full_where text;
+    orderby text;
+    query text;
+    token_type text := substr(_search->>'token',1,4);
+    _limit int := coalesce((_search->>'limit')::int, 10);
+    curs refcursor;
+    cntr int := 0;
+    iter_record items%ROWTYPE;
+    first_record jsonb;
+    first_item items%ROWTYPE;
+    last_item items%ROWTYPE;
+    last_record jsonb;
+    out_records jsonb := '[]'::jsonb;
+    prev_query text;
+    next text;
+    prev_id text;
+    has_next boolean := false;
+    has_prev boolean := false;
+    prev text;
+    total_count bigint;
+    context jsonb;
+    collection jsonb;
+    includes text[];
+    excludes text[];
+    exit_flag boolean := FALSE;
+    batches int := 0;
+    timer timestamptz := clock_timestamp();
+    pstart timestamptz;
+    pend timestamptz;
+    pcurs refcursor;
+    search_where search_wheres%ROWTYPE;
+    id text;
+BEGIN
+CREATE TEMP TABLE results (content jsonb) ON COMMIT DROP;
+-- if ids is set, short circuit and just use direct ids query for each id
+-- skip any paging or caching
+-- hard codes ordering in the same order as the array of ids
+IF _search ? 'ids' THEN
+    INSERT INTO results
+    SELECT
+        CASE WHEN _search->'conf'->>'nohydrate' IS NOT NULL AND (_search->'conf'->>'nohydrate')::boolean = true THEN
+            content_nonhydrated(items, _search->'fields')
+        ELSE
+            content_hydrate(items, _search->'fields')
+        END
+    FROM items WHERE
+        items.id = ANY(to_text_array(_search->'ids'))
+        AND
+            CASE WHEN _search ? 'collections' THEN
+                items.collection = ANY(to_text_array(_search->'collections'))
+            ELSE TRUE
+            END
+    ORDER BY items.datetime desc, items.id desc
+    ;
+    SELECT INTO total_count count(*) FROM results;
+ELSE
+    searches := search_query(_search);
+    _where := searches._where;
+    orderby := searches.orderby;
+    search_where := where_stats(_where);
+    total_count := coalesce(search_where.total_count, search_where.estimated_count);
+
+    IF token_type='prev' THEN
+        token_where := get_token_filter(_search, null::jsonb);
+        orderby := sort_sqlorderby(_search, TRUE);
+    END IF;
+    IF token_type='next' THEN
+        token_where := get_token_filter(_search, null::jsonb);
+    END IF;
+
+    full_where := concat_ws(' AND ', _where, token_where);
+    RAISE NOTICE 'FULL QUERY % %', full_where, clock_timestamp()-timer;
+    timer := clock_timestamp();
+
+    FOR query IN SELECT partition_queries(full_where, orderby, search_where.partitions) LOOP
+        timer := clock_timestamp();
+        query := format('%s LIMIT %s', query, _limit + 1);
+        RAISE NOTICE 'Partition Query: %', query;
+        batches := batches + 1;
+        -- curs = create_cursor(query);
+        RAISE NOTICE 'cursor_tuple_fraction: %', current_setting('cursor_tuple_fraction');
+        OPEN curs FOR EXECUTE query;
+        LOOP
+            FETCH curs into iter_record;
+            EXIT WHEN NOT FOUND;
+            cntr := cntr + 1;
+
+            IF _search->'conf'->>'nohydrate' IS NOT NULL AND (_search->'conf'->>'nohydrate')::boolean = true THEN
+                last_record := content_nonhydrated(iter_record, _search->'fields');
+            ELSE
+                last_record := content_hydrate(iter_record, _search->'fields');
+            END IF;
+            last_item := iter_record;
+            IF cntr = 1 THEN
+                first_item := last_item;
+                first_record := last_record;
+            END IF;
+            IF cntr <= _limit THEN
+                INSERT INTO results (content) VALUES (last_record);
+            ELSIF cntr > _limit THEN
+                has_next := true;
+                exit_flag := true;
+                EXIT;
+            END IF;
+        END LOOP;
+        CLOSE curs;
+        RAISE NOTICE 'Query took %.', clock_timestamp()-timer;
+        timer := clock_timestamp();
+        EXIT WHEN exit_flag;
+    END LOOP;
+    RAISE NOTICE 'Scanned through % partitions.', batches;
+END IF;
+
+SELECT jsonb_agg(content) INTO out_records FROM results WHERE content is not NULL;
+
+DROP TABLE results;
+
+
+-- Flip things around if this was the result of a prev token query
+IF token_type='prev' THEN
+    out_records := flip_jsonb_array(out_records);
+    first_item := last_item;
+    first_record := last_record;
+END IF;
+
+-- If this query has a token, see if there is data before the first record
+IF _search ? 'token' THEN
+    prev_query := format(
+        'SELECT 1 FROM items WHERE %s LIMIT 1',
+        concat_ws(
+            ' AND ',
+            _where,
+            trim(get_token_filter(_search, to_jsonb(first_item)))
+        )
+    );
+    RAISE NOTICE 'Query to get previous record: % --- %', prev_query, first_record;
+    EXECUTE prev_query INTO has_prev;
+    IF FOUND and has_prev IS NOT NULL THEN
+        RAISE NOTICE 'Query results from prev query: %', has_prev;
+        has_prev := TRUE;
+    END IF;
+END IF;
+has_prev := COALESCE(has_prev, FALSE);
+
+IF has_prev THEN
+    prev := out_records->0->>'id';
+END IF;
+IF has_next OR token_type='prev' THEN
+    next := out_records->-1->>'id';
+END IF;
+
+IF context(_search->'conf') != 'off' THEN
+    context := jsonb_strip_nulls(jsonb_build_object(
+        'limit', _limit,
+        'matched', total_count,
+        'returned', coalesce(jsonb_array_length(out_records), 0)
+    ));
+ELSE
+    context := jsonb_strip_nulls(jsonb_build_object(
+        'limit', _limit,
+        'returned', coalesce(jsonb_array_length(out_records), 0)
+    ));
+END IF;
+
+collection := jsonb_build_object(
+    'type', 'FeatureCollection',
+    'features', coalesce(out_records, '[]'::jsonb),
+    'next', next,
+    'prev', prev,
+    'context', context
+);
+
+RETURN collection;
+END;
+$function$
+;
+
+
+
+SELECT set_version('0.6.9');