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="postgres-iam-user@gmail.com",
- 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,
+ )