From 060acad30a8b9ee04d310e44bc2d788400f6c759 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Wed, 25 Jun 2025 14:53:25 -0700 Subject: [PATCH 01/10] Exclude libss/crypto from libpq On linux, we have the following dep chain: aws-dsql-auth -> aws-c-http -> aws-c-io -> aws-lc We only use aws-c-http for request signing, but we cannot make aws-c-io conditionally required. aws-c-io only works with aws-lc on linux, which means we land up in a situation where we have two libcryptos (postgres wants openssl). Our overall solution is to make fe-dsql-auth.o a self-contained file (hiding the aws-dsql-auth dep chain) and explicitly exclude the libs we don't want. --- src/interfaces/libpq/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 1d0b9eee226ed..4e016309f34ed 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -81,7 +81,7 @@ OBJS += \ endif AWS_DSQL_AUTH_CPPFLAGS := $(addprefix -I,$(shell find $(top_srcdir)/aws-dsql-auth/ -type d -name include -not -path '*/test*')) -AWS_DSQL_AUTH_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/ -type f -name '*.a' -not -path '*/test*') +AWS_DSQL_AUTH_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/ -type f -name '*.a' -not -path '*/test*' -not -name libssl.a -not -name libcrypto.a) # Custom rule for fe-dsql-auth.o that includes AWS libraries fe-dsql-auth.o: fe-dsql-auth.c From 2a8eaec4b39d5acfd6273f32e72577c23cf2d779 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Thu, 26 Jun 2025 14:49:44 -0700 Subject: [PATCH 02/10] pbbench support for dsql Doesn't handle occ yet --- scripts/build-dsql.sh | 19 ++++++- src/bin/pgbench/pgbench.c | 109 +++++++++++++++++++++++++++++--------- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/scripts/build-dsql.sh b/scripts/build-dsql.sh index 6deb2136aedc4..39fc3101436e6 100755 --- a/scripts/build-dsql.sh +++ b/scripts/build-dsql.sh @@ -81,10 +81,14 @@ echo "Step 4: Building psql..." make -C src/bin/psql echo " psql built successfully!" +# Step 5: Build pgbench +echo "Step 5: Building pgbench..." +make -C src/bin/pgbench +echo " pgbench built successfully!" + # Final instructions echo "" echo "Build completed successfully!" -echo "" echo "To run psql with DSQL authentication, use the following command:" echo "" echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" @@ -95,4 +99,17 @@ echo "" echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" echo " ./src/bin/psql/psql --dsql \"dbname=postgres user=admin host=your-dsql-endpoint.example.com\"" echo "" +echo "To run pgbench with DSQL authentication, use the following command:" +echo "" +echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" +echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres" +echo "" +echo "For example, to run a simple benchmark test:" +echo "" +echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" +echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --initialize --scale=1" +echo "" +echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" +echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --time=60 --client=10" +echo "" echo "Note: You need to have AWS credentials configured in your environment for DSQL authentication to work." diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c index 497a936c141f3..8c4617917bb34 100644 --- a/src/bin/pgbench/pgbench.c +++ b/src/bin/pgbench/pgbench.c @@ -771,6 +771,8 @@ static bool verbose_errors = false; /* print verbose messages of all errors */ static bool exit_on_abort = false; /* exit when any client is aborted */ +static bool dsql = false; /* --dsql command line option */ + /* Builtin test scripts */ typedef struct BuiltinScript { @@ -849,6 +851,10 @@ static const PsqlScanCallbacks pgbench_callbacks = { NULL, /* don't need get_variable functionality */ }; +static bool is_dsql() { + return strcmp(getenv("PGDSQL"), "1") == 0; +} + static char get_table_relkind(PGconn *con, const char *table) { @@ -4852,26 +4858,26 @@ initCreateTables(PGconn *con) static const struct ddlinfo DDLs[] = { { "pgbench_history", - "tid int,bid int,aid int,delta int,mtime timestamp,filler char(22)", - "tid int,bid int,aid bigint,delta int,mtime timestamp,filler char(22)", + "id uuid default gen_random_uuid(),tid int,bid int,aid int,delta int,mtime timestamp,filler char(22)", + "id uuid default gen_random_uuid(),tid int,bid int,aid bigint,delta int,mtime timestamp,filler char(22)", 0 }, { "pgbench_tellers", - "tid int not null,bid int,tbalance int,filler char(84)", - "tid int not null,bid int,tbalance int,filler char(84)", + "tid int primary key,bid int,tbalance int,filler char(84)", + "tid int primary key,bid int,tbalance int,filler char(84)", 1 }, { "pgbench_accounts", - "aid int not null,bid int,abalance int,filler char(84)", - "aid bigint not null,bid int,abalance int,filler char(84)", + "aid int primary key,bid int,abalance int,filler char(84)", + "aid bigint primary key,bid int,abalance int,filler char(84)", 1 }, { "pgbench_branches", - "bid int not null,bbalance int,filler char(88)", - "bid int not null,bbalance int,filler char(88)", + "bid int primary key,bbalance int,filler char(88)", + "bid int primary key,bbalance int,filler char(88)", 1 } }; @@ -4892,14 +4898,16 @@ initCreateTables(PGconn *con) ddl->table, (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols); - /* Partition pgbench_accounts table */ - if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0) - appendPQExpBuffer(&query, - " partition by %s (aid)", PARTITION_METHOD[partition_method]); - else if (ddl->declare_fillfactor) - { - /* fillfactor is only expected on actual tables */ - appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); + if (!is_dsql()) { + /* Partition pgbench_accounts table */ + if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0) + appendPQExpBuffer(&query, + " partition by %s (aid)", PARTITION_METHOD[partition_method]); + else if (ddl->declare_fillfactor) + { + /* fillfactor is only expected on actual tables */ + appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); + } } if (tablespace != NULL) @@ -4926,11 +4934,17 @@ initCreateTables(PGconn *con) static void initTruncateTables(PGconn *con) { - executeStatement(con, "truncate table " + if (!is_dsql()) { + executeStatement(con, "truncate table " "pgbench_accounts, " "pgbench_branches, " "pgbench_history, " "pgbench_tellers"); + } else { + // TODO: DSQL mode needs to do deletes in batches. + // This is not implemented because it's relatively easy to simply + // drop all tables and re-init with -i. + } } static void @@ -4984,8 +4998,9 @@ initPopulateTable(PGconn *con, const char *table, int64 base, initPQExpBuffer(&sql); /* Use COPY with FREEZE on v14 and later for all ordinary tables */ + // XXX: DSQL doesn't support (or need) freeze. if ((PQserverVersion(con) >= 140000) && - get_table_relkind(con, table) == RELKIND_RELATION) + get_table_relkind(con, table) == RELKIND_RELATION && !is_dsql()) copy_statement_fmt = "copy %s from stdin with (freeze on)"; @@ -4995,16 +5010,22 @@ initPopulateTable(PGconn *con, const char *table, int64 base, else if (n == -1) pg_fatal("invalid format string"); - res = PQexec(con, copy_statement); - - if (PQresultStatus(res) != PGRES_COPY_IN) - pg_fatal("unexpected copy in result: %s", PQerrorMessage(con)); - PQclear(res); - start = pg_time_now(); for (k = 0; k < total; k++) { + if (k == 0) { + if (is_dsql()) { + executeStatement(con, "begin"); + } + + res = PQexec(con, copy_statement); + + if (PQresultStatus(res) != PGRES_COPY_IN) + pg_fatal("unexpected copy in result: %s", PQerrorMessage(con)); + PQclear(res); + } + int64 j = k + 1; init_row(&sql, k); @@ -5014,6 +5035,21 @@ initPopulateTable(PGconn *con, const char *table, int64 base, if (CancelRequested) break; + // XXX: Start a new transaction every 1000 rows + if (is_dsql() && (k > 0 && k % 1000 == 0)) { + if (PQputline(con, "\\.\n")) + pg_fatal("very last PQputline failed"); + if (PQendcopy(con)) + pg_fatal("PQendcopy failed"); + + executeStatement(con, "commit"); + executeStatement(con, "begin"); + res = PQexec(con, copy_statement); + if (PQresultStatus(res) != PGRES_COPY_IN) + pg_fatal("unexpected copy in result: %s", PQerrorMessage(con)); + PQclear(res); + } + /* * If we want to stick with the original logging, print a message each * 100k inserted rows. @@ -5076,6 +5112,10 @@ initPopulateTable(PGconn *con, const char *table, int64 base, if (PQendcopy(con)) pg_fatal("PQendcopy failed"); + if (is_dsql()) { + executeStatement(con, "commit"); + } + termPQExpBuffer(&sql); } @@ -5090,11 +5130,14 @@ initGenerateDataClientSide(PGconn *con) { fprintf(stderr, "generating data (client-side)...\n"); + // XXX: On DSQL we do batch inserts within populate. + if (!is_dsql()) { /* * we do all of this in one transaction to enable the backend's * data-loading optimizations */ executeStatement(con, "begin"); + } /* truncate away any old data */ initTruncateTables(con); @@ -5107,7 +5150,9 @@ initGenerateDataClientSide(PGconn *con) initPopulateTable(con, "pgbench_tellers", ntellers, initTeller); initPopulateTable(con, "pgbench_accounts", naccounts, initAccount); + if (!is_dsql()) { executeStatement(con, "commit"); + } } /* @@ -5178,6 +5223,11 @@ initVacuum(PGconn *con) static void initCreatePKeys(PGconn *con) { + // Schema has been updated to always have PKs. + if (is_dsql()) { + return; + } + static const char *const DDLINDEXes[] = { "alter table pgbench_branches add primary key (bid)", "alter table pgbench_tellers add primary key (tid)", @@ -6705,6 +6755,7 @@ main(int argc, char **argv) {"verbose-errors", no_argument, NULL, 15}, {"exit-on-abort", no_argument, NULL, 16}, {"debug", no_argument, NULL, 17}, + {"dsql", no_argument, NULL, 18}, {NULL, 0, NULL, 0} }; @@ -7058,6 +7109,9 @@ main(int argc, char **argv) case 17: /* debug */ pg_logging_increase_verbosity(); break; + case 18: /* dsql */ + dsql = true; + break; default: /* getopt_long already emitted a complaint */ pg_log_error_hint("Try \"%s --help\" for more information.", progname); @@ -7065,6 +7119,13 @@ main(int argc, char **argv) } } + if (dsql) + { + setenv("PGDSQL", "1", 1); + is_no_vacuum = true; + foreign_keys = false; + } + /* set default script if none */ if (num_scripts == 0 && !is_init_mode) { From 6e4bdb319adb3261c4e0c7027a0545ff2307bedc Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Thu, 26 Jun 2025 16:11:01 -0700 Subject: [PATCH 03/10] pkg pgbench --- scripts/build-dsql.sh | 6 ++-- scripts/package.sh | 77 +++++++++++++++++++++++++++++++------------ 2 files changed, 59 insertions(+), 24 deletions(-) diff --git a/scripts/build-dsql.sh b/scripts/build-dsql.sh index 39fc3101436e6..0c225191ae4c9 100755 --- a/scripts/build-dsql.sh +++ b/scripts/build-dsql.sh @@ -102,14 +102,14 @@ echo "" echo "To run pgbench with DSQL authentication, use the following command:" echo "" echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" -echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres" +echo " ./src/bin/pgbench/pgbench --dsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres" echo "" echo "For example, to run a simple benchmark test:" echo "" echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" -echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --initialize --scale=1" +echo " ./src/bin/pgbench/pgbench --dsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --initialize --scale=1" echo "" echo " $LIBRARY_PATH_VAR=$(pwd)/src/interfaces/libpq \\" -echo " ./src/bin/pgbench/pgbench --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --time=60 --client=10" +echo " ./src/bin/pgbench/pgbench --dsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --time=60 --client=10" echo "" echo "Note: You need to have AWS credentials configured in your environment for DSQL authentication to work." diff --git a/scripts/package.sh b/scripts/package.sh index 851dd8ebf5c9a..262b88dd67ddd 100755 --- a/scripts/package.sh +++ b/scripts/package.sh @@ -2,8 +2,8 @@ set -e # Package PostgreSQL DSQL client for distribution -# This script creates a standalone distribution with psql (renamed to pdsql) -# and libpq that can be used without additional dependencies +# This script creates a standalone distribution with psql (renamed to pdsql), +# pgbench, and libpq that can be used without additional dependencies echo "Packaging PostgreSQL DSQL client" echo "================================" @@ -13,8 +13,10 @@ ROOT_DIR=$(pwd) BUILD_DIR="$ROOT_DIR/build" DIST_NAME="postgres-dsql" DIST_DIR="$BUILD_DIR/$DIST_NAME" -SRC_BIN="$ROOT_DIR/src/bin/psql/psql" -BINARY_NAME="pdsql" +SRC_PSQL_BIN="$ROOT_DIR/src/bin/psql/psql" +SRC_PGBENCH_BIN="$ROOT_DIR/src/bin/pgbench/pgbench" +PSQL_BINARY_NAME="pdsql" +PGBENCH_BINARY_NAME="pgbench" # Detect OS and set appropriate library paths if [[ "$OSTYPE" == "darwin"* ]]; then @@ -34,8 +36,14 @@ else fi # Check if the build artifacts exist -if [ ! -f "$SRC_BIN" ]; then - echo "Error: psql binary not found at $SRC_BIN" +if [ ! -f "$SRC_PSQL_BIN" ]; then + echo "Error: psql binary not found at $SRC_PSQL_BIN" + echo "Please run scripts/build-dsql.sh first" + exit 1 +fi + +if [ ! -f "$SRC_PGBENCH_BIN" ]; then + echo "Error: pgbench binary not found at $SRC_PGBENCH_BIN" echo "Please run scripts/build-dsql.sh first" exit 1 fi @@ -60,8 +68,11 @@ mkdir -p "$DIST_DIR/bin" mkdir -p "$DIST_DIR/lib" # Copy binaries and libraries -echo "Copying psql to $DIST_DIR/bin/$BINARY_NAME" -cp "$SRC_BIN" "$DIST_DIR/bin/$BINARY_NAME" +echo "Copying psql to $DIST_DIR/bin/$PSQL_BINARY_NAME" +cp "$SRC_PSQL_BIN" "$DIST_DIR/bin/$PSQL_BINARY_NAME" + +echo "Copying pgbench to $DIST_DIR/bin/$PGBENCH_BINARY_NAME" +cp "$SRC_PGBENCH_BIN" "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" echo "Copying libpq to $DIST_DIR/lib/" cp "$SRC_LIB" "$DIST_DIR/lib/" @@ -71,41 +82,52 @@ if [[ "$PLATFORM" == "macos" ]]; then # Copy additional dylib if it exists cp "$ROOT_DIR/src/interfaces/libpq/libpq.dylib" "$DIST_DIR/lib/" 2>/dev/null || true - # Set up correct library paths in the binary - echo "Updating library paths in $BINARY_NAME binary..." - LIBRARY_PATH=$(otool -L "$DIST_DIR/bin/$BINARY_NAME" | grep libpq | awk '{print $1}') - install_name_tool -change "$LIBRARY_PATH" "@loader_path/../lib/libpq.5.dylib" "$DIST_DIR/bin/$BINARY_NAME" + # Set up correct library paths in the binaries + echo "Updating library paths in $PSQL_BINARY_NAME binary..." + LIBRARY_PATH=$(otool -L "$DIST_DIR/bin/$PSQL_BINARY_NAME" | grep libpq | awk '{print $1}') + install_name_tool -change "$LIBRARY_PATH" "@loader_path/../lib/libpq.5.dylib" "$DIST_DIR/bin/$PSQL_BINARY_NAME" + + echo "Updating library paths in $PGBENCH_BINARY_NAME binary..." + LIBRARY_PATH=$(otool -L "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" | grep libpq | awk '{print $1}') + install_name_tool -change "$LIBRARY_PATH" "@loader_path/../lib/libpq.5.dylib" "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" # Fix library itself to refer to itself by relative path install_name_tool -id "@loader_path/libpq.5.dylib" "$DIST_DIR/lib/libpq.5.dylib" # Verify the changes echo "Verifying library path changes:" - otool -L "$DIST_DIR/bin/$BINARY_NAME" | grep libpq + otool -L "$DIST_DIR/bin/$PSQL_BINARY_NAME" | grep libpq + otool -L "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" | grep libpq otool -L "$DIST_DIR/lib/libpq.5.dylib" | grep libpq elif [[ "$PLATFORM" == "linux" ]]; then # Copy additional .so files if they exist cp "$ROOT_DIR/src/interfaces/libpq/libpq.so" "$DIST_DIR/lib/" 2>/dev/null || true - # Set up RPATH for the binary to find libraries in ../lib - echo "Setting RPATH for $BINARY_NAME binary..." - patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$BINARY_NAME" 2>/dev/null || { + # Set up RPATH for the binaries to find libraries in ../lib + echo "Setting RPATH for $PSQL_BINARY_NAME binary..." + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PSQL_BINARY_NAME" 2>/dev/null || { echo "Warning: patchelf not available. Installing patchelf..." if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y patchelf - patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$BINARY_NAME" + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PSQL_BINARY_NAME" + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" elif command -v yum >/dev/null 2>&1; then sudo yum install -y patchelf - patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$BINARY_NAME" + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PSQL_BINARY_NAME" + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" else echo "Warning: Could not install patchelf. The binary may not find libraries correctly." fi } + echo "Setting RPATH for $PGBENCH_BINARY_NAME binary..." + patchelf --set-rpath '$ORIGIN/../lib' "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" 2>/dev/null || echo "patchelf already handled above" + # Verify the changes echo "Verifying RPATH changes:" - ldd "$DIST_DIR/bin/$BINARY_NAME" | grep libpq || echo "libpq dependency check complete" + ldd "$DIST_DIR/bin/$PSQL_BINARY_NAME" | grep libpq || echo "libpq dependency check complete for pdsql" + ldd "$DIST_DIR/bin/$PGBENCH_BINARY_NAME" | grep libpq || echo "libpq dependency check complete for pgbench" fi # Create a ZIP archive @@ -171,20 +193,27 @@ mkdir -p %{buildroot}/usr/bin # Install binaries and libraries to /opt to avoid conflicts cp %{_sourcedir}/bin/pdsql %{buildroot}/opt/postgres-dsql/bin/ +cp %{_sourcedir}/bin/pgbench %{buildroot}/opt/postgres-dsql/bin/ cp %{_sourcedir}/lib/* %{buildroot}/opt/postgres-dsql/lib/ -# Create symlink in /usr/bin for easy access +# Create symlinks in /usr/bin for easy access ln -s /opt/postgres-dsql/bin/pdsql %{buildroot}/usr/bin/pdsql +ln -s /opt/postgres-dsql/bin/pgbench %{buildroot}/usr/bin/pgbench %files /opt/postgres-dsql/bin/pdsql +/opt/postgres-dsql/bin/pgbench /opt/postgres-dsql/lib/* /usr/bin/pdsql +/usr/bin/pgbench %post echo "PostgreSQL DSQL client installed successfully!" echo "Use 'pdsql' command to connect to AWS DSQL databases." echo "Example: pdsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres" +echo "" +echo "Use 'pgbench' command to run PostgreSQL benchmarks against DSQL databases." +echo "Example: pgbench --dsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --initialize --scale=1" %changelog * $(date +'%a %b %d %Y') Build System - 1.0.0-1 @@ -195,6 +224,7 @@ EOF mkdir -p "$RPM_BUILD_DIR/SOURCES/bin" mkdir -p "$RPM_BUILD_DIR/SOURCES/lib" cp "$DIST_DIR/bin/pdsql" "$RPM_BUILD_DIR/SOURCES/bin/" + cp "$DIST_DIR/bin/pgbench" "$RPM_BUILD_DIR/SOURCES/bin/" cp "$DIST_DIR/lib"/* "$RPM_BUILD_DIR/SOURCES/lib/" # Build the RPM @@ -271,10 +301,12 @@ EOF # Copy files cp "$DIST_DIR/bin/pdsql" "$DEB_PKG_DIR/opt/postgres-dsql/bin/" + cp "$DIST_DIR/bin/pgbench" "$DEB_PKG_DIR/opt/postgres-dsql/bin/" cp "$DIST_DIR/lib"/* "$DEB_PKG_DIR/opt/postgres-dsql/lib/" - # Create symlink + # Create symlinks ln -s /opt/postgres-dsql/bin/pdsql "$DEB_PKG_DIR/usr/bin/pdsql" + ln -s /opt/postgres-dsql/bin/pgbench "$DEB_PKG_DIR/usr/bin/pgbench" # Create control file cat > "$DEB_PKG_DIR/DEBIAN/control" << EOF @@ -297,6 +329,9 @@ EOF echo "PostgreSQL DSQL client installed successfully!" echo "Use 'pdsql' command to connect to AWS DSQL databases." echo "Example: pdsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres" +echo "" +echo "Use 'pgbench' command to run PostgreSQL benchmarks against DSQL databases." +echo "Example: pgbench --dsql --host=your-dsql-endpoint.example.com --user=admin --dbname=postgres --initialize --scale=1" EOF chmod 755 "$DEB_PKG_DIR/DEBIAN/postinst" From 0187de5adff741e7f6810eaa4eaa7eb41dd1e784 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Thu, 26 Jun 2025 22:47:38 -0700 Subject: [PATCH 04/10] Revert the tls filter It doesn't work because we end up with a few unnamed syms. --- src/interfaces/libpq/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 4e016309f34ed..1d0b9eee226ed 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -81,7 +81,7 @@ OBJS += \ endif AWS_DSQL_AUTH_CPPFLAGS := $(addprefix -I,$(shell find $(top_srcdir)/aws-dsql-auth/ -type d -name include -not -path '*/test*')) -AWS_DSQL_AUTH_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/ -type f -name '*.a' -not -path '*/test*' -not -name libssl.a -not -name libcrypto.a) +AWS_DSQL_AUTH_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/ -type f -name '*.a' -not -path '*/test*') # Custom rule for fe-dsql-auth.o that includes AWS libraries fe-dsql-auth.o: fe-dsql-auth.c From d61f60f5b9a0588b92694e22f2e9980e1d0abbc6 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 16:54:05 -0700 Subject: [PATCH 05/10] Bump aws-dsql-auth --- aws-dsql-auth | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws-dsql-auth b/aws-dsql-auth index 86a5f98749538..f33fce18f3864 160000 --- a/aws-dsql-auth +++ b/aws-dsql-auth @@ -1 +1 @@ -Subproject commit 86a5f9874953866409095f3f2391596ec51dd1a9 +Subproject commit f33fce18f3864e9428b10ae63c018b004b18436a From b501e2abe06135a7cc77c3d309bf40daeeb19b45 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 17:06:36 -0700 Subject: [PATCH 06/10] Use sym viz. to deal with dups --- scripts/build-dsql.sh | 2 +- src/interfaces/libpq/Makefile | 25 +++++++++++++++++++------ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/scripts/build-dsql.sh b/scripts/build-dsql.sh index 0c225191ae4c9..83173df1b8ada 100755 --- a/scripts/build-dsql.sh +++ b/scripts/build-dsql.sh @@ -51,7 +51,7 @@ else echo " aws-dsql-auth submodules already initialized." fi -if [ ! -f "aws-dsql-auth/build/aws-dsql-auth/libaws-dsql-auth.a" ]; then +if [ ! -f "aws-dsql-auth/build/install/lib64/libaws-dsql-auth.a" ]; then # Build aws-dsql-auth echo " Building aws-dsql-auth library..." cd aws-dsql-auth diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 1d0b9eee226ed..03556fcd12c49 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -80,18 +80,31 @@ OBJS += \ win32.o endif -AWS_DSQL_AUTH_CPPFLAGS := $(addprefix -I,$(shell find $(top_srcdir)/aws-dsql-auth/ -type d -name include -not -path '*/test*')) -AWS_DSQL_AUTH_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/ -type f -name '*.a' -not -path '*/test*') +AWS_DSQL_AUTH_CPPFLAGS := $(addprefix -I,$(shell find $(top_srcdir)/aws-dsql-auth/build/install -type d -name include)) +AWS_DSQL_AUTH_ALL_LIBS := $(shell find $(top_srcdir)/aws-dsql-auth/build/install/ -type f -name '*.a') -# Custom rule for fe-dsql-auth.o that includes AWS libraries +# Custom rule for fe-dsql-auth.o that hides all symbols except the public API fe-dsql-auth.o: fe-dsql-auth.c $(CC) $(CFLAGS) $(CPPFLAGS) $(AWS_DSQL_AUTH_CPPFLAGS) -c fe-dsql-auth.c -o fe-dsql-auth-temp.o ifeq ($(PORTNAME), linux) - $(LD) -r -o $@ fe-dsql-auth-temp.o --whole-archive $(AWS_DSQL_AUTH_LIBS) + # Link with AWS libraries using start-group to resolve dependencies + $(LD) -r -o fe-dsql-auth-with-aws.o fe-dsql-auth-temp.o --start-group $(AWS_DSQL_AUTH_ALL_LIBS) --end-group + # Create a list of symbols to keep (only the public API from fe-dsql-auth.h) + echo "generate_dsql_token" > keep-symbols.txt + echo "dsql_auth_cleanup" >> keep-symbols.txt + # Hide all symbols except the ones we want to keep + objcopy --keep-global-symbols=keep-symbols.txt fe-dsql-auth-with-aws.o $@ + rm -f fe-dsql-auth-temp.o fe-dsql-auth-with-aws.o keep-symbols.txt +else ifeq ($(PORTNAME), darwin) + # macOS: Use ld with exported symbols list + echo "_generate_dsql_token" > exported-symbols.txt + echo "_dsql_auth_cleanup" >> exported-symbols.txt + $(LD) -r -o $@ fe-dsql-auth-temp.o -exported_symbols_list exported-symbols.txt $(AWS_DSQL_AUTH_ALL_LIBS) + rm -f fe-dsql-auth-temp.o exported-symbols.txt else - $(LD) -r -o $@ fe-dsql-auth-temp.o $(AWS_DSQL_AUTH_LIBS) -endif + $(LD) -r -o $@ fe-dsql-auth-temp.o --whole-archive $(AWS_DSQL_AUTH_ALL_LIBS) rm -f fe-dsql-auth-temp.o +endif # Add libraries that libpq depends (or might depend) on into the # shared library link. (The order in which you list them here doesn't From 560789dfe23ba0b03a0ee3555f64b36de5639537 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 20:41:06 -0700 Subject: [PATCH 07/10] Add logging and imds support This commit wires up the aws-c-common logging infra. You can log to std{out,err} or a file, and set the level through env vars. The readme has examples. With this logging I realized that imds wasn't working due to a missing http client. That's been fixed now. --- .github/README.md | 62 ++++++++++++ src/interfaces/libpq/fe-dsql-auth.c | 144 +++++++++++++++++++++++++--- 2 files changed, 195 insertions(+), 11 deletions(-) diff --git a/.github/README.md b/.github/README.md index 2a06f9a525486..4f9cb6e604ac1 100644 --- a/.github/README.md +++ b/.github/README.md @@ -164,6 +164,68 @@ pdsql --host=your-endpoint.dsql.amazonaws.com \ - Or set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables - Verify your credentials work: `aws sts get-caller-identity` +## 🔍 Debug Logging + +For troubleshooting authentication and connection issues, enable detailed AWS SDK logging: + +### Environment Variables + +- **`AWS_LOG_LEVEL`**: Controls verbosity (NONE, FATAL, ERROR, WARN, INFO, DEBUG, TRACE) +- **`AWS_LOG_FILE`**: Controls output destination (stdout, stderr, or file path) + +### Basic Debugging + +Enable debug logging to stderr (default): +```bash +AWS_LOG_LEVEL=DEBUG pdsql --host=your-endpoint.dsql.amazonaws.com --user=admin +``` + +### Detailed Tracing + +Enable maximum verbosity for deep debugging: +```bash +AWS_LOG_LEVEL=TRACE pdsql --host=your-endpoint.dsql.amazonaws.com --user=admin +``` + +### Log to File + +Save logs to a file for analysis: +```bash +AWS_LOG_LEVEL=DEBUG AWS_LOG_FILE=/tmp/dsql-debug.log pdsql --host=your-endpoint.dsql.amazonaws.com --user=admin +``` + +### Log to stdout + +Send logs to stdout (useful for piping): +```bash +AWS_LOG_LEVEL=INFO AWS_LOG_FILE=stdout pdsql --host=your-endpoint.dsql.amazonaws.com --user=admin +``` + +### What the Logs Show + +The debug logs will reveal: +- **Token Generation**: Process of creating DSQL authentication tokens +- **AWS Region Detection**: How the region is determined from hostname or environment +- **Credentials Provider Chain**: Which credential sources are tried (environment, files, IAM roles, IMDS) +- **HTTP Infrastructure**: Event loops and network setup for IMDS on EC2 +- **Error Details**: Specific AWS SDK errors with error codes + +### Example Log Output + +``` +[INFO] [2025-06-28T03:33:55Z] Starting DSQL token generation for endpoint: your-endpoint.dsql.amazonaws.com +[DEBUG] [2025-06-28T03:33:55Z] Using AWS_REGION from environment: us-west-2 +[DEBUG] [2025-06-28T03:33:55Z] Creating credentials provider chain with bootstrap for IMDS +[INFO] [2025-06-28T03:33:55Z] Token generation successful +``` + +### Disable Logging + +To disable all logging: +```bash +AWS_LOG_LEVEL=NONE pdsql --host=your-endpoint.dsql.amazonaws.com --user=admin +``` + ### Getting Help For additional help: diff --git a/src/interfaces/libpq/fe-dsql-auth.c b/src/interfaces/libpq/fe-dsql-auth.c index 217e19c1c9c4f..c7ed75d300eaf 100644 --- a/src/interfaces/libpq/fe-dsql-auth.c +++ b/src/interfaces/libpq/fe-dsql-auth.c @@ -17,27 +17,76 @@ /* Include AWS DSQL Auth library functions */ #include +#include #include #include #include +#include +#include +#include +#include #include #include static bool aws_libs_initialized = false; +static struct aws_logger dsql_logger; +static bool dsql_logger_initialized = false; + +/* HTTP infrastructure for IMDS */ +static struct aws_event_loop_group *s_el_group = NULL; +static struct aws_host_resolver *s_host_resolver = NULL; +static struct aws_client_bootstrap *s_client_bootstrap = NULL; /* - * This function can be called to clean up AWS library resources + * Initialize DSQL logging */ static void -_dsql_auth_cleanup(void) +initialize_dsql_logging(void) { - if (aws_libs_initialized) + if (!dsql_logger_initialized) { - aws_sdkutils_library_clean_up(); - aws_auth_library_clean_up(); - aws_io_library_clean_up(); - aws_common_library_clean_up(); - aws_libs_initialized = false; + struct aws_allocator *allocator = aws_default_allocator(); + struct aws_logger_standard_options logger_options = { + .level = AWS_LOG_LEVEL_DEBUG, /* Can be controlled by environment variable */ + .file = stderr /* Log to stderr by default */ + }; + + /* Check for AWS_LOG_LEVEL environment variable */ + const char *log_level_str = getenv("AWS_LOG_LEVEL"); + if (log_level_str != NULL) + { + enum aws_log_level level; + if (aws_string_to_log_level(log_level_str, &level) == AWS_OP_SUCCESS) + { + logger_options.level = level; + } + } + + /* Check for AWS_LOG_FILE environment variable for file output */ + const char *log_file_str = getenv("AWS_LOG_FILE"); + if (log_file_str != NULL && strlen(log_file_str) > 0) + { + if (strcmp(log_file_str, "stdout") == 0) + { + logger_options.file = stdout; + } + else if (strcmp(log_file_str, "stderr") == 0) + { + logger_options.file = stderr; + } + else + { + /* Use as a filename */ + logger_options.filename = log_file_str; + logger_options.file = NULL; + } + } + + if (aws_logger_init_standard(&dsql_logger, allocator, &logger_options) == AWS_OP_SUCCESS) + { + aws_logger_set(&dsql_logger); + dsql_logger_initialized = true; + } } } @@ -47,7 +96,21 @@ _dsql_auth_cleanup(void) void dsql_auth_cleanup(void) { - _dsql_auth_cleanup(); + if (dsql_logger_initialized) + { + aws_logger_set(NULL); + aws_logger_clean_up(&dsql_logger); + dsql_logger_initialized = false; + } + + if (aws_libs_initialized) + { + aws_sdkutils_library_clean_up(); + aws_auth_library_clean_up(); + aws_io_library_clean_up(); + aws_common_library_clean_up(); + aws_libs_initialized = false; + } } /* @@ -61,13 +124,62 @@ initialize_aws_libs(void) struct aws_allocator *allocator = aws_default_allocator(); aws_common_library_init(allocator); aws_io_library_init(allocator); + aws_http_library_init(allocator); aws_auth_library_init(allocator); aws_sdkutils_library_init(allocator); + + /* Initialize HTTP infrastructure for IMDS */ + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Initializing HTTP infrastructure for IMDS"); + + s_el_group = aws_event_loop_group_new_default(allocator, 1, NULL); + if (!s_el_group) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create event loop group"); + goto error; + } + + struct aws_host_resolver_default_options resolver_options = { + .el_group = s_el_group, + .max_entries = 8, + }; + s_host_resolver = aws_host_resolver_new_default(allocator, &resolver_options); + if (!s_host_resolver) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create host resolver"); + goto error; + } + + struct aws_client_bootstrap_options bootstrap_options = { + .event_loop_group = s_el_group, + .host_resolver = s_host_resolver, + }; + s_client_bootstrap = aws_client_bootstrap_new(allocator, &bootstrap_options); + if (!s_client_bootstrap) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create client bootstrap"); + goto error; + } + aws_libs_initialized = true; + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "AWS libraries and HTTP infrastructure initialized successfully"); /* Note: We cannot use atexit() in libpq as it's not allowed to call exit-related functions. * The cleanup will be handled by explicit calls at application shutdown or by the OS. */ + return; + + error: + /* Clean up on error */ + if (s_client_bootstrap) { + aws_client_bootstrap_release(s_client_bootstrap); + s_client_bootstrap = NULL; + } + if (s_host_resolver) { + aws_host_resolver_release(s_host_resolver); + s_host_resolver = NULL; + } + if (s_el_group) { + aws_event_loop_group_release(s_el_group); + s_el_group = NULL; + } + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to initialize AWS libraries"); } } @@ -90,11 +202,14 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) const char *env_region; const char *token_str; - /* Initialize AWS libraries */ + /* Initialize AWS libraries and logging */ initialize_aws_libs(); + initialize_dsql_logging(); allocator = aws_default_allocator(); + AWS_LOGF_INFO(AWS_LS_AUTH_GENERAL, "Starting DSQL token generation for endpoint: %s", endpoint); + /* Initialize DSQL auth config */ if (aws_dsql_auth_config_init(&auth_config) != AWS_OP_SUCCESS) { if (err_msg) @@ -109,8 +224,10 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) env_region = getenv("AWS_REGION"); if (env_region != NULL && env_region[0] != '\0') { + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Using AWS_REGION from environment: %s", env_region); aws_region = aws_string_new_from_c_str(allocator, env_region); if (!aws_region) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create region string from AWS_REGION"); if (err_msg) *err_msg = strdup("Failed to create region string from AWS_REGION"); goto cleanup; @@ -118,20 +235,25 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) } else { + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "AWS_REGION not set, attempting to infer from hostname: %s", endpoint); /* Try to infer region from hostname */ if (aws_dsql_auth_config_infer_region(allocator, &auth_config, &aws_region) != AWS_OP_SUCCESS || aws_region == NULL) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to infer AWS region from hostname: %s", endpoint); if (err_msg) *err_msg = strdup("Failed to infer AWS region from hostname. Please set AWS_REGION environment variable."); goto cleanup; } + AWS_LOGF_INFO(AWS_LS_AUTH_GENERAL, "Inferred region: %s", aws_string_c_str(aws_region)); } aws_dsql_auth_config_set_region(&auth_config, aws_region); - /* Create default credentials provider */ + /* Create default credentials provider with client bootstrap for IMDS */ AWS_ZERO_STRUCT(credentials_options); + credentials_options.bootstrap = s_client_bootstrap; + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Creating credentials provider chain with bootstrap for IMDS"); credentials_provider = aws_credentials_provider_new_chain_default(allocator, &credentials_options); if (!credentials_provider) { aws_error = aws_last_error(); From 6c19265f447ed89bb2530fbbd6aaac15a4291569 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 22:26:59 -0700 Subject: [PATCH 08/10] Improve init The token gen infra (allocator, creds provider) are setup once, globally. This lets us cache creds and do a fast-fail if there are no creds. --- scripts/build-dsql.sh | 2 +- src/bin/pgbench/pgbench.c | 24 +++ src/bin/psql/startup.c | 26 +++ src/interfaces/libpq/Makefile | 12 +- src/interfaces/libpq/exports.txt | 4 + src/interfaces/libpq/fe-connect.c | 2 +- src/interfaces/libpq/fe-dsql-auth.c | 321 +++++++++++++++++++++------- src/interfaces/libpq/fe-dsql-auth.h | 10 +- 8 files changed, 320 insertions(+), 81 deletions(-) diff --git a/scripts/build-dsql.sh b/scripts/build-dsql.sh index 83173df1b8ada..2a2f9f6197cab 100755 --- a/scripts/build-dsql.sh +++ b/scripts/build-dsql.sh @@ -51,7 +51,7 @@ else echo " aws-dsql-auth submodules already initialized." fi -if [ ! -f "aws-dsql-auth/build/install/lib64/libaws-dsql-auth.a" ]; then +if [ ! -d "aws-dsql-auth/build/install/" ]; then # Build aws-dsql-auth echo " Building aws-dsql-auth library..." cd aws-dsql-auth diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c index 8c4617917bb34..967612a185828 100644 --- a/src/bin/pgbench/pgbench.c +++ b/src/bin/pgbench/pgbench.c @@ -68,6 +68,7 @@ #include "pgbench.h" #include "port/pg_bitutils.h" #include "portability/instr_time.h" +#include "fe-dsql-auth.h" /* X/Open (XSI) requires to provide M_PI, but core POSIX does not */ #ifndef M_PI @@ -7124,6 +7125,29 @@ main(int argc, char **argv) setenv("PGDSQL", "1", 1); is_no_vacuum = true; foreign_keys = false; + + /* Initialize DSQL token generator */ + if (dsql_initialize_token_generator() != 0) + { + pg_fatal("Failed to initialize DSQL token generator"); + } + + /* Validate AWS credentials */ + { + char *err_msg = NULL; + if (dsql_validate_aws_credentials(&err_msg) != 0) + { + if (err_msg) + { + pg_fatal("DSQL credential validation failed: %s", err_msg); + free(err_msg); + } + else + { + pg_fatal("DSQL credential validation failed"); + } + } + } } /* set default script if none */ diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c index 291fd317b024f..abdcb57abf589 100644 --- a/src/bin/psql/startup.c +++ b/src/bin/psql/startup.c @@ -6,6 +6,7 @@ * src/bin/psql/startup.c */ #include "postgres_fe.h" +#include "libpq-fe.h" #ifndef WIN32 #include @@ -26,6 +27,8 @@ #include "mainloop.h" #include "settings.h" +#include "fe-dsql-auth.h" + /* * Global psql options */ @@ -221,6 +224,29 @@ main(int argc, char *argv[]) { setenv("PGDSQL", "1", 1); pset.getPassword = TRI_NO; + + /* Initialize DSQL token generator */ + if (dsql_initialize_token_generator() != 0) + { + pg_fatal("Failed to initialize DSQL token generator"); + } + + /* Validate AWS credentials */ + { + char *err_msg = NULL; + if (dsql_validate_aws_credentials(&err_msg) != 0) + { + if (err_msg) + { + pg_fatal("DSQL credential validation failed: %s", err_msg); + free(err_msg); + } + else + { + pg_fatal("DSQL credential validation failed"); + } + } + } } /* diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 03556fcd12c49..883aaa78da118 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -90,15 +90,19 @@ ifeq ($(PORTNAME), linux) # Link with AWS libraries using start-group to resolve dependencies $(LD) -r -o fe-dsql-auth-with-aws.o fe-dsql-auth-temp.o --start-group $(AWS_DSQL_AUTH_ALL_LIBS) --end-group # Create a list of symbols to keep (only the public API from fe-dsql-auth.h) - echo "generate_dsql_token" > keep-symbols.txt - echo "dsql_auth_cleanup" >> keep-symbols.txt + echo "dsql_initialize_token_generator" > keep-symbols.txt + echo "dsql_generate_token" >> keep-symbols.txt + echo "dsql_validate_aws_credentials" >> keep-symbols.txt + echo "dsql_cleanup" >> keep-symbols.txt # Hide all symbols except the ones we want to keep objcopy --keep-global-symbols=keep-symbols.txt fe-dsql-auth-with-aws.o $@ rm -f fe-dsql-auth-temp.o fe-dsql-auth-with-aws.o keep-symbols.txt else ifeq ($(PORTNAME), darwin) # macOS: Use ld with exported symbols list - echo "_generate_dsql_token" > exported-symbols.txt - echo "_dsql_auth_cleanup" >> exported-symbols.txt + echo "_dsql_initialize_token_generator" > exported-symbols.txt + echo "_dsql_generate_token" >> exported-symbols.txt + echo "_dsql_validate_aws_credentials" >> exported-symbols.txt + echo "_dsql_cleanup" >> exported-symbols.txt $(LD) -r -o $@ fe-dsql-auth-temp.o -exported_symbols_list exported-symbols.txt $(AWS_DSQL_AUTH_ALL_LIBS) rm -f fe-dsql-auth-temp.o exported-symbols.txt else diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt index 0625cf39e9af3..7ccf0b80cbd03 100644 --- a/src/interfaces/libpq/exports.txt +++ b/src/interfaces/libpq/exports.txt @@ -211,3 +211,7 @@ PQgetAuthDataHook 208 PQdefaultAuthDataHook 209 PQfullProtocolVersion 210 appendPQExpBufferVA 211 +dsql_initialize_token_generator 212 +dsql_generate_token 213 +dsql_validate_aws_credentials 214 +dsql_cleanup 215 diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c index 682b3b064c406..55cd08889ca0e 100644 --- a/src/interfaces/libpq/fe-connect.c +++ b/src/interfaces/libpq/fe-connect.c @@ -1441,7 +1441,7 @@ pqConnectOptions2(PGconn *conn) pwhost = conn->connhost[i].hostaddr; is_admin = strcmp("admin", conn->pguser) == 0; - token = generate_dsql_token(pwhost, is_admin, &err_msg); + token = dsql_generate_token(pwhost, is_admin, &err_msg); if (!token) { libpq_append_conn_error(conn, "DSQL token generation failed for host=%s: %s", diff --git a/src/interfaces/libpq/fe-dsql-auth.c b/src/interfaces/libpq/fe-dsql-auth.c index c7ed75d300eaf..8fa042d1b31e7 100644 --- a/src/interfaces/libpq/fe-dsql-auth.c +++ b/src/interfaces/libpq/fe-dsql-auth.c @@ -18,6 +18,8 @@ /* Include AWS DSQL Auth library functions */ #include #include +#include +#include #include #include #include @@ -37,6 +39,17 @@ static struct aws_event_loop_group *s_el_group = NULL; static struct aws_host_resolver *s_host_resolver = NULL; static struct aws_client_bootstrap *s_client_bootstrap = NULL; +/* + * DSQL Token Generator - holds state for efficient token generation + */ +struct dsql_token_generator { + struct aws_allocator *allocator; + struct aws_credentials_provider *credentials_provider; +}; + +/* Global token generator instance */ +static struct dsql_token_generator s_token_generator = {0}; + /* * Initialize DSQL logging */ @@ -47,7 +60,7 @@ initialize_dsql_logging(void) { struct aws_allocator *allocator = aws_default_allocator(); struct aws_logger_standard_options logger_options = { - .level = AWS_LOG_LEVEL_DEBUG, /* Can be controlled by environment variable */ + .level = AWS_LOG_LEVEL_NONE, /* Can be controlled by environment variable */ .file = stderr /* Log to stderr by default */ }; @@ -90,29 +103,6 @@ initialize_dsql_logging(void) } } -/* - * Clean up DSQL authentication resources - */ -void -dsql_auth_cleanup(void) -{ - if (dsql_logger_initialized) - { - aws_logger_set(NULL); - aws_logger_clean_up(&dsql_logger); - dsql_logger_initialized = false; - } - - if (aws_libs_initialized) - { - aws_sdkutils_library_clean_up(); - aws_auth_library_clean_up(); - aws_io_library_clean_up(); - aws_common_library_clean_up(); - aws_libs_initialized = false; - } -} - /* * Initialize AWS libraries if not already initialized */ @@ -121,7 +111,11 @@ initialize_aws_libs(void) { if (!aws_libs_initialized) { - struct aws_allocator *allocator = aws_default_allocator(); + struct aws_allocator *allocator; + struct aws_host_resolver_default_options resolver_options; + struct aws_client_bootstrap_options bootstrap_options; + + allocator = aws_default_allocator(); aws_common_library_init(allocator); aws_io_library_init(allocator); aws_http_library_init(allocator); @@ -137,20 +131,16 @@ initialize_aws_libs(void) goto error; } - struct aws_host_resolver_default_options resolver_options = { - .el_group = s_el_group, - .max_entries = 8, - }; + resolver_options.el_group = s_el_group; + resolver_options.max_entries = 8; s_host_resolver = aws_host_resolver_new_default(allocator, &resolver_options); if (!s_host_resolver) { AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create host resolver"); goto error; } - struct aws_client_bootstrap_options bootstrap_options = { - .event_loop_group = s_el_group, - .host_resolver = s_host_resolver, - }; + bootstrap_options.event_loop_group = s_el_group; + bootstrap_options.host_resolver = s_host_resolver; s_client_bootstrap = aws_client_bootstrap_new(allocator, &bootstrap_options); if (!s_client_bootstrap) { AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create client bootstrap"); @@ -183,41 +173,130 @@ initialize_aws_libs(void) } } +/* + * Initialize the DSQL token generator with long-lived components + */ +static int +initialize_token_generator(void) +{ + struct aws_credentials_provider_chain_default_options credentials_options; + + if (s_token_generator.allocator != NULL) { + /* Already initialized */ + return AWS_OP_SUCCESS; + } + + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Initializing DSQL token generator"); + + s_token_generator.allocator = aws_default_allocator(); + + /* Create credentials provider with client bootstrap for IMDS */ + AWS_ZERO_STRUCT(credentials_options); + credentials_options.bootstrap = s_client_bootstrap; + + s_token_generator.credentials_provider = aws_credentials_provider_new_chain_default( + s_token_generator.allocator, &credentials_options); + + if (!s_token_generator.credentials_provider) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create credentials provider for token generator"); + s_token_generator.allocator = NULL; + return AWS_OP_ERR; + } + + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "DSQL token generator initialized successfully"); + return AWS_OP_SUCCESS; +} + +/* + * Public initialization function for the DSQL token generator + */ +int +dsql_initialize_token_generator(void) +{ + /* Initialize AWS libraries and logging */ + initialize_aws_libs(); + initialize_dsql_logging(); + + return initialize_token_generator(); +} + +/* + * Clean up the DSQL token generator + */ +static void +cleanup_token_generator(void) +{ + if (s_token_generator.allocator != NULL) { + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Cleaning up DSQL token generator"); + + if (s_token_generator.credentials_provider) { + aws_credentials_provider_release(s_token_generator.credentials_provider); + s_token_generator.credentials_provider = NULL; + } + + s_token_generator.allocator = NULL; + } +} + +/* + * Clean up DSQL authentication resources + */ +void +dsql_cleanup(void) +{ + cleanup_token_generator(); + + if (dsql_logger_initialized) + { + aws_logger_set(NULL); + aws_logger_clean_up(&dsql_logger); + dsql_logger_initialized = false; + } + + if (aws_libs_initialized) + { + aws_sdkutils_library_clean_up(); + aws_auth_library_clean_up(); + aws_io_library_clean_up(); + aws_common_library_clean_up(); + aws_libs_initialized = false; + } +} + /* * Generate a DSQL authentication token for the specified endpoint. - * Uses the AWS DSQL auth library to generate a real token. + * Uses a local auth_config for thread safety and cached credentials provider for efficiency. * Returns a newly allocated string containing the token. */ char * -generate_dsql_token(const char *endpoint, bool admin, char **err_msg) +dsql_generate_token(const char *endpoint, bool admin, char **err_msg) { - struct aws_allocator *allocator; - struct aws_dsql_auth_config auth_config; + struct aws_dsql_auth_config auth_config = {0}; struct aws_dsql_auth_token auth_token = {0}; struct aws_string *aws_region = NULL; - struct aws_credentials_provider *credentials_provider = NULL; - struct aws_credentials_provider_chain_default_options credentials_options; char *token = NULL; int aws_error; const char *env_region; const char *token_str; - /* Initialize AWS libraries and logging */ - initialize_aws_libs(); - initialize_dsql_logging(); - - allocator = aws_default_allocator(); + /* Check if token generator is initialized */ + if (s_token_generator.allocator == NULL) { + if (err_msg) + *err_msg = strdup("Token generator not initialized"); + return NULL; + } AWS_LOGF_INFO(AWS_LS_AUTH_GENERAL, "Starting DSQL token generation for endpoint: %s", endpoint); - /* Initialize DSQL auth config */ + /* Initialize a local auth config for thread safety */ if (aws_dsql_auth_config_init(&auth_config) != AWS_OP_SUCCESS) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to initialize local auth config"); if (err_msg) - *err_msg = strdup("Failed to initialize DSQL auth config"); - goto cleanup; + *err_msg = strdup("Failed to initialize auth config"); + return NULL; } - /* Set hostname */ + /* Set hostname on the local auth config */ aws_dsql_auth_config_set_hostname(&auth_config, endpoint); /* Try to get region from environment variable first */ @@ -225,7 +304,7 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) if (env_region != NULL && env_region[0] != '\0') { AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Using AWS_REGION from environment: %s", env_region); - aws_region = aws_string_new_from_c_str(allocator, env_region); + aws_region = aws_string_new_from_c_str(s_token_generator.allocator, env_region); if (!aws_region) { AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to create region string from AWS_REGION"); if (err_msg) @@ -237,7 +316,7 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) { AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "AWS_REGION not set, attempting to infer from hostname: %s", endpoint); /* Try to infer region from hostname */ - if (aws_dsql_auth_config_infer_region(allocator, &auth_config, &aws_region) != AWS_OP_SUCCESS || + if (aws_dsql_auth_config_infer_region(s_token_generator.allocator, &auth_config, &aws_region) != AWS_OP_SUCCESS || aws_region == NULL) { AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to infer AWS region from hostname: %s", endpoint); @@ -249,28 +328,15 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) } aws_dsql_auth_config_set_region(&auth_config, aws_region); - /* Create default credentials provider with client bootstrap for IMDS */ - AWS_ZERO_STRUCT(credentials_options); - credentials_options.bootstrap = s_client_bootstrap; - - AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Creating credentials provider chain with bootstrap for IMDS"); - credentials_provider = aws_credentials_provider_new_chain_default(allocator, &credentials_options); - if (!credentials_provider) { - aws_error = aws_last_error(); - if (err_msg) - *err_msg = strdup(aws_error_str(aws_error)); - goto cleanup; - } - - /* Set credentials provider */ - aws_dsql_auth_config_set_credentials_provider(&auth_config, credentials_provider); + /* Set the cached credentials provider */ + aws_dsql_auth_config_set_credentials_provider(&auth_config, s_token_generator.credentials_provider); /* Set expiration time to 5 seconds for shorter token lifetime */ aws_dsql_auth_config_set_expires_in(&auth_config, 5); /* 5 seconds */ - /* Generate the token */ + /* Generate the token using local auth config and cached components */ AWS_ZERO_STRUCT(auth_token); - if (aws_dsql_auth_token_generate(&auth_config, admin, allocator, &auth_token) != AWS_OP_SUCCESS) + if (aws_dsql_auth_token_generate(&auth_config, admin, s_token_generator.allocator, &auth_token) != AWS_OP_SUCCESS) { aws_error = aws_last_error(); if (err_msg) @@ -283,10 +349,11 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) if (token_str) { token = strdup(token_str); - /* Token generation successful */ + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "DSQL token generated successfully using local auth config and cached credentials"); } else { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to get token string from generated token"); if (err_msg) *err_msg = strdup("Failed to get token string"); } @@ -295,11 +362,6 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) aws_dsql_auth_token_clean_up(&auth_token); aws_dsql_auth_config_clean_up(&auth_config); - if (credentials_provider) - { - aws_credentials_provider_release(credentials_provider); - } - if (aws_region) { aws_string_destroy(aws_region); @@ -307,3 +369,116 @@ generate_dsql_token(const char *endpoint, bool admin, char **err_msg) return token; } + +/* Synchronous credential retrieval state */ +struct credential_validation_state { + struct aws_credentials *credentials; + int error_code; + bool completed; + struct aws_mutex mutex; + struct aws_condition_variable condition_variable; +}; + +/* Callback for synchronous credential retrieval */ +static void +s_on_credentials_acquired(struct aws_credentials *credentials, int error_code, void *user_data) +{ + struct credential_validation_state *state = (struct credential_validation_state *)user_data; + + aws_mutex_lock(&state->mutex); + + state->credentials = credentials; + state->error_code = error_code; + state->completed = true; + + if (credentials) { + aws_credentials_acquire(credentials); + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Credentials acquired successfully for validation"); + } else { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Credentials acquisition failed with error: %d", error_code); + } + + aws_condition_variable_notify_one(&state->condition_variable); + aws_mutex_unlock(&state->mutex); +} + +/* + * Validate AWS credentials early for DSQL authentication. + * This initializes the token generator and validates that credentials can be obtained. + * Returns AWS_OP_SUCCESS on success, AWS_OP_ERR on failure. + */ +int +dsql_validate_aws_credentials(char **err_msg) +{ + struct credential_validation_state state = {0}; + int result = AWS_OP_ERR; + + /* Check if token generator is initialized */ + if (s_token_generator.allocator == NULL) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Token generator not initialized during credential validation"); + if (err_msg) + *err_msg = strdup("Token generator not initialized"); + return AWS_OP_ERR; + } + + AWS_LOGF_INFO(AWS_LS_AUTH_GENERAL, "Validating AWS credentials for DSQL authentication"); + + /* Initialize synchronization primitives */ + if (aws_mutex_init(&state.mutex) != AWS_OP_SUCCESS) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to initialize mutex for credential validation"); + if (err_msg) + *err_msg = strdup("Failed to initialize synchronization"); + return AWS_OP_ERR; + } + + if (aws_condition_variable_init(&state.condition_variable) != AWS_OP_SUCCESS) { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to initialize condition variable for credential validation"); + aws_mutex_clean_up(&state.mutex); + if (err_msg) + *err_msg = strdup("Failed to initialize synchronization"); + return AWS_OP_ERR; + } + + /* Actually retrieve credentials to validate they exist and are accessible */ + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Attempting to retrieve AWS credentials for validation"); + + if (aws_credentials_provider_get_credentials( + s_token_generator.credentials_provider, + s_on_credentials_acquired, + &state) != AWS_OP_SUCCESS) { + + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to initiate credentials retrieval"); + if (err_msg) + *err_msg = strdup("Failed to initiate credentials retrieval"); + goto cleanup; + } + + /* Wait for credentials retrieval to complete */ + aws_mutex_lock(&state.mutex); + while (!state.completed) { + aws_condition_variable_wait(&state.condition_variable, &state.mutex); + } + aws_mutex_unlock(&state.mutex); + + /* Check if credentials were successfully retrieved */ + if (state.credentials && state.error_code == AWS_OP_SUCCESS) { + AWS_LOGF_INFO(AWS_LS_AUTH_GENERAL, "AWS credentials validation completed successfully"); + AWS_LOGF_DEBUG(AWS_LS_AUTH_GENERAL, "Token generator ready for DSQL authentication"); + result = AWS_OP_SUCCESS; + + aws_credentials_release(state.credentials); + } else { + AWS_LOGF_ERROR(AWS_LS_AUTH_GENERAL, "Failed to retrieve AWS credentials: %s", + aws_error_str(state.error_code)); + if (err_msg) { + const char *error_str = aws_error_str(state.error_code); + *err_msg = strdup(error_str ? error_str : "Unknown credential retrieval error"); + } + } + +cleanup: + aws_condition_variable_clean_up(&state.condition_variable); + aws_mutex_clean_up(&state.mutex); + + return result; +} diff --git a/src/interfaces/libpq/fe-dsql-auth.h b/src/interfaces/libpq/fe-dsql-auth.h index 23c65c2f0cd5a..d24c3a600c7d7 100644 --- a/src/interfaces/libpq/fe-dsql-auth.h +++ b/src/interfaces/libpq/fe-dsql-auth.h @@ -10,10 +10,16 @@ #include +/* Initialize the DSQL token generator */ +int dsql_initialize_token_generator(void); + /* Generate a DSQL authentication token for the specified endpoint */ -char *generate_dsql_token(const char *endpoint, bool admin, char **err_msg); +char *dsql_generate_token(const char *endpoint, bool admin, char **err_msg); + +/* Initialize and validate AWS credentials early (for startup validation) */ +int dsql_validate_aws_credentials(char **err_msg); /* Clean up DSQL authentication resources */ -void dsql_auth_cleanup(void); +void dsql_cleanup(void); #endif /* FE_DSQL_AUTH_H */ From 996390bcb56101767831b11fc68db0dac03537fd Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 22:39:31 -0700 Subject: [PATCH 09/10] Fix missing zeros --- src/interfaces/libpq/fe-dsql-auth.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/interfaces/libpq/fe-dsql-auth.c b/src/interfaces/libpq/fe-dsql-auth.c index 8fa042d1b31e7..6163302a38cab 100644 --- a/src/interfaces/libpq/fe-dsql-auth.c +++ b/src/interfaces/libpq/fe-dsql-auth.c @@ -131,6 +131,7 @@ initialize_aws_libs(void) goto error; } + AWS_ZERO_STRUCT(resolver_options); resolver_options.el_group = s_el_group; resolver_options.max_entries = 8; s_host_resolver = aws_host_resolver_new_default(allocator, &resolver_options); @@ -139,6 +140,7 @@ initialize_aws_libs(void) goto error; } + AWS_ZERO_STRUCT(bootstrap_options); bootstrap_options.event_loop_group = s_el_group; bootstrap_options.host_resolver = s_host_resolver; s_client_bootstrap = aws_client_bootstrap_new(allocator, &bootstrap_options); From d67c60c5b242cb9458b99fa2560a7f15fa019bf9 Mon Sep 17 00:00:00 2001 From: Marc Bowes <15209+marcbowes@users.noreply.github.com> Date: Fri, 27 Jun 2025 22:57:33 -0700 Subject: [PATCH 10/10] Support private endpoints, install pgbench --- aws-dsql-auth | 2 +- scripts/install.sh | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/aws-dsql-auth b/aws-dsql-auth index f33fce18f3864..3c53ec315ed61 160000 --- a/aws-dsql-auth +++ b/aws-dsql-auth @@ -1 +1 @@ -Subproject commit f33fce18f3864e9428b10ae63c018b004b18436a +Subproject commit 3c53ec315ed619715c56875e81476119c365ddac diff --git a/scripts/install.sh b/scripts/install.sh index 7959634c7d33b..0dcbba2d306d3 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -196,11 +196,14 @@ install_via_zip() { # Clean up temp files rm -rf "$TEMP_DIR" - # Make the binary executable + # Make the binaries executable chmod +x "$INSTALL_PATH/bin/pdsql" + chmod +x "$INSTALL_PATH/bin/pgbench" echo "ZIP installation completed successfully!" - echo "PostgreSQL DSQL (pdsql) installed to: $INSTALL_PATH/bin/pdsql" + echo "PostgreSQL DSQL tools installed to:" + echo " - pdsql (DSQL client): $INSTALL_PATH/bin/pdsql" + echo " - pgbench (benchmark tool): $INSTALL_PATH/bin/pgbench" # Check if installation path is in PATH if [[ ":$PATH:" != *":$INSTALL_PATH/bin:"* ]]; then