]> granicus.if.org Git - postgresql/blob - src/backend/utils/adt/jsonfuncs.c
Fix conversion of JSON strings to JSON output columns in json_to_record().
[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-2019, 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 is_json, bool have_record_arg);
425 static Datum populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
426                                                                         bool is_json, 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",
2300                                                                   false, true);
2301 }
2302
2303 Datum
2304 jsonb_to_record(PG_FUNCTION_ARGS)
2305 {
2306         return populate_record_worker(fcinfo, "jsonb_to_record",
2307                                                                   false, false);
2308 }
2309
2310 Datum
2311 json_populate_record(PG_FUNCTION_ARGS)
2312 {
2313         return populate_record_worker(fcinfo, "json_populate_record",
2314                                                                   true, true);
2315 }
2316
2317 Datum
2318 json_to_record(PG_FUNCTION_ARGS)
2319 {
2320         return populate_record_worker(fcinfo, "json_to_record",
2321                                                                   true, false);
2322 }
2323
2324 /* helper function for diagnostics */
2325 static void
2326 populate_array_report_expected_array(PopulateArrayContext *ctx, int ndim)
2327 {
2328         if (ndim <= 0)
2329         {
2330                 if (ctx->colname)
2331                         ereport(ERROR,
2332                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2333                                          errmsg("expected JSON array"),
2334                                          errhint("See the value of key \"%s\".", ctx->colname)));
2335                 else
2336                         ereport(ERROR,
2337                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2338                                          errmsg("expected JSON array")));
2339         }
2340         else
2341         {
2342                 StringInfoData indices;
2343                 int                     i;
2344
2345                 initStringInfo(&indices);
2346
2347                 Assert(ctx->ndims > 0 && ndim < ctx->ndims);
2348
2349                 for (i = 0; i < ndim; i++)
2350                         appendStringInfo(&indices, "[%d]", ctx->sizes[i]);
2351
2352                 if (ctx->colname)
2353                         ereport(ERROR,
2354                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2355                                          errmsg("expected JSON array"),
2356                                          errhint("See the array element %s of key \"%s\".",
2357                                                          indices.data, ctx->colname)));
2358                 else
2359                         ereport(ERROR,
2360                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2361                                          errmsg("expected JSON array"),
2362                                          errhint("See the array element %s.",
2363                                                          indices.data)));
2364         }
2365 }
2366
2367 /* set the number of dimensions of the populated array when it becomes known */
2368 static void
2369 populate_array_assign_ndims(PopulateArrayContext *ctx, int ndims)
2370 {
2371         int                     i;
2372
2373         Assert(ctx->ndims <= 0);
2374
2375         if (ndims <= 0)
2376                 populate_array_report_expected_array(ctx, ndims);
2377
2378         ctx->ndims = ndims;
2379         ctx->dims = palloc(sizeof(int) * ndims);
2380         ctx->sizes = palloc0(sizeof(int) * ndims);
2381
2382         for (i = 0; i < ndims; i++)
2383                 ctx->dims[i] = -1;              /* dimensions are unknown yet */
2384 }
2385
2386 /* check the populated subarray dimension */
2387 static void
2388 populate_array_check_dimension(PopulateArrayContext *ctx, int ndim)
2389 {
2390         int                     dim = ctx->sizes[ndim]; /* current dimension counter */
2391
2392         if (ctx->dims[ndim] == -1)
2393                 ctx->dims[ndim] = dim;  /* assign dimension if not yet known */
2394         else if (ctx->dims[ndim] != dim)
2395                 ereport(ERROR,
2396                                 (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
2397                                  errmsg("malformed JSON array"),
2398                                  errdetail("Multidimensional arrays must have "
2399                                                    "sub-arrays with matching dimensions.")));
2400
2401         /* reset the current array dimension size counter */
2402         ctx->sizes[ndim] = 0;
2403
2404         /* increment the parent dimension counter if it is a nested sub-array */
2405         if (ndim > 0)
2406                 ctx->sizes[ndim - 1]++;
2407 }
2408
2409 static void
2410 populate_array_element(PopulateArrayContext *ctx, int ndim, JsValue *jsv)
2411 {
2412         Datum           element;
2413         bool            element_isnull;
2414
2415         /* populate the array element */
2416         element = populate_record_field(ctx->aio->element_info,
2417                                                                         ctx->aio->element_type,
2418                                                                         ctx->aio->element_typmod,
2419                                                                         NULL, ctx->mcxt, PointerGetDatum(NULL),
2420                                                                         jsv, &element_isnull);
2421
2422         accumArrayResult(ctx->astate, element, element_isnull,
2423                                          ctx->aio->element_type, ctx->acxt);
2424
2425         Assert(ndim > 0);
2426         ctx->sizes[ndim - 1]++;         /* increment current dimension counter */
2427 }
2428
2429 /* json object start handler for populate_array_json() */
2430 static void
2431 populate_array_object_start(void *_state)
2432 {
2433         PopulateArrayState *state = (PopulateArrayState *) _state;
2434         int                     ndim = state->lex->lex_level;
2435
2436         if (state->ctx->ndims <= 0)
2437                 populate_array_assign_ndims(state->ctx, ndim);
2438         else if (ndim < state->ctx->ndims)
2439                 populate_array_report_expected_array(state->ctx, ndim);
2440 }
2441
2442 /* json array end handler for populate_array_json() */
2443 static void
2444 populate_array_array_end(void *_state)
2445 {
2446         PopulateArrayState *state = (PopulateArrayState *) _state;
2447         PopulateArrayContext *ctx = state->ctx;
2448         int                     ndim = state->lex->lex_level;
2449
2450         if (ctx->ndims <= 0)
2451                 populate_array_assign_ndims(ctx, ndim + 1);
2452
2453         if (ndim < ctx->ndims)
2454                 populate_array_check_dimension(ctx, ndim);
2455 }
2456
2457 /* json array element start handler for populate_array_json() */
2458 static void
2459 populate_array_element_start(void *_state, bool isnull)
2460 {
2461         PopulateArrayState *state = (PopulateArrayState *) _state;
2462         int                     ndim = state->lex->lex_level;
2463
2464         if (state->ctx->ndims <= 0 || ndim == state->ctx->ndims)
2465         {
2466                 /* remember current array element start */
2467                 state->element_start = state->lex->token_start;
2468                 state->element_type = state->lex->token_type;
2469                 state->element_scalar = NULL;
2470         }
2471 }
2472
2473 /* json array element end handler for populate_array_json() */
2474 static void
2475 populate_array_element_end(void *_state, bool isnull)
2476 {
2477         PopulateArrayState *state = (PopulateArrayState *) _state;
2478         PopulateArrayContext *ctx = state->ctx;
2479         int                     ndim = state->lex->lex_level;
2480
2481         Assert(ctx->ndims > 0);
2482
2483         if (ndim == ctx->ndims)
2484         {
2485                 JsValue         jsv;
2486
2487                 jsv.is_json = true;
2488                 jsv.val.json.type = state->element_type;
2489
2490                 if (isnull)
2491                 {
2492                         Assert(jsv.val.json.type == JSON_TOKEN_NULL);
2493                         jsv.val.json.str = NULL;
2494                         jsv.val.json.len = 0;
2495                 }
2496                 else if (state->element_scalar)
2497                 {
2498                         jsv.val.json.str = state->element_scalar;
2499                         jsv.val.json.len = -1;  /* null-terminated */
2500                 }
2501                 else
2502                 {
2503                         jsv.val.json.str = state->element_start;
2504                         jsv.val.json.len = (state->lex->prev_token_terminator -
2505                                                                 state->element_start) * sizeof(char);
2506                 }
2507
2508                 populate_array_element(ctx, ndim, &jsv);
2509         }
2510 }
2511
2512 /* json scalar handler for populate_array_json() */
2513 static void
2514 populate_array_scalar(void *_state, char *token, JsonTokenType tokentype)
2515 {
2516         PopulateArrayState *state = (PopulateArrayState *) _state;
2517         PopulateArrayContext *ctx = state->ctx;
2518         int                     ndim = state->lex->lex_level;
2519
2520         if (ctx->ndims <= 0)
2521                 populate_array_assign_ndims(ctx, ndim);
2522         else if (ndim < ctx->ndims)
2523                 populate_array_report_expected_array(ctx, ndim);
2524
2525         if (ndim == ctx->ndims)
2526         {
2527                 /* remember the scalar element token */
2528                 state->element_scalar = token;
2529                 /* element_type must already be set in populate_array_element_start() */
2530                 Assert(state->element_type == tokentype);
2531         }
2532 }
2533
2534 /* parse a json array and populate array */
2535 static void
2536 populate_array_json(PopulateArrayContext *ctx, char *json, int len)
2537 {
2538         PopulateArrayState state;
2539         JsonSemAction sem;
2540
2541         state.lex = makeJsonLexContextCstringLen(json, len, true);
2542         state.ctx = ctx;
2543
2544         memset(&sem, 0, sizeof(sem));
2545         sem.semstate = (void *) &state;
2546         sem.object_start = populate_array_object_start;
2547         sem.array_end = populate_array_array_end;
2548         sem.array_element_start = populate_array_element_start;
2549         sem.array_element_end = populate_array_element_end;
2550         sem.scalar = populate_array_scalar;
2551
2552         pg_parse_json(state.lex, &sem);
2553
2554         /* number of dimensions should be already known */
2555         Assert(ctx->ndims > 0 && ctx->dims);
2556
2557         pfree(state.lex);
2558 }
2559
2560 /*
2561  * populate_array_dim_jsonb() -- Iterate recursively through jsonb sub-array
2562  *              elements and accumulate result using given ArrayBuildState.
2563  */
2564 static void
2565 populate_array_dim_jsonb(PopulateArrayContext *ctx, /* context */
2566                                                  JsonbValue *jbv,       /* jsonb sub-array */
2567                                                  int ndim)      /* current dimension */
2568 {
2569         JsonbContainer *jbc = jbv->val.binary.data;
2570         JsonbIterator *it;
2571         JsonbIteratorToken tok;
2572         JsonbValue      val;
2573         JsValue         jsv;
2574
2575         check_stack_depth();
2576
2577         if (jbv->type != jbvBinary || !JsonContainerIsArray(jbc))
2578                 populate_array_report_expected_array(ctx, ndim - 1);
2579
2580         Assert(!JsonContainerIsScalar(jbc));
2581
2582         it = JsonbIteratorInit(jbc);
2583
2584         tok = JsonbIteratorNext(&it, &val, true);
2585         Assert(tok == WJB_BEGIN_ARRAY);
2586
2587         tok = JsonbIteratorNext(&it, &val, true);
2588
2589         /*
2590          * If the number of dimensions is not yet known and we have found end of
2591          * the array, or the first child element is not an array, then assign the
2592          * number of dimensions now.
2593          */
2594         if (ctx->ndims <= 0 &&
2595                 (tok == WJB_END_ARRAY ||
2596                  (tok == WJB_ELEM &&
2597                   (val.type != jbvBinary ||
2598                    !JsonContainerIsArray(val.val.binary.data)))))
2599                 populate_array_assign_ndims(ctx, ndim);
2600
2601         jsv.is_json = false;
2602         jsv.val.jsonb = &val;
2603
2604         /* process all the array elements */
2605         while (tok == WJB_ELEM)
2606         {
2607                 /*
2608                  * Recurse only if the dimensions of dimensions is still unknown or if
2609                  * it is not the innermost dimension.
2610                  */
2611                 if (ctx->ndims > 0 && ndim >= ctx->ndims)
2612                         populate_array_element(ctx, ndim, &jsv);
2613                 else
2614                 {
2615                         /* populate child sub-array */
2616                         populate_array_dim_jsonb(ctx, &val, ndim + 1);
2617
2618                         /* number of dimensions should be already known */
2619                         Assert(ctx->ndims > 0 && ctx->dims);
2620
2621                         populate_array_check_dimension(ctx, ndim);
2622                 }
2623
2624                 tok = JsonbIteratorNext(&it, &val, true);
2625         }
2626
2627         Assert(tok == WJB_END_ARRAY);
2628
2629         /* free iterator, iterating until WJB_DONE */
2630         tok = JsonbIteratorNext(&it, &val, true);
2631         Assert(tok == WJB_DONE && !it);
2632 }
2633
2634 /* recursively populate an array from json/jsonb */
2635 static Datum
2636 populate_array(ArrayIOData *aio,
2637                            const char *colname,
2638                            MemoryContext mcxt,
2639                            JsValue *jsv)
2640 {
2641         PopulateArrayContext ctx;
2642         Datum           result;
2643         int                *lbs;
2644         int                     i;
2645
2646         ctx.aio = aio;
2647         ctx.mcxt = mcxt;
2648         ctx.acxt = CurrentMemoryContext;
2649         ctx.astate = initArrayResult(aio->element_type, ctx.acxt, true);
2650         ctx.colname = colname;
2651         ctx.ndims = 0;                          /* unknown yet */
2652         ctx.dims = NULL;
2653         ctx.sizes = NULL;
2654
2655         if (jsv->is_json)
2656                 populate_array_json(&ctx, jsv->val.json.str,
2657                                                         jsv->val.json.len >= 0 ? jsv->val.json.len
2658                                                         : strlen(jsv->val.json.str));
2659         else
2660         {
2661                 populate_array_dim_jsonb(&ctx, jsv->val.jsonb, 1);
2662                 ctx.dims[0] = ctx.sizes[0];
2663         }
2664
2665         Assert(ctx.ndims > 0);
2666
2667         lbs = palloc(sizeof(int) * ctx.ndims);
2668
2669         for (i = 0; i < ctx.ndims; i++)
2670                 lbs[i] = 1;
2671
2672         result = makeMdArrayResult(ctx.astate, ctx.ndims, ctx.dims, lbs,
2673                                                            ctx.acxt, true);
2674
2675         pfree(ctx.dims);
2676         pfree(ctx.sizes);
2677         pfree(lbs);
2678
2679         return result;
2680 }
2681
2682 static void
2683 JsValueToJsObject(JsValue *jsv, JsObject *jso)
2684 {
2685         jso->is_json = jsv->is_json;
2686
2687         if (jsv->is_json)
2688         {
2689                 /* convert plain-text json into a hash table */
2690                 jso->val.json_hash =
2691                         get_json_object_as_hash(jsv->val.json.str,
2692                                                                         jsv->val.json.len >= 0
2693                                                                         ? jsv->val.json.len
2694                                                                         : strlen(jsv->val.json.str),
2695                                                                         "populate_composite");
2696         }
2697         else
2698         {
2699                 JsonbValue *jbv = jsv->val.jsonb;
2700
2701                 if (jbv->type == jbvBinary &&
2702                         JsonContainerIsObject(jbv->val.binary.data))
2703                 {
2704                         jso->val.jsonb_cont = jbv->val.binary.data;
2705                 }
2706                 else
2707                 {
2708                         bool            is_scalar;
2709
2710                         is_scalar = IsAJsonbScalar(jbv) ||
2711                                 (jbv->type == jbvBinary &&
2712                                  JsonContainerIsScalar(jbv->val.binary.data));
2713                         ereport(ERROR,
2714                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2715                                          is_scalar
2716                                          ? errmsg("cannot call %s on a scalar",
2717                                                           "populate_composite")
2718                                          : errmsg("cannot call %s on an array",
2719                                                           "populate_composite")));
2720                 }
2721         }
2722 }
2723
2724 /* acquire or update cached tuple descriptor for a composite type */
2725 static void
2726 update_cached_tupdesc(CompositeIOData *io, MemoryContext mcxt)
2727 {
2728         if (!io->tupdesc ||
2729                 io->tupdesc->tdtypeid != io->base_typid ||
2730                 io->tupdesc->tdtypmod != io->base_typmod)
2731         {
2732                 TupleDesc       tupdesc = lookup_rowtype_tupdesc(io->base_typid,
2733                                                                                                          io->base_typmod);
2734                 MemoryContext oldcxt;
2735
2736                 if (io->tupdesc)
2737                         FreeTupleDesc(io->tupdesc);
2738
2739                 /* copy tuple desc without constraints into cache memory context */
2740                 oldcxt = MemoryContextSwitchTo(mcxt);
2741                 io->tupdesc = CreateTupleDescCopy(tupdesc);
2742                 MemoryContextSwitchTo(oldcxt);
2743
2744                 ReleaseTupleDesc(tupdesc);
2745         }
2746 }
2747
2748 /* recursively populate a composite (row type) value from json/jsonb */
2749 static Datum
2750 populate_composite(CompositeIOData *io,
2751                                    Oid typid,
2752                                    const char *colname,
2753                                    MemoryContext mcxt,
2754                                    HeapTupleHeader defaultval,
2755                                    JsValue *jsv,
2756                                    bool isnull)
2757 {
2758         Datum           result;
2759
2760         /* acquire/update cached tuple descriptor */
2761         update_cached_tupdesc(io, mcxt);
2762
2763         if (isnull)
2764                 result = (Datum) 0;
2765         else
2766         {
2767                 HeapTupleHeader tuple;
2768                 JsObject        jso;
2769
2770                 /* prepare input value */
2771                 JsValueToJsObject(jsv, &jso);
2772
2773                 /* populate resulting record tuple */
2774                 tuple = populate_record(io->tupdesc, &io->record_io,
2775                                                                 defaultval, mcxt, &jso);
2776                 result = HeapTupleHeaderGetDatum(tuple);
2777
2778                 JsObjectFree(&jso);
2779         }
2780
2781         /*
2782          * If it's domain over composite, check domain constraints.  (This should
2783          * probably get refactored so that we can see the TYPECAT value, but for
2784          * now, we can tell by comparing typid to base_typid.)
2785          */
2786         if (typid != io->base_typid && typid != RECORDOID)
2787                 domain_check(result, isnull, typid, &io->domain_info, mcxt);
2788
2789         return result;
2790 }
2791
2792 /* populate non-null scalar value from json/jsonb value */
2793 static Datum
2794 populate_scalar(ScalarIOData *io, Oid typid, int32 typmod, JsValue *jsv)
2795 {
2796         Datum           res;
2797         char       *str = NULL;
2798         char       *json = NULL;
2799
2800         if (jsv->is_json)
2801         {
2802                 int                     len = jsv->val.json.len;
2803
2804                 json = jsv->val.json.str;
2805                 Assert(json);
2806                 if (len >= 0)
2807                 {
2808                         /* Need to copy non-null-terminated string */
2809                         str = palloc(len + 1 * sizeof(char));
2810                         memcpy(str, json, len);
2811                         str[len] = '\0';
2812                 }
2813                 else
2814                         str = json;                     /* string is already null-terminated */
2815
2816                 /* If converting to json/jsonb, make string into valid JSON literal */
2817                 if ((typid == JSONOID || typid == JSONBOID) &&
2818                         jsv->val.json.type == JSON_TOKEN_STRING)
2819                 {
2820                         StringInfoData buf;
2821
2822                         initStringInfo(&buf);
2823                         escape_json(&buf, str);
2824                         /* free temporary buffer */
2825                         if (str != json)
2826                                 pfree(str);
2827                         str = buf.data;
2828                 }
2829         }
2830         else
2831         {
2832                 JsonbValue *jbv = jsv->val.jsonb;
2833
2834                 if (typid == JSONBOID)
2835                 {
2836                         Jsonb      *jsonb = JsonbValueToJsonb(jbv); /* directly use jsonb */
2837
2838                         return JsonbPGetDatum(jsonb);
2839                 }
2840                 /* convert jsonb to string for typio call */
2841                 else if (typid == JSONOID && jbv->type != jbvBinary)
2842                 {
2843                         /*
2844                          * Convert scalar jsonb (non-scalars are passed here as jbvBinary)
2845                          * to json string, preserving quotes around top-level strings.
2846                          */
2847                         Jsonb      *jsonb = JsonbValueToJsonb(jbv);
2848
2849                         str = JsonbToCString(NULL, &jsonb->root, VARSIZE(jsonb));
2850                 }
2851                 else if (jbv->type == jbvString)        /* quotes are stripped */
2852                         str = pnstrdup(jbv->val.string.val, jbv->val.string.len);
2853                 else if (jbv->type == jbvBool)
2854                         str = pstrdup(jbv->val.boolean ? "true" : "false");
2855                 else if (jbv->type == jbvNumeric)
2856                         str = DatumGetCString(DirectFunctionCall1(numeric_out,
2857                                                                                                           PointerGetDatum(jbv->val.numeric)));
2858                 else if (jbv->type == jbvBinary)
2859                         str = JsonbToCString(NULL, jbv->val.binary.data,
2860                                                                  jbv->val.binary.len);
2861                 else
2862                         elog(ERROR, "unrecognized jsonb type: %d", (int) jbv->type);
2863         }
2864
2865         res = InputFunctionCall(&io->typiofunc, str, io->typioparam, typmod);
2866
2867         /* free temporary buffer */
2868         if (str != json)
2869                 pfree(str);
2870
2871         return res;
2872 }
2873
2874 static Datum
2875 populate_domain(DomainIOData *io,
2876                                 Oid typid,
2877                                 const char *colname,
2878                                 MemoryContext mcxt,
2879                                 JsValue *jsv,
2880                                 bool isnull)
2881 {
2882         Datum           res;
2883
2884         if (isnull)
2885                 res = (Datum) 0;
2886         else
2887         {
2888                 res = populate_record_field(io->base_io,
2889                                                                         io->base_typid, io->base_typmod,
2890                                                                         colname, mcxt, PointerGetDatum(NULL),
2891                                                                         jsv, &isnull);
2892                 Assert(!isnull);
2893         }
2894
2895         domain_check(res, isnull, typid, &io->domain_info, mcxt);
2896
2897         return res;
2898 }
2899
2900 /* prepare column metadata cache for the given type */
2901 static void
2902 prepare_column_cache(ColumnIOData *column,
2903                                          Oid typid,
2904                                          int32 typmod,
2905                                          MemoryContext mcxt,
2906                                          bool need_scalar)
2907 {
2908         HeapTuple       tup;
2909         Form_pg_type type;
2910
2911         column->typid = typid;
2912         column->typmod = typmod;
2913
2914         tup = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typid));
2915         if (!HeapTupleIsValid(tup))
2916                 elog(ERROR, "cache lookup failed for type %u", typid);
2917
2918         type = (Form_pg_type) GETSTRUCT(tup);
2919
2920         if (type->typtype == TYPTYPE_DOMAIN)
2921         {
2922                 /*
2923                  * We can move directly to the bottom base type; domain_check() will
2924                  * take care of checking all constraints for a stack of domains.
2925                  */
2926                 Oid                     base_typid;
2927                 int32           base_typmod = typmod;
2928
2929                 base_typid = getBaseTypeAndTypmod(typid, &base_typmod);
2930                 if (get_typtype(base_typid) == TYPTYPE_COMPOSITE)
2931                 {
2932                         /* domain over composite has its own code path */
2933                         column->typcat = TYPECAT_COMPOSITE_DOMAIN;
2934                         column->io.composite.record_io = NULL;
2935                         column->io.composite.tupdesc = NULL;
2936                         column->io.composite.base_typid = base_typid;
2937                         column->io.composite.base_typmod = base_typmod;
2938                         column->io.composite.domain_info = NULL;
2939                 }
2940                 else
2941                 {
2942                         /* domain over anything else */
2943                         column->typcat = TYPECAT_DOMAIN;
2944                         column->io.domain.base_typid = base_typid;
2945                         column->io.domain.base_typmod = base_typmod;
2946                         column->io.domain.base_io =
2947                                 MemoryContextAllocZero(mcxt, sizeof(ColumnIOData));
2948                         column->io.domain.domain_info = NULL;
2949                 }
2950         }
2951         else if (type->typtype == TYPTYPE_COMPOSITE || typid == RECORDOID)
2952         {
2953                 column->typcat = TYPECAT_COMPOSITE;
2954                 column->io.composite.record_io = NULL;
2955                 column->io.composite.tupdesc = NULL;
2956                 column->io.composite.base_typid = typid;
2957                 column->io.composite.base_typmod = typmod;
2958                 column->io.composite.domain_info = NULL;
2959         }
2960         else if (type->typlen == -1 && OidIsValid(type->typelem))
2961         {
2962                 column->typcat = TYPECAT_ARRAY;
2963                 column->io.array.element_info = MemoryContextAllocZero(mcxt,
2964                                                                                                                            sizeof(ColumnIOData));
2965                 column->io.array.element_type = type->typelem;
2966                 /* array element typemod stored in attribute's typmod */
2967                 column->io.array.element_typmod = typmod;
2968         }
2969         else
2970         {
2971                 column->typcat = TYPECAT_SCALAR;
2972                 need_scalar = true;
2973         }
2974
2975         /* caller can force us to look up scalar_io info even for non-scalars */
2976         if (need_scalar)
2977         {
2978                 Oid                     typioproc;
2979
2980                 getTypeInputInfo(typid, &typioproc, &column->scalar_io.typioparam);
2981                 fmgr_info_cxt(typioproc, &column->scalar_io.typiofunc, mcxt);
2982         }
2983
2984         ReleaseSysCache(tup);
2985 }
2986
2987 /* recursively populate a record field or an array element from a json/jsonb value */
2988 static Datum
2989 populate_record_field(ColumnIOData *col,
2990                                           Oid typid,
2991                                           int32 typmod,
2992                                           const char *colname,
2993                                           MemoryContext mcxt,
2994                                           Datum defaultval,
2995                                           JsValue *jsv,
2996                                           bool *isnull)
2997 {
2998         TypeCat         typcat;
2999
3000         check_stack_depth();
3001
3002         /*
3003          * Prepare column metadata cache for the given type.  Force lookup of the
3004          * scalar_io data so that the json string hack below will work.
3005          */
3006         if (col->typid != typid || col->typmod != typmod)
3007                 prepare_column_cache(col, typid, typmod, mcxt, true);
3008
3009         *isnull = JsValueIsNull(jsv);
3010
3011         typcat = col->typcat;
3012
3013         /* try to convert json string to a non-scalar type through input function */
3014         if (JsValueIsString(jsv) &&
3015                 (typcat == TYPECAT_ARRAY ||
3016                  typcat == TYPECAT_COMPOSITE ||
3017                  typcat == TYPECAT_COMPOSITE_DOMAIN))
3018                 typcat = TYPECAT_SCALAR;
3019
3020         /* we must perform domain checks for NULLs, otherwise exit immediately */
3021         if (*isnull &&
3022                 typcat != TYPECAT_DOMAIN &&
3023                 typcat != TYPECAT_COMPOSITE_DOMAIN)
3024                 return (Datum) 0;
3025
3026         switch (typcat)
3027         {
3028                 case TYPECAT_SCALAR:
3029                         return populate_scalar(&col->scalar_io, typid, typmod, jsv);
3030
3031                 case TYPECAT_ARRAY:
3032                         return populate_array(&col->io.array, colname, mcxt, jsv);
3033
3034                 case TYPECAT_COMPOSITE:
3035                 case TYPECAT_COMPOSITE_DOMAIN:
3036                         return populate_composite(&col->io.composite, typid,
3037                                                                           colname, mcxt,
3038                                                                           DatumGetPointer(defaultval)
3039                                                                           ? DatumGetHeapTupleHeader(defaultval)
3040                                                                           : NULL,
3041                                                                           jsv, *isnull);
3042
3043                 case TYPECAT_DOMAIN:
3044                         return populate_domain(&col->io.domain, typid, colname, mcxt,
3045                                                                    jsv, *isnull);
3046
3047                 default:
3048                         elog(ERROR, "unrecognized type category '%c'", typcat);
3049                         return (Datum) 0;
3050         }
3051 }
3052
3053 static RecordIOData *
3054 allocate_record_info(MemoryContext mcxt, int ncolumns)
3055 {
3056         RecordIOData *data = (RecordIOData *)
3057         MemoryContextAlloc(mcxt,
3058                                            offsetof(RecordIOData, columns) +
3059                                            ncolumns * sizeof(ColumnIOData));
3060
3061         data->record_type = InvalidOid;
3062         data->record_typmod = 0;
3063         data->ncolumns = ncolumns;
3064         MemSet(data->columns, 0, sizeof(ColumnIOData) * ncolumns);
3065
3066         return data;
3067 }
3068
3069 static bool
3070 JsObjectGetField(JsObject *obj, char *field, JsValue *jsv)
3071 {
3072         jsv->is_json = obj->is_json;
3073
3074         if (jsv->is_json)
3075         {
3076                 JsonHashEntry *hashentry = hash_search(obj->val.json_hash, field,
3077                                                                                            HASH_FIND, NULL);
3078
3079                 jsv->val.json.type = hashentry ? hashentry->type : JSON_TOKEN_NULL;
3080                 jsv->val.json.str = jsv->val.json.type == JSON_TOKEN_NULL ? NULL :
3081                         hashentry->val;
3082                 jsv->val.json.len = jsv->val.json.str ? -1 : 0; /* null-terminated */
3083
3084                 return hashentry != NULL;
3085         }
3086         else
3087         {
3088                 jsv->val.jsonb = !obj->val.jsonb_cont ? NULL :
3089                         findJsonbValueFromContainerLen(obj->val.jsonb_cont, JB_FOBJECT,
3090                                                                                    field, strlen(field));
3091
3092                 return jsv->val.jsonb != NULL;
3093         }
3094 }
3095
3096 /* populate a record tuple from json/jsonb value */
3097 static HeapTupleHeader
3098 populate_record(TupleDesc tupdesc,
3099                                 RecordIOData **record_p,
3100                                 HeapTupleHeader defaultval,
3101                                 MemoryContext mcxt,
3102                                 JsObject *obj)
3103 {
3104         RecordIOData *record = *record_p;
3105         Datum      *values;
3106         bool       *nulls;
3107         HeapTuple       res;
3108         int                     ncolumns = tupdesc->natts;
3109         int                     i;
3110
3111         /*
3112          * if the input json is empty, we can only skip the rest if we were passed
3113          * in a non-null record, since otherwise there may be issues with domain
3114          * nulls.
3115          */
3116         if (defaultval && JsObjectIsEmpty(obj))
3117                 return defaultval;
3118
3119         /* (re)allocate metadata cache */
3120         if (record == NULL ||
3121                 record->ncolumns != ncolumns)
3122                 *record_p = record = allocate_record_info(mcxt, ncolumns);
3123
3124         /* invalidate metadata cache if the record type has changed */
3125         if (record->record_type != tupdesc->tdtypeid ||
3126                 record->record_typmod != tupdesc->tdtypmod)
3127         {
3128                 MemSet(record, 0, offsetof(RecordIOData, columns) +
3129                            ncolumns * sizeof(ColumnIOData));
3130                 record->record_type = tupdesc->tdtypeid;
3131                 record->record_typmod = tupdesc->tdtypmod;
3132                 record->ncolumns = ncolumns;
3133         }
3134
3135         values = (Datum *) palloc(ncolumns * sizeof(Datum));
3136         nulls = (bool *) palloc(ncolumns * sizeof(bool));
3137
3138         if (defaultval)
3139         {
3140                 HeapTupleData tuple;
3141
3142                 /* Build a temporary HeapTuple control structure */
3143                 tuple.t_len = HeapTupleHeaderGetDatumLength(defaultval);
3144                 ItemPointerSetInvalid(&(tuple.t_self));
3145                 tuple.t_tableOid = InvalidOid;
3146                 tuple.t_data = defaultval;
3147
3148                 /* Break down the tuple into fields */
3149                 heap_deform_tuple(&tuple, tupdesc, values, nulls);
3150         }
3151         else
3152         {
3153                 for (i = 0; i < ncolumns; ++i)
3154                 {
3155                         values[i] = (Datum) 0;
3156                         nulls[i] = true;
3157                 }
3158         }
3159
3160         for (i = 0; i < ncolumns; ++i)
3161         {
3162                 Form_pg_attribute att = TupleDescAttr(tupdesc, i);
3163                 char       *colname = NameStr(att->attname);
3164                 JsValue         field = {0};
3165                 bool            found;
3166
3167                 /* Ignore dropped columns in datatype */
3168                 if (att->attisdropped)
3169                 {
3170                         nulls[i] = true;
3171                         continue;
3172                 }
3173
3174                 found = JsObjectGetField(obj, colname, &field);
3175
3176                 /*
3177                  * we can't just skip here if the key wasn't found since we might have
3178                  * a domain to deal with. If we were passed in a non-null record
3179                  * datum, we assume that the existing values are valid (if they're
3180                  * not, then it's not our fault), but if we were passed in a null,
3181                  * then every field which we don't populate needs to be run through
3182                  * the input function just in case it's a domain type.
3183                  */
3184                 if (defaultval && !found)
3185                         continue;
3186
3187                 values[i] = populate_record_field(&record->columns[i],
3188                                                                                   att->atttypid,
3189                                                                                   att->atttypmod,
3190                                                                                   colname,
3191                                                                                   mcxt,
3192                                                                                   nulls[i] ? (Datum) 0 : values[i],
3193                                                                                   &field,
3194                                                                                   &nulls[i]);
3195         }
3196
3197         res = heap_form_tuple(tupdesc, values, nulls);
3198
3199         pfree(values);
3200         pfree(nulls);
3201
3202         return res->t_data;
3203 }
3204
3205 /*
3206  * common worker for json{b}_populate_record() and json{b}_to_record()
3207  * is_json and have_record_arg identify the specific function
3208  */
3209 static Datum
3210 populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
3211                                            bool is_json, bool have_record_arg)
3212 {
3213         int                     json_arg_num = have_record_arg ? 1 : 0;
3214         JsValue         jsv = {0};
3215         HeapTupleHeader rec;
3216         Datum           rettuple;
3217         JsonbValue      jbv;
3218         MemoryContext fnmcxt = fcinfo->flinfo->fn_mcxt;
3219         PopulateRecordCache *cache = fcinfo->flinfo->fn_extra;
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 = is_json;
3307
3308         if (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",
3493                                                                          false, true);
3494 }
3495
3496 Datum
3497 jsonb_to_recordset(PG_FUNCTION_ARGS)
3498 {
3499         return populate_recordset_worker(fcinfo, "jsonb_to_recordset",
3500                                                                          false, false);
3501 }
3502
3503 Datum
3504 json_populate_recordset(PG_FUNCTION_ARGS)
3505 {
3506         return populate_recordset_worker(fcinfo, "json_populate_recordset",
3507                                                                          true, true);
3508 }
3509
3510 Datum
3511 json_to_recordset(PG_FUNCTION_ARGS)
3512 {
3513         return populate_recordset_worker(fcinfo, "json_to_recordset",
3514                                                                          true, false);
3515 }
3516
3517 static void
3518 populate_recordset_record(PopulateRecordsetState *state, JsObject *obj)
3519 {
3520         PopulateRecordsetCache *cache = state->cache;
3521         HeapTupleHeader tuphead;
3522         HeapTupleData tuple;
3523
3524         /* acquire/update cached tuple descriptor */
3525         update_cached_tupdesc(&cache->c.io.composite, cache->fn_mcxt);
3526
3527         /* replace record fields from json */
3528         tuphead = populate_record(cache->c.io.composite.tupdesc,
3529                                                           &cache->c.io.composite.record_io,
3530                                                           state->rec,
3531                                                           cache->fn_mcxt,
3532                                                           obj);
3533
3534         /* if it's domain over composite, check domain constraints */
3535         if (cache->c.typcat == TYPECAT_COMPOSITE_DOMAIN)
3536                 domain_check(HeapTupleHeaderGetDatum(tuphead), false,
3537                                          cache->argtype,
3538                                          &cache->c.io.composite.domain_info,
3539                                          cache->fn_mcxt);
3540
3541         /* ok, save into tuplestore */
3542         tuple.t_len = HeapTupleHeaderGetDatumLength(tuphead);
3543         ItemPointerSetInvalid(&(tuple.t_self));
3544         tuple.t_tableOid = InvalidOid;
3545         tuple.t_data = tuphead;
3546
3547         tuplestore_puttuple(state->tuple_store, &tuple);
3548 }
3549
3550 /*
3551  * common worker for json{b}_populate_recordset() and json{b}_to_recordset()
3552  * is_json and have_record_arg identify the specific function
3553  */
3554 static Datum
3555 populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
3556                                                   bool is_json, bool have_record_arg)
3557 {
3558         int                     json_arg_num = have_record_arg ? 1 : 0;
3559         ReturnSetInfo *rsi;
3560         MemoryContext old_cxt;
3561         HeapTupleHeader rec;
3562         PopulateRecordsetCache *cache = fcinfo->flinfo->fn_extra;
3563         PopulateRecordsetState *state;
3564
3565         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
3566
3567         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
3568                 (rsi->allowedModes & SFRM_Materialize) == 0)
3569                 ereport(ERROR,
3570                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
3571                                  errmsg("set-valued function called in context that "
3572                                                 "cannot accept a set")));
3573
3574         rsi->returnMode = SFRM_Materialize;
3575
3576         /*
3577          * If first time through, identify input/result record type.  Note that
3578          * this stanza looks only at fcinfo context, which can't change during the
3579          * query; so we may not be able to fully resolve a RECORD input type yet.
3580          */
3581         if (!cache)
3582         {
3583                 fcinfo->flinfo->fn_extra = cache =
3584                         MemoryContextAllocZero(fcinfo->flinfo->fn_mcxt, sizeof(*cache));
3585                 cache->fn_mcxt = fcinfo->flinfo->fn_mcxt;
3586
3587                 if (have_record_arg)
3588                 {
3589                         /*
3590                          * json{b}_populate_recordset case: result type will be same as
3591                          * first argument's.
3592                          */
3593                         cache->argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
3594                         prepare_column_cache(&cache->c,
3595                                                                  cache->argtype, -1,
3596                                                                  cache->fn_mcxt, false);
3597                         if (cache->c.typcat != TYPECAT_COMPOSITE &&
3598                                 cache->c.typcat != TYPECAT_COMPOSITE_DOMAIN)
3599                                 ereport(ERROR,
3600                                                 (errcode(ERRCODE_DATATYPE_MISMATCH),
3601                                                  errmsg("first argument of %s must be a row type",
3602                                                                 funcname)));
3603                 }
3604                 else
3605                 {
3606                         /*
3607                          * json{b}_to_recordset case: result type is specified by calling
3608                          * query.  Here it is syntactically impossible to specify the
3609                          * target type as domain-over-composite.
3610                          */
3611                         TupleDesc       tupdesc;
3612
3613                         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
3614                                 ereport(ERROR,
3615                                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
3616                                                  errmsg("function returning record called in context "
3617                                                                 "that cannot accept type record"),
3618                                                  errhint("Try calling the function in the FROM clause "
3619                                                                  "using a column definition list.")));
3620
3621                         Assert(tupdesc);
3622                         cache->argtype = tupdesc->tdtypeid;
3623
3624                         /* Save identified tupdesc */
3625                         old_cxt = MemoryContextSwitchTo(cache->fn_mcxt);
3626                         cache->c.io.composite.tupdesc = CreateTupleDescCopy(tupdesc);
3627                         cache->c.io.composite.base_typid = tupdesc->tdtypeid;
3628                         cache->c.io.composite.base_typmod = tupdesc->tdtypmod;
3629                         MemoryContextSwitchTo(old_cxt);
3630                 }
3631         }
3632
3633         /* Collect record arg if we have one */
3634         if (have_record_arg && !PG_ARGISNULL(0))
3635         {
3636                 rec = PG_GETARG_HEAPTUPLEHEADER(0);
3637
3638                 /*
3639                  * When declared arg type is RECORD, identify actual record type from
3640                  * the tuple itself.  Note the lookup_rowtype_tupdesc call in
3641                  * update_cached_tupdesc will fail if we're unable to do this.
3642                  */
3643                 if (cache->argtype == RECORDOID)
3644                 {
3645                         cache->c.io.composite.base_typid = HeapTupleHeaderGetTypeId(rec);
3646                         cache->c.io.composite.base_typmod = HeapTupleHeaderGetTypMod(rec);
3647                 }
3648         }
3649         else
3650                 rec = NULL;
3651
3652         /* if the json is null send back an empty set */
3653         if (PG_ARGISNULL(json_arg_num))
3654                 PG_RETURN_NULL();
3655
3656         /*
3657          * Forcibly update the cached tupdesc, to ensure we have the right tupdesc
3658          * to return even if the JSON contains no rows.
3659          */
3660         update_cached_tupdesc(&cache->c.io.composite, cache->fn_mcxt);
3661
3662         state = palloc0(sizeof(PopulateRecordsetState));
3663
3664         /* make tuplestore in a sufficiently long-lived memory context */
3665         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
3666         state->tuple_store = tuplestore_begin_heap(rsi->allowedModes &
3667                                                                                            SFRM_Materialize_Random,
3668                                                                                            false, work_mem);
3669         MemoryContextSwitchTo(old_cxt);
3670
3671         state->function_name = funcname;
3672         state->cache = cache;
3673         state->rec = rec;
3674
3675         if (is_json)
3676         {
3677                 text       *json = PG_GETARG_TEXT_PP(json_arg_num);
3678                 JsonLexContext *lex;
3679                 JsonSemAction *sem;
3680
3681                 sem = palloc0(sizeof(JsonSemAction));
3682
3683                 lex = makeJsonLexContext(json, true);
3684
3685                 sem->semstate = (void *) state;
3686                 sem->array_start = populate_recordset_array_start;
3687                 sem->array_element_start = populate_recordset_array_element_start;
3688                 sem->scalar = populate_recordset_scalar;
3689                 sem->object_field_start = populate_recordset_object_field_start;
3690                 sem->object_field_end = populate_recordset_object_field_end;
3691                 sem->object_start = populate_recordset_object_start;
3692                 sem->object_end = populate_recordset_object_end;
3693
3694                 state->lex = lex;
3695
3696                 pg_parse_json(lex, sem);
3697         }
3698         else
3699         {
3700                 Jsonb      *jb = PG_GETARG_JSONB_P(json_arg_num);
3701                 JsonbIterator *it;
3702                 JsonbValue      v;
3703                 bool            skipNested = false;
3704                 JsonbIteratorToken r;
3705
3706                 if (JB_ROOT_IS_SCALAR(jb) || !JB_ROOT_IS_ARRAY(jb))
3707                         ereport(ERROR,
3708                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3709                                          errmsg("cannot call %s on a non-array",
3710                                                         funcname)));
3711
3712                 it = JsonbIteratorInit(&jb->root);
3713
3714                 while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
3715                 {
3716                         skipNested = true;
3717
3718                         if (r == WJB_ELEM)
3719                         {
3720                                 JsObject        obj;
3721
3722                                 if (v.type != jbvBinary ||
3723                                         !JsonContainerIsObject(v.val.binary.data))
3724                                         ereport(ERROR,
3725                                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3726                                                          errmsg("argument of %s must be an array of objects",
3727                                                                         funcname)));
3728
3729                                 obj.is_json = false;
3730                                 obj.val.jsonb_cont = v.val.binary.data;
3731
3732                                 populate_recordset_record(state, &obj);
3733                         }
3734                 }
3735         }
3736
3737         /*
3738          * Note: we must copy the cached tupdesc because the executor will free
3739          * the passed-back setDesc, but we want to hang onto the cache in case
3740          * we're called again in the same query.
3741          */
3742         rsi->setResult = state->tuple_store;
3743         rsi->setDesc = CreateTupleDescCopy(cache->c.io.composite.tupdesc);
3744
3745         PG_RETURN_NULL();
3746 }
3747
3748 static void
3749 populate_recordset_object_start(void *state)
3750 {
3751         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3752         int                     lex_level = _state->lex->lex_level;
3753         HASHCTL         ctl;
3754
3755         /* Reject object at top level: we must have an array at level 0 */
3756         if (lex_level == 0)
3757                 ereport(ERROR,
3758                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3759                                  errmsg("cannot call %s on an object",
3760                                                 _state->function_name)));
3761
3762         /* Nested objects require no special processing */
3763         if (lex_level > 1)
3764                 return;
3765
3766         /* Object at level 1: set up a new hash table for this object */
3767         memset(&ctl, 0, sizeof(ctl));
3768         ctl.keysize = NAMEDATALEN;
3769         ctl.entrysize = sizeof(JsonHashEntry);
3770         ctl.hcxt = CurrentMemoryContext;
3771         _state->json_hash = hash_create("json object hashtable",
3772                                                                         100,
3773                                                                         &ctl,
3774                                                                         HASH_ELEM | HASH_CONTEXT);
3775 }
3776
3777 static void
3778 populate_recordset_object_end(void *state)
3779 {
3780         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3781         JsObject        obj;
3782
3783         /* Nested objects require no special processing */
3784         if (_state->lex->lex_level > 1)
3785                 return;
3786
3787         obj.is_json = true;
3788         obj.val.json_hash = _state->json_hash;
3789
3790         /* Otherwise, construct and return a tuple based on this level-1 object */
3791         populate_recordset_record(_state, &obj);
3792
3793         /* Done with hash for this object */
3794         hash_destroy(_state->json_hash);
3795         _state->json_hash = NULL;
3796 }
3797
3798 static void
3799 populate_recordset_array_element_start(void *state, bool isnull)
3800 {
3801         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3802
3803         if (_state->lex->lex_level == 1 &&
3804                 _state->lex->token_type != JSON_TOKEN_OBJECT_START)
3805                 ereport(ERROR,
3806                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3807                                  errmsg("argument of %s must be an array of objects",
3808                                                 _state->function_name)));
3809 }
3810
3811 static void
3812 populate_recordset_array_start(void *state)
3813 {
3814         /* nothing to do */
3815 }
3816
3817 static void
3818 populate_recordset_scalar(void *state, char *token, JsonTokenType tokentype)
3819 {
3820         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3821
3822         if (_state->lex->lex_level == 0)
3823                 ereport(ERROR,
3824                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3825                                  errmsg("cannot call %s on a scalar",
3826                                                 _state->function_name)));
3827
3828         if (_state->lex->lex_level == 2)
3829                 _state->saved_scalar = token;
3830 }
3831
3832 static void
3833 populate_recordset_object_field_start(void *state, char *fname, bool isnull)
3834 {
3835         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3836
3837         if (_state->lex->lex_level > 2)
3838                 return;
3839
3840         _state->saved_token_type = _state->lex->token_type;
3841
3842         if (_state->lex->token_type == JSON_TOKEN_ARRAY_START ||
3843                 _state->lex->token_type == JSON_TOKEN_OBJECT_START)
3844         {
3845                 _state->save_json_start = _state->lex->token_start;
3846         }
3847         else
3848         {
3849                 _state->save_json_start = NULL;
3850         }
3851 }
3852
3853 static void
3854 populate_recordset_object_field_end(void *state, char *fname, bool isnull)
3855 {
3856         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3857         JsonHashEntry *hashentry;
3858         bool            found;
3859
3860         /*
3861          * Ignore nested fields.
3862          */
3863         if (_state->lex->lex_level > 2)
3864                 return;
3865
3866         /*
3867          * Ignore field names >= NAMEDATALEN - they can't match a record field.
3868          * (Note: without this test, the hash code would truncate the string at
3869          * NAMEDATALEN-1, and could then match against a similarly-truncated
3870          * record field name.  That would be a reasonable behavior, but this code
3871          * has previously insisted on exact equality, so we keep this behavior.)
3872          */
3873         if (strlen(fname) >= NAMEDATALEN)
3874                 return;
3875
3876         hashentry = hash_search(_state->json_hash, fname, HASH_ENTER, &found);
3877
3878         /*
3879          * found being true indicates a duplicate. We don't do anything about
3880          * that, a later field with the same name overrides the earlier field.
3881          */
3882
3883         hashentry->type = _state->saved_token_type;
3884         Assert(isnull == (hashentry->type == JSON_TOKEN_NULL));
3885
3886         if (_state->save_json_start != NULL)
3887         {
3888                 int                     len = _state->lex->prev_token_terminator - _state->save_json_start;
3889                 char       *val = palloc((len + 1) * sizeof(char));
3890
3891                 memcpy(val, _state->save_json_start, len);
3892                 val[len] = '\0';
3893                 hashentry->val = val;
3894         }
3895         else
3896         {
3897                 /* must have had a scalar instead */
3898                 hashentry->val = _state->saved_scalar;
3899         }
3900 }
3901
3902 /*
3903  * findJsonbValueFromContainer() wrapper that sets up JsonbValue key string.
3904  */
3905 static JsonbValue *
3906 findJsonbValueFromContainerLen(JsonbContainer *container, uint32 flags,
3907                                                            char *key, uint32 keylen)
3908 {
3909         JsonbValue      k;
3910
3911         k.type = jbvString;
3912         k.val.string.val = key;
3913         k.val.string.len = keylen;
3914
3915         return findJsonbValueFromContainer(container, flags, &k);
3916 }
3917
3918 /*
3919  * Semantic actions for json_strip_nulls.
3920  *
3921  * Simply repeat the input on the output unless we encounter
3922  * a null object field. State for this is set when the field
3923  * is started and reset when the scalar action (which must be next)
3924  * is called.
3925  */
3926
3927 static void
3928 sn_object_start(void *state)
3929 {
3930         StripnullState *_state = (StripnullState *) state;
3931
3932         appendStringInfoCharMacro(_state->strval, '{');
3933 }
3934
3935 static void
3936 sn_object_end(void *state)
3937 {
3938         StripnullState *_state = (StripnullState *) state;
3939
3940         appendStringInfoCharMacro(_state->strval, '}');
3941 }
3942
3943 static void
3944 sn_array_start(void *state)
3945 {
3946         StripnullState *_state = (StripnullState *) state;
3947
3948         appendStringInfoCharMacro(_state->strval, '[');
3949 }
3950
3951 static void
3952 sn_array_end(void *state)
3953 {
3954         StripnullState *_state = (StripnullState *) state;
3955
3956         appendStringInfoCharMacro(_state->strval, ']');
3957 }
3958
3959 static void
3960 sn_object_field_start(void *state, char *fname, bool isnull)
3961 {
3962         StripnullState *_state = (StripnullState *) state;
3963
3964         if (isnull)
3965         {
3966                 /*
3967                  * The next thing must be a scalar or isnull couldn't be true, so
3968                  * there is no danger of this state being carried down into a nested
3969                  * object or array. The flag will be reset in the scalar action.
3970                  */
3971                 _state->skip_next_null = true;
3972                 return;
3973         }
3974
3975         if (_state->strval->data[_state->strval->len - 1] != '{')
3976                 appendStringInfoCharMacro(_state->strval, ',');
3977
3978         /*
3979          * Unfortunately we don't have the quoted and escaped string any more, so
3980          * we have to re-escape it.
3981          */
3982         escape_json(_state->strval, fname);
3983
3984         appendStringInfoCharMacro(_state->strval, ':');
3985 }
3986
3987 static void
3988 sn_array_element_start(void *state, bool isnull)
3989 {
3990         StripnullState *_state = (StripnullState *) state;
3991
3992         if (_state->strval->data[_state->strval->len - 1] != '[')
3993                 appendStringInfoCharMacro(_state->strval, ',');
3994 }
3995
3996 static void
3997 sn_scalar(void *state, char *token, JsonTokenType tokentype)
3998 {
3999         StripnullState *_state = (StripnullState *) state;
4000
4001         if (_state->skip_next_null)
4002         {
4003                 Assert(tokentype == JSON_TOKEN_NULL);
4004                 _state->skip_next_null = false;
4005                 return;
4006         }
4007
4008         if (tokentype == JSON_TOKEN_STRING)
4009                 escape_json(_state->strval, token);
4010         else
4011                 appendStringInfoString(_state->strval, token);
4012 }
4013
4014 /*
4015  * SQL function json_strip_nulls(json) -> json
4016  */
4017 Datum
4018 json_strip_nulls(PG_FUNCTION_ARGS)
4019 {
4020         text       *json = PG_GETARG_TEXT_PP(0);
4021         StripnullState *state;
4022         JsonLexContext *lex;
4023         JsonSemAction *sem;
4024
4025         lex = makeJsonLexContext(json, true);
4026         state = palloc0(sizeof(StripnullState));
4027         sem = palloc0(sizeof(JsonSemAction));
4028
4029         state->strval = makeStringInfo();
4030         state->skip_next_null = false;
4031         state->lex = lex;
4032
4033         sem->semstate = (void *) state;
4034         sem->object_start = sn_object_start;
4035         sem->object_end = sn_object_end;
4036         sem->array_start = sn_array_start;
4037         sem->array_end = sn_array_end;
4038         sem->scalar = sn_scalar;
4039         sem->array_element_start = sn_array_element_start;
4040         sem->object_field_start = sn_object_field_start;
4041
4042         pg_parse_json(lex, sem);
4043
4044         PG_RETURN_TEXT_P(cstring_to_text_with_len(state->strval->data,
4045                                                                                           state->strval->len));
4046
4047 }
4048
4049 /*
4050  * SQL function jsonb_strip_nulls(jsonb) -> jsonb
4051  */
4052 Datum
4053 jsonb_strip_nulls(PG_FUNCTION_ARGS)
4054 {
4055         Jsonb      *jb = PG_GETARG_JSONB_P(0);
4056         JsonbIterator *it;
4057         JsonbParseState *parseState = NULL;
4058         JsonbValue *res = NULL;
4059         JsonbValue      v,
4060                                 k;
4061         JsonbIteratorToken type;
4062         bool            last_was_key = false;
4063
4064         if (JB_ROOT_IS_SCALAR(jb))
4065                 PG_RETURN_POINTER(jb);
4066
4067         it = JsonbIteratorInit(&jb->root);
4068
4069         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
4070         {
4071                 Assert(!(type == WJB_KEY && last_was_key));
4072
4073                 if (type == WJB_KEY)
4074                 {
4075                         /* stash the key until we know if it has a null value */
4076                         k = v;
4077                         last_was_key = true;
4078                         continue;
4079                 }
4080
4081                 if (last_was_key)
4082                 {
4083                         /* if the last element was a key this one can't be */
4084                         last_was_key = false;
4085
4086                         /* skip this field if value is null */
4087                         if (type == WJB_VALUE && v.type == jbvNull)
4088                                 continue;
4089
4090                         /* otherwise, do a delayed push of the key */
4091                         (void) pushJsonbValue(&parseState, WJB_KEY, &k);
4092                 }
4093
4094                 if (type == WJB_VALUE || type == WJB_ELEM)
4095                         res = pushJsonbValue(&parseState, type, &v);
4096                 else
4097                         res = pushJsonbValue(&parseState, type, NULL);
4098         }
4099
4100         Assert(res != NULL);
4101
4102         PG_RETURN_POINTER(JsonbValueToJsonb(res));
4103 }
4104
4105 /*
4106  * Add values from the jsonb to the parse state.
4107  *
4108  * If the parse state container is an object, the jsonb is pushed as
4109  * a value, not a key.
4110  *
4111  * This needs to be done using an iterator because pushJsonbValue doesn't
4112  * like getting jbvBinary values, so we can't just push jb as a whole.
4113  */
4114 static void
4115 addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb)
4116 {
4117         JsonbIterator *it;
4118         JsonbValue *o = &(*jbps)->contVal;
4119         JsonbValue      v;
4120         JsonbIteratorToken type;
4121
4122         it = JsonbIteratorInit(&jb->root);
4123
4124         Assert(o->type == jbvArray || o->type == jbvObject);
4125
4126         if (JB_ROOT_IS_SCALAR(jb))
4127         {
4128                 (void) JsonbIteratorNext(&it, &v, false);       /* skip array header */
4129                 Assert(v.type == jbvArray);
4130                 (void) JsonbIteratorNext(&it, &v, false);       /* fetch scalar value */
4131
4132                 switch (o->type)
4133                 {
4134                         case jbvArray:
4135                                 (void) pushJsonbValue(jbps, WJB_ELEM, &v);
4136                                 break;
4137                         case jbvObject:
4138                                 (void) pushJsonbValue(jbps, WJB_VALUE, &v);
4139                                 break;
4140                         default:
4141                                 elog(ERROR, "unexpected parent of nested structure");
4142                 }
4143         }
4144         else
4145         {
4146                 while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
4147                 {
4148                         if (type == WJB_KEY || type == WJB_VALUE || type == WJB_ELEM)
4149                                 (void) pushJsonbValue(jbps, type, &v);
4150                         else
4151                                 (void) pushJsonbValue(jbps, type, NULL);
4152                 }
4153         }
4154
4155 }
4156
4157 /*
4158  * SQL function jsonb_pretty (jsonb)
4159  *
4160  * Pretty-printed text for the jsonb
4161  */
4162 Datum
4163 jsonb_pretty(PG_FUNCTION_ARGS)
4164 {
4165         Jsonb      *jb = PG_GETARG_JSONB_P(0);
4166         StringInfo      str = makeStringInfo();
4167
4168         JsonbToCStringIndent(str, &jb->root, VARSIZE(jb));
4169
4170         PG_RETURN_TEXT_P(cstring_to_text_with_len(str->data, str->len));
4171 }
4172
4173 /*
4174  * SQL function jsonb_concat (jsonb, jsonb)
4175  *
4176  * function for || operator
4177  */
4178 Datum
4179 jsonb_concat(PG_FUNCTION_ARGS)
4180 {
4181         Jsonb      *jb1 = PG_GETARG_JSONB_P(0);
4182         Jsonb      *jb2 = PG_GETARG_JSONB_P(1);
4183         JsonbParseState *state = NULL;
4184         JsonbValue *res;
4185         JsonbIterator *it1,
4186                            *it2;
4187
4188         /*
4189          * If one of the jsonb is empty, just return the other if it's not scalar
4190          * and both are of the same kind.  If it's a scalar or they are of
4191          * different kinds we need to perform the concatenation even if one is
4192          * empty.
4193          */
4194         if (JB_ROOT_IS_OBJECT(jb1) == JB_ROOT_IS_OBJECT(jb2))
4195         {
4196                 if (JB_ROOT_COUNT(jb1) == 0 && !JB_ROOT_IS_SCALAR(jb2))
4197                         PG_RETURN_JSONB_P(jb2);
4198                 else if (JB_ROOT_COUNT(jb2) == 0 && !JB_ROOT_IS_SCALAR(jb1))
4199                         PG_RETURN_JSONB_P(jb1);
4200         }
4201
4202         it1 = JsonbIteratorInit(&jb1->root);
4203         it2 = JsonbIteratorInit(&jb2->root);
4204
4205         res = IteratorConcat(&it1, &it2, &state);
4206
4207         Assert(res != NULL);
4208
4209         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4210 }
4211
4212
4213 /*
4214  * SQL function jsonb_delete (jsonb, text)
4215  *
4216  * return a copy of the jsonb with the indicated item
4217  * removed.
4218  */
4219 Datum
4220 jsonb_delete(PG_FUNCTION_ARGS)
4221 {
4222         Jsonb      *in = PG_GETARG_JSONB_P(0);
4223         text       *key = PG_GETARG_TEXT_PP(1);
4224         char       *keyptr = VARDATA_ANY(key);
4225         int                     keylen = VARSIZE_ANY_EXHDR(key);
4226         JsonbParseState *state = NULL;
4227         JsonbIterator *it;
4228         JsonbValue      v,
4229                            *res = NULL;
4230         bool            skipNested = false;
4231         JsonbIteratorToken r;
4232
4233         if (JB_ROOT_IS_SCALAR(in))
4234                 ereport(ERROR,
4235                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4236                                  errmsg("cannot delete from scalar")));
4237
4238         if (JB_ROOT_COUNT(in) == 0)
4239                 PG_RETURN_JSONB_P(in);
4240
4241         it = JsonbIteratorInit(&in->root);
4242
4243         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
4244         {
4245                 skipNested = true;
4246
4247                 if ((r == WJB_ELEM || r == WJB_KEY) &&
4248                         (v.type == jbvString && keylen == v.val.string.len &&
4249                          memcmp(keyptr, v.val.string.val, keylen) == 0))
4250                 {
4251                         /* skip corresponding value as well */
4252                         if (r == WJB_KEY)
4253                                 (void) JsonbIteratorNext(&it, &v, true);
4254
4255                         continue;
4256                 }
4257
4258                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4259         }
4260
4261         Assert(res != NULL);
4262
4263         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4264 }
4265
4266 /*
4267  * SQL function jsonb_delete (jsonb, variadic text[])
4268  *
4269  * return a copy of the jsonb with the indicated items
4270  * removed.
4271  */
4272 Datum
4273 jsonb_delete_array(PG_FUNCTION_ARGS)
4274 {
4275         Jsonb      *in = PG_GETARG_JSONB_P(0);
4276         ArrayType  *keys = PG_GETARG_ARRAYTYPE_P(1);
4277         Datum      *keys_elems;
4278         bool       *keys_nulls;
4279         int                     keys_len;
4280         JsonbParseState *state = NULL;
4281         JsonbIterator *it;
4282         JsonbValue      v,
4283                            *res = NULL;
4284         bool            skipNested = false;
4285         JsonbIteratorToken r;
4286
4287         if (ARR_NDIM(keys) > 1)
4288                 ereport(ERROR,
4289                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4290                                  errmsg("wrong number of array subscripts")));
4291
4292         if (JB_ROOT_IS_SCALAR(in))
4293                 ereport(ERROR,
4294                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4295                                  errmsg("cannot delete from scalar")));
4296
4297         if (JB_ROOT_COUNT(in) == 0)
4298                 PG_RETURN_JSONB_P(in);
4299
4300         deconstruct_array(keys, TEXTOID, -1, false, 'i',
4301                                           &keys_elems, &keys_nulls, &keys_len);
4302
4303         if (keys_len == 0)
4304                 PG_RETURN_JSONB_P(in);
4305
4306         it = JsonbIteratorInit(&in->root);
4307
4308         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
4309         {
4310                 skipNested = true;
4311
4312                 if ((r == WJB_ELEM || r == WJB_KEY) && v.type == jbvString)
4313                 {
4314                         int                     i;
4315                         bool            found = false;
4316
4317                         for (i = 0; i < keys_len; i++)
4318                         {
4319                                 char       *keyptr;
4320                                 int                     keylen;
4321
4322                                 if (keys_nulls[i])
4323                                         continue;
4324
4325                                 keyptr = VARDATA_ANY(keys_elems[i]);
4326                                 keylen = VARSIZE_ANY_EXHDR(keys_elems[i]);
4327                                 if (keylen == v.val.string.len &&
4328                                         memcmp(keyptr, v.val.string.val, keylen) == 0)
4329                                 {
4330                                         found = true;
4331                                         break;
4332                                 }
4333                         }
4334                         if (found)
4335                         {
4336                                 /* skip corresponding value as well */
4337                                 if (r == WJB_KEY)
4338                                         (void) JsonbIteratorNext(&it, &v, true);
4339
4340                                 continue;
4341                         }
4342                 }
4343
4344                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4345         }
4346
4347         Assert(res != NULL);
4348
4349         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4350 }
4351
4352 /*
4353  * SQL function jsonb_delete (jsonb, int)
4354  *
4355  * return a copy of the jsonb with the indicated item
4356  * removed. Negative int means count back from the
4357  * end of the items.
4358  */
4359 Datum
4360 jsonb_delete_idx(PG_FUNCTION_ARGS)
4361 {
4362         Jsonb      *in = PG_GETARG_JSONB_P(0);
4363         int                     idx = PG_GETARG_INT32(1);
4364         JsonbParseState *state = NULL;
4365         JsonbIterator *it;
4366         uint32          i = 0,
4367                                 n;
4368         JsonbValue      v,
4369                            *res = NULL;
4370         JsonbIteratorToken r;
4371
4372         if (JB_ROOT_IS_SCALAR(in))
4373                 ereport(ERROR,
4374                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4375                                  errmsg("cannot delete from scalar")));
4376
4377         if (JB_ROOT_IS_OBJECT(in))
4378                 ereport(ERROR,
4379                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4380                                  errmsg("cannot delete from object using integer index")));
4381
4382         if (JB_ROOT_COUNT(in) == 0)
4383                 PG_RETURN_JSONB_P(in);
4384
4385         it = JsonbIteratorInit(&in->root);
4386
4387         r = JsonbIteratorNext(&it, &v, false);
4388         Assert(r == WJB_BEGIN_ARRAY);
4389         n = v.val.array.nElems;
4390
4391         if (idx < 0)
4392         {
4393                 if (-idx > n)
4394                         idx = n;
4395                 else
4396                         idx = n + idx;
4397         }
4398
4399         if (idx >= n)
4400                 PG_RETURN_JSONB_P(in);
4401
4402         pushJsonbValue(&state, r, NULL);
4403
4404         while ((r = JsonbIteratorNext(&it, &v, true)) != WJB_DONE)
4405         {
4406                 if (r == WJB_ELEM)
4407                 {
4408                         if (i++ == idx)
4409                                 continue;
4410                 }
4411
4412                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4413         }
4414
4415         Assert(res != NULL);
4416
4417         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4418 }
4419
4420 /*
4421  * SQL function jsonb_set(jsonb, text[], jsonb, boolean)
4422  *
4423  */
4424 Datum
4425 jsonb_set(PG_FUNCTION_ARGS)
4426 {
4427         Jsonb      *in = PG_GETARG_JSONB_P(0);
4428         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4429         Jsonb      *newval = PG_GETARG_JSONB_P(2);
4430         bool            create = PG_GETARG_BOOL(3);
4431         JsonbValue *res = NULL;
4432         Datum      *path_elems;
4433         bool       *path_nulls;
4434         int                     path_len;
4435         JsonbIterator *it;
4436         JsonbParseState *st = NULL;
4437
4438         if (ARR_NDIM(path) > 1)
4439                 ereport(ERROR,
4440                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4441                                  errmsg("wrong number of array subscripts")));
4442
4443         if (JB_ROOT_IS_SCALAR(in))
4444                 ereport(ERROR,
4445                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4446                                  errmsg("cannot set path in scalar")));
4447
4448         if (JB_ROOT_COUNT(in) == 0 && !create)
4449                 PG_RETURN_JSONB_P(in);
4450
4451         deconstruct_array(path, TEXTOID, -1, false, 'i',
4452                                           &path_elems, &path_nulls, &path_len);
4453
4454         if (path_len == 0)
4455                 PG_RETURN_JSONB_P(in);
4456
4457         it = JsonbIteratorInit(&in->root);
4458
4459         res = setPath(&it, path_elems, path_nulls, path_len, &st,
4460                                   0, newval, create ? JB_PATH_CREATE : JB_PATH_REPLACE);
4461
4462         Assert(res != NULL);
4463
4464         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4465 }
4466
4467
4468 /*
4469  * SQL function jsonb_delete_path(jsonb, text[])
4470  */
4471 Datum
4472 jsonb_delete_path(PG_FUNCTION_ARGS)
4473 {
4474         Jsonb      *in = PG_GETARG_JSONB_P(0);
4475         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4476         JsonbValue *res = NULL;
4477         Datum      *path_elems;
4478         bool       *path_nulls;
4479         int                     path_len;
4480         JsonbIterator *it;
4481         JsonbParseState *st = NULL;
4482
4483         if (ARR_NDIM(path) > 1)
4484                 ereport(ERROR,
4485                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4486                                  errmsg("wrong number of array subscripts")));
4487
4488         if (JB_ROOT_IS_SCALAR(in))
4489                 ereport(ERROR,
4490                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4491                                  errmsg("cannot delete path in scalar")));
4492
4493         if (JB_ROOT_COUNT(in) == 0)
4494                 PG_RETURN_JSONB_P(in);
4495
4496         deconstruct_array(path, TEXTOID, -1, false, 'i',
4497                                           &path_elems, &path_nulls, &path_len);
4498
4499         if (path_len == 0)
4500                 PG_RETURN_JSONB_P(in);
4501
4502         it = JsonbIteratorInit(&in->root);
4503
4504         res = setPath(&it, path_elems, path_nulls, path_len, &st,
4505                                   0, NULL, JB_PATH_DELETE);
4506
4507         Assert(res != NULL);
4508
4509         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4510 }
4511
4512 /*
4513  * SQL function jsonb_insert(jsonb, text[], jsonb, boolean)
4514  *
4515  */
4516 Datum
4517 jsonb_insert(PG_FUNCTION_ARGS)
4518 {
4519         Jsonb      *in = PG_GETARG_JSONB_P(0);
4520         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
4521         Jsonb      *newval = PG_GETARG_JSONB_P(2);
4522         bool            after = PG_GETARG_BOOL(3);
4523         JsonbValue *res = NULL;
4524         Datum      *path_elems;
4525         bool       *path_nulls;
4526         int                     path_len;
4527         JsonbIterator *it;
4528         JsonbParseState *st = NULL;
4529
4530         if (ARR_NDIM(path) > 1)
4531                 ereport(ERROR,
4532                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
4533                                  errmsg("wrong number of array subscripts")));
4534
4535         if (JB_ROOT_IS_SCALAR(in))
4536                 ereport(ERROR,
4537                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4538                                  errmsg("cannot set path in scalar")));
4539
4540         deconstruct_array(path, TEXTOID, -1, false, 'i',
4541                                           &path_elems, &path_nulls, &path_len);
4542
4543         if (path_len == 0)
4544                 PG_RETURN_JSONB_P(in);
4545
4546         it = JsonbIteratorInit(&in->root);
4547
4548         res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, newval,
4549                                   after ? JB_PATH_INSERT_AFTER : JB_PATH_INSERT_BEFORE);
4550
4551         Assert(res != NULL);
4552
4553         PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
4554 }
4555
4556 /*
4557  * Iterate over all jsonb objects and merge them into one.
4558  * The logic of this function copied from the same hstore function,
4559  * except the case, when it1 & it2 represents jbvObject.
4560  * In that case we just append the content of it2 to it1 without any
4561  * verifications.
4562  */
4563 static JsonbValue *
4564 IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
4565                            JsonbParseState **state)
4566 {
4567         JsonbValue      v1,
4568                                 v2,
4569                            *res = NULL;
4570         JsonbIteratorToken r1,
4571                                 r2,
4572                                 rk1,
4573                                 rk2;
4574
4575         r1 = rk1 = JsonbIteratorNext(it1, &v1, false);
4576         r2 = rk2 = JsonbIteratorNext(it2, &v2, false);
4577
4578         /*
4579          * Both elements are objects.
4580          */
4581         if (rk1 == WJB_BEGIN_OBJECT && rk2 == WJB_BEGIN_OBJECT)
4582         {
4583                 /*
4584                  * Append the all tokens from v1 to res, except last WJB_END_OBJECT
4585                  * (because res will not be finished yet).
4586                  */
4587                 pushJsonbValue(state, r1, NULL);
4588                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_OBJECT)
4589                         pushJsonbValue(state, r1, &v1);
4590
4591                 /*
4592                  * Append the all tokens from v2 to res, include last WJB_END_OBJECT
4593                  * (the concatenation will be completed).
4594                  */
4595                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_DONE)
4596                         res = pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
4597         }
4598
4599         /*
4600          * Both elements are arrays (either can be scalar).
4601          */
4602         else if (rk1 == WJB_BEGIN_ARRAY && rk2 == WJB_BEGIN_ARRAY)
4603         {
4604                 pushJsonbValue(state, r1, NULL);
4605
4606                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_ARRAY)
4607                 {
4608                         Assert(r1 == WJB_ELEM);
4609                         pushJsonbValue(state, r1, &v1);
4610                 }
4611
4612                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_END_ARRAY)
4613                 {
4614                         Assert(r2 == WJB_ELEM);
4615                         pushJsonbValue(state, WJB_ELEM, &v2);
4616                 }
4617
4618                 res = pushJsonbValue(state, WJB_END_ARRAY, NULL /* signal to sort */ );
4619         }
4620         /* have we got array || object or object || array? */
4621         else if (((rk1 == WJB_BEGIN_ARRAY && !(*it1)->isScalar) && rk2 == WJB_BEGIN_OBJECT) ||
4622                          (rk1 == WJB_BEGIN_OBJECT && (rk2 == WJB_BEGIN_ARRAY && !(*it2)->isScalar)))
4623         {
4624
4625                 JsonbIterator **it_array = rk1 == WJB_BEGIN_ARRAY ? it1 : it2;
4626                 JsonbIterator **it_object = rk1 == WJB_BEGIN_OBJECT ? it1 : it2;
4627
4628                 bool            prepend = (rk1 == WJB_BEGIN_OBJECT);
4629
4630                 pushJsonbValue(state, WJB_BEGIN_ARRAY, NULL);
4631
4632                 if (prepend)
4633                 {
4634                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
4635                         while ((r1 = JsonbIteratorNext(it_object, &v1, true)) != WJB_DONE)
4636                                 pushJsonbValue(state, r1, r1 != WJB_END_OBJECT ? &v1 : NULL);
4637
4638                         while ((r2 = JsonbIteratorNext(it_array, &v2, true)) != WJB_DONE)
4639                                 res = pushJsonbValue(state, r2, r2 != WJB_END_ARRAY ? &v2 : NULL);
4640                 }
4641                 else
4642                 {
4643                         while ((r1 = JsonbIteratorNext(it_array, &v1, true)) != WJB_END_ARRAY)
4644                                 pushJsonbValue(state, r1, &v1);
4645
4646                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
4647                         while ((r2 = JsonbIteratorNext(it_object, &v2, true)) != WJB_DONE)
4648                                 pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
4649
4650                         res = pushJsonbValue(state, WJB_END_ARRAY, NULL);
4651                 }
4652         }
4653         else
4654         {
4655                 /*
4656                  * This must be scalar || object or object || scalar, as that's all
4657                  * that's left. Both of these make no sense, so error out.
4658                  */
4659                 ereport(ERROR,
4660                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4661                                  errmsg("invalid concatenation of jsonb objects")));
4662         }
4663
4664         return res;
4665 }
4666
4667 /*
4668  * Do most of the heavy work for jsonb_set/jsonb_insert
4669  *
4670  * If JB_PATH_DELETE bit is set in op_type, the element is to be removed.
4671  *
4672  * If any bit mentioned in JB_PATH_CREATE_OR_INSERT is set in op_type,
4673  * we create the new value if the key or array index does not exist.
4674  *
4675  * Bits JB_PATH_INSERT_BEFORE and JB_PATH_INSERT_AFTER in op_type
4676  * behave as JB_PATH_CREATE if new value is inserted in JsonbObject.
4677  *
4678  * All path elements before the last must already exist
4679  * whatever bits in op_type are set, or nothing is done.
4680  */
4681 static JsonbValue *
4682 setPath(JsonbIterator **it, Datum *path_elems,
4683                 bool *path_nulls, int path_len,
4684                 JsonbParseState **st, int level, Jsonb *newval, int op_type)
4685 {
4686         JsonbValue      v;
4687         JsonbIteratorToken r;
4688         JsonbValue *res;
4689
4690         check_stack_depth();
4691
4692         if (path_nulls[level])
4693                 ereport(ERROR,
4694                                 (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
4695                                  errmsg("path element at position %d is null",
4696                                                 level + 1)));
4697
4698         r = JsonbIteratorNext(it, &v, false);
4699
4700         switch (r)
4701         {
4702                 case WJB_BEGIN_ARRAY:
4703                         (void) pushJsonbValue(st, r, NULL);
4704                         setPathArray(it, path_elems, path_nulls, path_len, st, level,
4705                                                  newval, v.val.array.nElems, op_type);
4706                         r = JsonbIteratorNext(it, &v, false);
4707                         Assert(r == WJB_END_ARRAY);
4708                         res = pushJsonbValue(st, r, NULL);
4709                         break;
4710                 case WJB_BEGIN_OBJECT:
4711                         (void) pushJsonbValue(st, r, NULL);
4712                         setPathObject(it, path_elems, path_nulls, path_len, st, level,
4713                                                   newval, v.val.object.nPairs, op_type);
4714                         r = JsonbIteratorNext(it, &v, true);
4715                         Assert(r == WJB_END_OBJECT);
4716                         res = pushJsonbValue(st, r, NULL);
4717                         break;
4718                 case WJB_ELEM:
4719                 case WJB_VALUE:
4720                         res = pushJsonbValue(st, r, &v);
4721                         break;
4722                 default:
4723                         elog(ERROR, "unrecognized iterator result: %d", (int) r);
4724                         res = NULL;                     /* keep compiler quiet */
4725                         break;
4726         }
4727
4728         return res;
4729 }
4730
4731 /*
4732  * Object walker for setPath
4733  */
4734 static void
4735 setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
4736                           int path_len, JsonbParseState **st, int level,
4737                           Jsonb *newval, uint32 npairs, int op_type)
4738 {
4739         JsonbValue      v;
4740         int                     i;
4741         JsonbValue      k;
4742         bool            done = false;
4743
4744         if (level >= path_len || path_nulls[level])
4745                 done = true;
4746
4747         /* empty object is a special case for create */
4748         if ((npairs == 0) && (op_type & JB_PATH_CREATE_OR_INSERT) &&
4749                 (level == path_len - 1))
4750         {
4751                 JsonbValue      newkey;
4752
4753                 newkey.type = jbvString;
4754                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
4755                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
4756
4757                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
4758                 addJsonbToParseState(st, newval);
4759         }
4760
4761         for (i = 0; i < npairs; i++)
4762         {
4763                 JsonbIteratorToken r = JsonbIteratorNext(it, &k, true);
4764
4765                 Assert(r == WJB_KEY);
4766
4767                 if (!done &&
4768                         k.val.string.len == VARSIZE_ANY_EXHDR(path_elems[level]) &&
4769                         memcmp(k.val.string.val, VARDATA_ANY(path_elems[level]),
4770                                    k.val.string.len) == 0)
4771                 {
4772                         if (level == path_len - 1)
4773                         {
4774                                 /*
4775                                  * called from jsonb_insert(), it forbids redefining an
4776                                  * existing value
4777                                  */
4778                                 if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER))
4779                                         ereport(ERROR,
4780                                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4781                                                          errmsg("cannot replace existing key"),
4782                                                          errhint("Try using the function jsonb_set "
4783                                                                          "to replace key value.")));
4784
4785                                 r = JsonbIteratorNext(it, &v, true);    /* skip value */
4786                                 if (!(op_type & JB_PATH_DELETE))
4787                                 {
4788                                         (void) pushJsonbValue(st, WJB_KEY, &k);
4789                                         addJsonbToParseState(st, newval);
4790                                 }
4791                                 done = true;
4792                         }
4793                         else
4794                         {
4795                                 (void) pushJsonbValue(st, r, &k);
4796                                 setPath(it, path_elems, path_nulls, path_len,
4797                                                 st, level + 1, newval, op_type);
4798                         }
4799                 }
4800                 else
4801                 {
4802                         if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
4803                                 level == path_len - 1 && i == npairs - 1)
4804                         {
4805                                 JsonbValue      newkey;
4806
4807                                 newkey.type = jbvString;
4808                                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
4809                                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
4810
4811                                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
4812                                 addJsonbToParseState(st, newval);
4813                         }
4814
4815                         (void) pushJsonbValue(st, r, &k);
4816                         r = JsonbIteratorNext(it, &v, false);
4817                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4818                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4819                         {
4820                                 int                     walking_level = 1;
4821
4822                                 while (walking_level != 0)
4823                                 {
4824                                         r = JsonbIteratorNext(it, &v, false);
4825
4826                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4827                                                 ++walking_level;
4828                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
4829                                                 --walking_level;
4830
4831                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4832                                 }
4833                         }
4834                 }
4835         }
4836 }
4837
4838 /*
4839  * Array walker for setPath
4840  */
4841 static void
4842 setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
4843                          int path_len, JsonbParseState **st, int level,
4844                          Jsonb *newval, uint32 nelems, int op_type)
4845 {
4846         JsonbValue      v;
4847         int                     idx,
4848                                 i;
4849         bool            done = false;
4850
4851         /* pick correct index */
4852         if (level < path_len && !path_nulls[level])
4853         {
4854                 char       *c = TextDatumGetCString(path_elems[level]);
4855                 long            lindex;
4856                 char       *badp;
4857
4858                 errno = 0;
4859                 lindex = strtol(c, &badp, 10);
4860                 if (errno != 0 || badp == c || *badp != '\0' || lindex > INT_MAX ||
4861                         lindex < INT_MIN)
4862                         ereport(ERROR,
4863                                         (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
4864                                          errmsg("path element at position %d is not an integer: \"%s\"",
4865                                                         level + 1, c)));
4866                 idx = lindex;
4867         }
4868         else
4869                 idx = nelems;
4870
4871         if (idx < 0)
4872         {
4873                 if (-idx > nelems)
4874                         idx = INT_MIN;
4875                 else
4876                         idx = nelems + idx;
4877         }
4878
4879         if (idx > 0 && idx > nelems)
4880                 idx = nelems;
4881
4882         /*
4883          * if we're creating, and idx == INT_MIN, we prepend the new value to the
4884          * array also if the array is empty - in which case we don't really care
4885          * what the idx value is
4886          */
4887
4888         if ((idx == INT_MIN || nelems == 0) && (level == path_len - 1) &&
4889                 (op_type & JB_PATH_CREATE_OR_INSERT))
4890         {
4891                 Assert(newval != NULL);
4892                 addJsonbToParseState(st, newval);
4893                 done = true;
4894         }
4895
4896         /* iterate over the array elements */
4897         for (i = 0; i < nelems; i++)
4898         {
4899                 JsonbIteratorToken r;
4900
4901                 if (i == idx && level < path_len)
4902                 {
4903                         if (level == path_len - 1)
4904                         {
4905                                 r = JsonbIteratorNext(it, &v, true);    /* skip */
4906
4907                                 if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_CREATE))
4908                                         addJsonbToParseState(st, newval);
4909
4910                                 /*
4911                                  * We should keep current value only in case of
4912                                  * JB_PATH_INSERT_BEFORE or JB_PATH_INSERT_AFTER because
4913                                  * otherwise it should be deleted or replaced
4914                                  */
4915                                 if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_INSERT_BEFORE))
4916                                         (void) pushJsonbValue(st, r, &v);
4917
4918                                 if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_REPLACE))
4919                                         addJsonbToParseState(st, newval);
4920
4921                                 done = true;
4922                         }
4923                         else
4924                                 (void) setPath(it, path_elems, path_nulls, path_len,
4925                                                            st, level + 1, newval, op_type);
4926                 }
4927                 else
4928                 {
4929                         r = JsonbIteratorNext(it, &v, false);
4930
4931                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4932
4933                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4934                         {
4935                                 int                     walking_level = 1;
4936
4937                                 while (walking_level != 0)
4938                                 {
4939                                         r = JsonbIteratorNext(it, &v, false);
4940
4941                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
4942                                                 ++walking_level;
4943                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
4944                                                 --walking_level;
4945
4946                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
4947                                 }
4948                         }
4949
4950                         if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
4951                                 level == path_len - 1 && i == nelems - 1)
4952                         {
4953                                 addJsonbToParseState(st, newval);
4954                         }
4955                 }
4956         }
4957 }
4958
4959 /*
4960  * Parse information about what elements of a jsonb document we want to iterate
4961  * in functions iterate_json(b)_values. This information is presented in jsonb
4962  * format, so that it can be easily extended in the future.
4963  */
4964 uint32
4965 parse_jsonb_index_flags(Jsonb *jb)
4966 {
4967         JsonbIterator *it;
4968         JsonbValue      v;
4969         JsonbIteratorToken type;
4970         uint32          flags = 0;
4971
4972         it = JsonbIteratorInit(&jb->root);
4973
4974         type = JsonbIteratorNext(&it, &v, false);
4975
4976         /*
4977          * We iterate over array (scalar internally is represented as array, so,
4978          * we will accept it too) to check all its elements.  Flag names are
4979          * chosen the same as jsonb_typeof uses.
4980          */
4981         if (type != WJB_BEGIN_ARRAY)
4982                 ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4983                                                 errmsg("wrong flag type, only arrays and scalars are allowed")));
4984
4985         while ((type = JsonbIteratorNext(&it, &v, false)) == WJB_ELEM)
4986         {
4987                 if (v.type != jbvString)
4988                         ereport(ERROR,
4989                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
4990                                          errmsg("flag array element is not a string"),
4991                                          errhint("Possible values are: \"string\", \"numeric\", \"boolean\", \"key\", and \"all\".")));
4992
4993                 if (v.val.string.len == 3 &&
4994                         pg_strncasecmp(v.val.string.val, "all", 3) == 0)
4995                         flags |= jtiAll;
4996                 else if (v.val.string.len == 3 &&
4997                                  pg_strncasecmp(v.val.string.val, "key", 3) == 0)
4998                         flags |= jtiKey;
4999                 else if (v.val.string.len == 6 &&
5000                                  pg_strncasecmp(v.val.string.val, "string", 5) == 0)
5001                         flags |= jtiString;
5002                 else if (v.val.string.len == 7 &&
5003                                  pg_strncasecmp(v.val.string.val, "numeric", 7) == 0)
5004                         flags |= jtiNumeric;
5005                 else if (v.val.string.len == 7 &&
5006                                  pg_strncasecmp(v.val.string.val, "boolean", 7) == 0)
5007                         flags |= jtiBool;
5008                 else
5009                         ereport(ERROR,
5010                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
5011                                          errmsg("wrong flag in flag array: \"%s\"",
5012                                                         pnstrdup(v.val.string.val, v.val.string.len)),
5013                                          errhint("Possible values are: \"string\", \"numeric\", \"boolean\", \"key\", and \"all\".")));
5014         }
5015
5016         /* expect end of array now */
5017         if (type != WJB_END_ARRAY)
5018                 elog(ERROR, "unexpected end of flag array");
5019
5020         /* get final WJB_DONE and free iterator */
5021         type = JsonbIteratorNext(&it, &v, false);
5022         if (type != WJB_DONE)
5023                 elog(ERROR, "unexpected end of flag array");
5024
5025         return flags;
5026 }
5027
5028 /*
5029  * Iterate over jsonb values or elements, specified by flags, and pass them
5030  * together with an iteration state to a specified JsonIterateStringValuesAction.
5031  */
5032 void
5033 iterate_jsonb_values(Jsonb *jb, uint32 flags, void *state,
5034                                          JsonIterateStringValuesAction action)
5035 {
5036         JsonbIterator *it;
5037         JsonbValue      v;
5038         JsonbIteratorToken type;
5039
5040         it = JsonbIteratorInit(&jb->root);
5041
5042         /*
5043          * Just recursively iterating over jsonb and call callback on all
5044          * corresponding elements
5045          */
5046         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
5047         {
5048                 if (type == WJB_KEY)
5049                 {
5050                         if (flags & jtiKey)
5051                                 action(state, v.val.string.val, v.val.string.len);
5052
5053                         continue;
5054                 }
5055                 else if (!(type == WJB_VALUE || type == WJB_ELEM))
5056                 {
5057                         /* do not call callback for composite JsonbValue */
5058                         continue;
5059                 }
5060
5061                 /* JsonbValue is a value of object or element of array */
5062                 switch (v.type)
5063                 {
5064                         case jbvString:
5065                                 if (flags & jtiString)
5066                                         action(state, v.val.string.val, v.val.string.len);
5067                                 break;
5068                         case jbvNumeric:
5069                                 if (flags & jtiNumeric)
5070                                 {
5071                                         char       *val;
5072
5073                                         val = DatumGetCString(DirectFunctionCall1(numeric_out,
5074                                                                                                                           NumericGetDatum(v.val.numeric)));
5075
5076                                         action(state, val, strlen(val));
5077                                         pfree(val);
5078                                 }
5079                                 break;
5080                         case jbvBool:
5081                                 if (flags & jtiBool)
5082                                 {
5083                                         if (v.val.boolean)
5084                                                 action(state, "true", 4);
5085                                         else
5086                                                 action(state, "false", 5);
5087                                 }
5088                                 break;
5089                         default:
5090                                 /* do not call callback for composite JsonbValue */
5091                                 break;
5092                 }
5093         }
5094 }
5095
5096 /*
5097  * Iterate over json values and elements, specified by flags, and pass them
5098  * together with an iteration state to a specified JsonIterateStringValuesAction.
5099  */
5100 void
5101 iterate_json_values(text *json, uint32 flags, void *action_state,
5102                                         JsonIterateStringValuesAction action)
5103 {
5104         JsonLexContext *lex = makeJsonLexContext(json, true);
5105         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
5106         IterateJsonStringValuesState *state = palloc0(sizeof(IterateJsonStringValuesState));
5107
5108         state->lex = lex;
5109         state->action = action;
5110         state->action_state = action_state;
5111         state->flags = flags;
5112
5113         sem->semstate = (void *) state;
5114         sem->scalar = iterate_values_scalar;
5115         sem->object_field_start = iterate_values_object_field_start;
5116
5117         pg_parse_json(lex, sem);
5118 }
5119
5120 /*
5121  * An auxiliary function for iterate_json_values to invoke a specified
5122  * JsonIterateStringValuesAction for specified values.
5123  */
5124 static void
5125 iterate_values_scalar(void *state, char *token, JsonTokenType tokentype)
5126 {
5127         IterateJsonStringValuesState *_state = (IterateJsonStringValuesState *) state;
5128
5129         switch (tokentype)
5130         {
5131                 case JSON_TOKEN_STRING:
5132                         if (_state->flags & jtiString)
5133                                 _state->action(_state->action_state, token, strlen(token));
5134                         break;
5135                 case JSON_TOKEN_NUMBER:
5136                         if (_state->flags & jtiNumeric)
5137                                 _state->action(_state->action_state, token, strlen(token));
5138                         break;
5139                 case JSON_TOKEN_TRUE:
5140                 case JSON_TOKEN_FALSE:
5141                         if (_state->flags & jtiBool)
5142                                 _state->action(_state->action_state, token, strlen(token));
5143                         break;
5144                 default:
5145                         /* do not call callback for any other token */
5146                         break;
5147         }
5148 }
5149
5150 static void
5151 iterate_values_object_field_start(void *state, char *fname, bool isnull)
5152 {
5153         IterateJsonStringValuesState *_state = (IterateJsonStringValuesState *) state;
5154
5155         if (_state->flags & jtiKey)
5156         {
5157                 char       *val = pstrdup(fname);
5158
5159                 _state->action(_state->action_state, val, strlen(val));
5160         }
5161 }
5162
5163 /*
5164  * Iterate over a jsonb, and apply a specified JsonTransformStringValuesAction
5165  * to every string value or element. Any necessary context for a
5166  * JsonTransformStringValuesAction can be passed in the action_state variable.
5167  * Function returns a copy of an original jsonb object with transformed values.
5168  */
5169 Jsonb *
5170 transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
5171                                                           JsonTransformStringValuesAction transform_action)
5172 {
5173         JsonbIterator *it;
5174         JsonbValue      v,
5175                            *res = NULL;
5176         JsonbIteratorToken type;
5177         JsonbParseState *st = NULL;
5178         text       *out;
5179         bool            is_scalar = false;
5180
5181         it = JsonbIteratorInit(&jsonb->root);
5182         is_scalar = it->isScalar;
5183
5184         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
5185         {
5186                 if ((type == WJB_VALUE || type == WJB_ELEM) && v.type == jbvString)
5187                 {
5188                         out = transform_action(action_state, v.val.string.val, v.val.string.len);
5189                         v.val.string.val = VARDATA_ANY(out);
5190                         v.val.string.len = VARSIZE_ANY_EXHDR(out);
5191                         res = pushJsonbValue(&st, type, type < WJB_BEGIN_ARRAY ? &v : NULL);
5192                 }
5193                 else
5194                 {
5195                         res = pushJsonbValue(&st, type, (type == WJB_KEY ||
5196                                                                                          type == WJB_VALUE ||
5197                                                                                          type == WJB_ELEM) ? &v : NULL);
5198                 }
5199         }
5200
5201         if (res->type == jbvArray)
5202                 res->val.array.rawScalar = is_scalar;
5203
5204         return JsonbValueToJsonb(res);
5205 }
5206
5207 /*
5208  * Iterate over a json, and apply a specified JsonTransformStringValuesAction
5209  * to every string value or element. Any necessary context for a
5210  * JsonTransformStringValuesAction can be passed in the action_state variable.
5211  * Function returns a StringInfo, which is a copy of an original json with
5212  * transformed values.
5213  */
5214 text *
5215 transform_json_string_values(text *json, void *action_state,
5216                                                          JsonTransformStringValuesAction transform_action)
5217 {
5218         JsonLexContext *lex = makeJsonLexContext(json, true);
5219         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
5220         TransformJsonStringValuesState *state = palloc0(sizeof(TransformJsonStringValuesState));
5221
5222         state->lex = lex;
5223         state->strval = makeStringInfo();
5224         state->action = transform_action;
5225         state->action_state = action_state;
5226
5227         sem->semstate = (void *) state;
5228         sem->scalar = transform_string_values_scalar;
5229         sem->object_start = transform_string_values_object_start;
5230         sem->object_end = transform_string_values_object_end;
5231         sem->array_start = transform_string_values_array_start;
5232         sem->array_end = transform_string_values_array_end;
5233         sem->scalar = transform_string_values_scalar;
5234         sem->array_element_start = transform_string_values_array_element_start;
5235         sem->object_field_start = transform_string_values_object_field_start;
5236
5237         pg_parse_json(lex, sem);
5238
5239         return cstring_to_text_with_len(state->strval->data, state->strval->len);
5240 }
5241
5242 /*
5243  * Set of auxiliary functions for transform_json_string_values to invoke a
5244  * specified JsonTransformStringValuesAction for all values and left everything
5245  * else untouched.
5246  */
5247 static void
5248 transform_string_values_object_start(void *state)
5249 {
5250         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5251
5252         appendStringInfoCharMacro(_state->strval, '{');
5253 }
5254
5255 static void
5256 transform_string_values_object_end(void *state)
5257 {
5258         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5259
5260         appendStringInfoCharMacro(_state->strval, '}');
5261 }
5262
5263 static void
5264 transform_string_values_array_start(void *state)
5265 {
5266         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5267
5268         appendStringInfoCharMacro(_state->strval, '[');
5269 }
5270
5271 static void
5272 transform_string_values_array_end(void *state)
5273 {
5274         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5275
5276         appendStringInfoCharMacro(_state->strval, ']');
5277 }
5278
5279 static void
5280 transform_string_values_object_field_start(void *state, char *fname, bool isnull)
5281 {
5282         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5283
5284         if (_state->strval->data[_state->strval->len - 1] != '{')
5285                 appendStringInfoCharMacro(_state->strval, ',');
5286
5287         /*
5288          * Unfortunately we don't have the quoted and escaped string any more, so
5289          * we have to re-escape it.
5290          */
5291         escape_json(_state->strval, fname);
5292         appendStringInfoCharMacro(_state->strval, ':');
5293 }
5294
5295 static void
5296 transform_string_values_array_element_start(void *state, bool isnull)
5297 {
5298         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5299
5300         if (_state->strval->data[_state->strval->len - 1] != '[')
5301                 appendStringInfoCharMacro(_state->strval, ',');
5302 }
5303
5304 static void
5305 transform_string_values_scalar(void *state, char *token, JsonTokenType tokentype)
5306 {
5307         TransformJsonStringValuesState *_state = (TransformJsonStringValuesState *) state;
5308
5309         if (tokentype == JSON_TOKEN_STRING)
5310         {
5311                 text       *out = _state->action(_state->action_state, token, strlen(token));
5312
5313                 escape_json(_state->strval, text_to_cstring(out));
5314         }
5315         else
5316                 appendStringInfoString(_state->strval, token);
5317 }