Added prepare and queryPrepared functions

Author: dankmolot

source/async_postgres.hpp

@@ -75,8 +75,22 @@ namespace async_postgres {
         ParamValues param;
     };
 
+    struct CreatePreparedCommand {
+        std::string name;
+        std::string command;
+    };
+
+    struct PreparedCommand {
+        std::string name;
+        ParamValues param;
+    };
+
     struct Query {
-        std::variant<SimpleCommand, ParameterizedCommand> command;
+        using CommandVariant =
+            std::variant<SimpleCommand, ParameterizedCommand,
+                         CreatePreparedCommand, PreparedCommand>;
+
+        CommandVariant command;
         GLua::AutoReference callback;
         bool sent = false;
         bool flushed = false;

source/main.cpp

@@ -84,7 +84,54 @@ namespace async_postgres::lua {
         }
 
         state->query = std::move(query);
+        return 0;
+    }
+
+    lua_protected_fn(prepare) {
+        lua->CheckType(1, async_postgres::connection_meta);
+        lua->CheckType(2, GLua::Type::String);
+        lua->CheckType(3, GLua::Type::String);
+        lua->CheckType(4, GLua::Type::Function);
+
+        auto state = lua_connection_state();
+        if (state->query) {
+            throw std::runtime_error("query already in progress");
+        }
+
+        async_postgres::CreatePreparedCommand command = {lua->GetString(2),
+                                                         lua->GetString(3)};
+        async_postgres::Query query = {std::move(command)};
+
+        if (lua->IsType(4, GLua::Type::Function)) {
+            query.callback = GLua::AutoReference(lua, 4);
+        }
+
+        state->query = std::move(query);
+        return 0;
+    }
 
+    lua_protected_fn(queryPrepared) {
+        lua->CheckType(1, async_postgres::connection_meta);
+        lua->CheckType(2, GLua::Type::String);
+        lua->CheckType(3, GLua::Type::Table);
+        lua->CheckType(4, GLua::Type::Function);
+
+        auto state = lua_connection_state();
+        if (state->query) {
+            throw std::runtime_error("query already in progress");
+        }
+
+        async_postgres::PreparedCommand command = {
+            lua->GetString(2),
+            async_postgres::array_to_params(lua, 3),
+        };
+        async_postgres::Query query = {std::move(command)};
+
+        if (lua->IsType(4, GLua::Type::Function)) {
+            query.callback = GLua::AutoReference(lua, 4);
+        }
+
+        state->query = std::move(query);
         return 0;
     }
 
@@ -126,6 +173,8 @@ void register_connection_mt(GLua::ILuaInterface* lua) {
     register_lua_fn(__gc);
     register_lua_fn(query);
     register_lua_fn(queryParams);
+    register_lua_fn(prepare);
+    register_lua_fn(queryPrepared);
     register_lua_fn(reset);
     register_lua_fn(setNotifyCallback);
 

source/query.cpp

@@ -2,18 +2,28 @@
 
 using namespace async_postgres;
 
+#define get_if_command(type) \
+    const auto* command = std::get_if<type>(&query.command)
+
 // returns true if query was sent
 // returns false on error
 inline bool send_query(PGconn* conn, Query& query) {
-    if (const auto* command = std::get_if<SimpleCommand>(&query.command)) {
+    if (get_if_command(SimpleCommand)) {
         return PQsendQuery(conn, command->command.c_str()) == 1;
-    } else if (const auto* command =
-                   std::get_if<ParameterizedCommand>(&query.command)) {
+    } else if (get_if_command(ParameterizedCommand)) {
         return PQsendQueryParams(conn, command->command.c_str(),
                                  command->param.length(), nullptr,
                                  command->param.values.data(),
                                  command->param.lengths.data(),
                                  command->param.formats.data(), 0) == 1;
+    } else if (get_if_command(CreatePreparedCommand)) {
+        return PQsendPrepare(conn, command->name.c_str(),
+                             command->command.c_str(), 0, nullptr) == 1;
+    } else if (get_if_command(PreparedCommand)) {
+        return PQsendQueryPrepared(
+                   conn, command->name.c_str(), command->param.length(),
+                   command->param.values.data(), command->param.lengths.data(),
+                   command->param.formats.data(), 0) == 1;
     }
     return false;
 }