feat: Support callable kwargs for token callback authentication

[rcleveng]

Summary

• Add support for callable kwargs parameter to enable token-based authentication with psycopg-pool v3.3.0+
• Bump version to 1.0.3

Description

This PR adds support for the kwargs parameter to be a callable (function that returns a dict) in addition to a static dict. This enables integration with token-based authentication systems like AWS IAM, Azure AD, or GCP IAM where credentials need to be refreshed dynamically.

Background

In psycopg-pool v3.3.0, a new feature was introduced that allows the kwargs parameter to be a callable. When provided as a callable, it's invoked each time a new connection is created, allowing for dynamic credential refresh. This is essential for:

• GCP Cloud SQL IAM Authentication - Requires periodic token refresh

Should also work for:

• AWS RDS IAM Authentication - Tokens expire after 15 minutes
• Azure AD Authentication - OAuth tokens have limited lifetimes
• Any OAuth/OIDC based database authentication

Changes

• Modified PGMQueue.kwargs type hint to accept Union[dict, Callable[[], dict[str, Any]]]
• Updated __post_init__ to pass callable kwargs correctly to ConnectionPool

Usage Example

import boto3

def get_auth_token():
    token = "XXXXX" # instead of XXXXX, fetch token 
    return {'password': token}

queue = PGMQueue(
    host='MY-DB-HOST',
    database='mydb',
    username='myuser',
    password='NOTUSED',  # Not used when using token auth
    kwargs= get_auth_token  # Callable that returns fresh credentials
)

[gemini-code-assist]

Summary of Changes

Hello @rcleveng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant enhancement to the PGMQueue class by allowing the kwargs parameter to accept a callable. This change provides greater flexibility in how connection parameters are supplied, particularly benefiting scenarios requiring dynamic credentials like token authentication. Alongside this functional improvement, the project's version has been incremented.

Highlights

• Callable kwargs support: The PGMQueue class now supports passing a callable for its kwargs parameter, enabling dynamic connection configurations, such as those required for token-based authentication.
• Version Update: The project version has been updated from 1.0.2 to 1.0.3 in pyproject.toml and uv.lock.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature by allowing kwargs to be a callable, which is useful for token-based authentication where credentials need to be refreshed. The implementation is clean and correctly handles both dictionary and callable types for kwargs. I've added one suggestion to improve robustness by preventing a potential runtime error if the kwargs dictionary contains a reserved key. Overall, this is a great enhancement.