1 /*
2 ** LuaSQL, PostgreSQL driver
3 ** Authors: Pedro Rabinovitch, Roberto Ierusalimschy, Carlos Cassino
4 ** Tomas Guisasola, Eduardo Quintao
5 ** See Copyright Notice in license.html
6 */
7
8 #include <assert.h>
9 #include <stdio.h>
10 #include <stdlib.h>
11 #include <string.h>
12 #include <ctype.h>
13
14 #include "libpq-fe.h"
15
16 #include "lua.h"
17 #include "lauxlib.h"
18
19
20 #include "luasql.h"
21
22 #define LUASQL_ENVIRONMENT_PG "PostgreSQL environment"
23 #define LUASQL_CONNECTION_PG "PostgreSQL connection"
24 #define LUASQL_CURSOR_PG "PostgreSQL cursor"
25
26 typedef struct {
27 short closed;
28 } env_data;
29
30
31 typedef struct {
32 short closed;
33 int env; /* reference to environment */
34 int auto_commit; /* 0 for manual commit */
35 PGconn *pg_conn;
36 } conn_data;
37
38
39 typedef struct {
40 short closed;
41 int conn; /* reference to connection */
42 int numcols; /* number of columns */
43 int colnames, coltypes; /* reference to column information tables */
44 int curr_tuple; /* next tuple to be read */
45 PGresult *pg_res;
46 } cur_data;
47
48
49 typedef void (*creator) (lua_State *L, cur_data *cur);
50
51
52 /*
53 ** Check for valid environment.
54 */
getenvironment(lua_State * L)55 static env_data *getenvironment (lua_State *L) {
56 env_data *env = (env_data *)luaL_checkudata (L, 1, LUASQL_ENVIRONMENT_PG);
57 luaL_argcheck (L, env != NULL, 1, LUASQL_PREFIX"environment expected");
58 luaL_argcheck (L, !env->closed, 1, LUASQL_PREFIX"environment is closed");
59 return env;
60 }
61
62
63 /*
64 ** Check for valid connection.
65 */
getconnection(lua_State * L)66 static conn_data *getconnection (lua_State *L) {
67 conn_data *conn = (conn_data *)luaL_checkudata (L, 1, LUASQL_CONNECTION_PG);
68 luaL_argcheck (L, conn != NULL, 1, LUASQL_PREFIX"connection expected");
69 luaL_argcheck (L, !conn->closed, 1, LUASQL_PREFIX"connection is closed");
70 return conn;
71 }
72
73
74 /*
75 ** Check for valid cursor.
76 */
getcursor(lua_State * L)77 static cur_data *getcursor (lua_State *L) {
78 cur_data *cur = (cur_data *)luaL_checkudata (L, 1, LUASQL_CURSOR_PG);
79 luaL_argcheck (L, cur != NULL, 1, LUASQL_PREFIX"cursor expected");
80 luaL_argcheck (L, !cur->closed, 1, LUASQL_PREFIX"cursor is closed");
81 return cur;
82 }
83
84
85 /*
86 ** Push the value of #i field of #tuple row.
87 */
pushvalue(lua_State * L,PGresult * res,int tuple,int i)88 static void pushvalue (lua_State *L, PGresult *res, int tuple, int i) {
89 if (PQgetisnull (res, tuple, i-1))
90 lua_pushnil (L);
91 else
92 lua_pushstring (L, PQgetvalue (res, tuple, i-1));
93 }
94
95
96 /*
97 ** Closes the cursor and nullify all structure fields.
98 */
cur_nullify(lua_State * L,cur_data * cur)99 static void cur_nullify (lua_State *L, cur_data *cur) {
100 /* Nullify structure fields. */
101 cur->closed = 1;
102 PQclear(cur->pg_res);
103 luaL_unref (L, LUA_REGISTRYINDEX, cur->conn);
104 luaL_unref (L, LUA_REGISTRYINDEX, cur->colnames);
105 luaL_unref (L, LUA_REGISTRYINDEX, cur->coltypes);
106 }
107
108
109 /*
110 ** Get another row of the given cursor.
111 */
cur_fetch(lua_State * L)112 static int cur_fetch (lua_State *L) {
113 cur_data *cur = getcursor (L);
114 PGresult *res = cur->pg_res;
115 int tuple = cur->curr_tuple;
116
117 if (tuple >= PQntuples(cur->pg_res)) {
118 cur_nullify (L, cur);
119 lua_pushnil(L); /* no more results */
120 return 1;
121 }
122
123 cur->curr_tuple++;
124 if (lua_istable (L, 2)) {
125 int i;
126 const char *opts = luaL_optstring (L, 3, "n");
127 if (strchr (opts, 'n') != NULL)
128 /* Copy values to numerical indices */
129 for (i = 1; i <= cur->numcols; i++) {
130 pushvalue (L, res, tuple, i);
131 lua_rawseti (L, 2, i);
132 }
133 if (strchr (opts, 'a') != NULL)
134 /* Copy values to alphanumerical indices */
135 for (i = 1; i <= cur->numcols; i++) {
136 lua_pushstring (L, PQfname (res, i-1));
137 pushvalue (L, res, tuple, i);
138 lua_rawset (L, 2);
139 }
140 lua_pushvalue(L, 2);
141 return 1; /* return table */
142 }
143 else {
144 int i;
145 luaL_checkstack (L, cur->numcols, LUASQL_PREFIX"too many columns");
146 for (i = 1; i <= cur->numcols; i++)
147 pushvalue (L, res, tuple, i);
148 return cur->numcols; /* return #numcols values */
149 }
150 }
151
152
153 /*
154 ** Cursor object collector function
155 */
cur_gc(lua_State * L)156 static int cur_gc (lua_State *L) {
157 cur_data *cur = (cur_data *)luaL_checkudata (L, 1, LUASQL_CURSOR_PG);
158 if (cur != NULL && !(cur->closed))
159 cur_nullify (L, cur);
160 return 0;
161 }
162
163
164 /*
165 ** Closes the cursor on top of the stack.
166 ** Returns true in case of success, or false in case the cursor was
167 ** already closed.
168 ** Throws an error if the argument is not a cursor.
169 */
cur_close(lua_State * L)170 static int cur_close (lua_State *L) {
171 cur_data *cur = (cur_data *)luaL_checkudata (L, 1, LUASQL_CURSOR_PG);
172 luaL_argcheck (L, cur != NULL, 1, LUASQL_PREFIX"cursor expected");
173 if (cur->closed) {
174 lua_pushboolean (L, 0);
175 return 1;
176 }
177 cur_nullify (L, cur); /* == cur_gc (L); */
178 lua_pushboolean (L, 1);
179 return 1;
180 }
181
182
183 /*
184 ** Get the internal database type of the given column.
185 */
getcolumntype(PGconn * conn,PGresult * result,int i,char * buff)186 static char *getcolumntype (PGconn *conn, PGresult *result, int i, char *buff) {
187 Oid codigo = PQftype (result, i);
188 char stmt[100];
189 PGresult *res;
190
191 sprintf (stmt, "select typname from pg_type where oid = %d", codigo);
192 res = PQexec(conn, stmt);
193 strcpy (buff, "undefined");
194
195 if (PQresultStatus (res) == PGRES_TUPLES_OK) {
196 if (PQntuples(res) > 0) {
197 char *name = PQgetvalue(res, 0, 0);
198 if (strcmp (name, "bpchar")==0 || strcmp (name, "varchar")==0) {
199 int modifier = PQfmod (result, i) - 4;
200 sprintf (buff, "%.20s (%d)", name, modifier);
201 }
202 else
203 strncpy (buff, name, 20);
204 }
205 }
206 PQclear(res);
207 return buff;
208 }
209
210
211 /*
212 ** Creates the list of fields names and pushes it on top of the stack.
213 */
create_colnames(lua_State * L,cur_data * cur)214 static void create_colnames (lua_State *L, cur_data *cur) {
215 PGresult *result = cur->pg_res;
216 int i;
217 lua_newtable (L);
218 for (i = 1; i <= cur->numcols; i++) {
219 lua_pushstring (L, PQfname (result, i-1));
220 lua_rawseti (L, -2, i);
221 }
222 }
223
224
225 /*
226 ** Creates the list of fields types and pushes it on top of the stack.
227 */
create_coltypes(lua_State * L,cur_data * cur)228 static void create_coltypes (lua_State *L, cur_data *cur) {
229 PGresult *result = cur->pg_res;
230 conn_data *conn;
231 char typename[100];
232 int i;
233 lua_rawgeti (L, LUA_REGISTRYINDEX, cur->conn);
234 if (!lua_isuserdata (L, -1))
235 luaL_error (L, LUASQL_PREFIX"invalid connection");
236 conn = (conn_data *)lua_touserdata (L, -1);
237 lua_newtable (L);
238 for (i = 1; i <= cur->numcols; i++) {
239 lua_pushstring(L, getcolumntype (conn->pg_conn, result, i-1, typename));
240 lua_rawseti (L, -2, i);
241 }
242 }
243
244
245 /*
246 ** Pushes a column information table on top of the stack.
247 ** If the table isn't built yet, call the creator function and stores
248 ** a reference to it on the cursor structure.
249 */
_pushtable(lua_State * L,cur_data * cur,size_t off,creator func)250 static void _pushtable (lua_State *L, cur_data *cur, size_t off, creator func) {
251 int *ref = (int *)((char *)cur + off);
252 if (*ref != LUA_NOREF)
253 lua_rawgeti (L, LUA_REGISTRYINDEX, *ref);
254 else {
255 func (L, cur);
256 /* Stores a reference to it on the cursor structure */
257 lua_pushvalue (L, -1);
258 *ref = luaL_ref (L, LUA_REGISTRYINDEX);
259 }
260 }
261 #define pushtable(L,c,m,f) (_pushtable(L,c,offsetof(cur_data,m),f))
262
263
264 /*
265 ** Return the list of field names.
266 */
cur_getcolnames(lua_State * L)267 static int cur_getcolnames (lua_State *L) {
268 pushtable (L, getcursor(L), colnames, create_colnames);
269 return 1;
270 }
271
272
273 /*
274 ** Return the list of field types.
275 */
cur_getcoltypes(lua_State * L)276 static int cur_getcoltypes (lua_State *L) {
277 pushtable (L, getcursor(L), coltypes, create_coltypes);
278 return 1;
279 }
280
281
282 /*
283 ** Push the number of rows.
284 */
cur_numrows(lua_State * L)285 static int cur_numrows (lua_State *L) {
286 lua_pushnumber (L, PQntuples (getcursor(L)->pg_res));
287 return 1;
288 }
289
290
291 /*
292 ** Create a new Cursor object and push it on top of the stack.
293 */
create_cursor(lua_State * L,int conn,PGresult * result)294 static int create_cursor (lua_State *L, int conn, PGresult *result) {
295 cur_data *cur = (cur_data *)lua_newuserdata(L, sizeof(cur_data));
296 luasql_setmeta (L, LUASQL_CURSOR_PG);
297
298 /* fill in structure */
299 cur->closed = 0;
300 cur->conn = LUA_NOREF;
301 cur->numcols = PQnfields(result);
302 cur->colnames = LUA_NOREF;
303 cur->coltypes = LUA_NOREF;
304 cur->curr_tuple = 0;
305 cur->pg_res = result;
306 lua_pushvalue (L, conn);
307 cur->conn = luaL_ref (L, LUA_REGISTRYINDEX);
308
309 return 1;
310 }
311
312
sql_commit(conn_data * conn)313 static void sql_commit(conn_data *conn) {
314 PQclear(PQexec(conn->pg_conn, "COMMIT"));
315 }
316
317
sql_begin(conn_data * conn)318 static void sql_begin(conn_data *conn) {
319 PQclear(PQexec(conn->pg_conn, "BEGIN"));
320 }
321
322
sql_rollback(conn_data * conn)323 static void sql_rollback(conn_data *conn) {
324 PQclear(PQexec(conn->pg_conn, "ROLLBACK"));
325 }
326
327
328 /*
329 ** Connection object collector function
330 */
conn_gc(lua_State * L)331 static int conn_gc (lua_State *L) {
332 conn_data *conn = (conn_data *)luaL_checkudata (L, 1, LUASQL_CONNECTION_PG);
333 if (conn != NULL && !(conn->closed)) {
334 /* Nullify structure fields. */
335 conn->closed = 1;
336 luaL_unref (L, LUA_REGISTRYINDEX, conn->env);
337 PQfinish (conn->pg_conn);
338 }
339 return 0;
340 }
341
342
343 /*
344 ** Closes the connection on top of the stack.
345 ** Returns true in case of success, or false in case the connection was
346 ** already closed.
347 ** Throws an error if the argument is not a connection.
348 */
conn_close(lua_State * L)349 static int conn_close (lua_State *L) {
350 conn_data *conn = (conn_data *)luaL_checkudata (L, 1, LUASQL_CONNECTION_PG);
351 luaL_argcheck (L, conn != NULL, 1, LUASQL_PREFIX"connection expected");
352 if (conn->closed) {
353 lua_pushboolean (L, 0);
354 return 1;
355 }
356 conn_gc (L);
357 lua_pushboolean (L, 1);
358 return 1;
359 }
360
361
362 /*
363 ** Escapes a string for use within an SQL statement.
364 ** Returns a string with the escaped string.
365 */
conn_escape(lua_State * L)366 static int conn_escape (lua_State *L) {
367 conn_data *conn = getconnection (L);
368 size_t len;
369 const char *from = luaL_checklstring (L, 2, &len);
370 int error;
371 int ret = 1;
372 luaL_Buffer b;
373 char *to;
374 #if !defined(LUA_VERSION_NUM) || (LUA_VERSION_NUM == 501)
375 /* Lua 5.0 and 5.1 */
376 luaL_buffinit (L, &b);
377 do {
378 int max = LUAL_BUFFERSIZE / 2;
379 size_t bytes_copied;
380 size_t this_len = (len > max) ? max : len;
381 to = luaL_prepbuffer (&b);
382 bytes_copied = PQescapeStringConn (conn->pg_conn, to, from, this_len, &error);
383 if (error != 0) { /* failed ! */
384 return luasql_failmsg (L, "cannot escape string. PostgreSQL: ", PQerrorMessage (conn->pg_conn));
385 }
386 luaL_addsize (&b, bytes_copied);
387 from += this_len;
388 len -= this_len;
389 } while (len > 0);
390 luaL_pushresult (&b);
391 #else
392 /* Lua 5.2 and 5.3 */
393 to = luaL_buffinitsize (L, &b, 2*len+1);
394 len = PQescapeStringConn (conn->pg_conn, to, from, len, &error);
395 if (error == 0) { /* success ! */
396 luaL_pushresultsize (&b, len);
397 } else {
398 ret = luasql_failmsg (L, "cannot escape string. PostgreSQL: ", PQerrorMessage (conn->pg_conn));
399 }
400 #endif
401 return ret;
402 }
403
404
405 /*
406 ** Execute an SQL statement.
407 ** Return a Cursor object if the statement is a query, otherwise
408 ** return the number of tuples affected by the statement.
409 */
conn_execute(lua_State * L)410 static int conn_execute (lua_State *L) {
411 conn_data *conn = getconnection (L);
412 const char *statement = luaL_checkstring (L, 2);
413 PGresult *res = PQexec(conn->pg_conn, statement);
414 if (res && PQresultStatus(res)==PGRES_COMMAND_OK) {
415 /* no tuples returned */
416 lua_pushnumber(L, atof(PQcmdTuples(res)));
417 PQclear (res);
418 return 1;
419 }
420 else if (res && PQresultStatus(res)==PGRES_TUPLES_OK)
421 /* tuples returned */
422 return create_cursor (L, 1, res);
423 else {
424 /* error */
425 PQclear (res);
426 return luasql_failmsg(L, "error executing statement. PostgreSQL: ", PQerrorMessage(conn->pg_conn));
427 }
428 }
429
430
431 /*
432 ** Commit the current transaction.
433 */
conn_commit(lua_State * L)434 static int conn_commit (lua_State *L) {
435 conn_data *conn = getconnection (L);
436 sql_commit(conn);
437 if (conn->auto_commit == 0) {
438 sql_begin(conn);
439 lua_pushboolean (L, 1);
440 } else
441 lua_pushboolean (L, 0);
442 return 1;
443 }
444
445
446 /*
447 ** Rollback the current transaction.
448 */
conn_rollback(lua_State * L)449 static int conn_rollback (lua_State *L) {
450 conn_data *conn = getconnection (L);
451 sql_rollback(conn);
452 if (conn->auto_commit == 0) {
453 sql_begin(conn);
454 lua_pushboolean (L, 1);
455 } else
456 lua_pushboolean (L, 0);
457 return 1;
458 }
459
460
461 /*
462 ** Set "auto commit" property of the connection.
463 ** If 'true', then rollback current transaction.
464 ** If 'false', then start a new transaction.
465 */
conn_setautocommit(lua_State * L)466 static int conn_setautocommit (lua_State *L) {
467 conn_data *conn = getconnection (L);
468 if (lua_toboolean (L, 2)) {
469 conn->auto_commit = 1;
470 sql_rollback(conn); /* Undo active transaction. */
471 }
472 else {
473 conn->auto_commit = 0;
474 sql_begin(conn);
475 }
476 lua_pushboolean(L, 1);
477 return 1;
478 }
479
480
481 /*
482 ** Create a new Connection object and push it on top of the stack.
483 */
create_connection(lua_State * L,int env,PGconn * const pg_conn)484 static int create_connection (lua_State *L, int env, PGconn *const pg_conn) {
485 conn_data *conn = (conn_data *)lua_newuserdata(L, sizeof(conn_data));
486 luasql_setmeta (L, LUASQL_CONNECTION_PG);
487
488 /* fill in structure */
489 conn->closed = 0;
490 conn->env = LUA_NOREF;
491 conn->auto_commit = 1;
492 conn->pg_conn = pg_conn;
493 lua_pushvalue (L, env);
494 conn->env = luaL_ref (L, LUA_REGISTRYINDEX);
495 return 1;
496 }
497
498
notice_processor(void * arg,const char * message)499 static void notice_processor (void *arg, const char *message) {
500 (void)arg; (void)message;
501 /* arg == NULL */
502 }
503
504
505 /*
506 ** Connects to a data source.
507 */
env_connect(lua_State * L)508 static int env_connect (lua_State *L) {
509 const char *sourcename = luaL_checkstring(L, 2);
510 const char *username = luaL_optstring(L, 3, NULL);
511 const char *password = luaL_optstring(L, 4, NULL);
512 const char *pghost = luaL_optstring(L, 5, NULL);
513 const char *pgport = luaL_optstring(L, 6, NULL);
514 PGconn *conn;
515 getenvironment (L); /* validate environment */
516 conn = PQsetdbLogin(pghost, pgport, NULL, NULL, sourcename, username, password);
517
518 if (PQstatus(conn) == CONNECTION_BAD) {
519 int rc = luasql_failmsg(L, "error connecting to database. PostgreSQL: ", PQerrorMessage(conn));
520 PQfinish(conn);
521 return rc;
522 }
523 PQsetNoticeProcessor(conn, notice_processor, NULL);
524 return create_connection(L, 1, conn);
525 }
526
527
528 /*
529 ** Environment object collector function.
530 */
env_gc(lua_State * L)531 static int env_gc (lua_State *L) {
532 env_data *env = (env_data *)luaL_checkudata (L, 1, LUASQL_ENVIRONMENT_PG);
533 if (env != NULL && !(env->closed))
534 env->closed = 1;
535 return 0;
536 }
537
538
539 /*
540 ** Closes the environment on top of the stack.
541 ** Returns true in case of success, or false in case the environment was
542 ** already closed.
543 ** Throws an error if the argument is not an environment.
544 */
env_close(lua_State * L)545 static int env_close (lua_State *L) {
546 env_data *env = (env_data *)luaL_checkudata (L, 1, LUASQL_ENVIRONMENT_PG);
547 luaL_argcheck (L, env != NULL, 1, LUASQL_PREFIX"environment expected");
548 if (env->closed) {
549 lua_pushboolean (L, 0);
550 return 1;
551 }
552 env_gc (L);
553 lua_pushboolean (L, 1);
554 return 1;
555 }
556
557
558
559 /*
560 ** Create metatables for each class of object.
561 */
create_metatables(lua_State * L)562 static void create_metatables (lua_State *L) {
563 struct luaL_Reg environment_methods[] = {
564 {"__gc", env_gc},
565 {"close", env_close},
566 {"connect", env_connect},
567 {NULL, NULL},
568 };
569 struct luaL_Reg connection_methods[] = {
570 {"__gc", conn_gc},
571 {"close", conn_close},
572 {"escape", conn_escape},
573 {"execute", conn_execute},
574 {"commit", conn_commit},
575 {"rollback", conn_rollback},
576 {"setautocommit", conn_setautocommit},
577 {NULL, NULL},
578 };
579 struct luaL_Reg cursor_methods[] = {
580 {"__gc", cur_gc},
581 {"close", cur_close},
582 {"getcolnames", cur_getcolnames},
583 {"getcoltypes", cur_getcoltypes},
584 {"fetch", cur_fetch},
585 {"numrows", cur_numrows},
586 {NULL, NULL},
587 };
588 luasql_createmeta (L, LUASQL_ENVIRONMENT_PG, environment_methods);
589 luasql_createmeta (L, LUASQL_CONNECTION_PG, connection_methods);
590 luasql_createmeta (L, LUASQL_CURSOR_PG, cursor_methods);
591 lua_pop (L, 3);
592 }
593
594 /*
595 ** Creates an Environment and returns it.
596 */
create_environment(lua_State * L)597 static int create_environment (lua_State *L) {
598 env_data *env = (env_data *)lua_newuserdata(L, sizeof(env_data));
599 luasql_setmeta (L, LUASQL_ENVIRONMENT_PG);
600
601 /* fill in structure */
602 env->closed = 0;
603 return 1;
604 }
605
606
607 /*
608 ** Creates the metatables for the objects and registers the
609 ** driver open method.
610 */
luaopen_luasql_postgres(lua_State * L)611 LUASQL_API int luaopen_luasql_postgres (lua_State *L) {
612 struct luaL_Reg driver[] = {
613 {"postgres", create_environment},
614 {NULL, NULL},
615 };
616 create_metatables (L);
617 lua_newtable (L);
618 luaL_setfuncs (L, driver, 0);
619 luasql_set_info (L);
620 #if defined(PQlibVersion)
621 lua_pushliteral (L, "_CLIENTVERSION");
622 lua_pushinteger (L, PQlibVersion());
623 lua_settable (L, -3);
624 #endif
625 return 1;
626 }
627