Merge pull request #4 from MagicStack/master

Sync Fork from Upstream Repo

Author: sthagen

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'

asyncpg/connect_utils.py

@@ -21,6 +21,7 @@
 import typing
 import urllib.parse
 import warnings
+import inspect
 
 from . import compat
 from . import exceptions
@@ -601,6 +602,16 @@ async def _connect_addr(*, addr, loop, timeout, params, config,
         raise asyncio.TimeoutError
 
     connected = _create_future(loop)
+
+    params_input = params
+    if callable(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)
 
@@ -633,7 +644,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
 

asyncpg/connection.py

@@ -813,7 +813,7 @@ async def _copy_in(self, copy_stmt, source, timeout):
 
         if path is not None:
             # a path
-            f = await run_in_executor(None, open, path, 'wb')
+            f = await run_in_executor(None, open, path, 'rb')
             opened_by_us = True
         elif hasattr(source, 'read'):
             # file-like
@@ -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``
@@ -1646,6 +1650,9 @@ async def connect(dsn=None, *,
        Added ability to specify multiple hosts in the *dsn*
        and *host* arguments.
 
+    .. versionchanged:: 0.21.0
+       The *password* argument now accepts a callable or an async function.
+
     .. _SSLContext: https://docs.python.org/3/library/ssl.html#ssl.SSLContext
     .. _create_default_context:
         https://docs.python.org/3/library/ssl.html#ssl.create_default_context

tests/test_connect.py

@@ -204,6 +204,44 @@ 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_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')

tests/test_copy.py

@@ -8,6 +8,7 @@
 import asyncio
 import datetime
 import io
+import os
 import tempfile
 
 import asyncpg
@@ -582,6 +583,48 @@ async def __anext__(self):
         finally:
             await self.con.execute('DROP TABLE copytab')
 
+    async def test_copy_to_table_from_file_path(self):
+        await self.con.execute('''
+            CREATE TABLE copytab(a text, "b~" text, i int);
+        ''')
+
+        f = tempfile.NamedTemporaryFile(delete=False)
+        try:
+            f.write(
+                '\n'.join([
+                    'a1\tb1\t1',
+                    'a2\tb2\t2',
+                    'a3\tb3\t3',
+                    'a4\tb4\t4',
+                    'a5\tb5\t5',
+                    '*\t\\N\t\\N',
+                    ''
+                ]).encode('utf-8')
+            )
+            f.close()
+
+            res = await self.con.copy_to_table('copytab', source=f.name)
+            self.assertEqual(res, 'COPY 6')
+
+            output = await self.con.fetch("""
+                SELECT * FROM copytab ORDER BY a
+            """)
+            self.assertEqual(
+                output,
+                [
+                    ('*', None, None),
+                    ('a1', 'b1', 1),
+                    ('a2', 'b2', 2),
+                    ('a3', 'b3', 3),
+                    ('a4', 'b4', 4),
+                    ('a5', 'b5', 5),
+                ]
+            )
+
+        finally:
+            await self.con.execute('DROP TABLE public.copytab')
+            os.unlink(f.name)
+
     async def test_copy_records_to_table_1(self):
         await self.con.execute('''
             CREATE TABLE copytab(a text, b int, c timestamptz);