contrib/jsonb_plpython/jsonb_plpython.c

   1 #include "postgres.h"
   2
   3 #include "plpy_elog.h"
   4 #include "plpy_typeio.h"
   5 #include "plpython.h"
   6 #include "utils/fmgrprotos.h"
   7 #include "utils/jsonb.h"
   8 #include "utils/numeric.h"
   9
  10 PG_MODULE_MAGIC;
  11
  12 /* for PLyObject_AsString in plpy_typeio.c */
  13 typedef char *(*PLyObject_AsString_t) (PyObject *plrv);
  14 static PLyObject_AsString_t PLyObject_AsString_p;
  15
  16 typedef void (*PLy_elog_impl_t) (int elevel, const char *fmt,...);
  17 static PLy_elog_impl_t PLy_elog_impl_p;
  18
  19 /*
  20  * decimal_constructor is a function from python library and used
  21  * for transforming strings into python decimal type
  22  */
  23 static PyObject *decimal_constructor;
  24
  25 static PyObject *PLyObject_FromJsonbContainer(JsonbContainer *jsonb);
  26 static JsonbValue *PLyObject_ToJsonbValue(PyObject *obj,
  27                                                                                   JsonbParseState **jsonb_state, bool is_elem);
  28
  29 typedef PyObject *(*PLyUnicode_FromStringAndSize_t)
  30                         (const char *s, Py_ssize_t size);
  31 static PLyUnicode_FromStringAndSize_t PLyUnicode_FromStringAndSize_p;
  32
  33 /*
  34  * Module initialize function: fetch function pointers for cross-module calls.
  35  */
  36 void
  37 _PG_init(void)
  38 {
  39         /* Asserts verify that typedefs above match original declarations */
  40         AssertVariableIsOfType(&PLyObject_AsString, PLyObject_AsString_t);
  41         PLyObject_AsString_p = (PLyObject_AsString_t)
  42                 load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLyObject_AsString",
  43                                                            true, NULL);
  44         AssertVariableIsOfType(&PLyUnicode_FromStringAndSize, PLyUnicode_FromStringAndSize_t);
  45         PLyUnicode_FromStringAndSize_p = (PLyUnicode_FromStringAndSize_t)
  46                 load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLyUnicode_FromStringAndSize",
  47                                                            true, NULL);
  48         AssertVariableIsOfType(&PLy_elog_impl, PLy_elog_impl_t);
  49         PLy_elog_impl_p = (PLy_elog_impl_t)
  50                 load_external_function("$libdir/" PLPYTHON_LIBNAME, "PLy_elog_impl",
  51                                                            true, NULL);
  52 }
  53
  54 /* These defines must be after the _PG_init */
  55 #define PLyObject_AsString (PLyObject_AsString_p)
  56 #define PLyUnicode_FromStringAndSize (PLyUnicode_FromStringAndSize_p)
  57 #undef PLy_elog
  58 #define PLy_elog (PLy_elog_impl_p)
  59
  60 /*
  61  * PLyUnicode_FromJsonbValue
  62  *
  63  * Transform string JsonbValue to Python string.
  64  */
  65 static PyObject *
  66 PLyUnicode_FromJsonbValue(JsonbValue *jbv)
  67 {
  68         Assert(jbv->type == jbvString);
  69
  70         return PLyUnicode_FromStringAndSize(jbv->val.string.val, jbv->val.string.len);
  71 }
  72
  73 /*
  74  * PLyUnicode_ToJsonbValue
  75  *
  76  * Transform Python string to JsonbValue.
  77  */
  78 static void
  79 PLyUnicode_ToJsonbValue(PyObject *obj, JsonbValue *jbvElem)
  80 {
  81         jbvElem->type = jbvString;
  82         jbvElem->val.string.val = PLyObject_AsString(obj);
  83         jbvElem->val.string.len = strlen(jbvElem->val.string.val);
  84 }
  85
  86 /*
  87  * PLyObject_FromJsonbValue
  88  *
  89  * Transform JsonbValue to PyObject.
  90  */
  91 static PyObject *
  92 PLyObject_FromJsonbValue(JsonbValue *jsonbValue)
  93 {
  94         switch (jsonbValue->type)
  95         {
  96                 case jbvNull:
  97                         Py_RETURN_NONE;
  98
  99                 case jbvBinary:
 100                         return PLyObject_FromJsonbContainer(jsonbValue->val.binary.data);
 101
 102                 case jbvNumeric:
 103                         {
 104                                 Datum           num;
 105                                 char       *str;
 106
 107                                 num = NumericGetDatum(jsonbValue->val.numeric);
 108                                 str = DatumGetCString(DirectFunctionCall1(numeric_out, num));
 109
 110                                 return PyObject_CallFunction(decimal_constructor, "s", str);
 111                         }
 112
 113                 case jbvString:
 114                         return PLyUnicode_FromJsonbValue(jsonbValue);
 115
 116                 case jbvBool:
 117                         if (jsonbValue->val.boolean)
 118                                 Py_RETURN_TRUE;
 119                         else
 120                                 Py_RETURN_FALSE;
 121
 122                 default:
 123                         elog(ERROR, "unexpected jsonb value type: %d", jsonbValue->type);
 124                         return NULL;
 125         }
 126 }
 127
 128 /*
 129  * PLyObject_FromJsonbContainer
 130  *
 131  * Transform JsonbContainer to PyObject.
 132  */
 133 static PyObject *
 134 PLyObject_FromJsonbContainer(JsonbContainer *jsonb)
 135 {
 136         JsonbIteratorToken r;
 137         JsonbValue      v;
 138         JsonbIterator *it;
 139         PyObject   *result;
 140
 141         it = JsonbIteratorInit(jsonb);
 142         r = JsonbIteratorNext(&it, &v, true);
 143
 144         switch (r)
 145         {
 146                 case WJB_BEGIN_ARRAY:
 147                         if (v.val.array.rawScalar)
 148                         {
 149                                 JsonbValue      tmp;
 150
 151                                 if ((r = JsonbIteratorNext(&it, &v, true)) != WJB_ELEM ||
 152                                         (r = JsonbIteratorNext(&it, &tmp, true)) != WJB_END_ARRAY ||
 153                                         (r = JsonbIteratorNext(&it, &tmp, true)) != WJB_DONE)
 154                                         elog(ERROR, "unexpected jsonb token: %d", r);
 155
 156                                 result = PLyObject_FromJsonbValue(&v);
 157                         }
 158                         else
 159                         {
 160                                 PyObject   *volatile elem = NULL;
 161
 162                                 result = PyList_New(0);
 163                                 if (!result)
 164                                         return NULL;
 165
 166                                 PG_TRY();
 167                                 {
 168                                         while ((r = JsonbIteratorNext(&it, &v, true)) != WJB_DONE)
 169                                         {
 170                                                 if (r != WJB_ELEM)
 171                                                         continue;
 172
 173                                                 elem = PLyObject_FromJsonbValue(&v);
 174
 175                                                 PyList_Append(result, elem);
 176                                                 Py_XDECREF(elem);
 177                                                 elem = NULL;
 178                                         }
 179                                 }
 180                                 PG_CATCH();
 181                                 {
 182                                         Py_XDECREF(elem);
 183                                         Py_XDECREF(result);
 184                                         PG_RE_THROW();
 185                                 }
 186                                 PG_END_TRY();
 187                         }
 188                         break;
 189
 190                 case WJB_BEGIN_OBJECT:
 191                         {
 192                                 PyObject   *volatile result_v = PyDict_New();
 193                                 PyObject   *volatile key = NULL;
 194                                 PyObject   *volatile val = NULL;
 195
 196                                 if (!result_v)
 197                                         return NULL;
 198
 199                                 PG_TRY();
 200                                 {
 201                                         while ((r = JsonbIteratorNext(&it, &v, true)) != WJB_DONE)
 202                                         {
 203                                                 if (r != WJB_KEY)
 204                                                         continue;
 205
 206                                                 key = PLyUnicode_FromJsonbValue(&v);
 207                                                 if (!key)
 208                                                 {
 209                                                         Py_XDECREF(result_v);
 210                                                         result_v = NULL;
 211                                                         break;
 212                                                 }
 213
 214                                                 if ((r = JsonbIteratorNext(&it, &v, true)) != WJB_VALUE)
 215                                                         elog(ERROR, "unexpected jsonb token: %d", r);
 216
 217                                                 val = PLyObject_FromJsonbValue(&v);
 218                                                 if (!val)
 219                                                 {
 220                                                         Py_XDECREF(key);
 221                                                         key = NULL;
 222                                                         Py_XDECREF(result_v);
 223                                                         result_v = NULL;
 224                                                         break;
 225                                                 }
 226
 227                                                 PyDict_SetItem(result_v, key, val);
 228
 229                                                 Py_XDECREF(key);
 230                                                 key = NULL;
 231                                                 Py_XDECREF(val);
 232                                                 val = NULL;
 233                                         }
 234                                 }
 235                                 PG_CATCH();
 236                                 {
 237                                         Py_XDECREF(result_v);
 238                                         Py_XDECREF(key);
 239                                         Py_XDECREF(val);
 240                                         PG_RE_THROW();
 241                                 }
 242                                 PG_END_TRY();
 243
 244                                 result = result_v;
 245                         }
 246                         break;
 247
 248                 default:
 249                         elog(ERROR, "unexpected jsonb token: %d", r);
 250                         return NULL;
 251         }
 252
 253         return result;
 254 }
 255
 256 /*
 257  * PLyMapping_ToJsonbValue
 258  *
 259  * Transform Python dict to JsonbValue.
 260  */
 261 static JsonbValue *
 262 PLyMapping_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
 263 {
 264         Py_ssize_t      pcount;
 265         PyObject   *volatile items;
 266         JsonbValue *volatile out;
 267
 268         pcount = PyMapping_Size(obj);
 269         items = PyMapping_Items(obj);
 270
 271         PG_TRY();
 272         {
 273                 Py_ssize_t      i;
 274
 275                 pushJsonbValue(jsonb_state, WJB_BEGIN_OBJECT, NULL);
 276
 277                 for (i = 0; i < pcount; i++)
 278                 {
 279                         JsonbValue      jbvKey;
 280                         PyObject   *item = PyList_GetItem(items, i);
 281                         PyObject   *key = PyTuple_GetItem(item, 0);
 282                         PyObject   *value = PyTuple_GetItem(item, 1);
 283
 284                         /* Python dictionary can have None as key */
 285                         if (key == Py_None)
 286                         {
 287                                 jbvKey.type = jbvString;
 288                                 jbvKey.val.string.len = 0;
 289                                 jbvKey.val.string.val = "";
 290                         }
 291                         else
 292                         {
 293                                 /* All others types of keys we serialize to string */
 294                                 PLyUnicode_ToJsonbValue(key, &jbvKey);
 295                         }
 296
 297                         (void) pushJsonbValue(jsonb_state, WJB_KEY, &jbvKey);
 298                         (void) PLyObject_ToJsonbValue(value, jsonb_state, false);
 299                 }
 300
 301                 out = pushJsonbValue(jsonb_state, WJB_END_OBJECT, NULL);
 302         }
 303         PG_FINALLY();
 304         {
 305                 Py_DECREF(items);
 306         }
 307         PG_END_TRY();
 308
 309         return out;
 310 }
 311
 312 /*
 313  * PLySequence_ToJsonbValue
 314  *
 315  * Transform python list to JsonbValue. Expects transformed PyObject and
 316  * a state required for jsonb construction.
 317  */
 318 static JsonbValue *
 319 PLySequence_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
 320 {
 321         Py_ssize_t      i;
 322         Py_ssize_t      pcount;
 323         PyObject   *volatile value = NULL;
 324
 325         pcount = PySequence_Size(obj);
 326
 327         pushJsonbValue(jsonb_state, WJB_BEGIN_ARRAY, NULL);
 328
 329         PG_TRY();
 330         {
 331                 for (i = 0; i < pcount; i++)
 332                 {
 333                         value = PySequence_GetItem(obj, i);
 334                         Assert(value);
 335
 336                         (void) PLyObject_ToJsonbValue(value, jsonb_state, true);
 337                         Py_XDECREF(value);
 338                         value = NULL;
 339                 }
 340         }
 341         PG_CATCH();
 342         {
 343                 Py_XDECREF(value);
 344                 PG_RE_THROW();
 345         }
 346         PG_END_TRY();
 347
 348         return pushJsonbValue(jsonb_state, WJB_END_ARRAY, NULL);
 349 }
 350
 351 /*
 352  * PLyNumber_ToJsonbValue(PyObject *obj)
 353  *
 354  * Transform python number to JsonbValue.
 355  */
 356 static JsonbValue *
 357 PLyNumber_ToJsonbValue(PyObject *obj, JsonbValue *jbvNum)
 358 {
 359         Numeric         num;
 360         char       *str = PLyObject_AsString(obj);
 361
 362         PG_TRY();
 363         {
 364                 Datum           numd;
 365
 366                 numd = DirectFunctionCall3(numeric_in,
 367                                                                    CStringGetDatum(str),
 368                                                                    ObjectIdGetDatum(InvalidOid),
 369                                                                    Int32GetDatum(-1));
 370                 num = DatumGetNumeric(numd);
 371         }
 372         PG_CATCH();
 373         {
 374                 ereport(ERROR,
 375                                 (errcode(ERRCODE_DATATYPE_MISMATCH),
 376                                  errmsg("could not convert value \"%s\" to jsonb", str)));
 377         }
 378         PG_END_TRY();
 379
 380         pfree(str);
 381
 382         /*
 383          * jsonb doesn't allow NaN or infinity (per JSON specification), so we
 384          * have to reject those here explicitly.
 385          */
 386         if (numeric_is_nan(num))
 387                 ereport(ERROR,
 388                                 (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
 389                                  errmsg("cannot convert NaN to jsonb")));
 390         if (numeric_is_inf(num))
 391                 ereport(ERROR,
 392                                 (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
 393                                  errmsg("cannot convert infinity to jsonb")));
 394
 395         jbvNum->type = jbvNumeric;
 396         jbvNum->val.numeric = num;
 397
 398         return jbvNum;
 399 }
 400
 401 /*
 402  * PLyObject_ToJsonbValue(PyObject *obj)
 403  *
 404  * Transform python object to JsonbValue.
 405  */
 406 static JsonbValue *
 407 PLyObject_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state, bool is_elem)
 408 {
 409         JsonbValue *out;
 410
 411         if (!PyUnicode_Check(obj))
 412         {
 413                 if (PySequence_Check(obj))
 414                         return PLySequence_ToJsonbValue(obj, jsonb_state);
 415                 else if (PyMapping_Check(obj))
 416                         return PLyMapping_ToJsonbValue(obj, jsonb_state);
 417         }
 418
 419         out = palloc(sizeof(JsonbValue));
 420
 421         if (obj == Py_None)
 422                 out->type = jbvNull;
 423         else if (PyUnicode_Check(obj))
 424                 PLyUnicode_ToJsonbValue(obj, out);
 425
 426         /*
 427          * PyNumber_Check() returns true for booleans, so boolean check should
 428          * come first.
 429          */
 430         else if (PyBool_Check(obj))
 431         {
 432                 out->type = jbvBool;
 433                 out->val.boolean = (obj == Py_True);
 434         }
 435         else if (PyNumber_Check(obj))
 436                 out = PLyNumber_ToJsonbValue(obj, out);
 437         else
 438                 ereport(ERROR,
 439                                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 440                                  errmsg("Python type \"%s\" cannot be transformed to jsonb",
 441                                                 PLyObject_AsString((PyObject *) obj->ob_type))));
 442
 443         /* Push result into 'jsonb_state' unless it is raw scalar value. */
 444         return (*jsonb_state ?
 445                         pushJsonbValue(jsonb_state, is_elem ? WJB_ELEM : WJB_VALUE, out) :
 446                         out);
 447 }
 448
 449 /*
 450  * plpython_to_jsonb
 451  *
 452  * Transform python object to Jsonb datum
 453  */
 454 PG_FUNCTION_INFO_V1(plpython_to_jsonb);
 455 Datum
 456 plpython_to_jsonb(PG_FUNCTION_ARGS)
 457 {
 458         PyObject   *obj;
 459         JsonbValue *out;
 460         JsonbParseState *jsonb_state = NULL;
 461
 462         obj = (PyObject *) PG_GETARG_POINTER(0);
 463         out = PLyObject_ToJsonbValue(obj, &jsonb_state, true);
 464         PG_RETURN_POINTER(JsonbValueToJsonb(out));
 465 }
 466
 467 /*
 468  * jsonb_to_plpython
 469  *
 470  * Transform Jsonb datum to PyObject and return it as internal.
 471  */
 472 PG_FUNCTION_INFO_V1(jsonb_to_plpython);
 473 Datum
 474 jsonb_to_plpython(PG_FUNCTION_ARGS)
 475 {
 476         PyObject   *result;
 477         Jsonb      *in = PG_GETARG_JSONB_P(0);
 478
 479         /*
 480          * Initialize pointer to Decimal constructor. First we try "cdecimal", C
 481          * version of decimal library. In case of failure we use slower "decimal"
 482          * module.
 483          */
 484         if (!decimal_constructor)
 485         {
 486                 PyObject   *decimal_module = PyImport_ImportModule("cdecimal");
 487
 488                 if (!decimal_module)
 489                 {
 490                         PyErr_Clear();
 491                         decimal_module = PyImport_ImportModule("decimal");
 492                 }
 493                 Assert(decimal_module);
 494                 decimal_constructor = PyObject_GetAttrString(decimal_module, "Decimal");
 495         }
 496
 497         result = PLyObject_FromJsonbContainer(&in->root);
 498         if (!result)
 499                 PLy_elog(ERROR, "transformation from jsonb to Python failed");
 500
 501         return PointerGetDatum(result);
 502 }