Add Postgres automatic reconnection

src/drivers/pgsql/drv_pgsql.c

@@ -554,10 +554,53 @@ static db_error_t pgsql_check_status(db_conn_t *con, PGresult *pgres,
     xfree(con->sql_state);
     xfree(con->sql_errmsg);
 
-    con->sql_state = strdup(PQresultErrorField(pgres, PG_DIAG_SQLSTATE));
-    con->sql_errmsg = strdup(PQresultErrorField(pgres, PG_DIAG_MESSAGE_PRIMARY));
+    const char *sqlstate = PQresultErrorField(pgres, PG_DIAG_SQLSTATE);
+    const char *errmsg = PQresultErrorField(pgres, PG_DIAG_MESSAGE_PRIMARY);
+
+    /* For connection errors, PQresultErrorField might return NULL,
+       so check connection status first */
+    if (PQstatus(pgcon) == CONNECTION_BAD)
+    {
+      /* Connection is bad, get error from connection object */
+      con->sql_state = strdup("08006"); /* connection_failure */
+      const char *conn_errmsg = PQerrorMessage(pgcon);
+      con->sql_errmsg = conn_errmsg ? strdup(conn_errmsg) : strdup("Connection failure");
+      log_text(LOG_DEBUG, "Connection error: %s", 
+               con->sql_errmsg);
+    }
+    else
+    {
+      /* Normal query error, use result error fields */
+      con->sql_state = sqlstate ? strdup(sqlstate) : strdup("00000");
+      con->sql_errmsg = errmsg ? strdup(errmsg) : strdup("Unknown error");
+      log_text(LOG_DEBUG, "Query error: SQLSTATE=%s, error=%s", 
+               con->sql_state, con->sql_errmsg);
+    }
 
-    if (!strcmp(con->sql_state, "40P01") /* deadlock_detected */ ||
+    /* Check for connection-related errors that require reconnection */
+    if (!strcmp(con->sql_state, "08000") /* connection_exception */ ||
+        !strcmp(con->sql_state, "08003") /* connection_does_not_exist */ ||
+        !strcmp(con->sql_state, "08006") /* connection_failure */ ||
+        !strcmp(con->sql_state, "08001") /* sqlclient_unable_to_establish_sqlconnection */ ||
+        !strcmp(con->sql_state, "08004") /* sqlserver_rejected_establishment_of_sqlconnection */ ||
+        !strcmp(con->sql_state, "08007") /* transaction_resolution_unknown */ ||
+        PQstatus(pgcon) == CONNECTION_BAD)
+    {
+      log_text(LOG_DEBUG, "Connection error detected (SQLSTATE: %s), attempting reconnection", 
+               con->sql_state ? con->sql_state : "unknown");
+
+      rs->counter = SB_CNT_RECONNECT;
+
+      PQclear(pgres);
+      /* Attempt reconnection */
+      while (pgsql_drv_connect(con))
+      {
+        if (sb_globals.error)
+          return DB_ERROR_FATAL;
+      }
+      return DB_ERROR_IGNORABLE;
+    }
+    else if (!strcmp(con->sql_state, "40P01") /* deadlock_detected */ ||
         !strcmp(con->sql_state, "23505") /* unique violation */ ||
         !strcmp(con->sql_state, "40001"))/* serialization_failure */
     {
@@ -568,8 +611,10 @@ static db_error_t pgsql_check_status(db_conn_t *con, PGresult *pgres,
     }
     else
     {
-      log_text(LOG_FATAL, "%s() failed: %d %s", funcname, status,
-               con->sql_errmsg);
+      log_text(LOG_FATAL, "%s() failed: status=%d, SQLSTATE='%s', error='%s'", 
+               funcname, status, 
+               con->sql_state ? con->sql_state : "unknown",
+               con->sql_errmsg ? con->sql_errmsg : "unknown");
 
       if (query != NULL)
         log_text(LOG_FATAL, "failed query was: %s", query);

src/lua/internal/sysbench.sql.lua

@@ -256,6 +256,7 @@ function connection_methods.check_error(self, rs, query)
           sysbench.hooks.sql_error_ignorable(errdesc)) or
       self.error == sysbench.sql.error.IGNORABLE
    then
+      print("Detected ignorable error, attempting reconnect")
       -- Throw a 'restart event' exception that can be caught by the user script
       -- to do some extra steps to restart a transaction (e.g. reprepare
       -- statements after a reconnect). Otherwise it will be caught by

src/lua/oltp_common.lua

@@ -422,7 +422,6 @@ function execute_point_selects()
 
    for i = 1, sysbench.opt.point_selects do
       param[tnum].point_selects[1]:set(get_id())
-
       stmt[tnum].point_selects:execute()
    end
 end
@@ -497,15 +496,28 @@ function execute_delete_inserts()
 end
 
 -- Re-prepare statements if we have reconnected, which is possible when some of
--- the listed error codes are in the --mysql-ignore-errors list
+-- the listed error codes are in the --mysql-ignore-errors list for MySQL
+-- or when PostgreSQL connection errors occur
 function sysbench.hooks.before_restart_event(errdesc)
+   -- MySQL connection errors
    if errdesc.sql_errno == 2013 or -- CR_SERVER_LOST
       errdesc.sql_errno == 2055 or -- CR_SERVER_LOST_EXTENDED
       errdesc.sql_errno == 2006 or -- CR_SERVER_GONE_ERROR
       errdesc.sql_errno == 2011    -- CR_TCP_CONNECTION
    then
       close_statements()
       prepare_statements()
+   -- PostgreSQL connection errors (SQLSTATE codes)
+   elseif errdesc.sql_state and (
+      errdesc.sql_state == "00000" or -- connection_exception
+      errdesc.sql_state == "08003" or -- connection_does_not_exist
+      errdesc.sql_state == "08006" or -- connection_failure
+      errdesc.sql_state == "08001" or -- sqlclient_unable_to_establish_sqlconnection
+      errdesc.sql_state == "08004" or -- sqlserver_rejected_establishment_of_sqlconnection
+      errdesc.sql_state == "08007")   -- transaction_resolution_unknown
+   then
+      close_statements()
+      prepare_statements()
    end
 end
 

src/lua/oltp_read_write.lua

@@ -45,7 +45,6 @@ function event()
    if not sysbench.opt.skip_trx then
       begin()
    end
-
    execute_point_selects()
 
    if sysbench.opt.range_selects then