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 : }
|