granicus.if.org Git - postgresql/blob - src/backend/utils/adt/jsonfuncs.c

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