Connection Pool: Acquire Connections from Pool

[whilsman-tmt]

I am currently testing using asyncpg for connecting to AWS RDS Aurora specifically using IAM credentials as the password. Given the nature of IAM tokens with AWS they will expire after a duration of fifteen minutes and are no longer be usable to connect to a database. With the usage of these credentials though a user if still connected to the database after fifteen minutes will still be able to perform transactions.

What I am struggling to understand right now is how connections are created from the async with pool.acquire() as con. My assumption is that each time a new connection is needed using the mentioned method it needs to establish a new connection to the database. My test right now involves waiting sixteen minutes AFTER my first connection is acquired from the pool. I'd assume that given the nature of IAM credentials this would not work given that they have expired after fifteen minutes from their generation. So in short my questions are:

1. When the pool.acquire() method is called AND the current size of the pool is zero , and idle connections are not present is there some type of replication of connections occurring?
2. When the pool.acquire() method is called is a new connection created when there are no idle connections, or connections in the pool?
3. What's the best method for updating credentials in a pool?

Posting my spaghetti code below. Will be cleaning it up today as I've just been trying to force this.

class Db():

    def __init__(self, iam_creds):
        self.iam_creds = iam_creds
        self.token = self._generate_token()
        self.uri = self._generate_uri()
        self._pool = None


    def _generate_uri(self):
        f = f"postgres://{self.iam_creds['user']}:{self.token[0]}@{self.iam_creds['host']}:5432/{self.iam_creds['db-name']}"
        return f

    def _generate_token(self):
        """Simple function for generating a token for connecting to rds database"""
        current_time = datetime.datetime.now()
        session = boto3.Session()
        client = session.client('rds')
        token = client.generate_db_auth_token(DBHostname=self.iam_creds['host'],
                                              Port=self.iam_creds['port'],
                                              DBUsername=self.iam_creds['user'],
                                              Region=self.iam_creds['region']
                                              )
        return (token, current_time)
    

    def _renew_access_token(func):
        def wrapper(self, *args, **kwargs):
            if (datetime.datetime.now() - self.token[1]).seconds/60 > 15:
                logging.info(f"Refresh token: Expired at: {self.token[1] - datetime.timedelta(minutes=15)}")
                self.token = self._generate_token()
                ans = func(self, *args, **kwargs)
                return ans
            else:
                logging.info(f"Refresh token: Not Expired")
                ans = func(self, *args, **kwargs)
                return ans
        return wrapper
    

    @asynccontextmanager
    async def acquire(self, timeout=None):
        async with self._pool.acquire(timeout=timeout) as conn:
            yield conn


    async def get_size(self):
        return self._pool.get_size()


    async def initialize(self):
        self._pool = await asyncpg.pool.create_pool(user=self.iam_creds['user'],
                                     host=self.iam_creds['host'], 
                                     port=self.iam_creds['port'],
                                     database=self.iam_creds['db-name'],
                                     password=self.token[0],
                                     #max_inactive_connection_lifetime=0,
                                     min_size=5,
                                     max_size=25
                                     )

async def test(pool, i,run_id):
    if i == 1:
        tm = datetime.datetime.now()
        v = await pool.get_size()
        logging.info(f"{tm}: Loop Number {i}")
        logging.info(f"Token: {pool.token[0]}")
        logging.info(f"Pool Count: {v}")
        async with pool.acquire() as con:
            await con.fetchval('SELECT pg_sleep(30);')
            await con.execute('''INSERT INTO foobar.wh_test (loop_id, run_id, id, hit_time, token) VALUES($1, $2, $3, current_timestamp ,$4)''',1 ,run_id, i, pool.token[0])
        v = await pool.get_size()
        logging.info(f"Pool Count: {v}")
    
    elif 10 >= i and i >= 2:
        await sleep(1000)
        tm = datetime.datetime.now()
        v = await pool.get_size()
        logging.info(f"{tm}: Loop Number {i}")
        logging.info(f"Token: {pool.token[0]}")
        logging.info(f"Pool Count: {v}")
        async with pool.acquire() as con:
            await con.fetchval('SELECT pg_sleep(90);')
            await con.execute('''INSERT INTO foobar.wh_test (loop_id, run_id, id, hit_time, token) VALUES($1, $2, $3, current_timestamp,$4)''',2 ,run_id, i, pool.token[0])
        v = await pool.get_size()
        logging.info(f"Pool Count: {v}")

    else:
        await sleep(1050)
        tm = datetime.datetime.now()
        v = await pool.get_size()
        logging.info(f"{tm}: Loop Number {i}")
        logging.info(f"Token: {pool.token[0]}")
        logging.info(f"Pool Count: {v}")
        async with pool.acquire() as con:
            await con.execute('''INSERT INTO foobar.wh_test (loop_id, run_id, id, hit_time, token) VALUES($1, $2, $3, current_timestamp,$4)''',3 ,run_id, i, pool.token[0])
        v = await pool.get_size()
        logging.info(f"Pool Count: {v}")


async def create_pool(iam_creds=cr):
    db = Db(iam_creds)
    await db.initialize()
    return db


async def bench_asyncpg_pool(run_id):
    pool = await create_pool()
    tasks = []
    for i in range(1, 20):
        tasks.append(create_task(test(pool, i, run_id)))
    await gather(*tasks)


if __name__ == "__main__":
    run_id = 1
    while True:
        logging.info(f"{datetime.datetime.now()}: Script Start")
        run(bench_asyncpg_pool(run_id))
        logging.info(f"{datetime.datetime.now()}: Script End")
        run_id += 1