From 48868097548992965b6f0adbff6da97ae62468a2 Mon Sep 17 00:00:00 2001
From: izeau <izeau@users.noreply.github.com>
Date: Tue, 28 May 2024 21:14:50 +0200
Subject: [PATCH] allow user & database to be provided dynamically

---
 README.md         |  4 ++--
 src/connection.js | 26 ++++++++++++++------------
 types/index.d.ts  | 10 +++++++---
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 421d19a0..0b0c83df 100644
--- a/README.md
+++ b/README.md
@@ -994,9 +994,9 @@ const sql = postgres('postgres://username:password@host:port/database', {
 
 Note that `max_lifetime = 60 * (30 + Math.random() * 30)` by default. This resolves to an interval between 30 and 60 minutes to optimize for the benefits of prepared statements **and** working nicely with Linux's OOM killer.
 
-### Dynamic passwords
+### Dynamic credentials
 
-When clients need to use alternative authentication schemes such as access tokens or connections to databases with rotating passwords, provide either a synchronous or asynchronous function that will resolve the dynamic password value at connection time.
+When clients need to use alternative authentication schemes such as access tokens or connections to databases with rotating usernames, database names or passwords, provide either a synchronous or asynchronous function that will resolve the dynamic value at connection time.
 
 ```js
 const sql = postgres(url, {
diff --git a/src/connection.js b/src/connection.js
index 578a6a02..076db926 100644
--- a/src/connection.js
+++ b/src/connection.js
@@ -53,10 +53,8 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
   const {
     ssl,
     max,
-    user,
     host,
     port,
-    database,
     parsers,
     transform,
     onnotice,
@@ -71,6 +69,9 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
   const sent = Queue()
       , id = uid++
       , backend = { pid: null, secret: null }
+      , User = Resolve(options.user)
+      , Database = Resolve(options.database)
+      , Pass = Resolve(options.pass)
       , idleTimer = timer(end, options.idle_timeout)
       , lifeTimer = timer(end, options.max_lifetime)
       , connectTimer = timer(connectTimedOut, options.connect_timeout)
@@ -353,7 +354,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
     setTimeout(connect, closedDate ? closedDate + delay - performance.now() : 0)
   }
 
-  function connected() {
+  async function connected() {
     try {
       statements = {}
       needsTypes = options.fetch_types
@@ -362,7 +363,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
       lifeTimer.start()
       socket.on('data', data)
       keep_alive && socket.setKeepAlive && socket.setKeepAlive(true, 1000 * keep_alive)
-      const s = StartupMessage()
+      const s = await StartupMessage()
       write(s)
     } catch (err) {
       error(err)
@@ -665,10 +666,11 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
   }
 
   async function AuthenticationMD5Password(x) {
+    const [pass, user] = await Promise.all([Pass(), User()])
     const payload = 'md5' + (
-      await md5(
+      md5(
         Buffer.concat([
-          Buffer.from(await md5((await Pass()) + user)),
+          Buffer.from(md5(pass + user)),
           x.subarray(9)
         ])
       )
@@ -720,11 +722,10 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
     socket.destroy()
   }
 
-  function Pass() {
-    return Promise.resolve(typeof options.pass === 'function'
-      ? options.pass()
-      : options.pass
-    )
+  function Resolve(opt) {
+    return async () => typeof opt === 'function'
+      ? opt.apply(options, [options])
+      : opt
   }
 
   function NoData() {
@@ -967,7 +968,8 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
     ])
   }
 
-  function StartupMessage() {
+  async function StartupMessage() {
+    const [user, database] = await Promise.all([User(), Database()])
     return cancelMessage || b().inc(4).i16(3).z(2).str(
       Object.entries(Object.assign({
         user,
diff --git a/types/index.d.ts b/types/index.d.ts
index 8989ff47..5dc0f752 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -34,12 +34,12 @@ interface BaseOptions<T extends Record<string, postgres.PostgresType>> {
    * Name of database to connect to
    * @default process.env['PGDATABASE'] || options.user
    */
-  database: string;
+  database: Resolver<T, string>;
   /**
    * Username of database user
    * @default process.env['PGUSERNAME'] || process.env['PGUSER'] || require('os').userInfo().username
    */
-  user: string;
+  user: Resolver<T, string>;
   /**
    * How to deal with ssl (can be a tls.connect option object)
    * @default false
@@ -125,6 +125,10 @@ interface BaseOptions<T extends Record<string, postgres.PostgresType>> {
 }
 
 
+type Resolver<T extends Record<string, postgres.PostgresType>, U> =
+  | U
+  | ((this: postgres.Options<T>, options: postgres.Options<T>) => U)
+
 declare const PRIVATE: unique symbol;
 
 declare class NotAPromise {
@@ -356,7 +360,7 @@ declare namespace postgres {
      * Password of database user
      * @default process.env['PGPASSWORD']
      */
-    password?: string | (() => string | Promise<string>) | undefined;
+    password?: Resolver<T, string> | undefined;
     /** Name of database to connect to (an alias for `database`) */
     db?: Options<T>['database'] | undefined;
     /** Username of database user (an alias for `user`) */