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

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