-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add cursor_tuple_fraction setting to search function (#143)
- Loading branch information
Showing
6 changed files
with
3,178 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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'); |
Oops, something went wrong.