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.

Author: marcbowes

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

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 <math.h> 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 */

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 <unistd.h>
@@ -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");
+				}
+			}
+		}
 	}
 
 	/*

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

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

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",