src/lib-sql/driver-pgsql.c

	driver-pgsql.c revision 158bd06e896df0a709e2677e3a046658ca8a2efd
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano/* Copyright (c) 2004-2009 Dovecot authors, see the included COPYING file */
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include "lib.h"
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include "array.h"
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include "ioloop.h"
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include "ioloop-internal.h" /* kind of dirty, but it should be fine.. */
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include "sql-api-private.h"
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#ifdef BUILD_PGSQL
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include <stdlib.h>
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include <time.h>
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#include <libpq-fe.h>
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano#define QUERY_TIMEOUT_SECS 6
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostruct pgsql_db {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct sql_db api;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    pool_t pool;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    char *connect_string;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    PGconn *pg;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct io *io;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    enum io_condition io_dir;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
99e4008cad9e959b683c6f48411fcf15a92be3b5Michel Normand    struct pgsql_queue *queue, **queue_tail;
99e4008cad9e959b683c6f48411fcf15a92be3b5Michel Normand    struct timeout *queue_to;
99e4008cad9e959b683c6f48411fcf15a92be3b5Michel Normand
99e4008cad9e959b683c6f48411fcf15a92be3b5Michel Normand    struct ioloop *ioloop;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct sql_result *sync_result;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    char *error;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    time_t last_connect;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int connecting:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int connected:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int querying:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int query_finished:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano};
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostruct pgsql_binary_value {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned char *value;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    size_t size;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano};
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostruct pgsql_result {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct sql_result api;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    PGresult *pgres;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    char *query;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int rownum, rows;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int fields_count;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    const char **fields;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    const char **values;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    ARRAY_DEFINE(binary_values, struct pgsql_binary_value);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    sql_query_callback_t *callback;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    void *context;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int retry_query:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano};
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostruct pgsql_queue {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct pgsql_queue *next;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    time_t created;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    char *query;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct pgsql_result *result;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano};
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostruct pgsql_transaction_context {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    struct sql_transaction_context ctx;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    sql_commit_callback_t *callback;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    void *context;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    char *first_update;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    const char *error;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int opened:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    unsigned int failed:1;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano};
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanoextern struct sql_db driver_pgsql_db;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanoextern struct sql_result driver_pgsql_result;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanotransaction_update_callback(struct sql_result *result, void *context);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanodriver_pgsql_query_full(struct sql_db *db, const char *query,
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano            sql_query_callback_t *callback, void *context,
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano            bool retry_query);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void queue_send_next(struct pgsql_db *db);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void result_finish(struct pgsql_result *result);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void driver_pgsql_close(struct pgsql_db *db)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano{
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    db->io_dir = 0;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    PQfinish(db->pg);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    db->pg = NULL;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    if (db->io != NULL) {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        /* The fd may be closed before call to PQfinish() already,
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano           so use io_remove_closed(). */
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        io_remove_closed(&db->io);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    }
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    db->connecting = FALSE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    db->connected = FALSE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        db->querying = FALSE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano}
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic const char *last_error(struct pgsql_db *db)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano{
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    const char *msg;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    size_t len;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    msg = PQerrorMessage(db->pg);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    if (msg == NULL)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        return "(no error set)";
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    /* Error message should contain trailing \n, we don't want it */
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    len = strlen(msg);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    return len == 0 || msg[len-1] != '\n' ? msg :
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        t_strndup(msg, len-1);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano}
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcanostatic void connect_callback(struct pgsql_db *db)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano{
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    enum io_condition io_dir = 0;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    int ret;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    while ((ret = PQconnectPoll(db->pg)) == PGRES_POLLING_ACTIVE)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        ;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    switch (ret) {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    case PGRES_POLLING_READING:
99e4008cad9e959b683c6f48411fcf15a92be3b5Michel Normand        io_dir = IO_READ;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        break;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    case PGRES_POLLING_WRITING:
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        io_dir = IO_WRITE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        break;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    case PGRES_POLLING_OK:
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        i_info("pgsql: Connected to %s", PQdb(db->pg));
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        db->connecting = FALSE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        db->connected = TRUE;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        break;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    case PGRES_POLLING_FAILED:
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        i_error("pgsql: Connect failed to %s: %s",
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano            PQdb(db->pg), last_error(db));
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        driver_pgsql_close(db);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        return;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    }
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    if (db->io_dir != io_dir) {
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        if (db->io != NULL)
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano            io_remove(&db->io);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        db->io = io_dir == 0 ? NULL :
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano            io_add(PQsocket(db->pg), io_dir, connect_callback, db);
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano        db->io_dir = io_dir;
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano    }
d823d5b966f49d975a09a8512d084389d6d7ffc7dlezcano
    if (db->connected && db->queue != NULL)
        queue_send_next(db);
}

static int driver_pgsql_connect(struct sql_db *_db)
{
    struct pgsql_db *db = (struct pgsql_db *)_db;
    time_t now;

    /* don't try reconnecting more than once a second */
    now = time(NULL);
    if (db->connecting || db->last_connect == now)
        return db->connected ? 1 : (db->connecting ? 0 : -1);
    db->last_connect = now;

    db->pg = PQconnectStart(db->connect_string);
    if (db->pg == NULL)
        i_fatal("pgsql: PQconnectStart() failed (out of memory)");

    if (PQstatus(db->pg) == CONNECTION_BAD) {
        i_error("pgsql: Connect failed to %s: %s",
            PQdb(db->pg), last_error(db));
        driver_pgsql_close(db);
        return -1;
    } else {
        /* nonblocking connecting begins. */
        db->io = io_add(PQsocket(db->pg), IO_WRITE,
                connect_callback, db);
        db->io_dir = IO_WRITE;
        db->connecting = TRUE;
        return 0;
    }
}

static struct sql_db *driver_pgsql_init_v(const char *connect_string)
{
    struct pgsql_db *db;

    i_assert(connect_string != NULL);

    db = i_new(struct pgsql_db, 1);
    db->connect_string = i_strdup(connect_string);
    db->api = driver_pgsql_db;
    db->queue_tail = &db->queue;
    return &db->api;
}

static void driver_pgsql_deinit_v(struct sql_db *_db)
{
    struct pgsql_db *db = (struct pgsql_db *)_db;

    while (db->queue != NULL) {
        struct pgsql_queue *next = db->queue->next;

                result_finish(db->queue->result);
        i_free(db->queue->query);
        i_free(db->queue);

        db->queue = next;
    }

    if (db->queue_to != NULL)
        timeout_remove(&db->queue_to);
        driver_pgsql_close(db);
    i_free(db->error);
    i_free(db->connect_string);
    array_free(&_db->module_contexts);
    i_free(db);
}

static enum sql_db_flags
driver_pgsql_get_flags(struct sql_db *db ATTR_UNUSED)
{
    return 0;
}

static void consume_results(struct pgsql_db *db)
{
    do {
        if (!PQconsumeInput(db->pg))
            break;

        if (PQisBusy(db->pg))
            return;
    } while (PQgetResult(db->pg) != NULL);

    io_remove(&db->io);

    db->querying = FALSE;
    if (db->queue != NULL && db->connected)
        queue_send_next(db);
}

static void driver_pgsql_result_free(struct sql_result *_result)
{
    struct pgsql_db *db = (struct pgsql_db *)_result->db;
        struct pgsql_result *result = (struct pgsql_result *)_result;

    if (result->api.callback)
        return;

    if (_result == db->sync_result)
        db->sync_result = NULL;

    if (result->pgres != NULL) {
        PQclear(result->pgres);

        /* we'll have to read the rest of the results as well */
        i_assert(db->io == NULL);
        db->io = io_add(PQsocket(db->pg), IO_READ,
                consume_results, db);
        db->io_dir = IO_READ;
        consume_results(db);
    } else {
        db->querying = FALSE;
    }

    if (array_is_created(&result->binary_values)) {
        struct pgsql_binary_value *values;
        unsigned int i, count;

        values = array_get_modifiable(&result->binary_values, &count);
        for (i = 0; i < count; i++)
            PQfreemem(values[i].value);
        array_free(&result->binary_values);
    }

    i_free(result->fields);
    i_free(result->values);
    i_free(result->query);
    i_free(result);

    if (db->queue != NULL && !db->querying && db->connected)
        queue_send_next(db);
}

static void result_finish(struct pgsql_result *result)
{
    struct pgsql_db *db = (struct pgsql_db *)result->api.db;
    bool free_result = TRUE, retry = FALSE;
    bool disconnected;

    disconnected = PQstatus(db->pg) == CONNECTION_BAD;
    if (disconnected && result->pgres == NULL && result->retry_query) {
        /* retry the query */
        i_error("pgsql: Query failed, retrying: %s", last_error(db));
        retry = TRUE;
    } else if (result->callback != NULL) {
        result->api.callback = TRUE;
        T_BEGIN {
            result->callback(&result->api, result->context);
        } T_END;
        result->api.callback = FALSE;
        free_result = db->sync_result != &result->api;
    }

    if (disconnected) {
        /* disconnected */
        driver_pgsql_close(db);

        if (retry) {
            /* retry the query */
            driver_pgsql_query_full(&db->api, result->query,
                        result->callback,
                        result->context, FALSE);
        }
    }
    if (free_result)
        driver_pgsql_result_free(&result->api);
}

static void get_result(struct pgsql_result *result)
{
        struct pgsql_db *db = (struct pgsql_db *)result->api.db;

    if (!PQconsumeInput(db->pg)) {
        db->connected = FALSE;
        result_finish(result);
        return;
    }

    if (PQisBusy(db->pg)) {
        if (db->io == NULL) {
            db->io = io_add(PQsocket(db->pg), IO_READ,
                    get_result, result);
            db->io_dir = IO_READ;
        }
        return;
    }

    if (db->io != NULL)
        io_remove(&db->io);

    result->pgres = PQgetResult(db->pg);
    result_finish(result);
}

static void flush_callback(struct pgsql_result *result)
{
        struct pgsql_db *db = (struct pgsql_db *)result->api.db;
    int ret;

    ret = PQflush(db->pg);
    if (ret > 0)
        return;

    io_remove(&db->io);

    if (ret < 0) {
        db->connected = FALSE;
        result_finish(result);
    } else {
        /* all flushed */
        get_result(result);
    }
}

static void send_query(struct pgsql_result *result, const char *query)
{
        struct pgsql_db *db = (struct pgsql_db *)result->api.db;
    int ret;

    i_assert(db->io == NULL);
    i_assert(!db->querying);
    i_assert(db->connected);

    if (!PQsendQuery(db->pg, query)) {
        db->connected = FALSE;
        result_finish(result);
        return;
    }

    ret = PQflush(db->pg);
    if (ret < 0) {
        db->connected = FALSE;
        result_finish(result);
        return;
    }

    db->querying = TRUE;
    if (ret > 0) {
        /* write blocks */
        db->io = io_add(PQsocket(db->pg), IO_WRITE,
                flush_callback, result);
        db->io_dir = IO_WRITE;
    } else {
        get_result(result);
    }
}

static struct pgsql_queue *queue_unlink_first(struct pgsql_db *db)
{
    struct pgsql_queue *queue;

    queue = db->queue;
    db->queue = queue->next;

    if (db->queue == NULL)
        db->queue_tail = &db->queue;
    return queue;
}

static void queue_send_next(struct pgsql_db *db)
{
    struct pgsql_queue *queue;

    queue = queue_unlink_first(db);
    send_query(queue->result, queue->query);

    i_free(queue->query);
    i_free(queue);
}

static void queue_abort_next(struct pgsql_db *db)
{
    struct pgsql_queue *queue;

    queue = queue_unlink_first(db);

    queue->result->callback(&sql_not_connected_result,
                queue->result->context);
    i_free(queue->result);
    i_free(queue->query);
    i_free(queue);
}

static void queue_drop_timed_out_queries(struct pgsql_db *db)
{
    while (db->queue != NULL &&
           db->queue->created + QUERY_TIMEOUT_SECS < ioloop_time)
        queue_abort_next(db);

}

static void queue_timeout(struct pgsql_db *db)
{
    if (db->querying)
        return;

    if (!db->connected) {
        queue_drop_timed_out_queries(db);
        driver_pgsql_connect(&db->api);
        return;
    }

    if (db->queue != NULL)
        queue_send_next(db);

    if (db->queue == NULL)
        timeout_remove(&db->queue_to);
}

static void
driver_pgsql_queue_query(struct pgsql_result *result, const char *query)
{
        struct pgsql_db *db = (struct pgsql_db *)result->api.db;
    struct pgsql_queue *queue;

    queue = i_new(struct pgsql_queue, 1);
    queue->created = time(NULL);
    queue->query = i_strdup(query);
    queue->result = result;

    *db->queue_tail = queue;
    db->queue_tail = &queue->next;

    if (db->queue_to == NULL)
        db->queue_to = timeout_add(5000, queue_timeout, db);
}

static void do_query(struct pgsql_result *result, const char *query)
{
        struct pgsql_db *db = (struct pgsql_db *)result->api.db;

    i_assert(db->sync_result == NULL);

    if (db->querying) {
        /* only one query at a time */
        driver_pgsql_queue_query(result, query);
        return;
    }

    if (!db->connected) {
        /* try connecting again */
        driver_pgsql_connect(&db->api);
        driver_pgsql_queue_query(result, query);
        return;
    }

    if (db->queue == NULL)
        send_query(result, query);
    else {
        /* there's already queries queued, send them first */
        driver_pgsql_queue_query(result, query);
        queue_send_next(db);
    }
}

static void exec_callback(struct sql_result *_result,
              void *context ATTR_UNUSED)
{
    struct pgsql_db *db = (struct pgsql_db *)_result->db;

    i_error("pgsql: sql_exec() failed: %s", last_error(db));
}

static const char *
driver_pgsql_escape_string(struct sql_db *_db, const char *string)
{
    struct pgsql_db *db = (struct pgsql_db *)_db;
    size_t len = strlen(string);
    char *to;

#ifdef HAVE_PQESCAPE_STRING_CONN
    if (!db->connected) {
        /* try connecting again */
        (void)driver_pgsql_connect(&db->api);
    }
    to = t_buffer_get(len * 2 + 1);
    len = PQescapeStringConn(db->pg, to, string, len, NULL);
#else
    to = t_buffer_get(len * 2 + 1);
    len = PQescapeString(to, string, len);
#endif
    t_buffer_alloc(len + 1);
    return to;
}

static void driver_pgsql_exec_full(struct sql_db *db, const char *query,
                   bool retry_query)
{
    struct pgsql_result *result;

    result = i_new(struct pgsql_result, 1);
    result->api = driver_pgsql_result;
    result->api.db = db;
    result->callback = exec_callback;
    if (retry_query) {
        result->query = i_strdup(query);
        result->retry_query = TRUE;
    }
    do_query(result, query);
}

static void driver_pgsql_exec(struct sql_db *db, const char *query)
{
    driver_pgsql_exec_full(db, query, TRUE);
}

static void
driver_pgsql_query_full(struct sql_db *db, const char *query,
            sql_query_callback_t *callback, void *context,
            bool retry_query)
{
    struct pgsql_result *result;

    result = i_new(struct pgsql_result, 1);
    result->api = driver_pgsql_result;
    result->api.db = db;
    result->callback = callback;
    result->context = context;
    if (retry_query) {
        result->query = i_strdup(query);
        result->retry_query = TRUE;
    }
    do_query(result, query);
}

static void driver_pgsql_query(struct sql_db *db, const char *query,
                   sql_query_callback_t *callback, void *context)
{
    driver_pgsql_query_full(db, query, callback, context, TRUE);
}

static void pgsql_query_s_callback(struct sql_result *result, void *context)
{
        struct pgsql_db *db = context;

    db->query_finished = TRUE;
    db->sync_result = result;
    io_loop_stop(db->ioloop);
}

static struct sql_result *
driver_pgsql_query_s(struct sql_db *_db, const char *query)
{
        struct pgsql_db *db = (struct pgsql_db *)_db;
    struct sql_result *result;
    struct io old_io;

    if (db->queue_to != NULL) {
        /* we're creating a new ioloop, make sure the timeout gets
           added there. */
        timeout_remove(&db->queue_to);
    }

    if (db->io == NULL)
        db->ioloop = io_loop_create();
    else {
        /* have to move our existing I/O handler to new I/O loop */
        old_io = *db->io;
        io_remove(&db->io);

        db->ioloop = io_loop_create();

        db->io = io_add(PQsocket(db->pg), old_io.condition,
                old_io.callback, old_io.context);
    }

    db->query_finished = FALSE;
    driver_pgsql_query(_db, query, pgsql_query_s_callback, db);

    if (!db->query_finished) {
        if ((db->connected || db->connecting) && db->io != NULL)
            io_loop_run(db->ioloop);
        else
            queue_abort_next(db);

        if (db->io != NULL) {
            i_assert(db->sync_result == &sql_not_connected_result);
            io_remove(&db->io);
        }
        if (db->queue_to != NULL)
            timeout_remove(&db->queue_to);
    } else {
        i_assert(db->io == NULL);
        i_assert(db->queue_to == NULL);
    }
    io_loop_destroy(&db->ioloop);

    result = db->sync_result;
    if (result == &sql_not_connected_result) {
        /* we don't end up in pgsql's free function, so sync_result
           won't be set to NULL if we don't do it here. */
        db->sync_result = NULL;
    }

    i_assert(db->io == NULL);
    return result;
}

static int driver_pgsql_result_next_row(struct sql_result *_result)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;
    struct pgsql_db *db = (struct pgsql_db *)_result->db;

    if (result->rows != 0) {
        /* second time we're here */
        return ++result->rownum < result->rows;
    }

    if (result->pgres == NULL)
        return -1;

    switch (PQresultStatus(result->pgres)) {
    case PGRES_COMMAND_OK:
        /* no rows returned */
        return 0;
    case PGRES_TUPLES_OK:
        result->rows = PQntuples(result->pgres);
        return result->rows > 0;
    case PGRES_EMPTY_QUERY:
    case PGRES_NONFATAL_ERROR:
        /* nonfatal error */
        return -1;
    default:
        /* treat as fatal error */
        db->connected = FALSE;
        return -1;
    }
}

static void driver_pgsql_result_fetch_fields(struct pgsql_result *result)
{
    unsigned int i;

    if (result->fields != NULL)
        return;

    /* @UNSAFE */
    result->fields_count = PQnfields(result->pgres);
    result->fields = i_new(const char *, result->fields_count);
    for (i = 0; i < result->fields_count; i++)
        result->fields[i] = PQfname(result->pgres, i);
}

static unsigned int
driver_pgsql_result_get_fields_count(struct sql_result *_result)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;

        driver_pgsql_result_fetch_fields(result);
    return result->fields_count;
}

static const char *
driver_pgsql_result_get_field_name(struct sql_result *_result, unsigned int idx)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;

    driver_pgsql_result_fetch_fields(result);
    i_assert(idx < result->fields_count);
    return result->fields[idx];
}

static int driver_pgsql_result_find_field(struct sql_result *_result,
                      const char *field_name)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;
    unsigned int i;

    driver_pgsql_result_fetch_fields(result);
    for (i = 0; i < result->fields_count; i++) {
        if (strcmp(result->fields[i], field_name) == 0)
            return i;
    }
    return -1;
}

static const char *
driver_pgsql_result_get_field_value(struct sql_result *_result,
                    unsigned int idx)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;

    if (PQgetisnull(result->pgres, result->rownum, idx))
        return NULL;

    return PQgetvalue(result->pgres, result->rownum, idx);
}

static const unsigned char *
driver_pgsql_result_get_field_value_binary(struct sql_result *_result,
                       unsigned int idx, size_t *size_r)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;
    const char *value;
    struct pgsql_binary_value *binary_value;

    if (PQgetisnull(result->pgres, result->rownum, idx)) {
        *size_r = 0;
        return NULL;
    }

    value = PQgetvalue(result->pgres, result->rownum, idx);

    if (!array_is_created(&result->binary_values))
        i_array_init(&result->binary_values, idx + 1);

    binary_value = array_idx_modifiable(&result->binary_values, idx);
    if (binary_value->value == NULL) {
        binary_value->value =
            PQunescapeBytea((const unsigned char *)value,
                    &binary_value->size);
    }

    *size_r = binary_value->size;
    return binary_value->value;
}

static const char *
driver_pgsql_result_find_field_value(struct sql_result *result,
                     const char *field_name)
{
    int idx;

    idx = driver_pgsql_result_find_field(result, field_name);
    if (idx < 0)
        return NULL;
    return driver_pgsql_result_get_field_value(result, idx);
}

static const char *const *
driver_pgsql_result_get_values(struct sql_result *_result)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;
    unsigned int i;

    if (result->values == NULL) {
        driver_pgsql_result_fetch_fields(result);
        result->values = i_new(const char *, result->fields_count);
    }

    /* @UNSAFE */
    for (i = 0; i < result->fields_count; i++) {
        result->values[i] =
                        driver_pgsql_result_get_field_value(_result, i);
    }

    return result->values;
}

static const char *driver_pgsql_result_get_error(struct sql_result *_result)
{
    struct pgsql_result *result = (struct pgsql_result *)_result;
    struct pgsql_db *db = (struct pgsql_db *)_result->db;
    const char *msg;
    size_t len;

    i_free_and_null(db->error);

    if (result->pgres == NULL) {
        /* connection error */
        db->error = i_strdup(last_error(db));
    } else {
        msg = PQresultErrorMessage(result->pgres);
        if (msg == NULL)
            return "(no error set)";

        /* Error message should contain trailing \n, we don't want it */
        len = strlen(msg);
        db->error = len == 0 || msg[len-1] != '\n' ?
            i_strdup(msg) : i_strndup(msg, len-1);
    }
    return db->error;
}

static struct sql_transaction_context *
driver_pgsql_transaction_begin(struct sql_db *db)
{
    struct pgsql_transaction_context *ctx;

    ctx = i_new(struct pgsql_transaction_context, 1);
    ctx->ctx.db = db;
    return &ctx->ctx;
}

static void
transaction_commit_callback(struct sql_result *result, void *context)
{
    struct pgsql_transaction_context *ctx =
        (struct pgsql_transaction_context *)context;

    if (sql_result_next_row(result) < 0)
        ctx->callback(sql_result_get_error(result), ctx->context);
    else
        ctx->callback(NULL, ctx->context);
}

static void
driver_pgsql_transaction_commit(struct sql_transaction_context *_ctx,
                sql_commit_callback_t *callback, void *context)
{
    struct pgsql_transaction_context *ctx =
        (struct pgsql_transaction_context *)_ctx;

    if (ctx->failed) {
        callback(ctx->error, context);
        if (ctx->opened)
            sql_exec(_ctx->db, "ROLLBACK");
        i_free(ctx);
        return;
    }

    if (!ctx->opened) {
        /* nothing done */
        ctx->callback(NULL, ctx->context);
        return;
    }

    ctx->callback = callback;
    ctx->context = context;

    driver_pgsql_query_full(_ctx->db, "COMMIT",
                transaction_commit_callback, ctx, FALSE);
}

static int
driver_pgsql_transaction_commit_s(struct sql_transaction_context *_ctx,
                  const char **error_r)
{
    struct pgsql_transaction_context *ctx =
        (struct pgsql_transaction_context *)_ctx;
    struct sql_result *result;

    if (ctx->first_update != NULL) {
        sql_query(_ctx->db, ctx->first_update,
              transaction_update_callback, ctx);
        i_free_and_null(ctx->first_update);
    }

    if (ctx->failed) {
        *error_r = ctx->error;
        if (ctx->opened)
            sql_exec(_ctx->db, "ROLLBACK");
    } else if (!ctx->opened)
        *error_r = NULL;
    else {
        result = sql_query_s(_ctx->db, "COMMIT");
        if (ctx->failed)
            *error_r = ctx->error;
        else if (sql_result_next_row(result) < 0)
            *error_r = sql_result_get_error(result);
        else
            *error_r = NULL;
        sql_result_free(result);
    }

    i_free(ctx);
    return *error_r == NULL ? 0 : -1;
}

static void
driver_pgsql_transaction_rollback(struct sql_transaction_context *_ctx)
{
    struct pgsql_transaction_context *ctx =
        (struct pgsql_transaction_context *)_ctx;

    if (ctx->opened)
        sql_exec(_ctx->db, "ROLLBACK");
    i_free(ctx->first_update);
    i_free(ctx);
}

static void
transaction_update_callback(struct sql_result *result, void *context)
{
    struct pgsql_transaction_context *ctx = context;

    if (sql_result_next_row(result) < 0) {
        ctx->failed = TRUE;
        ctx->error = sql_result_get_error(result);
    }
}

static void
driver_pgsql_update(struct sql_transaction_context *_ctx, const char *query)
{
    struct pgsql_transaction_context *ctx =
        (struct pgsql_transaction_context *)_ctx;

    if (ctx->failed)
        return;

    if (!ctx->opened) {
        if (ctx->first_update == NULL) {
            /* delay sending the first update in case there is
               only one to be sent and we don't need BEGIN/COMMIT */
            ctx->first_update = i_strdup(query);
            return;
        }
        ctx->opened = TRUE;
        sql_query(_ctx->db, "BEGIN", transaction_update_callback, ctx);
        sql_query(_ctx->db, ctx->first_update,
              transaction_update_callback, ctx);
        i_free_and_null(ctx->first_update);
    }

    driver_pgsql_query_full(_ctx->db, query,
                transaction_update_callback, ctx, FALSE);
}

struct sql_db driver_pgsql_db = {
    "pgsql",

    MEMBER(v) {
        driver_pgsql_init_v,
        driver_pgsql_deinit_v,
        driver_pgsql_get_flags,
        driver_pgsql_connect,
        driver_pgsql_escape_string,
        driver_pgsql_exec,
        driver_pgsql_query,
        driver_pgsql_query_s,

        driver_pgsql_transaction_begin,
        driver_pgsql_transaction_commit,
        driver_pgsql_transaction_commit_s,
        driver_pgsql_transaction_rollback,

        driver_pgsql_update
    }
};

struct sql_result driver_pgsql_result = {
    MEMBER(v) {
        driver_pgsql_result_free,
        driver_pgsql_result_next_row,
        driver_pgsql_result_get_fields_count,
        driver_pgsql_result_get_field_name,
        driver_pgsql_result_find_field,
        driver_pgsql_result_get_field_value,
        driver_pgsql_result_get_field_value_binary,
        driver_pgsql_result_find_field_value,
        driver_pgsql_result_get_values,
        driver_pgsql_result_get_error
    }
};

void driver_pgsql_init(void);
void driver_pgsql_deinit(void);

void driver_pgsql_init(void)
{
    sql_driver_register(&driver_pgsql_db);
}

void driver_pgsql_deinit(void)
{
    sql_driver_unregister(&driver_pgsql_db);
}

#endif