backup module

#include "postgres.h"
#include "fmgr.h"
#include "miscadmin.h"
#include "commands/dbcommands.h"
#include "utils/builtins.h"

#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

// remember to pfree returned pointer
// the contents variable must not be used after being passed in
// only the return value is guaranteed to have a pointer to valid memory
char* read_to_end(int fd, char* contents)
{
    char buf[1024];
    char* tmp = NULL;
    int sz = 0;
    int n = 0;
    while ((n = read(fd, buf, 1024)) > 0)
    {
        tmp = contents;
        contents = (char*) palloc(sz+1024);
        memcpy(contents, tmp, sz);
        memcpy(contents+sz, buf, n);
        if (tmp != NULL)
          pfree(tmp);
        sz+=n;
    }
    return contents;
}

PG_FUNCTION_INFO_V1(replicon_pgdump);

Datum
replicon_pgdump(PG_FUNCTION_ARGS)
{
    pid_t pid, wpid;
    FunctionCallInfoData fcid;
    char port[16];
    char *targetPath, *sourceSchema, *dbname, *username, *format;

    if (PG_ARGISNULL(0))
      ereport(ERROR, (errmsg("parameter 0 (targetPath) must not be null")));
    if (PG_ARGISNULL(2))
      ereport(ERROR, (errmsg("parameter 2 (format) must not be null")));

    targetPath = TextDatumGetCString(PG_GETARG_TEXT_P(0));
    sourceSchema = NULL;
    if (!PG_ARGISNULL(1))
      sourceSchema = TextDatumGetCString(PG_GETARG_TEXT_P(1));
    format = TextDatumGetCString(PG_GETARG_TEXT_P(2));
    dbname = get_database_name(MyDatabaseId);
    username = GetUserNameFromId(GetUserId());

    MemSet(&fcid, 0, sizeof(fcid));
    snprintf(port, 16, "%i", (int) inet_server_port(&fcid));

    pid = fork();
    if (pid < 0)
    {
        // handle error
        ereport(ERROR, (errmsg("fork failed (%d)", pid)));
    }
    else if (pid == 0) // child process
    {
        char* argv[32];
        int i=0;
        argv[i++] = PG_BINDIR "/pg_dump";
        argv[i++] = "--no-password";
        argv[i++] = "-p";
        argv[i++] = &port[0];
        argv[i++] = "-f";
        argv[i++] = targetPath;
        argv[i++] = "-F";
        argv[i++] = format;
        argv[i++] = "-U";
        argv[i++] = username;

        if (sourceSchema != NULL && sourceSchema[0] != 0)
        {
            argv[i++] = "-n";
            argv[i++] = sourceSchema;
        }

        argv[i++] = dbname;
        argv[i++] = NULL;

        execv(PG_BINDIR "/pg_dump", argv);

        exit(-1);
    }
    else // parent process
    {
        int status;
        do
        {
            wpid = waitpid(pid, &status, WUNTRACED);
            if (wpid == -1)
                ereport(ERROR, (errmsg("waitpid failed")));
            else
            {
                if (WIFEXITED(status))
                {
                    int exitCode = WEXITSTATUS(status);
                    if (exitCode != 0)
                    {
                        ereport(ERROR, (errmsg("child exited with status=%d", exitCode)));
                    }
                }
                else if (WIFSIGNALED(status))
                {
                    ereport(ERROR, (errmsg("child killed with signal %d", WTERMSIG(status))));
                }
                else if (WIFSTOPPED(status))
                {
                    ereport(ERROR, (errmsg("child stopped with signal %d", WSTOPSIG(status))));
                }
                else
                {
                    ereport(ERROR, (errmsg("unexpected status (0x%x)", status)));
                }
            }
        }
        while (!WIFEXITED(status) && !WIFSIGNALED(status));
    }

    PG_RETURN_NULL();
}

PG_FUNCTION_INFO_V1(replicon_pgrestore);

Datum
replicon_pgrestore(PG_FUNCTION_ARGS)
{
    pid_t pid, wpid;
    FunctionCallInfoData fcid;
    char port[16];
    char *sourcePath, *sourceSchema, *dbname, *username;

    if (PG_ARGISNULL(0))
      ereport(ERROR, (errmsg("parameter 0 (sourcePath) must not be null")));

    sourcePath = TextDatumGetCString(PG_GETARG_TEXT_P(0));
    sourceSchema = NULL;
    if (!PG_ARGISNULL(1))
      sourceSchema = TextDatumGetCString(PG_GETARG_TEXT_P(1));
    dbname = get_database_name(MyDatabaseId);
    username = GetUserNameFromId(GetUserId());

    MemSet(&fcid, 0, sizeof(fcid));
    snprintf(port, 16, "%i", (int) inet_server_port(&fcid));

    pid = fork();
    if (pid < 0)
    {
        // handle error
        ereport(ERROR, (errmsg("fork failed (%d)", pid)));
    }
    else if (pid == 0) // child process
    {
        char* argv[32];
        int i=0;
        argv[i++] = PG_BINDIR "/pg_restore";
        argv[i++] = "--no-password";
        argv[i++] = "--no-owner";
        argv[i++] = "--exit-on-error";
        argv[i++] = "-p";
        argv[i++] = &port[0];
        argv[i++] = "-U";
        argv[i++] = username;
        argv[i++] = "-d";
        argv[i++] = dbname;

        if (sourceSchema != NULL && sourceSchema[0] != 0)
        {
            argv[i++] = "-n";
            argv[i++] = sourceSchema;
            argv[i++] = "--create";
        }
        else
        {
            argv[i++] = "--create";
        }

        argv[i++] = sourcePath;
        argv[i++] = NULL;

        execv(PG_BINDIR "/pg_restore", argv);

        exit(-1);
    }
    else // parent process
    {
        int status;
        do
        {
            wpid = waitpid(pid, &status, WUNTRACED);
            if (wpid == -1)
                ereport(ERROR, (errmsg("waitpid failed")));
            else
            {
                if (WIFEXITED(status))
                {
                    int exitCode = WEXITSTATUS(status);
                    if (exitCode != 0)
                    {
                        ereport(ERROR, (errmsg("child exited with status=%d", exitCode)));
                    }
                }
                else if (WIFSIGNALED(status))
                {
                    ereport(ERROR, (errmsg("child killed with signal %d", WTERMSIG(status))));
                }
                else if (WIFSTOPPED(status))
                {
                    ereport(ERROR, (errmsg("child stopped with signal %d", WSTOPSIG(status))));
                }
                else
                {
                    ereport(ERROR, (errmsg("unexpected status (0x%x)", status)));
                }
            }
        }
        while (!WIFEXITED(status) && !WIFSIGNALED(status));
    }

    PG_RETURN_NULL();
}

PG_FUNCTION_INFO_V1(replicon_backupcontainsschema);

Datum
replicon_backupcontainsschema(PG_FUNCTION_ARGS)
{
    pid_t pid, wpid;
    FunctionCallInfoData fcid;
    int fd[2];
    char port[16];
    char schemaSearch[1024];
    char *sourcePath, *sourceSchema, *dbname, *username, *contents = NULL;
    bool retval = false;

    if (PG_ARGISNULL(0))
      ereport(ERROR, (errmsg("parameter 0 (sourcePath) must not be null")));
    if (PG_ARGISNULL(1))
      ereport(ERROR, (errmsg("parameter 1 (schemaName) must not be null")));

    sourcePath = TextDatumGetCString(PG_GETARG_TEXT_P(0));
    sourceSchema = TextDatumGetCString(PG_GETARG_TEXT_P(1));
    dbname = get_database_name(MyDatabaseId);
    username = GetUserNameFromId(GetUserId());

    MemSet(&fcid, 0, sizeof(fcid));
    snprintf(port, 16, "%i", (int) inet_server_port(&fcid));

    snprintf(schemaSearch, 1024, "SCHEMA - %s ", sourceSchema);

    pipe(fd);

    pid = fork();
    if (pid < 0)
    {
        // handle error
        ereport(ERROR, (errmsg("fork failed (%d)", pid)));
    }
    else if (pid == 0) // child process
    {
        char* argv[32];
        int i=0;

        close(fd[0]); //close child read end of pipes

        dup2(fd[1], 1); //redirect stdout to pipe

        argv[i++] = PG_BINDIR "/pg_restore";
        argv[i++] = "--no-password";
        argv[i++] = "--list";
        argv[i++] = "--exit-on-error";
        argv[i++] = "-p";
        argv[i++] = &port[0];
        argv[i++] = "-U";
        argv[i++] = username;
        argv[i++] = sourcePath;
        argv[i++] = NULL;

        execv(PG_BINDIR "/pg_restore", argv);

        exit(-1);
    }
    else // parent process
    {
        int status;

        close(fd[1]); //close write end of pipe
        do
        {
            wpid = waitpid(pid, &status, WUNTRACED | WNOHANG);
            if (wpid == -1)
                ereport(ERROR, (errmsg("waitpid failed")));
            else if (wpid != 0)
            {
                if (WIFEXITED(status))
                {
                    int exitCode = WEXITSTATUS(status);
                    if (exitCode != 0)
                    {
                        ereport(ERROR, (errmsg("child exited with status=%d", exitCode)));
                    }
                }
                else if (WIFSIGNALED(status))
                {
                    ereport(ERROR, (errmsg("child killed with signal %d", WTERMSIG(status))));
                }
                else if (WIFSTOPPED(status))
                {
                    ereport(ERROR, (errmsg("child stopped with signal %d", WSTOPSIG(status))));
                }
                else
                {
                    ereport(ERROR, (errmsg("unexpected status (0x%x)", status)));
                }
            }
            else
            {
              // continue reading from pipe while process is running to ensure we
              // don't get hung by filling the buffer
              contents = read_to_end(fd[0], contents);
              sleep(0.01); //sleep for 10ms just to avoid a busy loop
            }
        }
        while (wpid == 0 || (!WIFEXITED(status) && !WIFSIGNALED(status)));

        // read from pipe
        contents = read_to_end(fd[0], contents);
        close(fd[0]); // close read end of pipe
        if (contents != NULL) {
          retval = strstr(contents, schemaSearch) != NULL;
          pfree(contents);
        }
    }

    PG_RETURN_BOOL(retval);
}