pqsql/pqsql.c

   1 /*
   2  * PQSQL module
   3  *
   4  * 99% of the handling is stolen from Q9.
   5  */
   6
   7 #include "../core/config.h"
   8 #include "../core/error.h"
   9 #include "../irc/irc_config.h"
  10 #include "../core/events.h"
  11 #include "../core/hooks.h"
  12 #include "../lib/irc_string.h"
  13 #include "../lib/version.h"
  14 #include "pqsql.h"
  15
  16 #include <stdlib.h>
  17 #include <sys/poll.h>
  18 #include <stdarg.h>
  19 #include <string.h>
  20
  21 MODULE_VERSION("");
  22
  23 /* It's possible that we might want to do a very long query, longer than the
  24  * IRC-oriented SSTRING_MAX value.  One option would be to increase
  25  * SSTRING_MAX, but the whole purpose of sstring's is to efficiently deal
  26  * with situations where the malloc() padding overhead is large compared to
  27  * string length and strings are frequently recycled.  Since neither of
  28  * these are necessarily true for longer strings it makes more sense to use
  29  * malloc() for them.
  30  *
  31  * So, query always points at the query string.  If it fitted in a sstring,
  32  * query_ss will point at the sstring for freeing purposes.  If query_ss is
  33  * NULL then it was malloc'd so should be free()'d directly.
  34  */
  35 typedef struct pqasyncquery_s {
  36   sstring *query_ss;
  37   char *query;
  38   void *tag;
  39   PQQueryHandler handler;
  40   int flags;
  41   struct pqasyncquery_s *next;
  42 } pqasyncquery_s;
  43
  44 typedef struct pqtableloaderinfo_s
  45 {
  46     sstring *tablename;
  47     PQQueryHandler init, data, fini;
  48 } pqtableloaderinfo_s;
  49
  50 pqasyncquery_s *queryhead = NULL, *querytail = NULL;
  51
  52 int dbconnected = 0;
  53 PGconn *dbconn;
  54
  55 void dbhandler(int fd, short revents);
  56 void pqstartloadtable(PGconn *dbconn, void *arg);
  57 void dbstatus(int hooknum, void *arg);
  58 void disconnectdb(void);
  59 void connectdb(void);
  60
  61 void _init(void) {
  62   connectdb();
  63 }
  64
  65 void _fini(void) {
  66   disconnectdb();
  67 }
  68
  69 void connectdb(void) {
  70   sstring *dbhost, *dbusername, *dbpassword, *dbdatabase, *dbport;
  71   char connectstr[1024];
  72
  73   if(pqconnected())
  74     return;
  75
  76   /* stolen from chanserv as I'm lazy */
  77   dbhost = getcopyconfigitem("pqsql", "host", "localhost", HOSTLEN);
  78   dbusername = getcopyconfigitem("pqsql", "username", "newserv", 20);
  79   dbpassword = getcopyconfigitem("pqsql", "password", "moo", 20);
  80   dbdatabase = getcopyconfigitem("pqsql", "database", "newserv", 20);
  81   dbport = getcopyconfigitem("pqsql", "port", "431", 8);
  82
  83   if(!dbhost || !dbusername || !dbpassword || !dbdatabase || !dbport) {
  84     /* freesstring allows NULL */
  85     freesstring(dbhost);
  86     freesstring(dbusername);
  87     freesstring(dbpassword);
  88     freesstring(dbdatabase);
  89     freesstring(dbport);
  90     return;
  91   }
  92
  93   snprintf(connectstr, sizeof(connectstr), "host=%s port=%s dbname=%s user=%s password=%s", dbhost->content, dbport->content, dbdatabase->content, dbusername->content, dbpassword->content);
  94
  95   freesstring(dbhost);
  96   freesstring(dbusername);
  97   freesstring(dbpassword);
  98   freesstring(dbdatabase);
  99   freesstring(dbport);
 100
 101   Error("pqsql", ERR_INFO, "Attempting database connection: %s", connectstr);
 102
 103   /* Blocking connect for now.. */
 104   dbconn = PQconnectdb(connectstr);
 105
 106   if (!dbconn || (PQstatus(dbconn) != CONNECTION_OK)) {
 107     Error("pqsql", ERR_ERROR, "Unable to connect to db: %s", PQerrorMessage(dbconn));
 108     return;
 109   }
 110   Error("pqsql", ERR_INFO, "Connected!");
 111
 112   dbconnected = 1;
 113
 114   PQsetnonblocking(dbconn, 1);
 115
 116   /* this kicks ass, thanks splidge! */
 117   registerhandler(PQsocket(dbconn), POLLIN, dbhandler);
 118   registerhook(HOOK_CORE_STATSREQUEST, dbstatus);
 119 }
 120
 121 void dbhandler(int fd, short revents) {
 122   PGresult *res;
 123   pqasyncquery_s *qqp;
 124
 125   if(revents & POLLIN) {
 126     PQconsumeInput(dbconn);
 127
 128     if(!PQisBusy(dbconn)) { /* query is complete */
 129       if(queryhead->handler)
 130         (queryhead->handler)(dbconn, queryhead->tag);
 131
 132       while((res = PQgetResult(dbconn))) {
 133         switch(PQresultStatus(res)) {
 134           case PGRES_TUPLES_OK:
 135             Error("pqsql", ERR_WARNING, "Unhandled tuples output (query: %s)", queryhead->query);
 136             break;
 137
 138           case PGRES_NONFATAL_ERROR:
 139           case PGRES_FATAL_ERROR:
 140             /* if a create query returns an error assume it went ok, paul will winge about this */
 141             if(!(queryhead->flags & QH_CREATE))
 142               Error("pqsql", ERR_WARNING, "Unhandled error response (query: %s)", queryhead->query);
 143             break;
 144
 145           default:
 146             break;
 147         }
 148
 149         PQclear(res);
 150       }
 151
 152       /* Free the query and advance */
 153       qqp = queryhead;
 154       if(queryhead == querytail)
 155         querytail = NULL;
 156
 157       queryhead = queryhead->next;
 158
 159       if (qqp->query_ss) {
 160         freesstring(qqp->query_ss);
 161         qqp->query_ss=NULL;
 162         qqp->query=NULL;
 163       } else if (qqp->query) {
 164         free(qqp->query);
 165         qqp->query=NULL;
 166       }
 167       free(qqp);
 168
 169       if(queryhead) { /* Submit the next query */
 170         PQsendQuery(dbconn, queryhead->query);
 171         PQflush(dbconn);
 172       }
 173     }
 174   }
 175 }
 176
 177 /* sorry Q9 */
 178 void pqasyncqueryf(PQQueryHandler handler, void *tag, int flags, char *format, ...) {
 179   char querybuf[8192];
 180   va_list va;
 181   int len;
 182   pqasyncquery_s *qp;
 183
 184   if(!pqconnected())
 185     return;
 186
 187   va_start(va, format);
 188   len = vsnprintf(querybuf, sizeof(querybuf), format, va);
 189   va_end(va);
 190
 191   /* PPA: no check here... */
 192   qp = (pqasyncquery_s *)malloc(sizeof(pqasyncquery_s));
 193
 194   if(!qp)
 195     Error("pqsql",ERR_STOP,"malloc() failed in pqsql.c");
 196
 197   /* Use sstring or allocate (see above rant) */
 198   if (len > SSTRING_MAX) {
 199     qp->query = (char *)malloc(len+1);
 200     strcpy(qp->query,querybuf);
 201     qp->query_ss=NULL;
 202   } else {
 203     qp->query_ss = getsstring(querybuf, len);
 204     qp->query = qp->query_ss->content;
 205   }
 206   qp->tag = tag;
 207   qp->handler = handler;
 208   qp->next = NULL; /* shove them at the end */
 209   qp->flags = flags;
 210
 211   if(querytail) {
 212     querytail->next = qp;
 213     querytail = qp;
 214   } else {
 215     querytail = queryhead = qp;
 216     PQsendQuery(dbconn, qp->query);
 217     PQflush(dbconn);
 218   }
 219 }
 220
 221 void pqloadtable(char *tablename, PQQueryHandler init, PQQueryHandler data, PQQueryHandler fini)
 222 {
 223   pqtableloaderinfo_s *tli;
 224
 225   tli=(pqtableloaderinfo_s *)malloc(sizeof(pqtableloaderinfo_s));
 226   tli->tablename=getsstring(tablename, 100);
 227   tli->init=init;
 228   tli->data=data;
 229   tli->fini=fini;
 230   pqasyncquery(pqstartloadtable, tli, "SELECT COUNT(*) FROM %s", tli->tablename->content);
 231 }
 232
 233 void pqstartloadtable(PGconn *dbconn, void *arg)
 234 {
 235   PGresult *res;
 236   unsigned long i, count, tablecrc;
 237   pqtableloaderinfo_s *tli = arg;
 238
 239   res = PQgetResult(dbconn);
 240
 241   if (PQresultStatus(res) != PGRES_TUPLES_OK && PQresultStatus(res) != PGRES_COMMAND_OK) {
 242     Error("pqsql", ERR_ERROR, "Error getting row count for %s.", tli->tablename->content);
 243     return;
 244   }
 245
 246   if (PQnfields(res) != 1) {
 247     Error("pqsql", ERR_ERROR, "Count query format error for %s.", tli->tablename->content);
 248     return;
 249   }
 250
 251   tablecrc=crc32(tli->tablename->content);
 252   count=strtoul(PQgetvalue(res, 0, 0), NULL, 10);
 253   PQclear(res);
 254
 255   Error("pqsql", ERR_INFO, "Found %lu entries in table %s, scheduling load.", count, tli->tablename->content);
 256
 257   pqasyncquery(tli->init, NULL, "BEGIN");
 258   pqasyncquery(NULL, NULL, "DECLARE table%lx%lx CURSOR FOR SELECT * FROM %s", tablecrc, count, tli->tablename->content);
 259
 260   for (i=0;(count - i) > 1000; i+=1000)
 261     pqasyncquery(tli->data, NULL, "FETCH 1000 FROM table%lx%lx", tablecrc, count);
 262
 263   pqasyncquery(tli->data, NULL, "FETCH ALL FROM table%lx%lx", tablecrc, count);
 264
 265   pqasyncquery(NULL, NULL, "CLOSE table%lx%lx", tablecrc, count);
 266   pqasyncquery(tli->fini, NULL, "COMMIT");
 267
 268   freesstring(tli->tablename);
 269   free(tli);
 270 }
 271
 272 void disconnectdb(void) {
 273   pqasyncquery_s *qqp = queryhead, *nqqp;
 274
 275   if(!pqconnected())
 276     return;
 277
 278   /* do this first else we may get conflicts */
 279   deregisterhandler(PQsocket(dbconn), 0);
 280
 281   /* Throw all the queued queries away, beware of data malloc()ed inside the query item.. */
 282   while(qqp) {
 283     nqqp = qqp->next;
 284     if (qqp->query_ss) {
 285       freesstring(qqp->query_ss);
 286       qqp->query_ss=NULL;
 287       qqp->query=NULL;
 288     } else if (qqp->query) {
 289       free(qqp->query);
 290       qqp->query=NULL;
 291     }
 292     free(qqp);
 293     qqp = nqqp;
 294   }
 295
 296   deregisterhook(HOOK_CORE_STATSREQUEST, dbstatus);
 297   PQfinish(dbconn);
 298   dbconn = NULL; /* hmm? */
 299
 300   dbconnected = 0;
 301 }
 302
 303 /* more stolen code from Q9 */
 304 void dbstatus(int hooknum, void *arg) {
 305   if ((long)arg > 10) {
 306     int i = 0;
 307     pqasyncquery_s *qqp;
 308     char message[100];
 309
 310     if(queryhead)
 311       for(qqp=queryhead;qqp;qqp=qqp->next)
 312         i++;
 313
 314     snprintf(message, sizeof(message), "PQSQL   : %6d queries queued.",i);
 315
 316     triggerhook(HOOK_CORE_STATSREPLY, message);
 317   }
 318 }
 319
 320 int pqconnected(void) {
 321   return dbconnected;
 322 }