Skip to content

Commit cd766cb

Browse files
d1mansonelprans
authored andcommitted
Add connect_fn kwarg to pool
1 parent cee97e1 commit cd766cb

File tree

2 files changed

+11
-4
lines changed

2 files changed

+11
-4
lines changed

Diff for: asyncpg/_testbase/__init__.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -268,14 +268,15 @@ def create_pool(dsn=None, *,
268268
pool_class=pg_pool.Pool,
269269
connection_class=pg_connection.Connection,
270270
record_class=asyncpg.Record,
271+
connect_fn=pg_connection.connect,
271272
**connect_kwargs):
272273
return pool_class(
273274
dsn,
274275
min_size=min_size, max_size=max_size,
275276
max_queries=max_queries, loop=loop, setup=setup, init=init,
276277
max_inactive_connection_lifetime=max_inactive_connection_lifetime,
277278
connection_class=connection_class,
278-
record_class=record_class,
279+
record_class=record_class, connect_fn=connect_fn,
279280
**connect_kwargs)
280281

281282

Diff for: asyncpg/pool.py

+9-3
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,7 @@ class Pool:
313313

314314
__slots__ = (
315315
'_queue', '_loop', '_minsize', '_maxsize',
316-
'_init', '_connect_args', '_connect_kwargs',
316+
'_init', '_connect_fn', '_connect_args', '_connect_kwargs',
317317
'_holders', '_initialized', '_initializing', '_closing',
318318
'_closed', '_connection_class', '_record_class', '_generation',
319319
'_setup', '_max_queries', '_max_inactive_connection_lifetime'
@@ -329,6 +329,7 @@ def __init__(self, *connect_args,
329329
loop,
330330
connection_class,
331331
record_class,
332+
connect_fn,
332333
**connect_kwargs):
333334

334335
if len(connect_args) > 1:
@@ -388,6 +389,7 @@ def __init__(self, *connect_args,
388389
self._init = init
389390
self._connect_args = connect_args
390391
self._connect_kwargs = connect_kwargs
392+
self._connect_fn = connect_fn
391393

392394
self._setup = setup
393395
self._max_queries = max_queries
@@ -503,7 +505,7 @@ def set_connect_args(self, dsn=None, **connect_kwargs):
503505
self._connect_kwargs = connect_kwargs
504506

505507
async def _get_new_connection(self):
506-
con = await connection.connect(
508+
con = await self._connect_fn(
507509
*self._connect_args,
508510
loop=self._loop,
509511
connection_class=self._connection_class,
@@ -1097,6 +1099,10 @@ def create_pool(dsn=None, *,
10971099
or :meth:`Connection.set_type_codec() <\
10981100
asyncpg.connection.Connection.set_type_codec>`.
10991101
1102+
:param coroutine connect_fn:
1103+
A coroutine with signature identical to :func:`~asyncpg.connection.connect`. This can be used to add custom
1104+
authentication or ssl logic when creating a connection, as is required by GCP's cloud-sql-python-connector.
1105+
11001106
:param loop:
11011107
An asyncio event loop instance. If ``None``, the default
11021108
event loop will be used.
@@ -1127,7 +1133,7 @@ def create_pool(dsn=None, *,
11271133
return Pool(
11281134
dsn,
11291135
connection_class=connection_class,
1130-
record_class=record_class,
1136+
record_class=record_class, connect_fn=connection.connect,
11311137
min_size=min_size, max_size=max_size,
11321138
max_queries=max_queries, loop=loop, setup=setup, init=init,
11331139
max_inactive_connection_lifetime=max_inactive_connection_lifetime,

0 commit comments

Comments
 (0)