]> granicus.if.org Git - postgresql/blob - src/backend/utils/adt/jsonfuncs.c
Refine error messages
[postgresql] / src / backend / utils / adt / jsonfuncs.c
1 /*-------------------------------------------------------------------------
2  *
3  * jsonfuncs.c
4  *              Functions to process JSON data types.
5  *
6  * Portions Copyright (c) 1996-2018, PostgreSQL Global Development Group
7  * Portions Copyright (c) 1994, Regents of the University of California
8  *
9  * IDENTIFICATION
10  *        src/backend/utils/adt/jsonfuncs.c
11  *
12  *-------------------------------------------------------------------------
13  */
14
15 #include "postgres.h"
16
17 #include <limits.h>
18
19 #include "access/htup_details.h"
20 #include "catalog/pg_type.h"
21 #include "fmgr.h"
22 #include "funcapi.h"
23 #include "lib/stringinfo.h"
24 #include "mb/pg_wchar.h"
25 #include "miscadmin.h"
26 #include "utils/array.h"
27 #include "utils/builtins.h"
28 #include "utils/hsearch.h"
29 #include "utils/json.h"
30 #include "utils/jsonapi.h"
31 #include "utils/jsonb.h"
32 #include "utils/lsyscache.h"
33 #include "utils/memutils.h"
34 #include "utils/syscache.h"
35 #include "utils/typcache.h"
36
37 /* Operations available for setPath */
38 #define JB_PATH_CREATE                                  0x0001
39 #define JB_PATH_DELETE                                  0x0002
40 #define JB_PATH_REPLACE                                 0x0004
41 #define JB_PATH_INSERT_BEFORE                   0x0008
42 #define JB_PATH_INSERT_AFTER                    0x0010
43 #define JB_PATH_CREATE_OR_INSERT \
44         (JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER | JB_PATH_CREATE)
45
46 /* state for json_object_keys */
47 typedef struct OkeysState
48 {
49         JsonLexContext *lex;
50         char      **result;
51         int                     result_size;
52         int                     result_count;
53         int                     sent_count;
54 } OkeysState;
55
56 /* state for iterate_json_string_values function */
57 typedef struct IterateJsonStringValuesState
58 {
59         JsonLexContext *lex;
60         JsonIterateStringValuesAction action;   /* an action that will be applied
61                                                                                          * to each json value */
62         void       *action_state;       /* any necessary context for iteration */
63         uint32          flags;                  /* what kind of elements from a json we want
64                                                                  * to iterate */
65 } IterateJsonStringValuesState;
66
67 /* state for transform_json_string_values function */
68 typedef struct TransformJsonStringValuesState
69 {
70         JsonLexContext *lex;
71         StringInfo      strval;                 /* resulting json */
72         JsonTransformStringValuesAction action; /* an action that will be applied
73                                                                                          * to each json value */
74         void       *action_state;       /* any necessary context for transformation */
75 } TransformJsonStringValuesState;
76
77 /* state for json_get* functions */
78 typedef struct GetState
79 {
80         JsonLexContext *lex;
81         text       *tresult;
82         char       *result_start;
83         bool            normalize_results;
84         bool            next_scalar;
85         int                     npath;                  /* length of each path-related array */
86         char      **path_names;         /* field name(s) being sought */
87         int                *path_indexes;       /* array index(es) being sought */
88         bool       *pathok;                     /* is path matched to current depth? */
89         int                *array_cur_index;    /* current element index at each path
90                                                                          * level */
91 } GetState;
92
93 /* state for json_array_length */
94 typedef struct AlenState
95 {
96         JsonLexContext *lex;
97         int                     count;
98 } AlenState;
99
100 /* state for json_each */
101 typedef struct EachState
102 {
103         JsonLexContext *lex;
104         Tuplestorestate *tuple_store;
105         TupleDesc       ret_tdesc;
106         MemoryContext tmp_cxt;
107         char       *result_start;
108         bool            normalize_results;
109         bool            next_scalar;
110         char       *normalized_scalar;
111 } EachState;
112
113 /* state for json_array_elements */
114 typedef struct ElementsState
115 {
116         JsonLexContext *lex;
117         const char *function_name;
118         Tuplestorestate *tuple_store;
119         TupleDesc       ret_tdesc;
120         MemoryContext tmp_cxt;
121         char       *result_start;
122         bool            normalize_results;
123         bool            next_scalar;
124         char       *normalized_scalar;
125 } ElementsState;
126
127 /* state for get_json_object_as_hash */
128 typedef struct JHashState
129 {
130         JsonLexContext *lex;
131         const char *function_name;
132         HTAB       *hash;
133         char       *saved_scalar;
134         char       *save_json_start;
135         JsonTokenType saved_token_type;
136 } JHashState;
137
138 /* hashtable element */
139 typedef struct JsonHashEntry
140 {
141         char            fname[NAMEDATALEN]; /* hash key (MUST BE FIRST) */
142         char       *val;
143         JsonTokenType type;
144 } JsonHashEntry;
145
146 /* structure to cache type I/O metadata needed for populate_scalar() */
147 typedef struct ScalarIOData
148 {
149         Oid                     typioparam;
150         FmgrInfo        typiofunc;
151 } ScalarIOData;
152
153 /* these two structures are used recursively */
154 typedef struct ColumnIOData ColumnIOData;
155 typedef struct RecordIOData RecordIOData;
156
157 /* structure to cache metadata needed for populate_array() */
158 typedef struct ArrayIOData
159 {
160         ColumnIOData *element_info; /* metadata cache */
161         Oid                     element_type;   /* array element type id */
162         int32           element_typmod; /* array element type modifier */
163 } ArrayIOData;
164
165 /* structure to cache metadata needed for populate_composite() */
166 typedef struct CompositeIOData
167 {
168         /*
169          * We use pointer to a RecordIOData here because variable-length struct
170          * RecordIOData can't be used directly in ColumnIOData.io union
171          */
172         RecordIOData *record_io;        /* metadata cache for populate_record() */
173         TupleDesc       tupdesc;                /* cached tuple descriptor */
174         /* these fields differ from target type only if domain over composite: */
175         Oid                     base_typid;             /* base type id */
176         int32           base_typmod;    /* base type modifier */
177         /* this field is used only if target type is domain over composite: */
178         void       *domain_info;        /* opaque cache for domain checks */
179 } CompositeIOData;
180
181 /* structure to cache metadata needed for populate_domain() */
182 typedef struct DomainIOData
183 {
184         ColumnIOData *base_io;          /* metadata cache */
185         Oid                     base_typid;             /* base type id */
186         int32           base_typmod;    /* base type modifier */
187         void       *domain_info;        /* opaque cache for domain checks */
188 } DomainIOData;
189
190 /* enumeration type categories */
191 typedef enum TypeCat
192 {
193         TYPECAT_SCALAR = 's',
194         TYPECAT_ARRAY = 'a',
195         TYPECAT_COMPOSITE = 'c',
196         TYPECAT_COMPOSITE_DOMAIN = 'C',
197         TYPECAT_DOMAIN = 'd'
198 } TypeCat;
199
200 /* these two are stolen from hstore / record_out, used in populate_record* */
201
202 /* structure to cache record metadata needed for populate_record_field() */
203 struct ColumnIOData
204 {
205         Oid                     typid;                  /* column type id */
206         int32           typmod;                 /* column type modifier */
207         TypeCat         typcat;                 /* column type category */
208         ScalarIOData scalar_io;         /* metadata cache for directi conversion
209                                                                  * through input function */
210         union
211         {
212                 ArrayIOData array;
213                 CompositeIOData composite;
214                 DomainIOData domain;
215         }                       io;                             /* metadata cache for various column type
216                                                                  * categories */
217 };
218
219 /* structure to cache record metadata needed for populate_record() */
220 struct RecordIOData
221 {
222         Oid                     record_type;
223         int32           record_typmod;
224         int                     ncolumns;
225         ColumnIOData columns[FLEXIBLE_ARRAY_MEMBER];
226 };
227
228 /* per-query cache for populate_recordset */
229 typedef struct PopulateRecordsetCache
230 {
231         Oid                     argtype;                /* declared type of the record argument */
232         ColumnIOData c;                         /* metadata cache for populate_composite() */
233         MemoryContext fn_mcxt;          /* where this is stored */
234 } PopulateRecordsetCache;
235
236 /* per-call state for populate_recordset */
237 typedef struct PopulateRecordsetState
238 {
239         JsonLexContext *lex;
240         const char *function_name;
241         HTAB       *json_hash;
242         char       *saved_scalar;
243         char       *save_json_start;
244         JsonTokenType saved_token_type;
245         Tuplestorestate *tuple_store;
246         HeapTupleHeader rec;
247         PopulateRecordsetCache *cache;
248 } PopulateRecordsetState;
249
250 /* structure to cache metadata needed for populate_record_worker() */
251 typedef struct PopulateRecordCache
252 {
253         Oid                     argtype;                /* declared type of the record argument */
254         ColumnIOData c;                         /* metadata cache for populate_composite() */
255 } PopulateRecordCache;
256
257 /* common data for populate_array_json() and populate_array_dim_jsonb() */
258 typedef struct PopulateArrayContext
259 {
260         ArrayBuildState *astate;        /* array build state */
261         ArrayIOData *aio;                       /* metadata cache */
262         MemoryContext acxt;                     /* array build memory context */
263         MemoryContext mcxt;                     /* cache memory context */
264         const char *colname;            /* for diagnostics only */
265         int                *dims;                       /* dimensions */
266         int                *sizes;                      /* current dimension counters */
267         int                     ndims;                  /* number of dimensions */
268 } PopulateArrayContext;
269
270 /* state for populate_array_json() */
271 typedef struct PopulateArrayState
272 {
273         JsonLexContext *lex;            /* json lexer */
274         PopulateArrayContext *ctx;      /* context */
275         char       *element_start;      /* start of the current array element */
276         char       *element_scalar; /* current array element token if it is a
277                                                                  * scalar */
278         JsonTokenType element_type; /* current array element type */
279 } PopulateArrayState;
280
281 /* state for json_strip_nulls */
282 typedef struct StripnullState
283 {
284         JsonLexContext *lex;
285         StringInfo      strval;
286         bool            skip_next_null;
287 } StripnullState;
288
289 /* structure for generalized json/jsonb value passing */
290 typedef struct JsValue
291 {
292         bool            is_json;                /* json/jsonb */
293         union
294         {
295                 struct
296                 {
297                         char       *str;        /* json string */
298                         int                     len;    /* json string length or -1 if null-terminated */
299                         JsonTokenType type; /* json type */
300                 }                       json;           /* json value */
301
302                 JsonbValue *jsonb;              /* jsonb value */
303         }                       val;
304 } JsValue;
305
306 typedef struct JsObject
307 {
308         bool            is_json;                /* json/jsonb */
309         union
310         {
311                 HTAB       *json_hash;
312                 JsonbContainer *jsonb_cont;
313         }                       val;
314 } JsObject;
315
316 /* useful macros for testing JsValue properties */
317 #define JsValueIsNull(jsv) \
318         ((jsv)->is_json ?  \
319                 (!(jsv)->val.json.str || (jsv)->val.json.type == JSON_TOKEN_NULL) : \
320                 (!(jsv)->val.jsonb || (jsv)->val.jsonb->type == jbvNull))
321
322 #define JsValueIsString(jsv) \
323         ((jsv)->is_json ? (jsv)->val.json.type == JSON_TOKEN_STRING \
324                 : ((jsv)->val.jsonb && (jsv)->val.jsonb->type == jbvString))
325
326 #define JsObjectIsEmpty(jso) \
327         ((jso)->is_json \
328                 ? hash_get_num_entries((jso)->val.json_hash) == 0 \
329                 : ((jso)->val.jsonb_cont == NULL || \
330                    JsonContainerSize((jso)->val.jsonb_cont) == 0))
331
332 #define JsObjectFree(jso) \
333         do { \
334                 if ((jso)->is_json) \
335                         hash_destroy((jso)->val.json_hash); \
336         } while (0)
337
338 /* semantic action functions for json_object_keys */
339 static void okeys_object_field_start(void *state, char *fname, bool isnull);
340 static void okeys_array_start(void *state);
341 static void okeys_scalar(void *state, char *token, JsonTokenType tokentype);
342
343 /* semantic action functions for json_get* functions */
344 static void get_object_start(void *state);
345 static void get_object_end(void *state);
346 static void get_object_field_start(void *state, char *fname, bool isnull);
347 static void get_object_field_end(void *state, char *fname, bool isnull);
348 static void get_array_start(void *state);
349 static void get_array_end(void *state);
350 static void get_array_element_start(void *state, bool isnull);
351 static void get_array_element_end(void *state, bool isnull);
352 static void get_scalar(void *state, char *token, JsonTokenType tokentype);
353
354 /* common worker function for json getter functions */
355 static Datum get_path_all(FunctionCallInfo fcinfo, bool as_text);
356 static text *get_worker(text *json, char **tpath, int *ipath, int npath,
357                    bool normalize_results);
358 static Datum get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text);
359
360 /* semantic action functions for json_array_length */
361 static void alen_object_start(void *state);
362 static void alen_scalar(void *state, char *token, JsonTokenType tokentype);
363 static void alen_array_element_start(void *state, bool isnull);
364
365 /* common workers for json{b}_each* functions */
366 static Datum each_worker(FunctionCallInfo fcinfo, bool as_text);
367 static Datum each_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
368                                   bool as_text);
369
370 /* semantic action functions for json_each */
371 static void each_object_field_start(void *state, char *fname, bool isnull);
372 static void each_object_field_end(void *state, char *fname, bool isnull);
373 static void each_array_start(void *state);
374 static void each_scalar(void *state, char *token, JsonTokenType tokentype);
375
376 /* common workers for json{b}_array_elements_* functions */
377 static Datum elements_worker(FunctionCallInfo fcinfo, const char *funcname,
378                                 bool as_text);
379 static Datum elements_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
380                                           bool as_text);
381
382 /* semantic action functions for json_array_elements */
383 static void elements_object_start(void *state);
384 static void elements_array_element_start(void *state, bool isnull);
385 static void elements_array_element_end(void *state, bool isnull);
386 static void elements_scalar(void *state, char *token, JsonTokenType tokentype);
387
388 /* turn a json object into a hash table */
389 static HTAB *get_json_object_as_hash(char *json, int len, const char *funcname);
390
391 /* semantic actions for populate_array_json */
392 static void populate_array_object_start(void *_state);
393 static void populate_array_array_end(void *_state);
394 static void populate_array_element_start(void *_state, bool isnull);
395 static void populate_array_element_end(void *_state, bool isnull);
396 static void populate_array_scalar(void *_state, char *token, JsonTokenType tokentype);
397
398 /* semantic action functions for get_json_object_as_hash */
399 static void hash_object_field_start(void *state, char *fname, bool isnull);
400 static void hash_object_field_end(void *state, char *fname, bool isnull);
401 static void hash_array_start(void *state);
402 static void hash_scalar(void *state, char *token, JsonTokenType tokentype);
403
404 /* semantic action functions for populate_recordset */
405 static void populate_recordset_object_field_start(void *state, char *fname, bool isnull);
406 static void populate_recordset_object_field_end(void *state, char *fname, bool isnull);
407 static void populate_recordset_scalar(void *state, char *token, JsonTokenType tokentype);
408 static void populate_recordset_object_start(void *state);
409 static void populate_recordset_object_end(void *state);
410 static void populate_recordset_array_start(void *state);
411 static void populate_recordset_array_element_start(void *state, bool isnull);
412
413 /* semantic action functions for json_strip_nulls */
414 static void sn_object_start(void *state);
415 static void sn_object_end(void *state);
416 static void sn_array_start(void *state);
417 static void sn_array_end(void *state);
418 static void sn_object_field_start(void *state, char *fname, bool isnull);
419 static void sn_array_element_start(void *state, bool isnull);
420 static void sn_scalar(void *state, char *token, JsonTokenType tokentype);
421
422 /* worker functions for populate_record, to_record, populate_recordset and to_recordset */
423 static Datum populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
424                                                   bool have_record_arg);
425 static Datum populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
426                                            bool have_record_arg);
427
428 /* helper functions for populate_record[set] */
429 static HeapTupleHeader populate_record(TupleDesc tupdesc, RecordIOData **record_p,
430                                 HeapTupleHeader defaultval, MemoryContext mcxt,
431                                 JsObject *obj);
432 static void JsValueToJsObject(JsValue *jsv, JsObject *jso);
433 static Datum populate_composite(CompositeIOData *io, Oid typid,
434                                    const char *colname, MemoryContext mcxt,
435                                    HeapTupleHeader defaultval, JsValue *jsv, bool isnull);
436 static Datum populate_scalar(ScalarIOData *io, Oid typid, int32 typmod, JsValue *jsv);
437 static void prepare_column_cache(ColumnIOData *column, Oid typid, int32 typmod,
438                                          MemoryContext mcxt, bool need_scalar);
439 static Datum populate_record_field(ColumnIOData *col, Oid typid, int32 typmod,
440                                           const char *colname, MemoryContext mcxt, Datum defaultval,
441                                           JsValue *jsv, bool *isnull);
442 static RecordIOData *allocate_record_info(MemoryContext mcxt, int ncolumns);
443 static bool JsObjectGetField(JsObject *obj, char *field, JsValue *jsv);
444 static void populate_recordset_record(PopulateRecordsetState *state, JsObject *obj);
445 static void populate_array_json(PopulateArrayContext *ctx, char *json, int len);
446 static void populate_array_dim_jsonb(PopulateArrayContext *ctx, JsonbValue *jbv,
447                                                  int ndim);
448 static void populate_array_report_expected_array(PopulateArrayContext *ctx, int ndim);
449 static void populate_array_assign_ndims(PopulateArrayContext *ctx, int ndims);
450 static void populate_array_check_dimension(PopulateArrayContext *ctx, int ndim);
451 static void populate_array_element(PopulateArrayContext *ctx, int ndim, JsValue *jsv);
452 static Datum populate_array(ArrayIOData *aio, const char *colname,
453                            MemoryContext mcxt, JsValue *jsv);
454 static Datum populate_domain(DomainIOData *io, Oid typid, const char *colname,
455                                 MemoryContext mcxt, JsValue *jsv, bool isnull);
456
457 /* Worker that takes care of common setup for us */
458 static JsonbValue *findJsonbValueFromContainerLen(JsonbContainer *container,
459                                                            uint32 flags,
460                                                            char *key,
461                                                            uint32 keylen);
462
463 /* functions supporting jsonb_delete, jsonb_set and jsonb_concat */
464 static JsonbValue *IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
465                            JsonbParseState **state);
466 static JsonbValue *setPath(JsonbIterator **it, Datum *path_elems,
467                 bool *path_nulls, int path_len,
468                 JsonbParseState **st, int level, Jsonb *newval,
469                 int op_type);
470 static void setPathObject(JsonbIterator **it, Datum *path_elems,
471                           bool *path_nulls, int path_len, JsonbParseState **st,
472                           int level,
473                           Jsonb *newval, uint32 npairs, int op_type);
474 static void setPathArray(JsonbIterator **it, Datum *path_elems,
475                          bool *path_nulls, int path_len, JsonbParseState **st,
476                          int level, Jsonb *newval, uint32 nelems, int op_type);
477 static void addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb);
478
479 /* function supporting iterate_json_values */
480 static void iterate_values_scalar(void *state, char *token, JsonTokenType tokentype);
481 static void iterate_values_object_field_start(void *state, char *fname, bool isnull);
482
483 /* functions supporting transform_json_string_values */
484 static void transform_string_values_object_start(void *state);
485 static void transform_string_values_object_end(void *state);
486 static void transform_string_values_array_start(void *state);
487 static void transform_string_values_array_end(void *state);
488 static void transform_string_values_object_field_start(void *state, char *fname, bool isnull);
489 static void transform_string_values_array_element_start(void *state, bool isnull);
490 static void transform_string_values_scalar(void *state, char *token, JsonTokenType tokentype);
491
492 /*
493  * SQL function json_object_keys
494  *
495  * Returns the set of keys for the object argument.
496  *
497  * This SRF operates in value-per-call mode. It processes the
498  * object during the first call, and the keys are simply stashed
499  * in an array, whose size is expanded as necessary. This is probably
500  * safe enough for a list of keys of a single object, since they are
501  * limited in size to NAMEDATALEN and the number of keys is unlikely to
502  * be so huge that it has major memory implications.
503  */
504 Datum
505 jsonb_object_keys(PG_FUNCTION_ARGS)
506 {
507         FuncCallContext *funcctx;
508         OkeysState *state;
509         int                     i;
510
511         if (SRF_IS_FIRSTCALL())
512         {
513                 MemoryContext oldcontext;
514                 Jsonb      *jb = PG_GETARG_JSONB_P(0);
515                 bool            skipNested = false;
516                 JsonbIterator *it;
517                 JsonbValue      v;
518                 JsonbIteratorToken r;
519
520                 if (JB_ROOT_IS_SCALAR(jb))
521                         ereport(ERROR,
522                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
523                                          errmsg("cannot call %s on a scalar",
524                                                         "jsonb_object_keys")));
525                 else if (JB_ROOT_IS_ARRAY(jb))
526                         ereport(ERROR,
527                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
528                                          errmsg("cannot call %s on an array",
529                                                         "jsonb_object_keys")));
530
531                 funcctx = SRF_FIRSTCALL_INIT();
532                 oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
533
534                 state = palloc(sizeof(OkeysState));
535
536                 state->result_size = JB_ROOT_COUNT(jb);
537                 state->result_count = 0;
538                 state->sent_count = 0;
539                 state->result = palloc(state->result_size * sizeof(char *));
540
541                 it = JsonbIteratorInit(&jb->root);
542
543                 while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
544                 {
545                         skipNested = true;
546
547                         if (r == WJB_KEY)
548                         {
549                                 char       *cstr;
550
551                                 cstr = palloc(v.val.string.len + 1 * sizeof(char));
552                                 memcpy(cstr, v.val.string.val, v.val.string.len);
553                                 cstr[v.val.string.len] = '\0';
554                                 state->result[state->result_count++] = cstr;
555                         }
556                 }
557
558                 MemoryContextSwitchTo(oldcontext);
559                 funcctx->user_fctx = (void *) state;
560         }
561
562         funcctx = SRF_PERCALL_SETUP();
563         state = (OkeysState *) funcctx->user_fctx;
564
565         if (state->sent_count < state->result_count)
566         {
567                 char       *nxt = state->result[state->sent_count++];
568
569                 SRF_RETURN_NEXT(funcctx, CStringGetTextDatum(nxt));
570         }
571
572         /* cleanup to reduce or eliminate memory leaks */
573         for (i = 0; i < state->result_count; i++)
574                 pfree(state->result[i]);
575         pfree(state->result);
576         pfree(state);
577
578         SRF_RETURN_DONE(funcctx);
579 }
580
581
582 Datum
583 json_object_keys(PG_FUNCTION_ARGS)
584 {
585         FuncCallContext *funcctx;
586         OkeysState *state;
587         int                     i;
588
589         if (SRF_IS_FIRSTCALL())
590         {
591                 text       *json = PG_GETARG_TEXT_PP(0);
592                 JsonLexContext *lex = makeJsonLexContext(json, true);
593                 JsonSemAction *sem;
594                 MemoryContext oldcontext;
595
596                 funcctx = SRF_FIRSTCALL_INIT();
597                 oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
598
599                 state = palloc(sizeof(OkeysState));
600                 sem = palloc0(sizeof(JsonSemAction));
601
602                 state->lex = lex;
603                 state->result_size = 256;
604                 state->result_count = 0;
605                 state->sent_count = 0;
606                 state->result = palloc(256 * sizeof(char *));
607
608                 sem->semstate = (void *) state;
609                 sem->array_start = okeys_array_start;
610                 sem->scalar = okeys_scalar;
611                 sem->object_field_start = okeys_object_field_start;
612                 /* remainder are all NULL, courtesy of palloc0 above */
613
614                 pg_parse_json(lex, sem);
615                 /* keys are now in state->result */
616
617                 pfree(lex->strval->data);
618                 pfree(lex->strval);
619                 pfree(lex);
620                 pfree(sem);
621
622                 MemoryContextSwitchTo(oldcontext);
623                 funcctx->user_fctx = (void *) state;
624         }
625
626         funcctx = SRF_PERCALL_SETUP();
627         state = (OkeysState *) funcctx->user_fctx;
628
629         if (state->sent_count < state->result_count)
630         {
631                 char       *nxt = state->result[state->sent_count++];
632
633                 SRF_RETURN_NEXT(funcctx, CStringGetTextDatum(nxt));
634         }
635
636         /* cleanup to reduce or eliminate memory leaks */
637         for (i = 0; i < state->result_count; i++)
638                 pfree(state->result[i]);
639         pfree(state->result);
640         pfree(state);
641
642         SRF_RETURN_DONE(funcctx);
643 }
644
645 static void
646 okeys_object_field_start(void *state, char *fname, bool isnull)
647 {
648         OkeysState *_state = (OkeysState *) state;
649
650         /* only collecting keys for the top level object */
651         if (_state->lex->lex_level != 1)
652                 return;
653
654         /* enlarge result array if necessary */
655         if (_state->result_count >= _state->result_size)
656         {
657                 _state->result_size *= 2;
658                 _state->result = (char **)
659                         repalloc(_state->result, sizeof(char *) * _state->result_size);
660         }
661
662         /* save a copy of the field name */
663         _state->result[_state->result_count++] = pstrdup(fname);
664 }
665
666 static void
667 okeys_array_start(void *state)
668 {
669         OkeysState *_state = (OkeysState *) state;
670
671         /* top level must be a json object */
672         if (_state->lex->lex_level == 0)
673                 ereport(ERROR,
674                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
675                                  errmsg("cannot call %s on an array",
676                                                 "json_object_keys")));
677 }
678
679 static void
680 okeys_scalar(void *state, char *token, JsonTokenType tokentype)
681 {
682         OkeysState *_state = (OkeysState *) state;
683
684         /* top level must be a json object */
685         if (_state->lex->lex_level == 0)
686                 ereport(ERROR,
687                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
688                                  errmsg("cannot call %s on a scalar",
689                                                 "json_object_keys")));
690 }
691
692 /*
693  * json and jsonb getter functions
694  * these implement the -> ->> #> and #>> operators
695  * and the json{b?}_extract_path*(json, text, ...) functions
696  */
697
698
699 Datum
700 json_object_field(PG_FUNCTION_ARGS)
701 {
702         text       *json = PG_GETARG_TEXT_PP(0);
703         text       *fname = PG_GETARG_TEXT_PP(1);
704         char       *fnamestr = text_to_cstring(fname);
705         text       *result;
706
707         result = get_worker(json, &fnamestr, NULL, 1, false);
708
709         if (result != NULL)
710                 PG_RETURN_TEXT_P(result);
711         else
712                 PG_RETURN_NULL();
713 }
714
715 Datum
716 jsonb_object_field(PG_FUNCTION_ARGS)
717 {
718         Jsonb      *jb = PG_GETARG_JSONB_P(0);
719         text       *key = PG_GETARG_TEXT_PP(1);
720         JsonbValue *v;
721
722         if (!JB_ROOT_IS_OBJECT(jb))
723                 PG_RETURN_NULL();
724
725         v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
726                                                                            VARDATA_ANY(key),
727                                                                            VARSIZE_ANY_EXHDR(key));
728
729         if (v != NULL)
730                 PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
731
732         PG_RETURN_NULL();
733 }
734
735 Datum
736 json_object_field_text(PG_FUNCTION_ARGS)
737 {
738         text       *json = PG_GETARG_TEXT_PP(0);
739         text       *fname = PG_GETARG_TEXT_PP(1);
740         char       *fnamestr = text_to_cstring(fname);
741         text       *result;
742
743         result = get_worker(json, &fnamestr, NULL, 1, true);
744
745         if (result != NULL)
746                 PG_RETURN_TEXT_P(result);
747         else
748                 PG_RETURN_NULL();
749 }
750
751 Datum
752 jsonb_object_field_text(PG_FUNCTION_ARGS)
753 {
754         Jsonb      *jb = PG_GETARG_JSONB_P(0);
755         text       *key = PG_GETARG_TEXT_PP(1);
756         JsonbValue *v;
757
758         if (!JB_ROOT_IS_OBJECT(jb))
759                 PG_RETURN_NULL();
760
761         v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
762                                                                            VARDATA_ANY(key),
763                                                                            VARSIZE_ANY_EXHDR(key));
764
765         if (v != NULL)
766         {
767                 text       *result = NULL;
768
769                 switch (v->type)
770                 {
771                         case jbvNull:
772                                 break;
773                         case jbvBool:
774                                 result = cstring_to_text(v->val.boolean ? "true" : "false");
775                                 break;
776                         case jbvString:
777                                 result = cstring_to_text_with_len(v->val.string.val, v->val.string.len);
778                                 break;
779                         case jbvNumeric:
780                                 result = cstring_to_text(DatumGetCString(DirectFunctionCall1(numeric_out,
781                                                                                                                                                          PointerGetDatum(v->val.numeric))));
782                                 break;
783                         case jbvBinary:
784                                 {
785                                         StringInfo      jtext = makeStringInfo();
786
787                                         (void) JsonbToCString(jtext, v->val.binary.data, -1);
788                                         result = cstring_to_text_with_len(jtext->data, jtext->len);
789                                 }
790                                 break;
791                         default:
792                                 elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
793                 }
794
795                 if (result)
796                         PG_RETURN_TEXT_P(result);
797         }
798
799         PG_RETURN_NULL();
800 }
801
802 Datum
803 json_array_element(PG_FUNCTION_ARGS)
804 {
805         text       *json = PG_GETARG_TEXT_PP(0);
806         int                     element = PG_GETARG_INT32(1);
807         text       *result;
808
809         result = get_worker(json, NULL, &element, 1, false);
810
811         if (result != NULL)
812                 PG_RETURN_TEXT_P(result);
813         else
814                 PG_RETURN_NULL();
815 }
816
817 Datum
818 jsonb_array_element(PG_FUNCTION_ARGS)
819 {
820         Jsonb      *jb = PG_GETARG_JSONB_P(0);
821         int                     element = PG_GETARG_INT32(1);
822         JsonbValue *v;
823
824         if (!JB_ROOT_IS_ARRAY(jb))
825                 PG_RETURN_NULL();
826
827         /* Handle negative subscript */
828         if (element < 0)
829         {
830                 uint32          nelements = JB_ROOT_COUNT(jb);
831
832                 if (-element > nelements)
833                         PG_RETURN_NULL();
834                 else
835                         element += nelements;
836         }
837
838         v = getIthJsonbValueFromContainer(&jb->root, element);
839         if (v != NULL)
840                 PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
841
842         PG_RETURN_NULL();
843 }
844
845 Datum
846 json_array_element_text(PG_FUNCTION_ARGS)
847 {
848         text       *json = PG_GETARG_TEXT_PP(0);
849         int                     element = PG_GETARG_INT32(1);
850         text       *result;
851
852         result = get_worker(json, NULL, &element, 1, true);
853
854         if (result != NULL)
855                 PG_RETURN_TEXT_P(result);
856         else
857                 PG_RETURN_NULL();
858 }
859
860 Datum
861 jsonb_array_element_text(PG_FUNCTION_ARGS)
862 {
863         Jsonb      *jb = PG_GETARG_JSONB_P(0);
864         int                     element = PG_GETARG_INT32(1);
865         JsonbValue *v;
866
867         if (!JB_ROOT_IS_ARRAY(jb))
868                 PG_RETURN_NULL();
869
870         /* Handle negative subscript */
871         if (element < 0)
872         {
873                 uint32          nelements = JB_ROOT_COUNT(jb);
874
875                 if (-element > nelements)
876                         PG_RETURN_NULL();
877                 else
878                         element += nelements;
879         }
880
881         v = getIthJsonbValueFromContainer(&jb->root, element);
882         if (v != NULL)
883         {
884                 text       *result = NULL;
885
886                 switch (v->type)
887                 {
888                         case jbvNull:
889                                 break;
890                         case jbvBool:
891                                 result = cstring_to_text(v->val.boolean ? "true" : "false");
892                                 break;
893                         case jbvString:
894                                 result = cstring_to_text_with_len(v->val.string.val, v->val.string.len);
895                                 break;
896                         case jbvNumeric:
897                                 result = cstring_to_text(DatumGetCString(DirectFunctionCall1(numeric_out,
898                                                                                                                                                          PointerGetDatum(v->val.numeric))));
899                                 break;
900                         case jbvBinary:
901                                 {
902                                         StringInfo      jtext = makeStringInfo();
903
904                                         (void) JsonbToCString(jtext, v->val.binary.data, -1);
905                                         result = cstring_to_text_with_len(jtext->data, jtext->len);
906                                 }
907                                 break;
908                         default:
909                                 elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
910                 }
911
912                 if (result)
913                         PG_RETURN_TEXT_P(result);
914         }
915
916         PG_RETURN_NULL();
917 }
918
919 Datum
920 json_extract_path(PG_FUNCTION_ARGS)
921 {
922         return get_path_all(fcinfo, false);
923 }
924
925 Datum
926 json_extract_path_text(PG_FUNCTION_ARGS)
927 {
928         return get_path_all(fcinfo, true);
929 }
930
931 /*
932  * common routine for extract_path functions
933  */
934 static Datum
935 get_path_all(FunctionCallInfo fcinfo, bool as_text)
936 {
937         text       *json = PG_GETARG_TEXT_PP(0);
938         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
939         text       *result;
940         Datum      *pathtext;
941         bool       *pathnulls;
942         int                     npath;
943         char      **tpath;
944         int                *ipath;
945         int                     i;
946
947         /*
948          * If the array contains any null elements, return NULL, on the grounds
949          * that you'd have gotten NULL if any RHS value were NULL in a nested
950          * series of applications of the -> operator.  (Note: because we also
951          * return NULL for error cases such as no-such-field, this is true
952          * regardless of the contents of the rest of the array.)
953          */
954         if (array_contains_nulls(path))
955                 PG_RETURN_NULL();
956
957         deconstruct_array(path, TEXTOID, -1, false, 'i',
958                                           &pathtext, &pathnulls, &npath);
959
960         tpath = palloc(npath * sizeof(char *));
961         ipath = palloc(npath * sizeof(int));
962
963         for (i = 0; i < npath; i++)
964         {
965                 Assert(!pathnulls[i]);
966                 tpath[i] = TextDatumGetCString(pathtext[i]);
967
968                 /*
969                  * we have no idea at this stage what structure the document is so
970                  * just convert anything in the path that we can to an integer and set
971                  * all the other integers to INT_MIN which will never match.
972                  */
973                 if (*tpath[i] != '\0')
974                 {
975                         long            ind;
976                         char       *endptr;
977
978                         errno = 0;
979                         ind = strtol(tpath[i], &endptr, 10);
980                         if (*endptr == '\0' && errno == 0 && ind <= INT_MAX && ind >= INT_MIN)
981                                 ipath[i] = (int) ind;
982                         else
983                                 ipath[i] = INT_MIN;
984                 }
985                 else
986                         ipath[i] = INT_MIN;
987         }
988
989         result = get_worker(json, tpath, ipath, npath, as_text);
990
991         if (result != NULL)
992                 PG_RETURN_TEXT_P(result);
993         else
994                 PG_RETURN_NULL();
995 }
996
997 /*
998  * get_worker
999  *
1000  * common worker for all the json getter functions
1001  *
1002  * json: JSON object (in text form)
1003  * tpath[]: field name(s) to extract
1004  * ipath[]: array index(es) (zero-based) to extract, accepts negatives
1005  * npath: length of tpath[] and/or ipath[]
1006  * normalize_results: true to de-escape string and null scalars
1007  *
1008  * tpath can be NULL, or any one tpath[] entry can be NULL, if an object
1009  * field is not to be matched at that nesting level.  Similarly, ipath can
1010  * be NULL, or any one ipath[] entry can be INT_MIN if an array element is
1011  * not to be matched at that nesting level (a json datum should never be
1012  * large enough to have -INT_MIN elements due to MaxAllocSize restriction).
1013  */
1014 static text *
1015 get_worker(text *json,
1016                    char **tpath,
1017                    int *ipath,
1018                    int npath,
1019                    bool normalize_results)
1020 {
1021         JsonLexContext *lex = makeJsonLexContext(json, true);
1022         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
1023         GetState   *state = palloc0(sizeof(GetState));
1024
1025         Assert(npath >= 0);
1026
1027         state->lex = lex;
1028         /* is it "_as_text" variant? */
1029         state->normalize_results = normalize_results;
1030         state->npath = npath;
1031         state->path_names = tpath;
1032         state->path_indexes = ipath;
1033         state->pathok = palloc0(sizeof(bool) * npath);
1034         state->array_cur_index = palloc(sizeof(int) * npath);
1035
1036         if (npath > 0)
1037                 state->pathok[0] = true;
1038
1039         sem->semstate = (void *) state;
1040
1041         /*
1042          * Not all variants need all the semantic routines. Only set the ones that
1043          * are actually needed for maximum efficiency.
1044          */
1045         sem->scalar = get_scalar;
1046         if (npath == 0)
1047         {
1048                 sem->object_start = get_object_start;
1049                 sem->object_end = get_object_end;
1050                 sem->array_start = get_array_start;
1051                 sem->array_end = get_array_end;
1052         }
1053         if (tpath != NULL)
1054         {
1055                 sem->object_field_start = get_object_field_start;
1056                 sem->object_field_end = get_object_field_end;
1057         }
1058         if (ipath != NULL)
1059         {
1060                 sem->array_start = get_array_start;
1061                 sem->array_element_start = get_array_element_start;
1062                 sem->array_element_end = get_array_element_end;
1063         }
1064
1065         pg_parse_json(lex, sem);
1066
1067         return state->tresult;
1068 }
1069
1070 static void
1071 get_object_start(void *state)
1072 {
1073         GetState   *_state = (GetState *) state;
1074         int                     lex_level = _state->lex->lex_level;
1075
1076         if (lex_level == 0 && _state->npath == 0)
1077         {
1078                 /*
1079                  * Special case: we should match the entire object.  We only need this
1080                  * at outermost level because at nested levels the match will have
1081                  * been started by the outer field or array element callback.
1082                  */
1083                 _state->result_start = _state->lex->token_start;
1084         }
1085 }
1086
1087 static void
1088 get_object_end(void *state)
1089 {
1090         GetState   *_state = (GetState *) state;
1091         int                     lex_level = _state->lex->lex_level;
1092
1093         if (lex_level == 0 && _state->npath == 0)
1094         {
1095                 /* Special case: return the entire object */
1096                 char       *start = _state->result_start;
1097                 int                     len = _state->lex->prev_token_terminator - start;
1098
1099                 _state->tresult = cstring_to_text_with_len(start, len);
1100         }
1101 }
1102
1103 static void
1104 get_object_field_start(void *state, char *fname, bool isnull)
1105 {
1106         GetState   *_state = (GetState *) state;
1107         bool            get_next = false;
1108         int                     lex_level = _state->lex->lex_level;
1109
1110         if (lex_level <= _state->npath &&
1111                 _state->pathok[lex_level - 1] &&
1112                 _state->path_names != NULL &&
1113                 _state->path_names[lex_level - 1] != NULL &&
1114                 strcmp(fname, _state->path_names[lex_level - 1]) == 0)
1115         {
1116                 if (lex_level < _state->npath)
1117                 {
1118                         /* if not at end of path just mark path ok */
1119                         _state->pathok[lex_level] = true;
1120                 }
1121                 else
1122                 {
1123                         /* end of path, so we want this value */
1124                         get_next = true;
1125                 }
1126         }
1127
1128         if (get_next)
1129         {
1130                 /* this object overrides any previous matching object */
1131                 _state->tresult = NULL;
1132                 _state->result_start = NULL;
1133
1134                 if (_state->normalize_results &&
1135                         _state->lex->token_type == JSON_TOKEN_STRING)
1136                 {
1137                         /* for as_text variants, tell get_scalar to set it for us */
1138                         _state->next_scalar = true;
1139                 }
1140                 else
1141                 {
1142                         /* for non-as_text variants, just note the json starting point */
1143                         _state->result_start = _state->lex->token_start;
1144                 }
1145         }
1146 }
1147
1148 static void
1149 get_object_field_end(void *state, char *fname, bool isnull)
1150 {
1151         GetState   *_state = (GetState *) state;
1152         bool            get_last = false;
1153         int                     lex_level = _state->lex->lex_level;
1154
1155         /* same tests as in get_object_field_start */
1156         if (lex_level <= _state->npath &&
1157                 _state->pathok[lex_level - 1] &&
1158                 _state->path_names != NULL &&
1159                 _state->path_names[lex_level - 1] != NULL &&
1160                 strcmp(fname, _state->path_names[lex_level - 1]) == 0)
1161         {
1162                 if (lex_level < _state->npath)
1163                 {
1164                         /* done with this field so reset pathok */
1165                         _state->pathok[lex_level] = false;
1166                 }
1167                 else
1168                 {
1169                         /* end of path, so we want this value */
1170                         get_last = true;
1171                 }
1172         }
1173
1174         /* for as_text scalar case, our work is already done */
1175         if (get_last && _state->result_start != NULL)
1176         {
1177                 /*
1178                  * make a text object from the string from the previously noted json
1179                  * start up to the end of the previous token (the lexer is by now
1180                  * ahead of us on whatever came after what we're interested in).
1181                  */
1182                 if (isnull && _state->normalize_results)
1183                         _state->tresult = (text *) NULL;
1184                 else
1185                 {
1186                         char       *start = _state->result_start;
1187                         int                     len = _state->lex->prev_token_terminator - start;
1188
1189                         _state->tresult = cstring_to_text_with_len(start, len);
1190                 }
1191
1192                 /* this should be unnecessary but let's do it for cleanliness: */
1193                 _state->result_start = NULL;
1194         }
1195 }
1196
1197 static void
1198 get_array_start(void *state)
1199 {
1200         GetState   *_state = (GetState *) state;
1201         int                     lex_level = _state->lex->lex_level;
1202
1203         if (lex_level < _state->npath)
1204         {
1205                 /* Initialize counting of elements in this array */
1206                 _state->array_cur_index[lex_level] = -1;
1207
1208                 /* INT_MIN value is reserved to represent invalid subscript */
1209                 if (_state->path_indexes[lex_level] < 0 &&
1210                         _state->path_indexes[lex_level] != INT_MIN)
1211                 {
1212                         /* Negative subscript -- convert to positive-wise subscript */
1213                         int                     nelements = json_count_array_elements(_state->lex);
1214
1215                         if (-_state->path_indexes[lex_level] <= nelements)
1216                                 _state->path_indexes[lex_level] += nelements;
1217                 }
1218         }
1219         else if (lex_level == 0 && _state->npath == 0)
1220         {
1221                 /*
1222                  * Special case: we should match the entire array.  We only need this
1223                  * at the outermost level because at nested levels the match will have
1224                  * been started by the outer field or array element callback.
1225                  */
1226                 _state->result_start = _state->lex->token_start;
1227         }
1228 }
1229
1230 static void
1231 get_array_end(void *state)
1232 {
1233         GetState   *_state = (GetState *) state;
1234         int                     lex_level = _state->lex->lex_level;
1235
1236         if (lex_level == 0 && _state->npath == 0)
1237         {
1238                 /* Special case: return the entire array */
1239                 char       *start = _state->result_start;
1240                 int                     len = _state->lex->prev_token_terminator - start;
1241
1242                 _state->tresult = cstring_to_text_with_len(start, len);
1243         }
1244 }
1245
1246 static void
1247 get_array_element_start(void *state, bool isnull)
1248 {
1249         GetState   *_state = (GetState *) state;
1250         bool            get_next = false;
1251         int                     lex_level = _state->lex->lex_level;
1252
1253         /* Update array element counter */
1254         if (lex_level <= _state->npath)
1255                 _state->array_cur_index[lex_level - 1]++;
1256
1257         if (lex_level <= _state->npath &&
1258                 _state->pathok[lex_level - 1] &&
1259                 _state->path_indexes != NULL &&
1260                 _state->array_cur_index[lex_level - 1] == _state->path_indexes[lex_level - 1])
1261         {
1262                 if (lex_level < _state->npath)
1263                 {
1264                         /* if not at end of path just mark path ok */
1265                         _state->pathok[lex_level] = true;
1266                 }
1267                 else
1268                 {
1269                         /* end of path, so we want this value */
1270                         get_next = true;
1271                 }
1272         }
1273
1274         /* same logic as for objects */
1275         if (get_next)
1276         {
1277                 _state->tresult = NULL;
1278                 _state->result_start = NULL;
1279
1280                 if (_state->normalize_results &&
1281                         _state->lex->token_type == JSON_TOKEN_STRING)
1282                 {
1283                         _state->next_scalar = true;
1284                 }
1285                 else
1286                 {
1287                         _state->result_start = _state->lex->token_start;
1288                 }
1289         }
1290 }
1291
1292 static void
1293 get_array_element_end(void *state, bool isnull)
1294 {
1295         GetState   *_state = (GetState *) state;
1296         bool            get_last = false;
1297         int                     lex_level = _state->lex->lex_level;
1298
1299         /* same tests as in get_array_element_start */
1300         if (lex_level <= _state->npath &&
1301                 _state->pathok[lex_level - 1] &&
1302                 _state->path_indexes != NULL &&
1303                 _state->array_cur_index[lex_level - 1] == _state->path_indexes[lex_level - 1])
1304         {
1305                 if (lex_level < _state->npath)
1306                 {
1307                         /* done with this element so reset pathok */
1308                         _state->pathok[lex_level] = false;
1309                 }
1310                 else
1311                 {
1312                         /* end of path, so we want this value */
1313                         get_last = true;
1314                 }
1315         }
1316
1317         /* same logic as for objects */
1318         if (get_last && _state->result_start != NULL)
1319         {
1320                 if (isnull && _state->normalize_results)
1321                         _state->tresult = (text *) NULL;
1322                 else
1323                 {
1324                         char       *start = _state->result_start;
1325                         int                     len = _state->lex->prev_token_terminator - start;
1326
1327                         _state->tresult = cstring_to_text_with_len(start, len);
1328                 }
1329
1330                 _state->result_start = NULL;
1331         }
1332 }
1333
1334 static void
1335 get_scalar(void *state, char *token, JsonTokenType tokentype)
1336 {
1337         GetState   *_state = (GetState *) state;
1338         int                     lex_level = _state->lex->lex_level;
1339
1340         /* Check for whole-object match */
1341         if (lex_level == 0 && _state->npath == 0)
1342         {
1343                 if (_state->normalize_results && tokentype == JSON_TOKEN_STRING)
1344                 {
1345                         /* we want the de-escaped string */
1346                         _state->next_scalar = true;
1347                 }
1348                 else if (_state->normalize_results && tokentype == JSON_TOKEN_NULL)
1349                 {
1350                         _state->tresult = (text *) NULL;
1351                 }
1352                 else
1353                 {
1354                         /*
1355                          * This is a bit hokey: we will suppress whitespace after the
1356                          * scalar token, but not whitespace before it.  Probably not worth
1357                          * doing our own space-skipping to avoid that.
1358                          */
1359                         char       *start = _state->lex->input;
1360                         int                     len = _state->lex->prev_token_terminator - start;
1361
1362                         _state->tresult = cstring_to_text_with_len(start, len);
1363                 }
1364         }
1365
1366         if (_state->next_scalar)
1367         {
1368                 /* a de-escaped text value is wanted, so supply it */
1369                 _state->tresult = cstring_to_text(token);
1370                 /* make sure the next call to get_scalar doesn't overwrite it */
1371                 _state->next_scalar = false;
1372         }
1373 }
1374
1375 Datum
1376 jsonb_extract_path(PG_FUNCTION_ARGS)
1377 {
1378         return get_jsonb_path_all(fcinfo, false);
1379 }
1380
1381 Datum
1382 jsonb_extract_path_text(PG_FUNCTION_ARGS)
1383 {
1384         return get_jsonb_path_all(fcinfo, true);
1385 }
1386
1387 static Datum
1388 get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
1389 {
1390         Jsonb      *jb = PG_GETARG_JSONB_P(0);
1391         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
1392         Jsonb      *res;
1393         Datum      *pathtext;
1394         bool       *pathnulls;
1395         int                     npath;
1396         int                     i;
1397         bool            have_object = false,
1398                                 have_array = false;
1399         JsonbValue *jbvp = NULL;
1400         JsonbValue      tv;
1401         JsonbContainer *container;
1402
1403         /*
1404          * If the array contains any null elements, return NULL, on the grounds
1405          * that you'd have gotten NULL if any RHS value were NULL in a nested
1406          * series of applications of the -> operator.  (Note: because we also
1407          * return NULL for error cases such as no-such-field, this is true
1408          * regardless of the contents of the rest of the array.)
1409          */
1410         if (array_contains_nulls(path))
1411                 PG_RETURN_NULL();
1412
1413         deconstruct_array(path, TEXTOID, -1, false, 'i',
1414                                           &pathtext, &pathnulls, &npath);
1415
1416         /* Identify whether we have object, array, or scalar at top-level */
1417         container = &jb->root;
1418
1419         if (JB_ROOT_IS_OBJECT(jb))
1420                 have_object = true;
1421         else if (JB_ROOT_IS_ARRAY(jb) && !JB_ROOT_IS_SCALAR(jb))
1422                 have_array = true;
1423         else
1424         {
1425                 Assert(JB_ROOT_IS_ARRAY(jb) && JB_ROOT_IS_SCALAR(jb));
1426                 /* Extract the scalar value, if it is what we'll return */
1427                 if (npath <= 0)
1428                         jbvp = getIthJsonbValueFromContainer(container, 0);
1429         }
1430
1431         /*
1432          * If the array is empty, return the entire LHS object, on the grounds
1433          * that we should do zero field or element extractions.  For the
1434          * non-scalar case we can just hand back the object without much work. For
1435          * the scalar case, fall through and deal with the value below the loop.
1436          * (This inconsistency arises because there's no easy way to generate a
1437          * JsonbValue directly for root-level containers.)
1438          */
1439         if (npath <= 0 && jbvp == NULL)
1440         {
1441                 if (as_text)
1442                 {
1443                         PG_RETURN_TEXT_P(cstring_to_text(JsonbToCString(NULL,
1444                                                                                                                         container,
1445                                                                                                                         VARSIZE(jb))));
1446                 }
1447                 else
1448                 {
1449                         /* not text mode - just hand back the jsonb */
1450                         PG_RETURN_JSONB_P(jb);
1451                 }
1452         }
1453
1454         for (i = 0; i < npath; i++)
1455         {
1456                 if (have_object)
1457                 {
1458                         jbvp = findJsonbValueFromContainerLen(container,
1459                                                                                                   JB_FOBJECT,
1460                                                                                                   VARDATA(pathtext[i]),
1461                                                                                                   VARSIZE(pathtext[i]) - VARHDRSZ);
1462                 }
1463                 else if (have_array)
1464                 {
1465                         long            lindex;
1466                         uint32          index;
1467                         char       *indextext = TextDatumGetCString(pathtext[i]);
1468                         char       *endptr;
1469
1470                         errno = 0;
1471                         lindex = strtol(indextext, &endptr, 10);
1472                         if (endptr == indextext || *endptr != '\0' || errno != 0 ||
1473                                 lindex > INT_MAX || lindex < INT_MIN)
1474                                 PG_RETURN_NULL();
1475
1476                         if (lindex >= 0)
1477                         {
1478                                 index = (uint32) lindex;
1479                         }
1480                         else
1481                         {
1482                                 /* Handle negative subscript */
1483                                 uint32          nelements;
1484
1485                                 /* Container must be array, but make sure */
1486                                 if (!JsonContainerIsArray(container))
1487                                         elog(ERROR, "not a jsonb array");
1488
1489                                 nelements = JsonContainerSize(container);
1490
1491                                 if (-lindex > nelements)
1492                                         PG_RETURN_NULL();
1493                                 else
1494                                         index = nelements + lindex;
1495                         }
1496
1497                         jbvp = getIthJsonbValueFromContainer(container, index);
1498                 }
1499                 else
1500                 {
1501                         /* scalar, extraction yields a null */
1502                         PG_RETURN_NULL();
1503                 }
1504
1505                 if (jbvp == NULL)
1506                         PG_RETURN_NULL();
1507                 else if (i == npath - 1)
1508                         break;
1509
1510                 if (jbvp->type == jbvBinary)
1511                 {
1512                         JsonbIterator *it = JsonbIteratorInit((JsonbContainer *) jbvp->val.binary.data);
1513                         JsonbIteratorToken r;
1514
1515                         r = JsonbIteratorNext(&it, &tv, true);
1516                         container = (JsonbContainer *) jbvp->val.binary.data;
1517                         have_object = r == WJB_BEGIN_OBJECT;
1518                         have_array = r == WJB_BEGIN_ARRAY;
1519                 }
1520                 else
1521                 {
1522                         have_object = jbvp->type == jbvObject;
1523                         have_array = jbvp->type == jbvArray;
1524                 }
1525         }
1526
1527         if (as_text)
1528         {
1529                 /* special-case outputs for string and null values */
1530                 if (jbvp->type == jbvString)
1531                         PG_RETURN_TEXT_P(cstring_to_text_with_len(jbvp->val.string.val,
1532                                                                                                           jbvp->val.string.len));
1533                 if (jbvp->type == jbvNull)
1534                         PG_RETURN_NULL();
1535         }
1536
1537         res = JsonbValueToJsonb(jbvp);
1538
1539         if (as_text)
1540         {
1541                 PG_RETURN_TEXT_P(cstring_to_text(JsonbToCString(NULL,
1542                                                                                                                 &res->root,
1543                                                                                                                 VARSIZE(res))));
1544         }
1545         else
1546         {
1547                 /* not text mode - just hand back the jsonb */
1548                 PG_RETURN_JSONB_P(res);
1549         }
1550 }
1551
1552 /*
1553  * SQL function json_array_length(json) -> int
1554  */
1555 Datum
1556 json_array_length(PG_FUNCTION_ARGS)
1557 {
1558         text       *json = PG_GETARG_TEXT_PP(0);
1559         AlenState  *state;
1560         JsonLexContext *lex;
1561         JsonSemAction *sem;
1562
1563         lex = makeJsonLexContext(json, false);
1564         state = palloc0(sizeof(AlenState));
1565         sem = palloc0(sizeof(JsonSemAction));
1566
1567         /* palloc0 does this for us */
1568 #if 0
1569         state->count = 0;
1570 #endif
1571         state->lex = lex;
1572
1573         sem->semstate = (void *) state;
1574         sem->object_start = alen_object_start;
1575         sem->scalar = alen_scalar;
1576         sem->array_element_start = alen_array_element_start;
1577
1578         pg_parse_json(lex, sem);
1579
1580         PG_RETURN_INT32(state->count);
1581 }
1582
1583 Datum
1584 jsonb_array_length(PG_FUNCTION_ARGS)
1585 {
1586         Jsonb      *jb = PG_GETARG_JSONB_P(0);
1587
1588         if (JB_ROOT_IS_SCALAR(jb))
1589                 ereport(ERROR,
1590                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1591                                  errmsg("cannot get array length of a scalar")));
1592         else if (!JB_ROOT_IS_ARRAY(jb))
1593                 ereport(ERROR,
1594                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1595                                  errmsg("cannot get array length of a non-array")));
1596
1597         PG_RETURN_INT32(JB_ROOT_COUNT(jb));
1598 }
1599
1600 /*
1601  * These next two checks ensure that the json is an array (since it can't be
1602  * a scalar or an object).
1603  */
1604
1605 static void
1606 alen_object_start(void *state)
1607 {
1608         AlenState  *_state = (AlenState *) state;
1609
1610         /* json structure check */
1611         if (_state->lex->lex_level == 0)
1612                 ereport(ERROR,
1613                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1614                                  errmsg("cannot get array length of a non-array")));
1615 }
1616
1617 static void
1618 alen_scalar(void *state, char *token, JsonTokenType tokentype)
1619 {
1620         AlenState  *_state = (AlenState *) state;
1621
1622         /* json structure check */
1623         if (_state->lex->lex_level == 0)
1624                 ereport(ERROR,
1625                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1626                                  errmsg("cannot get array length of a scalar")));
1627 }
1628
1629 static void
1630 alen_array_element_start(void *state, bool isnull)
1631 {
1632         AlenState  *_state = (AlenState *) state;
1633
1634         /* just count up all the level 1 elements */
1635         if (_state->lex->lex_level == 1)
1636                 _state->count++;
1637 }
1638
1639 /*
1640  * SQL function json_each and json_each_text
1641  *
1642  * decompose a json object into key value pairs.
1643  *
1644  * Unlike json_object_keys() these SRFs operate in materialize mode,
1645  * stashing results into a Tuplestore object as they go.
1646  * The construction of tuples is done using a temporary memory context
1647  * that is cleared out after each tuple is built.
1648  */
1649 Datum
1650 json_each(PG_FUNCTION_ARGS)
1651 {
1652         return each_worker(fcinfo, false);
1653 }
1654
1655 Datum
1656 jsonb_each(PG_FUNCTION_ARGS)
1657 {
1658         return each_worker_jsonb(fcinfo, "jsonb_each", false);
1659 }
1660
1661 Datum
1662 json_each_text(PG_FUNCTION_ARGS)
1663 {
1664         return each_worker(fcinfo, true);
1665 }
1666
1667 Datum
1668 jsonb_each_text(PG_FUNCTION_ARGS)
1669 {
1670         return each_worker_jsonb(fcinfo, "jsonb_each_text", true);
1671 }
1672
1673 static Datum
1674 each_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
1675 {
1676         Jsonb      *jb = PG_GETARG_JSONB_P(0);
1677         ReturnSetInfo *rsi;
1678         Tuplestorestate *tuple_store;
1679         TupleDesc       tupdesc;
1680         TupleDesc       ret_tdesc;
1681         MemoryContext old_cxt,
1682                                 tmp_cxt;
1683         bool            skipNested = false;
1684         JsonbIterator *it;
1685         JsonbValue      v;
1686         JsonbIteratorToken r;
1687
1688         if (!JB_ROOT_IS_OBJECT(jb))
1689                 ereport(ERROR,
1690                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1691                                  errmsg("cannot call %s on a non-object",
1692                                                 funcname)));
1693
1694         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1695
1696         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1697                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1698                 rsi->expectedDesc == NULL)
1699                 ereport(ERROR,
1700                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1701                                  errmsg("set-valued function called in context that "
1702                                                 "cannot accept a set")));
1703
1704         rsi->returnMode = SFRM_Materialize;
1705
1706         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
1707                 ereport(ERROR,
1708                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1709                                  errmsg("function returning record called in context "
1710                                                 "that cannot accept type record")));
1711
1712         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1713
1714         ret_tdesc = CreateTupleDescCopy(tupdesc);
1715         BlessTupleDesc(ret_tdesc);
1716         tuple_store =
1717                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1718                                                           false, work_mem);
1719
1720         MemoryContextSwitchTo(old_cxt);
1721
1722         tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1723                                                                         "jsonb_each temporary cxt",
1724                                                                         ALLOCSET_DEFAULT_SIZES);
1725
1726         it = JsonbIteratorInit(&jb->root);
1727
1728         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
1729         {
1730                 skipNested = true;
1731
1732                 if (r == WJB_KEY)
1733                 {
1734                         text       *key;
1735                         HeapTuple       tuple;
1736                         Datum           values[2];
1737                         bool            nulls[2] = {false, false};
1738
1739                         /* Use the tmp context so we can clean up after each tuple is done */
1740                         old_cxt = MemoryContextSwitchTo(tmp_cxt);
1741
1742                         key = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
1743
1744                         /*
1745                          * The next thing the iterator fetches should be the value, no
1746                          * matter what shape it is.
1747                          */
1748                         r = JsonbIteratorNext(&it, &v, skipNested);
1749                         Assert(r != WJB_DONE);
1750
1751                         values[0] = PointerGetDatum(key);
1752
1753                         if (as_text)
1754                         {
1755                                 if (v.type == jbvNull)
1756                                 {
1757                                         /* a json null is an sql null in text mode */
1758                                         nulls[1] = true;
1759                                         values[1] = (Datum) NULL;
1760                                 }
1761                                 else
1762                                 {
1763                                         text       *sv;
1764
1765                                         if (v.type == jbvString)
1766                                         {
1767                                                 /* In text mode, scalar strings should be dequoted */
1768                                                 sv = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
1769                                         }
1770                                         else
1771                                         {
1772                                                 /* Turn anything else into a json string */
1773                                                 StringInfo      jtext = makeStringInfo();
1774                                                 Jsonb      *jb = JsonbValueToJsonb(&v);
1775
1776                                                 (void) JsonbToCString(jtext, &jb->root, 0);
1777                                                 sv = cstring_to_text_with_len(jtext->data, jtext->len);
1778                                         }
1779
1780                                         values[1] = PointerGetDatum(sv);
1781                                 }
1782                         }
1783                         else
1784                         {
1785                                 /* Not in text mode, just return the Jsonb */
1786                                 Jsonb      *val = JsonbValueToJsonb(&v);
1787
1788                                 values[1] = PointerGetDatum(val);
1789                         }
1790
1791                         tuple = heap_form_tuple(ret_tdesc, values, nulls);
1792
1793                         tuplestore_puttuple(tuple_store, tuple);
1794
1795                         /* clean up and switch back */
1796                         MemoryContextSwitchTo(old_cxt);
1797                         MemoryContextReset(tmp_cxt);
1798                 }
1799         }
1800
1801         MemoryContextDelete(tmp_cxt);
1802
1803         rsi->setResult = tuple_store;
1804         rsi->setDesc = ret_tdesc;
1805
1806         PG_RETURN_NULL();
1807 }
1808
1809
1810 static Datum
1811 each_worker(FunctionCallInfo fcinfo, bool as_text)
1812 {
1813         text       *json = PG_GETARG_TEXT_PP(0);
1814         JsonLexContext *lex;
1815         JsonSemAction *sem;
1816         ReturnSetInfo *rsi;
1817         MemoryContext old_cxt;
1818         TupleDesc       tupdesc;
1819         EachState  *state;
1820
1821         lex = makeJsonLexContext(json, true);
1822         state = palloc0(sizeof(EachState));
1823         sem = palloc0(sizeof(JsonSemAction));
1824
1825         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1826
1827         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1828                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1829                 rsi->expectedDesc == NULL)
1830                 ereport(ERROR,
1831                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1832                                  errmsg("set-valued function called in context that "
1833                                                 "cannot accept a set")));
1834
1835         rsi->returnMode = SFRM_Materialize;
1836
1837         (void) get_call_result_type(fcinfo, NULL, &tupdesc);
1838
1839         /* make these in a sufficiently long-lived memory context */
1840         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1841
1842         state->ret_tdesc = CreateTupleDescCopy(tupdesc);
1843         BlessTupleDesc(state->ret_tdesc);
1844         state->tuple_store =
1845                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1846                                                           false, work_mem);
1847
1848         MemoryContextSwitchTo(old_cxt);
1849
1850         sem->semstate = (void *) state;
1851         sem->array_start = each_array_start;
1852         sem->scalar = each_scalar;
1853         sem->object_field_start = each_object_field_start;
1854         sem->object_field_end = each_object_field_end;
1855
1856         state->normalize_results = as_text;
1857         state->next_scalar = false;
1858         state->lex = lex;
1859         state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1860                                                                                    "json_each temporary cxt",
1861                                                                                    ALLOCSET_DEFAULT_SIZES);
1862
1863         pg_parse_json(lex, sem);
1864
1865         MemoryContextDelete(state->tmp_cxt);
1866
1867         rsi->setResult = state->tuple_store;
1868         rsi->setDesc = state->ret_tdesc;
1869
1870         PG_RETURN_NULL();
1871 }
1872
1873
1874 static void
1875 each_object_field_start(void *state, char *fname, bool isnull)
1876 {
1877         EachState  *_state = (EachState *) state;
1878
1879         /* save a pointer to where the value starts */
1880         if (_state->lex->lex_level == 1)
1881         {
1882                 /*
1883                  * next_scalar will be reset in the object_field_end handler, and
1884                  * since we know the value is a scalar there is no danger of it being
1885                  * on while recursing down the tree.
1886                  */
1887                 if (_state->normalize_results && _state->lex->token_type == JSON_TOKEN_STRING)
1888                         _state->next_scalar = true;
1889                 else
1890                         _state->result_start = _state->lex->token_start;
1891         }
1892 }
1893
1894 static void
1895 each_object_field_end(void *state, char *fname, bool isnull)
1896 {
1897         EachState  *_state = (EachState *) state;
1898         MemoryContext old_cxt;
1899         int                     len;
1900         text       *val;
1901         HeapTuple       tuple;
1902         Datum           values[2];
1903         bool            nulls[2] = {false, false};
1904
1905         /* skip over nested objects */
1906         if (_state->lex->lex_level != 1)
1907                 return;
1908
1909         /* use the tmp context so we can clean up after each tuple is done */
1910         old_cxt = MemoryContextSwitchTo(_state->tmp_cxt);
1911
1912         values[0] = CStringGetTextDatum(fname);
1913
1914         if (isnull && _state->normalize_results)
1915         {
1916                 nulls[1] = true;
1917                 values[1] = (Datum) 0;
1918         }
1919         else if (_state->next_scalar)
1920         {
1921                 values[1] = CStringGetTextDatum(_state->normalized_scalar);
1922                 _state->next_scalar = false;
1923         }
1924         else
1925         {
1926                 len = _state->lex->prev_token_terminator - _state->result_start;
1927                 val = cstring_to_text_with_len(_state->result_start, len);
1928                 values[1] = PointerGetDatum(val);
1929         }
1930
1931         tuple = heap_form_tuple(_state->ret_tdesc, values, nulls);
1932
1933         tuplestore_puttuple(_state->tuple_store, tuple);
1934
1935         /* clean up and switch back */
1936         MemoryContextSwitchTo(old_cxt);
1937         MemoryContextReset(_state->tmp_cxt);
1938 }
1939
1940 static void
1941 each_array_start(void *state)
1942 {
1943         EachState  *_state = (EachState *) state;
1944
1945         /* json structure check */
1946         if (_state->lex->lex_level == 0)
1947                 ereport(ERROR,
1948                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1949                                  errmsg("cannot deconstruct an array as an object")));
1950 }
1951
1952 static void
1953 each_scalar(void *state, char *token, JsonTokenType tokentype)
1954 {
1955         EachState  *_state = (EachState *) state;
1956
1957         /* json structure check */
1958         if (_state->lex->lex_level == 0)
1959                 ereport(ERROR,
1960                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1961                                  errmsg("cannot deconstruct a scalar")));
1962
1963         /* supply de-escaped value if required */
1964         if (_state->next_scalar)
1965                 _state->normalized_scalar = token;
1966 }
1967
1968 /*
1969  * SQL functions json_array_elements and json_array_elements_text
1970  *
1971  * get the elements from a json array
1972  *
1973  * a lot of this processing is similar to the json_each* functions
1974  */
1975
1976 Datum
1977 jsonb_array_elements(PG_FUNCTION_ARGS)
1978 {
1979         return elements_worker_jsonb(fcinfo, "jsonb_array_elements", false);
1980 }
1981
1982 Datum
1983 jsonb_array_elements_text(PG_FUNCTION_ARGS)
1984 {
1985         return elements_worker_jsonb(fcinfo, "jsonb_array_elements_text", true);
1986 }
1987
1988 static Datum
1989 elements_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
1990                                           bool as_text)
1991 {
1992         Jsonb      *jb = PG_GETARG_JSONB_P(0);
1993         ReturnSetInfo *rsi;
1994         Tuplestorestate *tuple_store;
1995         TupleDesc       tupdesc;
1996         TupleDesc       ret_tdesc;
1997         MemoryContext old_cxt,
1998                                 tmp_cxt;
1999         bool            skipNested = false;
2000         JsonbIterator *it;
2001         JsonbValue      v;
2002         JsonbIteratorToken r;
2003
2004         if (JB_ROOT_IS_SCALAR(jb))
2005                 ereport(ERROR,
2006                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2007                                  errmsg("cannot extract elements from a scalar")));
2008         else if (!JB_ROOT_IS_ARRAY(jb))
2009                 ereport(ERROR,
2010                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2011                                  errmsg("cannot extract elements from an object")));
2012
2013         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
2014
2015         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
2016                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
2017                 rsi->expectedDesc == NULL)
2018                 ereport(ERROR,
2019                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
2020                                  errmsg("set-valued function called in context that "
2021                                                 "cannot accept a set")));
2022
2023         rsi->returnMode = SFRM_Materialize;
2024
2025         /* it's a simple type, so don't use get_call_result_type() */
2026         tupdesc = rsi->expectedDesc;
2027
2028         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
2029
2030         ret_tdesc = CreateTupleDescCopy(tupdesc);
2031         BlessTupleDesc(ret_tdesc);
2032         tuple_store =
2033                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
2034                                                           false, work_mem);
2035
2036         MemoryContextSwitchTo(old_cxt);
2037
2038         tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
2039                                                                         "jsonb_array_elements temporary cxt",
2040                                                                         ALLOCSET_DEFAULT_SIZES);
2041
2042         it = JsonbIteratorInit(&jb->root);
2043
2044         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
2045         {
2046                 skipNested = true;
2047
2048                 if (r == WJB_ELEM)
2049                 {
2050                         HeapTuple       tuple;
2051                         Datum           values[1];
2052                         bool            nulls[1] = {false};
2053
2054                         /* use the tmp context so we can clean up after each tuple is done */
2055                         old_cxt = MemoryContextSwitchTo(tmp_cxt);
2056
2057                         if (!as_text)
2058                         {
2059                                 Jsonb      *val = JsonbValueToJsonb(&v);
2060
2061                                 values[0] = PointerGetDatum(val);
2062                         }
2063                         else
2064                         {
2065                                 if (v.type == jbvNull)
2066                                 {
2067                                         /* a json null is an sql null in text mode */
2068                                         nulls[0] = true;
2069                                         values[0] = (Datum) NULL;
2070                                 }
2071                                 else
2072                                 {
2073                                         text       *sv;
2074
2075                                         if (v.type == jbvString)
2076                                         {
2077                                                 /* in text mode scalar strings should be dequoted */
2078                                                 sv = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
2079                                         }
2080                                         else
2081                                         {
2082                                                 /* turn anything else into a json string */
2083                                                 StringInfo      jtext = makeStringInfo();
2084                                                 Jsonb      *jb = JsonbValueToJsonb(&v);
2085
2086                                                 (void) JsonbToCString(jtext, &jb->root, 0);
2087                                                 sv = cstring_to_text_with_len(jtext->data, jtext->len);
2088                                         }
2089
2090                                         values[0] = PointerGetDatum(sv);
2091                                 }
2092                         }
2093
2094                         tuple = heap_form_tuple(ret_tdesc, values, nulls);
2095
2096                         tuplestore_puttuple(tuple_store, tuple);
2097
2098                         /* clean up and switch back */
2099                         MemoryContextSwitchTo(old_cxt);
2100                         MemoryContextReset(tmp_cxt);
2101                 }
2102         }
2103
2104         MemoryContextDelete(tmp_cxt);
2105
2106         rsi->setResult = tuple_store;
2107         rsi->setDesc = ret_tdesc;
2108
2109         PG_RETURN_NULL();
2110 }
2111
2112 Datum
2113 json_array_elements(PG_FUNCTION_ARGS)
2114 {
2115         return elements_worker(fcinfo, "json_array_elements", false);
2116 }
2117
2118 Datum
2119 json_array_elements_text(PG_FUNCTION_ARGS)
2120 {
2121         return elements_worker(fcinfo, "json_array_elements_text", true);
2122 }
2123
2124 static Datum
2125 elements_worker(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
2126 {
2127         text       *json = PG_GETARG_TEXT_PP(0);
2128
2129         /* elements only needs escaped strings when as_text */
2130         JsonLexContext *lex = makeJsonLexContext(json, as_text);
2131         JsonSemAction *sem;
2132         ReturnSetInfo *rsi;
2133         MemoryContext old_cxt;
2134         TupleDesc       tupdesc;
2135         ElementsState *state;
2136
2137         state = palloc0(sizeof(ElementsState));
2138         sem = palloc0(sizeof(JsonSemAction));
2139
2140         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
2141
2142         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
2143                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
2144                 rsi->expectedDesc == NULL)
2145                 ereport(ERROR,
2146                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
2147                                  errmsg("set-valued function called in context that "
2148                                                 "cannot accept a set")));
2149
2150         rsi->returnMode = SFRM_Materialize;
2151
2152         /* it's a simple type, so don't use get_call_result_type() */
2153         tupdesc = rsi->expectedDesc;
2154
2155         /* make these in a sufficiently long-lived memory context */
2156         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
2157
2158         state->ret_tdesc = CreateTupleDescCopy(tupdesc);
2159         BlessTupleDesc(state->ret_tdesc);
2160         state->tuple_store =
2161                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
2162                                                           false, work_mem);
2163
2164         MemoryContextSwitchTo(old_cxt);
2165
2166         sem->semstate = (void *) state;
2167         sem->object_start = elements_object_start;
2168         sem->scalar = elements_scalar;
2169         sem->array_element_start = elements_array_element_start;
2170         sem->array_element_end = elements_array_element_end;
2171
2172         state->function_name = funcname;
2173         state->normalize_results = as_text;
2174         state->next_scalar = false;
2175         state->lex = lex;
2176         state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
2177                                                                                    "json_array_elements temporary cxt",
2178                                                                                    ALLOCSET_DEFAULT_SIZES);
2179
2180         pg_parse_json(lex, sem);
2181
2182         MemoryContextDelete(state->tmp_cxt);
2183
2184         rsi->setResult = state->tuple_store;
2185         rsi->setDesc = state->ret_tdesc;
2186
2187         PG_RETURN_NULL();
2188 }
2189
2190 static void
2191 elements_array_element_start(void *state, bool isnull)
2192 {
2193         ElementsState *_state = (ElementsState *) state;
2194
2195         /* save a pointer to where the value starts */
2196         if (_state->lex->lex_level == 1)
2197         {
2198                 /*
2199                  * next_scalar will be reset in the array_element_end handler, and
2200                  * since we know the value is a scalar there is no danger of it being
2201                  * on while recursing down the tree.
2202                  */
2203                 if (_state->normalize_results && _state->lex->token_type == JSON_TOKEN_STRING)
2204                         _state->next_scalar = true;
2205                 else
2206                         _state->result_start = _state->lex->token_start;
2207         }
2208 }
2209
2210 static void
2211 elements_array_element_end(void *state, bool isnull)
2212 {
2213         ElementsState *_state = (ElementsState *) state;
2214         MemoryContext old_cxt;
2215         int                     len;
2216         text       *val;
2217         HeapTuple       tuple;
2218         Datum           values[1];
2219         bool            nulls[1] = {false};
2220
2221         /* skip over nested objects */
2222         if (_state->lex->lex_level != 1)
2223                 return;
2224
2225         /* use the tmp context so we can clean up after each tuple is done */
2226         old_cxt = MemoryContextSwitchTo(_state->tmp_cxt);
2227
2228         if (isnull && _state->normalize_results)
2229         {
2230                 nulls[0] = true;
2231                 values[0] = (Datum) NULL;
2232         }
2233         else if (_state->next_scalar)
2234         {
2235                 values[0] = CStringGetTextDatum(_state->normalized_scalar);
2236                 _state->next_scalar = false;
2237         }
2238         else
2239         {
2240                 len = _state->lex->prev_token_terminator - _state->result_start;
2241                 val = cstring_to_text_with_len(_state->result_start, len);
2242                 values[0] = PointerGetDatum(val);
2243         }
2244
2245         tuple = heap_form_tuple(_state->ret_tdesc, values, nulls);
2246
2247         tuplestore_puttuple(_state->tuple_store, tuple);
2248
2249         /* clean up and switch back */
2250         MemoryContextSwitchTo(old_cxt);
2251         MemoryContextReset(_state->tmp_cxt);
2252 }
2253
2254 static void
2255 elements_object_start(void *state)
2256 {
2257         ElementsState *_state = (ElementsState *) state;
2258
2259         /* json structure check */
2260         if (_state->lex->lex_level == 0)
2261                 ereport(ERROR,
2262                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2263                                  errmsg("cannot call %s on a non-array",
2264                                                 _state->function_name)));
2265 }
2266
2267 static void
2268 elements_scalar(void *state, char *token, JsonTokenType tokentype)
2269 {
2270         ElementsState *_state = (ElementsState *) state;
2271
2272         /* json structure check */
2273         if (_state->lex->lex_level == 0)
2274                 ereport(ERROR,
2275                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2276                                  errmsg("cannot call %s on a scalar",
2277                                                 _state->function_name)));
2278
2279         /* supply de-escaped value if required */
2280         if (_state->next_scalar)
2281                 _state->normalized_scalar = token;
2282 }
2283
2284 /*
2285  * SQL function json_populate_record
2286  *
2287  * set fields in a record from the argument json
2288  *
2289  * Code adapted shamelessly from hstore's populate_record
2290  * which is in turn partly adapted from record_out.
2291  *
2292  * The json is decomposed into a hash table, in which each
2293  * field in the record is then looked up by name. For jsonb
2294  * we fetch the values direct from the object.
2295  */
2296 Datum
2297 jsonb_populate_record(PG_FUNCTION_ARGS)
2298 {
2299         return populate_record_worker(fcinfo, "jsonb_populate_record", true);
2300 }
2301
2302 Datum
2303 jsonb_to_record(PG_FUNCTION_ARGS)
2304 {
2305         return populate_record_worker(fcinfo, "jsonb_to_record", false);
2306 }
2307
2308 Datum
2309 json_populate_record(PG_FUNCTION_ARGS)
2310 {
2311         return populate_record_worker(fcinfo, "json_populate_record", true);
2312 }
2313
2314 Datum
2315 json_to_record(PG_FUNCTION_ARGS)
2316 {
2317         return populate_record_worker(fcinfo, "json_to_record", false);
2318 }
2319
2320 /* helper function for diagnostics */
2321 static void
2322 populate_array_report_expected_array(PopulateArrayContext *ctx, int ndim)
2323 {
2324         if (ndim <= 0)
2325         {
2326                 if (ctx->colname)
2327                         ereport(ERROR,
2328                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2329                                          errmsg("expected JSON array"),
2330                                          errhint("See the value of key \"%s\".", ctx->colname)));
2331                 else
2332                         ereport(ERROR,
2333                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2334                                          errmsg("expected JSON array")));
2335         }
2336         else
2337         {
2338                 StringInfoData indices;
2339                 int                     i;
2340
2341                 initStringInfo(&indices);
2342
2343                 Assert(ctx->ndims > 0 && ndim < ctx->ndims);
2344
2345                 for (i = 0; i < ndim; i++)
2346                         appendStringInfo(&indices, "[%d]", ctx->sizes[i]);
2347
2348                 if (ctx->colname)
2349                         ereport(ERROR,
2350                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2351                                          errmsg("expected JSON array"),
2352                                          errhint("See the array element %s of key \"%s\".",
2353                                                          indices.data, ctx->colname)));
2354                 else
2355                         ereport(ERROR,
2356                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2357                                          errmsg("expected JSON array"),
2358                                          errhint("See the array element %s.",
2359                                                          indices.data)));
2360         }
2361 }
2362
2363 /* set the number of dimensions of the populated array when it becomes known */
2364 static void
2365 populate_array_assign_ndims(PopulateArrayContext *ctx, int ndims)
2366 {
2367         int                     i;
2368
2369         Assert(ctx->ndims <= 0);
2370
2371         if (ndims <= 0)
2372                 populate_array_report_expected_array(ctx, ndims);
2373
2374         ctx->ndims = ndims;
2375         ctx->dims = palloc(sizeof(int) * ndims);
2376         ctx->sizes = palloc0(sizeof(int) * ndims);
2377
2378         for (i = 0; i < ndims; i++)
2379                 ctx->dims[i] = -1;              /* dimensions are unknown yet */
2380 }
2381
2382 /* check the populated subarray dimension */
2383 static void
2384 populate_array_check_dimension(PopulateArrayContext *ctx, int ndim)
2385 {
2386         int                     dim = ctx->sizes[ndim]; /* current dimension counter */
2387
2388         if (ctx->dims[ndim] == -1)
2389                 ctx->dims[ndim] = dim;  /* assign dimension if not yet known */
2390         else if (ctx->dims[ndim] != dim)
2391                 ereport(ERROR,
2392                                 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2393                                  errmsg("malformed JSON array"),
2394                                  errdetail("Multidimensional arrays must have "
2395                                                    "sub-arrays with matching dimensions.")));
2396
2397         /* reset the current array dimension size counter */
2398         ctx->sizes[ndim] = 0;
2399
2400         /* increment the parent dimension counter if it is a nested sub-array */
2401         if (ndim > 0)
2402                 ctx->sizes[ndim - 1]++;
2403 }
2404
2405 static void
2406 populate_array_element(PopulateArrayContext *ctx, int ndim, JsValue *jsv)
2407 {
2408         Datum           element;
2409         bool            element_isnull;
2410
2411         /* populate the array element */
2412         element = populate_record_field(ctx->aio->element_info,
2413                                                                         ctx->aio->element_type,
2414                                                                         ctx->aio->element_typmod,
2415                                                                         NULL, ctx->mcxt, PointerGetDatum(NULL),
2416                                                                         jsv, &element_isnull);
2417
2418         accumArrayResult(ctx->astate, element, element_isnull,
2419                                          ctx->aio->element_type, ctx->acxt);
2420
2421         Assert(ndim > 0);
2422         ctx->sizes[ndim - 1]++;         /* increment current dimension counter */
2423 }
2424
2425 /* json object start handler for populate_array_json() */
2426 static void
2427 populate_array_object_start(void *_state)
2428 {
2429         PopulateArrayState *state = (PopulateArrayState *) _state;
2430         int                     ndim = state->lex->lex_level;
2431
2432         if (state->ctx->ndims <= 0)
2433                 populate_array_assign_ndims(state->ctx, ndim);
2434         else if (ndim < state->ctx->ndims)
2435                 populate_array_report_expected_array(state->ctx, ndim);
2436 }
2437
2438 /* json array end handler for populate_array_json() */
2439 static void
2440 populate_array_array_end(void *_state)
2441 {
2442         PopulateArrayState *state = (PopulateArrayState *) _state;
2443         PopulateArrayContext *ctx = state->ctx;
2444         int                     ndim = state->lex->lex_level;
2445
2446         if (ctx->ndims <= 0)
2447                 populate_array_assign_ndims(ctx, ndim + 1);
2448
2449         if (ndim < ctx->ndims)
2450                 populate_array_check_dimension(ctx, ndim);
2451 }
2452
2453 /* json array element start handler for populate_array_json() */
2454 static void
2455 populate_array_element_start(void *_state, bool isnull)
2456 {
2457         PopulateArrayState *state = (PopulateArrayState *) _state;
2458         int                     ndim = state->lex->lex_level;
2459
2460         if (state->ctx->ndims <= 0 || ndim == state->ctx->ndims)
2461         {
2462                 /* remember current array element start */
2463                 state->element_start = state->lex->token_start;
2464                 state->element_type = state->lex->token_type;
2465                 state->element_scalar = NULL;
2466         }
2467 }
2468
2469 /* json array element end handler for populate_array_json() */
2470 static void
2471 populate_array_element_end(void *_state, bool isnull)
2472 {
2473         PopulateArrayState *state = (PopulateArrayState *) _state;
2474         PopulateArrayContext *ctx = state->ctx;
2475         int                     ndim = state->lex->lex_level;
2476
2477         Assert(ctx->ndims > 0);
2478
2479         if (ndim == ctx->ndims)
2480         {
2481                 JsValue         jsv;
2482
2483                 jsv.is_json = true;
2484                 jsv.val.json.type = state->element_type;
2485
2486                 if (isnull)
2487                 {
2488                         Assert(jsv.val.json.type == JSON_TOKEN_NULL);
2489                         jsv.val.json.str = NULL;
2490                         jsv.val.json.len = 0;
2491                 }
2492                 else if (state->element_scalar)
2493                 {
2494                         jsv.val.json.str = state->element_scalar;
2495                         jsv.val.json.len = -1;  /* null-terminated */
2496                 }
2497                 else
2498                 {
2499                         jsv.val.json.str = state->element_start;
2500                         jsv.val.json.len = (state->lex->prev_token_terminator -
2501                                                                 state->element_start) * sizeof(char);
2502                 }
2503
2504                 populate_array_element(ctx, ndim, &jsv);
2505         }
2506 }
2507
2508 /* json scalar handler for populate_array_json() */
2509 static void
2510 populate_array_scalar(void *_state, char *token, JsonTokenType tokentype)
2511 {
2512         PopulateArrayState *state = (PopulateArrayState *) _state;
2513         PopulateArrayContext *ctx = state->ctx;
2514         int                     ndim = state->lex->lex_level;
2515
2516         if (ctx->ndims <= 0)
2517                 populate_array_assign_ndims(ctx, ndim);
2518         else if (ndim < ctx->ndims)
2519                 populate_array_report_expected_array(ctx, ndim);
2520
2521         if (ndim == ctx->ndims)
2522         {
2523                 /* remember the scalar element token */
2524                 state->element_scalar = token;
2525                 /* element_type must already be set in populate_array_element_start() */
2526                 Assert(state->element_type == tokentype);
2527         }
2528 }
2529
2530 /* parse a json array and populate array */
2531 static void
2532 populate_array_json(PopulateArrayContext *ctx, char *json, int len)
2533 {
2534         PopulateArrayState state;
2535         JsonSemAction sem;
2536
2537         state.lex = makeJsonLexContextCstringLen(json, len, true);
2538         state.ctx = ctx;
2539
2540         memset(&sem, 0, sizeof(sem));
2541         sem.semstate = (void *) &state;
2542         sem.object_start = populate_array_object_start;
2543         sem.array_end = populate_array_array_end;
2544         sem.array_element_start = populate_array_element_start;
2545         sem.array_element_end = populate_array_element_end;
2546         sem.scalar = populate_array_scalar;
2547
2548         pg_parse_json(state.lex, &sem);
2549
2550         /* number of dimensions should be already known */
2551         Assert(ctx->ndims > 0 && ctx->dims);
2552
2553         pfree(state.lex);
2554 }
2555
2556 /*
2557  * populate_array_dim_jsonb() -- Iterate recursively through jsonb sub-array
2558  *              elements and accumulate result using given ArrayBuildState.
2559  */
2560 static void
2561 populate_array_dim_jsonb(PopulateArrayContext *ctx, /* context */
2562                                                  JsonbValue *jbv,       /* jsonb sub-array */
2563                                                  int ndim)      /* current dimension */
2564 {
2565         JsonbContainer *jbc = jbv->val.binary.data;
2566         JsonbIterator *it;
2567         JsonbIteratorToken tok;
2568         JsonbValue      val;
2569         JsValue         jsv;
2570
2571         check_stack_depth();
2572
2573         if (jbv->type != jbvBinary || !JsonContainerIsArray(jbc))
2574                 populate_array_report_expected_array(ctx, ndim - 1);
2575
2576         Assert(!JsonContainerIsScalar(jbc));
2577
2578         it = JsonbIteratorInit(jbc);
2579
2580         tok = JsonbIteratorNext(&it, &val, true);
2581         Assert(tok == WJB_BEGIN_ARRAY);
2582
2583         tok = JsonbIteratorNext(&it, &val, true);
2584
2585         /*
2586          * If the number of dimensions is not yet known and we have found end of
2587          * the array, or the first child element is not an array, then assign the
2588          * number of dimensions now.
2589          */
2590         if (ctx->ndims <= 0 &&
2591                 (tok == WJB_END_ARRAY ||
2592                  (tok == WJB_ELEM &&
2593                   (val.type != jbvBinary ||
2594                    !JsonContainerIsArray(val.val.binary.data)))))
2595                 populate_array_assign_ndims(ctx, ndim);
2596
2597         jsv.is_json = false;
2598         jsv.val.jsonb = &val;
2599
2600         /* process all the array elements */
2601         while (tok == WJB_ELEM)
2602         {
2603                 /*
2604                  * Recurse only if the dimensions of dimensions is still unknown or if
2605                  * it is not the innermost dimension.
2606                  */
2607                 if (ctx->ndims > 0 && ndim >= ctx->ndims)
2608                         populate_array_element(ctx, ndim, &jsv);
2609                 else
2610                 {
2611                         /* populate child sub-array */
2612                         populate_array_dim_jsonb(ctx, &val, ndim + 1);
2613
2614                         /* number of dimensions should be already known */
2615                         Assert(ctx->ndims > 0 && ctx->dims);
2616
2617                         populate_array_check_dimension(ctx, ndim);
2618                 }
2619
2620                 tok = JsonbIteratorNext(&it, &val, true);
2621         }
2622
2623         Assert(tok == WJB_END_ARRAY);
2624
2625         /* free iterator, iterating until WJB_DONE */
2626         tok = JsonbIteratorNext(&it, &val, true);
2627         Assert(tok == WJB_DONE && !it);
2628 }
2629
2630 /* recursively populate an array from json/jsonb */
2631 static Datum
2632 populate_array(ArrayIOData *aio,
2633                            const char *colname,
2634                            MemoryContext mcxt,
2635                            JsValue *jsv)
2636 {
2637         PopulateArrayContext ctx;
2638         Datum           result;
2639         int                *lbs;
2640         int                     i;
2641
2642         ctx.aio = aio;
2643         ctx.mcxt = mcxt;
2644         ctx.acxt = CurrentMemoryContext;
2645         ctx.astate = initArrayResult(aio->element_type, ctx.acxt, true);
2646         ctx.colname = colname;
2647         ctx.ndims = 0;                          /* unknown yet */
2648         ctx.dims = NULL;
2649         ctx.sizes = NULL;
2650
2651         if (jsv->is_json)
2652                 populate_array_json(&ctx, jsv->val.json.str,
2653                                                         jsv->val.json.len >= 0 ? jsv->val.json.len
2654                                                         : strlen(jsv->val.json.str));
2655         else
2656         {
2657                 populate_array_dim_jsonb(&ctx, jsv->val.jsonb, 1);
2658                 ctx.dims[0] = ctx.sizes[0];
2659         }
2660
2661         Assert(ctx.ndims > 0);
2662
2663         lbs = palloc(sizeof(int) * ctx.ndims);
2664
2665         for (i = 0; i < ctx.ndims; i++)
2666                 lbs[i] = 1;
2667
2668         result = makeMdArrayResult(ctx.astate, ctx.ndims, ctx.dims, lbs,
2669                                                            ctx.acxt, true);
2670
2671         pfree(ctx.dims);
2672         pfree(ctx.sizes);
2673         pfree(lbs);
2674
2675         return result;
2676 }
2677
2678 static void
2679 JsValueToJsObject(JsValue *jsv, JsObject *jso)
2680 {
2681         jso->is_json = jsv->is_json;
2682
2683         if (jsv->is_json)
2684         {
2685                 /* convert plain-text json into a hash table */
2686                 jso->val.json_hash =
2687                         get_json_object_as_hash(jsv->val.json.str,
2688                                                                         jsv->val.json.len >= 0
2689                                                                         ? jsv->val.json.len
2690                                                                         : strlen(jsv->val.json.str),
2691                                                                         "populate_composite");
2692         }
2693         else
2694         {
2695                 JsonbValue *jbv = jsv->val.jsonb;
2696
2697                 if (jbv->type == jbvBinary &&
2698                         JsonContainerIsObject(jbv->val.binary.data))
2699                 {
2700                         jso->val.jsonb_cont = jbv->val.binary.data;
2701                 }
2702                 else
2703                 {
2704                         bool            is_scalar;
2705
2706                         is_scalar = IsAJsonbScalar(jbv) ||
2707                                 (jbv->type == jbvBinary &&
2708                                  JsonContainerIsScalar(jbv->val.binary.data));
2709                         ereport(ERROR,
2710                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2711                                          is_scalar
2712                                          ? errmsg("cannot call %s on a scalar",
2713                                                           "populate_composite")
2714                                          : errmsg("cannot call %s on an array",
2715                                                           "populate_composite")));
2716                 }
2717         }
2718 }
2719
2720 /* acquire or update cached tuple descriptor for a composite type */
2721 static void
2722 update_cached_tupdesc(CompositeIOData *io, MemoryContext mcxt)
2723 {
2724         if (!io->tupdesc ||
2725                 io->tupdesc->tdtypeid != io->base_typid ||
2726                 io->tupdesc->tdtypmod != io->base_typmod)
2727         {
2728                 TupleDesc       tupdesc = lookup_rowtype_tupdesc(io->base_typid,
2729                                                                                                          io->base_typmod);
2730                 MemoryContext oldcxt;
2731
2732                 if (io->tupdesc)
2733                         FreeTupleDesc(io->tupdesc);
2734
2735                 /* copy tuple desc without constraints into cache memory context */
2736                 oldcxt = MemoryContextSwitchTo(mcxt);
2737                 io->tupdesc = CreateTupleDescCopy(tupdesc);
2738                 MemoryContextSwitchTo(oldcxt);
2739
2740                 ReleaseTupleDesc(tupdesc);
2741         }
2742 }
2743
2744 /* recursively populate a composite (row type) value from json/jsonb */
2745 static Datum
2746 populate_composite(CompositeIOData *io,
2747                                    Oid typid,
2748                                    const char *colname,
2749                                    MemoryContext mcxt,
2750                                    HeapTupleHeader defaultval,
2751                                    JsValue *jsv,
2752                                    bool isnull)
2753 {
2754         Datum           result;
2755
2756         /* acquire/update cached tuple descriptor */
2757         update_cached_tupdesc(io, mcxt);
2758
2759         if (isnull)
2760                 result = (Datum) 0;
2761         else
2762         {
2763                 HeapTupleHeader tuple;
2764                 JsObject        jso;
2765
2766                 /* prepare input value */
2767                 JsValueToJsObject(jsv, &jso);
2768
2769                 /* populate resulting record tuple */
2770                 tuple = populate_record(io->tupdesc, &io->record_io,
2771                                                                 defaultval, mcxt, &jso);
2772                 result = HeapTupleHeaderGetDatum(tuple);
2773
2774                 JsObjectFree(&jso);
2775         }
2776
2777         /*
2778          * If it's domain over composite, check domain constraints.  (This should
2779          * probably get refactored so that we can see the TYPECAT value, but for
2780          * now, we can tell by comparing typid to base_typid.)
2781          */
2782         if (typid != io->base_typid && typid != RECORDOID)
2783                 domain_check(result, isnull, typid, &io->domain_info, mcxt);
2784
2785         return result;
2786 }
2787
2788 /* populate non-null scalar value from json/jsonb value */
2789 static Datum
2790 populate_scalar(ScalarIOData *io, Oid typid, int32 typmod, JsValue *jsv)
2791 {
2792         Datum           res;
2793         char       *str = NULL;
2794         char       *json = NULL;
2795
2796         if (jsv->is_json)
2797         {
2798                 int                     len = jsv->val.json.len;
2799
2800                 json = jsv->val.json.str;
2801                 Assert(json);
2802
2803                 /* already done the hard work in the json case */
2804                 if ((typid == JSONOID || typid == JSONBOID) &&
2805                         jsv->val.json.type == JSON_TOKEN_STRING)
2806                 {
2807                         /*
2808                          * Add quotes around string value (should be already escaped) if
2809                          * converting to json/jsonb.
2810                          */
2811
2812                         if (len < 0)
2813                                 len = strlen(json);
2814
2815                         str = palloc(len + sizeof(char) * 3);
2816                         str[0] = '"';
2817                         memcpy(&str[1], json, len);
2818                         str[len + 1] = '"';
2819                         str[len + 2] = '\0';
2820                 }
2821                 else if (len >= 0)
2822                 {
2823                         /* Need to copy non-null-terminated string */
2824                         str = palloc(len + 1 * sizeof(char));
2825                         memcpy(str, json, len);
2826                         str[len] = '\0';
2827                 }
2828                 else
2829                         str = json;                     /* null-terminated string */
2830         }
2831         else
2832         {
2833                 JsonbValue *jbv = jsv->val.jsonb;
2834
2835                 if (typid == JSONBOID)
2836                 {
2837                         Jsonb      *jsonb = JsonbValueToJsonb(jbv); /* directly use jsonb */
2838
2839                         return JsonbPGetDatum(jsonb);
2840                 }
2841                 /* convert jsonb to string for typio call */
2842                 else if (typid == JSONOID && jbv->type != jbvBinary)
2843                 {
2844                         /*
2845                          * Convert scalar jsonb (non-scalars are passed here as jbvBinary)
2846                          * to json string, preserving quotes around top-level strings.
2847                          */
2848                         Jsonb      *jsonb = JsonbValueToJsonb(jbv);
2849
2850                         str = JsonbToCString(NULL, &jsonb->root, VARSIZE(jsonb));
2851                 }
2852                 else if (jbv->type == jbvString)        /* quotes are stripped */
2853                         str = pnstrdup(jbv->val.string.val, jbv->val.string.len);
2854                 else if (jbv->type == jbvBool)
2855                         str = pstrdup(jbv->val.boolean ? "true" : "false");
2856                 else if (jbv->type == jbvNumeric)
2857                         str = DatumGetCString(DirectFunctionCall1(numeric_out,
2858                                                                                                           PointerGetDatum(jbv->val.numeric)));
2859                 else if (jbv->type == jbvBinary)
2860                         str = JsonbToCString(NULL, jbv->val.binary.data,
2861                                                                  jbv->val.binary.len);
2862                 else
2863                         elog(ERROR, "unrecognized jsonb type: %d", (int) jbv->type);
2864         }
2865
2866         res = InputFunctionCall(&io->typiofunc, str, io->typioparam, typmod);
2867
2868         /* free temporary buffer */
2869         if (str != json)
2870                 pfree(str);
2871
2872         return res;
2873 }
2874
2875 static Datum
2876 populate_domain(DomainIOData *io,
2877                                 Oid typid,
2878                                 const char *colname,
2879                                 MemoryContext mcxt,
2880                                 JsValue *jsv,
2881                                 bool isnull)
2882 {
2883         Datum           res;
2884
2885         if (isnull)
2886                 res = (Datum) 0;
2887         else
2888         {
2889                 res = populate_record_field(io->base_io,
2890                                                                         io->base_typid, io->base_typmod,
2891                                                                         colname, mcxt, PointerGetDatum(NULL),
2892                                                                         jsv, &isnull);
2893                 Assert(!isnull);
2894         }
2895
2896         domain_check(res, isnull, typid, &io->domain_info, mcxt);
2897
2898         return res;
2899 }
2900
2901 /* prepare column metadata cache for the given type */
2902 static void
2903 prepare_column_cache(ColumnIOData *column,
2904                                          Oid typid,
2905                                          int32 typmod,
2906                                          MemoryContext mcxt,
2907                                          bool need_scalar)
2908 {
2909         HeapTuple       tup;
2910         Form_pg_type type;
2911
2912         column->typid = typid;
2913         column->typmod = typmod;
2914
2915         tup = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typid));
2916         if (!HeapTupleIsValid(tup))
2917                 elog(ERROR, "cache lookup failed for type %u", typid);
2918
2919         type = (Form_pg_type) GETSTRUCT(tup);
2920
2921         if (type->typtype == TYPTYPE_DOMAIN)
2922         {
2923                 /*
2924                  * We can move directly to the bottom base type; domain_check() will
2925                  * take care of checking all constraints for a stack of domains.
2926                  */
2927                 Oid                     base_typid;
2928                 int32           base_typmod = typmod;
2929
2930                 base_typid = getBaseTypeAndTypmod(typid, &base_typmod);
2931                 if (get_typtype(base_typid) == TYPTYPE_COMPOSITE)
2932                 {
2933                         /* domain over composite has its own code path */
2934                         column->typcat = TYPECAT_COMPOSITE_DOMAIN;
2935                         column->io.composite.record_io = NULL;
2936                         column->io.composite.tupdesc = NULL;
2937                         column->io.composite.base_typid = base_typid;
2938                         column->io.composite.base_typmod = base_typmod;
2939                         column->io.composite.domain_info = NULL;
2940                 }
2941                 else
2942                 {
2943                         /* domain over anything else */
2944                         column->typcat = TYPECAT_DOMAIN;
2945                         column->io.domain.base_typid = base_typid;
2946                         column->io.domain.base_typmod = base_typmod;
2947                         column->io.domain.base_io =
2948                                 MemoryContextAllocZero(mcxt, sizeof(ColumnIOData));
2949                         column->io.domain.domain_info = NULL;
2950                 }
2951         }
2952         else if (type->typtype == TYPTYPE_COMPOSITE || typid == RECORDOID)
2953         {
2954                 column->typcat = TYPECAT_COMPOSITE;
2955                 column->io.composite.record_io = NULL;
2956                 column->io.composite.tupdesc = NULL;
2957                 column->io.composite.base_typid = typid;
2958                 column->io.composite.base_typmod = typmod;
2959                 column->io.composite.domain_info = NULL;
2960         }
2961         else if (type->typlen == -1 && OidIsValid(type->typelem))
2962         {
2963                 column->typcat = TYPECAT_ARRAY;
2964                 column->io.array.element_info = MemoryContextAllocZero(mcxt,
2965                                                                                                                            sizeof(ColumnIOData));
2966                 column->io.array.element_type = type->typelem;
2967                 /* array element typemod stored in attribute's typmod */
2968                 column->io.array.element_typmod = typmod;
2969         }
2970         else
2971         {
2972                 column->typcat = TYPECAT_SCALAR;
2973                 need_scalar = true;
2974         }
2975
2976         /* caller can force us to look up scalar_io info even for non-scalars */
2977         if (need_scalar)
2978         {
2979                 Oid                     typioproc;
2980
2981                 getTypeInputInfo(typid, &typioproc, &column->scalar_io.typioparam);
2982                 fmgr_info_cxt(typioproc, &column->scalar_io.typiofunc, mcxt);
2983         }
2984
2985         ReleaseSysCache(tup);
2986 }
2987
2988 /* recursively populate a record field or an array element from a json/jsonb value */
2989 static Datum
2990 populate_record_field(ColumnIOData *col,
2991                                           Oid typid,
2992                                           int32 typmod,
2993                                           const char *colname,
2994                                           MemoryContext mcxt,
2995                                           Datum defaultval,
2996                                           JsValue *jsv,
2997                                           bool *isnull)
2998 {
2999         TypeCat         typcat;
3000
3001         check_stack_depth();
3002
3003         /*
3004          * Prepare column metadata cache for the given type.  Force lookup of the
3005          * scalar_io data so that the json string hack below will work.
3006          */
3007         if (col->typid != typid || col->typmod != typmod)
3008                 prepare_column_cache(col, typid, typmod, mcxt, true);
3009
3010         *isnull = JsValueIsNull(jsv);
3011
3012         typcat = col->typcat;
3013
3014         /* try to convert json string to a non-scalar type through input function */
3015         if (JsValueIsString(jsv) &&
3016                 (typcat == TYPECAT_ARRAY ||
3017                  typcat == TYPECAT_COMPOSITE ||
3018                  typcat == TYPECAT_COMPOSITE_DOMAIN))
3019                 typcat = TYPECAT_SCALAR;
3020
3021         /* we must perform domain checks for NULLs, otherwise exit immediately */
3022         if (*isnull &&
3023                 typcat != TYPECAT_DOMAIN &&
3024                 typcat != TYPECAT_COMPOSITE_DOMAIN)
3025                 return (Datum) 0;
3026
3027         switch (typcat)
3028         {
3029                 case TYPECAT_SCALAR:
3030                         return populate_scalar(&col->scalar_io, typid, typmod, jsv);
3031
3032                 case TYPECAT_ARRAY:
3033                         return populate_array(&col->io.array, colname, mcxt, jsv);
3034
3035                 case TYPECAT_COMPOSITE:
3036                 case TYPECAT_COMPOSITE_DOMAIN:
3037                         return populate_composite(&col->io.composite, typid,
3038                                                                           colname, mcxt,
3039                                                                           DatumGetPointer(defaultval)
3040                                                                           ? DatumGetHeapTupleHeader(defaultval)
3041                                                                           : NULL,
3042                                                                           jsv, *isnull);
3043
3044                 case TYPECAT_DOMAIN:
3045                         return populate_domain(&col->io.domain, typid, colname, mcxt,
3046                                                                    jsv, *isnull);
3047
3048                 default:
3049                         elog(ERROR, "unrecognized type category '%c'", typcat);
3050                         return (Datum) 0;
3051         }
3052 }
3053
3054 static RecordIOData *
3055 allocate_record_info(MemoryContext mcxt, int ncolumns)
3056 {
3057         RecordIOData *data = (RecordIOData *)
3058         MemoryContextAlloc(mcxt,
3059                                            offsetof(RecordIOData, columns) +
3060                                            ncolumns * sizeof(ColumnIOData));
3061
3062         data->record_type = InvalidOid;
3063         data->record_typmod = 0;
3064         data->ncolumns = ncolumns;
3065         MemSet(data->columns, 0, sizeof(ColumnIOData) * ncolumns);
3066
3067         return data;
3068 }
3069
3070 static bool
3071 JsObjectGetField(JsObject *obj, char *field, JsValue *jsv)
3072 {
3073         jsv->is_json = obj->is_json;
3074
3075         if (jsv->is_json)
3076         {
3077                 JsonHashEntry *hashentry = hash_search(obj->val.json_hash, field,
3078                                                                                            HASH_FIND, NULL);
3079
3080                 jsv->val.json.type = hashentry ? hashentry->type : JSON_TOKEN_NULL;
3081                 jsv->val.json.str = jsv->val.json.type == JSON_TOKEN_NULL ? NULL :
3082                         hashentry->val;
3083                 jsv->val.json.len = jsv->val.json.str ? -1 : 0; /* null-terminated */
3084
3085                 return hashentry != NULL;
3086         }
3087         else
3088         {
3089                 jsv->val.jsonb = !obj->val.jsonb_cont ? NULL :
3090                         findJsonbValueFromContainerLen(obj->val.jsonb_cont, JB_FOBJECT,
3091                                                                                    field, strlen(field));
3092
3093                 return jsv->val.jsonb != NULL;
3094         }
3095 }
3096
3097 /* populate a record tuple from json/jsonb value */
3098 static HeapTupleHeader
3099 populate_record(TupleDesc tupdesc,
3100                                 RecordIOData **record_p,
3101                                 HeapTupleHeader defaultval,
3102                                 MemoryContext mcxt,
3103                                 JsObject *obj)
3104 {
3105         RecordIOData *record = *record_p;
3106         Datum      *values;
3107         bool       *nulls;
3108         HeapTuple       res;
3109         int                     ncolumns = tupdesc->natts;
3110         int                     i;
3111
3112         /*
3113          * if the input json is empty, we can only skip the rest if we were passed
3114          * in a non-null record, since otherwise there may be issues with domain
3115          * nulls.
3116          */
3117         if (defaultval && JsObjectIsEmpty(obj))
3118                 return defaultval;
3119
3120         /* (re)allocate metadata cache */
3121         if (record == NULL ||
3122                 record->ncolumns != ncolumns)
3123                 *record_p = record = allocate_record_info(mcxt, ncolumns);
3124
3125         /* invalidate metadata cache if the record type has changed */
3126         if (record->record_type != tupdesc->tdtypeid ||
3127                 record->record_typmod != tupdesc->tdtypmod)
3128         {
3129                 MemSet(record, 0, offsetof(RecordIOData, columns) +
3130                            ncolumns * sizeof(ColumnIOData));
3131                 record->record_type = tupdesc->tdtypeid;
3132                 record->record_typmod = tupdesc->tdtypmod;
3133                 record->ncolumns = ncolumns;
3134         }
3135
3136         values = (Datum *) palloc(ncolumns * sizeof(Datum));
3137         nulls = (bool *) palloc(ncolumns * sizeof(bool));
3138
3139         if (defaultval)
3140         {
3141                 HeapTupleData tuple;
3142
3143                 /* Build a temporary HeapTuple control structure */
3144                 tuple.t_len = HeapTupleHeaderGetDatumLength(defaultval);
3145                 ItemPointerSetInvalid(&(tuple.t_self));
3146                 tuple.t_tableOid = InvalidOid;
3147                 tuple.t_data = defaultval;
3148
3149                 /* Break down the tuple into fields */
3150                 heap_deform_tuple(&tuple, tupdesc, values, nulls);
3151         }
3152         else
3153         {
3154                 for (i = 0; i < ncolumns; ++i)
3155                 {
3156                         values[i] = (Datum) 0;
3157                         nulls[i] = true;
3158                 }
3159         }
3160
3161         for (i = 0; i < ncolumns; ++i)
3162         {
3163                 Form_pg_attribute att = TupleDescAttr(tupdesc, i);
3164                 char       *colname = NameStr(att->attname);
3165                 JsValue         field = {0};
3166                 bool            found;
3167
3168                 /* Ignore dropped columns in datatype */
3169                 if (att->attisdropped)
3170                 {
3171                         nulls[i] = true;
3172                         continue;
3173                 }
3174
3175                 found = JsObjectGetField(obj, colname, &field);
3176
3177                 /*
3178                  * we can't just skip here if the key wasn't found since we might have
3179                  * a domain to deal with. If we were passed in a non-null record
3180                  * datum, we assume that the existing values are valid (if they're
3181                  * not, then it's not our fault), but if we were passed in a null,
3182                  * then every field which we don't populate needs to be run through
3183                  * the input function just in case it's a domain type.
3184                  */
3185                 if (defaultval && !found)
3186                         continue;
3187
3188                 values[i] = populate_record_field(&record->columns[i],
3189                                                                                   att->atttypid,
3190                                                                                   att->atttypmod,
3191                                                                                   colname,
3192                                                                                   mcxt,
3193                                                                                   nulls[i] ? (Datum) 0 : values[i],
3194                                                                                   &field,
3195                                                                                   &nulls[i]);
3196         }
3197
3198         res = heap_form_tuple(tupdesc, values, nulls);
3199
3200         pfree(values);
3201         pfree(nulls);
3202
3203         return res->t_data;
3204 }
3205
3206 static Datum
3207 populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
3208                                            bool have_record_arg)
3209 {
3210         int                     json_arg_num = have_record_arg ? 1 : 0;
3211         Oid                     jtype = get_fn_expr_argtype(fcinfo->flinfo, json_arg_num);
3212         JsValue         jsv = {0};
3213         HeapTupleHeader rec;
3214         Datum           rettuple;
3215         JsonbValue      jbv;
3216         MemoryContext fnmcxt = fcinfo->flinfo->fn_mcxt;
3217         PopulateRecordCache *cache = fcinfo->flinfo->fn_extra;
3218
3219         Assert(jtype == JSONOID || jtype == JSONBOID);
3220
3221         /*
3222          * If first time through, identify input/result record type.  Note that
3223          * this stanza looks only at fcinfo context, which can't change during the
3224          * query; so we may not be able to fully resolve a RECORD input type yet.
3225          */
3226         if (!cache)
3227         {
3228                 fcinfo->flinfo->fn_extra = cache =
3229                         MemoryContextAllocZero(fnmcxt, sizeof(*cache));
3230
3231                 if (have_record_arg)
3232                 {
3233                         /*
3234                          * json{b}_populate_record case: result type will be same as first
3235                          * argument's.
3236                          */
3237                         cache->argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
3238                         prepare_column_cache(&cache->c,
3239                                                                  cache->argtype, -1,
3240                                                                  fnmcxt, false);
3241                         if (cache->c.typcat != TYPECAT_COMPOSITE &&
3242                                 cache->c.typcat != TYPECAT_COMPOSITE_DOMAIN)
3243                                 ereport(ERROR,
3244                                                 (errcode(ERRCODE_DATATYPE_MISMATCH),
3245                                                  errmsg("first argument of %s must be a row type",
3246                                                                 funcname)));
3247                 }
3248                 else
3249                 {
3250                         /*
3251                          * json{b}_to_record case: result type is specified by calling
3252                          * query.  Here it is syntactically impossible to specify the
3253                          * target type as domain-over-composite.
3254                          */
3255                         TupleDesc       tupdesc;
3256                         MemoryContext old_cxt;
3257
3258                         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
3259                                 ereport(ERROR,
3260                                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
3261                                                  errmsg("function returning record called in context "
3262                                                                 "that cannot accept type record"),
3263                                                  errhint("Try calling the function in the FROM clause "
3264                                                                  "using a column definition list.")));
3265
3266                         Assert(tupdesc);
3267                         cache->argtype = tupdesc->tdtypeid;
3268
3269                         /* Save identified tupdesc */
3270                         old_cxt = MemoryContextSwitchTo(fnmcxt);
3271                         cache->c.io.composite.tupdesc = CreateTupleDescCopy(tupdesc);
3272                         cache->c.io.composite.base_typid = tupdesc->tdtypeid;
3273                         cache->c.io.composite.base_typmod = tupdesc->tdtypmod;
3274                         MemoryContextSwitchTo(old_cxt);
3275                 }
3276         }
3277
3278         /* Collect record arg if we have one */
3279         if (have_record_arg && !PG_ARGISNULL(0))
3280         {
3281                 rec = PG_GETARG_HEAPTUPLEHEADER(0);
3282
3283                 /*
3284                  * When declared arg type is RECORD, identify actual record type from
3285                  * the tuple itself.  Note the lookup_rowtype_tupdesc call in
3286                  * update_cached_tupdesc will fail if we're unable to do this.
3287                  */
3288                 if (cache->argtype == RECORDOID)
3289                 {
3290                         cache->c.io.composite.base_typid = HeapTupleHeaderGetTypeId(rec);
3291                         cache->c.io.composite.base_typmod = HeapTupleHeaderGetTypMod(rec);
3292                 }
3293         }
3294         else
3295                 rec = NULL;
3296
3297         /* If no JSON argument, just return the record (if any) unchanged */
3298         if (PG_ARGISNULL(json_arg_num))
3299         {
3300                 if (rec)
3301                         PG_RETURN_POINTER(rec);
3302                 else
3303                         PG_RETURN_NULL();
3304         }
3305
3306         jsv.is_json = jtype == JSONOID;
3307
3308         if (jsv.is_json)
3309         {
3310                 text       *json = PG_GETARG_TEXT_PP(json_arg_num);
3311
3312                 jsv.val.json.str = VARDATA_ANY(json);
3313                 jsv.val.json.len = VARSIZE_ANY_EXHDR(json);
3314                 jsv.val.json.type = JSON_TOKEN_INVALID; /* not used in
3315                                                                                                  * populate_composite() */
3316         }
3317         else
3318         {
3319                 Jsonb      *jb = PG_GETARG_JSONB_P(json_arg_num);
3320
3321                 jsv.val.jsonb = &jbv;
3322
3323                 /* fill binary jsonb value pointing to jb */
3324                 jbv.type = jbvBinary;
3325                 jbv.val.binary.data = &jb->root;
3326                 jbv.val.binary.len = VARSIZE(jb) - VARHDRSZ;
3327         }
3328
3329         rettuple = populate_composite(&cache->c.io.composite, cache->argtype,
3330                                                                   NULL, fnmcxt, rec, &jsv, false);
3331
3332         PG_RETURN_DATUM(rettuple);
3333 }
3334
3335 /*
3336  * get_json_object_as_hash
3337  *
3338  * decompose a json object into a hash table.
3339  */
3340 static HTAB *
3341 get_json_object_as_hash(char *json, int len, const char *funcname)
3342 {
3343         HASHCTL         ctl;
3344         HTAB       *tab;
3345         JHashState *state;
3346         JsonLexContext *lex = makeJsonLexContextCstringLen(json, len, true);
3347         JsonSemAction *sem;
3348
3349         memset(&ctl, 0, sizeof(ctl));
3350         ctl.keysize = NAMEDATALEN;
3351         ctl.entrysize = sizeof(JsonHashEntry);
3352         ctl.hcxt = CurrentMemoryContext;
3353         tab = hash_create("json object hashtable",
3354                                           100,
3355                                           &ctl,
3356                                           HASH_ELEM | HASH_CONTEXT);
3357
3358         state = palloc0(sizeof(JHashState));
3359         sem = palloc0(sizeof(JsonSemAction));
3360
3361         state->function_name = funcname;
3362         state->hash = tab;
3363         state->lex = lex;
3364
3365         sem->semstate = (void *) state;
3366         sem->array_start = hash_array_start;
3367         sem->scalar = hash_scalar;
3368         sem->object_field_start = hash_object_field_start;
3369         sem->object_field_end = hash_object_field_end;
3370
3371         pg_parse_json(lex, sem);
3372
3373         return tab;
3374 }
3375
3376 static void
3377 hash_object_field_start(void *state, char *fname, bool isnull)
3378 {
3379         JHashState *_state = (JHashState *) state;
3380
3381         if (_state->lex->lex_level > 1)
3382                 return;
3383
3384         /* remember token type */
3385         _state->saved_token_type = _state->lex->token_type;
3386
3387         if (_state->lex->token_type == JSON_TOKEN_ARRAY_START ||
3388                 _state->lex->token_type == JSON_TOKEN_OBJECT_START)
3389         {
3390                 /* remember start position of the whole text of the subobject */
3391                 _state->save_json_start = _state->lex->token_start;
3392         }
3393         else
3394         {
3395                 /* must be a scalar */
3396                 _state->save_json_start = NULL;
3397         }
3398 }
3399
3400 static void
3401 hash_object_field_end(void *state, char *fname, bool isnull)
3402 {
3403         JHashState *_state = (JHashState *) state;
3404         JsonHashEntry *hashentry;
3405         bool            found;
3406
3407         /*
3408          * Ignore nested fields.
3409          */
3410         if (_state->lex->lex_level > 1)
3411                 return;
3412
3413         /*
3414          * Ignore field names >= NAMEDATALEN - they can't match a record field.
3415          * (Note: without this test, the hash code would truncate the string at
3416          * NAMEDATALEN-1, and could then match against a similarly-truncated
3417          * record field name.  That would be a reasonable behavior, but this code
3418          * has previously insisted on exact equality, so we keep this behavior.)
3419          */
3420         if (strlen(fname) >= NAMEDATALEN)
3421                 return;
3422
3423         hashentry = hash_search(_state->hash, fname, HASH_ENTER, &found);
3424
3425         /*
3426          * found being true indicates a duplicate. We don't do anything about
3427          * that, a later field with the same name overrides the earlier field.
3428          */
3429
3430         hashentry->type = _state->saved_token_type;
3431         Assert(isnull == (hashentry->type == JSON_TOKEN_NULL));
3432
3433         if (_state->save_json_start != NULL)
3434         {
3435                 int                     len = _state->lex->prev_token_terminator - _state->save_json_start;
3436                 char       *val = palloc((len + 1) * sizeof(char));
3437
3438                 memcpy(val, _state->save_json_start, len);
3439                 val[len] = '\0';
3440                 hashentry->val = val;
3441         }
3442         else
3443         {
3444                 /* must have had a scalar instead */
3445                 hashentry->val = _state->saved_scalar;
3446         }
3447 }
3448
3449 static void
3450 hash_array_start(void *state)
3451 {
3452         JHashState *_state = (JHashState *) state;
3453
3454         if (_state->lex->lex_level == 0)
3455                 ereport(ERROR,
3456                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3457                                  errmsg("cannot call %s on an array", _state->function_name)));
3458 }
3459
3460 static void
3461 hash_scalar(void *state, char *token, JsonTokenType tokentype)
3462 {
3463         JHashState *_state = (JHashState *) state;
3464
3465         if (_state->lex->lex_level == 0)
3466                 ereport(ERROR,
3467                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3468                                  errmsg("cannot call %s on a scalar", _state->function_name)));
3469
3470         if (_state->lex->lex_level == 1)
3471         {
3472                 _state->saved_scalar = token;
3473                 /* saved_token_type must already be set in hash_object_field_start() */
3474                 Assert(_state->saved_token_type == tokentype);
3475         }
3476 }
3477
3478
3479 /*
3480  * SQL function json_populate_recordset
3481  *
3482  * set fields in a set of records from the argument json,
3483  * which must be an array of objects.
3484  *
3485  * similar to json_populate_record, but the tuple-building code
3486  * is pushed down into the semantic action handlers so it's done
3487  * per object in the array.
3488  */
3489 Datum
3490 jsonb_populate_recordset(PG_FUNCTION_ARGS)
3491 {
3492         return populate_recordset_worker(fcinfo, "jsonb_populate_recordset", true);
3493 }
3494
3495 Datum
3496 jsonb_to_recordset(PG_FUNCTION_ARGS)
3497 {
3498         return populate_recordset_worker(fcinfo, "jsonb_to_recordset", false);
3499 }
3500
3501 Datum
3502 json_populate_recordset(PG_FUNCTION_ARGS)
3503 {
3504         return populate_recordset_worker(fcinfo, "json_populate_recordset", true);
3505 }
3506
3507 Datum
3508 json_to_recordset(PG_FUNCTION_ARGS)
3509 {
3510         return populate_recordset_worker(fcinfo, "json_to_recordset", false);
3511 }
3512
3513 static void
3514 populate_recordset_record(PopulateRecordsetState *state, JsObject *obj)
3515 {
3516         PopulateRecordsetCache *cache = state->cache;
3517         HeapTupleHeader tuphead;
3518         HeapTupleData tuple;
3519
3520         /* acquire/update cached tuple descriptor */
3521         update_cached_tupdesc(&cache->c.io.composite, cache->fn_mcxt);
3522
3523         /* replace record fields from json */
3524         tuphead = populate_record(cache->c.io.composite.tupdesc,
3525                                                           &cache->c.io.composite.record_io,
3526                                                           state->rec,
3527                                                           cache->fn_mcxt,
3528                                                           obj);
3529
3530         /* if it's domain over composite, check domain constraints */
3531         if (cache->c.typcat == TYPECAT_COMPOSITE_DOMAIN)
3532                 domain_check(HeapTupleHeaderGetDatum(tuphead), false,
3533                                          cache->argtype,
3534                                          &cache->c.io.composite.domain_info,
3535                                          cache->fn_mcxt);
3536
3537         /* ok, save into tuplestore */
3538         tuple.t_len = HeapTupleHeaderGetDatumLength(tuphead);
3539         ItemPointerSetInvalid(&(tuple.t_self));
3540         tuple.t_tableOid = InvalidOid;
3541         tuple.t_data = tuphead;
3542
3543         tuplestore_puttuple(state->tuple_store, &tuple);
3544 }
3545
3546 /*
3547  * common worker for json_populate_recordset() and json_to_recordset()
3548  */
3549 static Datum
3550 populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
3551                                                   bool have_record_arg)
3552 {
3553         int                     json_arg_num = have_record_arg ? 1 : 0;
3554         Oid                     jtype = get_fn_expr_argtype(fcinfo->flinfo, json_arg_num);
3555         ReturnSetInfo *rsi;
3556         MemoryContext old_cxt;
3557         HeapTupleHeader rec;
3558         PopulateRecordsetCache *cache = fcinfo->flinfo->fn_extra;
3559         PopulateRecordsetState *state;
3560
3561         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
3562
3563         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
3564                 (rsi->allowedModes & SFRM_Materialize) == 0)
3565                 ereport(ERROR,
3566                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
3567                                  errmsg("set-valued function called in context that "
3568                                                 "cannot accept a set")));
3569
3570         rsi->returnMode = SFRM_Materialize;
3571
3572         /*
3573          * If first time through, identify input/result record type.  Note that
3574          * this stanza looks only at fcinfo context, which can't change during the
3575          * query; so we may not be able to fully resolve a RECORD input type yet.
3576          */
3577         if (!cache)
3578         {
3579                 fcinfo->flinfo->fn_extra = cache =
3580                         MemoryContextAllocZero(fcinfo->flinfo->fn_mcxt, sizeof(*cache));
3581                 cache->fn_mcxt = fcinfo->flinfo->fn_mcxt;
3582
3583                 if (have_record_arg)
3584                 {
3585                         /*
3586                          * json{b}_populate_recordset case: result type will be same as
3587                          * first argument's.
3588                          */
3589                         cache->argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
3590                         prepare_column_cache(&cache->c,
3591                                                                  cache->argtype, -1,
3592                                                                  cache->fn_mcxt, false);
3593                         if (cache->c.typcat != TYPECAT_COMPOSITE &&
3594                                 cache->c.typcat != TYPECAT_COMPOSITE_DOMAIN)
3595                                 ereport(ERROR,
3596                                                 (errcode(ERRCODE_DATATYPE_MISMATCH),
3597                                                  errmsg("first argument of %s must be a row type",
3598                                                                 funcname)));
3599                 }
3600                 else
3601                 {
3602                         /*
3603                          * json{b}_to_recordset case: result type is specified by calling
3604                          * query.  Here it is syntactically impossible to specify the
3605                          * target type as domain-over-composite.
3606                          */
3607                         TupleDesc       tupdesc;
3608
3609                         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
3610                                 ereport(ERROR,
3611                                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
3612                                                  errmsg("function returning record called in context "
3613                                                                 "that cannot accept type record"),
3614                                                  errhint("Try calling the function in the FROM clause "
3615                                                                  "using a column definition list.")));
3616
3617                         Assert(tupdesc);
3618                         cache->argtype = tupdesc->tdtypeid;
3619
3620                         /* Save identified tupdesc */
3621                         old_cxt = MemoryContextSwitchTo(cache->fn_mcxt);
3622                         cache->c.io.composite.tupdesc = CreateTupleDescCopy(tupdesc);
3623                         cache->c.io.composite.base_typid = tupdesc->tdtypeid;
3624                         cache->c.io.composite.base_typmod = tupdesc->tdtypmod;
3625                         MemoryContextSwitchTo(old_cxt);
3626                 }
3627         }
3628
3629         /* Collect record arg if we have one */
3630         if (have_record_arg && !PG_ARGISNULL(0))
3631         {
3632                 rec = PG_GETARG_HEAPTUPLEHEADER(0);
3633
3634                 /*
3635                  * When declared arg type is RECORD, identify actual record type from
3636                  * the tuple itself.  Note the lookup_rowtype_tupdesc call in
3637                  * update_cached_tupdesc will fail if we're unable to do this.
3638                  */
3639                 if (cache->argtype == RECORDOID)
3640                 {
3641                         cache->c.io.composite.base_typid = HeapTupleHeaderGetTypeId(rec);
3642                         cache->c.io.composite.base_typmod = HeapTupleHeaderGetTypMod(rec);
3643                 }
3644         }
3645         else
3646                 rec = NULL;
3647
3648         /* if the json is null send back an empty set */
3649         if (PG_ARGISNULL(json_arg_num))
3650                 PG_RETURN_NULL();
3651
3652         state = palloc0(sizeof(PopulateRecordsetState));
3653
3654         /* make tuplestore in a sufficiently long-lived memory context */
3655         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
3656         state->tuple_store = tuplestore_begin_heap(rsi->allowedModes &
3657                                                                                            SFRM_Materialize_Random,
3658                                                                                            false, work_mem);
3659         MemoryContextSwitchTo(old_cxt);
3660
3661         state->function_name = funcname;
3662         state->cache = cache;
3663         state->rec = rec;
3664
3665         if (jtype == JSONOID)
3666         {
3667                 text       *json = PG_GETARG_TEXT_PP(json_arg_num);
3668                 JsonLexContext *lex;
3669                 JsonSemAction *sem;
3670
3671                 sem = palloc0(sizeof(JsonSemAction));
3672
3673                 lex = makeJsonLexContext(json, true);
3674
3675                 sem->semstate = (void *) state;
3676                 sem->array_start = populate_recordset_array_start;
3677                 sem->array_element_start = populate_recordset_array_element_start;
3678                 sem->scalar = populate_recordset_scalar;
3679                 sem->object_field_start = populate_recordset_object_field_start;
3680                 sem->object_field_end = populate_recordset_object_field_end;
3681                 sem->object_start = populate_recordset_object_start;
3682                 sem->object_end = populate_recordset_object_end;
3683
3684                 state->lex = lex;
3685
3686                 pg_parse_json(lex, sem);
3687         }
3688         else
3689         {
3690                 Jsonb      *jb = PG_GETARG_JSONB_P(json_arg_num);
3691                 JsonbIterator *it;
3692                 JsonbValue      v;
3693                 bool            skipNested = false;
3694                 JsonbIteratorToken r;
3695
3696                 Assert(jtype == JSONBOID);
3697
3698                 if (JB_ROOT_IS_SCALAR(jb) || !JB_ROOT_IS_ARRAY(jb))
3699                         ereport(ERROR,
3700                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3701                                          errmsg("cannot call %s on a non-array",
3702                                                         funcname)));
3703
3704                 it = JsonbIteratorInit(&jb->root);
3705
3706                 while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
3707                 {
3708                         skipNested = true;
3709
3710                         if (r == WJB_ELEM)
3711                         {
3712                                 JsObject        obj;
3713
3714                                 if (v.type != jbvBinary ||
3715                                         !JsonContainerIsObject(v.val.binary.data))
3716                                         ereport(ERROR,
3717                                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3718                                                          errmsg("argument of %s must be an array of objects",
3719                                                                         funcname)));
3720
3721                                 obj.is_json = false;
3722                                 obj.val.jsonb_cont = v.val.binary.data;
3723
3724                                 populate_recordset_record(state, &obj);
3725                         }
3726                 }
3727         }
3728
3729         rsi->setResult = state->tuple_store;
3730         rsi->setDesc = cache->c.io.composite.tupdesc;
3731
3732         PG_RETURN_NULL();
3733 }
3734
3735 static void
3736 populate_recordset_object_start(void *state)
3737 {
3738         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3739         int                     lex_level = _state->lex->lex_level;
3740         HASHCTL         ctl;
3741
3742         /* Reject object at top level: we must have an array at level 0 */
3743         if (lex_level == 0)
3744                 ereport(ERROR,
3745                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3746                                  errmsg("cannot call %s on an object",
3747                                                 _state->function_name)));
3748
3749         /* Nested objects require no special processing */
3750         if (lex_level > 1)
3751                 return;
3752
3753         /* Object at level 1: set up a new hash table for this object */
3754         memset(&ctl, 0, sizeof(ctl));
3755         ctl.keysize = NAMEDATALEN;
3756         ctl.entrysize = sizeof(JsonHashEntry);
3757         ctl.hcxt = CurrentMemoryContext;
3758         _state->json_hash = hash_create("json object hashtable",
3759                                                                         100,
3760                                                                         &ctl,
3761                                                                         HASH_ELEM | HASH_CONTEXT);
3762 }
3763
3764 static void
3765 populate_recordset_object_end(void *state)
3766 {
3767         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3768         JsObject        obj;
3769
3770         /* Nested objects require no special processing */
3771         if (_state->lex->lex_level > 1)
3772                 return;
3773
3774         obj.is_json = true;
3775         obj.val.json_hash = _state->json_hash;
3776
3777         /* Otherwise, construct and return a tuple based on this level-1 object */
3778         populate_recordset_record(_state, &obj);
3779
3780         /* Done with hash for this object */
3781         hash_destroy(_state->json_hash);
3782         _state->json_hash = NULL;
3783 }
3784
3785 static void
3786 populate_recordset_array_element_start(void *state, bool isnull)
3787 {
3788         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3789
3790         if (_state->lex->lex_level == 1 &&
3791                 _state->lex->token_type != JSON_TOKEN_OBJECT_START)
3792                 ereport(ERROR,
3793                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3794                                  errmsg("argument of %s must be an array of objects",
3795                                                 _state->function_name)));
3796 }
3797
3798 static void
3799 populate_recordset_array_start(void *state)
3800 {
3801         /* nothing to do */
3802 }
3803
3804 static void
3805 populate_recordset_scalar(void *state, char *token, JsonTokenType tokentype)
3806 {
3807         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3808
3809         if (_state->lex->lex_level == 0)
3810                 ereport(ERROR,
3811                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3812                                  errmsg("cannot call %s on a scalar",
3813                                                 _state->function_name)));
3814
3815         if (_state->lex->lex_level == 2)
3816                 _state->saved_scalar = token;
3817 }
3818
3819 static void
3820 populate_recordset_object_field_start(void *state, char *fname, bool isnull)
3821 {
3822         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3823
3824         if (_state->lex->lex_level > 2)
3825                 return;
3826
3827         _state->saved_token_type = _state->lex->token_type;
3828
3829         if (_state->lex->token_type == JSON_TOKEN_ARRAY_START ||
3830                 _state->lex->token_type == JSON_TOKEN_OBJECT_START)
3831         {
3832                 _state->save_json_start = _state->lex->token_start;
3833         }
3834         else
3835         {
3836                 _state->save_json_start = NULL;
3837         }
3838 }
3839
3840 static void
3841 populate_recordset_object_field_end(void *state, char *fname, bool isnull)
3842 {
3843         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3844         JsonHashEntry *hashentry;
3845         bool            found;
3846
3847         /*
3848          * Ignore nested fields.
3849          */
3850         if (_state->lex->lex_level > 2)
3851                 return;
3852
3853         /*
3854          * Ignore field names >= NAMEDATALEN - they can't match a record field.
3855          * (Note: without this test, the hash code would truncate the string at
3856          * NAMEDATALEN-1, and could then match against a similarly-truncated
3857          * record field name.  That would be a reasonable behavior, but this code
3858          * has previously insisted on exact equality, so we keep this behavior.)
3859          */
3860         if (strlen(fname) >= NAMEDATALEN)
3861                 return;
3862
3863         hashentry = hash_search(_state->json_hash, fname, HASH_ENTER, &found);
3864
3865         /*
3866          * found being true indicates a duplicate. We don't do anything about
3867          * that, a later field with the same name overrides the earlier field.
3868          */
3869
3870         hashentry->type = _state->saved_token_type;
3871         Assert(isnull == (hashentry->type == JSON_TOKEN_NULL));
3872
3873         if (_state->save_json_start != NULL)
3874         {
3875                 int                     len = _state->lex->prev_token_terminator - _state->save_json_start;
3876                 char       *val = palloc((len + 1) * sizeof(char));
3877
3878                 memcpy(val, _state->save_json_start, len);
3879                 val[len] = '\0';
3880                 hashentry->val = val;
3881         }
3882         else
3883         {
3884                 /* must have had a scalar instead */
3885                 hashentry->val = _state->saved_scalar;
3886         }
3887 }
3888
3889 /*
3890  * findJsonbValueFromContainer() wrapper that sets up JsonbValue key string.
3891  */
3892 static JsonbValue *
3893 findJsonbValueFromContainerLen(JsonbContainer *container, uint32 flags,
3894                                                            char *key, uint32 keylen)
3895 {
3896         JsonbValue      k;
3897
3898         k.type = jbvString;
3899         k.val.string.val = key;
3900         k.val.string.len = keylen;
3901
3902         return findJsonbValueFromContainer(container, flags, &k);
3903 }
3904
3905 /*
3906  * Semantic actions for json_strip_nulls.
3907  *
3908  * Simply repeat the input on the output unless we encounter
3909  * a null object field. State for this is set when the field
3910  * is started and reset when the scalar action (which must be next)
3911  * is called.
3912  */
3913
3914 static void
3915 sn_object_start(void *state)
3916 {
3917         StripnullState *_state = (StripnullState *) state;
3918
3919         appendStringInfoCharMacro(_state->strval, '{');
3920 }
3921
3922 static void
3923 sn_object_end(void *state)
3924 {
3925         StripnullState *_state = (StripnullState *) state;
3926
3927         appendStringInfoCharMacro(_state->strval, '}');
3928 }
3929
3930 static void
3931 sn_array_start(void *state)
3932 {
3933         StripnullState *_state = (StripnullState *) state;
3934
3935         appendStringInfoCharMacro(_state->strval, '[');
3936 }
3937
3938 static void
3939 sn_array_end(void *state)
3940 {
3941         StripnullState *_state = (StripnullState *) state;
3942
3943         appendStringInfoCharMacro(_state->strval, ']');
3944 }
3945
3946 static void
3947 sn_object_field_start(void *state, char *fname, bool isnull)
3948 {
3949         StripnullState *_state = (StripnullState *) state;
3950
3951         if (isnull)
3952         {
3953                 /*
3954                  * The next thing must be a scalar or isnull couldn't be true, so
3955                  * there is no danger of this state being carried down into a nested
3956                  * object or array. The flag will be reset in the scalar action.
3957                  */
3958                 _state->skip_next_null = true;
3959                 return;
3960         }
3961
3962         if (_state->strval->data[_state->strval->len - 1] != '{')
3963                 appendStringInfoCharMacro(_state->strval, ',');
3964
3965         /*
3966          * Unfortunately we don't have the quoted and escaped string any more, so
3967          * we have to re-escape it.
3968          */
3969         escape_json(_state->strval, fname);
3970
3971         appendStringInfoCharMacro(_state->strval, ':');
3972 }
3973
3974 static void
3975 sn_array_element_start(void *state, bool isnull)
3976 {
3977         StripnullState *_state = (StripnullState *) state;
3978
3979         if (_state->strval->data[_state->strval->len - 1] != '[')
3980                 appendStringInfoCharMacro(_state->strval, ',');
3981 }
3982
3983 static void
3984 sn_scalar(void *state, char *token, JsonTokenType tokentype)
3985 {
3986         StripnullState *_state = (StripnullState *) state;
3987
3988         if (_state->skip_next_null)
3989         {
3990                 Assert(tokentype == JSON_TOKEN_NULL);
3991                 _state->skip_next_null = false;
3992                 return;
3993         }
3994
3995         if (tokentype == JSON_TOKEN_STRING)
3996                 escape_json(_state->strval, token);
3997         else
3998                 appendStringInfoString(_state->strval, token);
3999 }
4000
4001 /*
4002  * SQL function json_strip_nulls(json) -> json
4003  */
4004 Datum
4005 json_strip_nulls(PG_FUNCTION_ARGS)
4006 {
4007         text       *json = PG_GETARG_TEXT_PP(0);
4008         StripnullState *state;
4009         JsonLexContext *lex;
4010         JsonSemAction *sem;
4011
4012         lex = makeJsonLexContext(json, true);
4013         state = palloc0(sizeof(StripnullState));
4014         sem = palloc0(sizeof(JsonSemAction));
4015
4016         state->strval = makeStringInfo();
4017         state->skip_next_null = false;
4018         state->lex = lex;
4019
4020         sem->semstate = (void *) state;
4021         sem->object_start = sn_object_start;
4022         sem->object_end = sn_object_end;
4023         sem->array_start = sn_array_start;
4024         sem->array_end = sn_array_end;
4025         sem->scalar = sn_scalar;
4026         sem->array_element_start = sn_array_element_start;
4027         sem->object_field_start = sn_object_field_start;
4028
4029         pg_parse_json(lex, sem);
4030
4031         PG_RETURN_TEXT_P(cstring_to_text_with_len(state->strval->data,
4032                                                                                           state->strval->len));
4033
4034 }
4035
4036 /*
4037  * SQL function jsonb_strip_nulls(jsonb) -> jsonb
4038  */
4039 Datum
4040 jsonb_strip_nulls(PG_FUNCTION_ARGS)
4041 {
4042         Jsonb      *jb = PG_GETARG_JSONB_P(0);
4043         JsonbIterator *it;
4044         JsonbParseState *parseState = NULL;
4045         JsonbValue *res = NULL;
4046         JsonbValue      v,
4047                                 k;
4048         JsonbIteratorToken type;
4049         bool            last_was_key = false;
4050
4051         if (JB_ROOT_IS_SCALAR(jb))
4052                 PG_RETURN_POINTER(jb);
4053
4054         it = JsonbIteratorInit(&jb->root);
4055
4056         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
4057         {
4058                 Assert(!(type == WJB_KEY && last_was_key));
4059
4060                 if (type == WJB_KEY)
4061                 {
4062                         /* stash the key until we know if it has a null value */
4063                         k = v;
4064                         last_was_key = true;
4065                         continue;
4066                 }
4067
4068                 if (last_was_key)
4069                 {
4070                         /* if the last element was a key this one can't be */
4071                         last_was_key = false;
4072
4073                         /* skip this field if value is null */
4074                         if (type == WJB_VALUE && v.type == jbvNull)
4075                                 continue;
4076
4077                         /* otherwise, do a delayed push of the key */
4078                         (void) pushJsonbValue(&parseState, WJB_KEY, &k);
4079                 }
4080
4081                 if (type == WJB_VALUE || type == WJB_ELEM)
4082                         res = pushJsonbValue(&parseState, type, &v);
4083                 else
4084                         res = pushJsonbValue(&parseState, type, NULL);
4085         }
4086
4087         Assert(res != NULL);
4088
4089         PG_RETURN_POINTER(JsonbValueToJsonb(res));
4090 }
4091
4092 /*
4093  * Add values from the jsonb to the parse state.
4094  *
4095  * If the parse state container is an object, the jsonb is pushed as
4096  * a value, not a key.
4097  *
4098  * This needs to be done using an iterator because pushJsonbValue doesn't
4099  * like getting jbvBinary values, so we can't just push jb as a whole.
4100  */
4101 static void
4102 addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb)
4103 {
4104         JsonbIterator *it;
4105         JsonbValue *o = &(*jbps)->contVal;
4106         JsonbValue      v;
4107         JsonbIteratorToken type;
4108
4109         it = JsonbIteratorInit(&jb->root);
4110
4111         Assert(o->type == jbvArray || o->type == jbvObject);
4112
4113         if (JB_ROOT_IS_SCALAR(jb))
4114         {
4115                 (void) JsonbIteratorNext(&it, &v, false);       /* skip array header */
4116                 Assert(v.type == jbvArray);
4117                 (void) JsonbIteratorNext(&it, &v, false);       /* fetch scalar value */
4118
4119                 switch (o->type)
4120                 {
4121                         case jbvArray:
4122                                 (void) pushJsonbValue(jbps, WJB_ELEM, &v);
4123                                 break;
4124                         case jbvObject:
4125                                 (void) pushJsonbValue(jbps, WJB_VALUE, &v);
4126                                 break;
4127                         default:
4128                                 elog(ERROR, "unexpected parent of nested structure");
4129                 }
4130         }
4131         else
4132         {
4133                 while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
4134                 {
4135                         if (type == WJB_KEY || type == WJB_VALUE || type == WJB_ELEM)
4136                                 (void) pushJsonbValue(jbps, type, &v);
4137                         else
4138                                 (void) pushJsonbValue(jbps, type, NULL);
4139                 }
4140         }
4141
4142 }
4143
4144 /*
4145  * SQL function jsonb_pretty (jsonb)
4146  *
4147  * Pretty-printed text for the jsonb
4148  */
4149 Datum
4150 jsonb_pretty(PG_FUNCTION_ARGS)
4151 {
4152         Jsonb      *jb = PG_GETARG_JSONB_P(0);
4153         StringInfo      str = makeStringInfo();
4154
4155         JsonbToCStringIndent(str, &jb->root, VARSIZE(jb));
4156
4157         PG_RETURN_TEXT_P(cstring_to_text_with_len(str->data, str->len));
4158 }
4159
4160 /*
4161  * SQL function jsonb_concat (jsonb, jsonb)
4162  *
4163  * function for || operator
4164  */
4165 Datum
4166 jsonb_concat(PG_FUNCTION_ARGS)
4167 {
4168         Jsonb      *jb1 = PG_GETARG_JSONB_P(0);
4169         Jsonb      *jb2 = PG_GETARG_JSONB_P(1);
4170         JsonbParseState *state = NULL;
4171         JsonbValue *res;
4172         JsonbIterator *it1,
4173                            *it2;
4174
4175         /*
4176          * If one of the jsonb is empty, just return the other if it's not scalar
4177          * and both are of the same kind.  If it's a scalar or they are of
4178          * different kinds we need to perform the concatenation even if one is
4179          * empty.
4180          */
4181         if (JB_ROOT_IS_OBJECT(jb1) == JB_ROOT_IS_OBJECT(jb2))
4182         {
4183                 if (JB_ROOT_COUNT(jb1) == 0 && !JB_ROOT_IS_SCALAR(jb2))
4184                         PG_RETURN_JSONB_P(jb2);
4185                 else if (JB_ROOT_COUNT(jb2) == 0 && !JB_ROOT_IS_SCALAR(jb1))
4186                         PG_RETURN_JSONB_P(jb1);
4187         }
4188
4189         it1 = JsonbIteratorInit(&jb1->root);
4190         it2 = JsonbIteratorInit(&jb2->root);
4191
4192         res = IteratorConcat(&it1, &it2, &state);
4193
4194         Assert(res != NULL);
4195
4196         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4197 }
4198
4199
4200 /*
4201  * SQL function jsonb_delete (jsonb, text)
4202  *
4203  * return a copy of the jsonb with the indicated item
4204  * removed.
4205  */
4206 Datum
4207 jsonb_delete(PG_FUNCTION_ARGS)
4208 {
4209         Jsonb      *in = PG_GETARG_JSONB_P(0);
4210         text       *key = PG_GETARG_TEXT_PP(1);
4211         char       *keyptr = VARDATA_ANY(key);
4212         int                     keylen = VARSIZE_ANY_EXHDR(key);
4213         JsonbParseState *state = NULL;
4214         JsonbIterator *it;
4215         JsonbValue      v,
4216                            *res = NULL;
4217         bool            skipNested = false;
4218         JsonbIteratorToken r;
4219
4220         if (JB_ROOT_IS_SCALAR(in))
4221                 ereport(ERROR,
4222                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4223                                  errmsg("cannot delete from scalar")));
4224
4225         if (JB_ROOT_COUNT(in) == 0)
4226                 PG_RETURN_JSONB_P(in);
4227
4228         it = JsonbIteratorInit(&in->root);
4229
4230         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
4231         {
4232                 skipNested = true;
4233
4234                 if ((r == WJB_ELEM || r == WJB_KEY) &&
4235                         (v.type == jbvString && keylen == v.val.string.len &&
4236                          memcmp(keyptr, v.val.string.val, keylen) == 0))
4237                 {
4238                         /* skip corresponding value as well */
4239                         if (r == WJB_KEY)
4240                                 (void) JsonbIteratorNext(&it, &v, true);
4241
4242                         continue;
4243                 }
4244
4245                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4246         }
4247
4248         Assert(res != NULL);
4249
4250         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4251 }
4252
4253 /*
4254  * SQL function jsonb_delete (jsonb, variadic text[])
4255  *
4256  * return a copy of the jsonb with the indicated items
4257  * removed.
4258  */
4259 Datum
4260 jsonb_delete_array(PG_FUNCTION_ARGS)
4261 {
4262         Jsonb      *in = PG_GETARG_JSONB_P(0);
4263         ArrayType  *keys = PG_GETARG_ARRAYTYPE_P(1);
4264         Datum      *keys_elems;
4265         bool       *keys_nulls;
4266         int                     keys_len;
4267         JsonbParseState *state = NULL;
4268         JsonbIterator *it;
4269         JsonbValue      v,
4270                            *res = NULL;
4271         bool            skipNested = false;
4272         JsonbIteratorToken r;
4273
4274         if (ARR_NDIM(keys) > 1)
4275                 ereport(ERROR,
4276                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4277                                  errmsg("wrong number of array subscripts")));
4278
4279         if (JB_ROOT_IS_SCALAR(in))
4280                 ereport(ERROR,
4281                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4282                                  errmsg("cannot delete from scalar")));
4283
4284         if (JB_ROOT_COUNT(in) == 0)
4285                 PG_RETURN_JSONB_P(in);
4286
4287         deconstruct_array(keys, TEXTOID, -1, false, 'i',
4288                                           &keys_elems, &keys_nulls, &keys_len);
4289
4290         if (keys_len == 0)
4291                 PG_RETURN_JSONB_P(in);
4292
4293         it = JsonbIteratorInit(&in->root);
4294
4295         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
4296         {
4297                 skipNested = true;
4298
4299                 if ((r == WJB_ELEM || r == WJB_KEY) && v.type == jbvString)
4300                 {
4301                         int                     i;
4302                         bool            found = false;
4303
4304                         for (i = 0; i < keys_len; i++)
4305                         {
4306                                 char       *keyptr;
4307                                 int                     keylen;
4308
4309                                 if (keys_nulls[i])
4310                                         continue;
4311
4312                                 keyptr = VARDATA_ANY(keys_elems[i]);
4313                                 keylen = VARSIZE_ANY_EXHDR(keys_elems[i]);
4314                                 if (keylen == v.val.string.len &&
4315                                         memcmp(keyptr, v.val.string.val, keylen) == 0)
4316                                 {
4317                                         found = true;
4318                                         break;
4319                                 }
4320                         }
4321                         if (found)
4322                         {
4323                                 /* skip corresponding value as well */
4324                                 if (r == WJB_KEY)
4325                                         (void) JsonbIteratorNext(&it, &v, true);
4326
4327                                 continue;
4328                         }
4329                 }
4330
4331                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4332         }
4333
4334         Assert(res != NULL);
4335
4336         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4337 }
4338
4339 /*
4340  * SQL function jsonb_delete (jsonb, int)
4341  *
4342  * return a copy of the jsonb with the indicated item
4343  * removed. Negative int means count back from the
4344  * end of the items.
4345  */
4346 Datum
4347 jsonb_delete_idx(PG_FUNCTION_ARGS)
4348 {
4349         Jsonb      *in = PG_GETARG_JSONB_P(0);
4350         int                     idx = PG_GETARG_INT32(1);
4351         JsonbParseState *state = NULL;
4352         JsonbIterator *it;
4353         uint32          i = 0,
4354                                 n;
4355         JsonbValue      v,
4356                            *res = NULL;
4357         JsonbIteratorToken r;
4358
4359         if (JB_ROOT_IS_SCALAR(in))
4360                 ereport(ERROR,
4361                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4362                                  errmsg("cannot delete from scalar")));
4363
4364         if (JB_ROOT_IS_OBJECT(in))
4365                 ereport(ERROR,
4366                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4367                                  errmsg("cannot delete from object using integer index")));
4368
4369         if (JB_ROOT_COUNT(in) == 0)
4370                 PG_RETURN_JSONB_P(in);
4371
4372         it = JsonbIteratorInit(&in->root);
4373
4374         r = JsonbIteratorNext(&it, &v, false);
4375         Assert(r == WJB_BEGIN_ARRAY);
4376         n = v.val.array.nElems;
4377
4378         if (idx < 0)
4379         {
4380                 if (-idx > n)
4381                         idx = n;
4382                 else
4383                         idx = n + idx;
4384         }
4385
4386         if (idx >= n)
4387                 PG_RETURN_JSONB_P(in);
4388
4389         pushJsonbValue(&state, r, NULL);
4390
4391         while ((r = JsonbIteratorNext(&it, &v, true)) != WJB_DONE)
4392         {
4393                 if (r == WJB_ELEM)
4394                 {
4395                         if (i++ == idx)
4396                                 continue;
4397                 }
4398
4399                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4400         }
4401
4402         Assert(res != NULL);
4403
4404         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4405 }
4406
4407 /*
4408  * SQL function jsonb_set(jsonb, text[], jsonb, boolean)
4409  *
4410  */
4411 Datum
4412 jsonb_set(PG_FUNCTION_ARGS)
4413 {
4414         Jsonb      *in = PG_GETARG_JSONB_P(0);
4415         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4416         Jsonb      *newval = PG_GETARG_JSONB_P(2);
4417         bool            create = PG_GETARG_BOOL(3);
4418         JsonbValue *res = NULL;
4419         Datum      *path_elems;
4420         bool       *path_nulls;
4421         int                     path_len;
4422         JsonbIterator *it;
4423         JsonbParseState *st = NULL;
4424
4425         if (ARR_NDIM(path) > 1)
4426                 ereport(ERROR,
4427                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4428                                  errmsg("wrong number of array subscripts")));
4429
4430         if (JB_ROOT_IS_SCALAR(in))
4431                 ereport(ERROR,
4432                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4433                                  errmsg("cannot set path in scalar")));
4434
4435         if (JB_ROOT_COUNT(in) == 0 && !create)
4436                 PG_RETURN_JSONB_P(in);
4437
4438         deconstruct_array(path, TEXTOID, -1, false, 'i',
4439                                           &path_elems, &path_nulls, &path_len);
4440
4441         if (path_len == 0)
4442                 PG_RETURN_JSONB_P(in);
4443
4444         it = JsonbIteratorInit(&in->root);
4445
4446         res = setPath(&it, path_elems, path_nulls, path_len, &st,
4447                                   0, newval, create ? JB_PATH_CREATE : JB_PATH_REPLACE);
4448
4449         Assert(res != NULL);
4450
4451         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4452 }
4453
4454
4455 /*
4456  * SQL function jsonb_delete_path(jsonb, text[])
4457  */
4458 Datum
4459 jsonb_delete_path(PG_FUNCTION_ARGS)
4460 {
4461         Jsonb      *in = PG_GETARG_JSONB_P(0);
4462         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4463         JsonbValue *res = NULL;
4464         Datum      *path_elems;
4465         bool       *path_nulls;
4466         int                     path_len;
4467         JsonbIterator *it;
4468         JsonbParseState *st = NULL;
4469
4470         if (ARR_NDIM(path) > 1)
4471                 ereport(ERROR,
4472                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4473                                  errmsg("wrong number of array subscripts")));
4474
4475         if (JB_ROOT_IS_SCALAR(in))
4476                 ereport(ERROR,
4477                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4478                                  errmsg("cannot delete path in scalar")));
4479
4480         if (JB_ROOT_COUNT(in) == 0)
4481                 PG_RETURN_JSONB_P(in);
4482
4483         deconstruct_array(path, TEXTOID, -1, false, 'i',
4484                                           &path_elems, &path_nulls, &path_len);
4485
4486         if (path_len == 0)
4487                 PG_RETURN_JSONB_P(in);
4488
4489         it = JsonbIteratorInit(&in->root);
4490
4491         res = setPath(&it, path_elems, path_nulls, path_len, &st,
4492                                   0, NULL, JB_PATH_DELETE);
4493
4494         Assert(res != NULL);
4495
4496         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4497 }
4498
4499 /*
4500  * SQL function jsonb_insert(jsonb, text[], jsonb, boolean)
4501  *
4502  */
4503 Datum
4504 jsonb_insert(PG_FUNCTION_ARGS)
4505 {
4506         Jsonb      *in = PG_GETARG_JSONB_P(0);
4507         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4508         Jsonb      *newval = PG_GETARG_JSONB_P(2);
4509         bool            after = PG_GETARG_BOOL(3);
4510         JsonbValue *res = NULL;
4511         Datum      *path_elems;
4512         bool       *path_nulls;
4513         int                     path_len;
4514         JsonbIterator *it;
4515         JsonbParseState *st = NULL;
4516
4517         if (ARR_NDIM(path) > 1)
4518                 ereport(ERROR,
4519                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4520                                  errmsg("wrong number of array subscripts")));
4521
4522         if (JB_ROOT_IS_SCALAR(in))
4523                 ereport(ERROR,
4524                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4525                                  errmsg("cannot set path in scalar")));
4526
4527         deconstruct_array(path, TEXTOID, -1, false, 'i',
4528                                           &path_elems, &path_nulls, &path_len);
4529
4530         if (path_len == 0)
4531                 PG_RETURN_JSONB_P(in);
4532
4533         it = JsonbIteratorInit(&in->root);
4534
4535         res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, newval,
4536                                   after ? JB_PATH_INSERT_AFTER : JB_PATH_INSERT_BEFORE);
4537
4538         Assert(res != NULL);
4539
4540         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4541 }
4542
4543 /*
4544  * Iterate over all jsonb objects and merge them into one.
4545  * The logic of this function copied from the same hstore function,
4546  * except the case, when it1 & it2 represents jbvObject.
4547  * In that case we just append the content of it2 to it1 without any
4548  * verifications.
4549  */
4550 static JsonbValue *
4551 IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
4552                            JsonbParseState **state)
4553 {
4554         JsonbValue      v1,
4555                                 v2,
4556                            *res = NULL;
4557         JsonbIteratorToken r1,
4558                                 r2,
4559                                 rk1,
4560                                 rk2;
4561
4562         r1 = rk1 = JsonbIteratorNext(it1, &v1, false);
4563         r2 = rk2 = JsonbIteratorNext(it2, &v2, false);
4564
4565         /*
4566          * Both elements are objects.
4567          */
4568         if (rk1 == WJB_BEGIN_OBJECT && rk2 == WJB_BEGIN_OBJECT)
4569         {
4570                 /*
4571                  * Append the all tokens from v1 to res, except last WJB_END_OBJECT
4572                  * (because res will not be finished yet).
4573                  */
4574                 pushJsonbValue(state, r1, NULL);
4575                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_OBJECT)
4576                         pushJsonbValue(state, r1, &v1);
4577
4578                 /*
4579                  * Append the all tokens from v2 to res, include last WJB_END_OBJECT
4580                  * (the concatenation will be completed).
4581                  */
4582                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_DONE)
4583                         res = pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
4584         }
4585
4586         /*
4587          * Both elements are arrays (either can be scalar).
4588          */
4589         else if (rk1 == WJB_BEGIN_ARRAY && rk2 == WJB_BEGIN_ARRAY)
4590         {
4591                 pushJsonbValue(state, r1, NULL);
4592
4593                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_ARRAY)
4594                 {
4595                         Assert(r1 == WJB_ELEM);
4596                         pushJsonbValue(state, r1, &v1);
4597                 }
4598
4599                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_END_ARRAY)
4600                 {
4601                         Assert(r2 == WJB_ELEM);
4602                         pushJsonbValue(state, WJB_ELEM, &v2);
4603                 }
4604
4605                 res = pushJsonbValue(state, WJB_END_ARRAY, NULL /* signal to sort */ );
4606         }
4607         /* have we got array || object or object || array? */
4608         else if (((rk1 == WJB_BEGIN_ARRAY && !(*it1)->isScalar) && rk2 == WJB_BEGIN_OBJECT) ||
4609                          (rk1 == WJB_BEGIN_OBJECT && (rk2 == WJB_BEGIN_ARRAY && !(*it2)->isScalar)))
4610         {
4611
4612                 JsonbIterator **it_array = rk1 == WJB_BEGIN_ARRAY ? it1 : it2;
4613                 JsonbIterator **it_object = rk1 == WJB_BEGIN_OBJECT ? it1 : it2;
4614
4615                 bool            prepend = (rk1 == WJB_BEGIN_OBJECT);
4616
4617                 pushJsonbValue(state, WJB_BEGIN_ARRAY, NULL);
4618
4619                 if (prepend)
4620                 {
4621                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
4622                         while ((r1 = JsonbIteratorNext(it_object, &v1, true)) != WJB_DONE)
4623                                 pushJsonbValue(state, r1, r1 != WJB_END_OBJECT ? &v1 : NULL);
4624
4625                         while ((r2 = JsonbIteratorNext(it_array, &v2, true)) != WJB_DONE)
4626                                 res = pushJsonbValue(state, r2, r2 != WJB_END_ARRAY ? &v2 : NULL);
4627                 }
4628                 else
4629                 {
4630                         while ((r1 = JsonbIteratorNext(it_array, &v1, true)) != WJB_END_ARRAY)
4631                                 pushJsonbValue(state, r1, &v1);
4632
4633                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
4634                         while ((r2 = JsonbIteratorNext(it_object, &v2, true)) != WJB_DONE)
4635                                 pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
4636
4637                         res = pushJsonbValue(state, WJB_END_ARRAY, NULL);
4638                 }
4639         }
4640         else
4641         {
4642                 /*
4643                  * This must be scalar || object or object || scalar, as that's all
4644                  * that's left. Both of these make no sense, so error out.
4645                  */
4646                 ereport(ERROR,
4647                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4648                                  errmsg("invalid concatenation of jsonb objects")));
4649         }
4650
4651         return res;
4652 }
4653
4654 /*
4655  * Do most of the heavy work for jsonb_set/jsonb_insert
4656  *
4657  * If JB_PATH_DELETE bit is set in op_type, the element is to be removed.
4658  *
4659  * If any bit mentioned in JB_PATH_CREATE_OR_INSERT is set in op_type,
4660  * we create the new value if the key or array index does not exist.
4661  *
4662  * Bits JB_PATH_INSERT_BEFORE and JB_PATH_INSERT_AFTER in op_type
4663  * behave as JB_PATH_CREATE if new value is inserted in JsonbObject.
4664  *
4665  * All path elements before the last must already exist
4666  * whatever bits in op_type are set, or nothing is done.
4667  */
4668 static JsonbValue *
4669 setPath(JsonbIterator **it, Datum *path_elems,
4670                 bool *path_nulls, int path_len,
4671                 JsonbParseState **st, int level, Jsonb *newval, int op_type)
4672 {
4673         JsonbValue      v;
4674         JsonbIteratorToken r;
4675         JsonbValue *res;
4676
4677         check_stack_depth();
4678
4679         if (path_nulls[level])
4680                 ereport(ERROR,
4681                                 (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
4682                                  errmsg("path element at position %d is null",
4683                                                 level + 1)));
4684
4685         r = JsonbIteratorNext(it, &v, false);
4686
4687         switch (r)
4688         {
4689                 case WJB_BEGIN_ARRAY:
4690                         (void) pushJsonbValue(st, r, NULL);
4691                         setPathArray(it, path_elems, path_nulls, path_len, st, level,
4692                                                  newval, v.val.array.nElems, op_type);
4693                         r = JsonbIteratorNext(it, &v, false);
4694                         Assert(r == WJB_END_ARRAY);
4695                         res = pushJsonbValue(st, r, NULL);
4696                         break;
4697                 case WJB_BEGIN_OBJECT:
4698                         (void) pushJsonbValue(st, r, NULL);
4699                         setPathObject(it, path_elems, path_nulls, path_len, st, level,
4700                                                   newval, v.val.object.nPairs, op_type);
4701                         r = JsonbIteratorNext(it, &v, true);
4702                         Assert(r == WJB_END_OBJECT);
4703                         res = pushJsonbValue(st, r, NULL);
4704                         break;
4705                 case WJB_ELEM:
4706                 case WJB_VALUE:
4707                         res = pushJsonbValue(st, r, &v);
4708                         break;
4709                 default:
4710                         elog(ERROR, "unrecognized iterator result: %d", (int) r);
4711                         res = NULL;                     /* keep compiler quiet */
4712                         break;
4713         }
4714
4715         return res;
4716 }
4717
4718 /*
4719  * Object walker for setPath
4720  */
4721 static void
4722 setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
4723                           int path_len, JsonbParseState **st, int level,
4724                           Jsonb *newval, uint32 npairs, int op_type)
4725 {
4726         JsonbValue      v;
4727         int                     i;
4728         JsonbValue      k;
4729         bool            done = false;
4730
4731         if (level >= path_len || path_nulls[level])
4732                 done = true;
4733
4734         /* empty object is a special case for create */
4735         if ((npairs == 0) && (op_type & JB_PATH_CREATE_OR_INSERT) &&
4736                 (level == path_len - 1))
4737         {
4738                 JsonbValue      newkey;
4739
4740                 newkey.type = jbvString;
4741                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
4742                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
4743
4744                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
4745                 addJsonbToParseState(st, newval);
4746         }
4747
4748         for (i = 0; i < npairs; i++)
4749         {
4750                 JsonbIteratorToken r = JsonbIteratorNext(it, &k, true);
4751
4752                 Assert(r == WJB_KEY);
4753
4754                 if (!done &&
4755                         k.val.string.len == VARSIZE_ANY_EXHDR(path_elems[level]) &&
4756                         memcmp(k.val.string.val, VARDATA_ANY(path_elems[level]),
4757                                    k.val.string.len) == 0)
4758                 {
4759                         if (level == path_len - 1)
4760                         {
4761                                 /*
4762                                  * called from jsonb_insert(), it forbids redefining an
4763                                  * existing value
4764                                  */
4765                                 if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER))
4766                                         ereport(ERROR,
4767                                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4768                                                          errmsg("cannot replace existing key"),
4769                                                          errhint("Try using the function jsonb_set "
4770                                                                          "to replace key value.")));
4771
4772                                 r = JsonbIteratorNext(it, &v, true);    /* skip value */
4773                                 if (!(op_type & JB_PATH_DELETE))
4774                                 {
4775                                         (void) pushJsonbValue(st, WJB_KEY, &k);
4776                                         addJsonbToParseState(st, newval);
4777                                 }
4778                                 done = true;
4779                         }
4780                         else
4781                         {
4782                                 (void) pushJsonbValue(st, r, &k);
4783                                 setPath(it, path_elems, path_nulls, path_len,
4784                                                 st, level + 1, newval, op_type);
4785                         }
4786                 }
4787                 else
4788                 {
4789                         if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
4790                                 level == path_len - 1 && i == npairs - 1)
4791                         {
4792                                 JsonbValue      newkey;
4793
4794                                 newkey.type = jbvString;
4795                                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
4796                                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
4797
4798                                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
4799                                 addJsonbToParseState(st, newval);
4800                         }
4801
4802                         (void) pushJsonbValue(st, r, &k);
4803                         r = JsonbIteratorNext(it, &v, false);
4804                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4805                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4806                         {
4807                                 int                     walking_level = 1;
4808
4809                                 while (walking_level != 0)
4810                                 {
4811                                         r = JsonbIteratorNext(it, &v, false);
4812
4813                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4814                                                 ++walking_level;
4815                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
4816                                                 --walking_level;
4817
4818                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4819                                 }
4820                         }
4821                 }
4822         }
4823 }
4824
4825 /*
4826  * Array walker for setPath
4827  */
4828 static void
4829 setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
4830                          int path_len, JsonbParseState **st, int level,
4831                          Jsonb *newval, uint32 nelems, int op_type)
4832 {
4833         JsonbValue      v;
4834         int                     idx,
4835                                 i;
4836         bool            done = false;
4837
4838         /* pick correct index */
4839         if (level < path_len && !path_nulls[level])
4840         {
4841                 char       *c = TextDatumGetCString(path_elems[level]);
4842                 long            lindex;
4843                 char       *badp;
4844
4845                 errno = 0;
4846                 lindex = strtol(c, &badp, 10);
4847                 if (errno != 0 || badp == c || *badp != '\0' || lindex > INT_MAX ||
4848                         lindex < INT_MIN)
4849                         ereport(ERROR,
4850                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
4851                                          errmsg("path element at position %d is not an integer: \"%s\"",
4852                                                         level + 1, c)));
4853                 idx = lindex;
4854         }
4855         else
4856                 idx = nelems;
4857
4858         if (idx < 0)
4859         {
4860                 if (-idx > nelems)
4861                         idx = INT_MIN;
4862                 else
4863                         idx = nelems + idx;
4864         }
4865
4866         if (idx > 0 && idx > nelems)
4867                 idx = nelems;
4868
4869         /*
4870          * if we're creating, and idx == INT_MIN, we prepend the new value to the
4871          * array also if the array is empty - in which case we don't really care
4872          * what the idx value is
4873          */
4874
4875         if ((idx == INT_MIN || nelems == 0) && (level == path_len - 1) &&
4876                 (op_type & JB_PATH_CREATE_OR_INSERT))
4877         {
4878                 Assert(newval != NULL);
4879                 addJsonbToParseState(st, newval);
4880                 done = true;
4881         }
4882
4883         /* iterate over the array elements */
4884         for (i = 0; i < nelems; i++)
4885         {
4886                 JsonbIteratorToken r;
4887
4888                 if (i == idx && level < path_len)
4889                 {
4890                         if (level == path_len - 1)
4891                         {
4892                                 r = JsonbIteratorNext(it, &v, true);    /* skip */
4893
4894                                 if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_CREATE))
4895                                         addJsonbToParseState(st, newval);
4896
4897                                 /*
4898                                  * We should keep current value only in case of
4899                                  * JB_PATH_INSERT_BEFORE or JB_PATH_INSERT_AFTER because
4900                                  * otherwise it should be deleted or replaced
4901                                  */
4902                                 if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_INSERT_BEFORE))
4903                                         (void) pushJsonbValue(st, r, &v);
4904
4905                                 if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_REPLACE))
4906                                         addJsonbToParseState(st, newval);
4907
4908                                 done = true;
4909                         }
4910                         else
4911                                 (void) setPath(it, path_elems, path_nulls, path_len,
4912                                                            st, level + 1, newval, op_type);
4913                 }
4914                 else
4915                 {
4916                         r = JsonbIteratorNext(it, &v, false);
4917
4918                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4919
4920                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4921                         {
4922                                 int                     walking_level = 1;
4923
4924                                 while (walking_level != 0)
4925                                 {
4926                                         r = JsonbIteratorNext(it, &v, false);
4927
4928                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4929                                                 ++walking_level;
4930                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
4931                                                 --walking_level;
4932
4933                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4934                                 }
4935                         }
4936
4937                         if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
4938                                 level == path_len - 1 && i == nelems - 1)
4939                         {
4940                                 addJsonbToParseState(st, newval);
4941                         }
4942                 }
4943         }
4944 }
4945
4946 /*
4947  * Parse information about what elements of a jsonb document we want to iterate
4948  * in functions iterate_json(b)_values. This information is presented in jsonb
4949  * format, so that it can be easily extended in the future.
4950  */
4951 uint32
4952 parse_jsonb_index_flags(Jsonb *jb)
4953 {
4954         JsonbIterator *it;
4955         JsonbValue      v;
4956         JsonbIteratorToken type;
4957         uint32          flags = 0;
4958
4959         it = JsonbIteratorInit(&jb->root);
4960
4961         type = JsonbIteratorNext(&it, &v, false);
4962
4963         /*
4964          * We iterate over array (scalar internally is represented as array, so,
4965          * we will accept it too) to check all its elements.  Flag names are
4966          * chosen the same as jsonb_typeof uses.
4967          */
4968         if (type != WJB_BEGIN_ARRAY)
4969                 ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4970                                                 errmsg("wrong flag type, only arrays and scalars are allowed")));
4971
4972         while ((type = JsonbIteratorNext(&it, &v, false)) == WJB_ELEM)
4973         {
4974                 if (v.type != jbvString)
4975                         ereport(ERROR,
4976                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4977                                          errmsg("flag array element is not a string"),
4978                                          errhint("Possible values are: \"string\", \"numeric\", \"boolean\", \"key\" and \"all\"")));
4979
4980                 if (v.val.string.len == 3 &&
4981                         pg_strncasecmp(v.val.string.val, "all", 3) == 0)
4982                         flags |= jtiAll;
4983                 else if (v.val.string.len == 3 &&
4984                                  pg_strncasecmp(v.val.string.val, "key", 3) == 0)
4985                         flags |= jtiKey;
4986                 else if (v.val.string.len == 6 &&
4987                                  pg_strncasecmp(v.val.string.val, "string", 5) == 0)
4988                         flags |= jtiString;
4989                 else if (v.val.string.len == 7 &&
4990                                  pg_strncasecmp(v.val.string.val, "numeric", 7) == 0)
4991                         flags |= jtiNumeric;
4992                 else if (v.val.string.len == 7 &&
4993                                  pg_strncasecmp(v.val.string.val, "boolean", 7) == 0)
4994                         flags |= jtiBool;
4995                 else
4996                         ereport(ERROR,
4997                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4998                                          errmsg("wrong flag in flag array: \"%s\"",
4999                                                         pnstrdup(v.val.string.val, v.val.string.len)),
5000                                          errhint("Possible values are: \"string\", \"numeric\", \"boolean\", \"key\" and \"all\"")));
5001         }
5002
5003         /* expect end of array now */
5004         if (type != WJB_END_ARRAY)
5005                 elog(ERROR, "unexpected end of flag array");
5006
5007         /* get final WJB_DONE and free iterator */
5008         type = JsonbIteratorNext(&it, &v, false);
5009         if (type != WJB_DONE)
5010                 elog(ERROR, "unexpected end of flag array");
5011
5012         return flags;
5013 }
5014
5015 /*
5016  * Iterate over jsonb values or elements, specified by flags, and pass them
5017  * together with an iteration state to a specified JsonIterateStringValuesAction.
5018  */
5019 void
5020 iterate_jsonb_values(Jsonb *jb, uint32 flags, void *state,
5021                                          JsonIterateStringValuesAction action)
5022 {
5023         JsonbIterator *it;
5024         JsonbValue      v;
5025         JsonbIteratorToken type;
5026
5027         it = JsonbIteratorInit(&jb->root);
5028
5029         /*
5030          * Just recursively iterating over jsonb and call callback on all
5031          * correspoding elements
5032          */
5033         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
5034         {
5035                 if (type == WJB_KEY)
5036                 {
5037                         if (flags & jtiKey)
5038                                 action(state, v.val.string.val, v.val.string.len);
5039
5040                         continue;
5041                 }
5042                 else if (!(type == WJB_VALUE || type == WJB_ELEM))
5043                 {
5044                         /* do not call callback for composite JsonbValue */
5045                         continue;
5046                 }
5047
5048                 /* JsonbValue is a value of object or element of array */
5049                 switch (v.type)
5050                 {
5051                         case jbvString:
5052                                 if (flags & jtiString)
5053                                         action(state, v.val.string.val, v.val.string.len);
5054                                 break;
5055                         case jbvNumeric:
5056                                 if (flags & jtiNumeric)
5057                                 {
5058                                         char       *val;
5059
5060                                         val = DatumGetCString(DirectFunctionCall1(numeric_out,
5061                                                                                                                           NumericGetDatum(v.val.numeric)));
5062
5063                                         action(state, val, strlen(val));
5064                                         pfree(val);
5065                                 }
5066                                 break;
5067                         case jbvBool:
5068                                 if (flags & jtiBool)
5069                                 {
5070                                         if (v.val.boolean)
5071                                                 action(state, "true", 4);
5072                                         else
5073                                                 action(state, "false", 5);
5074                                 }
5075                                 break;
5076                         default:
5077                                 /* do not call callback for composite JsonbValue */
5078                                 break;
5079                 }
5080         }
5081 }
5082
5083 /*
5084  * Iterate over json values and elements, specified by flags, and pass them
5085  * together with an iteration state to a specified JsonIterateStringValuesAction.
5086  */
5087 void
5088 iterate_json_values(text *json, uint32 flags, void *action_state,
5089                                         JsonIterateStringValuesAction action)
5090 {
5091         JsonLexContext *lex = makeJsonLexContext(json, true);
5092         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
5093         IterateJsonStringValuesState *state = palloc0(sizeof(IterateJsonStringValuesState));
5094
5095         state->lex = lex;
5096         state->action = action;
5097         state->action_state = action_state;
5098         state->flags = flags;
5099
5100         sem->semstate = (void *) state;
5101         sem->scalar = iterate_values_scalar;
5102         sem->object_field_start = iterate_values_object_field_start;
5103
5104         pg_parse_json(lex, sem);
5105 }
5106
5107 /*
5108  * An auxiliary function for iterate_json_values to invoke a specified
5109  * JsonIterateStringValuesAction for specified values.
5110  */
5111 static void
5112 iterate_values_scalar(void *state, char *token, JsonTokenType tokentype)
5113 {
5114         IterateJsonStringValuesState *_state = (IterateJsonStringValuesState *) state;
5115
5116         switch (tokentype)
5117         {
5118                 case JSON_TOKEN_STRING:
5119                         if (_state->flags & jtiString)
5120                                 _state->action(_state->action_state, token, strlen(token));
5121                         break;
5122                 case JSON_TOKEN_NUMBER:
5123                         if (_state->flags & jtiNumeric)
5124                                 _state->action(_state->action_state, token, strlen(token));
5125                         break;
5126                 case JSON_TOKEN_TRUE:
5127                 case JSON_TOKEN_FALSE:
5128                         if (_state->flags & jtiBool)
5129                                 _state->action(_state->action_state, token, strlen(token));
5130                         break;
5131                 default:
5132                         /* do not call callback for any other token */
5133                         break;
5134         }
5135 }
5136
5137 static void
5138 iterate_values_object_field_start(void *state, char *fname, bool isnull)
5139 {
5140         IterateJsonStringValuesState *_state = (IterateJsonStringValuesState *) state;
5141
5142         if (_state->flags & jtiKey)
5143         {
5144                 char       *val = pstrdup(fname);
5145
5146                 _state->action(_state->action_state, val, strlen(val));
5147         }
5148 }
5149
5150 /*
5151  * Iterate over a jsonb, and apply a specified JsonTransformStringValuesAction
5152  * to every string value or element. Any necessary context for a
5153  * JsonTransformStringValuesAction can be passed in the action_state variable.
5154  * Function returns a copy of an original jsonb object with transformed values.
5155  */
5156 Jsonb *
5157 transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
5158                                                           JsonTransformStringValuesAction transform_action)
5159 {
5160         JsonbIterator *it;
5161         JsonbValue      v,
5162                            *res = NULL;
5163         JsonbIteratorToken type;
5164         JsonbParseState *st = NULL;
5165         text       *out;
5166         bool            is_scalar = false;
5167
5168         it = JsonbIteratorInit(&jsonb->root);
5169         is_scalar = it->isScalar;
5170
5171         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
5172         {
5173                 if ((type == WJB_VALUE || type == WJB_ELEM) && v.type == jbvString)
5174                 {
5175                         out = transform_action(action_state, v.val.string.val, v.val.string.len);
5176                         v.val.string.val = VARDATA_ANY(out);
5177                         v.val.string.len = VARSIZE_ANY_EXHDR(out);
5178                         res = pushJsonbValue(&st, type, type < WJB_BEGIN_ARRAY ? &v : NULL);
5179                 }
5180                 else
5181                 {
5182                         res = pushJsonbValue(&st, type, (type == WJB_KEY ||
5183                                                                                          type == WJB_VALUE ||
5184                                                                                          type == WJB_ELEM) ? &v : NULL);
5185                 }
5186         }
5187
5188         if (res->type == jbvArray)
5189                 res->val.array.rawScalar = is_scalar;
5190
5191         return JsonbValueToJsonb(res);
5192 }
5193
5194 /*
5195  * Iterate over a json, and apply a specified JsonTransformStringValuesAction
5196  * to every string value or element. Any necessary context for a
5197  * JsonTransformStringValuesAction can be passed in the action_state variable.
5198  * Function returns a StringInfo, which is a copy of an original json with
5199  * transformed values.
5200  */
5201 text *
5202 transform_json_string_values(text *json, void *action_state,
5203                                                          JsonTransformStringValuesAction transform_action)
5204 {
5205         JsonLexContext *lex = makeJsonLexContext(json, true);
5206         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
5207         TransformJsonStringValuesState *state = palloc0(sizeof(TransformJsonStringValuesState));
5208
5209         state->lex = lex;
5210         state->strval = makeStringInfo();
5211         state->action = transform_action;
5212         state->action_state = action_state;
5213
5214         sem->semstate = (void *) state;
5215         sem->scalar = transform_string_values_scalar;
5216         sem->object_start = transform_string_values_object_start;
5217         sem->object_end = transform_string_values_object_end;
5218         sem->array_start = transform_string_values_array_start;
5219         sem->array_end = transform_string_values_array_end;
5220         sem->scalar = transform_string_values_scalar;
5221         sem->array_element_start = transform_string_values_array_element_start;
5222         sem->object_field_start = transform_string_values_object_field_start;
5223
5224         pg_parse_json(lex, sem);
5225
5226         return cstring_to_text_with_len(state->strval->data, state->strval->len);
5227 }
5228
5229 /*
5230  * Set of auxiliary functions for transform_json_string_values to invoke a
5231  * specified JsonTransformStringValuesAction for all values and left everything
5232  * else untouched.
5233  */
5234 static void
5235 transform_string_values_object_start(void *state)
5236 {
5237         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5238
5239         appendStringInfoCharMacro(_state->strval, '{');
5240 }
5241
5242 static void
5243 transform_string_values_object_end(void *state)
5244 {
5245         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5246
5247         appendStringInfoCharMacro(_state->strval, '}');
5248 }
5249
5250 static void
5251 transform_string_values_array_start(void *state)
5252 {
5253         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5254
5255         appendStringInfoCharMacro(_state->strval, '[');
5256 }
5257
5258 static void
5259 transform_string_values_array_end(void *state)
5260 {
5261         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5262
5263         appendStringInfoCharMacro(_state->strval, ']');
5264 }
5265
5266 static void
5267 transform_string_values_object_field_start(void *state, char *fname, bool isnull)
5268 {
5269         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5270
5271         if (_state->strval->data[_state->strval->len - 1] != '{')
5272                 appendStringInfoCharMacro(_state->strval, ',');
5273
5274         /*
5275          * Unfortunately we don't have the quoted and escaped string any more, so
5276          * we have to re-escape it.
5277          */
5278         escape_json(_state->strval, fname);
5279         appendStringInfoCharMacro(_state->strval, ':');
5280 }
5281
5282 static void
5283 transform_string_values_array_element_start(void *state, bool isnull)
5284 {
5285         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5286
5287         if (_state->strval->data[_state->strval->len - 1] != '[')
5288                 appendStringInfoCharMacro(_state->strval, ',');
5289 }
5290
5291 static void
5292 transform_string_values_scalar(void *state, char *token, JsonTokenType tokentype)
5293 {
5294         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5295
5296         if (tokentype == JSON_TOKEN_STRING)
5297         {
5298                 text       *out = _state->action(_state->action_state, token, strlen(token));
5299
5300                 escape_json(_state->strval, text_to_cstring(out));
5301         }
5302         else
5303                 appendStringInfoString(_state->strval, token);
5304 }