replaced libpq pointers with safe unique ptrs

Author: dankmolot

source/async_postgres.hpp

@@ -1,5 +1,4 @@
 #pragma once
-
 #include <GarrysMod/Lua/AutoReference.h>
 #include <GarrysMod/Lua/Interface.h>
 #include <GarrysMod/Lua/LuaInterface.h>
@@ -15,6 +14,8 @@
 #include <variant>
 #include <vector>
 
+#include "safe_pg.hpp"
+
 namespace GLua = GarrysMod::Lua;
 
 #define lua_interface()                                 \
@@ -68,17 +69,15 @@ namespace async_postgres {
         PostgresPollingStatusType status = PGRES_POLLING_WRITING;
     };
 
-    using PGconnPtr = std::unique_ptr<PGconn, decltype(&PQfinish)>;
-
     struct Connection {
-        PGconnPtr conn;
+        pg::conn conn;
         GLua::AutoReference lua_table;
         std::queue<Query> queries;
         std::optional<ResetEvent> reset_event;
         bool receive_notifications =
             false;  // enabled if on_notify lua field is set
 
-        Connection(GLua::ILuaInterface* lua, PGconnPtr&& conn);
+        Connection(GLua::ILuaInterface* lua, pg::conn&& conn);
         ~Connection();
     };
 

source/connection.cpp

@@ -4,7 +4,7 @@ using namespace async_postgres;
 
 std::vector<Connection*> async_postgres::connections = {};
 
-Connection::Connection(GLua::ILuaInterface* lua, PGconnPtr&& conn)
+Connection::Connection(GLua::ILuaInterface* lua, pg::conn&& conn)
     : conn(std::move(conn)) {
     lua->CreateTable();
     this->lua_table = GLua::AutoReference(lua);
@@ -20,7 +20,7 @@ Connection::~Connection() {
 }
 
 struct ConnectionEvent {
-    PGconnPtr conn;
+    pg::conn conn;
     GLua::AutoReference callback;
     PostgresPollingStatusType status = PGRES_POLLING_WRITING;
     bool is_reset = false;
@@ -40,17 +40,16 @@ inline bool socket_is_ready(PGconn* conn, PostgresPollingStatusType status) {
 
 void async_postgres::connect(GLua::ILuaInterface* lua, std::string_view url,
                              GLua::AutoReference&& callback) {
-    auto conn = PGconnPtr(PQconnectStart(url.data()), &PQfinish);
-    auto conn_ptr = conn.get();
+    auto conn = pg::connectStart(url);
 
     if (!conn) {
         // funnily enough, this probably will instead throw a std::bad_alloc
         throw std::runtime_error("failed to allocate connection");
     }
 
-    if (PQstatus(conn_ptr) == CONNECTION_BAD ||
-        PQsetnonblocking(conn_ptr, 1) != 0) {
-        throw std::runtime_error(PQerrorMessage(conn_ptr));
+    if (PQstatus(conn.get()) == CONNECTION_BAD ||
+        PQsetnonblocking(conn.get(), 1) != 0) {
+        throw std::runtime_error(PQerrorMessage(conn.get()));
     }
 
     auto event = ConnectionEvent{std::move(conn), std::move(callback)};
@@ -65,8 +64,6 @@ inline bool poll_pending_connection(GLua::ILuaInterface* lua,
         return false;
     }
 
-    // TODO: handle reset
-
     event.status = PQconnectPoll(event.conn.get());
     if (event.status == PGRES_POLLING_OK) {
         auto state = new Connection(lua, std::move(event.conn));

source/notifications.cpp

@@ -33,16 +33,13 @@ void async_postgres::process_notifications(GLua::ILuaInterface* lua,
         return;
     }
 
-    PGnotify* notify;
-    while ((notify = PQnotifies(state->conn.get()))) {
+    while (auto notify = pg::getNotify(state->conn)) {
         if (push_on_notify(lua, state)) {
             lua->PushString(notify->relname);  // arg 1 channel name
             lua->PushString(notify->extra);    // arg 2 payload
             lua->PushNumber(notify->be_pid);   // arg 3 backend pid
 
             pcall(lua, 3, 0);
         }
-
-        PQfreemem(notify);
     }
 }

source/query.cpp

@@ -99,7 +99,7 @@ void async_postgres::process_queries(GLua::ILuaInterface* lua,
     }
 
     while (PQisBusy(state->conn.get()) == 0) {
-        auto* result = PQgetResult(state->conn.get());
+        auto result = pg::getResult(state->conn);
         if (!result) {
             // query is done
             state->queries.pop();
@@ -108,16 +108,15 @@ void async_postgres::process_queries(GLua::ILuaInterface* lua,
 
         if (query.callback) {
             query.callback.Push();
-            if (!bad_result(result)) {
+            if (!bad_result(result.get())) {
                 lua->PushBool(true);
-                create_result_table(lua, result);
+                create_result_table(lua, result.get());
             } else {
                 lua->PushBool(false);
-                lua->PushString(PQresultErrorMessage(result));
+                lua->PushString(PQresultErrorMessage(result.get()));
             }
 
             pcall(lua, 2, 0);
         }
-        PQclear(result);
     }
 }

source/safe_pg.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <libpq-fe.h>
+
+#include <memory>
+#include <string_view>
+
+/// Just small wrappers around libpq types to not have to deal with freeing them
+/// manually.
+namespace async_postgres::pg {
+    using conn = std::unique_ptr<PGconn, decltype(&PQfinish)>;
+    using result = std::unique_ptr<PGresult, decltype(&PQclear)>;
+    using notify = std::unique_ptr<PGnotify, decltype(&PQfreemem)>;
+
+    inline conn connectStart(std::string_view conninfo) {
+        return conn(PQconnectStart(conninfo.data()), &PQfinish);
+    }
+
+    inline result getResult(conn& conn) {
+        return result(PQgetResult(conn.get()), &PQclear);
+    }
+
+    inline notify getNotify(conn& conn) {
+        return notify(PQnotifies(conn.get()), &PQfreemem);
+    }
+}  // namespace async_postgres::pg