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