dbd/postgresql/statement.c

#include "dbd_postgresql.h"

#define MAX_PLACEHOLDERS	9999 
#define MAX_PLACEHOLDER_SIZE	(1+4) /* $\d{4} */

static lua_push_type_t postgresql_to_lua_push(unsigned int postgresql_type) {
    lua_push_type_t lua_type;

    switch(postgresql_type) {
    case INT2OID:
    case INT4OID:
        lua_type =  LUA_PUSH_INTEGER;
        break;

    case FLOAT4OID:
    case FLOAT8OID:
        lua_type = LUA_PUSH_NUMBER;
        break;

    case BOOLOID:
	lua_type = LUA_PUSH_BOOLEAN;
	break;

    default:
        lua_type = LUA_PUSH_STRING;
    }

    return lua_type;
}

/*
 * replace '?' placeholders with $\d+ placeholders
 * to be compatible with PSQL API
 */
static char *replace_placeholders(lua_State *L, const char *sql) {
    size_t len = strlen(sql);
    int num_placeholders = 0;
    int extra_space = 0;
    int i;
    char *newsql;
    int newpos = 1;
    int ph_num = 1;
    int in_quote = 0;

    /*
     * dumb count of all '?'
     * this will match more placeholders than necessesary
     * but it's safer to allocate more placeholders at the
     * cost of a few bytes than risk a buffer overflow
     */ 
    for (i = 1; i < len; i++) {
	if (sql[i] == '?') {
	    num_placeholders++;
	}
    }
    
    /*
     * this is MAX_PLACEHOLDER_SIZE-1 because the '?' is 
     * replaced with '$'
     */ 
    extra_space = num_placeholders * (MAX_PLACEHOLDER_SIZE-1); 

    /*
     * allocate a new string for the converted SQL statement
     */
    newsql = malloc(sizeof(char) * (len+extra_space+1));
    memset(newsql, 0, sizeof(char) * (len+extra_space+1));
    
    /* 
     * copy first char. In valid SQL this cannot be a placeholder
     */
    newsql[0] = sql[0];

    /* 
     * only replace '?' not in a single quoted string
     */
    for (i = 1; i < len; i++) {
	/*
	 * don't change the quote flag if the ''' is preceded 
	 * bt a '\' to account for escaping
	 */
	if (sql[i] == '\'' && sql[i-1] != '\\') {
	    in_quote = !in_quote;
	}

	if (sql[i] == '?' && !in_quote) {
	    size_t n;

	    if (ph_num > MAX_PLACEHOLDERS) {
		luaL_error(L, "Sorry, you are using more than %d placeholders. Use ${num} format instead", MAX_PLACEHOLDERS);
	    }

	    n = snprintf(&newsql[newpos], MAX_PLACEHOLDER_SIZE, "$%u", ph_num++);

	    newpos += n;
	} else {
	    newsql[newpos] = sql[i];
	    newpos++;
	}
    }

    /* 
     * terminate string on the last position 
     */
    newsql[newpos] = '\0';

    /* fprintf(stderr, "[%s]\n", newsql); */
    return newsql;
}

/*
 * success = statement:close()
 */
static int statement_close(lua_State *L) {
    statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT);

    if (statement->result) {
	PQclear(statement->result);
	statement->result = NULL;
    }

    return 0;    
}

/*
 * success = statement:execute(...)
 */
static int statement_execute(lua_State *L) {
    int n = lua_gettop(L);
    statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT);
    int num_bind_params = n - 1;   
    ExecStatusType status;
    int p;
    const char *errstr = NULL;

    const char **params;
    PGresult *result = NULL;

    statement->tuple = 0;

    params = malloc(num_bind_params * sizeof(params));
    memset(params, 0, num_bind_params * sizeof(params));

    /*
     * convert and copy parameters into a string array
     */ 
    for (p = 2; p <= n; p++) {
	int i = p - 2;	
	int type = lua_type(L, p);
	char err[64];

	switch(type) {
	case LUA_TNIL:
	    params[i] = NULL;
	    break;
	case LUA_TBOOLEAN:
	    /*
	     * boolean values in postgresql can either be
	     * t/f or 1/0. Pass integer values rather than
	     * strings to maintain semantic compatibility
	     * with other DBD drivers that pass booleans
	     * as integers.
	     */
	    params[i] = lua_toboolean(L, p) ?  "1" : "0";
	    break;
	case LUA_TNUMBER:
	case LUA_TSTRING:
	    params[i] = lua_tostring(L, p);
	    break;
	default:
	    snprintf(err, sizeof(err)-1, DBI_ERR_BINDING_TYPE_ERR, lua_typename(L, type));
	    errstr = err;
	    goto cleanup;
	}
    }

    result = PQexecPrepared(
	statement->postgresql,
        statement->name,
        num_bind_params,
        (const char **)params,
	NULL,
        NULL,
        0
    );

cleanup:
    free(params);

    if (errstr) {
	lua_pushboolean(L, 0);
	lua_pushfstring(L, DBI_ERR_BINDING_PARAMS, errstr);
	return 2;
    }

    if (!result) {
	lua_pushboolean(L, 0);
	lua_pushfstring(L, DBI_ERR_ALLOC_RESULT,  PQerrorMessage(statement->postgresql));
	return 2;
    }
    
    status = PQresultStatus(result);
    if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK) {
	lua_pushboolean(L, 0);
	lua_pushfstring(L, DBI_ERR_BINDING_EXEC, PQresultErrorMessage(result));
	return 2;
    }
    
    statement->result = result;

    lua_pushboolean(L, 1);
    return 1;
}

/*
 * must be called after an execute
 */
static int statement_fetch_impl(lua_State *L, statement_t *statement, int named_columns) {
    int tuple = statement->tuple++;
    int i;
    int num_columns;
    int d = 1;

    if (!statement->result) {
	luaL_error(L, DBI_ERR_FETCH_INVALID);
	return 0;
    }

    if (PQresultStatus(statement->result) != PGRES_TUPLES_OK) {
	lua_pushnil(L);
	return 1;
    }

    if (tuple >= PQntuples(statement->result)) {
	lua_pushnil(L);  /* no more results */
	return 1;
    }

    num_columns = PQnfields(statement->result);
    lua_newtable(L);
    for (i = 0; i < num_columns; i++) {
	const char *name = PQfname(statement->result, i);

	if (PQgetisnull(statement->result, tuple, i)) {
	    if (named_columns) {
		LUA_PUSH_ATTRIB_NIL(name);
            } else {
                LUA_PUSH_ARRAY_NIL(d);
            }	    
	} else {
	    const char *value = PQgetvalue(statement->result, tuple, i);
	    lua_push_type_t lua_push = postgresql_to_lua_push(PQftype(statement->result, i));

	    /*
             * data is returned as strings from PSQL
             * convert them here into Lua types
             */ 

            if (lua_push == LUA_PUSH_NIL) {
                if (named_columns) {
                    LUA_PUSH_ATTRIB_NIL(name);
                } else {
                    LUA_PUSH_ARRAY_NIL(d);
                }
            } else if (lua_push == LUA_PUSH_INTEGER) {
                int val = atoi(value);

                if (named_columns) {
                    LUA_PUSH_ATTRIB_INT(name, val);
                } else {
                    LUA_PUSH_ARRAY_INT(d, val);
                }
            } else if (lua_push == LUA_PUSH_NUMBER) {
                double val = strtod(value, NULL);

                if (named_columns) {
                    LUA_PUSH_ATTRIB_FLOAT(name, val);
                } else {
                    LUA_PUSH_ARRAY_FLOAT(d, val);
                }
            } else if (lua_push == LUA_PUSH_STRING) {
                if (named_columns) {
                    LUA_PUSH_ATTRIB_STRING(name, value);
                } else {
                    LUA_PUSH_ARRAY_STRING(d, value);
                }
            } else if (lua_push == LUA_PUSH_BOOLEAN) {
		/* 
                 * booleans are returned as a string
                 * either 't' or 'f'
                 */
                int val = value[0] == 't' ? 1 : 0;

                if (named_columns) {
                    LUA_PUSH_ATTRIB_BOOL(name, val);
                } else {
                    LUA_PUSH_ARRAY_BOOL(d, val);
                }
            } else {
                luaL_error(L, DBI_ERR_UNKNOWN_PUSH);
            }
	}
    }

    return 1;    
}


static int next_iterator(lua_State *L) {
    statement_t *statement = (statement_t *)luaL_checkudata(L, lua_upvalueindex(1), DBD_POSTGRESQL_STATEMENT);
    int named_columns = lua_toboolean(L, lua_upvalueindex(2));

    return statement_fetch_impl(L, statement, named_columns);
}

/*
 * table = statement:fetch(named_indexes)
 */
static int statement_fetch(lua_State *L) {
    statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT);
    int named_columns = lua_toboolean(L, 2);

    return statement_fetch_impl(L, statement, named_columns);
}

/*
 * iterfunc = statement:rows(named_indexes)
 */
static int statement_rows(lua_State *L) {
    if (lua_gettop(L) == 1) {
        lua_pushvalue(L, 1);
        lua_pushboolean(L, 0);
    } else {
        lua_pushvalue(L, 1);
        lua_pushboolean(L, lua_toboolean(L, 2));
    }

    lua_pushcclosure(L, next_iterator, 2);
    return 1;
}

/*
 * __gc
 */
static int statement_gc(lua_State *L) {
    /* always free the handle */
    statement_close(L);

    return 0;
}

int dbd_postgresql_statement_create(lua_State *L, connection_t *conn, const char *sql_query) { 
    statement_t *statement = NULL;
    ExecStatusType status;
    PGresult *result = NULL;
    char *new_sql;
    char name[IDLEN];

    /*
     * convert SQL string into a PSQL API compatible SQL statement
     */ 
    new_sql = replace_placeholders(L, sql_query);

    snprintf(name, IDLEN, "%017u", ++conn->statement_id);

    result = PQprepare(conn->postgresql, name, new_sql, 0, NULL);

    /*
     * free converted statement after use
     */
    free(new_sql);

    if (!result) {
	lua_pushnil(L);
	lua_pushfstring(L, DBI_ERR_ALLOC_STATEMENT, PQerrorMessage(statement->postgresql));
	return 2;
    }
    
    status = PQresultStatus(result);
    if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK) {
	const char *err_string = PQresultErrorMessage(result);
	PQclear(result);

	lua_pushnil(L);
	lua_pushfstring(L, DBI_ERR_PREP_STATEMENT, err_string);
	return 2;
    }

    PQclear(result);

    statement = (statement_t *)lua_newuserdata(L, sizeof(statement_t));
    statement->postgresql = conn->postgresql;
    statement->result = NULL;
    statement->tuple = 0;
    strncpy(statement->name, name, IDLEN-1);
    statement->name[IDLEN-1] = '\0';

    luaL_getmetatable(L, DBD_POSTGRESQL_STATEMENT);
    lua_setmetatable(L, -2);

    return 1;
} 

int dbd_postgresql_statement(lua_State *L) {
    static const luaL_Reg statement_methods[] = {
	{"close", statement_close},
	{"execute", statement_execute},
	{"fetch", statement_fetch},
	{"rows", statement_rows},
	{NULL, NULL}
    };

    static const luaL_Reg statement_class_methods[] = {
	{NULL, NULL}
    };

    luaL_newmetatable(L, DBD_POSTGRESQL_STATEMENT);
    luaL_register(L, 0, statement_methods);
    lua_pushvalue(L,-1);
    lua_setfield(L, -2, "__index");

    lua_pushcfunction(L, statement_gc);
    lua_setfield(L, -2, "__gc");

    luaL_register(L, DBD_POSTGRESQL_STATEMENT, statement_class_methods);

    return 1;    
}