From 73d58dfb71c14109402d396d0a388135be42bb6c Mon Sep 17 00:00:00 2001
From: Jacob Burroughs <jburroughs@instructure.com>
Date: Mon, 18 Dec 2023 14:15:18 -0600
Subject: [PATCH v3 4/5] Add basic test for compression functionality

---
 src/Makefile.global.in                    |  2 +
 src/interfaces/libpq/Makefile             |  7 +-
 src/interfaces/libpq/meson.build          |  8 +-
 src/interfaces/libpq/t/005_compression.pl | 94 +++++++++++++++++++++++
 4 files changed, 108 insertions(+), 3 deletions(-)
 create mode 100644 src/interfaces/libpq/t/005_compression.pl

diff --git a/src/Makefile.global.in b/src/Makefile.global.in
index f8e461cbad..e8fad07936 100644
--- a/src/Makefile.global.in
+++ b/src/Makefile.global.in
@@ -195,9 +195,11 @@ with_ldap	= @with_ldap@
 with_libxml	= @with_libxml@
 with_libxslt	= @with_libxslt@
 with_llvm	= @with_llvm@
+with_lz4	= @with_lz4@
 with_system_tzdata = @with_system_tzdata@
 with_uuid	= @with_uuid@
 with_zlib	= @with_zlib@
+with_zstd	= @with_zstd@
 enable_rpath	= @enable_rpath@
 enable_nls	= @enable_nls@
 enable_debug	= @enable_debug@
diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile
index df8a1f31b2..12dc5868b6 100644
--- a/src/interfaces/libpq/Makefile
+++ b/src/interfaces/libpq/Makefile
@@ -14,6 +14,9 @@ top_builddir = ../../..
 include $(top_builddir)/src/Makefile.global
 
 export with_ssl
+export with_zlib
+export with_lz4
+export with_zstd
 
 PGFILEDESC = "PostgreSQL Access Library"
 
@@ -78,9 +81,9 @@ endif
 # that are built correctly for use in a shlib.
 SHLIB_LINK_INTERNAL = -lpgcommon_shlib -lpgport_shlib
 ifneq ($(PORTNAME), win32)
-SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -lsocket -lnsl -lresolv -lintl -lm, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS)
+SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -lsocket -lnsl -lresolv -lintl -lm -lz -llz4 -lzstd, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS)
 else
-SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi32 -lssl -lsocket -lnsl -lresolv -lintl -lm $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE)
+SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi32 -lssl -lsocket -lnsl -lresolv -lintl -lm -lz -llz4 -lzstd $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE)
 endif
 ifeq ($(PORTNAME), win32)
 SHLIB_LINK += -lshell32 -lws2_32 -lsecur32 $(filter -leay32 -lssleay32 -lcomerr32 -lkrb5_32, $(LIBS))
diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build
index 80e6a15adf..e9b3c22a52 100644
--- a/src/interfaces/libpq/meson.build
+++ b/src/interfaces/libpq/meson.build
@@ -118,8 +118,14 @@ tests += {
       't/002_api.pl',
       't/003_load_balance_host_list.pl',
       't/004_load_balance_dns.pl',
+      't/005_compression.pl',
     ],
-    'env': {'with_ssl': ssl_library},
+    'env': {
+      'with_ssl': ssl_library,
+      'with_zlib': zlib.found() ? 'yes' : 'no',
+      'with_lz4': lz4.found() ? 'yes' : 'no',
+      'with_zstd': zstd.found() ? 'yes' : 'no',
+    },
   },
 }
 
diff --git a/src/interfaces/libpq/t/005_compression.pl b/src/interfaces/libpq/t/005_compression.pl
new file mode 100644
index 0000000000..95c8575d80
--- /dev/null
+++ b/src/interfaces/libpq/t/005_compression.pl
@@ -0,0 +1,94 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+use strict;
+use warnings FATAL => 'all';
+use Config;
+use PostgreSQL::Test::Utils;
+use PostgreSQL::Test::Cluster;
+use Test::More;
+
+my $withZlib = $ENV{with_zlib} eq 'yes';
+my $withLz4 = $ENV{with_lz4} eq 'yes';
+my $withZstd = $ENV{with_zstd} eq 'yes';
+if (!$withZlib && !$withLz4 && !$withZstd)
+{
+	plan skip_all => 'no compression methods available';
+}
+
+my $node = PostgreSQL::Test::Cluster->new('primary');
+$node->init;
+$node->append_conf('postgresql.conf', "libpq_compression = off");
+$node->start;
+
+# Use a string that any reasonable compression algorithm will compress
+my $compressableString = "a" x 1000;
+
+my $result;
+$result = $node->safe_psql("postgres",
+	"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+	connstr => $node->connstr . " compression=on");
+is( (split "\n", $result)[-1], "f|f", 'successfully does not compress if server-side compression is disabled');
+
+$node->append_conf('postgresql.conf', "libpq_compression = on");
+$node->reload;
+
+$result = $node->safe_psql("postgres",
+	"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+	connstr => $node->connstr . " compression=on");
+is( (split "\n", $result)[-1], "t|t", 'successfully compresses bidirectionally with default algorithms');
+
+my $test_algorithm;
+if ($withZlib)
+{
+	$test_algorithm = "gzip";
+}
+elsif($withLz4)
+{
+	$test_algorithm = "lz4";
+}
+elsif($withZstd)
+{
+	$test_algorithm = "zstd";
+}
+$result = $node->safe_psql("postgres",
+	"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+	connstr => $node->connstr . " compression=$test_algorithm:compress=off");
+is( (split "\n", $result)[-1], "f|t", 'successfully compresses unidirectionally with client compression disabled');
+
+$node->append_conf('postgresql.conf', "libpq_compression = '$test_algorithm:compress=off'");
+$node->reload;
+$result = $node->safe_psql("postgres",
+	"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+	connstr => "compression=on");
+is( (split "\n", $result)[-1], "t|f", 'successfully compresses unidirectionally with server compression algorithm');
+
+$node->append_conf('postgresql.conf', 'libpq_compression = on');
+$node->reload;
+
+SKIP: {
+	skip "gzip not available", 1 unless $ENV{with_zlib};
+
+	$result = $node->safe_psql("postgres",
+		"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+		connstr => $node->connstr . " compression=gzip");
+	is( (split "\n", $result)[-1], "t|t", 'successfully compresses bidirectionally with gzip');
+}
+
+SKIP: {
+	skip "lz4 not available", 1 unless $ENV{with_lz4};
+
+	$result = $node->safe_psql("postgres",
+		"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+		connstr => $node->connstr . " compression=lz4");
+	is( (split "\n", $result)[-1], "t|t", 'successfully compresses bidirectionally with lz4');
+}
+
+SKIP: {
+	skip "zstd not available", 1 unless $ENV{with_zstd};
+
+	$result = $node->safe_psql("postgres",
+		"SELECT '$compressableString'; SELECT rx_socket_bytes < rx_pq_bytes, tx_socket_bytes < tx_pq_bytes FROM pg_stat_network_traffic;",
+		connstr => $node->connstr . " compression=zstd");
+	is( (split "\n", $result)[-1], "t|t", 'successfully compresses bidirectionally with zstd');
+}
+
+done_testing();
-- 
2.42.0

