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-2015, 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/typcache.h"
  35
  36 /* semantic action functions for json_object_keys */
  37 static void okeys_object_field_start(void *state, char *fname, bool isnull);
  38 static void okeys_array_start(void *state);
  39 static void okeys_scalar(void *state, char *token, JsonTokenType tokentype);
  40
  41 /* semantic action functions for json_get* functions */
  42 static void get_object_start(void *state);
  43 static void get_object_end(void *state);
  44 static void get_object_field_start(void *state, char *fname, bool isnull);
  45 static void get_object_field_end(void *state, char *fname, bool isnull);
  46 static void get_array_start(void *state);
  47 static void get_array_end(void *state);
  48 static void get_array_element_start(void *state, bool isnull);
  49 static void get_array_element_end(void *state, bool isnull);
  50 static void get_scalar(void *state, char *token, JsonTokenType tokentype);
  51
  52 /* common worker function for json getter functions */
  53 static Datum get_path_all(FunctionCallInfo fcinfo, bool as_text);
  54 static text *get_worker(text *json, char **tpath, int *ipath, int npath,
  55                    bool normalize_results);
  56 static Datum get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text);
  57
  58 /* semantic action functions for json_array_length */
  59 static void alen_object_start(void *state);
  60 static void alen_scalar(void *state, char *token, JsonTokenType tokentype);
  61 static void alen_array_element_start(void *state, bool isnull);
  62
  63 /* common workers for json{b}_each* functions */
  64 static Datum each_worker(FunctionCallInfo fcinfo, bool as_text);
  65 static Datum each_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
  66                                   bool as_text);
  67
  68 /* semantic action functions for json_each */
  69 static void each_object_field_start(void *state, char *fname, bool isnull);
  70 static void each_object_field_end(void *state, char *fname, bool isnull);
  71 static void each_array_start(void *state);
  72 static void each_scalar(void *state, char *token, JsonTokenType tokentype);
  73
  74 /* common workers for json{b}_array_elements_* functions */
  75 static Datum elements_worker(FunctionCallInfo fcinfo, const char *funcname,
  76                                 bool as_text);
  77 static Datum elements_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
  78                                           bool as_text);
  79
  80 /* semantic action functions for json_array_elements */
  81 static void elements_object_start(void *state);
  82 static void elements_array_element_start(void *state, bool isnull);
  83 static void elements_array_element_end(void *state, bool isnull);
  84 static void elements_scalar(void *state, char *token, JsonTokenType tokentype);
  85
  86 /* turn a json object into a hash table */
  87 static HTAB *get_json_object_as_hash(text *json, const char *funcname);
  88
  89 /* common worker for populate_record and to_record */
  90 static Datum populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
  91                                            bool have_record_arg);
  92
  93 /* semantic action functions for get_json_object_as_hash */
  94 static void hash_object_field_start(void *state, char *fname, bool isnull);
  95 static void hash_object_field_end(void *state, char *fname, bool isnull);
  96 static void hash_array_start(void *state);
  97 static void hash_scalar(void *state, char *token, JsonTokenType tokentype);
  98
  99 /* semantic action functions for populate_recordset */
 100 static void populate_recordset_object_field_start(void *state, char *fname, bool isnull);
 101 static void populate_recordset_object_field_end(void *state, char *fname, bool isnull);
 102 static void populate_recordset_scalar(void *state, char *token, JsonTokenType tokentype);
 103 static void populate_recordset_object_start(void *state);
 104 static void populate_recordset_object_end(void *state);
 105 static void populate_recordset_array_start(void *state);
 106 static void populate_recordset_array_element_start(void *state, bool isnull);
 107
 108 /* semantic action functions for json_strip_nulls */
 109 static void sn_object_start(void *state);
 110 static void sn_object_end(void *state);
 111 static void sn_array_start(void *state);
 112 static void sn_array_end(void *state);
 113 static void sn_object_field_start(void *state, char *fname, bool isnull);
 114 static void sn_array_element_start(void *state, bool isnull);
 115 static void sn_scalar(void *state, char *token, JsonTokenType tokentype);
 116
 117 /* worker function for populate_recordset and to_recordset */
 118 static Datum populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
 119                                                   bool have_record_arg);
 120
 121 /* Worker that takes care of common setup for us */
 122 static JsonbValue *findJsonbValueFromContainerLen(JsonbContainer *container,
 123                                                            uint32 flags,
 124                                                            char *key,
 125                                                            uint32 keylen);
 126
 127 /* functions supporting jsonb_delete, jsonb_set and jsonb_concat */
 128 static JsonbValue *IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 129                            JsonbParseState **state);
 130 static JsonbValue *setPath(JsonbIterator **it, Datum *path_elems,
 131                 bool *path_nulls, int path_len,
 132                 JsonbParseState **st, int level, Jsonb *newval,
 133                 bool create);
 134 static void setPathObject(JsonbIterator **it, Datum *path_elems,
 135                           bool *path_nulls, int path_len, JsonbParseState **st,
 136                           int level,
 137                           Jsonb *newval, uint32 npairs, bool create);
 138 static void setPathArray(JsonbIterator **it, Datum *path_elems,
 139                          bool *path_nulls, int path_len, JsonbParseState **st,
 140                          int level, Jsonb *newval, uint32 nelems, bool create);
 141 static void addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb);
 142
 143 /* state for json_object_keys */
 144 typedef struct OkeysState
 145 {
 146         JsonLexContext *lex;
 147         char      **result;
 148         int                     result_size;
 149         int                     result_count;
 150         int                     sent_count;
 151 } OkeysState;
 152
 153 /* state for json_get* functions */
 154 typedef struct GetState
 155 {
 156         JsonLexContext *lex;
 157         text       *tresult;
 158         char       *result_start;
 159         bool            normalize_results;
 160         bool            next_scalar;
 161         int                     npath;                  /* length of each path-related array */
 162         char      **path_names;         /* field name(s) being sought */
 163         int                *path_indexes;       /* array index(es) being sought */
 164         bool       *pathok;                     /* is path matched to current depth? */
 165         int                *array_cur_index;    /* current element index at each path level */
 166 } GetState;
 167
 168 /* state for json_array_length */
 169 typedef struct AlenState
 170 {
 171         JsonLexContext *lex;
 172         int                     count;
 173 } AlenState;
 174
 175 /* state for json_each */
 176 typedef struct EachState
 177 {
 178         JsonLexContext *lex;
 179         Tuplestorestate *tuple_store;
 180         TupleDesc       ret_tdesc;
 181         MemoryContext tmp_cxt;
 182         char       *result_start;
 183         bool            normalize_results;
 184         bool            next_scalar;
 185         char       *normalized_scalar;
 186 } EachState;
 187
 188 /* state for json_array_elements */
 189 typedef struct ElementsState
 190 {
 191         JsonLexContext *lex;
 192         const char *function_name;
 193         Tuplestorestate *tuple_store;
 194         TupleDesc       ret_tdesc;
 195         MemoryContext tmp_cxt;
 196         char       *result_start;
 197         bool            normalize_results;
 198         bool            next_scalar;
 199         char       *normalized_scalar;
 200 } ElementsState;
 201
 202 /* state for get_json_object_as_hash */
 203 typedef struct JhashState
 204 {
 205         JsonLexContext *lex;
 206         const char *function_name;
 207         HTAB       *hash;
 208         char       *saved_scalar;
 209         char       *save_json_start;
 210 } JHashState;
 211
 212 /* hashtable element */
 213 typedef struct JsonHashEntry
 214 {
 215         char            fname[NAMEDATALEN];             /* hash key (MUST BE FIRST) */
 216         char       *val;
 217         char       *json;
 218         bool            isnull;
 219 } JsonHashEntry;
 220
 221 /* these two are stolen from hstore / record_out, used in populate_record* */
 222 typedef struct ColumnIOData
 223 {
 224         Oid                     column_type;
 225         Oid                     typiofunc;
 226         Oid                     typioparam;
 227         FmgrInfo        proc;
 228 } ColumnIOData;
 229
 230 typedef struct RecordIOData
 231 {
 232         Oid                     record_type;
 233         int32           record_typmod;
 234         int                     ncolumns;
 235         ColumnIOData columns[FLEXIBLE_ARRAY_MEMBER];
 236 } RecordIOData;
 237
 238 /* state for populate_recordset */
 239 typedef struct PopulateRecordsetState
 240 {
 241         JsonLexContext *lex;
 242         const char *function_name;
 243         HTAB       *json_hash;
 244         char       *saved_scalar;
 245         char       *save_json_start;
 246         Tuplestorestate *tuple_store;
 247         TupleDesc       ret_tdesc;
 248         HeapTupleHeader rec;
 249         RecordIOData *my_extra;
 250         MemoryContext fn_mcxt;          /* used to stash IO funcs */
 251 } PopulateRecordsetState;
 252
 253 /* state for json_strip_nulls */
 254 typedef struct StripnullState
 255 {
 256         JsonLexContext *lex;
 257         StringInfo      strval;
 258         bool            skip_next_null;
 259 } StripnullState;
 260
 261 /* Turn a jsonb object into a record */
 262 static void make_row_from_rec_and_jsonb(Jsonb *element,
 263                                                         PopulateRecordsetState *state);
 264
 265 /*
 266  * SQL function json_object_keys
 267  *
 268  * Returns the set of keys for the object argument.
 269  *
 270  * This SRF operates in value-per-call mode. It processes the
 271  * object during the first call, and the keys are simply stashed
 272  * in an array, whose size is expanded as necessary. This is probably
 273  * safe enough for a list of keys of a single object, since they are
 274  * limited in size to NAMEDATALEN and the number of keys is unlikely to
 275  * be so huge that it has major memory implications.
 276  */
 277 Datum
 278 jsonb_object_keys(PG_FUNCTION_ARGS)
 279 {
 280         FuncCallContext *funcctx;
 281         OkeysState *state;
 282         int                     i;
 283
 284         if (SRF_IS_FIRSTCALL())
 285         {
 286                 MemoryContext oldcontext;
 287                 Jsonb      *jb = PG_GETARG_JSONB(0);
 288                 bool            skipNested = false;
 289                 JsonbIterator *it;
 290                 JsonbValue      v;
 291                 int                     r;
 292
 293                 if (JB_ROOT_IS_SCALAR(jb))
 294                         ereport(ERROR,
 295                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 296                                          errmsg("cannot call %s on a scalar",
 297                                                         "jsonb_object_keys")));
 298                 else if (JB_ROOT_IS_ARRAY(jb))
 299                         ereport(ERROR,
 300                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 301                                          errmsg("cannot call %s on an array",
 302                                                         "jsonb_object_keys")));
 303
 304                 funcctx = SRF_FIRSTCALL_INIT();
 305                 oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 306
 307                 state = palloc(sizeof(OkeysState));
 308
 309                 state->result_size = JB_ROOT_COUNT(jb);
 310                 state->result_count = 0;
 311                 state->sent_count = 0;
 312                 state->result = palloc(state->result_size * sizeof(char *));
 313
 314                 it = JsonbIteratorInit(&jb->root);
 315
 316                 while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
 317                 {
 318                         skipNested = true;
 319
 320                         if (r == WJB_KEY)
 321                         {
 322                                 char       *cstr;
 323
 324                                 cstr = palloc(v.val.string.len + 1 * sizeof(char));
 325                                 memcpy(cstr, v.val.string.val, v.val.string.len);
 326                                 cstr[v.val.string.len] = '\0';
 327                                 state->result[state->result_count++] = cstr;
 328                         }
 329                 }
 330
 331                 MemoryContextSwitchTo(oldcontext);
 332                 funcctx->user_fctx = (void *) state;
 333         }
 334
 335         funcctx = SRF_PERCALL_SETUP();
 336         state = (OkeysState *) funcctx->user_fctx;
 337
 338         if (state->sent_count < state->result_count)
 339         {
 340                 char       *nxt = state->result[state->sent_count++];
 341
 342                 SRF_RETURN_NEXT(funcctx, CStringGetTextDatum(nxt));
 343         }
 344
 345         /* cleanup to reduce or eliminate memory leaks */
 346         for (i = 0; i < state->result_count; i++)
 347                 pfree(state->result[i]);
 348         pfree(state->result);
 349         pfree(state);
 350
 351         SRF_RETURN_DONE(funcctx);
 352 }
 353
 354
 355 Datum
 356 json_object_keys(PG_FUNCTION_ARGS)
 357 {
 358         FuncCallContext *funcctx;
 359         OkeysState *state;
 360         int                     i;
 361
 362         if (SRF_IS_FIRSTCALL())
 363         {
 364                 text       *json = PG_GETARG_TEXT_P(0);
 365                 JsonLexContext *lex = makeJsonLexContext(json, true);
 366                 JsonSemAction *sem;
 367                 MemoryContext oldcontext;
 368
 369                 funcctx = SRF_FIRSTCALL_INIT();
 370                 oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 371
 372                 state = palloc(sizeof(OkeysState));
 373                 sem = palloc0(sizeof(JsonSemAction));
 374
 375                 state->lex = lex;
 376                 state->result_size = 256;
 377                 state->result_count = 0;
 378                 state->sent_count = 0;
 379                 state->result = palloc(256 * sizeof(char *));
 380
 381                 sem->semstate = (void *) state;
 382                 sem->array_start = okeys_array_start;
 383                 sem->scalar = okeys_scalar;
 384                 sem->object_field_start = okeys_object_field_start;
 385                 /* remainder are all NULL, courtesy of palloc0 above */
 386
 387                 pg_parse_json(lex, sem);
 388                 /* keys are now in state->result */
 389
 390                 pfree(lex->strval->data);
 391                 pfree(lex->strval);
 392                 pfree(lex);
 393                 pfree(sem);
 394
 395                 MemoryContextSwitchTo(oldcontext);
 396                 funcctx->user_fctx = (void *) state;
 397         }
 398
 399         funcctx = SRF_PERCALL_SETUP();
 400         state = (OkeysState *) funcctx->user_fctx;
 401
 402         if (state->sent_count < state->result_count)
 403         {
 404                 char       *nxt = state->result[state->sent_count++];
 405
 406                 SRF_RETURN_NEXT(funcctx, CStringGetTextDatum(nxt));
 407         }
 408
 409         /* cleanup to reduce or eliminate memory leaks */
 410         for (i = 0; i < state->result_count; i++)
 411                 pfree(state->result[i]);
 412         pfree(state->result);
 413         pfree(state);
 414
 415         SRF_RETURN_DONE(funcctx);
 416 }
 417
 418 static void
 419 okeys_object_field_start(void *state, char *fname, bool isnull)
 420 {
 421         OkeysState *_state = (OkeysState *) state;
 422
 423         /* only collecting keys for the top level object */
 424         if (_state->lex->lex_level != 1)
 425                 return;
 426
 427         /* enlarge result array if necessary */
 428         if (_state->result_count >= _state->result_size)
 429         {
 430                 _state->result_size *= 2;
 431                 _state->result = (char **)
 432                         repalloc(_state->result, sizeof(char *) * _state->result_size);
 433         }
 434
 435         /* save a copy of the field name */
 436         _state->result[_state->result_count++] = pstrdup(fname);
 437 }
 438
 439 static void
 440 okeys_array_start(void *state)
 441 {
 442         OkeysState *_state = (OkeysState *) state;
 443
 444         /* top level must be a json object */
 445         if (_state->lex->lex_level == 0)
 446                 ereport(ERROR,
 447                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 448                                  errmsg("cannot call %s on an array",
 449                                                 "json_object_keys")));
 450 }
 451
 452 static void
 453 okeys_scalar(void *state, char *token, JsonTokenType tokentype)
 454 {
 455         OkeysState *_state = (OkeysState *) state;
 456
 457         /* top level must be a json object */
 458         if (_state->lex->lex_level == 0)
 459                 ereport(ERROR,
 460                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 461                                  errmsg("cannot call %s on a scalar",
 462                                                 "json_object_keys")));
 463 }
 464
 465 /*
 466  * json and jsonb getter functions
 467  * these implement the -> ->> #> and #>> operators
 468  * and the json{b?}_extract_path*(json, text, ...) functions
 469  */
 470
 471
 472 Datum
 473 json_object_field(PG_FUNCTION_ARGS)
 474 {
 475         text       *json = PG_GETARG_TEXT_P(0);
 476         text       *fname = PG_GETARG_TEXT_PP(1);
 477         char       *fnamestr = text_to_cstring(fname);
 478         text       *result;
 479
 480         result = get_worker(json, &fnamestr, NULL, 1, false);
 481
 482         if (result != NULL)
 483                 PG_RETURN_TEXT_P(result);
 484         else
 485                 PG_RETURN_NULL();
 486 }
 487
 488 Datum
 489 jsonb_object_field(PG_FUNCTION_ARGS)
 490 {
 491         Jsonb      *jb = PG_GETARG_JSONB(0);
 492         text       *key = PG_GETARG_TEXT_PP(1);
 493         JsonbValue *v;
 494
 495         if (!JB_ROOT_IS_OBJECT(jb))
 496                 PG_RETURN_NULL();
 497
 498         v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
 499                                                                            VARDATA_ANY(key),
 500                                                                            VARSIZE_ANY_EXHDR(key));
 501
 502         if (v != NULL)
 503                 PG_RETURN_JSONB(JsonbValueToJsonb(v));
 504
 505         PG_RETURN_NULL();
 506 }
 507
 508 Datum
 509 json_object_field_text(PG_FUNCTION_ARGS)
 510 {
 511         text       *json = PG_GETARG_TEXT_P(0);
 512         text       *fname = PG_GETARG_TEXT_PP(1);
 513         char       *fnamestr = text_to_cstring(fname);
 514         text       *result;
 515
 516         result = get_worker(json, &fnamestr, NULL, 1, true);
 517
 518         if (result != NULL)
 519                 PG_RETURN_TEXT_P(result);
 520         else
 521                 PG_RETURN_NULL();
 522 }
 523
 524 Datum
 525 jsonb_object_field_text(PG_FUNCTION_ARGS)
 526 {
 527         Jsonb      *jb = PG_GETARG_JSONB(0);
 528         text       *key = PG_GETARG_TEXT_PP(1);
 529         JsonbValue *v;
 530
 531         if (!JB_ROOT_IS_OBJECT(jb))
 532                 PG_RETURN_NULL();
 533
 534         v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
 535                                                                            VARDATA_ANY(key),
 536                                                                            VARSIZE_ANY_EXHDR(key));
 537
 538         if (v != NULL)
 539         {
 540                 text       *result = NULL;
 541
 542                 switch (v->type)
 543                 {
 544                         case jbvNull:
 545                                 break;
 546                         case jbvBool:
 547                                 result = cstring_to_text(v->val.boolean ? "true" : "false");
 548                                 break;
 549                         case jbvString:
 550                                 result = cstring_to_text_with_len(v->val.string.val, v->val.string.len);
 551                                 break;
 552                         case jbvNumeric:
 553                                 result = cstring_to_text(DatumGetCString(DirectFunctionCall1(numeric_out,
 554                                                                                   PointerGetDatum(v->val.numeric))));
 555                                 break;
 556                         case jbvBinary:
 557                                 {
 558                                         StringInfo      jtext = makeStringInfo();
 559
 560                                         (void) JsonbToCString(jtext, v->val.binary.data, -1);
 561                                         result = cstring_to_text_with_len(jtext->data, jtext->len);
 562                                 }
 563                                 break;
 564                         default:
 565                                 elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
 566                 }
 567
 568                 if (result)
 569                         PG_RETURN_TEXT_P(result);
 570         }
 571
 572         PG_RETURN_NULL();
 573 }
 574
 575 Datum
 576 json_array_element(PG_FUNCTION_ARGS)
 577 {
 578         text       *json = PG_GETARG_TEXT_P(0);
 579         int                     element = PG_GETARG_INT32(1);
 580         text       *result;
 581
 582         result = get_worker(json, NULL, &element, 1, false);
 583
 584         if (result != NULL)
 585                 PG_RETURN_TEXT_P(result);
 586         else
 587                 PG_RETURN_NULL();
 588 }
 589
 590 Datum
 591 jsonb_array_element(PG_FUNCTION_ARGS)
 592 {
 593         Jsonb      *jb = PG_GETARG_JSONB(0);
 594         int                     element = PG_GETARG_INT32(1);
 595         JsonbValue *v;
 596
 597         if (!JB_ROOT_IS_ARRAY(jb))
 598                 PG_RETURN_NULL();
 599
 600         /* Handle negative subscript */
 601         if (element < 0)
 602         {
 603                 uint32  nelements = JB_ROOT_COUNT(jb);
 604
 605                 if (-element > nelements)
 606                         PG_RETURN_NULL();
 607                 else
 608                         element += nelements;
 609         }
 610
 611         v = getIthJsonbValueFromContainer(&jb->root, element);
 612         if (v != NULL)
 613                 PG_RETURN_JSONB(JsonbValueToJsonb(v));
 614
 615         PG_RETURN_NULL();
 616 }
 617
 618 Datum
 619 json_array_element_text(PG_FUNCTION_ARGS)
 620 {
 621         text       *json = PG_GETARG_TEXT_P(0);
 622         int                     element = PG_GETARG_INT32(1);
 623         text       *result;
 624
 625         result = get_worker(json, NULL, &element, 1, true);
 626
 627         if (result != NULL)
 628                 PG_RETURN_TEXT_P(result);
 629         else
 630                 PG_RETURN_NULL();
 631 }
 632
 633 Datum
 634 jsonb_array_element_text(PG_FUNCTION_ARGS)
 635 {
 636         Jsonb      *jb = PG_GETARG_JSONB(0);
 637         int                     element = PG_GETARG_INT32(1);
 638         JsonbValue *v;
 639
 640         if (!JB_ROOT_IS_ARRAY(jb))
 641                 PG_RETURN_NULL();
 642
 643         /* Handle negative subscript */
 644         if (element < 0)
 645         {
 646                 uint32  nelements = JB_ROOT_COUNT(jb);
 647
 648                 if (-element > nelements)
 649                         PG_RETURN_NULL();
 650                 else
 651                         element += nelements;
 652         }
 653
 654         v = getIthJsonbValueFromContainer(&jb->root, element);
 655         if (v != NULL)
 656         {
 657                 text       *result = NULL;
 658
 659                 switch (v->type)
 660                 {
 661                         case jbvNull:
 662                                 break;
 663                         case jbvBool:
 664                                 result = cstring_to_text(v->val.boolean ? "true" : "false");
 665                                 break;
 666                         case jbvString:
 667                                 result = cstring_to_text_with_len(v->val.string.val, v->val.string.len);
 668                                 break;
 669                         case jbvNumeric:
 670                                 result = cstring_to_text(DatumGetCString(DirectFunctionCall1(numeric_out,
 671                                                                                   PointerGetDatum(v->val.numeric))));
 672                                 break;
 673                         case jbvBinary:
 674                                 {
 675                                         StringInfo      jtext = makeStringInfo();
 676
 677                                         (void) JsonbToCString(jtext, v->val.binary.data, -1);
 678                                         result = cstring_to_text_with_len(jtext->data, jtext->len);
 679                                 }
 680                                 break;
 681                         default:
 682                                 elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
 683                 }
 684
 685                 if (result)
 686                         PG_RETURN_TEXT_P(result);
 687         }
 688
 689         PG_RETURN_NULL();
 690 }
 691
 692 Datum
 693 json_extract_path(PG_FUNCTION_ARGS)
 694 {
 695         return get_path_all(fcinfo, false);
 696 }
 697
 698 Datum
 699 json_extract_path_text(PG_FUNCTION_ARGS)
 700 {
 701         return get_path_all(fcinfo, true);
 702 }
 703
 704 /*
 705  * common routine for extract_path functions
 706  */
 707 static Datum
 708 get_path_all(FunctionCallInfo fcinfo, bool as_text)
 709 {
 710         text       *json = PG_GETARG_TEXT_P(0);
 711         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
 712         text       *result;
 713         Datum      *pathtext;
 714         bool       *pathnulls;
 715         int                     npath;
 716         char      **tpath;
 717         int                *ipath;
 718         int                     i;
 719
 720         /*
 721          * If the array contains any null elements, return NULL, on the grounds
 722          * that you'd have gotten NULL if any RHS value were NULL in a nested
 723          * series of applications of the -> operator.  (Note: because we also
 724          * return NULL for error cases such as no-such-field, this is true
 725          * regardless of the contents of the rest of the array.)
 726          */
 727         if (array_contains_nulls(path))
 728                 PG_RETURN_NULL();
 729
 730         deconstruct_array(path, TEXTOID, -1, false, 'i',
 731                                           &pathtext, &pathnulls, &npath);
 732
 733         tpath = palloc(npath * sizeof(char *));
 734         ipath = palloc(npath * sizeof(int));
 735
 736         for (i = 0; i < npath; i++)
 737         {
 738                 Assert(!pathnulls[i]);
 739                 tpath[i] = TextDatumGetCString(pathtext[i]);
 740
 741                 /*
 742                  * we have no idea at this stage what structure the document is so
 743                  * just convert anything in the path that we can to an integer and set
 744                  * all the other integers to INT_MIN which will never match.
 745                  */
 746                 if (*tpath[i] != '\0')
 747                 {
 748                         long            ind;
 749                         char       *endptr;
 750
 751                         errno = 0;
 752                         ind = strtol(tpath[i], &endptr, 10);
 753                         if (*endptr == '\0' && errno == 0 && ind <= INT_MAX && ind >= INT_MIN)
 754                                 ipath[i] = (int) ind;
 755                         else
 756                                 ipath[i] = INT_MIN;
 757                 }
 758                 else
 759                         ipath[i] = INT_MIN;
 760         }
 761
 762         result = get_worker(json, tpath, ipath, npath, as_text);
 763
 764         if (result != NULL)
 765                 PG_RETURN_TEXT_P(result);
 766         else
 767                 PG_RETURN_NULL();
 768 }
 769
 770 /*
 771  * get_worker
 772  *
 773  * common worker for all the json getter functions
 774  *
 775  * json: JSON object (in text form)
 776  * tpath[]: field name(s) to extract
 777  * ipath[]: array index(es) (zero-based) to extract, accepts negatives
 778  * npath: length of tpath[] and/or ipath[]
 779  * normalize_results: true to de-escape string and null scalars
 780  *
 781  * tpath can be NULL, or any one tpath[] entry can be NULL, if an object
 782  * field is not to be matched at that nesting level.  Similarly, ipath can
 783  * be NULL, or any one ipath[] entry can be INT_MIN if an array element is
 784  * not to be matched at that nesting level (a json datum should never be
 785  * large enough to have -INT_MIN elements due to MaxAllocSize restriction).
 786  */
 787 static text *
 788 get_worker(text *json,
 789                    char **tpath,
 790                    int *ipath,
 791                    int npath,
 792                    bool normalize_results)
 793 {
 794         JsonLexContext *lex = makeJsonLexContext(json, true);
 795         JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
 796         GetState   *state = palloc0(sizeof(GetState));
 797
 798         Assert(npath >= 0);
 799
 800         state->lex = lex;
 801         /* is it "_as_text" variant? */
 802         state->normalize_results = normalize_results;
 803         state->npath = npath;
 804         state->path_names = tpath;
 805         state->path_indexes = ipath;
 806         state->pathok = palloc0(sizeof(bool) * npath);
 807         state->array_cur_index = palloc(sizeof(int) * npath);
 808
 809         if (npath > 0)
 810                 state->pathok[0] = true;
 811
 812         sem->semstate = (void *) state;
 813
 814         /*
 815          * Not all variants need all the semantic routines. Only set the ones that
 816          * are actually needed for maximum efficiency.
 817          */
 818         sem->scalar = get_scalar;
 819         if (npath == 0)
 820         {
 821                 sem->object_start = get_object_start;
 822                 sem->object_end = get_object_end;
 823                 sem->array_start = get_array_start;
 824                 sem->array_end = get_array_end;
 825         }
 826         if (tpath != NULL)
 827         {
 828                 sem->object_field_start = get_object_field_start;
 829                 sem->object_field_end = get_object_field_end;
 830         }
 831         if (ipath != NULL)
 832         {
 833                 sem->array_start = get_array_start;
 834                 sem->array_element_start = get_array_element_start;
 835                 sem->array_element_end = get_array_element_end;
 836         }
 837
 838         pg_parse_json(lex, sem);
 839
 840         return state->tresult;
 841 }
 842
 843 static void
 844 get_object_start(void *state)
 845 {
 846         GetState   *_state = (GetState *) state;
 847         int                     lex_level = _state->lex->lex_level;
 848
 849         if (lex_level == 0 && _state->npath == 0)
 850         {
 851                 /*
 852                  * Special case: we should match the entire object.  We only need this
 853                  * at outermost level because at nested levels the match will have
 854                  * been started by the outer field or array element callback.
 855                  */
 856                 _state->result_start = _state->lex->token_start;
 857         }
 858 }
 859
 860 static void
 861 get_object_end(void *state)
 862 {
 863         GetState   *_state = (GetState *) state;
 864         int                     lex_level = _state->lex->lex_level;
 865
 866         if (lex_level == 0 && _state->npath == 0)
 867         {
 868                 /* Special case: return the entire object */
 869                 char       *start = _state->result_start;
 870                 int                     len = _state->lex->prev_token_terminator - start;
 871
 872                 _state->tresult = cstring_to_text_with_len(start, len);
 873         }
 874 }
 875
 876 static void
 877 get_object_field_start(void *state, char *fname, bool isnull)
 878 {
 879         GetState   *_state = (GetState *) state;
 880         bool            get_next = false;
 881         int                     lex_level = _state->lex->lex_level;
 882
 883         if (lex_level <= _state->npath &&
 884                 _state->pathok[lex_level - 1] &&
 885                 _state->path_names != NULL &&
 886                 _state->path_names[lex_level - 1] != NULL &&
 887                 strcmp(fname, _state->path_names[lex_level - 1]) == 0)
 888         {
 889                 if (lex_level < _state->npath)
 890                 {
 891                         /* if not at end of path just mark path ok */
 892                         _state->pathok[lex_level] = true;
 893                 }
 894                 else
 895                 {
 896                         /* end of path, so we want this value */
 897                         get_next = true;
 898                 }
 899         }
 900
 901         if (get_next)
 902         {
 903                 /* this object overrides any previous matching object */
 904                 _state->tresult = NULL;
 905                 _state->result_start = NULL;
 906
 907                 if (_state->normalize_results &&
 908                         _state->lex->token_type == JSON_TOKEN_STRING)
 909                 {
 910                         /* for as_text variants, tell get_scalar to set it for us */
 911                         _state->next_scalar = true;
 912                 }
 913                 else
 914                 {
 915                         /* for non-as_text variants, just note the json starting point */
 916                         _state->result_start = _state->lex->token_start;
 917                 }
 918         }
 919 }
 920
 921 static void
 922 get_object_field_end(void *state, char *fname, bool isnull)
 923 {
 924         GetState   *_state = (GetState *) state;
 925         bool            get_last = false;
 926         int                     lex_level = _state->lex->lex_level;
 927
 928         /* same tests as in get_object_field_start */
 929         if (lex_level <= _state->npath &&
 930                 _state->pathok[lex_level - 1] &&
 931                 _state->path_names != NULL &&
 932                 _state->path_names[lex_level - 1] != NULL &&
 933                 strcmp(fname, _state->path_names[lex_level - 1]) == 0)
 934         {
 935                 if (lex_level < _state->npath)
 936                 {
 937                         /* done with this field so reset pathok */
 938                         _state->pathok[lex_level] = false;
 939                 }
 940                 else
 941                 {
 942                         /* end of path, so we want this value */
 943                         get_last = true;
 944                 }
 945         }
 946
 947         /* for as_text scalar case, our work is already done */
 948         if (get_last && _state->result_start != NULL)
 949         {
 950                 /*
 951                  * make a text object from the string from the prevously noted json
 952                  * start up to the end of the previous token (the lexer is by now
 953                  * ahead of us on whatever came after what we're interested in).
 954                  */
 955                 if (isnull && _state->normalize_results)
 956                         _state->tresult = (text *) NULL;
 957                 else
 958                 {
 959                         char       *start = _state->result_start;
 960                         int                     len = _state->lex->prev_token_terminator - start;
 961
 962                         _state->tresult = cstring_to_text_with_len(start, len);
 963                 }
 964
 965                 /* this should be unnecessary but let's do it for cleanliness: */
 966                 _state->result_start = NULL;
 967         }
 968 }
 969
 970 static void
 971 get_array_start(void *state)
 972 {
 973         GetState   *_state = (GetState *) state;
 974         int                     lex_level = _state->lex->lex_level;
 975
 976         if (lex_level < _state->npath)
 977         {
 978                 /* Initialize counting of elements in this array */
 979                 _state->array_cur_index[lex_level] = -1;
 980
 981                 /* INT_MIN value is reserved to represent invalid subscript */
 982                 if (_state->path_indexes[lex_level] < 0 &&
 983                         _state->path_indexes[lex_level] != INT_MIN)
 984                 {
 985                         /* Negative subscript -- convert to positive-wise subscript */
 986                         int             nelements = json_count_array_elements(_state->lex);
 987
 988                         if (-_state->path_indexes[lex_level] <= nelements)
 989                                 _state->path_indexes[lex_level] += nelements;
 990                 }
 991         }
 992         else if (lex_level == 0 && _state->npath == 0)
 993         {
 994                 /*
 995                  * Special case: we should match the entire array.  We only need this
 996                  * at the outermost level because at nested levels the match will
 997                  * have been started by the outer field or array element callback.
 998                  */
 999                 _state->result_start = _state->lex->token_start;
1000         }
1001 }
1002
1003 static void
1004 get_array_end(void *state)
1005 {
1006         GetState   *_state = (GetState *) state;
1007         int                     lex_level = _state->lex->lex_level;
1008
1009         if (lex_level == 0 && _state->npath == 0)
1010         {
1011                 /* Special case: return the entire array */
1012                 char       *start = _state->result_start;
1013                 int                     len = _state->lex->prev_token_terminator - start;
1014
1015                 _state->tresult = cstring_to_text_with_len(start, len);
1016         }
1017 }
1018
1019 static void
1020 get_array_element_start(void *state, bool isnull)
1021 {
1022         GetState   *_state = (GetState *) state;
1023         bool            get_next = false;
1024         int                     lex_level = _state->lex->lex_level;
1025
1026         /* Update array element counter */
1027         if (lex_level <= _state->npath)
1028                 _state->array_cur_index[lex_level - 1]++;
1029
1030         if (lex_level <= _state->npath &&
1031                 _state->pathok[lex_level - 1] &&
1032                 _state->path_indexes != NULL &&
1033                 _state->array_cur_index[lex_level - 1] == _state->path_indexes[lex_level - 1])
1034         {
1035                 if (lex_level < _state->npath)
1036                 {
1037                         /* if not at end of path just mark path ok */
1038                         _state->pathok[lex_level] = true;
1039                 }
1040                 else
1041                 {
1042                         /* end of path, so we want this value */
1043                         get_next = true;
1044                 }
1045         }
1046
1047         /* same logic as for objects */
1048         if (get_next)
1049         {
1050                 _state->tresult = NULL;
1051                 _state->result_start = NULL;
1052
1053                 if (_state->normalize_results &&
1054                         _state->lex->token_type == JSON_TOKEN_STRING)
1055                 {
1056                         _state->next_scalar = true;
1057                 }
1058                 else
1059                 {
1060                         _state->result_start = _state->lex->token_start;
1061                 }
1062         }
1063 }
1064
1065 static void
1066 get_array_element_end(void *state, bool isnull)
1067 {
1068         GetState   *_state = (GetState *) state;
1069         bool            get_last = false;
1070         int                     lex_level = _state->lex->lex_level;
1071
1072         /* same tests as in get_array_element_start */
1073         if (lex_level <= _state->npath &&
1074                 _state->pathok[lex_level - 1] &&
1075                 _state->path_indexes != NULL &&
1076                 _state->array_cur_index[lex_level - 1] == _state->path_indexes[lex_level - 1])
1077         {
1078                 if (lex_level < _state->npath)
1079                 {
1080                         /* done with this element so reset pathok */
1081                         _state->pathok[lex_level] = false;
1082                 }
1083                 else
1084                 {
1085                         /* end of path, so we want this value */
1086                         get_last = true;
1087                 }
1088         }
1089
1090         /* same logic as for objects */
1091         if (get_last && _state->result_start != NULL)
1092         {
1093                 if (isnull && _state->normalize_results)
1094                         _state->tresult = (text *) NULL;
1095                 else
1096                 {
1097                         char       *start = _state->result_start;
1098                         int                     len = _state->lex->prev_token_terminator - start;
1099
1100                         _state->tresult = cstring_to_text_with_len(start, len);
1101                 }
1102
1103                 _state->result_start = NULL;
1104         }
1105 }
1106
1107 static void
1108 get_scalar(void *state, char *token, JsonTokenType tokentype)
1109 {
1110         GetState   *_state = (GetState *) state;
1111         int                     lex_level = _state->lex->lex_level;
1112
1113         /* Check for whole-object match */
1114         if (lex_level == 0 && _state->npath == 0)
1115         {
1116                 if (_state->normalize_results && tokentype == JSON_TOKEN_STRING)
1117                 {
1118                         /* we want the de-escaped string */
1119                         _state->next_scalar = true;
1120                 }
1121                 else if (_state->normalize_results && tokentype == JSON_TOKEN_NULL)
1122                 {
1123                         _state->tresult = (text *) NULL;
1124                 }
1125                 else
1126                 {
1127                         /*
1128                          * This is a bit hokey: we will suppress whitespace after the
1129                          * scalar token, but not whitespace before it.  Probably not worth
1130                          * doing our own space-skipping to avoid that.
1131                          */
1132                         char       *start = _state->lex->input;
1133                         int                     len = _state->lex->prev_token_terminator - start;
1134
1135                         _state->tresult = cstring_to_text_with_len(start, len);
1136                 }
1137         }
1138
1139         if (_state->next_scalar)
1140         {
1141                 /* a de-escaped text value is wanted, so supply it */
1142                 _state->tresult = cstring_to_text(token);
1143                 /* make sure the next call to get_scalar doesn't overwrite it */
1144                 _state->next_scalar = false;
1145         }
1146 }
1147
1148 Datum
1149 jsonb_extract_path(PG_FUNCTION_ARGS)
1150 {
1151         return get_jsonb_path_all(fcinfo, false);
1152 }
1153
1154 Datum
1155 jsonb_extract_path_text(PG_FUNCTION_ARGS)
1156 {
1157         return get_jsonb_path_all(fcinfo, true);
1158 }
1159
1160 static Datum
1161 get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
1162 {
1163         Jsonb      *jb = PG_GETARG_JSONB(0);
1164         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
1165         Jsonb      *res;
1166         Datum      *pathtext;
1167         bool       *pathnulls;
1168         int                     npath;
1169         int                     i;
1170         bool            have_object = false,
1171                                 have_array = false;
1172         JsonbValue *jbvp = NULL;
1173         JsonbValue      tv;
1174         JsonbContainer *container;
1175
1176         /*
1177          * If the array contains any null elements, return NULL, on the grounds
1178          * that you'd have gotten NULL if any RHS value were NULL in a nested
1179          * series of applications of the -> operator.  (Note: because we also
1180          * return NULL for error cases such as no-such-field, this is true
1181          * regardless of the contents of the rest of the array.)
1182          */
1183         if (array_contains_nulls(path))
1184                 PG_RETURN_NULL();
1185
1186         deconstruct_array(path, TEXTOID, -1, false, 'i',
1187                                           &pathtext, &pathnulls, &npath);
1188
1189         /* Identify whether we have object, array, or scalar at top-level */
1190         container = &jb->root;
1191
1192         if (JB_ROOT_IS_OBJECT(jb))
1193                 have_object = true;
1194         else if (JB_ROOT_IS_ARRAY(jb) && !JB_ROOT_IS_SCALAR(jb))
1195                 have_array = true;
1196         else
1197         {
1198                 Assert(JB_ROOT_IS_ARRAY(jb) && JB_ROOT_IS_SCALAR(jb));
1199                 /* Extract the scalar value, if it is what we'll return */
1200                 if (npath <= 0)
1201                         jbvp = getIthJsonbValueFromContainer(container, 0);
1202         }
1203
1204         /*
1205          * If the array is empty, return the entire LHS object, on the grounds
1206          * that we should do zero field or element extractions.  For the
1207          * non-scalar case we can just hand back the object without much work. For
1208          * the scalar case, fall through and deal with the value below the loop.
1209          * (This inconsistency arises because there's no easy way to generate a
1210          * JsonbValue directly for root-level containers.)
1211          */
1212         if (npath <= 0 && jbvp == NULL)
1213         {
1214                 if (as_text)
1215                 {
1216                         PG_RETURN_TEXT_P(cstring_to_text(JsonbToCString(NULL,
1217                                                                                                                         container,
1218                                                                                                                         VARSIZE(jb))));
1219                 }
1220                 else
1221                 {
1222                         /* not text mode - just hand back the jsonb */
1223                         PG_RETURN_JSONB(jb);
1224                 }
1225         }
1226
1227         for (i = 0; i < npath; i++)
1228         {
1229                 if (have_object)
1230                 {
1231                         jbvp = findJsonbValueFromContainerLen(container,
1232                                                                                                   JB_FOBJECT,
1233                                                                                                   VARDATA_ANY(pathtext[i]),
1234                                                                                          VARSIZE_ANY_EXHDR(pathtext[i]));
1235                 }
1236                 else if (have_array)
1237                 {
1238                         long            lindex;
1239                         uint32          index;
1240                         char       *indextext = TextDatumGetCString(pathtext[i]);
1241                         char       *endptr;
1242
1243                         errno = 0;
1244                         lindex = strtol(indextext, &endptr, 10);
1245                         if (endptr == indextext || *endptr != '\0' || errno != 0 ||
1246                                 lindex > INT_MAX || lindex < INT_MIN)
1247                                 PG_RETURN_NULL();
1248
1249                         if (lindex >= 0)
1250                         {
1251                                 index = (uint32) lindex;
1252                         }
1253                         else
1254                         {
1255                                 /* Handle negative subscript */
1256                                 uint32          nelements;
1257
1258                                 /* Container must be array, but make sure */
1259                                 if ((container->header & JB_FARRAY) == 0)
1260                                         elog(ERROR, "not a jsonb array");
1261
1262                                 nelements = container->header & JB_CMASK;
1263
1264                                 if (-lindex > nelements)
1265                                         PG_RETURN_NULL();
1266                                 else
1267                                         index = nelements + lindex;
1268                         }
1269
1270                         jbvp = getIthJsonbValueFromContainer(container, index);
1271                 }
1272                 else
1273                 {
1274                         /* scalar, extraction yields a null */
1275                         PG_RETURN_NULL();
1276                 }
1277
1278                 if (jbvp == NULL)
1279                         PG_RETURN_NULL();
1280                 else if (i == npath - 1)
1281                         break;
1282
1283                 if (jbvp->type == jbvBinary)
1284                 {
1285                         JsonbIterator *it = JsonbIteratorInit((JsonbContainer *) jbvp->val.binary.data);
1286                         int                     r;
1287
1288                         r = JsonbIteratorNext(&it, &tv, true);
1289                         container = (JsonbContainer *) jbvp->val.binary.data;
1290                         have_object = r == WJB_BEGIN_OBJECT;
1291                         have_array = r == WJB_BEGIN_ARRAY;
1292                 }
1293                 else
1294                 {
1295                         have_object = jbvp->type == jbvObject;
1296                         have_array = jbvp->type == jbvArray;
1297                 }
1298         }
1299
1300         if (as_text)
1301         {
1302                 /* special-case outputs for string and null values */
1303                 if (jbvp->type == jbvString)
1304                         PG_RETURN_TEXT_P(cstring_to_text_with_len(jbvp->val.string.val,
1305                                                                                                           jbvp->val.string.len));
1306                 if (jbvp->type == jbvNull)
1307                         PG_RETURN_NULL();
1308         }
1309
1310         res = JsonbValueToJsonb(jbvp);
1311
1312         if (as_text)
1313         {
1314                 PG_RETURN_TEXT_P(cstring_to_text(JsonbToCString(NULL,
1315                                                                                                                 &res->root,
1316                                                                                                                 VARSIZE(res))));
1317         }
1318         else
1319         {
1320                 /* not text mode - just hand back the jsonb */
1321                 PG_RETURN_JSONB(res);
1322         }
1323 }
1324
1325 /*
1326  * SQL function json_array_length(json) -> int
1327  */
1328 Datum
1329 json_array_length(PG_FUNCTION_ARGS)
1330 {
1331         text       *json = PG_GETARG_TEXT_P(0);
1332         AlenState  *state;
1333         JsonLexContext *lex;
1334         JsonSemAction *sem;
1335
1336         lex = makeJsonLexContext(json, false);
1337         state = palloc0(sizeof(AlenState));
1338         sem = palloc0(sizeof(JsonSemAction));
1339
1340         /* palloc0 does this for us */
1341 #if 0
1342         state->count = 0;
1343 #endif
1344         state->lex = lex;
1345
1346         sem->semstate = (void *) state;
1347         sem->object_start = alen_object_start;
1348         sem->scalar = alen_scalar;
1349         sem->array_element_start = alen_array_element_start;
1350
1351         pg_parse_json(lex, sem);
1352
1353         PG_RETURN_INT32(state->count);
1354 }
1355
1356 Datum
1357 jsonb_array_length(PG_FUNCTION_ARGS)
1358 {
1359         Jsonb      *jb = PG_GETARG_JSONB(0);
1360
1361         if (JB_ROOT_IS_SCALAR(jb))
1362                 ereport(ERROR,
1363                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1364                                  errmsg("cannot get array length of a scalar")));
1365         else if (!JB_ROOT_IS_ARRAY(jb))
1366                 ereport(ERROR,
1367                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1368                                  errmsg("cannot get array length of a non-array")));
1369
1370         PG_RETURN_INT32(JB_ROOT_COUNT(jb));
1371 }
1372
1373 /*
1374  * These next two checks ensure that the json is an array (since it can't be
1375  * a scalar or an object).
1376  */
1377
1378 static void
1379 alen_object_start(void *state)
1380 {
1381         AlenState  *_state = (AlenState *) state;
1382
1383         /* json structure check */
1384         if (_state->lex->lex_level == 0)
1385                 ereport(ERROR,
1386                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1387                                  errmsg("cannot get array length of a non-array")));
1388 }
1389
1390 static void
1391 alen_scalar(void *state, char *token, JsonTokenType tokentype)
1392 {
1393         AlenState  *_state = (AlenState *) state;
1394
1395         /* json structure check */
1396         if (_state->lex->lex_level == 0)
1397                 ereport(ERROR,
1398                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1399                                  errmsg("cannot get array length of a scalar")));
1400 }
1401
1402 static void
1403 alen_array_element_start(void *state, bool isnull)
1404 {
1405         AlenState  *_state = (AlenState *) state;
1406
1407         /* just count up all the level 1 elements */
1408         if (_state->lex->lex_level == 1)
1409                 _state->count++;
1410 }
1411
1412 /*
1413  * SQL function json_each and json_each_text
1414  *
1415  * decompose a json object into key value pairs.
1416  *
1417  * Unlike json_object_keys() these SRFs operate in materialize mode,
1418  * stashing results into a Tuplestore object as they go.
1419  * The construction of tuples is done using a temporary memory context
1420  * that is cleared out after each tuple is built.
1421  */
1422 Datum
1423 json_each(PG_FUNCTION_ARGS)
1424 {
1425         return each_worker(fcinfo, false);
1426 }
1427
1428 Datum
1429 jsonb_each(PG_FUNCTION_ARGS)
1430 {
1431         return each_worker_jsonb(fcinfo, "jsonb_each", false);
1432 }
1433
1434 Datum
1435 json_each_text(PG_FUNCTION_ARGS)
1436 {
1437         return each_worker(fcinfo, true);
1438 }
1439
1440 Datum
1441 jsonb_each_text(PG_FUNCTION_ARGS)
1442 {
1443         return each_worker_jsonb(fcinfo, "jsonb_each_text", true);
1444 }
1445
1446 static Datum
1447 each_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
1448 {
1449         Jsonb      *jb = PG_GETARG_JSONB(0);
1450         ReturnSetInfo *rsi;
1451         Tuplestorestate *tuple_store;
1452         TupleDesc       tupdesc;
1453         TupleDesc       ret_tdesc;
1454         MemoryContext old_cxt,
1455                                 tmp_cxt;
1456         bool            skipNested = false;
1457         JsonbIterator *it;
1458         JsonbValue      v;
1459         int                     r;
1460
1461         if (!JB_ROOT_IS_OBJECT(jb))
1462                 ereport(ERROR,
1463                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1464                                  errmsg("cannot call %s on a non-object",
1465                                                 funcname)));
1466
1467         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1468
1469         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1470                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1471                 rsi->expectedDesc == NULL)
1472                 ereport(ERROR,
1473                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1474                                  errmsg("set-valued function called in context that "
1475                                                 "cannot accept a set")));
1476
1477         rsi->returnMode = SFRM_Materialize;
1478
1479         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
1480                 ereport(ERROR,
1481                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1482                                  errmsg("function returning record called in context "
1483                                                 "that cannot accept type record")));
1484
1485         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1486
1487         ret_tdesc = CreateTupleDescCopy(tupdesc);
1488         BlessTupleDesc(ret_tdesc);
1489         tuple_store =
1490                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1491                                                           false, work_mem);
1492
1493         MemoryContextSwitchTo(old_cxt);
1494
1495         tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1496                                                                         "jsonb_each temporary cxt",
1497                                                                         ALLOCSET_DEFAULT_MINSIZE,
1498                                                                         ALLOCSET_DEFAULT_INITSIZE,
1499                                                                         ALLOCSET_DEFAULT_MAXSIZE);
1500
1501         it = JsonbIteratorInit(&jb->root);
1502
1503         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
1504         {
1505                 skipNested = true;
1506
1507                 if (r == WJB_KEY)
1508                 {
1509                         text       *key;
1510                         HeapTuple       tuple;
1511                         Datum           values[2];
1512                         bool            nulls[2] = {false, false};
1513
1514                         /* Use the tmp context so we can clean up after each tuple is done */
1515                         old_cxt = MemoryContextSwitchTo(tmp_cxt);
1516
1517                         key = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
1518
1519                         /*
1520                          * The next thing the iterator fetches should be the value, no
1521                          * matter what shape it is.
1522                          */
1523                         r = JsonbIteratorNext(&it, &v, skipNested);
1524
1525                         values[0] = PointerGetDatum(key);
1526
1527                         if (as_text)
1528                         {
1529                                 if (v.type == jbvNull)
1530                                 {
1531                                         /* a json null is an sql null in text mode */
1532                                         nulls[1] = true;
1533                                         values[1] = (Datum) NULL;
1534                                 }
1535                                 else
1536                                 {
1537                                         text       *sv;
1538
1539                                         if (v.type == jbvString)
1540                                         {
1541                                                 /* In text mode, scalar strings should be dequoted */
1542                                                 sv = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
1543                                         }
1544                                         else
1545                                         {
1546                                                 /* Turn anything else into a json string */
1547                                                 StringInfo      jtext = makeStringInfo();
1548                                                 Jsonb      *jb = JsonbValueToJsonb(&v);
1549
1550                                                 (void) JsonbToCString(jtext, &jb->root, 0);
1551                                                 sv = cstring_to_text_with_len(jtext->data, jtext->len);
1552                                         }
1553
1554                                         values[1] = PointerGetDatum(sv);
1555                                 }
1556                         }
1557                         else
1558                         {
1559                                 /* Not in text mode, just return the Jsonb */
1560                                 Jsonb      *val = JsonbValueToJsonb(&v);
1561
1562                                 values[1] = PointerGetDatum(val);
1563                         }
1564
1565                         tuple = heap_form_tuple(ret_tdesc, values, nulls);
1566
1567                         tuplestore_puttuple(tuple_store, tuple);
1568
1569                         /* clean up and switch back */
1570                         MemoryContextSwitchTo(old_cxt);
1571                         MemoryContextReset(tmp_cxt);
1572                 }
1573         }
1574
1575         MemoryContextDelete(tmp_cxt);
1576
1577         rsi->setResult = tuple_store;
1578         rsi->setDesc = ret_tdesc;
1579
1580         PG_RETURN_NULL();
1581 }
1582
1583
1584 static Datum
1585 each_worker(FunctionCallInfo fcinfo, bool as_text)
1586 {
1587         text       *json = PG_GETARG_TEXT_P(0);
1588         JsonLexContext *lex;
1589         JsonSemAction *sem;
1590         ReturnSetInfo *rsi;
1591         MemoryContext old_cxt;
1592         TupleDesc       tupdesc;
1593         EachState  *state;
1594
1595         lex = makeJsonLexContext(json, true);
1596         state = palloc0(sizeof(EachState));
1597         sem = palloc0(sizeof(JsonSemAction));
1598
1599         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1600
1601         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1602                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1603                 rsi->expectedDesc == NULL)
1604                 ereport(ERROR,
1605                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1606                                  errmsg("set-valued function called in context that "
1607                                                 "cannot accept a set")));
1608
1609         rsi->returnMode = SFRM_Materialize;
1610
1611         (void) get_call_result_type(fcinfo, NULL, &tupdesc);
1612
1613         /* make these in a sufficiently long-lived memory context */
1614         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1615
1616         state->ret_tdesc = CreateTupleDescCopy(tupdesc);
1617         BlessTupleDesc(state->ret_tdesc);
1618         state->tuple_store =
1619                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1620                                                           false, work_mem);
1621
1622         MemoryContextSwitchTo(old_cxt);
1623
1624         sem->semstate = (void *) state;
1625         sem->array_start = each_array_start;
1626         sem->scalar = each_scalar;
1627         sem->object_field_start = each_object_field_start;
1628         sem->object_field_end = each_object_field_end;
1629
1630         state->normalize_results = as_text;
1631         state->next_scalar = false;
1632         state->lex = lex;
1633         state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1634                                                                                    "json_each temporary cxt",
1635                                                                                    ALLOCSET_DEFAULT_MINSIZE,
1636                                                                                    ALLOCSET_DEFAULT_INITSIZE,
1637                                                                                    ALLOCSET_DEFAULT_MAXSIZE);
1638
1639         pg_parse_json(lex, sem);
1640
1641         MemoryContextDelete(state->tmp_cxt);
1642
1643         rsi->setResult = state->tuple_store;
1644         rsi->setDesc = state->ret_tdesc;
1645
1646         PG_RETURN_NULL();
1647 }
1648
1649
1650 static void
1651 each_object_field_start(void *state, char *fname, bool isnull)
1652 {
1653         EachState  *_state = (EachState *) state;
1654
1655         /* save a pointer to where the value starts */
1656         if (_state->lex->lex_level == 1)
1657         {
1658                 /*
1659                  * next_scalar will be reset in the object_field_end handler, and
1660                  * since we know the value is a scalar there is no danger of it being
1661                  * on while recursing down the tree.
1662                  */
1663                 if (_state->normalize_results && _state->lex->token_type == JSON_TOKEN_STRING)
1664                         _state->next_scalar = true;
1665                 else
1666                         _state->result_start = _state->lex->token_start;
1667         }
1668 }
1669
1670 static void
1671 each_object_field_end(void *state, char *fname, bool isnull)
1672 {
1673         EachState  *_state = (EachState *) state;
1674         MemoryContext old_cxt;
1675         int                     len;
1676         text       *val;
1677         HeapTuple       tuple;
1678         Datum           values[2];
1679         bool            nulls[2] = {false, false};
1680
1681         /* skip over nested objects */
1682         if (_state->lex->lex_level != 1)
1683                 return;
1684
1685         /* use the tmp context so we can clean up after each tuple is done */
1686         old_cxt = MemoryContextSwitchTo(_state->tmp_cxt);
1687
1688         values[0] = CStringGetTextDatum(fname);
1689
1690         if (isnull && _state->normalize_results)
1691         {
1692                 nulls[1] = true;
1693                 values[1] = (Datum) 0;
1694         }
1695         else if (_state->next_scalar)
1696         {
1697                 values[1] = CStringGetTextDatum(_state->normalized_scalar);
1698                 _state->next_scalar = false;
1699         }
1700         else
1701         {
1702                 len = _state->lex->prev_token_terminator - _state->result_start;
1703                 val = cstring_to_text_with_len(_state->result_start, len);
1704                 values[1] = PointerGetDatum(val);
1705         }
1706
1707         tuple = heap_form_tuple(_state->ret_tdesc, values, nulls);
1708
1709         tuplestore_puttuple(_state->tuple_store, tuple);
1710
1711         /* clean up and switch back */
1712         MemoryContextSwitchTo(old_cxt);
1713         MemoryContextReset(_state->tmp_cxt);
1714 }
1715
1716 static void
1717 each_array_start(void *state)
1718 {
1719         EachState  *_state = (EachState *) state;
1720
1721         /* json structure check */
1722         if (_state->lex->lex_level == 0)
1723                 ereport(ERROR,
1724                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1725                                  errmsg("cannot deconstruct an array as an object")));
1726 }
1727
1728 static void
1729 each_scalar(void *state, char *token, JsonTokenType tokentype)
1730 {
1731         EachState  *_state = (EachState *) state;
1732
1733         /* json structure check */
1734         if (_state->lex->lex_level == 0)
1735                 ereport(ERROR,
1736                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1737                                  errmsg("cannot deconstruct a scalar")));
1738
1739         /* supply de-escaped value if required */
1740         if (_state->next_scalar)
1741                 _state->normalized_scalar = token;
1742 }
1743
1744 /*
1745  * SQL functions json_array_elements and json_array_elements_text
1746  *
1747  * get the elements from a json array
1748  *
1749  * a lot of this processing is similar to the json_each* functions
1750  */
1751
1752 Datum
1753 jsonb_array_elements(PG_FUNCTION_ARGS)
1754 {
1755         return elements_worker_jsonb(fcinfo, "jsonb_array_elements", false);
1756 }
1757
1758 Datum
1759 jsonb_array_elements_text(PG_FUNCTION_ARGS)
1760 {
1761         return elements_worker_jsonb(fcinfo, "jsonb_array_elements_text", true);
1762 }
1763
1764 static Datum
1765 elements_worker_jsonb(FunctionCallInfo fcinfo, const char *funcname,
1766                                           bool as_text)
1767 {
1768         Jsonb      *jb = PG_GETARG_JSONB(0);
1769         ReturnSetInfo *rsi;
1770         Tuplestorestate *tuple_store;
1771         TupleDesc       tupdesc;
1772         TupleDesc       ret_tdesc;
1773         MemoryContext old_cxt,
1774                                 tmp_cxt;
1775         bool            skipNested = false;
1776         JsonbIterator *it;
1777         JsonbValue      v;
1778         int                     r;
1779
1780         if (JB_ROOT_IS_SCALAR(jb))
1781                 ereport(ERROR,
1782                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1783                                  errmsg("cannot extract elements from a scalar")));
1784         else if (!JB_ROOT_IS_ARRAY(jb))
1785                 ereport(ERROR,
1786                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
1787                                  errmsg("cannot extract elements from an object")));
1788
1789         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1790
1791         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1792                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1793                 rsi->expectedDesc == NULL)
1794                 ereport(ERROR,
1795                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1796                                  errmsg("set-valued function called in context that "
1797                                                 "cannot accept a set")));
1798
1799         rsi->returnMode = SFRM_Materialize;
1800
1801         /* it's a simple type, so don't use get_call_result_type() */
1802         tupdesc = rsi->expectedDesc;
1803
1804         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1805
1806         ret_tdesc = CreateTupleDescCopy(tupdesc);
1807         BlessTupleDesc(ret_tdesc);
1808         tuple_store =
1809                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1810                                                           false, work_mem);
1811
1812         MemoryContextSwitchTo(old_cxt);
1813
1814         tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1815                                                                         "jsonb_array_elements temporary cxt",
1816                                                                         ALLOCSET_DEFAULT_MINSIZE,
1817                                                                         ALLOCSET_DEFAULT_INITSIZE,
1818                                                                         ALLOCSET_DEFAULT_MAXSIZE);
1819
1820         it = JsonbIteratorInit(&jb->root);
1821
1822         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
1823         {
1824                 skipNested = true;
1825
1826                 if (r == WJB_ELEM)
1827                 {
1828                         HeapTuple       tuple;
1829                         Datum           values[1];
1830                         bool            nulls[1] = {false};
1831
1832                         /* use the tmp context so we can clean up after each tuple is done */
1833                         old_cxt = MemoryContextSwitchTo(tmp_cxt);
1834
1835                         if (!as_text)
1836                         {
1837                                 Jsonb      *val = JsonbValueToJsonb(&v);
1838
1839                                 values[0] = PointerGetDatum(val);
1840                         }
1841                         else
1842                         {
1843                                 if (v.type == jbvNull)
1844                                 {
1845                                         /* a json null is an sql null in text mode */
1846                                         nulls[0] = true;
1847                                         values[0] = (Datum) NULL;
1848                                 }
1849                                 else
1850                                 {
1851                                         text       *sv;
1852
1853                                         if (v.type == jbvString)
1854                                         {
1855                                                 /* in text mode scalar strings should be dequoted */
1856                                                 sv = cstring_to_text_with_len(v.val.string.val, v.val.string.len);
1857                                         }
1858                                         else
1859                                         {
1860                                                 /* turn anything else into a json string */
1861                                                 StringInfo      jtext = makeStringInfo();
1862                                                 Jsonb      *jb = JsonbValueToJsonb(&v);
1863
1864                                                 (void) JsonbToCString(jtext, &jb->root, 0);
1865                                                 sv = cstring_to_text_with_len(jtext->data, jtext->len);
1866                                         }
1867
1868                                         values[0] = PointerGetDatum(sv);
1869                                 }
1870                         }
1871
1872                         tuple = heap_form_tuple(ret_tdesc, values, nulls);
1873
1874                         tuplestore_puttuple(tuple_store, tuple);
1875
1876                         /* clean up and switch back */
1877                         MemoryContextSwitchTo(old_cxt);
1878                         MemoryContextReset(tmp_cxt);
1879                 }
1880         }
1881
1882         MemoryContextDelete(tmp_cxt);
1883
1884         rsi->setResult = tuple_store;
1885         rsi->setDesc = ret_tdesc;
1886
1887         PG_RETURN_NULL();
1888 }
1889
1890 Datum
1891 json_array_elements(PG_FUNCTION_ARGS)
1892 {
1893         return elements_worker(fcinfo, "json_array_elements", false);
1894 }
1895
1896 Datum
1897 json_array_elements_text(PG_FUNCTION_ARGS)
1898 {
1899         return elements_worker(fcinfo, "json_array_elements_text", true);
1900 }
1901
1902 static Datum
1903 elements_worker(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
1904 {
1905         text       *json = PG_GETARG_TEXT_P(0);
1906
1907         /* elements only needs escaped strings when as_text */
1908         JsonLexContext *lex = makeJsonLexContext(json, as_text);
1909         JsonSemAction *sem;
1910         ReturnSetInfo *rsi;
1911         MemoryContext old_cxt;
1912         TupleDesc       tupdesc;
1913         ElementsState *state;
1914
1915         state = palloc0(sizeof(ElementsState));
1916         sem = palloc0(sizeof(JsonSemAction));
1917
1918         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
1919
1920         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
1921                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
1922                 rsi->expectedDesc == NULL)
1923                 ereport(ERROR,
1924                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
1925                                  errmsg("set-valued function called in context that "
1926                                                 "cannot accept a set")));
1927
1928         rsi->returnMode = SFRM_Materialize;
1929
1930         /* it's a simple type, so don't use get_call_result_type() */
1931         tupdesc = rsi->expectedDesc;
1932
1933         /* make these in a sufficiently long-lived memory context */
1934         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
1935
1936         state->ret_tdesc = CreateTupleDescCopy(tupdesc);
1937         BlessTupleDesc(state->ret_tdesc);
1938         state->tuple_store =
1939                 tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random,
1940                                                           false, work_mem);
1941
1942         MemoryContextSwitchTo(old_cxt);
1943
1944         sem->semstate = (void *) state;
1945         sem->object_start = elements_object_start;
1946         sem->scalar = elements_scalar;
1947         sem->array_element_start = elements_array_element_start;
1948         sem->array_element_end = elements_array_element_end;
1949
1950         state->function_name = funcname;
1951         state->normalize_results = as_text;
1952         state->next_scalar = false;
1953         state->lex = lex;
1954         state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
1955                                                                                  "json_array_elements temporary cxt",
1956                                                                                    ALLOCSET_DEFAULT_MINSIZE,
1957                                                                                    ALLOCSET_DEFAULT_INITSIZE,
1958                                                                                    ALLOCSET_DEFAULT_MAXSIZE);
1959
1960         pg_parse_json(lex, sem);
1961
1962         MemoryContextDelete(state->tmp_cxt);
1963
1964         rsi->setResult = state->tuple_store;
1965         rsi->setDesc = state->ret_tdesc;
1966
1967         PG_RETURN_NULL();
1968 }
1969
1970 static void
1971 elements_array_element_start(void *state, bool isnull)
1972 {
1973         ElementsState *_state = (ElementsState *) state;
1974
1975         /* save a pointer to where the value starts */
1976         if (_state->lex->lex_level == 1)
1977         {
1978                 /*
1979                  * next_scalar will be reset in the array_element_end handler, and
1980                  * since we know the value is a scalar there is no danger of it being
1981                  * on while recursing down the tree.
1982                  */
1983                 if (_state->normalize_results && _state->lex->token_type == JSON_TOKEN_STRING)
1984                         _state->next_scalar = true;
1985                 else
1986                         _state->result_start = _state->lex->token_start;
1987         }
1988 }
1989
1990 static void
1991 elements_array_element_end(void *state, bool isnull)
1992 {
1993         ElementsState *_state = (ElementsState *) state;
1994         MemoryContext old_cxt;
1995         int                     len;
1996         text       *val;
1997         HeapTuple       tuple;
1998         Datum           values[1];
1999         bool            nulls[1] = {false};
2000
2001         /* skip over nested objects */
2002         if (_state->lex->lex_level != 1)
2003                 return;
2004
2005         /* use the tmp context so we can clean up after each tuple is done */
2006         old_cxt = MemoryContextSwitchTo(_state->tmp_cxt);
2007
2008         if (isnull && _state->normalize_results)
2009         {
2010                 nulls[0] = true;
2011                 values[0] = (Datum) NULL;
2012         }
2013         else if (_state->next_scalar)
2014         {
2015                 values[0] = CStringGetTextDatum(_state->normalized_scalar);
2016                 _state->next_scalar = false;
2017         }
2018         else
2019         {
2020                 len = _state->lex->prev_token_terminator - _state->result_start;
2021                 val = cstring_to_text_with_len(_state->result_start, len);
2022                 values[0] = PointerGetDatum(val);
2023         }
2024
2025         tuple = heap_form_tuple(_state->ret_tdesc, values, nulls);
2026
2027         tuplestore_puttuple(_state->tuple_store, tuple);
2028
2029         /* clean up and switch back */
2030         MemoryContextSwitchTo(old_cxt);
2031         MemoryContextReset(_state->tmp_cxt);
2032 }
2033
2034 static void
2035 elements_object_start(void *state)
2036 {
2037         ElementsState *_state = (ElementsState *) state;
2038
2039         /* json structure check */
2040         if (_state->lex->lex_level == 0)
2041                 ereport(ERROR,
2042                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2043                                  errmsg("cannot call %s on a non-array",
2044                                                 _state->function_name)));
2045 }
2046
2047 static void
2048 elements_scalar(void *state, char *token, JsonTokenType tokentype)
2049 {
2050         ElementsState *_state = (ElementsState *) state;
2051
2052         /* json structure check */
2053         if (_state->lex->lex_level == 0)
2054                 ereport(ERROR,
2055                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2056                                  errmsg("cannot call %s on a scalar",
2057                                                 _state->function_name)));
2058
2059         /* supply de-escaped value if required */
2060         if (_state->next_scalar)
2061                 _state->normalized_scalar = token;
2062 }
2063
2064 /*
2065  * SQL function json_populate_record
2066  *
2067  * set fields in a record from the argument json
2068  *
2069  * Code adapted shamelessly from hstore's populate_record
2070  * which is in turn partly adapted from record_out.
2071  *
2072  * The json is decomposed into a hash table, in which each
2073  * field in the record is then looked up by name. For jsonb
2074  * we fetch the values direct from the object.
2075  */
2076 Datum
2077 jsonb_populate_record(PG_FUNCTION_ARGS)
2078 {
2079         return populate_record_worker(fcinfo, "jsonb_populate_record", true);
2080 }
2081
2082 Datum
2083 jsonb_to_record(PG_FUNCTION_ARGS)
2084 {
2085         return populate_record_worker(fcinfo, "jsonb_to_record", false);
2086 }
2087
2088 Datum
2089 json_populate_record(PG_FUNCTION_ARGS)
2090 {
2091         return populate_record_worker(fcinfo, "json_populate_record", true);
2092 }
2093
2094 Datum
2095 json_to_record(PG_FUNCTION_ARGS)
2096 {
2097         return populate_record_worker(fcinfo, "json_to_record", false);
2098 }
2099
2100 static Datum
2101 populate_record_worker(FunctionCallInfo fcinfo, const char *funcname,
2102                                            bool have_record_arg)
2103 {
2104         int                     json_arg_num = have_record_arg ? 1 : 0;
2105         Oid                     jtype = get_fn_expr_argtype(fcinfo->flinfo, json_arg_num);
2106         text       *json;
2107         Jsonb      *jb = NULL;
2108         HTAB       *json_hash = NULL;
2109         HeapTupleHeader rec = NULL;
2110         Oid                     tupType = InvalidOid;
2111         int32           tupTypmod = -1;
2112         TupleDesc       tupdesc;
2113         HeapTupleData tuple;
2114         HeapTuple       rettuple;
2115         RecordIOData *my_extra;
2116         int                     ncolumns;
2117         int                     i;
2118         Datum      *values;
2119         bool       *nulls;
2120
2121         Assert(jtype == JSONOID || jtype == JSONBOID);
2122
2123         if (have_record_arg)
2124         {
2125                 Oid                     argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
2126
2127                 if (!type_is_rowtype(argtype))
2128                         ereport(ERROR,
2129                                         (errcode(ERRCODE_DATATYPE_MISMATCH),
2130                                          errmsg("first argument of %s must be a row type",
2131                                                         funcname)));
2132
2133                 if (PG_ARGISNULL(0))
2134                 {
2135                         if (PG_ARGISNULL(1))
2136                                 PG_RETURN_NULL();
2137
2138                         /*
2139                          * have no tuple to look at, so the only source of type info is
2140                          * the argtype. The lookup_rowtype_tupdesc call below will error
2141                          * out if we don't have a known composite type oid here.
2142                          */
2143                         tupType = argtype;
2144                         tupTypmod = -1;
2145                 }
2146                 else
2147                 {
2148                         rec = PG_GETARG_HEAPTUPLEHEADER(0);
2149
2150                         if (PG_ARGISNULL(1))
2151                                 PG_RETURN_POINTER(rec);
2152
2153                         /* Extract type info from the tuple itself */
2154                         tupType = HeapTupleHeaderGetTypeId(rec);
2155                         tupTypmod = HeapTupleHeaderGetTypMod(rec);
2156                 }
2157
2158                 tupdesc = lookup_rowtype_tupdesc(tupType, tupTypmod);
2159         }
2160         else
2161         {
2162                 /* json{b}_to_record case */
2163                 if (PG_ARGISNULL(0))
2164                         PG_RETURN_NULL();
2165
2166                 if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
2167                         ereport(ERROR,
2168                                         (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
2169                                          errmsg("function returning record called in context "
2170                                                         "that cannot accept type record"),
2171                                          errhint("Try calling the function in the FROM clause "
2172                                                          "using a column definition list.")));
2173         }
2174
2175         if (jtype == JSONOID)
2176         {
2177                 /* just get the text */
2178                 json = PG_GETARG_TEXT_P(json_arg_num);
2179
2180                 json_hash = get_json_object_as_hash(json, funcname);
2181
2182                 /*
2183                  * if the input json is empty, we can only skip the rest if we were
2184                  * passed in a non-null record, since otherwise there may be issues
2185                  * with domain nulls.
2186                  */
2187                 if (hash_get_num_entries(json_hash) == 0 && rec)
2188                 {
2189                         hash_destroy(json_hash);
2190                         ReleaseTupleDesc(tupdesc);
2191                         PG_RETURN_POINTER(rec);
2192                 }
2193         }
2194         else
2195         {
2196                 jb = PG_GETARG_JSONB(json_arg_num);
2197
2198                 /* same logic as for json */
2199                 if (JB_ROOT_COUNT(jb) == 0 && rec)
2200                 {
2201                         ReleaseTupleDesc(tupdesc);
2202                         PG_RETURN_POINTER(rec);
2203                 }
2204         }
2205
2206         ncolumns = tupdesc->natts;
2207
2208         if (rec)
2209         {
2210                 /* Build a temporary HeapTuple control structure */
2211                 tuple.t_len = HeapTupleHeaderGetDatumLength(rec);
2212                 ItemPointerSetInvalid(&(tuple.t_self));
2213                 tuple.t_tableOid = InvalidOid;
2214                 tuple.t_data = rec;
2215         }
2216
2217         /*
2218          * We arrange to look up the needed I/O info just once per series of
2219          * calls, assuming the record type doesn't change underneath us.
2220          */
2221         my_extra = (RecordIOData *) fcinfo->flinfo->fn_extra;
2222         if (my_extra == NULL ||
2223                 my_extra->ncolumns != ncolumns)
2224         {
2225                 fcinfo->flinfo->fn_extra =
2226                         MemoryContextAlloc(fcinfo->flinfo->fn_mcxt,
2227                                                            offsetof(RecordIOData, columns) +
2228                                                            ncolumns * sizeof(ColumnIOData));
2229                 my_extra = (RecordIOData *) fcinfo->flinfo->fn_extra;
2230                 my_extra->record_type = InvalidOid;
2231                 my_extra->record_typmod = 0;
2232                 my_extra->ncolumns = ncolumns;
2233                 MemSet(my_extra->columns, 0, sizeof(ColumnIOData) * ncolumns);
2234         }
2235
2236         if (have_record_arg && (my_extra->record_type != tupType ||
2237                                                         my_extra->record_typmod != tupTypmod))
2238         {
2239                 MemSet(my_extra, 0,
2240                            offsetof(RecordIOData, columns) +
2241                            ncolumns * sizeof(ColumnIOData));
2242                 my_extra->record_type = tupType;
2243                 my_extra->record_typmod = tupTypmod;
2244                 my_extra->ncolumns = ncolumns;
2245         }
2246
2247         values = (Datum *) palloc(ncolumns * sizeof(Datum));
2248         nulls = (bool *) palloc(ncolumns * sizeof(bool));
2249
2250         if (rec)
2251         {
2252                 /* Break down the tuple into fields */
2253                 heap_deform_tuple(&tuple, tupdesc, values, nulls);
2254         }
2255         else
2256         {
2257                 for (i = 0; i < ncolumns; ++i)
2258                 {
2259                         values[i] = (Datum) 0;
2260                         nulls[i] = true;
2261                 }
2262         }
2263
2264         for (i = 0; i < ncolumns; ++i)
2265         {
2266                 ColumnIOData *column_info = &my_extra->columns[i];
2267                 Oid                     column_type = tupdesc->attrs[i]->atttypid;
2268                 JsonbValue *v = NULL;
2269                 JsonHashEntry *hashentry = NULL;
2270
2271                 /* Ignore dropped columns in datatype */
2272                 if (tupdesc->attrs[i]->attisdropped)
2273                 {
2274                         nulls[i] = true;
2275                         continue;
2276                 }
2277
2278                 if (jtype == JSONOID)
2279                 {
2280                         hashentry = hash_search(json_hash,
2281                                                                         NameStr(tupdesc->attrs[i]->attname),
2282                                                                         HASH_FIND, NULL);
2283                 }
2284                 else
2285                 {
2286                         char       *key = NameStr(tupdesc->attrs[i]->attname);
2287
2288                         v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT, key,
2289                                                                                            strlen(key));
2290                 }
2291
2292                 /*
2293                  * we can't just skip here if the key wasn't found since we might have
2294                  * a domain to deal with. If we were passed in a non-null record
2295                  * datum, we assume that the existing values are valid (if they're
2296                  * not, then it's not our fault), but if we were passed in a null,
2297                  * then every field which we don't populate needs to be run through
2298                  * the input function just in case it's a domain type.
2299                  */
2300                 if (((jtype == JSONOID && hashentry == NULL) ||
2301                          (jtype == JSONBOID && v == NULL)) && rec)
2302                         continue;
2303
2304                 /*
2305                  * Prepare to convert the column value from text
2306                  */
2307                 if (column_info->column_type != column_type)
2308                 {
2309                         getTypeInputInfo(column_type,
2310                                                          &column_info->typiofunc,
2311                                                          &column_info->typioparam);
2312                         fmgr_info_cxt(column_info->typiofunc, &column_info->proc,
2313                                                   fcinfo->flinfo->fn_mcxt);
2314                         column_info->column_type = column_type;
2315                 }
2316                 if ((jtype == JSONOID && (hashentry == NULL || hashentry->isnull)) ||
2317                         (jtype == JSONBOID && (v == NULL || v->type == jbvNull)))
2318                 {
2319                         /*
2320                          * need InputFunctionCall to happen even for nulls, so that domain
2321                          * checks are done
2322                          */
2323                         values[i] = InputFunctionCall(&column_info->proc, NULL,
2324                                                                                   column_info->typioparam,
2325                                                                                   tupdesc->attrs[i]->atttypmod);
2326                         nulls[i] = true;
2327                 }
2328                 else
2329                 {
2330                         char       *s = NULL;
2331
2332                         if (jtype == JSONOID)
2333                         {
2334                                 /* already done the hard work in the json case */
2335                                 s = hashentry->val;
2336                         }
2337                         else
2338                         {
2339                                 if (v->type == jbvString)
2340                                         s = pnstrdup(v->val.string.val, v->val.string.len);
2341                                 else if (v->type == jbvBool)
2342                                         s = pnstrdup((v->val.boolean) ? "t" : "f", 1);
2343                                 else if (v->type == jbvNumeric)
2344                                         s = DatumGetCString(DirectFunctionCall1(numeric_out,
2345                                                                                    PointerGetDatum(v->val.numeric)));
2346                                 else if (v->type == jbvBinary)
2347                                         s = JsonbToCString(NULL, (JsonbContainer *) v->val.binary.data, v->val.binary.len);
2348                                 else
2349                                         elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
2350                         }
2351
2352                         values[i] = InputFunctionCall(&column_info->proc, s,
2353                                                                                   column_info->typioparam,
2354                                                                                   tupdesc->attrs[i]->atttypmod);
2355                         nulls[i] = false;
2356                 }
2357         }
2358
2359         rettuple = heap_form_tuple(tupdesc, values, nulls);
2360
2361         ReleaseTupleDesc(tupdesc);
2362
2363         if (json_hash)
2364                 hash_destroy(json_hash);
2365
2366         PG_RETURN_DATUM(HeapTupleGetDatum(rettuple));
2367 }
2368
2369 /*
2370  * get_json_object_as_hash
2371  *
2372  * decompose a json object into a hash table.
2373  */
2374 static HTAB *
2375 get_json_object_as_hash(text *json, const char *funcname)
2376 {
2377         HASHCTL         ctl;
2378         HTAB       *tab;
2379         JHashState *state;
2380         JsonLexContext *lex = makeJsonLexContext(json, true);
2381         JsonSemAction *sem;
2382
2383         memset(&ctl, 0, sizeof(ctl));
2384         ctl.keysize = NAMEDATALEN;
2385         ctl.entrysize = sizeof(JsonHashEntry);
2386         ctl.hcxt = CurrentMemoryContext;
2387         tab = hash_create("json object hashtable",
2388                                           100,
2389                                           &ctl,
2390                                           HASH_ELEM | HASH_CONTEXT);
2391
2392         state = palloc0(sizeof(JHashState));
2393         sem = palloc0(sizeof(JsonSemAction));
2394
2395         state->function_name = funcname;
2396         state->hash = tab;
2397         state->lex = lex;
2398
2399         sem->semstate = (void *) state;
2400         sem->array_start = hash_array_start;
2401         sem->scalar = hash_scalar;
2402         sem->object_field_start = hash_object_field_start;
2403         sem->object_field_end = hash_object_field_end;
2404
2405         pg_parse_json(lex, sem);
2406
2407         return tab;
2408 }
2409
2410 static void
2411 hash_object_field_start(void *state, char *fname, bool isnull)
2412 {
2413         JHashState *_state = (JHashState *) state;
2414
2415         if (_state->lex->lex_level > 1)
2416                 return;
2417
2418         if (_state->lex->token_type == JSON_TOKEN_ARRAY_START ||
2419                 _state->lex->token_type == JSON_TOKEN_OBJECT_START)
2420         {
2421                 /* remember start position of the whole text of the subobject */
2422                 _state->save_json_start = _state->lex->token_start;
2423         }
2424         else
2425         {
2426                 /* must be a scalar */
2427                 _state->save_json_start = NULL;
2428         }
2429 }
2430
2431 static void
2432 hash_object_field_end(void *state, char *fname, bool isnull)
2433 {
2434         JHashState *_state = (JHashState *) state;
2435         JsonHashEntry *hashentry;
2436         bool            found;
2437
2438         /*
2439          * Ignore nested fields.
2440          */
2441         if (_state->lex->lex_level > 2)
2442                 return;
2443
2444         /*
2445          * Ignore field names >= NAMEDATALEN - they can't match a record field.
2446          * (Note: without this test, the hash code would truncate the string at
2447          * NAMEDATALEN-1, and could then match against a similarly-truncated
2448          * record field name.  That would be a reasonable behavior, but this code
2449          * has previously insisted on exact equality, so we keep this behavior.)
2450          */
2451         if (strlen(fname) >= NAMEDATALEN)
2452                 return;
2453
2454         hashentry = hash_search(_state->hash, fname, HASH_ENTER, &found);
2455
2456         /*
2457          * found being true indicates a duplicate. We don't do anything about
2458          * that, a later field with the same name overrides the earlier field.
2459          */
2460
2461         hashentry->isnull = isnull;
2462         if (_state->save_json_start != NULL)
2463         {
2464                 int                     len = _state->lex->prev_token_terminator - _state->save_json_start;
2465                 char       *val = palloc((len + 1) * sizeof(char));
2466
2467                 memcpy(val, _state->save_json_start, len);
2468                 val[len] = '\0';
2469                 hashentry->val = val;
2470         }
2471         else
2472         {
2473                 /* must have had a scalar instead */
2474                 hashentry->val = _state->saved_scalar;
2475         }
2476 }
2477
2478 static void
2479 hash_array_start(void *state)
2480 {
2481         JHashState *_state = (JHashState *) state;
2482
2483         if (_state->lex->lex_level == 0)
2484                 ereport(ERROR,
2485                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2486                            errmsg("cannot call %s on an array", _state->function_name)));
2487 }
2488
2489 static void
2490 hash_scalar(void *state, char *token, JsonTokenType tokentype)
2491 {
2492         JHashState *_state = (JHashState *) state;
2493
2494         if (_state->lex->lex_level == 0)
2495                 ereport(ERROR,
2496                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2497                            errmsg("cannot call %s on a scalar", _state->function_name)));
2498
2499         if (_state->lex->lex_level == 1)
2500                 _state->saved_scalar = token;
2501 }
2502
2503
2504 /*
2505  * SQL function json_populate_recordset
2506  *
2507  * set fields in a set of records from the argument json,
2508  * which must be an array of objects.
2509  *
2510  * similar to json_populate_record, but the tuple-building code
2511  * is pushed down into the semantic action handlers so it's done
2512  * per object in the array.
2513  */
2514 Datum
2515 jsonb_populate_recordset(PG_FUNCTION_ARGS)
2516 {
2517         return populate_recordset_worker(fcinfo, "jsonb_populate_recordset", true);
2518 }
2519
2520 Datum
2521 jsonb_to_recordset(PG_FUNCTION_ARGS)
2522 {
2523         return populate_recordset_worker(fcinfo, "jsonb_to_recordset", false);
2524 }
2525
2526 Datum
2527 json_populate_recordset(PG_FUNCTION_ARGS)
2528 {
2529         return populate_recordset_worker(fcinfo, "json_populate_recordset", true);
2530 }
2531
2532 Datum
2533 json_to_recordset(PG_FUNCTION_ARGS)
2534 {
2535         return populate_recordset_worker(fcinfo, "json_to_recordset", false);
2536 }
2537
2538 static void
2539 make_row_from_rec_and_jsonb(Jsonb *element, PopulateRecordsetState *state)
2540 {
2541         Datum      *values;
2542         bool       *nulls;
2543         int                     i;
2544         RecordIOData *my_extra = state->my_extra;
2545         int                     ncolumns = my_extra->ncolumns;
2546         TupleDesc       tupdesc = state->ret_tdesc;
2547         HeapTupleHeader rec = state->rec;
2548         HeapTuple       rettuple;
2549
2550         values = (Datum *) palloc(ncolumns * sizeof(Datum));
2551         nulls = (bool *) palloc(ncolumns * sizeof(bool));
2552
2553         if (state->rec)
2554         {
2555                 HeapTupleData tuple;
2556
2557                 /* Build a temporary HeapTuple control structure */
2558                 tuple.t_len = HeapTupleHeaderGetDatumLength(state->rec);
2559                 ItemPointerSetInvalid(&(tuple.t_self));
2560                 tuple.t_tableOid = InvalidOid;
2561                 tuple.t_data = state->rec;
2562
2563                 /* Break down the tuple into fields */
2564                 heap_deform_tuple(&tuple, tupdesc, values, nulls);
2565         }
2566         else
2567         {
2568                 for (i = 0; i < ncolumns; ++i)
2569                 {
2570                         values[i] = (Datum) 0;
2571                         nulls[i] = true;
2572                 }
2573         }
2574
2575         for (i = 0; i < ncolumns; ++i)
2576         {
2577                 ColumnIOData *column_info = &my_extra->columns[i];
2578                 Oid                     column_type = tupdesc->attrs[i]->atttypid;
2579                 JsonbValue *v = NULL;
2580                 char       *key;
2581
2582                 /* Ignore dropped columns in datatype */
2583                 if (tupdesc->attrs[i]->attisdropped)
2584                 {
2585                         nulls[i] = true;
2586                         continue;
2587                 }
2588
2589                 key = NameStr(tupdesc->attrs[i]->attname);
2590
2591                 v = findJsonbValueFromContainerLen(&element->root, JB_FOBJECT,
2592                                                                                    key, strlen(key));
2593
2594                 /*
2595                  * We can't just skip here if the key wasn't found since we might have
2596                  * a domain to deal with. If we were passed in a non-null record
2597                  * datum, we assume that the existing values are valid (if they're
2598                  * not, then it's not our fault), but if we were passed in a null,
2599                  * then every field which we don't populate needs to be run through
2600                  * the input function just in case it's a domain type.
2601                  */
2602                 if (v == NULL && rec)
2603                         continue;
2604
2605                 /*
2606                  * Prepare to convert the column value from text
2607                  */
2608                 if (column_info->column_type != column_type)
2609                 {
2610                         getTypeInputInfo(column_type,
2611                                                          &column_info->typiofunc,
2612                                                          &column_info->typioparam);
2613                         fmgr_info_cxt(column_info->typiofunc, &column_info->proc,
2614                                                   state->fn_mcxt);
2615                         column_info->column_type = column_type;
2616                 }
2617                 if (v == NULL || v->type == jbvNull)
2618                 {
2619                         /*
2620                          * Need InputFunctionCall to happen even for nulls, so that domain
2621                          * checks are done
2622                          */
2623                         values[i] = InputFunctionCall(&column_info->proc, NULL,
2624                                                                                   column_info->typioparam,
2625                                                                                   tupdesc->attrs[i]->atttypmod);
2626                         nulls[i] = true;
2627                 }
2628                 else
2629                 {
2630                         char       *s = NULL;
2631
2632                         if (v->type == jbvString)
2633                                 s = pnstrdup(v->val.string.val, v->val.string.len);
2634                         else if (v->type == jbvBool)
2635                                 s = pnstrdup((v->val.boolean) ? "t" : "f", 1);
2636                         else if (v->type == jbvNumeric)
2637                                 s = DatumGetCString(DirectFunctionCall1(numeric_out,
2638                                                                                    PointerGetDatum(v->val.numeric)));
2639                         else if (v->type == jbvBinary)
2640                                 s = JsonbToCString(NULL, (JsonbContainer *) v->val.binary.data, v->val.binary.len);
2641                         else
2642                                 elog(ERROR, "unrecognized jsonb type: %d", (int) v->type);
2643
2644                         values[i] = InputFunctionCall(&column_info->proc, s,
2645                                                                                   column_info->typioparam,
2646                                                                                   tupdesc->attrs[i]->atttypmod);
2647                         nulls[i] = false;
2648                 }
2649         }
2650
2651         rettuple = heap_form_tuple(tupdesc, values, nulls);
2652
2653         tuplestore_puttuple(state->tuple_store, rettuple);
2654 }
2655
2656 /*
2657  * common worker for json_populate_recordset() and json_to_recordset()
2658  */
2659 static Datum
2660 populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
2661                                                   bool have_record_arg)
2662 {
2663         int                     json_arg_num = have_record_arg ? 1 : 0;
2664         Oid                     jtype = get_fn_expr_argtype(fcinfo->flinfo, json_arg_num);
2665         ReturnSetInfo *rsi;
2666         MemoryContext old_cxt;
2667         Oid                     tupType;
2668         int32           tupTypmod;
2669         HeapTupleHeader rec;
2670         TupleDesc       tupdesc;
2671         RecordIOData *my_extra;
2672         int                     ncolumns;
2673         PopulateRecordsetState *state;
2674
2675         if (have_record_arg)
2676         {
2677                 Oid                     argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
2678
2679                 if (!type_is_rowtype(argtype))
2680                         ereport(ERROR,
2681                                         (errcode(ERRCODE_DATATYPE_MISMATCH),
2682                                          errmsg("first argument of %s must be a row type",
2683                                                         funcname)));
2684         }
2685
2686         rsi = (ReturnSetInfo *) fcinfo->resultinfo;
2687
2688         if (!rsi || !IsA(rsi, ReturnSetInfo) ||
2689                 (rsi->allowedModes & SFRM_Materialize) == 0 ||
2690                 rsi->expectedDesc == NULL)
2691                 ereport(ERROR,
2692                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
2693                                  errmsg("set-valued function called in context that "
2694                                                 "cannot accept a set")));
2695
2696         rsi->returnMode = SFRM_Materialize;
2697
2698         /*
2699          * get the tupdesc from the result set info - it must be a record type
2700          * because we already checked that arg1 is a record type, or we're in a
2701          * to_record function which returns a setof record.
2702          */
2703         if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
2704                 ereport(ERROR,
2705                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
2706                                  errmsg("function returning record called in context "
2707                                                 "that cannot accept type record")));
2708
2709         /* if the json is null send back an empty set */
2710         if (PG_ARGISNULL(json_arg_num))
2711                 PG_RETURN_NULL();
2712
2713         if (!have_record_arg || PG_ARGISNULL(0))
2714                 rec = NULL;
2715         else
2716                 rec = PG_GETARG_HEAPTUPLEHEADER(0);
2717
2718         tupType = tupdesc->tdtypeid;
2719         tupTypmod = tupdesc->tdtypmod;
2720         ncolumns = tupdesc->natts;
2721
2722         /*
2723          * We arrange to look up the needed I/O info just once per series of
2724          * calls, assuming the record type doesn't change underneath us.
2725          */
2726         my_extra = (RecordIOData *) fcinfo->flinfo->fn_extra;
2727         if (my_extra == NULL ||
2728                 my_extra->ncolumns != ncolumns)
2729         {
2730                 fcinfo->flinfo->fn_extra =
2731                         MemoryContextAlloc(fcinfo->flinfo->fn_mcxt,
2732                                                            offsetof(RecordIOData, columns) +
2733                                                            ncolumns * sizeof(ColumnIOData));
2734                 my_extra = (RecordIOData *) fcinfo->flinfo->fn_extra;
2735                 my_extra->record_type = InvalidOid;
2736                 my_extra->record_typmod = 0;
2737         }
2738
2739         if (my_extra->record_type != tupType ||
2740                 my_extra->record_typmod != tupTypmod)
2741         {
2742                 MemSet(my_extra, 0,
2743                            offsetof(RecordIOData, columns) +
2744                            ncolumns * sizeof(ColumnIOData));
2745                 my_extra->record_type = tupType;
2746                 my_extra->record_typmod = tupTypmod;
2747                 my_extra->ncolumns = ncolumns;
2748         }
2749
2750         state = palloc0(sizeof(PopulateRecordsetState));
2751
2752         /* make these in a sufficiently long-lived memory context */
2753         old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory);
2754         state->ret_tdesc = CreateTupleDescCopy(tupdesc);
2755         BlessTupleDesc(state->ret_tdesc);
2756         state->tuple_store = tuplestore_begin_heap(rsi->allowedModes &
2757                                                                                            SFRM_Materialize_Random,
2758                                                                                            false, work_mem);
2759         MemoryContextSwitchTo(old_cxt);
2760
2761         state->function_name = funcname;
2762         state->my_extra = my_extra;
2763         state->rec = rec;
2764         state->fn_mcxt = fcinfo->flinfo->fn_mcxt;
2765
2766         if (jtype == JSONOID)
2767         {
2768                 text       *json = PG_GETARG_TEXT_P(json_arg_num);
2769                 JsonLexContext *lex;
2770                 JsonSemAction *sem;
2771
2772                 sem = palloc0(sizeof(JsonSemAction));
2773
2774                 lex = makeJsonLexContext(json, true);
2775
2776                 sem->semstate = (void *) state;
2777                 sem->array_start = populate_recordset_array_start;
2778                 sem->array_element_start = populate_recordset_array_element_start;
2779                 sem->scalar = populate_recordset_scalar;
2780                 sem->object_field_start = populate_recordset_object_field_start;
2781                 sem->object_field_end = populate_recordset_object_field_end;
2782                 sem->object_start = populate_recordset_object_start;
2783                 sem->object_end = populate_recordset_object_end;
2784
2785                 state->lex = lex;
2786
2787                 pg_parse_json(lex, sem);
2788         }
2789         else
2790         {
2791                 Jsonb      *jb = PG_GETARG_JSONB(json_arg_num);
2792                 JsonbIterator *it;
2793                 JsonbValue      v;
2794                 bool            skipNested = false;
2795                 int                     r;
2796
2797                 Assert(jtype == JSONBOID);
2798
2799                 if (JB_ROOT_IS_SCALAR(jb) || !JB_ROOT_IS_ARRAY(jb))
2800                         ereport(ERROR,
2801                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2802                                          errmsg("cannot call %s on a non-array",
2803                                                         funcname)));
2804
2805                 it = JsonbIteratorInit(&jb->root);
2806
2807                 while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
2808                 {
2809                         skipNested = true;
2810
2811                         if (r == WJB_ELEM)
2812                         {
2813                                 Jsonb      *element = JsonbValueToJsonb(&v);
2814
2815                                 if (!JB_ROOT_IS_OBJECT(element))
2816                                         ereport(ERROR,
2817                                                         (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2818                                                  errmsg("argument of %s must be an array of objects",
2819                                                                 funcname)));
2820                                 make_row_from_rec_and_jsonb(element, state);
2821                         }
2822                 }
2823         }
2824
2825         rsi->setResult = state->tuple_store;
2826         rsi->setDesc = state->ret_tdesc;
2827
2828         PG_RETURN_NULL();
2829 }
2830
2831 static void
2832 populate_recordset_object_start(void *state)
2833 {
2834         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
2835         int                     lex_level = _state->lex->lex_level;
2836         HASHCTL         ctl;
2837
2838         /* Reject object at top level: we must have an array at level 0 */
2839         if (lex_level == 0)
2840                 ereport(ERROR,
2841                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2842                                  errmsg("cannot call %s on an object",
2843                                                 _state->function_name)));
2844
2845         /* Nested objects require no special processing */
2846         if (lex_level > 1)
2847                 return;
2848
2849         /* Object at level 1: set up a new hash table for this object */
2850         memset(&ctl, 0, sizeof(ctl));
2851         ctl.keysize = NAMEDATALEN;
2852         ctl.entrysize = sizeof(JsonHashEntry);
2853         ctl.hcxt = CurrentMemoryContext;
2854         _state->json_hash = hash_create("json object hashtable",
2855                                                                         100,
2856                                                                         &ctl,
2857                                                                         HASH_ELEM | HASH_CONTEXT);
2858 }
2859
2860 static void
2861 populate_recordset_object_end(void *state)
2862 {
2863         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
2864         HTAB       *json_hash = _state->json_hash;
2865         Datum      *values;
2866         bool       *nulls;
2867         int                     i;
2868         RecordIOData *my_extra = _state->my_extra;
2869         int                     ncolumns = my_extra->ncolumns;
2870         TupleDesc       tupdesc = _state->ret_tdesc;
2871         JsonHashEntry *hashentry;
2872         HeapTupleHeader rec = _state->rec;
2873         HeapTuple       rettuple;
2874
2875         /* Nested objects require no special processing */
2876         if (_state->lex->lex_level > 1)
2877                 return;
2878
2879         /* Otherwise, construct and return a tuple based on this level-1 object */
2880         values = (Datum *) palloc(ncolumns * sizeof(Datum));
2881         nulls = (bool *) palloc(ncolumns * sizeof(bool));
2882
2883         if (_state->rec)
2884         {
2885                 HeapTupleData tuple;
2886
2887                 /* Build a temporary HeapTuple control structure */
2888                 tuple.t_len = HeapTupleHeaderGetDatumLength(_state->rec);
2889                 ItemPointerSetInvalid(&(tuple.t_self));
2890                 tuple.t_tableOid = InvalidOid;
2891                 tuple.t_data = _state->rec;
2892
2893                 /* Break down the tuple into fields */
2894                 heap_deform_tuple(&tuple, tupdesc, values, nulls);
2895         }
2896         else
2897         {
2898                 for (i = 0; i < ncolumns; ++i)
2899                 {
2900                         values[i] = (Datum) 0;
2901                         nulls[i] = true;
2902                 }
2903         }
2904
2905         for (i = 0; i < ncolumns; ++i)
2906         {
2907                 ColumnIOData *column_info = &my_extra->columns[i];
2908                 Oid                     column_type = tupdesc->attrs[i]->atttypid;
2909                 char       *value;
2910
2911                 /* Ignore dropped columns in datatype */
2912                 if (tupdesc->attrs[i]->attisdropped)
2913                 {
2914                         nulls[i] = true;
2915                         continue;
2916                 }
2917
2918                 hashentry = hash_search(json_hash,
2919                                                                 NameStr(tupdesc->attrs[i]->attname),
2920                                                                 HASH_FIND, NULL);
2921
2922                 /*
2923                  * we can't just skip here if the key wasn't found since we might have
2924                  * a domain to deal with. If we were passed in a non-null record
2925                  * datum, we assume that the existing values are valid (if they're
2926                  * not, then it's not our fault), but if we were passed in a null,
2927                  * then every field which we don't populate needs to be run through
2928                  * the input function just in case it's a domain type.
2929                  */
2930                 if (hashentry == NULL && rec)
2931                         continue;
2932
2933                 /*
2934                  * Prepare to convert the column value from text
2935                  */
2936                 if (column_info->column_type != column_type)
2937                 {
2938                         getTypeInputInfo(column_type,
2939                                                          &column_info->typiofunc,
2940                                                          &column_info->typioparam);
2941                         fmgr_info_cxt(column_info->typiofunc, &column_info->proc,
2942                                                   _state->fn_mcxt);
2943                         column_info->column_type = column_type;
2944                 }
2945                 if (hashentry == NULL || hashentry->isnull)
2946                 {
2947                         /*
2948                          * need InputFunctionCall to happen even for nulls, so that domain
2949                          * checks are done
2950                          */
2951                         values[i] = InputFunctionCall(&column_info->proc, NULL,
2952                                                                                   column_info->typioparam,
2953                                                                                   tupdesc->attrs[i]->atttypmod);
2954                         nulls[i] = true;
2955                 }
2956                 else
2957                 {
2958                         value = hashentry->val;
2959
2960                         values[i] = InputFunctionCall(&column_info->proc, value,
2961                                                                                   column_info->typioparam,
2962                                                                                   tupdesc->attrs[i]->atttypmod);
2963                         nulls[i] = false;
2964                 }
2965         }
2966
2967         rettuple = heap_form_tuple(tupdesc, values, nulls);
2968
2969         tuplestore_puttuple(_state->tuple_store, rettuple);
2970
2971         /* Done with hash for this object */
2972         hash_destroy(json_hash);
2973         _state->json_hash = NULL;
2974 }
2975
2976 static void
2977 populate_recordset_array_element_start(void *state, bool isnull)
2978 {
2979         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
2980
2981         if (_state->lex->lex_level == 1 &&
2982                 _state->lex->token_type != JSON_TOKEN_OBJECT_START)
2983                 ereport(ERROR,
2984                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
2985                                  errmsg("argument of %s must be an array of objects",
2986                                                 _state->function_name)));
2987 }
2988
2989 static void
2990 populate_recordset_array_start(void *state)
2991 {
2992         /* nothing to do */
2993 }
2994
2995 static void
2996 populate_recordset_scalar(void *state, char *token, JsonTokenType tokentype)
2997 {
2998         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
2999
3000         if (_state->lex->lex_level == 0)
3001                 ereport(ERROR,
3002                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3003                                  errmsg("cannot call %s on a scalar",
3004                                                 _state->function_name)));
3005
3006         if (_state->lex->lex_level == 2)
3007                 _state->saved_scalar = token;
3008 }
3009
3010 static void
3011 populate_recordset_object_field_start(void *state, char *fname, bool isnull)
3012 {
3013         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3014
3015         if (_state->lex->lex_level > 2)
3016                 return;
3017
3018         if (_state->lex->token_type == JSON_TOKEN_ARRAY_START ||
3019                 _state->lex->token_type == JSON_TOKEN_OBJECT_START)
3020         {
3021                 _state->save_json_start = _state->lex->token_start;
3022         }
3023         else
3024         {
3025                 _state->save_json_start = NULL;
3026         }
3027 }
3028
3029 static void
3030 populate_recordset_object_field_end(void *state, char *fname, bool isnull)
3031 {
3032         PopulateRecordsetState *_state = (PopulateRecordsetState *) state;
3033         JsonHashEntry *hashentry;
3034         bool            found;
3035
3036         /*
3037          * Ignore nested fields.
3038          */
3039         if (_state->lex->lex_level > 2)
3040                 return;
3041
3042         /*
3043          * Ignore field names >= NAMEDATALEN - they can't match a record field.
3044          * (Note: without this test, the hash code would truncate the string at
3045          * NAMEDATALEN-1, and could then match against a similarly-truncated
3046          * record field name.  That would be a reasonable behavior, but this code
3047          * has previously insisted on exact equality, so we keep this behavior.)
3048          */
3049         if (strlen(fname) >= NAMEDATALEN)
3050                 return;
3051
3052         hashentry = hash_search(_state->json_hash, fname, HASH_ENTER, &found);
3053
3054         /*
3055          * found being true indicates a duplicate. We don't do anything about
3056          * that, a later field with the same name overrides the earlier field.
3057          */
3058
3059         hashentry->isnull = isnull;
3060         if (_state->save_json_start != NULL)
3061         {
3062                 int                     len = _state->lex->prev_token_terminator - _state->save_json_start;
3063                 char       *val = palloc((len + 1) * sizeof(char));
3064
3065                 memcpy(val, _state->save_json_start, len);
3066                 val[len] = '\0';
3067                 hashentry->val = val;
3068         }
3069         else
3070         {
3071                 /* must have had a scalar instead */
3072                 hashentry->val = _state->saved_scalar;
3073         }
3074 }
3075
3076 /*
3077  * findJsonbValueFromContainer() wrapper that sets up JsonbValue key string.
3078  */
3079 static JsonbValue *
3080 findJsonbValueFromContainerLen(JsonbContainer *container, uint32 flags,
3081                                                            char *key, uint32 keylen)
3082 {
3083         JsonbValue      k;
3084
3085         k.type = jbvString;
3086         k.val.string.val = key;
3087         k.val.string.len = keylen;
3088
3089         return findJsonbValueFromContainer(container, flags, &k);
3090 }
3091
3092 /*
3093  * Semantic actions for json_strip_nulls.
3094  *
3095  * Simply repeat the input on the output unless we encounter
3096  * a null object field. State for this is set when the field
3097  * is started and reset when the scalar action (which must be next)
3098  * is called.
3099  */
3100
3101 static void
3102 sn_object_start(void *state)
3103 {
3104         StripnullState *_state = (StripnullState *) state;
3105
3106         appendStringInfoCharMacro(_state->strval, '{');
3107 }
3108
3109 static void
3110 sn_object_end(void *state)
3111 {
3112         StripnullState *_state = (StripnullState *) state;
3113
3114         appendStringInfoCharMacro(_state->strval, '}');
3115 }
3116
3117 static void
3118 sn_array_start(void *state)
3119 {
3120         StripnullState *_state = (StripnullState *) state;
3121
3122         appendStringInfoCharMacro(_state->strval, '[');
3123 }
3124
3125 static void
3126 sn_array_end(void *state)
3127 {
3128         StripnullState *_state = (StripnullState *) state;
3129
3130         appendStringInfoCharMacro(_state->strval, ']');
3131 }
3132
3133 static void
3134 sn_object_field_start(void *state, char *fname, bool isnull)
3135 {
3136         StripnullState *_state = (StripnullState *) state;
3137
3138         if (isnull)
3139         {
3140                 /*
3141                  * The next thing must be a scalar or isnull couldn't be true, so
3142                  * there is no danger of this state being carried down into a nested
3143                  * object or array. The flag will be reset in the scalar action.
3144                  */
3145                 _state->skip_next_null = true;
3146                 return;
3147         }
3148
3149         if (_state->strval->data[_state->strval->len - 1] != '{')
3150                 appendStringInfoCharMacro(_state->strval, ',');
3151
3152         /*
3153          * Unfortunately we don't have the quoted and escaped string any more, so
3154          * we have to re-escape it.
3155          */
3156         escape_json(_state->strval, fname);
3157
3158         appendStringInfoCharMacro(_state->strval, ':');
3159 }
3160
3161 static void
3162 sn_array_element_start(void *state, bool isnull)
3163 {
3164         StripnullState *_state = (StripnullState *) state;
3165
3166         if (_state->strval->data[_state->strval->len - 1] != '[')
3167                 appendStringInfoCharMacro(_state->strval, ',');
3168 }
3169
3170 static void
3171 sn_scalar(void *state, char *token, JsonTokenType tokentype)
3172 {
3173         StripnullState *_state = (StripnullState *) state;
3174
3175         if (_state->skip_next_null)
3176         {
3177                 Assert(tokentype == JSON_TOKEN_NULL);
3178                 _state->skip_next_null = false;
3179                 return;
3180         }
3181
3182         if (tokentype == JSON_TOKEN_STRING)
3183                 escape_json(_state->strval, token);
3184         else
3185                 appendStringInfoString(_state->strval, token);
3186 }
3187
3188 /*
3189  * SQL function json_strip_nulls(json) -> json
3190  */
3191 Datum
3192 json_strip_nulls(PG_FUNCTION_ARGS)
3193 {
3194         text       *json = PG_GETARG_TEXT_P(0);
3195         StripnullState *state;
3196         JsonLexContext *lex;
3197         JsonSemAction *sem;
3198
3199         lex = makeJsonLexContext(json, true);
3200         state = palloc0(sizeof(StripnullState));
3201         sem = palloc0(sizeof(JsonSemAction));
3202
3203         state->strval = makeStringInfo();
3204         state->skip_next_null = false;
3205         state->lex = lex;
3206
3207         sem->semstate = (void *) state;
3208         sem->object_start = sn_object_start;
3209         sem->object_end = sn_object_end;
3210         sem->array_start = sn_array_start;
3211         sem->array_end = sn_array_end;
3212         sem->scalar = sn_scalar;
3213         sem->array_element_start = sn_array_element_start;
3214         sem->object_field_start = sn_object_field_start;
3215
3216         pg_parse_json(lex, sem);
3217
3218         PG_RETURN_TEXT_P(cstring_to_text_with_len(state->strval->data,
3219                                                                                           state->strval->len));
3220
3221 }
3222
3223 /*
3224  * SQL function jsonb_strip_nulls(jsonb) -> jsonb
3225  */
3226 Datum
3227 jsonb_strip_nulls(PG_FUNCTION_ARGS)
3228 {
3229         Jsonb      *jb = PG_GETARG_JSONB(0);
3230         JsonbIterator *it;
3231         JsonbParseState *parseState = NULL;
3232         JsonbValue *res = NULL;
3233         int                     type;
3234         JsonbValue      v,
3235                                 k;
3236         bool            last_was_key = false;
3237
3238         if (JB_ROOT_IS_SCALAR(jb))
3239                 PG_RETURN_POINTER(jb);
3240
3241         it = JsonbIteratorInit(&jb->root);
3242
3243         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
3244         {
3245                 Assert(!(type == WJB_KEY && last_was_key));
3246
3247                 if (type == WJB_KEY)
3248                 {
3249                         /* stash the key until we know if it has a null value */
3250                         k = v;
3251                         last_was_key = true;
3252                         continue;
3253                 }
3254
3255                 if (last_was_key)
3256                 {
3257                         /* if the last element was a key this one can't be */
3258                         last_was_key = false;
3259
3260                         /* skip this field if value is null */
3261                         if (type == WJB_VALUE && v.type == jbvNull)
3262                                 continue;
3263
3264                         /* otherwise, do a delayed push of the key */
3265                         (void) pushJsonbValue(&parseState, WJB_KEY, &k);
3266                 }
3267
3268                 if (type == WJB_VALUE || type == WJB_ELEM)
3269                         res = pushJsonbValue(&parseState, type, &v);
3270                 else
3271                         res = pushJsonbValue(&parseState, type, NULL);
3272         }
3273
3274         Assert(res != NULL);
3275
3276         PG_RETURN_POINTER(JsonbValueToJsonb(res));
3277 }
3278
3279 /*
3280  * Add values from the jsonb to the parse state.
3281  *
3282  * If the parse state container is an object, the jsonb is pushed as
3283  * a value, not a key.
3284  *
3285  * This needs to be done using an iterator because pushJsonbValue doesn't
3286  * like getting jbvBinary values, so we can't just push jb as a whole.
3287  */
3288 static void
3289 addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb)
3290 {
3291         JsonbIterator *it;
3292         JsonbValue *o = &(*jbps)->contVal;
3293         int                     type;
3294         JsonbValue      v;
3295
3296         it = JsonbIteratorInit(&jb->root);
3297
3298         Assert(o->type == jbvArray || o->type == jbvObject);
3299
3300         if (JB_ROOT_IS_SCALAR(jb))
3301         {
3302                 (void) JsonbIteratorNext(&it, &v, false);               /* skip array header */
3303                 (void) JsonbIteratorNext(&it, &v, false);               /* fetch scalar value */
3304
3305                 switch (o->type)
3306                 {
3307                         case jbvArray:
3308                                 (void) pushJsonbValue(jbps, WJB_ELEM, &v);
3309                                 break;
3310                         case jbvObject:
3311                                 (void) pushJsonbValue(jbps, WJB_VALUE, &v);
3312                                 break;
3313                         default:
3314                                 elog(ERROR, "unexpected parent of nested structure");
3315                 }
3316         }
3317         else
3318         {
3319                 while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
3320                 {
3321                         if (type == WJB_KEY || type == WJB_VALUE || type == WJB_ELEM)
3322                                 (void) pushJsonbValue(jbps, type, &v);
3323                         else
3324                                 (void) pushJsonbValue(jbps, type, NULL);
3325                 }
3326         }
3327
3328 }
3329
3330 /*
3331  * SQL function jsonb_pretty (jsonb)
3332  *
3333  * Pretty-printed text for the jsonb
3334  */
3335 Datum
3336 jsonb_pretty(PG_FUNCTION_ARGS)
3337 {
3338         Jsonb      *jb = PG_GETARG_JSONB(0);
3339         StringInfo      str = makeStringInfo();
3340
3341         JsonbToCStringIndent(str, &jb->root, VARSIZE(jb));
3342
3343         PG_RETURN_TEXT_P(cstring_to_text_with_len(str->data, str->len));
3344 }
3345
3346 /*
3347  * SQL function jsonb_concat (jsonb, jsonb)
3348  *
3349  * function for || operator
3350  */
3351 Datum
3352 jsonb_concat(PG_FUNCTION_ARGS)
3353 {
3354         Jsonb      *jb1 = PG_GETARG_JSONB(0);
3355         Jsonb      *jb2 = PG_GETARG_JSONB(1);
3356         JsonbParseState *state = NULL;
3357         JsonbValue *res;
3358         JsonbIterator *it1,
3359                            *it2;
3360
3361         /*
3362          * If one of the jsonb is empty, just return the other if it's not
3363          * scalar and both are of the same kind.  If it's a scalar or they are
3364          * of different kinds we need to perform the concatenation even if one is
3365          * empty.
3366          */
3367         if (JB_ROOT_IS_OBJECT(jb1) == JB_ROOT_IS_OBJECT(jb2))
3368         {
3369                 if (JB_ROOT_COUNT(jb1) == 0 && !JB_ROOT_IS_SCALAR(jb2))
3370                         PG_RETURN_JSONB(jb2);
3371                 else if (JB_ROOT_COUNT(jb2) == 0 && !JB_ROOT_IS_SCALAR(jb1))
3372                         PG_RETURN_JSONB(jb1);
3373         }
3374
3375         it1 = JsonbIteratorInit(&jb1->root);
3376         it2 = JsonbIteratorInit(&jb2->root);
3377
3378         res = IteratorConcat(&it1, &it2, &state);
3379
3380         Assert(res != NULL);
3381
3382         PG_RETURN_JSONB(JsonbValueToJsonb(res));
3383 }
3384
3385
3386 /*
3387  * SQL function jsonb_delete (jsonb, text)
3388  *
3389  * return a copy of the jsonb with the indicated item
3390  * removed.
3391  */
3392 Datum
3393 jsonb_delete(PG_FUNCTION_ARGS)
3394 {
3395         Jsonb      *in = PG_GETARG_JSONB(0);
3396         text       *key = PG_GETARG_TEXT_PP(1);
3397         char       *keyptr = VARDATA_ANY(key);
3398         int                     keylen = VARSIZE_ANY_EXHDR(key);
3399         JsonbParseState *state = NULL;
3400         JsonbIterator *it;
3401         uint32          r;
3402         JsonbValue      v,
3403                            *res = NULL;
3404         bool            skipNested = false;
3405
3406         if (JB_ROOT_IS_SCALAR(in))
3407                 ereport(ERROR,
3408                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3409                                  errmsg("cannot delete from scalar")));
3410
3411         if (JB_ROOT_COUNT(in) == 0)
3412                 PG_RETURN_JSONB(in);
3413
3414         it = JsonbIteratorInit(&in->root);
3415
3416         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != 0)
3417         {
3418                 skipNested = true;
3419
3420                 if ((r == WJB_ELEM || r == WJB_KEY) &&
3421                         (v.type == jbvString && keylen == v.val.string.len &&
3422                          memcmp(keyptr, v.val.string.val, keylen) == 0))
3423                 {
3424                         /* skip corresponding value as well */
3425                         if (r == WJB_KEY)
3426                                 JsonbIteratorNext(&it, &v, true);
3427
3428                         continue;
3429                 }
3430
3431                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3432         }
3433
3434         Assert(res != NULL);
3435
3436         PG_RETURN_JSONB(JsonbValueToJsonb(res));
3437 }
3438
3439 /*
3440  * SQL function jsonb_delete (jsonb, int)
3441  *
3442  * return a copy of the jsonb with the indicated item
3443  * removed. Negative int means count back from the
3444  * end of the items.
3445  */
3446 Datum
3447 jsonb_delete_idx(PG_FUNCTION_ARGS)
3448 {
3449         Jsonb      *in = PG_GETARG_JSONB(0);
3450         int                     idx = PG_GETARG_INT32(1);
3451         JsonbParseState *state = NULL;
3452         JsonbIterator *it;
3453         uint32          r,
3454                                 i = 0,
3455                                 n;
3456         JsonbValue      v,
3457                            *res = NULL;
3458
3459         if (JB_ROOT_IS_SCALAR(in))
3460                 ereport(ERROR,
3461                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3462                                  errmsg("cannot delete from scalar")));
3463
3464         if (JB_ROOT_IS_OBJECT(in))
3465                 ereport(ERROR,
3466                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3467                                  errmsg("cannot delete from object using integer subscript")));
3468
3469         if (JB_ROOT_COUNT(in) == 0)
3470                 PG_RETURN_JSONB(in);
3471
3472         it = JsonbIteratorInit(&in->root);
3473
3474         r = JsonbIteratorNext(&it, &v, false);
3475         Assert (r == WJB_BEGIN_ARRAY);
3476         n = v.val.array.nElems;
3477
3478         if (idx < 0)
3479         {
3480                 if (-idx > n)
3481                         idx = n;
3482                 else
3483                         idx = n + idx;
3484         }
3485
3486         if (idx >= n)
3487                 PG_RETURN_JSONB(in);
3488
3489         pushJsonbValue(&state, r, NULL);
3490
3491         while ((r = JsonbIteratorNext(&it, &v, true)) != 0)
3492         {
3493                 if (r == WJB_ELEM)
3494                 {
3495                         if (i++ == idx)
3496                                 continue;
3497                 }
3498
3499                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3500         }
3501
3502         Assert(res != NULL);
3503
3504         PG_RETURN_JSONB(JsonbValueToJsonb(res));
3505 }
3506
3507 /*
3508  * SQL function jsonb_set(jsonb, text[], jsonb, boolean)
3509  *
3510  */
3511 Datum
3512 jsonb_set(PG_FUNCTION_ARGS)
3513 {
3514         Jsonb      *in = PG_GETARG_JSONB(0);
3515         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
3516         Jsonb      *newval = PG_GETARG_JSONB(2);
3517         bool            create = PG_GETARG_BOOL(3);
3518         JsonbValue *res = NULL;
3519         Datum      *path_elems;
3520         bool       *path_nulls;
3521         int                     path_len;
3522         JsonbIterator *it;
3523         JsonbParseState *st = NULL;
3524
3525         if (ARR_NDIM(path) > 1)
3526                 ereport(ERROR,
3527                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
3528                                  errmsg("wrong number of array subscripts")));
3529
3530         if (JB_ROOT_IS_SCALAR(in))
3531                 ereport(ERROR,
3532                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3533                                  errmsg("cannot set path in scalar")));
3534
3535         if (JB_ROOT_COUNT(in) == 0 && !create)
3536                 PG_RETURN_JSONB(in);
3537
3538         deconstruct_array(path, TEXTOID, -1, false, 'i',
3539                                           &path_elems, &path_nulls, &path_len);
3540
3541         if (path_len == 0)
3542                 PG_RETURN_JSONB(in);
3543
3544         it = JsonbIteratorInit(&in->root);
3545
3546         res = setPath(&it, path_elems, path_nulls, path_len, &st,
3547                                   0, newval, create);
3548
3549         Assert(res != NULL);
3550
3551         PG_RETURN_JSONB(JsonbValueToJsonb(res));
3552 }
3553
3554
3555 /*
3556  * SQL function jsonb_delete(jsonb, text[])
3557  */
3558 Datum
3559 jsonb_delete_path(PG_FUNCTION_ARGS)
3560 {
3561         Jsonb      *in = PG_GETARG_JSONB(0);
3562         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
3563         JsonbValue *res = NULL;
3564         Datum      *path_elems;
3565         bool       *path_nulls;
3566         int                     path_len;
3567         JsonbIterator *it;
3568         JsonbParseState *st = NULL;
3569
3570         if (ARR_NDIM(path) > 1)
3571                 ereport(ERROR,
3572                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
3573                                  errmsg("wrong number of array subscripts")));
3574
3575         if (JB_ROOT_IS_SCALAR(in))
3576                 ereport(ERROR,
3577                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3578                                  errmsg("cannot delete path in scalar")));
3579
3580         if (JB_ROOT_COUNT(in) == 0)
3581                 PG_RETURN_JSONB(in);
3582
3583         deconstruct_array(path, TEXTOID, -1, false, 'i',
3584                                           &path_elems, &path_nulls, &path_len);
3585
3586         if (path_len == 0)
3587                 PG_RETURN_JSONB(in);
3588
3589         it = JsonbIteratorInit(&in->root);
3590
3591         res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, NULL, false);
3592
3593         Assert(res != NULL);
3594
3595         PG_RETURN_JSONB(JsonbValueToJsonb(res));
3596 }
3597
3598 /*
3599  * Iterate over all jsonb objects and merge them into one.
3600  * The logic of this function copied from the same hstore function,
3601  * except the case, when it1 & it2 represents jbvObject.
3602  * In that case we just append the content of it2 to it1 without any
3603  * verifications.
3604  */
3605 static JsonbValue *
3606 IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
3607                            JsonbParseState **state)
3608 {
3609         uint32          r1,
3610                                 r2,
3611                                 rk1,
3612                                 rk2;
3613         JsonbValue      v1,
3614                                 v2,
3615                            *res = NULL;
3616
3617         r1 = rk1 = JsonbIteratorNext(it1, &v1, false);
3618         r2 = rk2 = JsonbIteratorNext(it2, &v2, false);
3619
3620         /*
3621          * Both elements are objects.
3622          */
3623         if (rk1 == WJB_BEGIN_OBJECT && rk2 == WJB_BEGIN_OBJECT)
3624         {
3625                 /*
3626                  * Append the all tokens from v1 to res, except last WJB_END_OBJECT
3627                  * (because res will not be finished yet).
3628                  */
3629                 pushJsonbValue(state, r1, NULL);
3630                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_OBJECT)
3631                         pushJsonbValue(state, r1, &v1);
3632
3633                 /*
3634                  * Append the all tokens from v2 to res, include last WJB_END_OBJECT
3635                  * (the concatenation will be completed).
3636                  */
3637                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != 0)
3638                         res = pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
3639         }
3640
3641         /*
3642          * Both elements are arrays (either can be scalar).
3643          */
3644         else if (rk1 == WJB_BEGIN_ARRAY && rk2 == WJB_BEGIN_ARRAY)
3645         {
3646                 pushJsonbValue(state, r1, NULL);
3647
3648                 while ((r1 = JsonbIteratorNext(it1, &v1, true)) != WJB_END_ARRAY)
3649                 {
3650                         Assert(r1 == WJB_ELEM);
3651                         pushJsonbValue(state, r1, &v1);
3652                 }
3653
3654                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_END_ARRAY)
3655                 {
3656                         Assert(r2 == WJB_ELEM);
3657                         pushJsonbValue(state, WJB_ELEM, &v2);
3658                 }
3659
3660                 res = pushJsonbValue(state, WJB_END_ARRAY, NULL /* signal to sort */ );
3661         }
3662         /* have we got array || object or object || array? */
3663         else if (((rk1 == WJB_BEGIN_ARRAY && !(*it1)->isScalar) && rk2 == WJB_BEGIN_OBJECT) ||
3664                          (rk1 == WJB_BEGIN_OBJECT && (rk2 == WJB_BEGIN_ARRAY && !(*it2)->isScalar)))
3665         {
3666
3667                 JsonbIterator **it_array = rk1 == WJB_BEGIN_ARRAY ? it1 : it2;
3668                 JsonbIterator **it_object = rk1 == WJB_BEGIN_OBJECT ? it1 : it2;
3669
3670                 bool            prepend = (rk1 == WJB_BEGIN_OBJECT);
3671
3672                 pushJsonbValue(state, WJB_BEGIN_ARRAY, NULL);
3673
3674                 if (prepend)
3675                 {
3676                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
3677                         while ((r1 = JsonbIteratorNext(it_object, &v1, true)) != 0)
3678                                 pushJsonbValue(state, r1, r1 != WJB_END_OBJECT ? &v1 : NULL);
3679
3680                         while ((r2 = JsonbIteratorNext(it_array, &v2, true)) != 0)
3681                                 res = pushJsonbValue(state, r2, r2 != WJB_END_ARRAY ? &v2 : NULL);
3682                 }
3683                 else
3684                 {
3685                         while ((r1 = JsonbIteratorNext(it_array, &v1, true)) != WJB_END_ARRAY)
3686                                 pushJsonbValue(state, r1, &v1);
3687
3688                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
3689                         while ((r2 = JsonbIteratorNext(it_object, &v2, true)) != 0)
3690                                 pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
3691
3692                         res = pushJsonbValue(state, WJB_END_ARRAY, NULL);
3693                 }
3694         }
3695         else
3696         {
3697                 /*
3698                  * This must be scalar || object or object || scalar, as that's all
3699                  * that's left. Both of these make no sense, so error out.
3700                  */
3701                 ereport(ERROR,
3702                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3703                                  errmsg("invalid concatenation of jsonb objects")));
3704         }
3705
3706         return res;
3707 }
3708
3709 /*
3710  * Do most of the heavy work for jsonb_set
3711  *
3712  * If newval is null, the element is to be removed.
3713  *
3714  * If create is true, we create the new value if the key or array index
3715  * does not exist. All path elements before the last must already exist
3716  * whether or not create is true, or nothing is done.
3717  */
3718 static JsonbValue *
3719 setPath(JsonbIterator **it, Datum *path_elems,
3720                 bool *path_nulls, int path_len,
3721                 JsonbParseState **st, int level, Jsonb *newval, bool create)
3722 {
3723         JsonbValue      v;
3724         JsonbValue *res = NULL;
3725         int                     r;
3726
3727         check_stack_depth();
3728
3729         if (path_nulls[level])
3730                 elog(ERROR, "path element at the position %d is NULL", level + 1);
3731
3732         r = JsonbIteratorNext(it, &v, false);
3733
3734         switch (r)
3735         {
3736                 case WJB_BEGIN_ARRAY:
3737                         (void) pushJsonbValue(st, r, NULL);
3738                         setPathArray(it, path_elems, path_nulls, path_len, st, level,
3739                                                  newval, v.val.array.nElems, create);
3740                         r = JsonbIteratorNext(it, &v, false);
3741                         Assert(r == WJB_END_ARRAY);
3742                         res = pushJsonbValue(st, r, NULL);
3743
3744                         break;
3745                 case WJB_BEGIN_OBJECT:
3746                         (void) pushJsonbValue(st, r, NULL);
3747                         setPathObject(it, path_elems, path_nulls, path_len, st, level,
3748                                                   newval, v.val.object.nPairs, create);
3749                         r = JsonbIteratorNext(it, &v, true);
3750                         Assert(r == WJB_END_OBJECT);
3751                         res = pushJsonbValue(st, r, NULL);
3752
3753                         break;
3754                 case WJB_ELEM:
3755                 case WJB_VALUE:
3756                         res = pushJsonbValue(st, r, &v);
3757                         break;
3758                 default:
3759                         elog(ERROR, "impossible state");
3760         }
3761
3762         return res;
3763 }
3764
3765 /*
3766  * Object walker for setPath
3767  */
3768 static void
3769 setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
3770                           int path_len, JsonbParseState **st, int level,
3771                           Jsonb *newval, uint32 npairs, bool create)
3772 {
3773         JsonbValue      v;
3774         int                     i;
3775         JsonbValue      k;
3776         bool            done = false;
3777
3778         if (level >= path_len || path_nulls[level])
3779                 done = true;
3780
3781         /* empty object is a special case for create */
3782         if ((npairs == 0) && create && (level == path_len - 1))
3783         {
3784                 JsonbValue      newkey;
3785
3786                 newkey.type = jbvString;
3787                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
3788                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
3789
3790                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
3791                 addJsonbToParseState(st, newval);
3792         }
3793
3794         for (i = 0; i < npairs; i++)
3795         {
3796                 int                     r = JsonbIteratorNext(it, &k, true);
3797
3798                 Assert(r == WJB_KEY);
3799
3800                 if (!done &&
3801                         k.val.string.len == VARSIZE_ANY_EXHDR(path_elems[level]) &&
3802                         memcmp(k.val.string.val, VARDATA_ANY(path_elems[level]),
3803                                    k.val.string.len) == 0)
3804                 {
3805                         if (level == path_len - 1)
3806                         {
3807                                 r = JsonbIteratorNext(it, &v, true);    /* skip */
3808                                 if (newval != NULL)
3809                                 {
3810                                         (void) pushJsonbValue(st, WJB_KEY, &k);
3811                                         addJsonbToParseState(st, newval);
3812                                 }
3813                                 done = true;
3814                         }
3815                         else
3816                         {
3817                                 (void) pushJsonbValue(st, r, &k);
3818                                 setPath(it, path_elems, path_nulls, path_len,
3819                                                 st, level + 1, newval, create);
3820                         }
3821                 }
3822                 else
3823                 {
3824                         if (create && !done && level == path_len - 1 && i == npairs - 1)
3825                         {
3826                                 JsonbValue      newkey;
3827
3828                                 newkey.type = jbvString;
3829                                 newkey.val.string.len = VARSIZE_ANY_EXHDR(path_elems[level]);
3830                                 newkey.val.string.val = VARDATA_ANY(path_elems[level]);
3831
3832                                 (void) pushJsonbValue(st, WJB_KEY, &newkey);
3833                                 addJsonbToParseState(st, newval);
3834                         }
3835
3836                         (void) pushJsonbValue(st, r, &k);
3837                         r = JsonbIteratorNext(it, &v, false);
3838                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3839                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3840                         {
3841                                 int                     walking_level = 1;
3842
3843                                 while (walking_level != 0)
3844                                 {
3845                                         r = JsonbIteratorNext(it, &v, false);
3846
3847                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3848                                                 ++walking_level;
3849                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
3850                                                 --walking_level;
3851
3852                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3853                                 }
3854                         }
3855                 }
3856         }
3857 }
3858
3859 /*
3860  * Array walker for setPath
3861  */
3862 static void
3863 setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
3864                          int path_len, JsonbParseState **st, int level,
3865                          Jsonb *newval, uint32 nelems, bool create)
3866 {
3867         JsonbValue      v;
3868         int                     idx,
3869                                 i;
3870         char       *badp;
3871         bool            done = false;
3872
3873         /* pick correct index */
3874         if (level < path_len && !path_nulls[level])
3875         {
3876                 char       *c = VARDATA_ANY(path_elems[level]);
3877                 long            lindex;
3878
3879                 errno = 0;
3880                 lindex = strtol(c, &badp, 10);
3881                 if (errno != 0 || badp == c || *badp != '\0' || lindex > INT_MAX ||
3882                         lindex < INT_MIN)
3883                         elog(ERROR, "path element at the position %d is not an integer", level + 1);
3884                 else
3885                         idx = lindex;
3886         }
3887         else
3888                 idx = nelems;
3889
3890         if (idx < 0)
3891         {
3892                 if (-idx > nelems)
3893                         idx = INT_MIN;
3894                 else
3895                         idx = nelems + idx;
3896         }
3897
3898         if (idx > 0 && idx > nelems)
3899                 idx = nelems;
3900
3901         /*
3902          * if we're creating, and idx == INT_MIN, we prepend the new value to the
3903          * array also if the array is empty - in which case we don't really care
3904          * what the idx value is
3905          */
3906
3907         if ((idx == INT_MIN || nelems == 0) && create && (level == path_len - 1))
3908         {
3909                 Assert(newval != NULL);
3910                 addJsonbToParseState(st, newval);
3911                 done = true;
3912         }
3913
3914         /* iterate over the array elements */
3915         for (i = 0; i < nelems; i++)
3916         {
3917                 int                     r;
3918
3919                 if (i == idx && level < path_len)
3920                 {
3921                         if (level == path_len - 1)
3922                         {
3923                                 r = JsonbIteratorNext(it, &v, true);    /* skip */
3924                                 if (newval != NULL)
3925                                         addJsonbToParseState(st, newval);
3926
3927                                 done = true;
3928                         }
3929                         else
3930                                 (void) setPath(it, path_elems, path_nulls, path_len,
3931                                                            st, level + 1, newval, create);
3932                 }
3933                 else
3934                 {
3935                         r = JsonbIteratorNext(it, &v, false);
3936
3937                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3938
3939                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3940                         {
3941                                 int                     walking_level = 1;
3942
3943                                 while (walking_level != 0)
3944                                 {
3945                                         r = JsonbIteratorNext(it, &v, false);
3946
3947                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3948                                                 ++walking_level;
3949                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
3950                                                 --walking_level;
3951
3952                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3953                                 }
3954                         }
3955
3956                         if (create && !done && level == path_len - 1 && i == nelems - 1)
3957                         {
3958                                 addJsonbToParseState(st, newval);
3959                         }
3960
3961                 }
3962         }
3963 }