Skip to content

Commit 1c0694e

Browse files
Make sure pgtap errors get caught (#46)
* Make sure features is at least [] * echo the entire pgtap log * Fixes for bugs with incorrectly referred parameters sort vs sortby, collection vs collections, and id vs ids * Add tests for fixes * bump version to 0.3.3
1 parent 77d8a71 commit 1c0694e

File tree

9 files changed

+2102
-117
lines changed

9 files changed

+2102
-117
lines changed

pypgstac/pypgstac/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
"""PyPGStac Version."""
2-
__version__ = "0.3.2"
2+
__version__ = "0.3.3"
Lines changed: 361 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,361 @@
1+
SET SEARCH_PATH to pgstac, public;
2+
set check_function_bodies = off;
3+
4+
CREATE OR REPLACE FUNCTION pgstac.add_filters_to_cql(j jsonb)
5+
RETURNS jsonb
6+
LANGUAGE plpgsql
7+
AS $function$
8+
DECLARE
9+
newprop jsonb;
10+
newprops jsonb := '[]'::jsonb;
11+
BEGIN
12+
IF j ? 'ids' THEN
13+
newprop := jsonb_build_object(
14+
'in',
15+
jsonb_build_array(
16+
'{"property":"id"}'::jsonb,
17+
j->'ids'
18+
)
19+
);
20+
newprops := jsonb_insert(newprops, '{1}', newprop);
21+
END IF;
22+
IF j ? 'collections' THEN
23+
newprop := jsonb_build_object(
24+
'in',
25+
jsonb_build_array(
26+
'{"property":"collection"}'::jsonb,
27+
j->'collections'
28+
)
29+
);
30+
newprops := jsonb_insert(newprops, '{1}', newprop);
31+
END IF;
32+
33+
IF j ? 'datetime' THEN
34+
newprop := format(
35+
'{"anyinteracts":[{"property":"datetime"}, %s]}',
36+
j->'datetime'
37+
);
38+
newprops := jsonb_insert(newprops, '{1}', newprop);
39+
END IF;
40+
41+
IF j ? 'bbox' THEN
42+
newprop := format(
43+
'{"intersects":[{"property":"geometry"}, %s]}',
44+
j->'bbox'
45+
);
46+
newprops := jsonb_insert(newprops, '{1}', newprop);
47+
END IF;
48+
49+
IF j ? 'intersects' THEN
50+
newprop := format(
51+
'{"intersects":[{"property":"geometry"}, %s]}',
52+
j->'intersects'
53+
);
54+
newprops := jsonb_insert(newprops, '{1}', newprop);
55+
END IF;
56+
57+
RAISE NOTICE 'newprops: %', newprops;
58+
59+
IF newprops IS NOT NULL AND jsonb_array_length(newprops) > 0 THEN
60+
return jsonb_set(
61+
j,
62+
'{filter}',
63+
cql_and_append(j, jsonb_build_object('and', newprops))
64+
) - '{ids,collections,datetime,bbox,intersects}'::text[];
65+
END IF;
66+
67+
return j;
68+
END;
69+
$function$
70+
;
71+
72+
CREATE OR REPLACE FUNCTION pgstac.get_token_filter(_search jsonb DEFAULT '{}'::jsonb, token_rec jsonb DEFAULT NULL::jsonb)
73+
RETURNS text
74+
LANGUAGE plpgsql
75+
AS $function$
76+
DECLARE
77+
token_id text;
78+
filters text[] := '{}'::text[];
79+
prev boolean := TRUE;
80+
field text;
81+
dir text;
82+
sort record;
83+
orfilters text[] := '{}'::text[];
84+
andfilters text[] := '{}'::text[];
85+
output text;
86+
token_where text;
87+
BEGIN
88+
-- If no token provided return NULL
89+
IF token_rec IS NULL THEN
90+
IF NOT (_search ? 'token' AND
91+
(
92+
(_search->>'token' ILIKE 'prev:%')
93+
OR
94+
(_search->>'token' ILIKE 'next:%')
95+
)
96+
) THEN
97+
RETURN NULL;
98+
END IF;
99+
prev := (_search->>'token' ILIKE 'prev:%');
100+
token_id := substr(_search->>'token', 6);
101+
SELECT to_jsonb(items) INTO token_rec FROM items WHERE id=token_id;
102+
END IF;
103+
RAISE NOTICE 'TOKEN ID: %', token_rec->'id';
104+
105+
CREATE TEMP TABLE sorts (
106+
_row int GENERATED ALWAYS AS IDENTITY NOT NULL,
107+
_field text PRIMARY KEY,
108+
_dir text NOT NULL,
109+
_val text
110+
) ON COMMIT DROP;
111+
112+
-- Make sure we only have distinct columns to sort with taking the first one we get
113+
INSERT INTO sorts (_field, _dir)
114+
SELECT
115+
(items_path(value->>'field')).path,
116+
get_sort_dir(value)
117+
FROM
118+
jsonb_array_elements(coalesce(_search->'sortby','[{"field":"datetime","direction":"desc"}]'))
119+
ON CONFLICT DO NOTHING
120+
;
121+
RAISE NOTICE 'sorts 1: %', (SELECT jsonb_agg(to_json(sorts)) FROM sorts);
122+
-- Get the first sort direction provided. As the id is a primary key, if there are any
123+
-- sorts after id they won't do anything, so make sure that id is the last sort item.
124+
SELECT _dir INTO dir FROM sorts ORDER BY _row ASC LIMIT 1;
125+
IF EXISTS (SELECT 1 FROM sorts WHERE _field = 'id') THEN
126+
DELETE FROM sorts WHERE _row > (SELECT _row FROM sorts WHERE _field = 'id' ORDER BY _row ASC);
127+
ELSE
128+
INSERT INTO sorts (_field, _dir) VALUES ('id', dir);
129+
END IF;
130+
131+
-- Add value from looked up item to the sorts table
132+
UPDATE sorts SET _val=quote_literal(token_rec->>_field);
133+
134+
-- Check if all sorts are the same direction and use row comparison
135+
-- to filter
136+
RAISE NOTICE 'sorts 2: %', (SELECT jsonb_agg(to_json(sorts)) FROM sorts);
137+
138+
IF (SELECT count(DISTINCT _dir) FROM sorts) = 1 THEN
139+
SELECT format(
140+
'(%s) %s (%s)',
141+
concat_ws(', ', VARIADIC array_agg(quote_ident(_field))),
142+
CASE WHEN (prev AND dir = 'ASC') OR (NOT prev AND dir = 'DESC') THEN '<' ELSE '>' END,
143+
concat_ws(', ', VARIADIC array_agg(_val))
144+
) INTO output FROM sorts
145+
WHERE token_rec ? _field
146+
;
147+
ELSE
148+
FOR sort IN SELECT * FROM sorts ORDER BY _row asc LOOP
149+
RAISE NOTICE 'SORT: %', sort;
150+
IF sort._row = 1 THEN
151+
orfilters := orfilters || format('(%s %s %s)',
152+
quote_ident(sort._field),
153+
CASE WHEN (prev AND sort._dir = 'ASC') OR (NOT prev AND sort._dir = 'DESC') THEN '<' ELSE '>' END,
154+
sort._val
155+
);
156+
ELSE
157+
orfilters := orfilters || format('(%s AND %s %s %s)',
158+
array_to_string(andfilters, ' AND '),
159+
quote_ident(sort._field),
160+
CASE WHEN (prev AND sort._dir = 'ASC') OR (NOT prev AND sort._dir = 'DESC') THEN '<' ELSE '>' END,
161+
sort._val
162+
);
163+
164+
END IF;
165+
andfilters := andfilters || format('%s = %s',
166+
quote_ident(sort._field),
167+
sort._val
168+
);
169+
END LOOP;
170+
output := array_to_string(orfilters, ' OR ');
171+
END IF;
172+
DROP TABLE IF EXISTS sorts;
173+
token_where := concat('(',coalesce(output,'true'),')');
174+
IF trim(token_where) = '' THEN
175+
token_where := NULL;
176+
END IF;
177+
RAISE NOTICE 'TOKEN_WHERE: |%|',token_where;
178+
RETURN token_where;
179+
END;
180+
$function$
181+
;
182+
183+
CREATE OR REPLACE FUNCTION pgstac.search(_search jsonb DEFAULT '{}'::jsonb)
184+
RETURNS jsonb
185+
LANGUAGE plpgsql
186+
SET jit TO 'off'
187+
AS $function$
188+
DECLARE
189+
searches searches%ROWTYPE;
190+
_where text;
191+
token_where text;
192+
full_where text;
193+
orderby text;
194+
query text;
195+
token_type text := substr(_search->>'token',1,4);
196+
_limit int := coalesce((_search->>'limit')::int, 10);
197+
curs refcursor;
198+
cntr int := 0;
199+
iter_record items%ROWTYPE;
200+
first_record items%ROWTYPE;
201+
last_record items%ROWTYPE;
202+
out_records jsonb := '[]'::jsonb;
203+
prev_query text;
204+
next text;
205+
prev_id text;
206+
has_next boolean := false;
207+
has_prev boolean := false;
208+
prev text;
209+
total_count bigint;
210+
context jsonb;
211+
collection jsonb;
212+
includes text[];
213+
excludes text[];
214+
exit_flag boolean := FALSE;
215+
batches int := 0;
216+
timer timestamptz := clock_timestamp();
217+
BEGIN
218+
searches := search_query(_search);
219+
_where := searches._where;
220+
orderby := searches.orderby;
221+
total_count := coalesce(searches.total_count, searches.estimated_count);
222+
223+
224+
IF token_type='prev' THEN
225+
token_where := get_token_filter(_search, null::jsonb);
226+
orderby := sort_sqlorderby(_search, TRUE);
227+
END IF;
228+
IF token_type='next' THEN
229+
token_where := get_token_filter(_search, null::jsonb);
230+
END IF;
231+
232+
full_where := concat_ws(' AND ', _where, token_where);
233+
RAISE NOTICE 'FULL QUERY % %', full_where, clock_timestamp()-timer;
234+
timer := clock_timestamp();
235+
236+
FOR query IN SELECT partition_queries(full_where, orderby) LOOP
237+
timer := clock_timestamp();
238+
query := format('%s LIMIT %L', query, _limit + 1);
239+
RAISE NOTICE 'Partition Query: %', query;
240+
batches := batches + 1;
241+
curs = create_cursor(query);
242+
LOOP
243+
FETCH curs into iter_record;
244+
EXIT WHEN NOT FOUND;
245+
cntr := cntr + 1;
246+
last_record := iter_record;
247+
IF cntr = 1 THEN
248+
first_record := last_record;
249+
END IF;
250+
IF cntr <= _limit THEN
251+
out_records := out_records || last_record.content;
252+
ELSIF cntr > _limit THEN
253+
has_next := true;
254+
exit_flag := true;
255+
EXIT;
256+
END IF;
257+
END LOOP;
258+
RAISE NOTICE 'Query took %', clock_timestamp()-timer;
259+
timer := clock_timestamp();
260+
EXIT WHEN exit_flag;
261+
END LOOP;
262+
RAISE NOTICE 'Scanned through % partitions.', batches;
263+
264+
265+
-- Flip things around if this was the result of a prev token query
266+
IF token_type='prev' THEN
267+
out_records := flip_jsonb_array(out_records);
268+
first_record := last_record;
269+
END IF;
270+
271+
-- If this query has a token, see if there is data before the first record
272+
IF _search ? 'token' THEN
273+
prev_query := format(
274+
'SELECT 1 FROM items WHERE %s LIMIT 1',
275+
concat_ws(
276+
' AND ',
277+
_where,
278+
trim(get_token_filter(_search, to_jsonb(first_record)))
279+
)
280+
);
281+
RAISE NOTICE 'Query to get previous record: % --- %', prev_query, first_record;
282+
EXECUTE prev_query INTO has_prev;
283+
IF FOUND and has_prev IS NOT NULL THEN
284+
RAISE NOTICE 'Query results from prev query: %', has_prev;
285+
has_prev := TRUE;
286+
END IF;
287+
END IF;
288+
has_prev := COALESCE(has_prev, FALSE);
289+
290+
RAISE NOTICE 'token_type: %, has_next: %, has_prev: %', token_type, has_next, has_prev;
291+
IF has_prev THEN
292+
prev := out_records->0->>'id';
293+
END IF;
294+
IF has_next OR token_type='prev' THEN
295+
next := out_records->-1->>'id';
296+
END IF;
297+
298+
299+
-- include/exclude any fields following fields extension
300+
IF _search ? 'fields' THEN
301+
IF _search->'fields' ? 'exclude' THEN
302+
excludes=textarr(_search->'fields'->'exclude');
303+
END IF;
304+
IF _search->'fields' ? 'include' THEN
305+
includes=textarr(_search->'fields'->'include');
306+
IF array_length(includes, 1)>0 AND NOT 'id' = ANY (includes) THEN
307+
includes = includes || '{id}';
308+
END IF;
309+
END IF;
310+
SELECT jsonb_agg(filter_jsonb(row, includes, excludes)) INTO out_records FROM jsonb_array_elements(out_records) row;
311+
END IF;
312+
313+
context := jsonb_strip_nulls(jsonb_build_object(
314+
'limit', _limit,
315+
'matched', total_count,
316+
'returned', coalesce(jsonb_array_length(out_records), 0)
317+
));
318+
319+
collection := jsonb_build_object(
320+
'type', 'FeatureCollection',
321+
'features', coalesce(out_records, '[]'::jsonb),
322+
'next', next,
323+
'prev', prev,
324+
'context', context
325+
);
326+
327+
RETURN collection;
328+
END;
329+
$function$
330+
;
331+
332+
CREATE OR REPLACE FUNCTION pgstac.sort_sqlorderby(_search jsonb DEFAULT NULL::jsonb, reverse boolean DEFAULT false)
333+
RETURNS text
334+
LANGUAGE sql
335+
AS $function$
336+
WITH sortby AS (
337+
SELECT coalesce(_search->'sortby','[{"field":"datetime", "direction":"desc"}]') as sort
338+
), withid AS (
339+
SELECT CASE
340+
WHEN sort @? '$[*] ? (@.field == "id")' THEN sort
341+
ELSE sort || '[{"field":"id", "direction":"desc"}]'::jsonb
342+
END as sort
343+
FROM sortby
344+
), withid_rows AS (
345+
SELECT jsonb_array_elements(sort) as value FROM withid
346+
),sorts AS (
347+
SELECT
348+
(items_path(value->>'field')).path as key,
349+
parse_sort_dir(value->>'direction', reverse) as dir
350+
FROM withid_rows
351+
)
352+
SELECT array_to_string(
353+
array_agg(concat(key, ' ', dir)),
354+
', '
355+
) FROM sorts;
356+
$function$
357+
;
358+
359+
360+
361+
INSERT INTO migrations (version) VALUES ('0.3.3');

0 commit comments

Comments
 (0)