<commit_before>#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;

    char **params;
    PGresult *result = NULL;

    statement->tuple = 0;

    params = malloc(num_bind_params * sizeof(params));

    /*
     * convert and copy parameters into a string array
     */ 
    for (p = 2; p <= n; p++) {
	int i = p - 2;	

	if (lua_isnil(L, p)) {
	    params[i] = NULL;
	} else {
	    const char *param = lua_tostring(L, p);
	    size_t len = strlen(param) + 1;

	    params[i] = malloc(len * sizeof(char));
	    memset(params[i], 0, len);

	    strncpy(params[i], param, len);
	    params[i][len] = '\0';
	}
    }

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

    /*
     * free string array
     */
    for (p = 0; p < num_bind_params; p++) {
	if (params[p]) {
	    free(params[p]);
	}
    }
    free(params);

    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);
    lua_pushnil(L);
    return 1;
}

/*
 * must be called after an execute
 */
static int statement_fetch_impl(lua_State *L, int named_columns) {
    statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT);
    int tuple = statement->tuple++;<commit_msg>Bugfix - execute was returning nil on success<commit_after>#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;

    char **params;
    PGresult *result = NULL;

    statement->tuple = 0;

    params = malloc(num_bind_params * sizeof(params));

    /*
     * convert and copy parameters into a string array
     */ 
    for (p = 2; p <= n; p++) {
	int i = p - 2;	

	if (lua_isnil(L, p)) {
	    params[i] = NULL;
	} else {
	    const char *param = lua_tostring(L, p);
	    size_t len = strlen(param) + 1;

	    params[i] = malloc(len * sizeof(char));
	    memset(params[i], 0, len);

	    strncpy(params[i], param, len);
	    params[i][len] = '\0';
	}
    }

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

    /*
     * free string array
     */
    for (p = 0; p < num_bind_params; p++) {
	if (params[p]) {
	    free(params[p]);
	}
    }
    free(params);

    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, int named_columns) {
    statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT);
    int tuple = statement->tuple++;
{'fraction_non_alphanumeric': 0.1338194295517907, 'fraction_numerical': 0.009864893845164058, 'mean_word_length': 2.609520123839009, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4008063', 'n_tokens_mistral': 3590, 'n_tokens_neox': 3371, 'n_words': 828}