feat: add support for asyncpg driver (#390)

Author: jackwotherspoon

.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]"
@@ -111,9 +119,9 @@ def getconn() -> pymysql.connections.Connection:
     conn: pymysql.connections.Connection = connector.connect(
         "project:region:instance",
         "pymysql",
-        user="root",
-        password="shhh",
-        db="your-db-name"
+        user="my-user",
+        password="my-password",
+        db="my-db-name"
     )
     return conn
 
@@ -188,9 +196,9 @@ def getconn() -> pymysql.connections.Connection:
         conn = connector.connect(
             "project:region:instance",
             "pymysql",
-            user="root",
-            password="shhh",
-            db="your-db-name"
+            user="my-user",
+            password="my-password",
+            db="my-db-name"
         )
     return conn
 
@@ -245,7 +253,7 @@ connector.connect(
      "project:region:instance",
      "pg8000",
      user="[email protected]",
-     db="my_database",
+     db="my-db-name",
      enable_iam_auth=True,
  )
 ```
@@ -258,7 +266,7 @@ Once you have followed the steps linked above, you can run the following code to
 connector.connect(
     "project:region:instance",
     "pytds",
-    db="my_database",
+    db="my-db-name",
     active_directory_auth=True,
     server_name="public.[instance].[location].[project].cloudsql.[domain]",
 )
@@ -268,13 +276,111 @@ Or, if using Private IP:
 connector.connect(
     "project:region:instance",
     "pytds",
-    db="my_database",
+    db="my-db-name",
     active_directory_auth=True,
     server_name="private.[instance].[location].[project].cloudsql.[domain]",
     ip_type=IPTypes.PRIVATE
 )
 ``` 
 
+### Async Driver Usage
+The Cloud SQL Connector is compatible with
+[asyncio](https://docs.python.org/3/library/asyncio.html) to improve the speed
+and efficiency of database connections through concurrency. You can use all
+non-asyncio drivers through the `Connector.connect_async` function, in addition
+to the following asyncio database drivers:
+- [asyncpg](https://magicstack.github.io/asyncpg) (Postgres)
+
+The Cloud SQL Connector has a helper `create_async_connector` function that is
+recommended for asyncio database connections. It returns a `Connector`
+object that uses the current thread's running event loop. This is different
+than `Connector()` which by default initializes a new event loop in a
+background thread.
+
+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="my-user",
+        password="my-password",
+        db="my-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 (example)
+    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).
+
+### Async Context Manager
+
+An alternative to using the `create_async_connector` function is initializing
+a `Connector` as an async context manager, removing the need for explicit
+calls to `connector.close_async()` to cleanup resources. 
+
+**Note:** This alternative requires that the running event loop be
+passed in as the `loop` argument to `Connector()`.
+
+```python
+import asyncio
+import asyncpg
+from google.cloud.sql.connector import Connector
+
+async def main():
+    # get current running event loop to be used with Connector
+    loop = asyncio.get_running_loop()
+    # intialize Connector object as async context manager
+    async with Connector(loop=loop) as connector:
+
+        # create connection to Cloud SQL database
+        conn: asyncpg.Connection = await connector.connect_async(
+            "project:region:instance", # Cloud SQL instance connection name
+            "asyncpg",
+            user="my-user",
+            password="my-password",
+            db="my-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 (example)
+        results = await conn.fetch("SELECT * from ratings")
+        for row in results:
+            # ... do something with results
+        
+        # close asyncpg connection
+        await conn.close
+```
+
 ## 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,
+    )