#include <assert.h>

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <time.h>

#if defined(_WIN32)

#include <windows.h>

#include <sqlext.h>

#elif defined(INFORMIX)

#include "infxcli.h"

#elif defined(UNIXODBC)

#include "sql.h"

#include "sqltypes.h"

#include "sqlext.h"

#include "lua.h"

#include "lauxlib.h"

#if ! defined (LUA_VERSION_NUM) || LUA_VERSION_NUM < 501

#include "compat-5.1.h"

#include "luasql.h"

#define LUASQL_ENVIRONMENT_ODBC "ODBC environment"

#define LUASQL_CONNECTION_ODBC "ODBC connection"

#define LUASQL_CURSOR_ODBC "ODBC cursor"

typedef struct {

short closed;

int conn_counter;

SQLHENV henv; /* environment handle */

} env_data;

typedef struct {

short closed;

int cur_counter;

int env; /* reference to environment */

SQLHDBC hdbc; /* database connection handle */

} conn_data;

typedef struct {

short closed;

int conn; /* reference to connection */

int numcols; /* number of columns */

int coltypes, colnames; /* reference to column information tables */

SQLHSTMT hstmt; /* statement handle */

} cur_data;

#define hENV SQL_HANDLE_ENV

#define hSTMT SQL_HANDLE_STMT

#define hDBC SQL_HANDLE_DBC

#define error(a) ((a) != SQL_SUCCESS && (a) != SQL_SUCCESS_WITH_INFO)

LUASQL_API int luaopen_luasql_odbc (lua_State *L);

static env_data *getenvironment (lua_State *L) {

env_data *env = (env_data *)luaL_checkudata (L, 1, LUASQL_ENVIRONMENT_ODBC);

luaL_argcheck (L, env != NULL, 1, LUASQL_PREFIX"environment expected");

luaL_argcheck (L, !env->closed, 1, LUASQL_PREFIX"environment is closed");

return env;

}

static conn_data *getconnection (lua_State *L) {

conn_data *conn = (conn_data *)luaL_checkudata (L, 1, LUASQL_CONNECTION_ODBC);

luaL_argcheck (L, conn != NULL, 1, LUASQL_PREFIX"connection expected");

luaL_argcheck (L, !conn->closed, 1, LUASQL_PREFIX"connection is closed");

return conn;

}

static cur_data *getcursor (lua_State *L) {

cur_data *cursor = (cur_data *)luaL_checkudata (L, 1, LUASQL_CURSOR_ODBC);

luaL_argcheck (L, cursor != NULL, 1, LUASQL_PREFIX"cursor expected");

luaL_argcheck (L, !cursor->closed, 1, LUASQL_PREFIX"cursor is closed");

return cursor;

}

static int pass(lua_State *L) {

lua_pushboolean (L, 1);

return 1;

}

static int fail(lua_State *L, const SQLSMALLINT type, const SQLHANDLE handle) {

SQLCHAR State[6];

SQLINTEGER NativeError;

SQLSMALLINT MsgSize, i;

SQLRETURN ret;

char Msg[SQL_MAX_MESSAGE_LENGTH];

luaL_Buffer b;

lua_pushnil(L);

luaL_buffinit(L, &b);

i = 1;

while (1) {

ret = SQLGetDiagRec(type, handle, i, State, &NativeError, Msg,

sizeof(Msg), &MsgSize);

if (ret == SQL_NO_DATA) break;

luaL_addlstring(&b, Msg, MsgSize);

luaL_putchar(&b, '\n');

i++;

}

luaL_pushresult(&b);

return 2;

}

static const char *sqltypetolua (const SQLSMALLINT type) {

switch (type) {

case SQL_UNKNOWN_TYPE: case SQL_CHAR: case SQL_VARCHAR:

case SQL_TYPE_DATE: case SQL_TYPE_TIME: case SQL_TYPE_TIMESTAMP:

case SQL_DATE: case SQL_INTERVAL: case SQL_TIMESTAMP:

case SQL_LONGVARCHAR:

return "string";

case SQL_BIGINT: case SQL_TINYINT: case SQL_NUMERIC:

case SQL_DECIMAL: case SQL_INTEGER: case SQL_SMALLINT:

case SQL_FLOAT: case SQL_REAL: case SQL_DOUBLE:

return "number";

case SQL_BINARY: case SQL_VARBINARY: case SQL_LONGVARBINARY:

return "binary"; /* !!!!!! nao seria string? */

case SQL_BIT:

return "boolean";

default:

assert(0);

return NULL;

}

}

static int push_column(lua_State *L, int coltypes, const SQLHSTMT hstmt,

SQLUSMALLINT i) {

const char *tname;

char type;

/* get column type from types table */

lua_rawgeti (L, LUA_REGISTRYINDEX, coltypes);

lua_rawgeti (L, -1, i); /* typename of the column */

tname = lua_tostring(L, -1);

if (!tname)

return luasql_faildirect(L, LUASQL_PREFIX"Invalid type in table.");

type = tname[1];

lua_pop(L, 2); /* pops type name and coltypes table */

/* deal with data according to type */

switch (type) {

/* nUmber */

case 'u': {

double num;

SQLINTEGER got;

SQLRETURN rc = SQLGetData(hstmt, i, SQL_C_DOUBLE, &num, 0, &got);

if (error(rc))

return fail(L, hSTMT, hstmt);

if (got == SQL_NULL_DATA)

lua_pushnil(L);

else

lua_pushnumber(L, num);

return 0;

}

/* bOol */

case 'o': {

char b;

SQLINTEGER got;

SQLRETURN rc = SQLGetData(hstmt, i, SQL_C_BIT, &b, 0, &got);

if (error(rc))

return fail(L, hSTMT, hstmt);

if (got == SQL_NULL_DATA)

lua_pushnil(L);

else

lua_pushboolean(L, b);

return 0;

}

/* sTring */

case 't':

/* bInary */

case 'i': {

SQLSMALLINT stype = (type == 't') ? SQL_C_CHAR : SQL_C_BINARY;

SQLINTEGER got;

char *buffer;

luaL_Buffer b;

SQLRETURN rc;

luaL_buffinit(L, &b);

buffer = luaL_prepbuffer(&b);

rc = SQLGetData(hstmt, i, stype, buffer, LUAL_BUFFERSIZE, &got);

if (got == SQL_NULL_DATA) {

lua_pushnil(L);

return 0;

}

/* concat intermediary chunks */

while (rc == SQL_SUCCESS_WITH_INFO) {

if (got >= LUAL_BUFFERSIZE || got == SQL_NO_TOTAL) {

got = LUAL_BUFFERSIZE;

/* get rid of null termination in string block */

if (stype == SQL_C_CHAR) got--;

}

luaL_addsize(&b, got);

buffer = luaL_prepbuffer(&b);

rc = SQLGetData(hstmt, i, stype, buffer,

LUAL_BUFFERSIZE, &got);

}

/* concat last chunk */

if (rc == SQL_SUCCESS) {

if (got >= LUAL_BUFFERSIZE || got == SQL_NO_TOTAL) {

got = LUAL_BUFFERSIZE;

/* get rid of null termination in string block */

if (stype == SQL_C_CHAR) got--;

}

luaL_addsize(&b, got);

}

if (rc == SQL_ERROR) return fail(L, hSTMT, hstmt);

/* return everything we got */

luaL_pushresult(&b);

return 0;

}

}

return 0;

}

static int cur_fetch (lua_State *L) {

cur_data *cur = (cur_data *) getcursor (L);

SQLHSTMT hstmt = cur->hstmt;

int ret;

SQLRETURN rc = SQLFetch(cur->hstmt);

if (rc == SQL_NO_DATA) {

lua_pushnil(L);

return 1;

} else if (error(rc)) return fail(L, hSTMT, hstmt);

if (lua_istable (L, 2)) {

SQLUSMALLINT i;

const char *opts = luaL_optstring (L, 3, "n");

int num = strchr (opts, 'n') != NULL;

int alpha = strchr (opts, 'a') != NULL;

for (i = 1; i <= cur->numcols; i++) {

ret = push_column (L, cur->coltypes, hstmt, i);

if (ret)

return ret;

if (alpha) {

lua_rawgeti (L, LUA_REGISTRYINDEX, cur->colnames);

lua_rawgeti (L, -1, i); /* gets column name */

lua_pushvalue (L, -3); /* duplicates column value */

lua_rawset (L, 2); /* table[name] = value */

lua_pop (L, 1); /* pops colnames table */

}

if (num)

lua_rawseti (L, 2, i);

else

lua_pop (L, 1); /* pops value */

}

lua_pushvalue (L, 2);

return 1; /* return table */

}

else {

SQLUSMALLINT i;

luaL_checkstack (L, cur->numcols, LUASQL_PREFIX"too many columns");

for (i = 1; i <= cur->numcols; i++) {

ret = push_column (L, cur->coltypes, hstmt, i);

if (ret)

return ret;

}

return cur->numcols;

}

}

static int cur_close (lua_State *L) {

conn_data *conn;

cur_data *cur = (cur_data *) luaL_checkudata (L, 1, LUASQL_CURSOR_ODBC);

SQLRETURN ret;

luaL_argcheck (L, cur != NULL, 1, LUASQL_PREFIX"cursor expected");

if (cur->closed) {

lua_pushboolean (L, 0);

return 1;

}

/* Nullify structure fields. */

cur->closed = 1;

ret = SQLCloseCursor(cur->hstmt);

if (error(ret))

return fail(L, hSTMT, cur->hstmt);

ret = SQLFreeHandle(hSTMT, cur->hstmt);

if (error(ret))

return fail(L, hSTMT, cur->hstmt);

/* Decrement cursor counter on connection object */

lua_rawgeti (L, LUA_REGISTRYINDEX, cur->conn);

conn = lua_touserdata (L, -1);

conn->cur_counter--;

luaL_unref (L, LUA_REGISTRYINDEX, cur->conn);

luaL_unref (L, LUA_REGISTRYINDEX, cur->colnames);

luaL_unref (L, LUA_REGISTRYINDEX, cur->coltypes);

return pass(L);

}

static int cur_colnames (lua_State *L) {

cur_data *cur = (cur_data *) getcursor (L);

lua_rawgeti (L, LUA_REGISTRYINDEX, cur->colnames);

return 1;

}

static int cur_coltypes (lua_State *L) {

cur_data *cur = (cur_data *) getcursor (L);

lua_rawgeti (L, LUA_REGISTRYINDEX, cur->coltypes);

return 1;

}

static void create_colinfo (lua_State *L, cur_data *cur) {

SQLCHAR buffer[256];

SQLSMALLINT namelen, datatype, i;

SQLRETURN ret;

int types, names;

lua_newtable(L);

types = lua_gettop (L);

lua_newtable(L);

names = lua_gettop (L);

for (i = 1; i <= cur->numcols; i++) {

ret = SQLDescribeCol(cur->hstmt, i, buffer, sizeof(buffer),

&namelen, &datatype, NULL, NULL, NULL);

/*if (ret == SQL_ERROR) return fail(L, hSTMT, cur->hstmt);*/

lua_pushstring (L, buffer);

lua_rawseti (L, names, i);

lua_pushstring(L, sqltypetolua(datatype));

lua_rawseti (L, types, i);

}

cur->colnames = luaL_ref (L, LUA_REGISTRYINDEX);

cur->coltypes = luaL_ref (L, LUA_REGISTRYINDEX);

}

static int create_cursor (lua_State *L, int o, conn_data *conn,

const SQLHSTMT hstmt, const SQLSMALLINT numcols) {

cur_data *cur = (cur_data *) lua_newuserdata(L, sizeof(cur_data));

luasql_setmeta (L, LUASQL_CURSOR_ODBC);

conn->cur_counter++;

/* fill in structure */

cur->closed = 0;

cur->conn = LUA_NOREF;

cur->numcols = numcols;

cur->colnames = LUA_NOREF;

cur->coltypes = LUA_NOREF;

cur->hstmt = hstmt;

lua_pushvalue (L, o);

cur->conn = luaL_ref (L, LUA_REGISTRYINDEX);

/* make and store column information table */

create_colinfo (L, cur);

return 1;

}

static int conn_close (lua_State *L) {

SQLRETURN ret;

env_data *env;

conn_data *conn = (conn_data *)luaL_checkudata(L,1,LUASQL_CONNECTION_ODBC);

luaL_argcheck (L, conn != NULL, 1, LUASQL_PREFIX"connection expected");

if (conn->closed) {

lua_pushboolean (L, 0);

return 1;

}

if (conn->cur_counter > 0)

return luaL_error (L, LUASQL_PREFIX"there are open cursors");

/* Decrement connection counter on environment object */

lua_rawgeti (L, LUA_REGISTRYINDEX, conn->env);

env = lua_touserdata (L, -1);

env->conn_counter--;

/* Nullify structure fields. */

conn->closed = 1;

luaL_unref (L, LUA_REGISTRYINDEX, conn->env);

ret = SQLDisconnect(conn->hdbc);

if (error(ret))

return fail(L, hDBC, conn->hdbc);

ret = SQLFreeHandle(hDBC, conn->hdbc);

if (error(ret))

return fail(L, hDBC, conn->hdbc);

return pass(L);

}

static int conn_execute (lua_State *L) {

conn_data *conn = (conn_data *) getconnection (L);

const char *statement = luaL_checkstring(L, 2);

SQLHDBC hdbc = conn->hdbc;

SQLHSTMT hstmt;

SQLSMALLINT numcols;

SQLRETURN ret;

ret = SQLAllocHandle(hSTMT, hdbc, &hstmt);

if (error(ret))

return fail(L, hDBC, hdbc);

ret = SQLPrepare(hstmt, (char *) statement, SQL_NTS);

if (error(ret)) {

ret = fail(L, hSTMT, hstmt);

SQLFreeHandle(hSTMT, hstmt);

return ret;

}

/* execute the statement */

ret = SQLExecute (hstmt);

if (error(ret)) {

ret = fail(L, hSTMT, hstmt);

SQLFreeHandle(hSTMT, hstmt);

return ret;

}

/* determine the number of results */

ret = SQLNumResultCols (hstmt, &numcols);

if (error(ret)) {

ret = fail(L, hSTMT, hstmt);

SQLFreeHandle(hSTMT, hstmt);

return ret;

}

if (numcols > 0)

/* if there is a results table (e.g., SELECT) */

return create_cursor (L, 1, conn, hstmt, numcols);

else {

/* if action has no results (e.g., UPDATE) */

SQLINTEGER numrows;

ret = SQLRowCount(hstmt, &numrows);

if (error(ret)) {

ret = fail(L, hSTMT, hstmt);

SQLFreeHandle(hSTMT, hstmt);

return ret;

}

lua_pushnumber(L, numrows);

SQLFreeHandle(hSTMT, hstmt);

return 1;

}

}

static int conn_commit (lua_State *L) {

conn_data *conn = (conn_data *) getconnection (L);

SQLRETURN ret = SQLEndTran(hDBC, conn->hdbc, SQL_COMMIT);

if (error(ret))

return fail(L, hSTMT, conn->hdbc);

else

return pass(L);

}

static int conn_rollback (lua_State *L) {

conn_data *conn = (conn_data *) getconnection (L);

SQLRETURN ret = SQLEndTran(hDBC, conn->hdbc, SQL_ROLLBACK);

if (error(ret))

return fail(L, hSTMT, conn->hdbc);

else

return pass(L);

}

static int conn_setautocommit (lua_State *L) {

conn_data *conn = (conn_data *) getconnection (L);

SQLRETURN ret;

if (lua_toboolean (L, 2)) {

ret = SQLSetConnectAttr(conn->hdbc, SQL_ATTR_AUTOCOMMIT,

(SQLPOINTER) SQL_AUTOCOMMIT_ON, 0);

} else {

ret = SQLSetConnectAttr(conn->hdbc, SQL_ATTR_AUTOCOMMIT,

(SQLPOINTER) SQL_AUTOCOMMIT_OFF, 0);

}

if (error(ret))

return fail(L, hSTMT, conn->hdbc);

else

return pass(L);

}

static int create_connection (lua_State *L, int o, env_data *env, SQLHDBC hdbc) {

conn_data *conn = (conn_data *) lua_newuserdata(L, sizeof(conn_data));

/* set auto commit mode */

SQLRETURN ret = SQLSetConnectAttr(hdbc, SQL_ATTR_AUTOCOMMIT,

(SQLPOINTER) SQL_AUTOCOMMIT_ON, 0);

if (error(ret))

return fail(L, hDBC, hdbc);

luasql_setmeta (L, LUASQL_CONNECTION_ODBC);

/* fill in structure */

conn->closed = 0;

conn->cur_counter = 0;

conn->env = LUA_NOREF;

conn->hdbc = hdbc;

lua_pushvalue (L, o);

conn->env = luaL_ref (L, LUA_REGISTRYINDEX);

env->conn_counter++;

return 1;

}

static int env_connect (lua_State *L) {

env_data *env = (env_data *) getenvironment (L);

const char *sourcename = luaL_checkstring (L, 2);

const char *username = luaL_optstring (L, 3, NULL);

const char *password = luaL_optstring (L, 4, NULL);

SQLHDBC hdbc;

SQLRETURN ret;

/* tries to allocate connection handle */

ret = SQLAllocHandle (hDBC, env->henv, &hdbc);

if (error(ret))

return luasql_faildirect (L, LUASQL_PREFIX"connection allocation error.");

/* tries to connect handle */

ret = SQLConnect (hdbc, (char *) sourcename, SQL_NTS,

(char *) username, SQL_NTS, (char *) password, SQL_NTS);

if (error(ret)) {

ret = fail(L, hDBC, hdbc);

SQLFreeHandle(hDBC, hdbc);

return ret;

}

/* success, return connection object */

return create_connection (L, 1, env, hdbc);

}

static int env_close (lua_State *L) {

SQLRETURN ret;

env_data *env = (env_data *)luaL_checkudata(L, 1, LUASQL_ENVIRONMENT_ODBC);

luaL_argcheck (L, env != NULL, 1, LUASQL_PREFIX"environment expected");

if (env->closed) {

lua_pushboolean (L, 0);

return 1;

}

if (env->conn_counter > 0)

return luaL_error (L, LUASQL_PREFIX"there are open connections");

env->closed = 1;

ret = SQLFreeHandle (hENV, env->henv);

if (error(ret)) {

int ret = fail(L, hENV, env->henv);

env->henv = NULL;

return ret;

}

return pass(L);

}

static void create_metatables (lua_State *L) {

struct luaL_reg environment_methods[] = {

{"close", env_close},

{"connect", env_connect},

{NULL, NULL},

};

struct luaL_reg connection_methods[] = {

{"close", conn_close},

{"execute", conn_execute},

{"commit", conn_commit},

{"rollback", conn_rollback},

{"setautocommit", conn_setautocommit},

{NULL, NULL},

};

struct luaL_reg cursor_methods[] = {

{"close", cur_close},

{"fetch", cur_fetch},

{"getcoltypes", cur_coltypes},

{"getcolnames", cur_colnames},

{NULL, NULL},

};

luasql_createmeta (L, LUASQL_ENVIRONMENT_ODBC, environment_methods);

luasql_createmeta (L, LUASQL_CONNECTION_ODBC, connection_methods);

luasql_createmeta (L, LUASQL_CURSOR_ODBC, cursor_methods);

lua_pop (L, 3);

}

static int create_environment (lua_State *L) {

env_data *env;

SQLHENV henv;

SQLRETURN ret = SQLAllocHandle(hENV, SQL_NULL_HANDLE, &henv);

if (error(ret))

return luasql_faildirect(L,LUASQL_PREFIX"error creating environment.");

ret = SQLSetEnvAttr (henv, SQL_ATTR_ODBC_VERSION,

(void*)SQL_OV_ODBC3, 0);

if (error(ret)) {

ret = luasql_faildirect (L, LUASQL_PREFIX"error setting SQL version.");

SQLFreeHandle (hENV, henv);

return ret;

}

env = (env_data *)lua_newuserdata (L, sizeof (env_data));

luasql_setmeta (L, LUASQL_ENVIRONMENT_ODBC);

/* fill in structure */

env->closed = 0;

env->conn_counter = 0;

env->henv = henv;

return 1;

}

LUASQL_API int luaopen_luasql_odbc (lua_State *L) {

struct luaL_reg driver[] = {

{"odbc", create_environment},

{NULL, NULL},

};

create_metatables (L);

luaL_openlib (L, LUASQL_TABLENAME, driver, 0);

luasql_set_info (L);

return 1;

}