Sun, 23 Nov 2008 01:29:09 +0000
Initial import.
1 | 1 | #include "dbd_postgresql.h" |
2 | ||
3 | #define MAX_PLACEHOLDERS 9999 | |
4 | #define MAX_PLACEHOLDER_SIZE (1+4) /* $\d{4} */ | |
5 | ||
6 | static lua_push_type_t postgresql_to_lua_push(unsigned int postgresql_type) { | |
7 | lua_push_type_t lua_type; | |
8 | ||
9 | switch(postgresql_type) { | |
10 | case INT2OID: | |
11 | case INT4OID: | |
12 | lua_type = LUA_PUSH_INTEGER; | |
13 | break; | |
14 | ||
15 | case FLOAT4OID: | |
16 | case FLOAT8OID: | |
17 | lua_type = LUA_PUSH_NUMBER; | |
18 | break; | |
19 | ||
20 | case BOOLOID: | |
21 | lua_type = LUA_PUSH_BOOLEAN; | |
22 | break; | |
23 | ||
24 | default: | |
25 | lua_type = LUA_PUSH_STRING; | |
26 | } | |
27 | ||
28 | return lua_type; | |
29 | } | |
30 | ||
31 | static char *replace_placeholders(lua_State *L, const char *sql) { | |
32 | size_t len = strlen(sql); | |
33 | int num_placeholders = 0; | |
34 | int extra_space = 0; | |
35 | int i; | |
36 | char *newsql; | |
37 | int newpos = 1; | |
38 | int ph_num = 1; | |
39 | int in_quote = 0; | |
40 | ||
41 | for (i = 1; i < len; i++) { | |
42 | if (sql[i] == '?') { | |
43 | num_placeholders++; | |
44 | } | |
45 | } | |
46 | ||
47 | extra_space = num_placeholders * (MAX_PLACEHOLDER_SIZE-1); | |
48 | ||
49 | newsql = malloc(sizeof(char) * (len+extra_space+1)); | |
50 | memset(newsql, 0, sizeof(char) * (len+extra_space+1)); | |
51 | ||
52 | /* copy first char */ | |
53 | newsql[0] = sql[0]; | |
54 | for (i = 1; i < len; i++) { | |
55 | if (sql[i] == '\'' && sql[i-1] != '\\') { | |
56 | in_quote = !in_quote; | |
57 | } | |
58 | ||
59 | if (sql[i] == '?' && !in_quote) { | |
60 | size_t n; | |
61 | ||
62 | if (ph_num > MAX_PLACEHOLDERS) { | |
63 | luaL_error(L, "Sorry, you are using more than %d placeholders. Use ${num} format instead", MAX_PLACEHOLDERS); | |
64 | } | |
65 | ||
66 | n = snprintf(&newsql[newpos], MAX_PLACEHOLDER_SIZE, "$%u", ph_num++); | |
67 | ||
68 | newpos += n; | |
69 | } else { | |
70 | newsql[newpos] = sql[i]; | |
71 | newpos++; | |
72 | } | |
73 | } | |
74 | ||
75 | /* terminate string on the last position */ | |
76 | newsql[newpos] = '\0'; | |
77 | ||
78 | /* fprintf(stderr, "[%s]\n", newsql); */ | |
79 | return newsql; | |
80 | } | |
81 | ||
82 | static int statement_close(lua_State *L) { | |
83 | statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT); | |
84 | ||
85 | if (statement->result) { | |
86 | PQclear(statement->result); | |
87 | statement->result = NULL; | |
88 | } | |
89 | ||
90 | return 0; | |
91 | } | |
92 | ||
93 | static int statement_execute(lua_State *L) { | |
94 | int n = lua_gettop(L); | |
95 | statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT); | |
96 | int num_bind_params = n - 1; | |
97 | ExecStatusType status; | |
98 | int p; | |
99 | ||
100 | char **params; | |
101 | PGresult *result = NULL; | |
102 | ||
103 | statement->tuple = 0; | |
104 | ||
105 | params = malloc(num_bind_params * sizeof(params)); | |
106 | ||
107 | for (p = 2; p <= n; p++) { | |
108 | int i = p - 2; | |
109 | ||
110 | if (lua_isnil(L, p)) { | |
111 | params[i] = NULL; | |
112 | } else { | |
113 | const char *param = lua_tostring(L, p); | |
114 | size_t len = strlen(param) + 1; | |
115 | ||
116 | params[i] = malloc(len * sizeof(char)); | |
117 | memset(params[i], 0, len); | |
118 | ||
119 | strncpy(params[i], param, len); | |
120 | params[i][len] = '\0'; | |
121 | } | |
122 | } | |
123 | ||
124 | result = PQexecPrepared( | |
125 | statement->postgresql, | |
126 | statement->name, | |
127 | num_bind_params, | |
128 | (const char **)params, | |
129 | NULL, | |
130 | NULL, | |
131 | 0 | |
132 | ); | |
133 | ||
134 | for (p = 0; p < num_bind_params; p++) { | |
135 | if (params[p]) { | |
136 | free(params[p]); | |
137 | } | |
138 | } | |
139 | free(params); | |
140 | ||
141 | if (!result) { | |
142 | luaL_error(L, "Unable to allocate result handle"); | |
143 | lua_pushboolean(L, 0); | |
144 | return 1; | |
145 | } | |
146 | ||
147 | status = PQresultStatus(result); | |
148 | if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK) { | |
149 | luaL_error(L, "Unable to execute statment: %s", PQresultErrorMessage(result)); | |
150 | lua_pushboolean(L, 0); | |
151 | return 1; | |
152 | } | |
153 | ||
154 | statement->result = result; | |
155 | ||
156 | lua_pushboolean(L, 1); | |
157 | return 1; | |
158 | } | |
159 | ||
160 | static int statement_fetch_impl(lua_State *L, int named_columns) { | |
161 | statement_t *statement = (statement_t *)luaL_checkudata(L, 1, DBD_POSTGRESQL_STATEMENT); | |
162 | int tuple = statement->tuple++; | |
163 | int i; | |
164 | int num_columns; | |
165 | ||
166 | if (PQresultStatus(statement->result) != PGRES_TUPLES_OK) { | |
167 | lua_pushnil(L); | |
168 | return 1; | |
169 | } | |
170 | ||
171 | if (tuple >= PQntuples(statement->result)) { | |
172 | lua_pushnil(L); /* no more results */ | |
173 | return 1; | |
174 | } | |
175 | ||
176 | num_columns = PQnfields(statement->result); | |
177 | lua_newtable(L); | |
178 | for (i = 0; i < num_columns; i++) { | |
179 | int d = 1; | |
180 | const char *name = PQfname(statement->result, i); | |
181 | ||
182 | if (PQgetisnull(statement->result, tuple, i)) { | |
183 | if (named_columns) { | |
184 | LUA_PUSH_ATTRIB_NIL(name); | |
185 | } else { | |
186 | LUA_PUSH_ARRAY_NIL(d); | |
187 | } | |
188 | } else { | |
189 | const char *value = PQgetvalue(statement->result, tuple, i); | |
190 | lua_push_type_t lua_push = postgresql_to_lua_push(PQftype(statement->result, i)); | |
191 | ||
192 | if (lua_push == LUA_PUSH_NIL) { | |
193 | if (named_columns) { | |
194 | LUA_PUSH_ATTRIB_NIL(name); | |
195 | } else { | |
196 | LUA_PUSH_ARRAY_NIL(d); | |
197 | } | |
198 | } else if (lua_push == LUA_PUSH_INTEGER) { | |
199 | int val = atoi(value); | |
200 | ||
201 | if (named_columns) { | |
202 | LUA_PUSH_ATTRIB_INT(name, val); | |
203 | } else { | |
204 | LUA_PUSH_ARRAY_INT(d, val); | |
205 | } | |
206 | } else if (lua_push == LUA_PUSH_NUMBER) { | |
207 | double val = strtod(value, NULL); | |
208 | ||
209 | if (named_columns) { | |
210 | LUA_PUSH_ATTRIB_FLOAT(name, val); | |
211 | } else { | |
212 | LUA_PUSH_ARRAY_FLOAT(d, val); | |
213 | } | |
214 | } else if (lua_push == LUA_PUSH_STRING) { | |
215 | if (named_columns) { | |
216 | LUA_PUSH_ATTRIB_STRING(name, value); | |
217 | } else { | |
218 | LUA_PUSH_ARRAY_STRING(d, value); | |
219 | } | |
220 | } else if (lua_push == LUA_PUSH_BOOLEAN) { | |
221 | int val = value[0] == 't' ? 1 : 0; | |
222 | ||
223 | if (named_columns) { | |
224 | LUA_PUSH_ATTRIB_BOOL(name, val); | |
225 | } else { | |
226 | LUA_PUSH_ARRAY_BOOL(d, val); | |
227 | } | |
228 | } else { | |
229 | luaL_error(L, "Unknown push type in result set"); | |
230 | } | |
231 | } | |
232 | } | |
233 | ||
234 | return 1; | |
235 | } | |
236 | ||
237 | ||
238 | static int statement_fetch(lua_State *L) { | |
239 | return statement_fetch_impl(L, 0); | |
240 | } | |
241 | ||
242 | static int statement_fetchtable(lua_State *L) { | |
243 | return statement_fetch_impl(L, 1); | |
244 | } | |
245 | ||
246 | static int statement_gc(lua_State *L) { | |
247 | /* always free the handle */ | |
248 | statement_close(L); | |
249 | ||
250 | return 0; | |
251 | } | |
252 | ||
253 | ||
254 | static const luaL_Reg statement_methods[] = { | |
255 | {"close", statement_close}, | |
256 | {"execute", statement_execute}, | |
257 | {"fetch", statement_fetch}, | |
258 | {"fetchtable", statement_fetchtable}, | |
259 | {NULL, NULL} | |
260 | }; | |
261 | ||
262 | static const luaL_Reg statement_class_methods[] = { | |
263 | {NULL, NULL} | |
264 | }; | |
265 | ||
266 | int dbd_postgresql_statement_create(lua_State *L, connection_t *conn, const char *sql_query) { | |
267 | statement_t *statement = NULL; | |
268 | ExecStatusType status; | |
269 | PGresult *result = NULL; | |
270 | char *new_sql; | |
271 | char name[IDLEN]; | |
272 | ||
273 | new_sql = replace_placeholders(L, sql_query); | |
274 | ||
275 | snprintf(name, IDLEN, "%017u", ++conn->statement_id); | |
276 | ||
277 | result = PQprepare(conn->postgresql, name, new_sql, 0, NULL); | |
278 | free(new_sql); | |
279 | ||
280 | if (!result) { | |
281 | luaL_error(L, "Unable to allocate prepare result handle"); | |
282 | } | |
283 | ||
284 | status = PQresultStatus(result); | |
285 | if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK) { | |
286 | const char *err_string = PQresultErrorMessage(result); | |
287 | PQclear(result); | |
288 | luaL_error(L, "Unable to prepare statment: %s", err_string); | |
289 | return 0; | |
290 | } | |
291 | ||
292 | PQclear(result); | |
293 | ||
294 | statement = (statement_t *)lua_newuserdata(L, sizeof(statement_t)); | |
295 | statement->postgresql = conn->postgresql; | |
296 | statement->result = NULL; | |
297 | statement->tuple = 0; | |
298 | strncpy(statement->name, name, IDLEN-1); | |
299 | statement->name[IDLEN-1] = '\0'; | |
300 | ||
301 | luaL_getmetatable(L, DBD_POSTGRESQL_STATEMENT); | |
302 | lua_setmetatable(L, -2); | |
303 | ||
304 | return 1; | |
305 | } | |
306 | ||
307 | int dbd_postgresql_statement(lua_State *L) { | |
308 | luaL_newmetatable(L, DBD_POSTGRESQL_STATEMENT); | |
309 | luaL_register(L, 0, statement_methods); | |
310 | lua_pushvalue(L,-1); | |
311 | lua_setfield(L, -2, "__index"); | |
312 | ||
313 | lua_pushcfunction(L, statement_gc); | |
314 | lua_setfield(L, -2, "__gc"); | |
315 | ||
316 | luaL_register(L, DBD_POSTGRESQL_STATEMENT, statement_class_methods); | |
317 | ||
318 | return 1; | |
319 | } |