optinal pg decimal to tnt decimal cast

Author: parihaaraka

pg/driver.c

@@ -49,6 +49,12 @@
 #undef PACKAGE_VERSION
 #include <module.h>
 
+struct dec_opt {
+	char cast;
+	int dnew_index;
+};
+typedef struct dec_opt dec_opt_t;
+
 /**
  * Infinity timeout from tarantool_ev.c. I mean, this should be in
  * a module.h file.
@@ -97,7 +103,7 @@ lua_push_error(struct lua_State *L)
  * Parse pg values to lua
  */
 static int
-parse_pg_value(struct lua_State *L, char dec_cast, PGresult *res, int row, int col)
+parse_pg_value(struct lua_State *L, dec_opt_t *dopt, PGresult *res, int row, int col)
 {
 	if (PQgetisnull(res, row, col))
 		return false;
@@ -108,9 +114,18 @@ parse_pg_value(struct lua_State *L, char dec_cast, PGresult *res, int row, int c
 
 	switch (PQftype(res, col)) {
 		case NUMERICOID: {
-			if (dec_cast == 's')
-			{
+			if (dopt->cast == 's') {
+				lua_pushlstring(L, val, len);
+				break;
+			}
+			else if (dopt->cast == 'd' && dopt->dnew_index != -1) {
+				lua_rawgeti(L, LUA_REGISTRYINDEX, dopt->dnew_index);
 				lua_pushlstring(L, val, len);
+				int fail = lua_pcall(L, 1, 1, 0);
+				if (fail) {
+					lua_pop(L, 2);
+					return false;
+				}
 				break;
 			}
 			// else fallthrough
@@ -148,15 +163,15 @@ static int
 safe_pg_parsetuples(struct lua_State *L)
 {
 	PGresult *res = (PGresult *)lua_topointer(L, 1);
-	const char dec_cast = (char)lua_tointeger(L, 2);
+	dec_opt_t *dopt = (dec_opt_t *)lua_topointer(L, 2);
 	int row, rows = PQntuples(res);
 	int col, cols = PQnfields(res);
 	lua_newtable(L);
 	for (row = 0; row < rows; ++row) {
 		lua_pushnumber(L, row + 1);
 		lua_newtable(L);
 		for (col = 0; col < cols; ++col)
-			parse_pg_value(L, dec_cast, res, row, col);
+			parse_pg_value(L, dopt, res, row, col);
 		lua_settable(L, -3);
 	}
 	return 1;
@@ -213,7 +228,7 @@ pg_wait_for_result(PGconn *conn)
  * Appends result fom postgres to lua table
  */
 static int
-pg_resultget(struct lua_State *L, const char dec_cast, PGconn *conn, int *res_no, int status_ok)
+pg_resultget(struct lua_State *L, dec_opt_t *dopt, PGconn *conn, int *res_no, int status_ok)
 {
 	int wait_res = pg_wait_for_result(conn);
 	if (wait_res != 1)
@@ -243,7 +258,7 @@ pg_resultget(struct lua_State *L, const char dec_cast, PGconn *conn, int *res_no
 			lua_pushinteger(L, (*res_no)++);
 			lua_pushcfunction(L, safe_pg_parsetuples);
 			lua_pushlightuserdata(L, pg_res);
-			lua_pushinteger(L, dec_cast);
+			lua_pushlightuserdata(L, dopt);
 			fail = lua_pcall(L, 2, 1, 0);
 			if (!fail)
 				lua_settable(L, -3);
@@ -338,19 +353,25 @@ lua_pg_execute(struct lua_State *L)
 {
 	PGconn *conn = lua_check_pgconn(L, 1);
 
-	char dec_cast = 'n';
+	dec_opt_t dopt = {'n', -1};
 	if (lua_isstring(L, 2)) {
 		const char *tmp = lua_tostring(L, 2);
-		if (*tmp == 'n' || *tmp == 's') // TODO 'd' - decimal
-			dec_cast = *tmp;
+		if (*tmp == 'n' || *tmp == 's' || *tmp == 'd')
+			dopt.cast = *tmp;
 	}
 
-	if (!lua_isstring(L, 3)) {
+	if (!lua_isstring(L, 4)) {
 		safe_pushstring(L, "Second param should be a sql command");
 		return lua_push_error(L);
 	}
-	const char *sql = lua_tostring(L, 3);
-	int paramCount = lua_gettop(L) - 3;
+
+	if (lua_isfunction(L, 3)) {
+		lua_pushvalue(L, 3);
+		dopt.dnew_index = luaL_ref(L, LUA_REGISTRYINDEX);
+	}
+
+	const char *sql = lua_tostring(L, 4);
+	int paramCount = lua_gettop(L) - 4;
 
 	const char **paramValues = NULL;
 	int  *paramLengths = NULL;
@@ -371,7 +392,7 @@ lua_pg_execute(struct lua_State *L)
 
 		int idx;
 		for (idx = 0; idx < paramCount; ++idx) {
-			lua_parse_param(L, idx + 4, paramValues + idx,
+			lua_parse_param(L, idx + 5, paramValues + idx,
 				paramLengths + idx, paramTypes + idx);
 		}
 		res = PQsendQueryParams(conn, sql, paramCount, paramTypes,
@@ -383,14 +404,19 @@ lua_pg_execute(struct lua_State *L)
 	if (res == -1) {
 		lua_pushinteger(L, PQstatus(conn) == CONNECTION_BAD ? -1: 0);
 		lua_pushstring(L, PQerrorMessage(conn));
+		if (dopt.dnew_index != -1)
+			luaL_unref(L, LUA_REGISTRYINDEX, dopt.dnew_index);
 		return 2;
 	}
 	lua_pushinteger(L, 0);
 	lua_newtable(L);
 
 	int res_no = 1;
 	int status_ok = 1;
-	while ((status_ok = pg_resultget(L, dec_cast, conn, &res_no, status_ok)));
+	while ((status_ok = pg_resultget(L, &dopt, conn, &res_no, status_ok)));
+
+	if (dopt.dnew_index != -1)
+		luaL_unref(L, LUA_REGISTRYINDEX, dopt.dnew_index);
 
 	return 2;
 }

pg/init.lua

@@ -4,6 +4,14 @@ local fiber = require('fiber')
 local driver = require('pg.driver')
 local ffi = require('ffi')
 
+local dnew
+do
+    local ok, dec = pcall(require, "decimal")
+    if ok then
+        dnew = dec.new
+    end
+end
+
 local pool_mt
 local conn_mt
 
@@ -61,7 +69,7 @@ conn_mt = {
                 self.queue:put(false)
                 return get_error(self.raise.pool, 'Connection is broken')
             end
-            local status, datas = self.conn:execute(self.dec_cast, sql, ...)
+            local status, datas = self.conn:execute(self.dec_cast, dnew, sql, ...)
             if status ~= 0 then
                 self.queue:put(status > 0)
                 return error(datas)