<commit_before>/* A client for the Postgres logical replication protocol, comparable to
 * pg_recvlogical but adapted to our needs. The protocol is documented here:
 * http://www.postgresql.org/docs/9.4/static/protocol-replication.html */

#include "replication.h"

#include <stdarg.h>
#include <string.h>
#include <arpa/inet.h>
#include <sys/time.h>

#include <datatype/timestamp.h>
#include <internal/pqexpbuffer.h>

#define CHECKPOINT_INTERVAL_SEC 10

// #define DEBUG 1

int replication_stream_finish(replication_stream_t stream);
int parse_keepalive_message(replication_stream_t stream, char *buf, int buflen);
int parse_xlogdata_message(replication_stream_t stream, char *buf, int buflen);
int send_checkpoint(replication_stream_t stream, int64 now);
void repl_error(replication_stream_t stream, char *fmt, ...) __attribute__ ((format (printf, 2, 3)));
int64 current_time(void);
void sendint64(int64 i64, char *buf);
int64 recvint64(char *buf);


/* Send a CREATE_REPLICATION_SLOT ... LOGICAL command to the server. This is similar to
 * the pg_create_logical_replication_slot() function you can call from SQL, but with a
 * crucial addition: it exports a consistent snapshot which we can use to dump a copy
 * of the database contents at the start of the replication slot. Note that the snapshot
 * is deleted when the next command is sent to the server on this replication connection,
 * so the snapshot name should be used immediately after the replication slot has been
 * created.
 *
 * The server response to the CREATE_REPLICATION_SLOT doesn't seem to be documented
 * anywhere, but based on reading the code, it's a tuple with the following fields:
 *
 *   1. "slot_name": name of the slot that was created, as requested
 *   2. "consistent_point": LSN at which we became consistent
 *   3. "snapshot_name": exported snapshot's name
 *   4. "output_plugin": name of the output plugin, as requested
 */
int replication_slot_create(replication_stream_t stream) {
    if (!stream->slot_name || stream->slot_name[0] == '\0') {
        repl_error(stream, "slot_name must be set in replication stream");
        return EINVAL;
    }
    if (!stream->output_plugin || stream->output_plugin[0] == '\0') {
        repl_error(stream, "output_plugin must be set in replication stream");
        return EINVAL;
    }

    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "CREATE_REPLICATION_SLOT \"%s\" LOGICAL \"%s\"",
            stream->slot_name, stream->output_plugin);

    PGresult *res = PQexec(stream->conn, query->data);
    if (PQresultStatus(res) != PGRES_TUPLES_OK) {
        repl_error(stream, "Command failed: %s: %s", query->data, PQerrorMessage(stream->conn));
        goto error;
    }

    if (PQntuples(res) != 1 || PQnfields(res) != 4) {
        repl_error(stream, "Unexpected CREATE_REPLICATION_SLOT result (%d rows, %d fields)",
                PQntuples(res), PQnfields(res));
        goto error;
    }

    if (PQgetisnull(res, 0, 1) || PQgetisnull(res, 0, 2)) {
        repl_error(stream, "Unexpected null value in CREATE_REPLICATION_SLOT response");
        goto error;
    }

    uint32 h32, l32;
    if (sscanf(PQgetvalue(res, 0, 1), "%X/%X", &h32, &l32) != 2) {
        repl_error(stream, "Could not parse LSN: \"%s\"", PQgetvalue(res, 0, 1));
        goto error;
    }

    stream->start_lsn = ((uint64) h32) << 32 | l32;
    stream->snapshot_name = strdup(PQgetvalue(res, 0, 2));

    destroyPQExpBuffer(query);
    PQclear(res);
    return 0;

error:
    destroyPQExpBuffer(query);
    PQclear(res);
    return EIO;
}


/* Drops the replication slot, like the SQL function pg_drop_replication_slot(). */
int replication_slot_drop(replication_stream_t stream) {
    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "DROP_REPLICATION_SLOT \"%s\"", stream->slot_name);

    PGresult *res = PQexec(stream->conn, query->data);
    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
        repl_error(stream, "Command failed: %s: %s", query->data, PQerrorMessage(stream->conn));
        destroyPQExpBuffer(query);
        PQclear(res);
        return EIO;
    }

    destroyPQExpBuffer(query);
    PQclear(res);
    return 0;
}


/* Checks that the connection to the database server supports logical replication. */
int replication_stream_check(replication_stream_t stream) {
    PGresult *res = PQexec(stream->conn, "IDENTIFY_SYSTEM");
    if (PQresultStatus(res) != PGRES_TUPLES_OK) {
        repl_error(stream, "IDENTIFY_SYSTEM failed: %s", PQerrorMessage(stream->conn));
        PQclear(res);
        return EIO;
    }

    if (PQntuples(res) != 1 || PQnfields(res) < 4) {
        repl_error(stream, "Unexpected IDENTIFY_SYSTEM result (%d rows, %d fields).",
                PQntuples(res), PQnfields(res));
        PQclear(res);
        return EIO;
    }

    /* Check that the database name (fourth column of the result tuple) is non-null,
     * implying a database-specific connection. */
    if (PQgetisnull(res, 0, 3)) {
        repl_error(stream, "Not using a database-specific replication connection.");
        PQclear(res);
        return EIO;
    }

    PQclear(res);
    return 0;
}


/* Starts streaming logical changes from replication slot stream->slot_name,
 * starting from position stream->start_lsn. */
int replication_stream_start(replication_stream_t stream, const char *error_policy) {
    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "START_REPLICATION SLOT \"%s\" LOGICAL %X/%X (error_policy '%s')",
            stream->slot_name,
            (uint32) (stream->start_lsn >> 32), (uint32) stream->start_lsn,
            error_policy);

    PGresult *res = PQexec(stream->conn, query->data);

    if (PQresultStatus(res) != PGRES_COPY_BOTH) {
        repl_error(stream, "Could not send replication command \"%s\": %s",
                query->data, PQresultErrorMessage(res));
        PQclear(res);
        destroyPQExpBuffer(query);
        return EIO;
    }

    PQclear(res);
    destroyPQExpBuffer(query);
    return 0;
}

<commit_msg>Quote output plugin param name for consistency<commit_after>/* A client for the Postgres logical replication protocol, comparable to
 * pg_recvlogical but adapted to our needs. The protocol is documented here:
 * http://www.postgresql.org/docs/9.4/static/protocol-replication.html */

#include "replication.h"

#include <stdarg.h>
#include <string.h>
#include <arpa/inet.h>
#include <sys/time.h>

#include <datatype/timestamp.h>
#include <internal/pqexpbuffer.h>

#define CHECKPOINT_INTERVAL_SEC 10

// #define DEBUG 1

int replication_stream_finish(replication_stream_t stream);
int parse_keepalive_message(replication_stream_t stream, char *buf, int buflen);
int parse_xlogdata_message(replication_stream_t stream, char *buf, int buflen);
int send_checkpoint(replication_stream_t stream, int64 now);
void repl_error(replication_stream_t stream, char *fmt, ...) __attribute__ ((format (printf, 2, 3)));
int64 current_time(void);
void sendint64(int64 i64, char *buf);
int64 recvint64(char *buf);


/* Send a CREATE_REPLICATION_SLOT ... LOGICAL command to the server. This is similar to
 * the pg_create_logical_replication_slot() function you can call from SQL, but with a
 * crucial addition: it exports a consistent snapshot which we can use to dump a copy
 * of the database contents at the start of the replication slot. Note that the snapshot
 * is deleted when the next command is sent to the server on this replication connection,
 * so the snapshot name should be used immediately after the replication slot has been
 * created.
 *
 * The server response to the CREATE_REPLICATION_SLOT doesn't seem to be documented
 * anywhere, but based on reading the code, it's a tuple with the following fields:
 *
 *   1. "slot_name": name of the slot that was created, as requested
 *   2. "consistent_point": LSN at which we became consistent
 *   3. "snapshot_name": exported snapshot's name
 *   4. "output_plugin": name of the output plugin, as requested
 */
int replication_slot_create(replication_stream_t stream) {
    if (!stream->slot_name || stream->slot_name[0] == '\0') {
        repl_error(stream, "slot_name must be set in replication stream");
        return EINVAL;
    }
    if (!stream->output_plugin || stream->output_plugin[0] == '\0') {
        repl_error(stream, "output_plugin must be set in replication stream");
        return EINVAL;
    }

    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "CREATE_REPLICATION_SLOT \"%s\" LOGICAL \"%s\"",
            stream->slot_name, stream->output_plugin);

    PGresult *res = PQexec(stream->conn, query->data);
    if (PQresultStatus(res) != PGRES_TUPLES_OK) {
        repl_error(stream, "Command failed: %s: %s", query->data, PQerrorMessage(stream->conn));
        goto error;
    }

    if (PQntuples(res) != 1 || PQnfields(res) != 4) {
        repl_error(stream, "Unexpected CREATE_REPLICATION_SLOT result (%d rows, %d fields)",
                PQntuples(res), PQnfields(res));
        goto error;
    }

    if (PQgetisnull(res, 0, 1) || PQgetisnull(res, 0, 2)) {
        repl_error(stream, "Unexpected null value in CREATE_REPLICATION_SLOT response");
        goto error;
    }

    uint32 h32, l32;
    if (sscanf(PQgetvalue(res, 0, 1), "%X/%X", &h32, &l32) != 2) {
        repl_error(stream, "Could not parse LSN: \"%s\"", PQgetvalue(res, 0, 1));
        goto error;
    }

    stream->start_lsn = ((uint64) h32) << 32 | l32;
    stream->snapshot_name = strdup(PQgetvalue(res, 0, 2));

    destroyPQExpBuffer(query);
    PQclear(res);
    return 0;

error:
    destroyPQExpBuffer(query);
    PQclear(res);
    return EIO;
}


/* Drops the replication slot, like the SQL function pg_drop_replication_slot(). */
int replication_slot_drop(replication_stream_t stream) {
    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "DROP_REPLICATION_SLOT \"%s\"", stream->slot_name);

    PGresult *res = PQexec(stream->conn, query->data);
    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
        repl_error(stream, "Command failed: %s: %s", query->data, PQerrorMessage(stream->conn));
        destroyPQExpBuffer(query);
        PQclear(res);
        return EIO;
    }

    destroyPQExpBuffer(query);
    PQclear(res);
    return 0;
}


/* Checks that the connection to the database server supports logical replication. */
int replication_stream_check(replication_stream_t stream) {
    PGresult *res = PQexec(stream->conn, "IDENTIFY_SYSTEM");
    if (PQresultStatus(res) != PGRES_TUPLES_OK) {
        repl_error(stream, "IDENTIFY_SYSTEM failed: %s", PQerrorMessage(stream->conn));
        PQclear(res);
        return EIO;
    }

    if (PQntuples(res) != 1 || PQnfields(res) < 4) {
        repl_error(stream, "Unexpected IDENTIFY_SYSTEM result (%d rows, %d fields).",
                PQntuples(res), PQnfields(res));
        PQclear(res);
        return EIO;
    }

    /* Check that the database name (fourth column of the result tuple) is non-null,
     * implying a database-specific connection. */
    if (PQgetisnull(res, 0, 3)) {
        repl_error(stream, "Not using a database-specific replication connection.");
        PQclear(res);
        return EIO;
    }

    PQclear(res);
    return 0;
}


/* Starts streaming logical changes from replication slot stream->slot_name,
 * starting from position stream->start_lsn. */
int replication_stream_start(replication_stream_t stream, const char *error_policy) {
    PQExpBuffer query = createPQExpBuffer();
    appendPQExpBuffer(query, "START_REPLICATION SLOT \"%s\" LOGICAL %X/%X (\"error_policy\" '%s')",
            stream->slot_name,
            (uint32) (stream->start_lsn >> 32), (uint32) stream->start_lsn,
            error_policy);

    PGresult *res = PQexec(stream->conn, query->data);

    if (PQresultStatus(res) != PGRES_COPY_BOTH) {
        repl_error(stream, "Could not send replication command \"%s\": %s",
                query->data, PQresultErrorMessage(res));
        PQclear(res);
        destroyPQExpBuffer(query);
        return EIO;
    }

    PQclear(res);
    destroyPQExpBuffer(query);
    return 0;
}


{'fraction_non_alphanumeric': 0.11752119167146736, 'fraction_numerical': 0.012180067484157682, 'mean_word_length': 3.5839305922293474, 'pattern_counts': {'":': 10, '<': 21, '<?xml version=': 0, '>': 73, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6095642', 'n_tokens_mistral': 4060, 'n_tokens_neox': 3803, 'n_words': 1084}