src/server/plugins/postgresql/resource.c

/*
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
 *
 * Copyright 2022 Olaf Wintermann. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1. Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *
 *   2. Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "resource.h"

static ResourceType pg_resource_type = {
    (resource_pool_init_func)pg_resourcepool_init,
    (resource_pool_destroy_func)pg_resourcepool_destroy,
    (resource_pool_createresource_func)pg_resourcepool_createresource,
    (resource_pool_freeresource_func)pg_resourcepool_freeresource,
    (resource_pool_prepare_func)pg_resourcepool_prepare,
    (resource_pool_finish_func)pg_resourcepool_finish,
    (resource_pool_getresourcedata_func)pg_resourcepool_getresourcedata
};

ResourceType* pg_get_resource_type(void) {
    return &pg_resource_type;
}

void * pg_resourcepool_init(pool_handle_t *pool, const char *rpname, pblock *pb) {
    char *connection = pblock_findval("Connection", pb);
    if(!connection) {
        log_ereport(LOG_MISCONFIG, "Resource pool %s: Missing Connection parameter", rpname);
        return NULL;
    }
    
    // test connection
    PGconn *test_connection = PQconnectdb(connection);
    if(pg_check_connection(LOG_WARN, rpname, test_connection)) {
        log_ereport(LOG_WARN, "Resource pool %s: Connection check failed", rpname);
    }
    if(test_connection) PQfinish(test_connection);
    
    PgResourcePool *pg = pool_malloc(pool, sizeof(PgResourcePool));
    if(!pg) {
        return NULL;
    }
    pg->pool = pool;
    pg->name = rpname;
    pg->connection = connection;
    
    return pg;
    
}

int pg_check_connection(int loglevel, const char *rpname, PGconn *connection) {
    if(!connection) {
        log_ereport(loglevel, "Resource pool %s: Cannot create PQ connection", rpname);
        return 1;
    }
    if(PQstatus(connection) != CONNECTION_OK) {
        char *err = PQerrorMessage(connection);
        int errlen = 0;
        if(err) {
            errlen = strlen(err);
            if(errlen > 0 && err[errlen-1] == '\n') {
                errlen--;
            }
        }
        log_ereport(loglevel, "Resource pool %s: Failed to connect to database: %.*s", rpname, errlen, err);
        PQfinish(connection);
        return 1;
    }
    return 0;
}


void pg_resourcepool_destroy(PgResourcePool *pg) {
    // unused
}
    
void * pg_resourcepool_createresource(PgResourcePool *pg) {
    PGconn *connection = PQconnectdb(pg->connection);
    if(pg_check_connection(LOG_FAILURE, pg->name, connection)) {
        return NULL;
    }
    
    PgResource *res = pool_malloc(pg->pool, sizeof(PgResource));
    if(!res) {
        PQfinish(connection);
        log_ereport(LOG_CATASTROPHE, "pg_resourcepool_createresource: OOM");
        return NULL;
    }
    res->connection = connection;
    
    return res;
}

void pg_resourcepool_freeresource(PgResourcePool *pg, PgResource *res) {
    if(res->connection) {
        PQfinish(res->connection);
    }
}

int pg_resourcepool_prepare(PgResourcePool *pg, PgResource *res) {
    PGresult *result = PQexec(res->connection, "BEGIN");
    PQclear(result); // TODO: handle error
    return 0;
}

int pg_resourcepool_finish(PgResourcePool *pg, PgResource *res) {
    PGresult *result = PQexec(res->connection, "COMMIT");
    PQclear(result); // TODO: handle error
    return 0;
}

void * pg_resourcepool_getresourcedata(PgResource *res) {
    return res->connection;
}