admin: construct query results in a pcall

diff --git a/src/admin.cpp b/src/admin.cpp
index 4b84c90..f7f2b27 100644 (file)
--- a/src/admin.cpp
+++ b/src/admin.cpp
@@ -390,36 +390,45 @@ static int fn_getluafile(lua_State *L) {
        return 1;
 }
 static int query(lua_State *L, db::Database &db) {
-       const char *sql = luaL_checkstring(L, 1);
-       auto stmt = db(sql);
-       for (int i = 2; i <= lua_gettop(L); i++) {
-               if (lua_isinteger(L, i)) {
-                       stmt.bind(i-1, lua_tointeger(L, i));
-               } else {
-                       size_t len = 0;
-                       const char *p = lua_tolstring(L, i, &len);
-                       stmt.bind(i-1, std::string_view(p, len));
-               }
-       }
-       int cols = stmt.columns();
-       int row = 0;
-       // FIXME: createtable/newtable/rawseti can fail on memory allocation
-       // if we link to lua using C ABI, this will cause a Statement leak
-       lua_newtable(L);
-       while (stmt.step()) {
-               lua_createtable(L, cols, 0);
-               for (int i = 0; i < cols; i++) {
-                       if (stmt.is_num(i)) {
-                               lua_pushinteger(L, stmt.num(i));
+       {
+               const char *sql = luaL_checkstring(L, 1);
+               auto stmt = db(sql);
+               for (int i = 2; i <= lua_gettop(L); i++) {
+                       if (lua_isinteger(L, i)) {
+                               stmt.bind(i-1, lua_tointeger(L, i));
                        } else {
-                               auto sv = stmt.text(i);
-                               sv.data() ? push(L, sv) : lua_pushnil(L);
+                               size_t len = 0;
+                               const char *p = lua_tolstring(L, i, &len);
+                               stmt.bind(i-1, std::string_view(p, len));
                        }
-                       lua_rawseti(L, -2, i+1);
                }
-               lua_rawseti(L, -2, ++row);
+               // createtable/newtable/rawseti can fail on memory allocation, so do
+               // this stuff in a pcall.
+               lua_pushcfunction(L, [](lua_State *L) {
+                       auto &stmt = *static_cast<db::Statement *>(lua_touserdata(L, 1));
+                       int cols = stmt.columns();
+                       int row = 0;
+                       lua_newtable(L);
+                       while (stmt.step()) {
+                               lua_createtable(L, cols, 0);
+                               for (int i = 0; i < cols; i++) {
+                                       if (stmt.is_num(i)) {
+                                               lua_pushinteger(L, stmt.num(i));
+                                       } else {
+                                               auto sv = stmt.text(i);
+                                               sv.data() ? push(L, sv) : lua_pushnil(L);
+                                       }
+                                       lua_rawseti(L, -2, i+1);
+                               }
+                               lua_rawseti(L, -2, ++row);
+                       }
+                       return 1;
+               });
+               lua_pushlightuserdata(L, &stmt);
+               if (lua_pcall(L, 1, 1, 0) == LUA_OK)
+                       return 1;
        }
-       return 1;
+       return lua_error(L);
 }
 static int fn_querydisk(lua_State *L) {
        return query(L, db::disk);