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