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 *walkJsonb(JsonbIterator **it, JsonbParseState **state, bool stop_at_level_zero);
 131 static JsonbValue *replacePath(JsonbIterator **it, Datum *path_elems,
 132                                                            bool *path_nulls, int path_len,
 133                                                            JsonbParseState **st, int level, Jsonb *newval);
 134 static void replacePathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 135                                                           int path_len, JsonbParseState **st, int level,
 136                                                           Jsonb *newval, uint32 nelems);
 137 static void replacePathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 138                                                          int path_len, JsonbParseState **st, int level,
 139                                                          Jsonb *newval, uint32 npairs);
 140 static void addJsonbToParseState(JsonbParseState **jbps, Jsonb * jb);
 141
 142 /* state for json_object_keys */
 143 typedef struct OkeysState
 144 {
 145         JsonLexContext *lex;
 146         char      **result;
 147         int                     result_size;
 148         int                     result_count;
 149         int                     sent_count;
 150 } OkeysState;
 151
 152 /* state for json_get* functions */
 153 typedef struct GetState
 154 {
 155         JsonLexContext *lex;
 156         text       *tresult;
 157         char       *result_start;
 158         bool            normalize_results;
 159         bool            next_scalar;
 160         int                     npath;                  /* length of each path-related array */
 161         char      **path_names;         /* field name(s) being sought */
 162         int                *path_indexes;       /* array index(es) being sought */
 163         bool       *pathok;                     /* is path matched to current depth? */
 164         int                *array_cur_index;    /* current element index at each path level */
 165 } GetState;
 166
 167 /* state for json_array_length */
 168 typedef struct AlenState
 169 {
 170         JsonLexContext *lex;
 171         int                     count;
 172 } AlenState;
 173
 174 /* state for json_each */
 175 typedef struct EachState
 176 {
 177         JsonLexContext *lex;
 178         Tuplestorestate *tuple_store;
 179         TupleDesc       ret_tdesc;
 180         MemoryContext tmp_cxt;
 181         char       *result_start;
 182         bool            normalize_results;
 183         bool            next_scalar;
 184         char       *normalized_scalar;
 185 } EachState;
 186
 187 /* state for json_array_elements */
 188 typedef struct ElementsState
 189 {
 190         JsonLexContext *lex;
 191         const char *function_name;
 192         Tuplestorestate *tuple_store;
 193         TupleDesc       ret_tdesc;
 194         MemoryContext tmp_cxt;
 195         char       *result_start;
 196         bool            normalize_results;
 197         bool            next_scalar;
 198         char       *normalized_scalar;
 199 } ElementsState;
 200
 201 /* state for get_json_object_as_hash */
 202 typedef struct JhashState
 203 {
 204         JsonLexContext *lex;
 205         const char *function_name;
 206         HTAB       *hash;
 207         char       *saved_scalar;
 208         char       *save_json_start;
 209 } JHashState;
 210
 211 /* hashtable element */
 212 typedef struct JsonHashEntry
 213 {
 214         char            fname[NAMEDATALEN];             /* hash key (MUST BE FIRST) */
 215         char       *val;
 216         char       *json;
 217         bool            isnull;
 218 } JsonHashEntry;
 219
 220 /* these two are stolen from hstore / record_out, used in populate_record* */
 221 typedef struct ColumnIOData
 222 {
 223         Oid                     column_type;
 224         Oid                     typiofunc;
 225         Oid                     typioparam;
 226         FmgrInfo        proc;
 227 } ColumnIOData;
 228
 229 typedef struct RecordIOData
 230 {
 231         Oid                     record_type;
 232         int32           record_typmod;
 233         int                     ncolumns;
 234         ColumnIOData columns[FLEXIBLE_ARRAY_MEMBER];
 235 } RecordIOData;
 236
 237 /* state for populate_recordset */
 238 typedef struct PopulateRecordsetState
 239 {
 240         JsonLexContext *lex;
 241         const char *function_name;
 242         HTAB       *json_hash;
 243         char       *saved_scalar;
 244         char       *save_json_start;
 245         Tuplestorestate *tuple_store;
 246         TupleDesc       ret_tdesc;
 247         HeapTupleHeader rec;
 248         RecordIOData *my_extra;
 249         MemoryContext fn_mcxt;          /* used to stash IO funcs */
 250 } PopulateRecordsetState;
 251
 252 /* state for json_strip_nulls */
 253 typedef struct StripnullState{
 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         appendStringInfoCharMacro(_state->strval, '{');
3049 }
3050
3051 static void
3052 sn_object_end(void *state)
3053 {
3054         StripnullState *_state = (StripnullState *) state;
3055         appendStringInfoCharMacro(_state->strval, '}');
3056 }
3057
3058 static void
3059 sn_array_start(void *state)
3060 {
3061         StripnullState *_state = (StripnullState *) state;
3062         appendStringInfoCharMacro(_state->strval, '[');
3063 }
3064
3065 static void
3066 sn_array_end(void *state)
3067 {
3068         StripnullState *_state = (StripnullState *) state;
3069         appendStringInfoCharMacro(_state->strval, ']');
3070 }
3071
3072 static void
3073 sn_object_field_start (void *state, char *fname, bool isnull)
3074 {
3075         StripnullState *_state = (StripnullState *) state;
3076
3077         if (isnull)
3078         {
3079                 /*
3080                  * The next thing must be a scalar or isnull couldn't be true,
3081                  * so there is no danger of this state being carried down
3082                  * into a nested  object or array. The flag will be reset in the
3083                  * scalar action.
3084                  */
3085                 _state->skip_next_null = true;
3086                 return;
3087         }
3088
3089         if (_state->strval->data[_state->strval->len - 1] != '{')
3090                 appendStringInfoCharMacro(_state->strval, ',');
3091
3092         /*
3093          * Unfortunately we don't have the quoted and escaped string any more,
3094          * so we have to re-escape it.
3095          */
3096         escape_json(_state->strval,fname);
3097
3098         appendStringInfoCharMacro(_state->strval, ':');
3099 }
3100
3101 static void
3102 sn_array_element_start (void *state, bool isnull)
3103 {
3104         StripnullState *_state = (StripnullState *) state;
3105
3106         if (_state->strval->data[_state->strval->len - 1] != '[')
3107                 appendStringInfoCharMacro(_state->strval, ',');
3108 }
3109
3110 static void
3111 sn_scalar(void *state, char *token, JsonTokenType tokentype)
3112 {
3113         StripnullState *_state = (StripnullState *) state;
3114
3115         if (_state->skip_next_null)
3116         {
3117                 Assert (tokentype == JSON_TOKEN_NULL);
3118                 _state->skip_next_null = false;
3119                 return;
3120         }
3121
3122         if (tokentype == JSON_TOKEN_STRING)
3123                 escape_json(_state->strval, token);
3124         else
3125                 appendStringInfoString(_state->strval, token);
3126 }
3127
3128 /*
3129  * SQL function json_strip_nulls(json) -> json
3130  */
3131 Datum
3132 json_strip_nulls(PG_FUNCTION_ARGS)
3133 {
3134         text       *json = PG_GETARG_TEXT_P(0);
3135         StripnullState  *state;
3136         JsonLexContext *lex;
3137         JsonSemAction *sem;
3138
3139         lex = makeJsonLexContext(json, true);
3140         state = palloc0(sizeof(StripnullState));
3141         sem = palloc0(sizeof(JsonSemAction));
3142
3143         state->strval = makeStringInfo();
3144         state->skip_next_null = false;
3145         state->lex = lex;
3146
3147         sem->semstate = (void *) state;
3148         sem->object_start = sn_object_start;
3149         sem->object_end = sn_object_end;
3150         sem->array_start = sn_array_start;
3151         sem->array_end = sn_array_end;
3152         sem->scalar = sn_scalar;
3153         sem->array_element_start = sn_array_element_start;
3154         sem->object_field_start = sn_object_field_start;
3155
3156         pg_parse_json(lex, sem);
3157
3158         PG_RETURN_TEXT_P(cstring_to_text_with_len(state->strval->data,
3159                                                                                           state->strval->len));
3160
3161 }
3162
3163 /*
3164  * SQL function jsonb_strip_nulls(jsonb) -> jsonb
3165  */
3166 Datum
3167 jsonb_strip_nulls(PG_FUNCTION_ARGS)
3168 {
3169         Jsonb * jb = PG_GETARG_JSONB(0);
3170         JsonbIterator *it;
3171         JsonbParseState *parseState = NULL;
3172         JsonbValue *res = NULL;
3173         int type;
3174         JsonbValue v,k;
3175         bool last_was_key = false;
3176
3177         if (JB_ROOT_IS_SCALAR(jb))
3178                 PG_RETURN_POINTER(jb);
3179
3180         it = JsonbIteratorInit(&jb->root);
3181
3182         while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
3183         {
3184                 Assert( ! (type == WJB_KEY && last_was_key));
3185
3186                 if (type == WJB_KEY)
3187                 {
3188                         /* stash the key until we know if it has a null value */
3189                         k = v;
3190                         last_was_key = true;
3191                         continue;
3192                 }
3193
3194                 if (last_was_key)
3195                 {
3196                         /* if the last element was a key this one can't be */
3197                         last_was_key = false;
3198
3199                         /* skip this field if value is null */
3200                         if (type == WJB_VALUE && v.type == jbvNull)
3201                                 continue;
3202
3203                         /* otherwise, do a delayed push of the key */
3204                         (void) pushJsonbValue(&parseState, WJB_KEY, &k);
3205                 }
3206
3207                 if (type == WJB_VALUE || type == WJB_ELEM)
3208                         res = pushJsonbValue(&parseState, type, &v);
3209                 else
3210                         res = pushJsonbValue(&parseState, type, NULL);
3211         }
3212
3213         Assert(res != NULL);
3214
3215         PG_RETURN_POINTER(JsonbValueToJsonb(res));
3216 }
3217
3218 /*
3219  * Add values from the jsonb to the parse state.
3220  *
3221  * If the parse state container is an object, the jsonb is pushed as
3222  * a value, not a key.
3223  *
3224  * This needs to be done using an iterator because pushJsonbValue doesn't
3225  * like getting jbvBinary values, so we can't just push jb as a whole.
3226  */
3227 static void
3228 addJsonbToParseState(JsonbParseState **jbps, Jsonb * jb)
3229 {
3230
3231         JsonbIterator *it;
3232         JsonbValue    *o = &(*jbps)->contVal;
3233         int            type;
3234         JsonbValue     v;
3235
3236         it = JsonbIteratorInit(&jb->root);
3237
3238         Assert(o->type == jbvArray || o->type == jbvObject);
3239
3240         if (JB_ROOT_IS_SCALAR(jb))
3241         {
3242                 (void) JsonbIteratorNext(&it, &v, false); /* skip array header */
3243                 (void) JsonbIteratorNext(&it, &v, false); /* fetch scalar value */
3244
3245                 switch (o->type)
3246                 {
3247                         case jbvArray:
3248                                 (void) pushJsonbValue(jbps, WJB_ELEM, &v);
3249                                 break;
3250                         case jbvObject:
3251                                 (void) pushJsonbValue(jbps, WJB_VALUE, &v);
3252                                 break;
3253                         default:
3254                                 elog(ERROR, "unexpected parent of nested structure");
3255                 }
3256         }
3257         else
3258         {
3259                 while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
3260                 {
3261                         if (type == WJB_KEY || type == WJB_VALUE || type == WJB_ELEM)
3262                                 (void) pushJsonbValue(jbps, type, &v);
3263                         else
3264                                 (void) pushJsonbValue(jbps, type, NULL);
3265                 }
3266         }
3267
3268 }
3269
3270 /*
3271  * SQL function jsonb_pretty (jsonb)
3272  *
3273  * Pretty-printed text for the jsonb
3274  */
3275 Datum
3276 jsonb_pretty(PG_FUNCTION_ARGS)
3277 {
3278         Jsonb      *jb = PG_GETARG_JSONB(0);
3279         StringInfo      str = makeStringInfo();
3280
3281         JsonbToCStringIndent(str, &jb->root, VARSIZE(jb));
3282
3283         PG_RETURN_TEXT_P(cstring_to_text_with_len(str->data, str->len));
3284 }
3285
3286
3287 /*
3288  * SQL function jsonb_concat (jsonb, jsonb)
3289  *
3290  * function for || operator
3291  */
3292 Datum
3293 jsonb_concat(PG_FUNCTION_ARGS)
3294 {
3295         Jsonb      *jb1 = PG_GETARG_JSONB(0);
3296         Jsonb      *jb2 = PG_GETARG_JSONB(1);
3297         Jsonb      *out = palloc(VARSIZE(jb1) + VARSIZE(jb2));
3298         JsonbParseState *state = NULL;
3299         JsonbValue *res;
3300         JsonbIterator  *it1,
3301                                    *it2;
3302
3303         /*
3304          * If one of the jsonb is empty, just return other.
3305          */
3306         if (JB_ROOT_COUNT(jb1) == 0)
3307         {
3308                 memcpy(out, jb2, VARSIZE(jb2));
3309                 PG_RETURN_POINTER(out);
3310         }
3311         else if (JB_ROOT_COUNT(jb2) == 0)
3312         {
3313                 memcpy(out, jb1, VARSIZE(jb1));
3314                 PG_RETURN_POINTER(out);
3315         }
3316
3317         it1 = JsonbIteratorInit(&jb1->root);
3318         it2 = JsonbIteratorInit(&jb2->root);
3319
3320         res = IteratorConcat(&it1, &it2, &state);
3321
3322         if (res == NULL || (res->type == jbvArray && res->val.array.nElems == 0) ||
3323                 (res->type == jbvObject && res->val.object.nPairs == 0))
3324         {
3325                 SET_VARSIZE(out, VARHDRSZ);
3326         }
3327         else
3328         {
3329                 if (res->type == jbvArray && res->val.array.nElems > 1)
3330                         res->val.array.rawScalar = false;
3331
3332                 out = JsonbValueToJsonb(res);
3333         }
3334
3335         PG_RETURN_POINTER(out);
3336 }
3337
3338
3339 /*
3340  * SQL function jsonb_delete (jsonb, text)
3341  *
3342  * return a copy of the jsonb with the indicated item
3343  * removed.
3344  */
3345 Datum
3346 jsonb_delete(PG_FUNCTION_ARGS)
3347 {
3348         Jsonb      *in = PG_GETARG_JSONB(0);
3349         text       *key = PG_GETARG_TEXT_PP(1);
3350         char       *keyptr = VARDATA_ANY(key);
3351         int                     keylen = VARSIZE_ANY_EXHDR(key);
3352         Jsonb      *out = palloc(VARSIZE(in));
3353         JsonbParseState *state = NULL;
3354         JsonbIterator *it;
3355         uint32          r;
3356         JsonbValue      v,
3357                            *res = NULL;
3358         bool            skipNested = false;
3359
3360         SET_VARSIZE(out, VARSIZE(in));
3361
3362         if (JB_ROOT_COUNT(in) == 0)
3363                 PG_RETURN_POINTER(out);
3364
3365         it = JsonbIteratorInit(&in->root);
3366
3367         while ((r = JsonbIteratorNext(&it, &v, skipNested)) != 0)
3368         {
3369                 skipNested = true;
3370
3371                 if ((r == WJB_ELEM || r == WJB_KEY) &&
3372                         (v.type == jbvString && keylen == v.val.string.len &&
3373                          memcmp(keyptr, v.val.string.val, keylen) == 0))
3374                 {
3375                         /* skip corresponding value as well */
3376                         if (r == WJB_KEY)
3377                                 JsonbIteratorNext(&it, &v, true);
3378
3379                         continue;
3380                 }
3381
3382                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3383         }
3384
3385         if (res == NULL || (res->type == jbvArray && res->val.array.nElems == 0) ||
3386                 (res->type == jbvObject && res->val.object.nPairs == 0))
3387                 SET_VARSIZE(out, VARHDRSZ);
3388         else
3389                 out = JsonbValueToJsonb(res);
3390
3391         PG_RETURN_POINTER(out);
3392 }
3393
3394 /*
3395  * SQL function jsonb_delete (jsonb, int)
3396  *
3397  * return a copy of the jsonb with the indicated item
3398  * removed. Negative int means count back from the
3399  * end of the items.
3400  */
3401 Datum
3402 jsonb_delete_idx(PG_FUNCTION_ARGS)
3403 {
3404         Jsonb      *in = PG_GETARG_JSONB(0);
3405         int                     idx = PG_GETARG_INT32(1);
3406         Jsonb      *out = palloc(VARSIZE(in));
3407         JsonbParseState *state = NULL;
3408         JsonbIterator *it;
3409         uint32          r,
3410                                 i = 0,
3411                                 n;
3412         JsonbValue      v,
3413                            *res = NULL;
3414
3415         if (JB_ROOT_COUNT(in) == 0)
3416         {
3417                 memcpy(out, in, VARSIZE(in));
3418                 PG_RETURN_POINTER(out);
3419         }
3420
3421         it = JsonbIteratorInit(&in->root);
3422
3423         r = JsonbIteratorNext(&it, &v, false);
3424         if (r == WJB_BEGIN_ARRAY)
3425                 n = v.val.array.nElems;
3426         else
3427                 n = v.val.object.nPairs;
3428
3429         if (idx < 0)
3430         {
3431                 if (-idx > n)
3432                         idx = n;
3433                 else
3434                         idx = n + idx;
3435         }
3436
3437         if (idx >= n)
3438         {
3439                 memcpy(out, in, VARSIZE(in));
3440                 PG_RETURN_POINTER(out);
3441         }
3442
3443         pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3444
3445         while ((r = JsonbIteratorNext(&it, &v, true)) != 0)
3446         {
3447                 if (r == WJB_ELEM || r == WJB_KEY)
3448                 {
3449                         if (i++ == idx)
3450                         {
3451                                 if (r == WJB_KEY)
3452                                         JsonbIteratorNext(&it, &v, true);       /* skip value */
3453                                 continue;
3454                         }
3455                 }
3456
3457                 res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3458         }
3459
3460         if (res == NULL || (res->type == jbvArray && res->val.array.nElems == 0) ||
3461                 (res->type == jbvObject && res->val.object.nPairs == 0))
3462                 SET_VARSIZE(out, VARHDRSZ);
3463         else
3464                 out = JsonbValueToJsonb(res);
3465
3466         PG_RETURN_POINTER(out);
3467 }
3468
3469 /*
3470  * SQL function jsonb_replace(jsonb, text[], jsonb)
3471  */
3472 Datum
3473 jsonb_replace(PG_FUNCTION_ARGS)
3474 {
3475         Jsonb      *in = PG_GETARG_JSONB(0);
3476         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
3477         Jsonb      *newval = PG_GETARG_JSONB(2);
3478         Jsonb      *out = palloc(VARSIZE(in) + VARSIZE(newval));
3479         JsonbValue *res = NULL;
3480         Datum      *path_elems;
3481         bool       *path_nulls;
3482         int                     path_len;
3483         JsonbIterator *it;
3484         JsonbParseState *st = NULL;
3485
3486         if (ARR_NDIM(path) > 1)
3487                 ereport(ERROR,
3488                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
3489                                  errmsg("wrong number of array subscripts")));
3490
3491         if (JB_ROOT_COUNT(in) == 0)
3492         {
3493                 memcpy(out, in, VARSIZE(in));
3494                 PG_RETURN_POINTER(out);
3495         }
3496
3497         deconstruct_array(path, TEXTOID, -1, false, 'i',
3498                                           &path_elems, &path_nulls, &path_len);
3499
3500         if (path_len == 0)
3501         {
3502                 memcpy(out, in, VARSIZE(in));
3503                 PG_RETURN_POINTER(out);
3504         }
3505
3506         it = JsonbIteratorInit(&in->root);
3507
3508         res = replacePath(&it, path_elems, path_nulls, path_len, &st, 0, newval);
3509
3510         if (res == NULL)
3511                 SET_VARSIZE(out, VARHDRSZ);
3512         else
3513                 out = JsonbValueToJsonb(res);
3514
3515         PG_RETURN_POINTER(out);
3516 }
3517
3518
3519 /*
3520  * SQL function jsonb_delete(jsonb, text[])
3521  */
3522 Datum
3523 jsonb_delete_path(PG_FUNCTION_ARGS)
3524 {
3525         Jsonb      *in = PG_GETARG_JSONB(0);
3526         ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
3527         Jsonb      *out = palloc(VARSIZE(in));
3528         JsonbValue *res = NULL;
3529         Datum      *path_elems;
3530         bool       *path_nulls;
3531         int                     path_len;
3532         JsonbIterator *it;
3533         JsonbParseState *st = NULL;
3534
3535         if (ARR_NDIM(path) > 1)
3536                 ereport(ERROR,
3537                                 (errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
3538                                  errmsg("wrong number of array subscripts")));
3539
3540         if (JB_ROOT_COUNT(in) == 0)
3541         {
3542                 memcpy(out, in, VARSIZE(in));
3543                 PG_RETURN_POINTER(out);
3544         }
3545
3546         deconstruct_array(path, TEXTOID, -1, false, 'i',
3547                                           &path_elems, &path_nulls, &path_len);
3548
3549         if (path_len == 0)
3550         {
3551                 memcpy(out, in, VARSIZE(in));
3552                 PG_RETURN_POINTER(out);
3553         }
3554
3555         it = JsonbIteratorInit(&in->root);
3556
3557         res = replacePath(&it, path_elems, path_nulls, path_len, &st, 0, NULL);
3558
3559         if (res == NULL)
3560                 SET_VARSIZE(out, VARHDRSZ);
3561         else
3562                 out = JsonbValueToJsonb(res);
3563
3564         PG_RETURN_POINTER(out);
3565 }
3566
3567
3568 /*
3569  * Iterate over all jsonb objects and merge them into one.
3570  * The logic of this function copied from the same hstore function,
3571  * except the case, when it1 & it2 represents jbvObject.
3572  * In that case we just append the content of it2 to it1 without any
3573  * verifications.
3574  */
3575 static JsonbValue *
3576 IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
3577                            JsonbParseState **state)
3578 {
3579         uint32          r1,
3580                                 r2,
3581                                 rk1,
3582                                 rk2;
3583         JsonbValue      v1,
3584                                 v2,
3585                            *res = NULL;
3586
3587         r1 = rk1 = JsonbIteratorNext(it1, &v1, false);
3588         r2 = rk2 = JsonbIteratorNext(it2, &v2, false);
3589
3590         /*
3591          * Both elements are objects.
3592          */
3593         if (rk1 == WJB_BEGIN_OBJECT && rk2 == WJB_BEGIN_OBJECT)
3594         {
3595                 int                     level = 1;
3596
3597                 /*
3598                  * Append the all tokens from v1 to res, exept last WJB_END_OBJECT
3599                  * (because res will not be finished yet).
3600                  */
3601                 (void) pushJsonbValue(state, r1, NULL);
3602                 while ((r1 = JsonbIteratorNext(it1, &v1, false)) != 0)
3603                 {
3604                         if (r1 == WJB_BEGIN_OBJECT)
3605                                 ++level;
3606                         else if (r1 == WJB_END_OBJECT)
3607                                 --level;
3608
3609                         if (level != 0)
3610                                 res = pushJsonbValue(state, r1, r1 < WJB_BEGIN_ARRAY ? &v1 : NULL);
3611                 }
3612
3613                 /*
3614                  * Append the all tokens from v2 to res, include last WJB_END_OBJECT
3615                  * (the concatenation will be completed).
3616                  */
3617                 while ((r2 = JsonbIteratorNext(it2, &v2, false)) != 0)
3618                         res = pushJsonbValue(state, r2, r2 < WJB_BEGIN_ARRAY ? &v2 : NULL);
3619         }
3620
3621         /*
3622          * Both elements are arrays (either can be scalar).
3623          */
3624         else if (rk1 == WJB_BEGIN_ARRAY && rk2 == WJB_BEGIN_ARRAY)
3625         {
3626                 res = pushJsonbValue(state, r1, NULL);
3627                 for (;;)
3628                 {
3629                         r1 = JsonbIteratorNext(it1, &v1, true);
3630                         if (r1 == WJB_END_OBJECT || r1 == WJB_END_ARRAY)
3631                                 break;
3632
3633                         Assert(r1 == WJB_KEY || r1 == WJB_VALUE || r1 == WJB_ELEM);
3634                         pushJsonbValue(state, r1, &v1);
3635                 }
3636
3637                 while ((r2 = JsonbIteratorNext(it2, &v2, true)) != 0)
3638                 {
3639                         if (!(r2 == WJB_END_OBJECT || r2 == WJB_END_ARRAY))
3640                         {
3641                                 if (rk1 == WJB_BEGIN_OBJECT)
3642                                 {
3643                                         pushJsonbValue(state, WJB_KEY, NULL);
3644                                         r2 = JsonbIteratorNext(it2, &v2, true);
3645                                         Assert(r2 == WJB_ELEM);
3646                                         pushJsonbValue(state, WJB_VALUE, &v2);
3647                                 }
3648                                 else
3649                                         pushJsonbValue(state, WJB_ELEM, &v2);
3650                         }
3651                 }
3652
3653                 res = pushJsonbValue(state,
3654                                   (rk1 == WJB_BEGIN_OBJECT) ? WJB_END_OBJECT : WJB_END_ARRAY,
3655                                                          NULL /* signal to sort */ );
3656         }
3657         /* have we got array || object or object || array? */
3658         else if (((rk1 == WJB_BEGIN_ARRAY && !(*it1)->isScalar) && rk2 == WJB_BEGIN_OBJECT) ||
3659                          (rk1 == WJB_BEGIN_OBJECT && (rk2 == WJB_BEGIN_ARRAY && !(*it2)->isScalar)))
3660         {
3661
3662                 JsonbIterator **it_array = rk1 == WJB_BEGIN_ARRAY ? it1 : it2;
3663                 JsonbIterator **it_object = rk1 == WJB_BEGIN_OBJECT ? it1 : it2;
3664
3665                 bool            prepend = (rk1 == WJB_BEGIN_OBJECT) ? true : false;
3666
3667                 pushJsonbValue(state, WJB_BEGIN_ARRAY, NULL);
3668                 if (prepend)
3669                 {
3670                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
3671                         walkJsonb(it_object, state, false);
3672
3673                         res = walkJsonb(it_array, state, false);
3674                 }
3675                 else
3676                 {
3677                         walkJsonb(it_array, state, true);
3678
3679                         pushJsonbValue(state, WJB_BEGIN_OBJECT, NULL);
3680                         walkJsonb(it_object, state, false);
3681
3682                         res = pushJsonbValue(state, WJB_END_ARRAY, NULL);
3683                 }
3684         }
3685         else
3686         {
3687                 /*
3688                  * This must be scalar || object or object || scalar, as that's all
3689                  * that's left. Both of these make no sense, so error out.
3690                  */
3691                 ereport(ERROR,
3692                                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
3693                                  errmsg("invalid concatenation of jsonb objects")));
3694         }
3695
3696         return res;
3697 }
3698
3699 /*
3700  * copy elements from the iterator to the parse state
3701  * stopping at level zero if required.
3702  */
3703 static JsonbValue *
3704 walkJsonb(JsonbIterator **it, JsonbParseState **state, bool stop_at_level_zero)
3705 {
3706         uint32          r,
3707                                 level = 1;
3708         JsonbValue  v;
3709         JsonbValue *res = NULL;
3710
3711         while ((r = JsonbIteratorNext(it, &v, false)) != WJB_DONE)
3712         {
3713                 if (r == WJB_BEGIN_OBJECT || r == WJB_BEGIN_ARRAY)
3714                         ++level;
3715                 else if (r == WJB_END_OBJECT || r == WJB_END_ARRAY)
3716                         --level;
3717
3718                 if (stop_at_level_zero && level == 0)
3719                         break;
3720
3721                 res = pushJsonbValue(state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3722         }
3723
3724         return res;
3725 }
3726
3727
3728 /*
3729  * do most of the heavy work for jsonb_replace
3730  */
3731 static JsonbValue *
3732 replacePath(JsonbIterator **it, Datum *path_elems,
3733                         bool *path_nulls, int path_len,
3734                         JsonbParseState **st, int level, Jsonb *newval)
3735 {
3736         JsonbValue      v;
3737         JsonbValue *res = NULL;
3738         int                     r;
3739
3740         r = JsonbIteratorNext(it, &v, false);
3741
3742         switch (r)
3743         {
3744                 case WJB_BEGIN_ARRAY:
3745                         (void) pushJsonbValue(st, r, NULL);
3746                         replacePathArray(it, path_elems, path_nulls, path_len, st, level,
3747                                                          newval, v.val.array.nElems);
3748                         r = JsonbIteratorNext(it, &v, false);
3749                         Assert(r == WJB_END_ARRAY);
3750                         res = pushJsonbValue(st, r, NULL);
3751
3752                         break;
3753                 case WJB_BEGIN_OBJECT:
3754                         (void) pushJsonbValue(st, r, NULL);
3755                         replacePathObject(it, path_elems, path_nulls, path_len, st, level,
3756                                                           newval, v.val.object.nPairs);
3757                         r = JsonbIteratorNext(it, &v, true);
3758                         Assert(r == WJB_END_OBJECT);
3759                         res = pushJsonbValue(st, r, NULL);
3760
3761                         break;
3762                 case WJB_ELEM:
3763                 case WJB_VALUE:
3764                         res = pushJsonbValue(st, r, &v);
3765                         break;
3766                 default:
3767                         elog(ERROR, "impossible state");
3768         }
3769
3770         return res;
3771 }
3772
3773 /*
3774  * Object walker for replacePath
3775  */
3776 static void
3777 replacePathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
3778                                   int path_len, JsonbParseState **st, int level,
3779                                   Jsonb *newval, uint32 nelems)
3780 {
3781         JsonbValue      v;
3782         int                     i;
3783         JsonbValue      k;
3784         bool            done = false;
3785
3786         if (level >= path_len || path_nulls[level])
3787                 done = true;
3788
3789         for (i = 0; i < nelems; i++)
3790         {
3791                 int             r = JsonbIteratorNext(it, &k, true);
3792                 Assert(r == WJB_KEY);
3793
3794                 if (!done &&
3795                         k.val.string.len == VARSIZE_ANY_EXHDR(path_elems[level]) &&
3796                         memcmp(k.val.string.val, VARDATA_ANY(path_elems[level]),
3797                                    k.val.string.len) == 0)
3798                 {
3799                         if (level == path_len - 1)
3800                         {
3801                                 r = JsonbIteratorNext(it, &v, true);            /* skip */
3802                                 if (newval != NULL)
3803                                 {
3804                                         (void) pushJsonbValue(st, WJB_KEY, &k);
3805                                         addJsonbToParseState(st, newval);
3806                                 }
3807                         }
3808                         else
3809                         {
3810                                 (void) pushJsonbValue(st, r, &k);
3811                                 replacePath(it, path_elems, path_nulls, path_len,
3812                                                         st, level + 1, newval);
3813                         }
3814                 }
3815                 else
3816                 {
3817                         (void) pushJsonbValue(st, r, &k);
3818                         r = JsonbIteratorNext(it, &v, false);
3819                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3820                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3821                         {
3822                                 int             walking_level = 1;
3823
3824                                 while (walking_level != 0)
3825                                 {
3826                                         r = JsonbIteratorNext(it, &v, false);
3827
3828                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3829                                                 ++walking_level;
3830                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
3831                                                 --walking_level;
3832
3833                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3834                                 }
3835                         }
3836                 }
3837         }
3838 }
3839
3840 /*
3841  * Array walker for replacePath
3842  */
3843 static void
3844 replacePathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
3845                                  int path_len, JsonbParseState **st, int level,
3846                                  Jsonb *newval, uint32 npairs)
3847 {
3848         JsonbValue      v;
3849         int                     idx,
3850                                 i;
3851         char       *badp;
3852
3853         /* pick correct index */
3854         if (level < path_len && !path_nulls[level])
3855         {
3856                 char       *c = VARDATA_ANY(path_elems[level]);
3857
3858                 errno = 0;
3859                 idx = (int) strtol(c, &badp, 10);
3860                 if (errno != 0 || badp == c)
3861                         idx = npairs;
3862         }
3863         else
3864                 idx = npairs;
3865
3866         if (idx < 0)
3867         {
3868                 if (-idx > npairs)
3869                         idx = npairs;
3870                 else
3871                         idx = npairs + idx;
3872         }
3873
3874         if (idx > npairs)
3875                 idx = npairs;
3876
3877         /* iterate over the array elements */
3878         for (i = 0; i < npairs; i++)
3879         {
3880                 int             r;
3881
3882                 if (i == idx && level < path_len)
3883                 {
3884                         if (level == path_len - 1)
3885                         {
3886                                 r = JsonbIteratorNext(it, &v, true);            /* skip */
3887                                 if (newval != NULL)
3888                                         addJsonbToParseState(st, newval);
3889                         }
3890                         else
3891                                 (void) replacePath(it, path_elems, path_nulls, path_len,
3892                                                                    st, level + 1, newval);
3893                 }
3894                 else
3895                 {
3896                         r = JsonbIteratorNext(it, &v, false);
3897
3898                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3899
3900                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3901                         {
3902                                 int             walking_level = 1;
3903
3904                                 while (walking_level != 0)
3905                                 {
3906                                         r = JsonbIteratorNext(it, &v, false);
3907
3908                                         if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
3909                                                 ++walking_level;
3910                                         if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
3911                                                 --walking_level;
3912
3913                                         (void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
3914                                 }
3915                         }
3916                 }
3917         }
3918 }