From 81e88ef3a63d1304e14134be14f41c2865c61226 Mon Sep 17 00:00:00 2001
From: Justin Pryzby <pryzbyj@telsasoft.com>
Date: Sat, 12 Dec 2020 00:40:39 -0600
Subject: [PATCH 4/7] cmdline parser for compression alg/level/opts

---
 src/bin/pg_dump/compress_io.c         | 153 +++++++++++---------------
 src/bin/pg_dump/compress_io.h         |  17 +--
 src/bin/pg_dump/core                  | Bin 0 -> 1257472 bytes
 src/bin/pg_dump/kwlookup.c            |   1 +
 src/bin/pg_dump/pg_backup.h           |  19 +++-
 src/bin/pg_dump/pg_backup_archiver.c  |  46 ++++----
 src/bin/pg_dump/pg_backup_archiver.h  |  11 +-
 src/bin/pg_dump/pg_backup_custom.c    |   6 +-
 src/bin/pg_dump/pg_backup_directory.c |  19 ++--
 src/bin/pg_dump/pg_backup_tar.c       |  48 ++++----
 src/bin/pg_dump/pg_dump.c             | 144 ++++++++++++++++++++----
 11 files changed, 278 insertions(+), 186 deletions(-)
 create mode 100644 src/bin/pg_dump/core
 create mode 120000 src/bin/pg_dump/kwlookup.c

diff --git a/src/bin/pg_dump/compress_io.c b/src/bin/pg_dump/compress_io.c
index b51ba680a2..10db22ff88 100644
--- a/src/bin/pg_dump/compress_io.c
+++ b/src/bin/pg_dump/compress_io.c
@@ -86,12 +86,9 @@ struct CompressorState
 
 };
 
-static void ParseCompressionOption(int compression, CompressionAlgorithm *alg,
-								   int *level);
-
 /* Routines that support zlib compressed data I/O */
 #ifdef HAVE_LIBZ
-static void InitCompressorZlib(CompressorState *cs, int level);
+static void InitCompressorZlib(CompressorState *cs, Compress *compress);
 static void DeflateCompressorZlib(ArchiveHandle *AH, CompressorState *cs,
 								  bool flush);
 static void ReadDataFromArchiveZlib(ArchiveHandle *AH, ReadFunc readF);
@@ -101,7 +98,7 @@ static void EndCompressorZlib(ArchiveHandle *AH, CompressorState *cs);
 #endif
 
 #ifdef HAVE_LIBZSTD
-static void InitCompressorZstd(CompressorState *cs, int level);
+static void InitCompressorZstd(CompressorState *cs, Compress *compress);
 static void EndCompressorZstd(ArchiveHandle *AH, CompressorState *cs);
 static void DeflateCompressorZstd(ArchiveHandle *AH, CompressorState *cs);
 static void WriteDataToArchiveZstd(ArchiveHandle *AH, CompressorState *cs,
@@ -114,80 +111,40 @@ static void ReadDataFromArchiveNone(ArchiveHandle *AH, ReadFunc readF);
 static void WriteDataToArchiveNone(ArchiveHandle *AH, CompressorState *cs,
 								   const char *data, size_t dLen);
 
-/*
- * Interprets a numeric 'compression' value. The algorithm implied by the
- * value (zlib or none at the moment), is returned in *alg, and the
- * zlib compression level in *level.
- */
-static void
-ParseCompressionOption(int compression, CompressionAlgorithm *alg, int *level)
-{
-	switch (compression)
-	{
-#ifdef HAVE_LIBZSTD
-		case ZSTD_COMPRESSION:
-			*alg = COMPR_ALG_ZSTD;
-			break;
-#endif
-#ifdef HAVE_ZLIB
-		case Z_DEFAULT_COMPRESSION:
-		case 1..9:
-			*alg = COMPR_ALG_LIBZ;
-			break;
-#endif
-		case 0:
-			*alg = COMPR_ALG_NONE;
-			break;
-		default:
-			fatal("invalid compression code: %d", compression);
-			*alg = COMPR_ALG_NONE;	/* keep compiler quiet */
-	}
-
-	/* The level is just the passed-in value. */
-	if (level)
-		*level = compression;
-}
-
 /* Public interface routines */
 
 /* Allocate a new compressor */
 CompressorState *
-AllocateCompressor(int compression, WriteFunc writeF)
+AllocateCompressor(Compress *compression, WriteFunc writeF)
 {
 	CompressorState *cs;
-	CompressionAlgorithm alg;
-	int			level;
-
-	ParseCompressionOption(compression, &alg, &level);
-
-#ifndef HAVE_LIBZ
-	if (alg == COMPR_ALG_LIBZ)
-		fatal("not built with zlib support");
-#endif
 
 	cs = (CompressorState *) pg_malloc0(sizeof(CompressorState));
 	cs->writeF = writeF;
-	cs->comprAlg = alg;
+	cs->comprAlg = compression->alg;
 
 	/*
 	 * Perform compression algorithm specific initialization.
 	 */
-	switch (alg)
+	Assert (compression->alg != COMPR_ALG_DEFAULT);
+	switch (compression->alg)
 	{
 #ifdef HAVE_LIBZ
 	case COMPR_ALG_LIBZ:
-		InitCompressorZlib(cs, level);
+		InitCompressorZlib(cs, compression);
 		break;
 #endif
 #ifdef HAVE_LIBZSTD
 	case COMPR_ALG_ZSTD:
-		InitCompressorZstd(cs, level);
+		InitCompressorZstd(cs, compression);
 		break;
 #endif
 	case COMPR_ALG_NONE:
 		/* Do nothing */
 		break;
-	// default:
+	default:
+		/* Should not happen */
+		fatal("requested compression not available in this installation");
 	}
 
 	return cs;
@@ -198,12 +155,9 @@ AllocateCompressor(int compression, WriteFunc writeF)
  * out with ahwrite().
  */
 void
-ReadDataFromArchive(ArchiveHandle *AH, int compression, ReadFunc readF)
+ReadDataFromArchive(ArchiveHandle *AH, ReadFunc readF)
 {
-	CompressionAlgorithm alg;
-
-	ParseCompressionOption(compression, &alg, NULL);
-
+	CompressionAlgorithm alg = AH->compression.alg;
 	if (alg == COMPR_ALG_NONE)
 		ReadDataFromArchiveNone(AH, readF);
 	else if (alg == COMPR_ALG_LIBZ)
@@ -233,25 +187,25 @@ WriteDataToArchive(ArchiveHandle *AH, CompressorState *cs,
 {
 	switch (cs->comprAlg)
 	{
-		case COMPR_ALG_LIBZ:
 #ifdef HAVE_LIBZ
+		case COMPR_ALG_LIBZ:
 			WriteDataToArchiveZlib(AH, cs, data, dLen);
-#else
-			fatal("not built with zlib support");
-#endif
 			break;
+#endif
 
-		case COMPR_ALG_ZSTD:
 #ifdef HAVE_LIBZSTD
+		case COMPR_ALG_ZSTD:
 			WriteDataToArchiveZstd(AH, cs, data, dLen);
-#else
-			fatal("not built with zstd support");
-#endif
 			break;
+#endif
 
 		case COMPR_ALG_NONE:
 			WriteDataToArchiveNone(AH, cs, data, dLen);
 			break;
+
+		default:
+			/* Should not happen */
+			fatal("requested compression not available in this installation");
 	}
 }
 
@@ -278,9 +232,30 @@ EndCompressor(ArchiveHandle *AH, CompressorState *cs)
 
 #ifdef HAVE_LIBZSTD
 static void
-InitCompressorZstd(CompressorState *cs, int level)
+InitCompressorZstd(CompressorState *cs, Compress *compress)
 {
+	size_t res;
+
 	cs->u.zstd.cstream = ZSTD_createCStream();
+
+	if (compress->level != 0) // XXX
+	{
+		res = ZSTD_CCtx_setParameter(cs->u.zstd.cstream,
+				ZSTD_c_compressionLevel, compress->level);
+		if (ZSTD_isError(res))
+			fatal("could not set compression level: %s",
+					ZSTD_getErrorName(res));
+	}
+
+	if (compress->zstdlong)
+	{
+		res = ZSTD_CCtx_setParameter(cs->u.zstd.cstream,
+				ZSTD_c_enableLongDistanceMatching, 1); // XXX: allow to disable it
+		if (ZSTD_isError(res))
+			fatal("could not set compression parameter: %s",
+					ZSTD_getErrorName(res));
+	}
+
 	if (cs->u.zstd.cstream == NULL)
 		fatal("could not initialize compression library");
 
@@ -436,7 +411,7 @@ ReadDataFromArchiveZstd(ArchiveHandle *AH, ReadFunc readF)
  */
 
 static void
-InitCompressorZlib(CompressorState *cs, int level)
+InitCompressorZlib(CompressorState *cs, Compress *compress)
 {
 	z_streamp	zp;
 
@@ -453,7 +428,7 @@ InitCompressorZlib(CompressorState *cs, int level)
 	cs->zlibOut = (char *) pg_malloc(ZLIB_OUT_SIZE + 1);
 	cs->zlibOutSize = ZLIB_OUT_SIZE;
 
-	if (deflateInit(zp, level) != Z_OK)
+	if (deflateInit(zp, compress->level) != Z_OK)
 		fatal("could not initialize compression library: %s",
 			  zp->msg);
 
@@ -683,7 +658,7 @@ free_keep_errno(void *p)
  * On failure, return NULL with an error code in errno.
  */
 cfp *
-cfopen_read(const char *path, const char *mode, int compression)
+cfopen_read(const char *path, const char *mode, Compress *compression)
 {
 	cfp		   *fp;
 
@@ -698,7 +673,7 @@ cfopen_read(const char *path, const char *mode, int compression)
 		if (fp == NULL)
 		{
 			char	   *fname;
-			char	   *suffix = compression == ZSTD_COMPRESSION ? "zst" : "gz";
+			char	   *suffix = compression->alg == COMPR_ALG_ZSTD ? "zst" : "gz";
 
 			fname = psprintf("%s.%s", path, suffix);
 			fp = cfopen(fname, mode, compression);
@@ -714,24 +689,23 @@ cfopen_read(const char *path, const char *mode, int compression)
  * be a filemode as accepted by fopen() and gzopen() that indicates writing
  * ("w", "wb", "a", or "ab").
  *
- * If 'compression' is non-zero, a gzip compressed stream is opened, and
- * 'compression' indicates the compression level used. The ".gz" suffix
- * is automatically added to 'path' in that case.
+ * Use compression if specified.
+ * The appropriate suffix is automatically added to 'path' in that case.
  *
  * On failure, return NULL with an error code in errno.
  */
 cfp *
-cfopen_write(const char *path, const char *mode, int compression)
+cfopen_write(const char *path, const char *mode, Compress *compression)
 {
 	cfp		   *fp;
 
-	if (compression == 0)
+	if (compression->alg == COMPR_ALG_NONE)
 		fp = cfopen(path, mode, compression);
 	else
 	{
 #ifdef HAVE_LIBZ // XXX
 		char	   *fname;
-		char	   *suffix = compression == ZSTD_COMPRESSION ? "zst" : "gz";
+		char	   *suffix = compression->alg == COMPR_ALG_ZSTD ? "zst" : "gz";
 
 		fname = psprintf("%s.%s", path, suffix);
 		fp = cfopen(fname, mode, compression);
@@ -752,7 +726,7 @@ cfopen_write(const char *path, const char *mode, int compression)
  * On failure, return NULL with an error code in errno.
  */
 cfp *
-cfopen(const char *path, const char *mode, int compression)
+cfopen(const char *path, const char *mode, Compress *compression)
 {
 	cfp		   *fp = pg_malloc(sizeof(cfp));
 
@@ -764,18 +738,17 @@ cfopen(const char *path, const char *mode, int compression)
 	fp->zstd.fp = NULL;
 #endif
 
-	switch (compression)
+	switch (compression->alg)
 	{
 #ifdef HAVE_LIBZ
-	case 1 ... 9: // XXX: nonportable
-	case Z_DEFAULT_COMPRESSION:
-		if (compression != Z_DEFAULT_COMPRESSION)
+	case COMPR_ALG_LIBZ:
+		if (compression->level != Z_DEFAULT_COMPRESSION)
 		{
 			/* user has specified a compression level, so tell zlib to use it */
 			char		mode_compression[32];
 
 			snprintf(mode_compression, sizeof(mode_compression), "%s%d",
-					 mode, compression);
+					 mode, compression->level);
 			fp->compressedfp = gzopen(path, mode_compression);
 		}
 		else
@@ -795,7 +768,7 @@ cfopen(const char *path, const char *mode, int compression)
 #endif
 
 #ifdef HAVE_LIBZSTD
-	case ZSTD_COMPRESSION:
+	case COMPR_ALG_ZSTD:
 		fp->zstd.fp = fopen(path, mode);
 		// XXX: save the compression params
 		if (fp->zstd.fp == NULL)
@@ -825,15 +798,19 @@ cfopen(const char *path, const char *mode, int compression)
 		break;
 #endif
 
-	default:
+	case COMPR_ALG_NONE:
 		fp->uncompressedfp = fopen(path, mode);
 		if (fp->uncompressedfp == NULL)
 		{
 			free_keep_errno(fp);
 			fp = NULL;
 		}
-
 		return fp;
+		break;
+
+	default:
+		/* should not happen */
+		fatal("requested compression not available in this installation");
 	}
 }
 
@@ -1005,7 +982,7 @@ cfgets(cfp *fp, char *buf, int len)
 		buf[res] = 0;
 		if (strchr(buf, '\n'))
 			*strchr(buf, '\n') = '\0';
-		return buf;
+		return res > 0 ? buf : 0;
 	}
 	else
 #endif
diff --git a/src/bin/pg_dump/compress_io.h b/src/bin/pg_dump/compress_io.h
index f0ce06f176..ae37be08d9 100644
--- a/src/bin/pg_dump/compress_io.h
+++ b/src/bin/pg_dump/compress_io.h
@@ -21,13 +21,6 @@
 #define ZLIB_OUT_SIZE	4096
 #define ZLIB_IN_SIZE	4096
 
-typedef enum
-{
-	COMPR_ALG_NONE,
-	COMPR_ALG_LIBZ,
-	COMPR_ALG_ZSTD,
-} CompressionAlgorithm;
-
 /* Prototype for callback function to WriteDataToArchive() */
 typedef void (*WriteFunc) (ArchiveHandle *AH, const char *buf, size_t len);
 
@@ -47,8 +40,8 @@ typedef size_t (*ReadFunc) (ArchiveHandle *AH, char **buf, size_t *buflen);
 /* struct definition appears in compress_io.c */
 typedef struct CompressorState CompressorState;
 
-extern CompressorState *AllocateCompressor(int compression, WriteFunc writeF);
-extern void ReadDataFromArchive(ArchiveHandle *AH, int compression,
+extern CompressorState *AllocateCompressor(Compress *compression, WriteFunc writeF);
+extern void ReadDataFromArchive(ArchiveHandle *AH,
 								ReadFunc readF);
 extern void WriteDataToArchive(ArchiveHandle *AH, CompressorState *cs,
 							   const void *data, size_t dLen);
@@ -57,9 +50,9 @@ extern void EndCompressor(ArchiveHandle *AH, CompressorState *cs);
 
 typedef struct cfp cfp;
 
-extern cfp *cfopen(const char *path, const char *mode, int compression);
-extern cfp *cfopen_read(const char *path, const char *mode, int compression);
-extern cfp *cfopen_write(const char *path, const char *mode, int compression);
+extern cfp *cfopen(const char *path, const char *mode, Compress *compression);
+extern cfp *cfopen_read(const char *path, const char *mode, Compress *compression);
+extern cfp *cfopen_write(const char *path, const char *mode, Compress *compression);
 extern int	cfread(void *ptr, int size, cfp *fp);
 extern int	cfwrite(const void *ptr, int size, cfp *fp);
 extern int	cfgetc(cfp *fp);
diff --git a/src/bin/pg_dump/core b/src/bin/pg_dump/core
new file mode 100644
index 0000000000000000000000000000000000000000..bdb556ec717e50744ccc2a7256ba9f90cb41427c
GIT binary patch
diff --git a/src/bin/pg_dump/kwlookup.c b/src/bin/pg_dump/kwlookup.c
new file mode 120000
index 0000000000..73c3abc1c7
--- /dev/null
+++ b/src/bin/pg_dump/kwlookup.c
@@ -0,0 +1 @@
+../../../src/backend/parser/kwlookup.c
\ No newline at end of file
diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index 9d0056a569..00d77ac16a 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -72,6 +72,21 @@ typedef struct _connParams
 	char	   *override_dbname;
 } ConnParams;
 
+typedef enum
+{
+		COMPR_ALG_DEFAULT = -1,
+		COMPR_ALG_NONE,
+		COMPR_ALG_LIBZ,
+		COMPR_ALG_ZSTD,
+} CompressionAlgorithm;
+
+typedef struct Compress {
+	CompressionAlgorithm	alg;
+	int			level;
+	bool		zstdlong;
+} Compress;
+
+
 typedef struct _restoreOptions
 {
 	int			createDB;		/* Issue commands to create the database */
@@ -125,7 +140,7 @@ typedef struct _restoreOptions
 
 	int			noDataForFailedTables;
 	int			exit_on_error;
-	int			compression;
+	Compress	compression;
 	int			suppressDumpWarnings;	/* Suppress output of WARNING entries
 										 * to stderr */
 	bool		single_txn;
@@ -281,7 +296,7 @@ extern Archive *OpenArchive(const char *FileSpec, const ArchiveFormat fmt);
 
 /* Create a new archive */
 extern Archive *CreateArchive(const char *FileSpec, const ArchiveFormat fmt,
-							  const int compression, bool dosync, ArchiveMode mode,
+							  Compress *compression, bool dosync, ArchiveMode mode,
 							  SetupWorkerPtrType setupDumpWorker);
 
 /* The --list option */
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index 1f82c6499b..be2b0882d9 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -70,7 +70,7 @@ typedef struct _parallelReadyList
 
 
 static ArchiveHandle *_allocAH(const char *FileSpec, const ArchiveFormat fmt,
-							   const int compression, bool dosync, ArchiveMode mode,
+							   Compress *compression, bool dosync, ArchiveMode mode,
 							   SetupWorkerPtrType setupWorkerPtr);
 static void _getObjectDescription(PQExpBuffer buf, TocEntry *te);
 static void _printTocEntry(ArchiveHandle *AH, TocEntry *te, bool isData);
@@ -98,7 +98,7 @@ static int	_discoverArchiveFormat(ArchiveHandle *AH);
 static int	RestoringToDB(ArchiveHandle *AH);
 static void dump_lo_buf(ArchiveHandle *AH);
 static void dumpTimestamp(ArchiveHandle *AH, const char *msg, time_t tim);
-static void SetOutput(ArchiveHandle *AH, const char *filename, int compression);
+static void SetOutput(ArchiveHandle *AH, const char *filename, Compress *compress);
 static OutputContext SaveOutput(ArchiveHandle *AH);
 static void RestoreOutput(ArchiveHandle *AH, OutputContext savedContext);
 
@@ -238,7 +238,7 @@ setupRestoreWorker(Archive *AHX)
 /* Public */
 Archive *
 CreateArchive(const char *FileSpec, const ArchiveFormat fmt,
-			  const int compression, bool dosync, ArchiveMode mode,
+			  Compress *compression, bool dosync, ArchiveMode mode,
 			  SetupWorkerPtrType setupDumpWorker)
 
 {
@@ -253,7 +253,9 @@ CreateArchive(const char *FileSpec, const ArchiveFormat fmt,
 Archive *
 OpenArchive(const char *FileSpec, const ArchiveFormat fmt)
 {
-	ArchiveHandle *AH = _allocAH(FileSpec, fmt, 0, true, archModeRead, setupRestoreWorker);
+	Compress compress = {0};
+	ArchiveHandle *AH = _allocAH(FileSpec, fmt, &compress, true, archModeRead,
+			setupRestoreWorker);
 
 	return (Archive *) AH;
 }
@@ -382,7 +384,7 @@ RestoreArchive(Archive *AHX)
 	 * Make sure we won't need (de)compression we haven't got
 	 */
 #ifndef HAVE_LIBZ
-	if (AH->compression != 0 && AH->PrintTocDataPtr != NULL)
+	if (AH->compression.alg != COMPR_ALG_NONE && AH->PrintTocDataPtr != NULL)
 	{
 		for (te = AH->toc->next; te != AH->toc; te = te->next)
 		{
@@ -457,8 +459,8 @@ RestoreArchive(Archive *AHX)
 	 * Setup the output file if necessary.
 	 */
 	sav = SaveOutput(AH);
-	if (ropt->filename || ropt->compression)
-		SetOutput(AH, ropt->filename, ropt->compression);
+	if (ropt->filename || ropt->compression.alg != COMPR_ALG_NONE)
+		SetOutput(AH, ropt->filename, &ropt->compression);
 
 	ahprintf(AH, "--\n-- PostgreSQL database dump\n--\n\n");
 
@@ -738,7 +740,7 @@ RestoreArchive(Archive *AHX)
 	 */
 	AH->stage = STAGE_FINALIZING;
 
-	if (ropt->filename || ropt->compression)
+	if (ropt->filename || ropt->compression.alg != COMPR_ALG_NONE)
 		RestoreOutput(AH, sav);
 
 	if (ropt->useDB)
@@ -1121,10 +1123,11 @@ PrintTOCSummary(Archive *AHX)
 	OutputContext sav;
 	const char *fmtName;
 	char		stamp_str[64];
+	Compress	nocompression = {0};
 
 	sav = SaveOutput(AH);
 	if (ropt->filename)
-		SetOutput(AH, ropt->filename, 0 /* no compression */ );
+		SetOutput(AH, ropt->filename, &nocompression);
 
 	if (strftime(stamp_str, sizeof(stamp_str), PGDUMP_STRFTIME_FMT,
 				 localtime(&AH->createDate)) == 0)
@@ -1133,7 +1136,7 @@ PrintTOCSummary(Archive *AHX)
 	ahprintf(AH, ";\n; Archive created at %s\n", stamp_str);
 	ahprintf(AH, ";     dbname: %s\n;     TOC Entries: %d\n;     Compression: %d\n",
 			 sanitize_line(AH->archdbname, false),
-			 AH->tocCount, AH->compression);
+			 AH->tocCount, AH->compression.alg);
 
 	switch (AH->format)
 	{
@@ -1487,7 +1490,7 @@ archprintf(Archive *AH, const char *fmt,...)
  *******************************/
 
 static void
-SetOutput(ArchiveHandle *AH, const char *filename, int compression)
+SetOutput(ArchiveHandle *AH, const char *filename, Compress *compression)
 {
 	int			fn;
 
@@ -1510,12 +1513,12 @@ SetOutput(ArchiveHandle *AH, const char *filename, int compression)
 
 	/* If compression explicitly requested, use gzopen */
 #ifdef HAVE_LIBZ
-	if (compression != 0)
+	if (compression->alg != COMPR_ALG_NONE)
 	{
 		char		fmode[14];
 
 		/* Don't use PG_BINARY_x since this is zlib */
-		sprintf(fmode, "wb%d", compression);
+		sprintf(fmode, "wb%d", compression->level);
 		if (fn >= 0)
 			AH->OF = gzdopen(dup(fn), fmode);
 		else
@@ -2259,7 +2262,7 @@ _discoverArchiveFormat(ArchiveHandle *AH)
  */
 static ArchiveHandle *
 _allocAH(const char *FileSpec, const ArchiveFormat fmt,
-		 const int compression, bool dosync, ArchiveMode mode,
+		 Compress *compression, bool dosync, ArchiveMode mode,
 		 SetupWorkerPtrType setupWorkerPtr)
 {
 	ArchiveHandle *AH;
@@ -2310,9 +2313,8 @@ _allocAH(const char *FileSpec, const ArchiveFormat fmt,
 	AH->toc->prev = AH->toc;
 
 	AH->mode = mode;
-	AH->compression = compression;
+	AH->compression = *compression;
 	AH->dosync = dosync;
-
 	memset(&(AH->sqlparse), 0, sizeof(AH->sqlparse));
 
 	/* Open stdout with no compression for AH output handle */
@@ -2325,7 +2327,7 @@ _allocAH(const char *FileSpec, const ArchiveFormat fmt,
 	 * Force stdin/stdout into binary mode if that is what we are using.
 	 */
 #ifdef WIN32
-	if ((fmt != archNull || compression != 0) &&
+	if ((fmt != archNull || compression->alg != COMPR_ALG_NONE) &&
 		(AH->fSpec == NULL || strcmp(AH->fSpec, "") == 0))
 	{
 		if (mode == archModeWrite)
@@ -3741,7 +3743,7 @@ WriteHead(ArchiveHandle *AH)
 	AH->WriteBytePtr(AH, AH->intSize);
 	AH->WriteBytePtr(AH, AH->offSize);
 	AH->WriteBytePtr(AH, AH->format);
-	WriteInt(AH, AH->compression);
+	WriteInt(AH, AH->compression.alg);
 	crtm = *localtime(&AH->createDate);
 	WriteInt(AH, crtm.tm_sec);
 	WriteInt(AH, crtm.tm_min);
@@ -3816,15 +3818,15 @@ ReadHead(ArchiveHandle *AH)
 	if (AH->version >= K_VERS_1_2)
 	{
 		if (AH->version < K_VERS_1_4)
-			AH->compression = AH->ReadBytePtr(AH);
+			AH->compression.alg = AH->ReadBytePtr(AH);
 		else
-			AH->compression = ReadInt(AH);
+			AH->compression.alg = ReadInt(AH);
 	}
 	else
-		AH->compression = Z_DEFAULT_COMPRESSION;
+		AH->compression.alg = Z_DEFAULT_COMPRESSION;
 
 #ifndef HAVE_LIBZ
-	if (AH->compression != 0)
+	if (AH->compression.alg != COMPR_ALG_NONE)
 		pg_log_warning("archive is compressed, but this installation does not support compression -- no data will be available");
 #endif
 
diff --git a/src/bin/pg_dump/pg_backup_archiver.h b/src/bin/pg_dump/pg_backup_archiver.h
index 1a229ebedb..641ba2d043 100644
--- a/src/bin/pg_dump/pg_backup_archiver.h
+++ b/src/bin/pg_dump/pg_backup_archiver.h
@@ -47,8 +47,6 @@
 #define GZWRITE(p, s, n, fh) (fwrite(p, s, n, fh) * (s))
 #define GZREAD(p, s, n, fh) fread(p, s, n, fh)
 #define GZEOF(fh)	feof(fh)
-/* this is just the redefinition of a libz constant */
-#define Z_DEFAULT_COMPRESSION (-1)
 
 typedef struct _z_stream
 {
@@ -62,7 +60,6 @@ typedef z_stream *z_streamp;
 
 #ifdef HAVE_LIBZSTD
 #include <zstd.h>
-#define ZSTD_COMPRESSION	-2
 #endif	/* HAVE_LIBZSTD */
 
 /* Data block types */
@@ -334,13 +331,7 @@ struct _archiveHandle
 	DumpId	   *tableDataId;	/* TABLE DATA ids, indexed by table dumpId */
 
 	struct _tocEntry *currToc;	/* Used when dumping data */
-	int			compression;	/*---------
-								 * Compression requested on open
-								 * Possible values for compression:
-								 * -2	ZSTD_COMPRESSION
-								 * -1	Z_DEFAULT_COMPRESSION
-								 *  0	COMPRESSION_NONE
-								 * 1-9 levels for gzip compression */
+	Compress	compression;	/* Compression requested on open */
 	bool		dosync;			/* data requested to be synced on sight */
 	ArchiveMode mode;			/* File mode - r or w */
 	void	   *formatData;		/* Header data specific to file format */
diff --git a/src/bin/pg_dump/pg_backup_custom.c b/src/bin/pg_dump/pg_backup_custom.c
index 77d402c323..55a887a236 100644
--- a/src/bin/pg_dump/pg_backup_custom.c
+++ b/src/bin/pg_dump/pg_backup_custom.c
@@ -298,7 +298,7 @@ _StartData(ArchiveHandle *AH, TocEntry *te)
 	_WriteByte(AH, BLK_DATA);	/* Block type */
 	WriteInt(AH, te->dumpId);	/* For sanity check */
 
-	ctx->cs = AllocateCompressor(AH->compression, _CustomWriteFunc);
+	ctx->cs = AllocateCompressor(&AH->compression, _CustomWriteFunc);
 }
 
 /*
@@ -377,7 +377,7 @@ _StartBlob(ArchiveHandle *AH, TocEntry *te, Oid oid)
 
 	WriteInt(AH, oid);
 
-	ctx->cs = AllocateCompressor(AH->compression, _CustomWriteFunc);
+	ctx->cs = AllocateCompressor(&AH->compression, _CustomWriteFunc);
 }
 
 /*
@@ -566,7 +566,7 @@ _PrintTocData(ArchiveHandle *AH, TocEntry *te)
 static void
 _PrintData(ArchiveHandle *AH)
 {
-	ReadDataFromArchive(AH, AH->compression, _CustomReadFunc);
+	ReadDataFromArchive(AH, _CustomReadFunc);
 }
 
 static void
diff --git a/src/bin/pg_dump/pg_backup_directory.c b/src/bin/pg_dump/pg_backup_directory.c
index 6e55cb425e..dfa36d5dc1 100644
--- a/src/bin/pg_dump/pg_backup_directory.c
+++ b/src/bin/pg_dump/pg_backup_directory.c
@@ -202,7 +202,7 @@ InitArchiveFmt_Directory(ArchiveHandle *AH)
 
 		setFilePath(AH, fname, "toc.dat");
 
-		tocFH = cfopen_read(fname, PG_BINARY_R, AH->compression);
+		tocFH = cfopen_read(fname, PG_BINARY_R, &AH->compression);
 		if (tocFH == NULL)
 			fatal("could not open input file \"%s\": %m", fname);
 
@@ -327,7 +327,7 @@ _StartData(ArchiveHandle *AH, TocEntry *te)
 
 	setFilePath(AH, fname, tctx->filename);
 
-	ctx->dataFH = cfopen_write(fname, PG_BINARY_W, AH->compression);
+	ctx->dataFH = cfopen_write(fname, PG_BINARY_W, &AH->compression);
 	if (ctx->dataFH == NULL)
 		fatal("could not open output file \"%s\": %m", fname);
 }
@@ -388,12 +388,12 @@ _PrintFileData(ArchiveHandle *AH, char *filename)
 	if (!filename)
 		return;
 
-	cfp = cfopen_read(filename, PG_BINARY_R, AH->compression);
+	cfp = cfopen_read(filename, PG_BINARY_R, &AH->compression);
 
 	if (!cfp)
 		fatal("could not open input file \"%s\": %m", filename);
 
-	buf = pg_malloc(ZLIB_OUT_SIZE);
+	buf = pg_malloc(ZLIB_OUT_SIZE); // XXX
 	buflen = ZLIB_OUT_SIZE;
 
 	while ((cnt = cfread(buf, buflen, cfp)))
@@ -435,12 +435,13 @@ _LoadBlobs(ArchiveHandle *AH)
 	lclContext *ctx = (lclContext *) AH->formatData;
 	char		fname[MAXPGPATH];
 	char		line[MAXPGPATH];
+	Compress	nocompression = {0};
 
 	StartRestoreBlobs(AH);
 
 	setFilePath(AH, fname, "blobs.toc");
 
-	ctx->blobsTocFH = cfopen_read(fname, PG_BINARY_R, AH->compression);
+	ctx->blobsTocFH = cfopen_read(fname, PG_BINARY_R, &nocompression);
 
 	if (ctx->blobsTocFH == NULL)
 		fatal("could not open large object TOC file \"%s\" for input: %m",
@@ -573,6 +574,7 @@ _CloseArchive(ArchiveHandle *AH)
 	{
 		cfp		   *tocFH;
 		char		fname[MAXPGPATH];
+		Compress	nocompression = {0};
 
 		setFilePath(AH, fname, "toc.dat");
 
@@ -580,7 +582,7 @@ _CloseArchive(ArchiveHandle *AH)
 		ctx->pstate = ParallelBackupStart(AH);
 
 		/* The TOC is always created uncompressed */
-		tocFH = cfopen_write(fname, PG_BINARY_W, 0);
+		tocFH = cfopen_write(fname, PG_BINARY_W, &nocompression);
 		if (tocFH == NULL)
 			fatal("could not open output file \"%s\": %m", fname);
 		ctx->dataFH = tocFH;
@@ -639,11 +641,12 @@ _StartBlobs(ArchiveHandle *AH, TocEntry *te)
 {
 	lclContext *ctx = (lclContext *) AH->formatData;
 	char		fname[MAXPGPATH];
+	Compress	nocompression = {0};
 
 	setFilePath(AH, fname, "blobs.toc");
 
 	/* The blob TOC file is never compressed */
-	ctx->blobsTocFH = cfopen_write(fname, "ab", 0);
+	ctx->blobsTocFH = cfopen_write(fname, "ab", &nocompression);
 	if (ctx->blobsTocFH == NULL)
 		fatal("could not open output file \"%s\": %m", fname);
 }
@@ -661,7 +664,7 @@ _StartBlob(ArchiveHandle *AH, TocEntry *te, Oid oid)
 
 	snprintf(fname, MAXPGPATH, "%s/blob_%u.dat", ctx->directory, oid);
 
-	ctx->dataFH = cfopen_write(fname, PG_BINARY_W, AH->compression);
+	ctx->dataFH = cfopen_write(fname, PG_BINARY_W, &AH->compression);
 
 	if (ctx->dataFH == NULL)
 		fatal("could not open output file \"%s\": %m", fname);
diff --git a/src/bin/pg_dump/pg_backup_tar.c b/src/bin/pg_dump/pg_backup_tar.c
index 54e708875c..79e528ea89 100644
--- a/src/bin/pg_dump/pg_backup_tar.c
+++ b/src/bin/pg_dump/pg_backup_tar.c
@@ -196,10 +196,10 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 
 		/*
 		 * We don't support compression because reading the files back is not
-		 * possible since gzdopen uses buffered IO which totally screws file
+		 * possible since gzdopen uses buffered IO which totally screws file XXX
 		 * positioning.
 		 */
-		if (AH->compression != 0)
+		if (AH->compression.alg != COMPR_ALG_NONE)
 			fatal("compression is not supported by tar archive format");
 	}
 	else
@@ -241,6 +241,22 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 	}
 }
 
+/*
+ * Return a string for the given AH's compression.
+ * The string is statically allocated.
+ */
+static const char *
+compress_suffix(ArchiveHandle *AH)
+{
+	CompressionAlgorithm alg;
+	alg = AH->compression.alg;
+	if (alg == COMPR_ALG_LIBZ)
+		return ".gz";
+	if (alg == COMPR_ALG_ZSTD)
+		return ".zstd";
+	return "";
+}
+
 /*
  * - Start a new TOC entry
  *	 Setup the output file name.
@@ -254,14 +270,8 @@ _ArchiveEntry(ArchiveHandle *AH, TocEntry *te)
 	ctx = (lclTocEntry *) pg_malloc0(sizeof(lclTocEntry));
 	if (te->dataDumper != NULL)
 	{
-#ifdef HAVE_LIBZ
-		if (AH->compression == 0)
-			sprintf(fn, "%d.dat", te->dumpId);
-		else
-			sprintf(fn, "%d.dat.gz", te->dumpId);
-#else
-		sprintf(fn, "%d.dat", te->dumpId);
-#endif
+		const char *suffix = compress_suffix(AH);
+		sprintf(fn, "%d.dat%s", te->dumpId, suffix);
 		ctx->filename = pg_strdup(fn);
 	}
 	else
@@ -352,7 +362,7 @@ tarOpen(ArchiveHandle *AH, const char *filename, char mode)
 
 #ifdef HAVE_LIBZ
 
-		if (AH->compression == 0)
+		if (AH->compression.alg == COMPR_ALG_NONE)
 			tm->nFH = ctx->tarFH;
 		else
 			fatal("compression is not supported by tar archive format");
@@ -413,9 +423,9 @@ tarOpen(ArchiveHandle *AH, const char *filename, char mode)
 
 #ifdef HAVE_LIBZ
 
-		if (AH->compression != 0)
+		if (AH->compression.alg != COMPR_ALG_NONE)
 		{
-			sprintf(fmode, "wb%d", AH->compression);
+			sprintf(fmode, "wb%d", AH->compression.level);
 			tm->zFH = gzdopen(dup(fileno(tm->tmpFH)), fmode);
 			if (tm->zFH == NULL)
 				fatal("could not open temporary file");
@@ -443,7 +453,7 @@ tarClose(ArchiveHandle *AH, TAR_MEMBER *th)
 	/*
 	 * Close the GZ file since we dup'd. This will flush the buffers.
 	 */
-	if (AH->compression != 0)
+	if (AH->compression.alg != COMPR_ALG_NONE)
 		if (GZCLOSE(th->zFH) != 0)
 			fatal("could not close tar member");
 
@@ -868,7 +878,7 @@ _CloseArchive(ArchiveHandle *AH)
 		memcpy(ropt, AH->public.ropt, sizeof(RestoreOptions));
 		ropt->filename = NULL;
 		ropt->dropSchema = 1;
-		ropt->compression = 0;
+		ropt->compression.alg = COMPR_ALG_NONE;
 		ropt->superuser = NULL;
 		ropt->suppressDumpWarnings = true;
 
@@ -952,16 +962,12 @@ _StartBlob(ArchiveHandle *AH, TocEntry *te, Oid oid)
 	lclContext *ctx = (lclContext *) AH->formatData;
 	lclTocEntry *tctx = (lclTocEntry *) te->formatData;
 	char		fname[255];
-	char	   *sfx;
+	const char *sfx;
 
 	if (oid == 0)
 		fatal("invalid OID for large object (%u)", oid);
 
-	if (AH->compression != 0)
-		sfx = ".gz";
-	else
-		sfx = "";
-
+	sfx = compress_suffix(AH);
 	sprintf(fname, "blob_%u.dat%s", oid, sfx);
 
 	tarPrintf(ctx->blobToc, "%u %s\n", oid, fname);
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index f21eb021c7..c8662f731b 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -59,6 +59,7 @@
 #include "getopt_long.h"
 #include "libpq/libpq-fs.h"
 #include "parallel.h"
+#include "compress_io.h"
 #include "pg_backup_db.h"
 #include "pg_backup_utils.h"
 #include "pg_dump.h"
@@ -297,6 +298,99 @@ static void setupDumpWorker(Archive *AHX);
 static TableInfo *getRootTableInfo(TableInfo *tbinfo);
 
 
+/* Parse the string into compression options */
+static void
+parse_compression(const char *optarg, Compress *compress)
+{
+	if (optarg[0] == '0' && optarg[1] == '\0')
+		compress->alg = COMPR_ALG_NONE;
+	else if ((optarg[0] > '0' && optarg[0] <= '9') ||
+		optarg[0] == '-')
+	{
+		compress->alg = COMPR_ALG_LIBZ;
+		compress->level = atoi(optarg);
+		if (optarg[1] != '\0')
+		{
+			pg_log_error("compression level must be in range 0..9");
+			exit_nicely(1);
+		}
+	}
+	else
+	{
+		/* Parse a more flexible string like level=3 alg=zlib opts=long */
+		for (;;)
+		{
+			char *eq = strchr(optarg, '=');
+			int len;
+
+			if (eq == NULL)
+			{
+				pg_log_error("compression options must be key=value: %s", optarg);
+				exit_nicely(1);
+			}
+
+			len = eq - optarg;
+			if (strncmp(optarg, "alg", len) == 0)
+			{
+				if (strchr(eq, ' '))
+					len = strchr(eq, ' ') - eq - 1;
+				else
+					len = strlen(eq) - len;
+				if (strncmp(1+eq, "zlib", len) == 0 ||
+						strncmp(1+eq, "libz", len) == 0)
+					compress->alg = COMPR_ALG_LIBZ;
+				else if (strncmp(1+eq, "zstd", len) == 0)
+					compress->alg = COMPR_ALG_ZSTD;
+				else
+				{
+					pg_log_error("unknown compression algorithm: %s", 1+eq);
+					exit_nicely(1);
+				}
+			}
+			else if (strncmp(optarg, "level", len) == 0)
+				compress->level = atoi(1+eq);
+			else if (strncmp(optarg, "opt", len) == 0)
+			{
+				if (strchr(eq, ' '))
+					len = strchr(eq, ' ') - eq - 1;
+				else
+					len = strlen(eq) - len;
+				if (strncmp(1+eq, "zstdlong", len) == 0)
+					compress->zstdlong = true;
+				else
+				{
+					pg_log_error("unknown compression option: %s", 1+eq);
+					exit_nicely(1);
+				}
+			}
+			else
+			{
+				pg_log_error("unknown compression setting: %s", optarg);
+				exit_nicely(1);
+			}
+
+			optarg = strchr(eq, ' ');
+			if (!optarg++)
+				break;
+		}
+
+		/* zstd will check its own compression level later */
+		if (compress->alg != COMPR_ALG_ZSTD)
+		{
+			if (compress->level < 0 || compress->level > 9)
+			{
+				pg_log_error("compression level must be in range 0..9");
+				exit_nicely(1);
+			}
+			if (compress->zstdlong)
+			{
+				pg_log_error("compression option not supported with this algorithm");
+				exit_nicely(1);
+			}
+		}
+	}
+}
+
 int
 main(int argc, char **argv)
 {
@@ -319,7 +413,7 @@ main(int argc, char **argv)
 	char	   *use_role = NULL;
 	long		rowsPerInsert;
 	int			numWorkers = 1;
-	int			compressLevel = -1;
+	Compress	compress = { .alg = COMPR_ALG_DEFAULT };
 	int			plainText = 0;
 	ArchiveFormat archiveFormat = archUnknown;
 	ArchiveMode archiveMode;
@@ -532,13 +626,7 @@ main(int argc, char **argv)
 				break;
 
 			case 'Z':			/* Compression Level */
-				compressLevel = atoi(optarg);
-				if ((compressLevel < 0 || compressLevel > 9) &&
-					compressLevel != -2)
-				{
-					pg_log_error("compression level must be in range 0..9");
-					exit_nicely(1);
-				}
+				parse_compression(optarg, &compress);
 				break;
 
 			case 0:
@@ -680,20 +768,39 @@ main(int argc, char **argv)
 		plainText = 1;
 
 	/* Custom and directory formats are compressed by default, others not */
-	if (compressLevel == -1)
+	if (compress.alg == COMPR_ALG_DEFAULT)
 	{
-#ifdef HAVE_LIBZ
 		if (archiveFormat == archCustom || archiveFormat == archDirectory)
-			compressLevel = Z_DEFAULT_COMPRESSION;
-		else
+		{
+#ifdef HAVE_LIBZ
+			compress.alg = COMPR_ALG_LIBZ;
 #endif
-			compressLevel = 0;
+#ifdef HAVE_LIBZSTD
+			compress.alg = COMPR_ALG_ZSTD; // Set default for testing purposes
+#endif
+			compress.level = 3;
+		}
+		else
+		{
+			compress.alg = COMPR_ALG_NONE;
+			compress.level = 0;
+		}
 	}
 
 #ifndef HAVE_LIBZ
-	if (compressLevel != 0)
+	if (compress.alg == COMPR_ALG_LIBZ)
+	{
 		pg_log_warning("requested compression not available in this installation -- archive will be uncompressed");
-	compressLevel = 0;
+		compress.alg = 0;
+	}
+#endif
+
+#ifndef HAVE_LIBZ
+	if (compress.alg == COMPR_ALG_ZSTD)
+	{
+		pg_log_warning("requested compression not available in this installation -- archive will be uncompressed");
+		compress.alg = 0;
+	}
 #endif
 
 	/*
@@ -724,7 +831,7 @@ main(int argc, char **argv)
 		fatal("option --index-collation-versions-unknown only works in binary upgrade mode");
 
 	/* Open the output file */
-	fout = CreateArchive(filename, archiveFormat, compressLevel, dosync,
+	fout = CreateArchive(filename, archiveFormat, &compress, dosync,
 						 archiveMode, setupDumpWorker);
 
 	/* Make dump options accessible right away */
@@ -958,10 +1065,7 @@ main(int argc, char **argv)
 	ropt->sequence_data = dopt.sequence_data;
 	ropt->binary_upgrade = dopt.binary_upgrade;
 
-	if (compressLevel == -1)
-		ropt->compression = 0;
-	else
-		ropt->compression = compressLevel;
+	ropt->compression = compress;
 
 	ropt->suppressDumpWarnings = true;	/* We've already shown them */
 
-- 
2.17.0

