LCOV - code coverage report
Current view: top level - contrib/jsonb_plpython - jsonb_plpython.c (source / functions) Hit Total Coverage
Test: PostgreSQL 12beta1 Lines: 150 183 82.0 %
Date: 2019-06-15 23:07:03 Functions: 14 14 100.0 %
Legend: Lines: hit not hit

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

Generated by: LCOV version 1.13