LCOV - code coverage report
Current view: top level - contrib/jsonb_plpython - jsonb_plpython.c (source / functions) Hit Total Coverage
Test: PostgreSQL 18devel Lines: 159 193 82.4 %
Date: 2025-04-01 14:15:22 Functions: 14 14 100.0 %
Legend: Lines: hit not hit

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

Generated by: LCOV version 1.14