Skip to content

Commit

Permalink
Add support to use awaitable object in password function. (MagicStack…
Browse files Browse the repository at this point in the history
…#889)

Add support to use awaitable object in password function.  This will allow to pass lambda or `functools.partial` that returns `Future`.
  • Loading branch information
kwarunek authored and rohitsanj committed May 8, 2023
1 parent 54ec77c commit fc06eb7
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 4 deletions.
7 changes: 3 additions & 4 deletions asyncpg/connect_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -757,10 +757,9 @@ async def _connect_addr(

params_input = params
if callable(params.password):
if inspect.iscoroutinefunction(params.password):
password = await params.password()
else:
password = params.password()
password = params.password()
if inspect.isawaitable(password):
password = await password

params = params._replace(password=password)
args = (addr, loop, config, connection_class, record_class, params_input)
Expand Down
19 changes: 19 additions & 0 deletions tests/test_connect.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,25 @@ async def get_wrongpassword():
user='password_user',
password=get_wrongpassword)

async def test_auth_password_cleartext_callable_awaitable(self):
async def get_correctpassword():
return 'correctpassword'

async def get_wrongpassword():
return 'wrongpassword'

conn = await self.connect(
user='password_user',
password=lambda: get_correctpassword())
await conn.close()

with self.assertRaisesRegex(
asyncpg.InvalidPasswordError,
'password authentication failed for user "password_user"'):
await self._try_connect(
user='password_user',
password=lambda: get_wrongpassword())

async def test_auth_password_md5(self):
conn = await self.connect(
user='md5_user', password='correctpassword')
Expand Down

0 comments on commit fc06eb7

Please sign in to comment.