This adds a simple struct and open and close functions to abstract and
isolate the zlib vs. stdio output logic and allow it to be reused.
---
 src/bin/pg_basebackup/pg_basebackup.c | 300 +++++++++++++++++-----------------
 1 file changed, 154 insertions(+), 146 deletions(-)

diff --git a/src/bin/pg_basebackup/pg_basebackup.c 
b/src/bin/pg_basebackup/pg_basebackup.c
index 80de882..fa942ab 100644
--- a/src/bin/pg_basebackup/pg_basebackup.c
+++ b/src/bin/pg_basebackup/pg_basebackup.c
@@ -51,6 +51,15 @@ typedef struct TablespaceList
        TablespaceListCell *tail;
 } TablespaceList;
 
+typedef struct TarStream {
+       char             path[MAXPGPATH];
+       FILE            *file;
+#ifdef HAVE_LIBZ
+       gzFile           zfile;
+#endif
+       bool             keepopen;
+} TarStream;
+
 /* Global options */
 static char *basedir = NULL;
 static TablespaceList tablespace_dirs = {NULL, NULL};
@@ -100,6 +109,8 @@ static void disconnect_and_exit(int code);
 static void verify_dir_is_empty_or_create(char *dirname);
 static void progress_report(int tablespacenum, const char *filename, bool 
force);
 
+static void OpenTarFile(TarStream *tarfile, const char *path);
+static void CloseTarFile(TarStream *tarfile);
 static void ReceiveTarFile(PGconn *conn, PGresult *res, int rownum);
 static void ReceiveAndUnpackTarFile(PGconn *conn, PGresult *res, int rownum);
 static void GenerateRecoveryConf(PGconn *conn);
@@ -725,169 +736,194 @@ parse_max_rate(char *src)
  * Write a piece of tar data
  */
 static void
-writeTarData(
-#ifdef HAVE_LIBZ
-                        gzFile ztarfile,
-#endif
-                        FILE *tarfile, char *buf, int r, char *current_file)
+writeTarData(TarStream *stream, char *buf, int r)
 {
 #ifdef HAVE_LIBZ
-       if (ztarfile != NULL)
+       if (stream->zfile != NULL)
        {
-               if (gzwrite(ztarfile, buf, r) != r)
+               if (gzwrite(stream->zfile, buf, r) != r)
                {
                        fprintf(stderr,
                                        _("%s: could not write to compressed 
file \"%s\": %s\n"),
-                                       progname, current_file, 
get_gz_error(ztarfile));
+                                       progname, stream->path, 
get_gz_error(stream->zfile));
                        disconnect_and_exit(1);
                }
        }
        else
 #endif
        {
-               if (fwrite(buf, r, 1, tarfile) != 1)
+               if (fwrite(buf, r, 1, stream->file) != 1)
                {
                        fprintf(stderr, _("%s: could not write to file \"%s\": 
%s\n"),
-                                       progname, current_file, 
strerror(errno));
+                                       progname, stream->path, 
strerror(errno));
                        disconnect_and_exit(1);
                }
        }
 }
 
-#ifdef HAVE_LIBZ
-#define WRITE_TAR_DATA(buf, sz) writeTarData(ztarfile, tarfile, buf, sz, 
filename)
-#else
-#define WRITE_TAR_DATA(buf, sz) writeTarData(tarfile, buf, sz, filename)
-#endif
 
 /*
- * Receive a tar format file from the connection to the server, and write
- * the data from this file directly into a tar file. If compression is
- * enabled, the data will be compressed while written to the file.
- *
- * The file will be named base.tar[.gz] if it's for the main data directory
- * or <tablespaceoid>.tar[.gz] if it's for another tablespace.
- *
- * No attempt to inspect or validate the contents of the file is done.
+ * Open a (possibly zlib-compressed) tar file for writing.  A filename of -
+ * will cause stdout to be used.
  */
 static void
-ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
+OpenTarFile(TarStream *tarfile, const char *path)
 {
-       char            filename[MAXPGPATH];
-       char       *copybuf = NULL;
-       FILE       *tarfile = NULL;
-       char            tarhdr[512];
-       bool            basetablespace = PQgetisnull(res, rownum, 0);
-       bool            in_tarhdr = true;
-       bool            skip_file = false;
-       size_t          tarhdrsz = 0;
-       size_t          filesz = 0;
+       bool            use_stdout;
 
-#ifdef HAVE_LIBZ
-       gzFile          ztarfile = NULL;
-#endif
+       MemSet(tarfile, 0, sizeof(*tarfile));
+       use_stdout = (strcmp(path, "-") == 0);
+       strlcpy(tarfile->path, path, sizeof(tarfile->path));
 
-       if (basetablespace)
+#ifdef HAVE_LIBZ
+       if (compresslevel != 0)
        {
-               /*
-                * Base tablespaces
-                */
-               if (strcmp(basedir, "-") == 0)
+               /* Compression is in use */
+               if (use_stdout)
+                       tarfile->zfile = gzdopen(dup(fileno(stdout)), "wb");
+               else
+                       tarfile->zfile = gzopen(tarfile->path, "wb");
+               if (!tarfile->zfile)
                {
-#ifdef HAVE_LIBZ
-                       if (compresslevel != 0)
-                       {
-                               ztarfile = gzdopen(dup(fileno(stdout)), "wb");
-                               if (gzsetparams(ztarfile, compresslevel,
-                                                               
Z_DEFAULT_STRATEGY) != Z_OK)
-                               {
-                                       fprintf(stderr,
-                                                       _("%s: could not set 
compression level %d: %s\n"),
-                                                       progname, 
compresslevel, get_gz_error(ztarfile));
-                                       disconnect_and_exit(1);
-                               }
-                       }
-                       else
-#endif
-                               tarfile = stdout;
-                       strcpy(filename, "-");
+                       fprintf(stderr,
+                                       _("%s: could not create compressed file 
\"%s\": %s\n"),
+                                       progname, tarfile->path, 
strerror(errno));
+                       disconnect_and_exit(1);
                }
-               else
+               if (gzsetparams(tarfile->zfile, compresslevel,
+                                               Z_DEFAULT_STRATEGY) != Z_OK)
                {
-#ifdef HAVE_LIBZ
-                       if (compresslevel != 0)
-                       {
-                               snprintf(filename, sizeof(filename), 
"%s/base.tar.gz", basedir);
-                               ztarfile = gzopen(filename, "wb");
-                               if (gzsetparams(ztarfile, compresslevel,
-                                                               
Z_DEFAULT_STRATEGY) != Z_OK)
-                               {
-                                       fprintf(stderr,
-                                                       _("%s: could not set 
compression level %d: %s\n"),
-                                                       progname, 
compresslevel, get_gz_error(ztarfile));
-                                       disconnect_and_exit(1);
-                               }
-                       }
-                       else
-#endif
-                       {
-                               snprintf(filename, sizeof(filename), 
"%s/base.tar", basedir);
-                               tarfile = fopen(filename, "wb");
-                       }
+                       fprintf(stderr,
+                                       _("%s: could not set compression level 
%d: %s\n"),
+                                       progname, compresslevel, 
get_gz_error(tarfile->zfile));
+                       disconnect_and_exit(1);
                }
        }
+
        else
+#endif
        {
-               /*
-                * Specific tablespace
-                */
-#ifdef HAVE_LIBZ
-               if (compresslevel != 0)
-               {
-                       snprintf(filename, sizeof(filename), "%s/%s.tar.gz", 
basedir,
-                                        PQgetvalue(res, rownum, 0));
-                       ztarfile = gzopen(filename, "wb");
-                       if (gzsetparams(ztarfile, compresslevel,
-                                                       Z_DEFAULT_STRATEGY) != 
Z_OK)
-                       {
-                               fprintf(stderr,
-                                               _("%s: could not set 
compression level %d: %s\n"),
-                                               progname, compresslevel, 
get_gz_error(ztarfile));
-                               disconnect_and_exit(1);
-                       }
+               /* Either no zlib support, or zlib support but compresslevel = 
0 */
+               if (use_stdout) {
+                       tarfile->file = stdout;
+                       tarfile->keepopen = true;
                }
                else
-#endif
+                       tarfile->file = fopen(tarfile->path, "wb");
+               if (!tarfile->file)
+               {
+                       fprintf(stderr, _("%s: could not create file \"%s\": 
%s\n"),
+                                       progname, tarfile->path, 
strerror(errno));
+                       disconnect_and_exit(1);
+               }
+       }
+}
+
+
+/*
+ * Close a tar file opened by OpenTarFile. The end-of-archive marker of two
+ * zero blocks will be written first.
+ */
+static void
+CloseTarFile(TarStream *tarfile)
+{
+       /* 2 * 512 bytes empty data */
+       static char             zerobuf[1024];
+
+       /*
+        * Write two completely empty blocks at the end of the tar file, as
+        * required by some tar programs.
+        */
+       writeTarData(tarfile, zerobuf, sizeof(zerobuf));
+
+       if (tarfile->keepopen)
+       {
+               if (fflush(tarfile->file) != 0)
                {
-                       snprintf(filename, sizeof(filename), "%s/%s.tar", 
basedir,
-                                        PQgetvalue(res, rownum, 0));
-                       tarfile = fopen(filename, "wb");
+                       fprintf(stderr,
+                                       _("%s: could not write to file \"%s\": 
%s\n"),
+                                       progname, tarfile->path, 
strerror(errno));
+                       disconnect_and_exit(1);
                }
+               return;
        }
 
 #ifdef HAVE_LIBZ
-       if (compresslevel != 0)
+       if (tarfile->zfile != NULL)
        {
-               if (!ztarfile)
+               if (gzclose(tarfile->zfile) != 0)
                {
-                       /* Compression is in use */
                        fprintf(stderr,
-                                       _("%s: could not create compressed file 
\"%s\": %s\n"),
-                                       progname, filename, 
get_gz_error(ztarfile));
+                                       _("%s: could not close compressed file 
\"%s\": %s\n"),
+                                       progname, tarfile->path, 
get_gz_error(tarfile->zfile));
                        disconnect_and_exit(1);
                }
        }
        else
 #endif
        {
-               /* Either no zlib support, or zlib support but compresslevel = 
0 */
-               if (!tarfile)
+               if (fclose(tarfile->file) != 0)
                {
-                       fprintf(stderr, _("%s: could not create file \"%s\": 
%s\n"),
-                                       progname, filename, strerror(errno));
+                       fprintf(stderr,
+                                       _("%s: could not close file \"%s\": 
%s\n"),
+                                       progname, tarfile->path, 
strerror(errno));
                        disconnect_and_exit(1);
                }
        }
+}
+
+
+/*
+ * Open a (possibly zlib-compressed) tar file for writing. The filebase
+ * argument should be the desired filename relative to basedir, without a .tar
+ * or .tar.gz file extension. If the user specified a basedir of - then stdout
+ * will be used.
+ */
+static void
+openRelativeTarFile(TarStream *tarfile, const char *filebase)
+{
+       char            path[MAXPGPATH];
+
+       if (strcmp(basedir, "-") == 0)
+               strlcpy(path, "-", sizeof(path));
+#ifdef HAVE_LIBZ
+       else if (compresslevel != 0)
+               snprintf(path, sizeof(path), "%s/%s.tar.gz", basedir, filebase);
+#endif
+       else
+               snprintf(path, sizeof(path), "%s/%s.tar", basedir, filebase);
+       OpenTarFile(tarfile, path);
+}
+
+
+/*
+ * Receive a tar format file from the connection to the server, and write
+ * the data from this file directly into a tar file. If compression is
+ * enabled, the data will be compressed while written to the file.
+ *
+ * The file will be named base.tar[.gz] if it's for the main data directory
+ * or <tablespaceoid>.tar[.gz] if it's for another tablespace.
+ *
+ * No attempt to inspect or validate the contents of the file is done.
+ */
+static void
+ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
+{
+       TarStream       stream;
+       char       *copybuf = NULL;
+       char            tarhdr[512];
+       bool            basetablespace = PQgetisnull(res, rownum, 0);
+       bool            in_tarhdr = true;
+       bool            skip_file = false;
+       size_t          tarhdrsz = 0;
+       size_t          filesz = 0;
+
+       if (basetablespace)
+               /* Base tablespace */
+               openRelativeTarFile(&stream, "base");
+       else
+               /* Specific tablespace */
+               openRelativeTarFile(&stream, PQgetvalue(res, rownum, 0));
 
        /*
         * Get the COPY data stream
@@ -937,41 +973,13 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
 
                                padding = ((recoveryconfcontents->len + 511) & 
~511) - recoveryconfcontents->len;
 
-                               WRITE_TAR_DATA(header, sizeof(header));
-                               WRITE_TAR_DATA(recoveryconfcontents->data, 
recoveryconfcontents->len);
+                               writeTarData(&stream, header, sizeof(header));
+                               writeTarData(&stream, 
recoveryconfcontents->data, recoveryconfcontents->len);
                                if (padding)
-                                       WRITE_TAR_DATA(zerobuf, padding);
-                       }
-
-                       /* 2 * 512 bytes empty data at end of file */
-                       WRITE_TAR_DATA(zerobuf, sizeof(zerobuf));
-
-#ifdef HAVE_LIBZ
-                       if (ztarfile != NULL)
-                       {
-                               if (gzclose(ztarfile) != 0)
-                               {
-                                       fprintf(stderr,
-                                          _("%s: could not close compressed 
file \"%s\": %s\n"),
-                                                       progname, filename, 
get_gz_error(ztarfile));
-                                       disconnect_and_exit(1);
-                               }
-                       }
-                       else
-#endif
-                       {
-                               if (strcmp(basedir, "-") != 0)
-                               {
-                                       if (fclose(tarfile) != 0)
-                                       {
-                                               fprintf(stderr,
-                                                               _("%s: could 
not close file \"%s\": %s\n"),
-                                                               progname, 
filename, strerror(errno));
-                                               disconnect_and_exit(1);
-                                       }
-                               }
+                                       writeTarData(&stream, zerobuf, padding);
                        }
 
+                       CloseTarFile(&stream);
                        break;
                }
                else if (r == -2)
@@ -988,7 +996,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
                         * tablespace, we never have to look for an existing 
recovery.conf
                         * file in the stream.
                         */
-                       WRITE_TAR_DATA(copybuf, r);
+                       writeTarData(&stream, copybuf, r);
                }
                else
                {
@@ -1059,7 +1067,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
                                                 * header unmodified.
                                                 */
                                                if (!skip_file)
-                                                       WRITE_TAR_DATA(tarhdr, 
512);
+                                                       writeTarData(&stream, 
tarhdr, 512);
                                        }
                                }
                                else
@@ -1077,7 +1085,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
                                                bytes2write = (filesz > rr ? rr 
: filesz);
 
                                                if (!skip_file)
-                                                       WRITE_TAR_DATA(copybuf 
+ pos, bytes2write);
+                                                       writeTarData(&stream, 
copybuf + pos, bytes2write);
 
                                                rr -= bytes2write;
                                                pos += bytes2write;
@@ -1098,9 +1106,9 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum)
                        }
                }
                totaldone += r;
-               progress_report(rownum, filename, false);
+               progress_report(rownum, stream.path, false);
        }                                                       /* while (1) */
-       progress_report(rownum, filename, true);
+       progress_report(rownum, stream.path, true);
 
        if (copybuf != NULL)
                PQfreemem(copybuf);
-- 
2.3.0



-- 
Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org)
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers

Reply via email to