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