feat: add support for asyncpg driver

.mypy.ini

@@ -12,6 +12,9 @@ ignore_missing_imports = True
 [mypy-pg8000]
 ignore_missing_imports = True
 
+[mypy-asyncpg]
+ignore_missing_imports = True
+
 [mypy-pytds]
 ignore_missing_imports = True
 

README.md

@@ -24,6 +24,7 @@ The Cloud SQL Python Connector is a package to be used alongside a database driv
 Currently supported drivers are:
  - [`pymysql`](https://github.com/PyMySQL/PyMySQL) (MySQL)
  - [`pg8000`](https://github.com/tlocke/pg8000) (PostgreSQL)
+ - [`asyncpg`](https://github.com/MagicStack/asyncpg) (PostgreSQL)
  - [`pytds`](https://github.com/denisenkom/pytds) (SQL Server)
 
 
@@ -37,9 +38,16 @@ based on your database dialect.
 pip install "cloud-sql-python-connector[pymysql]"
 ```
 ### Postgres
+There are two different database drivers that are supported for the Postgres dialect:
+
+#### pg8000
 ```
 pip install "cloud-sql-python-connector[pg8000]"
 ```
+#### asyncpg
+```
+pip install "cloud-sql-python-connector[asyncpg]"
+```
 ### SQL Server
 ```
 pip install "cloud-sql-python-connector[pytds]"
@@ -275,6 +283,58 @@ connector.connect(
 )
 ``` 
 
+### Async Driver Usage
+The Cloud SQL Connector for Python currently supports the
+[asyncpg](https://magicstack.github.io/asyncpg) Postgres database driver.
+This driver leverages [asyncio](https://docs.python.org/3/library/asyncio.html)
+to improve the speed and efficiency of database connections through concurrency.
+
+The Cloud SQL Connector has an async `create_async_connector` that can be used
+and is recommended for async drivers as it returns a `Connector` object
+that uses the current thread's running event loop automatically.
+
+The `create_async_connector` allows all the same input arguments as the
+[Connector](#configuring-the-connector) object.
+
+Once a `Connector` object is returned by `create_async_connector` you can call
+its `connect_async` method, just as you would the `connect` method:
+
+```python
+import asyncpg
+from google.cloud.sql.connector import create_async_connector
+
+async def main():
+    # intialize Connector object using 'create_async_connector'
+    connector = await create_async_connector()
+
+    # create connection to Cloud SQL database
+    conn: asyncpg.Connection = await connector.connect_async(
+        "project:region:instance", # Cloud SQL instance connection name
+        "asyncpg",
+        user="root",
+        password="shhh",
+        db="your-db-name"
+        # ... additional database driver args 
+    )
+
+    # insert into Cloud SQL database (example)
+    await conn.execute("INSERT INTO ratings (title, genre, rating) VALUES ('Batman', 'Action', 8.2)")
+
+    # query Cloud SQL database (examples)
+    results = await conn.fetch("SELECT * from ratings")
+    for row in results:
+        # ... do something with results
+
+    # close asyncpg connection
+    await conn.close
+
+    # close Cloud SQL Connector
+    await connector.close_async()
+```
+
+For more details on interacting with an `asyncpg.Connection`, please visit
+the [official documentation](https://magicstack.github.io/asyncpg/current/api/index.html).
+
 ## Support policy
 
 ### Major version lifecycle

google/cloud/sql/connector/__init__.py

@@ -13,11 +13,11 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from .connector import Connector
+from .connector import Connector, create_async_connector
 from .instance import IPTypes
 
 
-__ALL__ = [Connector, IPTypes]
+__ALL__ = [create_async_connector, Connector, IPTypes]
 
 try:
     import pkg_resources

google/cloud/sql/connector/asyncpg.py

@@ -0,0 +1,64 @@
+"""
+Copyright 2022 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import ssl
+from typing import Any, TYPE_CHECKING
+
+SERVER_PROXY_PORT = 3307
+
+if TYPE_CHECKING:
+    import asyncpg
+
+
+async def connect(
+    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+) -> "asyncpg.Connection":
+    """Helper function to create an asyncpg DB-API connection object.
+
+    :type ip_address: str
+    :param ip_address: A string containing an IP address for the Cloud SQL
+        instance.
+
+    :type ctx: ssl.SSLContext
+    :param ctx: An SSLContext object created from the Cloud SQL server CA
+        cert and ephemeral cert.
+
+    :type kwargs: Any
+    :param kwargs: Keyword arguments for establishing asyncpg connection
+        object to Cloud SQL instance.
+
+    :rtype: asyncpg.Connection
+    :returns: An asyncpg.Connection object to a Cloud SQL instance.
+    """
+    try:
+        import asyncpg
+    except ImportError:
+        raise ImportError(
+            'Unable to import module "asyncpg." Please install and try again.'
+        )
+    user = kwargs.pop("user")
+    db = kwargs.pop("db")
+    passwd = kwargs.pop("password", None)
+
+    return await asyncpg.connect(
+        user=user,
+        database=db,
+        password=passwd,
+        host=ip_address,
+        port=SERVER_PROXY_PORT,
+        ssl=ctx,
+        direct_tls=True,
+        **kwargs,
+    )

google/cloud/sql/connector/connector.py

@@ -14,7 +14,6 @@
 limitations under the License.
 """
 import asyncio
-import concurrent
 import logging
 from types import TracebackType
 from google.cloud.sql.connector.instance import (
@@ -24,6 +23,7 @@
 import google.cloud.sql.connector.pymysql as pymysql
 import google.cloud.sql.connector.pg8000 as pg8000
 import google.cloud.sql.connector.pytds as pytds
+import google.cloud.sql.connector.asyncpg as asyncpg
 from google.cloud.sql.connector.utils import generate_keys
 from google.auth.credentials import Credentials
 from threading import Thread
@@ -32,6 +32,18 @@
 
 logger = logging.getLogger(name=__name__)
 
+ASYNC_DRIVERS = ["asyncpg"]
+
+
+class ConnectorLoopError(Exception):
+    """
+    Raised when Connector.connect is called with Connector._loop
+        in an invalid state (event loop in current thread).
+    """
+
+    def __init__(self, *args: Any) -> None:
+        super(ConnectorLoopError, self).__init__(self, *args)
+
 
 class Connector:
     """A class to configure and create connections to Cloud SQL instances.
@@ -53,6 +65,11 @@ class Connector:
     :param credentials
         Credentials object used to authenticate connections to Cloud SQL server.
         If not specified, Application Default Credentials are used.
+
+    :type loop: asyncio.AbstractEventLoop
+    :param loop
+        Event loop to run asyncio tasks, if not specified, defaults to
+        creating new event loop on background thread.
     """
 
     def __init__(
@@ -61,13 +78,22 @@ def __init__(
         enable_iam_auth: bool = False,
         timeout: int = 30,
         credentials: Optional[Credentials] = None,
+        loop: asyncio.AbstractEventLoop = None,
     ) -> None:
-        self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
-        self._thread: Thread = Thread(target=self._loop.run_forever, daemon=True)
-        self._thread.start()
-        self._keys: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(
-            generate_keys(), self._loop
-        )
+        # if event loop is given, use for background tasks
+        if loop:
+            self._loop: asyncio.AbstractEventLoop = loop
+            self._thread: Optional[Thread] = None
+            self._keys: asyncio.Future = loop.create_task(generate_keys())
+        # if no event loop is given, spin up new loop in background thread
+        else:
+            self._loop = asyncio.new_event_loop()
+            self._thread = Thread(target=self._loop.run_forever, daemon=True)
+            self._thread.start()
+            self._keys = asyncio.wrap_future(
+                asyncio.run_coroutine_threadsafe(generate_keys(), self._loop),
+                loop=self._loop,
+            )
         self._instances: Dict[str, Instance] = {}
 
         # set default params for connections
@@ -102,6 +128,18 @@ def connect(
         :returns:
             A DB-API connection to the specified Cloud SQL instance.
         """
+        try:
+            # check if event loop is running in current thread
+            if self._loop == asyncio.get_running_loop():
+                raise ConnectorLoopError(
+                    "Connector event loop is running in current thread!"
+                    "Event loop must be attached to a different thread to prevent blocking code!"
+                )
+        # asyncio.get_running_loop will throw RunTimeError if no running loop is present
+        except RuntimeError:
+            pass
+
+        # if event loop is not in current thread, proceed with connection
         connect_task = asyncio.run_coroutine_threadsafe(
             self.connect_async(instance_connection_string, driver, **kwargs), self._loop
         )
@@ -123,7 +161,7 @@ async def connect_async(
         :type driver: str
         :param: driver:
             A string representing the driver to connect with. Supported drivers are
-            pymysql, pg8000, and pytds.
+            pymysql, pg8000, asyncpg, and pytds.
 
         :param kwargs:
             Pass in any driver-specific arguments needed to connect to the Cloud
@@ -133,7 +171,6 @@ async def connect_async(
         :returns:
             A DB-API connection to the specified Cloud SQL instance.
         """
-
         # Create an Instance object from the connection string.
         # The Instance should verify arguments.
         #
@@ -164,6 +201,7 @@ async def connect_async(
         connect_func = {
             "pymysql": pymysql.connect,
             "pg8000": pg8000.connect,
+            "asyncpg": asyncpg.connect,
             "pytds": pytds.connect,
         }
 
@@ -194,6 +232,10 @@ async def connect_async(
         # helper function to wrap in timeout
         async def get_connection() -> Any:
             instance_data, ip_address = await instance.connect_info(ip_type)
+            # async drivers are unblocking and can be awaited directly
+            if driver in ASYNC_DRIVERS:
+                return await connector(ip_address, instance_data.context, **kwargs)
+            # synchronous drivers are blocking and run using executor
             connect_partial = partial(
                 connector, ip_address, instance_data.context, **kwargs
             )
@@ -224,13 +266,55 @@ def __exit__(
 
     def close(self) -> None:
         """Close Connector by stopping tasks and releasing resources."""
-        close_future = asyncio.run_coroutine_threadsafe(self._close(), loop=self._loop)
+        close_future = asyncio.run_coroutine_threadsafe(
+            self.close_async(), loop=self._loop
+        )
         # Will attempt to safely shut down tasks for 5s
         close_future.result(timeout=5)
 
-    async def _close(self) -> None:
+    async def close_async(self) -> None:
         """Helper function to cancel Instances' tasks
         and close aiohttp.ClientSession."""
         await asyncio.gather(
             *[instance.close() for instance in self._instances.values()]
         )
+
+
+async def create_async_connector(
+    ip_type: IPTypes = IPTypes.PUBLIC,
+    enable_iam_auth: bool = False,
+    timeout: int = 30,
+    credentials: Optional[Credentials] = None,
+    loop: asyncio.AbstractEventLoop = None,
+) -> Connector:
+    """
+    Create Connector object for asyncio connections that can auto-detect
+    and use current thread's running event loop.
+
+    :type ip_type: IPTypes
+    :param ip_type
+        The IP type (public or private)  used to connect. IP types
+        can be either IPTypes.PUBLIC or IPTypes.PRIVATE.
+
+    :type enable_iam_auth: bool
+    :param enable_iam_auth
+        Enables IAM based authentication (Postgres only).
+
+    :type timeout: int
+    :param timeout
+        The time limit for a connection before raising a TimeoutError.
+
+    :type credentials: google.auth.credentials.Credentials
+    :param credentials
+        Credentials object used to authenticate connections to Cloud SQL server.
+        If not specified, Application Default Credentials are used.
+
+    :type loop: asyncio.AbstractEventLoop
+    :param loop
+        Event loop to run asyncio tasks, if not specified, defaults
+        to current thread's running event loop.
+    """
+    # if no loop given, automatically detect running event loop
+    if loop is None:
+        loop = asyncio.get_running_loop()
+    return Connector(ip_type, enable_iam_auth, timeout, credentials, loop)