storage_pgsql.c

来自「这是一个完全开放的」· C语言 代码 · 共 639 行 · 第 1/2 页

/*
 * jabberd - Jabber Open Source Server
 * Copyright (c) 2002-2003 Jeremie Miller, Thomas Muldowney,
 *                         Ryan Eatmon, Robert Norris
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA02111-1307USA
 */

/** @file sm/storage_pgsql.c
  * @brief postgresql storage module
  * @author Robert Norris
  * $Date: 2004/11/13 16:08:33 $
  * $Revision: 1.18.2.6 $
  */

#include "sm.h"

#ifdef STORAGE_PGSQL

#include <libpq-fe.h>

/** internal structure, holds our data */
typedef struct drvdata_st {
    PGconn *conn;

    char *prefix;

    xht filters;
    
    int txn;
} *drvdata_t;

#define FALLBACK_BLOCKSIZE (4096)

/** internal: do and return the math and ensure it gets realloc'd */
static size_t _st_pgsql_realloc(void **oblocks, size_t len) {
    void *nblocks;
    size_t nlen;
    static size_t block_size = 0;

    if (block_size == 0) {
#ifdef HAVE_GETPAGESIZE
        block_size = getpagesize();
#elif defined(_SC_PAGESIZE)
        block_size = sysconf(_SC_PAGESIZE);
#elif defined(_SC_PAGE_SIZE)
        block_size = sysconf(_SC_PAGE_SIZE);    
#else
        block_size = FALLBACK_BLOCKSIZE;
#endif
    }
    /* round up to standard block sizes */
    nlen = (((len-1)/block_size)+1)*block_size;

    /* keep trying till we get it */
    while((nblocks = realloc(*oblocks, nlen)) == NULL) sleep(1);
    *oblocks = nblocks;
    return nlen;
}

/** this is the safety check used to make sure there's always enough mem */
#define PGSQL_SAFE(blocks, size, len) if((size) >= len) len = _st_pgsql_realloc((void**)&(blocks),(size + 1));

static void _st_pgsql_convert_filter_recursive(st_driver_t drv, st_filter_t f, char **buf, int *buflen, int *nbuf) {
    drvdata_t data = (drvdata_t) drv->private;
    st_filter_t scan;
    int vlen;
    char *cval;

    switch(f->type) {
        case st_filter_type_PAIR:
            /* do sql escaping for apostrophes */
            cval = (char *) malloc(sizeof(char) * ((strlen(f->val) * 2) + 1));
            vlen = PQescapeString(cval, f->val, strlen(f->val));

            PGSQL_SAFE((*buf), *buflen + 12 + vlen - strlen(f->val), *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), "( \"%s\" = \'%s\' ) ", f->key, f->val);
            free(cval);

            break;

        case st_filter_type_AND:
            PGSQL_SAFE((*buf), *buflen + 2, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), "( ");

            for(scan = f->sub; scan != NULL; scan = scan->next) {
                _st_pgsql_convert_filter_recursive(drv, scan, buf, buflen, nbuf);

                if(scan->next != NULL) {
                    PGSQL_SAFE((*buf), *buflen + 4, *buflen);
                    *nbuf += sprintf(&((*buf)[*nbuf]), "AND ");
                }
            }

            PGSQL_SAFE((*buf), *buflen + 2, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), ") ");

            return;

        case st_filter_type_OR:
            PGSQL_SAFE((*buf), *buflen + 2, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), "( ");

            for(scan = f->sub; scan != NULL; scan = scan->next) {
                _st_pgsql_convert_filter_recursive(drv, scan, buf, buflen, nbuf);

                if(scan->next != NULL) {
                    PGSQL_SAFE((*buf), *buflen + 3, *buflen);
                    *nbuf += sprintf(&((*buf)[*nbuf]), "OR ");
                }
            }

            PGSQL_SAFE((*buf), *buflen + 2, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), ") ");

            return;

        case st_filter_type_NOT:
            PGSQL_SAFE((*buf), *buflen + 6, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), "( NOT ");

            _st_pgsql_convert_filter_recursive(drv, f->sub, buf, buflen, nbuf);

            PGSQL_SAFE((*buf), *buflen + 2, *buflen);
            *nbuf += sprintf(&((*buf)[*nbuf]), ") ");

            return;
    }
}

static char *_st_pgsql_convert_filter(st_driver_t drv, const char *owner, const char *filter) {
    drvdata_t data = (drvdata_t) drv->private;
    char *buf = NULL, *sbuf = NULL, *cfilter;
    int buflen = 0, nbuf = 0, fbuf;
    st_filter_t f;

    PGSQL_SAFE(buf, 24 + strlen(owner), buflen);

    nbuf = sprintf(buf, "\"collection-owner\" = '%s'", owner);

    sbuf = xhash_get(data->filters, filter);
    if(sbuf != NULL) {
        PGSQL_SAFE(buf, buflen + strlen(sbuf) + 7, buflen);
        nbuf += sprintf(&buf[nbuf], " AND %s", sbuf);
        return buf;
    }

    cfilter = pstrdup(xhash_pool(data->filters), filter);

    f = storage_filter(filter);
    if(f == NULL)
        return buf;

    PGSQL_SAFE(buf, buflen + 5, buflen);
    nbuf += sprintf(&buf[nbuf], " AND ");

    fbuf = nbuf;

    _st_pgsql_convert_filter_recursive(drv, f, &buf, &buflen, &nbuf);

    xhash_put(data->filters, cfilter, pstrdup(xhash_pool(data->filters), &buf[fbuf]));

    pool_free(f->p);

    return buf;
}

static st_ret_t _st_pgsql_add_type(st_driver_t drv, const char *type) {
    return st_SUCCESS;
}

static st_ret_t _st_pgsql_put_guts(st_driver_t drv, const char *type, const char *owner, os_t os) {
    drvdata_t data = (drvdata_t) drv->private;
    char *left = NULL, *right = NULL;
    int lleft = 0, lright = 0, nleft, nright;
    os_object_t o;
    char *key, *cval = NULL;
    void *val;
    int vlen;
    os_type_t ot;
    char *xml;
    int xlen;
    PGresult *res;
    char tbuf[128];

    if(os_count(os) == 0)
        return st_SUCCESS;

    if(data->prefix != NULL) {
        snprintf(tbuf, sizeof(tbuf), "%s%s", data->prefix, type);
        type = tbuf;
    }

    if(os_iter_first(os))
        do {
            PGSQL_SAFE(left, strlen(type) + 55, lleft);
            nleft = sprintf(left, "INSERT INTO \"%s\" ( \"collection-owner\", \"object-sequence\"", type);

            PGSQL_SAFE(right, strlen(owner) + 43, lright);
            nright = sprintf(right, " ) VALUES ( '%s', nextval('object-sequence')", owner);

            o = os_iter_object(os);
            if(os_object_iter_first(o))
                do {
                    os_object_iter_get(o, &key, &val, &ot);

                    switch(ot) {
                        case os_type_BOOLEAN:
                            cval = val ? strdup("t") : strdup("f");
                            vlen = 1;
                            break;

                        case os_type_INTEGER:
                            cval = (char *) malloc(sizeof(char) * 20);
                            sprintf(cval, "%d", (int) val);
                            vlen = strlen(cval);
                            break;

                        case os_type_STRING:
                            cval = (char *) malloc(sizeof(char) * ((strlen((char *) val) * 2) + 1));
                            vlen = PQescapeString(cval, (char *) val, strlen((char *) val));
                            break;

                        /* !!! might not be a good idea to mark nads this way */
                        case os_type_NAD:
                            nad_print((nad_t) val, 0, &xml, &xlen);
                            cval = (char *) malloc(sizeof(char) * ((xlen * 2) + 4));
                            vlen = PQescapeString(&cval[3], xml, xlen) + 3;
                            strncpy(cval, "NAD", 3);
                            break;
                    }

                    log_debug(ZONE, "key %s val %s", key, cval);

                    PGSQL_SAFE(left, lleft + strlen(key) + 4, lleft);
                    nleft += sprintf(&left[nleft], ", \"%s\"", key);

                    PGSQL_SAFE(right, lright + strlen(cval) + 4, lright);
                    nright += sprintf(&right[nright], ", '%s'", cval);

                    free(cval);
                } while(os_object_iter_next(o));

            PGSQL_SAFE(left, lleft + strlen(right) + 3, lleft);
            sprintf(&left[nleft], "%s );", right);
    
            log_debug(ZONE, "prepared sql: %s", left);

            res = PQexec(data->conn, left);

            if(PQresultStatus(res) != PGRES_COMMAND_OK && PQstatus(data->conn) != CONNECTION_OK) {
                log_write(drv->st->sm->log, LOG_ERR, "pgsql: lost connection to database, attempting reconnect");
                PQclear(res);
                PQreset(data->conn);
                res = PQexec(data->conn, left);
            }
            if(PQresultStatus(res) != PGRES_COMMAND_OK) {
                log_write(drv->st->sm->log, LOG_ERR, "pgsql: sql insert failed: %s", PQresultErrorMessage(res));
                free(left);
                free(right);
                PQclear(res);
                return st_FAILED;
            }

            PQclear(res);

        } while(os_iter_next(os));

    free(left);
    free(right);

    return st_SUCCESS;
}

static st_ret_t _st_pgsql_put(st_driver_t drv, const char *type, const char *owner, os_t os) {
    drvdata_t data = (drvdata_t) drv->private;
    PGresult *res;

    if(os_count(os) == 0)
        return st_SUCCESS;

    if(data->txn) {
        res = PQexec(data->conn, "BEGIN;");
        if(PQresultStatus(res) != PGRES_COMMAND_OK && PQstatus(data->conn) != CONNECTION_OK) {
            log_write(drv->st->sm->log, LOG_ERR, "pgsql: lost connection to database, attempting reconnect");
            PQclear(res);
            PQreset(data->conn);
            res = PQexec(data->conn, "BEGIN;");
        }
        if(PQresultStatus(res) != PGRES_COMMAND_OK) {
            log_write(drv->st->sm->log, LOG_ERR, "pgsql: sql transaction begin failed: %s", PQresultErrorMessage(res));
            PQclear(res);
            return st_FAILED;
        }
        PQclear(res);

        res = PQexec(data->conn, "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;");
        if(PQresultStatus(res) != PGRES_COMMAND_OK && PQstatus(data->conn) != CONNECTION_OK) {
            log_write(drv->st->sm->log, LOG_ERR, "pgsql: lost connection to database, attempting reconnect");
            PQclear(res);
            PQreset(data->conn);
            res = PQexec(data->conn, "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;");
        }
        if(PQresultStatus(res) != PGRES_COMMAND_OK) {
            log_write(drv->st->sm->log, LOG_ERR, "pgsql: sql transaction setup failed: %s", PQresultErrorMessage(res));
            PQclear(res);
            PQclear(PQexec(data->conn, "ROLLBACK;"));
            return st_FAILED;