WIP: asyncpg socket callback

google/cloud/sql/connector/asyncpg.py

@@ -13,6 +13,10 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import asyncio
+import contextvars
+import functools
+import socket
 import ssl
 from typing import Any, TYPE_CHECKING
 
@@ -22,6 +26,10 @@
     import asyncpg
 
 
+def sock_func(ip_address: str) -> socket.socket:
+    return socket.create_connection((ip_address, SERVER_PROXY_PORT))
+
+
 async def connect(
     ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
 ) -> "asyncpg.Connection":
@@ -48,6 +56,14 @@ async def connect(
         raise ImportError(
             'Unable to import module "asyncpg." Please install and try again.'
         )
+
+    async def async_sock_func() -> socket.socket:
+        # TODO: switch to asyncio.to_thread once Python 3.8 support is dropped
+        loop = asyncio.get_running_loop()
+        ctx_vars = contextvars.copy_context()
+        func_call = functools.partial(ctx_vars.run, sock_func, ip_address)
+        return await loop.run_in_executor(None, func_call)  # type: ignore
+
     user = kwargs.pop("user")
     db = kwargs.pop("db")
     passwd = kwargs.pop("password", None)
@@ -56,9 +72,7 @@ async def connect(
         user=user,
         database=db,
         password=passwd,
-        host=ip_address,
-        port=SERVER_PROXY_PORT,
         ssl=ctx,
-        direct_tls=True,
+        socket_callback=async_sock_func,
         **kwargs,
     )

requirements-test.txt

@@ -14,6 +14,6 @@ types-mock==5.1.0.3
 twine==4.0.2
 PyMySQL==1.1.0
 pg8000==1.30.3
-asyncpg==0.29.0
+git+https://github.com/jackwotherspoon/asyncpg.git@ff82aab23435d86ddccb0b553c2f116be11013a9
 python-tds==1.13.0
 aioresponses==0.7.6