copy_test.c

text/x-csrc

Download

/* Debugging libpq SSL connection crashes. */
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>

#include "libpq-fe.h"

#define BUFSIZE 100000000
#define CONNINFO "host=localhost sslmode=require"

static char buf[BUFSIZE];


void die(char *fmt, ...)
{
    va_list ap;

    va_start(ap, fmt);
    vfprintf(stderr, fmt, ap);
    va_end(ap);
    exit(1);
}

void exec_sql(PGconn *conn, const char *sql, int expected_status)
{
    PGresult *res;

    fprintf(stderr, "sql: %s\n", sql);
    res = PQexec(conn, sql);
    if (PQresultStatus(res) != expected_status)
        die("failed: %s\n", sql, PQerrorMessage(conn));
}

int main(void)
{
    PGconn *conn;
    PGresult *res;
    unsigned long bytes = 0;
    int count = 0;

    conn = PQconnectdb(CONNINFO);
    if (PQstatus(conn) != CONNECTION_OK)
        die("connection to database failed: %s", PQerrorMessage(conn));
    fprintf(stderr, "Connected to: %s\n", CONNINFO);

    exec_sql(conn, "BEGIN", PGRES_COMMAND_OK);
    exec_sql(conn, "DROP TABLE IF exists test_ssl_copy", PGRES_COMMAND_OK);
    exec_sql(conn, "CREATE TABLE test_ssl_copy(t text)", PGRES_COMMAND_OK);
    exec_sql(conn, "COPY test_ssl_copy(t) FROM stdin", PGRES_COPY_IN);

    /* Set the connection to non-blocking to enable crashing */
    if (PQsetnonblocking(conn, 1) != 0)
        die("PQsetnonblocking failed: %s", PQerrorMessage(conn));

    while (1) {
        int rc;

        fprintf(stderr, "PQputCopyData #%d buf:%lu total:%lu\n", count, sizeof(buf), bytes);
        rc = PQputCopyData(conn, buf, sizeof(buf));

        if (rc == 0) {
            fprintf(stderr, "PQputCopyData needs retry\n");
        } if (rc < 0) {
            die("PQputCopyData() = %d: %s", rc, PQerrorMessage(conn));
        }

        bytes += sizeof(buf);
        ++count;
    }

    if (PQputCopyEnd(conn, NULL) != 1)
        die("PQputCopyEnd failed: %s", PQerrorMessage(conn));

    while ((res = PQgetResult(conn)) != NULL)
        if (PQresultStatus(res) != PGRES_COMMAND_OK)
        fprintf(stderr, "COPY completed with error, status=%d: %s", 
                PQresultStatus(res), PQerrorMessage(conn));

    exec_sql(conn, "ROLLBACK", PGRES_COMMAND_OK);
    PQfinish(conn);

    return 0;
}