/*

pgfunc works as a wrapper for postgres functions and as a loader as module for

pllua (internal)->internal functions.

pgfunc(text - function singnature)

pgfunc(text - function singnature, table - options)

options =

{

only_internal = true --false,

--[[if only_internal is false, pgfunc accepts any function]]

throwable = true --false

--[[ throwable makes PG_TRY PG_CATCH for non internal functions]]

}

Note:

Set returning functions are not supported.

No checks if function is strict.

*/

#include "pllua_pgfunc.h"

#include "plluacommon.h"

#include "pllua.h"

#include "pllua_xact_cleanup.h"

#include <catalog/pg_language.h>

#include <lib/stringinfo.h>

#include "pllua_errors.h"

#define RESET_CONTEXT_AFTER 1000

static const char pg_func_type_name[] = "pg_func";

static Oid

find_lang_oids(const char* lang)

{

HeapTuple tuple;

tuple = SearchSysCache(LANGNAME, CStringGetDatum(lang), 0, 0, 0);

if (HeapTupleIsValid(tuple))

{

Oid langtupoid = HeapTupleGetOid(tuple);

ReleaseSysCache(tuple);

return langtupoid;

}

return 0;

}

static Oid pllua_oid = 0;

static Oid plluau_oid = 0;

static Oid

get_pllua_oid()

{

if (pllua_oid !=0)

return pllua_oid;

return find_lang_oids("pllua");

}

static Oid

get_plluau_oid()

{

if (plluau_oid !=0)

return plluau_oid;

return find_lang_oids("plluau");

}

typedef struct{

bool only_internal;

bool throwable;

} Pgfunc_options;

typedef struct{

Oid funcid;

int numargs;

Oid *argtypes;

lua_CFunction callfunc;

Oid prorettype;

FmgrInfo fi;

Pgfunc_options options;

} PgFuncInfo, Lua_pgfunc;

typedef struct{

FmgrInfo fi;

ExprContext econtext;

ReturnSetInfo rsinfo;

FunctionCallInfoData fcinfo;

Oid prorettype;

} Lua_pgfunc_srf;

#define freeandnil(p) do{ if (p){\

pfree(p);\

p = NULL;\

}}while(0)

static void

clean_pgfuncinfo(Lua_pgfunc *data)

{

freeandnil (data->argtypes);

}

#ifdef PGFUNC_CLEANUP

static MemoryContext

get_tmpcontext()

{

MemoryContext mc;

mc = AllocSetContextCreate(TopMemoryContext,

"pgfunc temporary context",

ALLOCSET_DEFAULT_MINSIZE,

ALLOCSET_DEFAULT_INITSIZE,

ALLOCSET_DEFAULT_MAXSIZE);

return mc;

}

#endif

static MemoryContext tmpcontext;

#ifdef PGFUNC_CLEANUP

static int tmpcontext_usage = 0;

#endif

static int

pg_callable_func(lua_State *L)

{

MemoryContext m;

int i;

FunctionCallInfoData fcinfo;

Lua_pgfunc *fi;

#ifndef PGFUNC_CLEANUP

tmpcontext = CurTransactionContext;

#endif

fi = (Lua_pgfunc *) lua_touserdata(L, lua_upvalueindex(1));

InitFunctionCallInfoData(fcinfo, &fi->fi, fi->numargs, InvalidOid, NULL, NULL);

#ifdef PGFUNC_CLEANUP

if(tmpcontext_usage> RESET_CONTEXT_AFTER ){

MemoryContextReset(tmpcontext);

tmpcontext_usage = 0;

}

++tmpcontext_usage;

#endif

m = MemoryContextSwitchTo(tmpcontext);

for (i=0; i<fi->numargs; ++i){

fcinfo.arg[i] = luaP_todatum(L, fi->argtypes[i], 0, &fcinfo.argnull[i], i+1);

}

if(!fi->options.only_internal && fi->options.throwable){

SPI_push();

PG_TRY();

{

Datum d = FunctionCallInvoke(&fcinfo);

MemoryContextSwitchTo(m);

if (fcinfo.isnull) {

lua_pushnil(L);

} else {

luaP_pushdatum(L, d, fi->prorettype);

}

SPI_pop();

}

PG_CATCH();

{

lua_pop(L, lua_gettop(L));

push_spi_error(L, m); /*context switch to m inside push_spi_error*/

SPI_pop();

return lua_error(L);

}PG_END_TRY();

}else{

Datum d = FunctionCallInvoke(&fcinfo);

MemoryContextSwitchTo(m);

if (fcinfo.isnull) {

lua_pushnil(L);

} else {

luaP_pushdatum(L, d, fi->prorettype);

}

}

return 1;

}

static void

parse_options(lua_State *L, Pgfunc_options *opt){

lua_pushnil(L);

while (lua_next(L, -2) != 0) {

if (lua_type(L, -2) == LUA_TSTRING){

const char *key = lua_tostring(L, -2);

if (strcmp(key, "only_internal") == 0){

opt->only_internal = lua_toboolean(L, -1);

} else if (strcmp(key, "throwable") == 0){

opt->throwable = lua_toboolean(L, -1);

} else {

luaL_error(L, "pgfunc unknown option \"%s\"", key);

}

}

lua_pop(L, 1);

}

}

static int pgfunc_rowsaux (lua_State *L) {

ReturnSetInfo *rsinfo;

FunctionCallInfoData *fcinfo;

Datum d;

Oid prorettype;

Lua_pgfunc_srf *srfi;

srfi = (Lua_pgfunc_srf *) lua_touserdata(L, lua_upvalueindex(1));

rsinfo = &srfi->rsinfo;

fcinfo = &srfi->fcinfo;

prorettype = srfi->prorettype;

d = FunctionCallInvoke(fcinfo);

if ((!fcinfo->isnull)&&(rsinfo->isDone != ExprEndResult)){

luaP_pushdatum(L, d, prorettype);

return 1;

}

lua_pushnil(L);

return 1;

}

static int

pgfunc_rows (lua_State *L) {

int i;

Lua_pgfunc *fi;

ReturnSetInfo *rsinfo;

ExprContext *econtext;

FunctionCallInfoData *fcinfo;

Lua_pgfunc_srf *srfi;

int argc;

BEGINLUA;

argc = lua_gettop(L);

fi = (Lua_pgfunc *) lua_touserdata(L, lua_upvalueindex(1));

srfi = (Lua_pgfunc_srf *)lua_newuserdata(L, sizeof(Lua_pgfunc_srf));

econtext = &srfi->econtext;

rsinfo = &srfi->rsinfo;

fcinfo = &srfi->fcinfo;

srfi->prorettype = fi->prorettype;

fmgr_info(fi->funcid, &srfi->fi);

memset(econtext, 0, sizeof(ExprContext));

econtext->ecxt_per_query_memory = CurrentMemoryContext;

rsinfo->type = T_ReturnSetInfo;

rsinfo->econtext = econtext;

rsinfo->allowedModes = (int)(SFRM_ValuePerCall /*| SFRM_Materialize*/);

rsinfo->returnMode = SFRM_ValuePerCall;//SFRM_Materialize;

rsinfo->setResult = NULL;

rsinfo->setDesc = NULL;

InitFunctionCallInfoData((*fcinfo), &srfi->fi, fi->numargs, InvalidOid, NULL, (fmNodePtr)rsinfo);

for (i=0; i<fi->numargs; ++i){

if(i>=argc){

for (i = argc; i<fi->numargs; ++i){

fcinfo->arg[i] = 0 ;//Datum;

fcinfo->argnull[i] = true;

}

break;

}

fcinfo->arg[i] = luaP_todatum(L, fi->argtypes[i], 0, &fcinfo->argnull[i], i+1);

}

lua_pushcclosure(L, pgfunc_rowsaux, 1);

ENDLUAV(1);

return 1;

}

int

get_pgfunc(lua_State *L)

{

Lua_pgfunc *lf;

Pgfunc_options opt;

MemoryContext m;

const char* reg_name = NULL;

HeapTuple proctup;

Form_pg_proc proc;

int luasrc = 0;

Oid funcid = 0;

BEGINLUA;

#ifndef PGFUNC_CLEANUP

tmpcontext = CurTransactionContext;

#endif

opt.only_internal = true;

opt.throwable = true;

if (lua_gettop(L) == 2){

luaL_checktype(L, 2, LUA_TTABLE);

parse_options(L, &opt);

}else if (lua_gettop(L) != 1){

return luaL_error(L, "pgfunc(text): wrong arguments");

}

if(lua_type(L, 1) == LUA_TSTRING){

reg_name = luaL_checkstring(L, 1);

m = MemoryContextSwitchTo(tmpcontext);

PG_TRY();

{

funcid = DatumGetObjectId(DirectFunctionCall1(regprocedurein, CStringGetDatum(reg_name)));

}

PG_CATCH();{}

PG_END_TRY();

MemoryContextSwitchTo(m);

#ifdef PGFUNC_CLEANUP

MemoryContextReset(tmpcontext);

#endif

}else if (lua_type(L, 1) == LUA_TNUMBER){

funcid = luaL_checkinteger(L, 1);

}

if (!OidIsValid(funcid)){

if (reg_name)

return luaL_error(L,"failed to register %s", reg_name);

return luaL_error(L,"failed to register function with oid %d", funcid);

}

proctup = SearchSysCache1(PROCOID, ObjectIdGetDatum(funcid));

if (!HeapTupleIsValid(proctup)){

return luaL_error(L,"cache lookup failed for function %d", funcid);

}

proc = (Form_pg_proc) GETSTRUCT(proctup);

luasrc = ((proc->prolang == get_pllua_oid())

|| (proc->prolang == get_plluau_oid()));

if ( opt.only_internal

&&(proc->prolang != INTERNALlanguageId)

&&(!luasrc) ){

ReleaseSysCache(proctup);

return luaL_error(L, "supported only SQL/internal functions");

}

lf = (Lua_pgfunc *)lua_newuserdata(L, sizeof(Lua_pgfunc));

/*make it g/collected*/

luaP_getfield(L, pg_func_type_name);

lua_setmetatable(L, -2);

lf->prorettype = proc->prorettype;

lf->funcid = funcid;

lf->options = opt;

{

Oid *argtypes;

char **argnames;

char *argmodes;

int argc;

MemoryContext cur = CurrentMemoryContext;

MemoryContextSwitchTo(tmpcontext);

argc = get_func_arg_info(proctup,

&argtypes, &argnames, &argmodes);

MemoryContextSwitchTo(get_common_ctx());

lf->numargs = argc;

lf->argtypes = (Oid*)palloc(argc * sizeof(Oid));

memcpy(lf->argtypes, argtypes, argc * sizeof(Oid));

MemoryContextSwitchTo(cur);

#ifdef PGFUNC_CLEANUP

MemoryContextReset(tmpcontext);

#endif

}

if (luasrc){

bool isnull;

text *t;

const char *s;

luaL_Buffer b;

int pcall_result;

Datum prosrc;

if((lf->numargs != 1)

|| (lf->argtypes[0] != INTERNALOID)

|| (lf->prorettype != INTERNALOID)){

luaL_error(L, "pgfunc accepts only 'internal' pllua/u functions with internal argument");

}

prosrc = SysCacheGetAttr(PROCOID, proctup, Anum_pg_proc_prosrc, &isnull);

if (isnull) elog(ERROR, "[pgfunc]: null lua prosrc");

luaL_buffinit(L, &b);

luaL_addstring(&b,"do ");

t = DatumGetTextP(prosrc);

luaL_addlstring(&b, VARDATA(t), VARSIZE(t) - VARHDRSZ);

luaL_addstring(&b, " end");

luaL_pushresult(&b);

s = lua_tostring(L, -1);

ReleaseSysCache(proctup);

clean_pgfuncinfo(lf);

if (luaL_loadbuffer(L, s, strlen(s), "pgfunc chunk"))

luaL_error(L, "compile");

lua_remove(L, -2); /*delete source element*/

pcall_result = lua_pcall(L, 0, 1, 0);

lua_remove(L, -2); /*delete chunk*/

if(pcall_result == 0){

ENDLUAV(1);

return 1;

}

if( pcall_result == LUA_ERRRUN)

luaL_error(L,"%s %s","Runtime error:",lua_tostring(L, -1));

else if(pcall_result == LUA_ERRMEM)

luaL_error(L,"%s %s","Memory error:",lua_tostring(L, -1));

else if(pcall_result == LUA_ERRERR)

luaL_error(L,"%s %s","Error:",lua_tostring(L, -1));

return luaL_error(L, "pgfunc unknown error");

}

if(proc->proretset) {

lua_pushcclosure(L, pgfunc_rows, 1);

} else {

fmgr_info(funcid, &lf->fi);

lua_pushcclosure(L, pg_callable_func, 1);

}

ReleaseSysCache(proctup);

ENDLUAV(1);

return 1;

}

static void

__newmetatable (lua_State *L, const char *tname)

{

lua_newtable(L);

lua_pushlightuserdata(L, (void *) tname);

lua_pushvalue(L, -2);

lua_rawset(L, LUA_REGISTRYINDEX);

}

static int

gc_pg_func(lua_State *L)

{

Lua_pgfunc *lf = lua_touserdata(L, 1);

clean_pgfuncinfo(lf);

return 0;

}

static luaL_Reg regs[] = {

{"__gc", gc_pg_func},

{ NULL, NULL }

};

#ifdef PGFUNC_CLEANUP

static bool cxt_initialized = false;

#endif

void

register_funcinfo_mt(lua_State *L)

{

__newmetatable(L, pg_func_type_name);

luaP_register(L, regs);

lua_pop(L, 1);

#ifdef PGFUNC_CLEANUP

if (!cxt_initialized){

tmpcontext = get_tmpcontext();

cxt_initialized = true;

}

#endif

}