From 0153f11a9c4fb441ee4d436b14e0e0d6fa80c443 Mon Sep 17 00:00:00 2001 From: Harvey Frye Date: Tue, 7 Apr 2020 14:42:10 -0500 Subject: [PATCH 1/4] Add support for password functions (useful for RDS IAM auth) --- asyncpg/__init__.py | 2 +- asyncpg/connect_utils.py | 7 ++++++- tests/test_connect.py | 20 ++++++++++++++++++++ 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/asyncpg/__init__.py b/asyncpg/__init__.py index 96b41c89..b0cef053 100644 --- a/asyncpg/__init__.py +++ b/asyncpg/__init__.py @@ -31,4 +31,4 @@ # snapshots will automatically include the git revision # in __version__, for example: '0.16.0.dev0+ge06ad03' -__version__ = '0.20.1' +__version__ = '0.21.0.dev0' diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index ec3d1090..83f4caaf 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -601,6 +601,11 @@ async def _connect_addr(*, addr, loop, timeout, params, config, raise asyncio.TimeoutError connected = _create_future(loop) + + params_input = params + if callable(params.password): + params = params._replace(password=params.password()) + proto_factory = lambda: protocol.Protocol( addr, connected, params, loop) @@ -633,7 +638,7 @@ async def _connect_addr(*, addr, loop, timeout, params, config, tr.close() raise - con = connection_class(pr, tr, loop, addr, config, params) + con = connection_class(pr, tr, loop, addr, config, params_input) pr.set_connection(con) return con diff --git a/tests/test_connect.py b/tests/test_connect.py index f0767827..d1db17e7 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -204,6 +204,26 @@ async def test_auth_password_cleartext(self): user='password_user', password='wrongpassword') + async def test_auth_password_cleartext_callable(self): + def get_correctpassword(): + return 'correctpassword' + + def get_wrongpassword(): + return 'wrongpassword' + + conn = await self.connect( + user='password_user', + password=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=get_wrongpassword) + + async def test_auth_password_md5(self): conn = await self.connect( user='md5_user', password='correctpassword') From d28e36cbcba139af0e2408b3ab3480d25ac7612a Mon Sep 17 00:00:00 2001 From: Harvey Frye Date: Tue, 7 Apr 2020 18:10:06 -0500 Subject: [PATCH 2/4] Added support for async password callbacks --- asyncpg/connect_utils.py | 8 +++++++- tests/test_connect.py | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index 83f4caaf..2678b358 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -21,6 +21,7 @@ import typing import urllib.parse import warnings +import inspect from . import compat from . import exceptions @@ -604,7 +605,12 @@ async def _connect_addr(*, addr, loop, timeout, params, config, params_input = params if callable(params.password): - params = params._replace(password=params.password()) + if inspect.iscoroutinefunction(params.password): + password = await params.password() + else: + password = params.password() + + params = params._replace(password=password) proto_factory = lambda: protocol.Protocol( addr, connected, params, loop) diff --git a/tests/test_connect.py b/tests/test_connect.py index d1db17e7..89cb0a76 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -224,6 +224,27 @@ def get_wrongpassword(): password=get_wrongpassword) + async def test_auth_password_cleartext_callable_coroutine(self): + async def get_correctpassword(): + return 'correctpassword' + + async def get_wrongpassword(): + return 'wrongpassword' + + conn = await self.connect( + user='password_user', + password=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=get_wrongpassword) + + + async def test_auth_password_md5(self): conn = await self.connect( user='md5_user', password='correctpassword') From 7a92ded6fd043bc4feb10757249822dd51d96b3c Mon Sep 17 00:00:00 2001 From: Harvey Frye Date: Wed, 8 Apr 2020 09:59:25 -0500 Subject: [PATCH 3/4] Cleaned up extra blank lines in test_connect.py --- tests/test_connect.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_connect.py b/tests/test_connect.py index 89cb0a76..116b8ad9 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -223,7 +223,6 @@ def get_wrongpassword(): user='password_user', password=get_wrongpassword) - async def test_auth_password_cleartext_callable_coroutine(self): async def get_correctpassword(): return 'correctpassword' @@ -243,8 +242,6 @@ async def get_wrongpassword(): user='password_user', password=get_wrongpassword) - - async def test_auth_password_md5(self): conn = await self.connect( user='md5_user', password='correctpassword') From 0e40859e41c1ead60d06ec8ea9c2184a699846cb Mon Sep 17 00:00:00 2001 From: Harvey Frye Date: Wed, 8 Apr 2020 10:14:35 -0500 Subject: [PATCH 4/4] Updated docstring for connect() --- asyncpg/connection.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/asyncpg/connection.py b/asyncpg/connection.py index ef1b595d..76a5a9cf 100644 --- a/asyncpg/connection.py +++ b/asyncpg/connection.py @@ -1566,6 +1566,10 @@ async def connect(dsn=None, *, other users and applications may be able to read it without needing specific privileges. It is recommended to use *passfile* instead. + Password may be either a string, or a callable that returns a string. + If a callable is provided, it will be called each time a new connection + is established. + :param passfile: The name of the file used to store passwords (defaults to ``~/.pgpass``, or ``%APPDATA%\postgresql\pgpass.conf``