feat: add support for IAM auth with pg8000 driver (#101)

* feat: add support for IAM auth with pg8000 driver

* Update google/cloud/sql/connector/instance_connection_manager.py

Co-authored-by: Kurtis Van Gent <[email protected]>

* Apply suggestions from code review

Co-authored-by: Kurtis Van Gent <[email protected]>

* calculate seconds to refresh before returning metadata object

* address review comments

Co-authored-by: Kurtis Van Gent <[email protected]>

Author: shubha-rajan

.mypy.ini

@@ -13,4 +13,7 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy-pytest]
+ignore_missing_imports = True
+
+[mypy-OpenSSL]
 ignore_missing_imports = True

README.md

@@ -82,9 +82,9 @@ Note: If specifying Private IP, your application must already be in the same VPC
 Tests can be run with `nox`. Change directory into the `cloud-sql-python-connector` and just run `nox` to run the tests.
 
 1. Create a MySQL instance on Google Cloud SQL. Make sure to note your root password when creating the MySQL instance. 
-2. When the MySQL instance has finished creating, go to the overview page and set the instance’s connection string to the environment variable INSTANCE_CONNECTION_NAME using the following command:
+2. When the MySQL instance has finished creating, go to the overview page and set the instance’s connection string to the environment variable MYSQL_CONNECTION_NAME using the following command:
 ```
-export INSTANCE_CONNECTION_NAME=your:connection:string
+export MYSQL_CONNECTION_NAME=your:connection:string
 ```
 3. Enable SSL for your Cloud SQL instance by following [these instructions](https://cloud.google.com/sql/docs/mysql/configure-ssl-instance).
 4. Create a service account with Cloud SQL Admin and Cloud SQL Client roles, then download the key and save it in a safe location. Set the path to the json file to the environment variable GOOGLE_APPLICATION_CREDENTIALS using the following command:

google/cloud/sql/connector/connector.py

@@ -52,6 +52,7 @@ def connect(
     instance_connection_string: str,
     driver: str,
     ip_types: IPTypes = IPTypes.PUBLIC,
+    enable_iam_auth: bool = False,
     **kwargs: Any
 ) -> Any:
     """Prepares and returns a database connection object and starts a
@@ -73,6 +74,10 @@ def connect(
         The IP type (public or private)  used to connect. IP types
         can be either IPTypes.PUBLIC or IPTypes.PRIVATE.
 
+    :param enable_iam_auth
+        Enables IAM based authentication for Postgres instances.
+    :type enable_iam_auth: bool
+
     :param kwargs:
         Pass in any driver-specific arguments needed to connect to the Cloud
         SQL instance.
@@ -96,7 +101,9 @@ def connect(
         icm = _instances[instance_connection_string]
     else:
         keys = _get_keys(loop)
-        icm = InstanceConnectionManager(instance_connection_string, driver, keys, loop)
+        icm = InstanceConnectionManager(
+            instance_connection_string, driver, keys, loop, enable_iam_auth
+        )
         _instances[instance_connection_string] = icm
 
     if "timeout" in kwargs:

google/cloud/sql/connector/instance_connection_manager.py

@@ -23,10 +23,12 @@
 import asyncio
 import aiohttp
 import concurrent
+import datetime
 from enum import Enum
 import google.auth
 from google.auth.credentials import Credentials
 import google.auth.transport.requests
+import OpenSSL
 import ssl
 import socket
 from tempfile import TemporaryDirectory
@@ -52,9 +54,15 @@
 APPLICATION_NAME = "cloud-sql-python-connector"
 SERVER_PROXY_PORT = 3307
 
-# The default delay is set to 55 minutes since each ephemeral certificate is only
-# valid for an hour. This gives five minutes of buffer time.
-_delay: int = 55 * 60
+# default_refresh_buffer is the amount of time before a refresh's result expires
+# that a new refresh operation begins.
+_default_refresh_buffer: int = 5 * 60  # 5 minutes
+
+# _iam_auth_refresh_buffer is the amount of time before a refresh's result expires
+# that a new refresh operation begins when IAM DB AuthN is enabled. Because token
+# sources may be cached until ~60 seconds before expiration, this value must be smaller
+# than default_refresh_buffer.
+_iam_auth_refresh_buffer: int = 55  # seconds
 
 
 class IPTypes(Enum):
@@ -94,16 +102,19 @@ def __init__(self, *args: Any) -> None:
 class InstanceMetadata:
     ip_addrs: Dict[str, Any]
     context: ssl.SSLContext
+    expiration: datetime.datetime
 
     def __init__(
         self,
         ephemeral_cert: str,
         ip_addrs: Dict[str, Any],
         private_key: bytes,
         server_ca_cert: str,
+        expiration: datetime.datetime,
     ) -> None:
         self.ip_addrs = ip_addrs
         self.context = ConnectionSSLContext()
+        self.expiration = expiration
 
         # tmpdir and its contents are automatically deleted after the CA cert
         # and ephemeral cert are loaded into the SSLcontext. The values
@@ -140,6 +151,10 @@ class InstanceConnectionManager:
         The user agent string to append to SQLAdmin API requests
     :type user_agent_string: str
 
+    :param enable_iam_auth
+        Enables IAM based authentication for Postgres instances.
+    :type enable_iam_auth: bool
+
     :param loop:
         A new event loop for the refresh function to run in.
     :type loop: asyncio.AbstractEventLoop
@@ -153,6 +168,8 @@ class InstanceConnectionManager:
     # https://github.com/GoogleCloudPlatform/cloud-sql-python-connector/issues/22
     _loop: asyncio.AbstractEventLoop
 
+    _enable_iam_auth: bool
+
     __client_session: Optional[aiohttp.ClientSession] = None
 
     @property
@@ -185,6 +202,7 @@ def __init__(
         driver_name: str,
         keys: concurrent.futures.Future,
         loop: asyncio.AbstractEventLoop,
+        enable_iam_auth: bool = False,
     ) -> None:
         # Validate connection string
         connection_string_split = instance_connection_string.split(":")
@@ -200,6 +218,8 @@ def __init__(
                 + "format: project:region:instance."
             )
 
+        self._enable_iam_auth = enable_iam_auth
+
         self._user_agent_string = f"{APPLICATION_NAME}/{version}+{driver_name}"
         self._loop = loop
         self._keys = asyncio.wrap_future(keys, loop=self._loop)
@@ -261,16 +281,30 @@ async def _get_instance_data(self) -> InstanceMetadata:
                 self._project,
                 self._instance,
                 pub_key,
+                self._enable_iam_auth,
             )
         )
 
         metadata, ephemeral_cert = await asyncio.gather(metadata_task, ephemeral_task)
 
+        x509 = OpenSSL.crypto.load_certificate(
+            OpenSSL.crypto.FILETYPE_PEM, ephemeral_cert
+        )
+        expiration = datetime.datetime.strptime(
+            x509.get_notAfter().decode("ascii"), "%Y%m%d%H%M%SZ"
+        )
+        if self._credentials is not None:
+            token_expiration: datetime.datetime = self._credentials.expiry
+
+        if expiration > token_expiration:
+            expiration = token_expiration
+
         return InstanceMetadata(
             ephemeral_cert,
             metadata["ip_addresses"],
             priv_key,
             metadata["server_ca_cert"],
+            expiration,
         )
 
     def _auth_init(self) -> None:
@@ -287,6 +321,27 @@ def _auth_init(self) -> None:
 
         self._credentials = credentials
 
+    async def seconds_until_refresh(self) -> int:
+        expiration = (await self._current).expiration
+
+        if self._enable_iam_auth:
+            refresh_buffer = _iam_auth_refresh_buffer
+        else:
+            refresh_buffer = _default_refresh_buffer
+
+        delay = (expiration - datetime.datetime.now()) - datetime.timedelta(
+            seconds=refresh_buffer
+        )
+
+        if delay.total_seconds() < 0:
+            # If the time until the certificate expires is less than the buffer,
+            # schedule the refresh closer to the expiration time
+            delay = (expiration - datetime.datetime.now()) - datetime.timedelta(
+                seconds=5
+            )
+
+        return int(delay.total_seconds())
+
     async def _perform_refresh(self) -> asyncio.Task:
         """Retrieves instance metadata and ephemeral certificate from the
         Cloud SQL Instance.
@@ -299,22 +354,23 @@ async def _perform_refresh(self) -> asyncio.Task:
 
         self._current = self._loop.create_task(self._get_instance_data())
         # Ephemeral certificate expires in 1 hour, so we schedule a refresh to happen in 55 minutes.
-        self._next = self._loop.create_task(self._schedule_refresh(_delay))
+
+        self._next = self._loop.create_task(self._schedule_refresh())
 
         return self._current
 
-    async def _schedule_refresh(self, delay: int) -> asyncio.Task:
+    async def _schedule_refresh(self, delay: Optional[int] = None) -> asyncio.Task:
         """A coroutine that sleeps for the specified amount of time before
         running _perform_refresh.
 
-        :type delay: int
-        :param delay: An integer representing the number of seconds for delay.
-
         :rtype: asyncio.Task
         :returns: A Task representing _get_instance_data.
         """
         logger.debug("Entering sleep")
 
+        if delay is None:
+            delay = await self.seconds_until_refresh()
+
         try:
             await asyncio.sleep(delay)
         except asyncio.CancelledError as e:
@@ -454,7 +510,7 @@ def _connect_with_pg8000(
             )
         user = kwargs.pop("user")
         db = kwargs.pop("db")
-        passwd = kwargs.pop("password")
+        passwd = kwargs.pop("password", None)
         setattr(ctx, "request_ssl", False)
         return pg8000.dbapi.connect(
             user,

google/cloud/sql/connector/refresh_utils.py

@@ -102,6 +102,7 @@ async def _get_ephemeral(
     project: str,
     instance: str,
     pub_key: str,
+    enable_iam_auth: bool = False,
 ) -> str:
     """Asynchronously requests an ephemeral certificate from the Cloud SQL Instance.
 
@@ -120,6 +121,10 @@ async def _get_ephemeral(
     :type pub_key:
     :param str: A string representing PEM-encoded RSA public key.
 
+    :type enable_iam_auth: bool
+    :param enable_iam_auth
+        Enables IAM based authentication for Postgres instances.
+
     :rtype: str
     :returns: An ephemeral certificate from the Cloud SQL instance that allows
           authorized connections to the instance.
@@ -141,7 +146,7 @@ async def _get_ephemeral(
     elif not isinstance(pub_key, str):
         raise TypeError(f"pub_key must be of type str, got {type(pub_key)}")
 
-    if not credentials.valid:
+    if not credentials.valid or enable_iam_auth:
         request = google.auth.transport.requests.Request()
         credentials.refresh(request)
 
@@ -155,6 +160,9 @@ async def _get_ephemeral(
 
     data = {"public_key": pub_key}
 
+    if enable_iam_auth:
+        data["access_token"] = credentials.token
+
     resp = await client_session.post(
         url, headers=headers, json=data, raise_for_status=True
     )

tests/conftest.py

@@ -65,7 +65,7 @@ def connect_string() -> str:
     returns it.
     """
     try:
-        connect_string = os.environ["INSTANCE_CONNECTION_NAME"]
+        connect_string = os.environ["MYSQL_CONNECTION_NAME"]
     except KeyError:
         raise KeyError(
             "Please set environment variable 'INSTANCE_CONNECTION"

tests/system/test_pg8000_iam_auth.py

@@ -0,0 +1,76 @@
+""""
+Copyright 2021 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 os
+import uuid
+from typing import Generator
+
+import pg8000
+import pytest
+import sqlalchemy
+from google.cloud.sql.connector import connector
+
+table_name = f"books_{uuid.uuid4().hex}"
+
+
+def init_connection_engine() -> sqlalchemy.engine.Engine:
+    def getconn() -> pg8000.dbapi.Connection:
+        conn: pg8000.dbapi.Connection = connector.connect(
+            os.environ["POSTGRES_IAM_CONNECTION_NAME"],
+            "pg8000",
+            user=os.environ["POSTGRES_IAM_USER"],
+            db=os.environ["POSTGRES_DB"],
+            enable_iam_auth=True,
+        )
+        return conn
+
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=getconn,
+    )
+    engine.dialect.description_encoding = None
+    return engine
+
+
+@pytest.fixture(name="pool")
+def setup() -> Generator:
+    pool = init_connection_engine()
+
+    with pool.connect() as conn:
+        conn.execute(
+            f"CREATE TABLE IF NOT EXISTS {table_name}"
+            " ( id CHAR(20) NOT NULL, title TEXT NOT NULL );"
+        )
+
+    yield pool
+
+    with pool.connect() as conn:
+        conn.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def test_pooled_connection_with_pg8000_iam_auth(pool: sqlalchemy.engine.Engine) -> None:
+    insert_stmt = sqlalchemy.text(
+        f"INSERT INTO {table_name} (id, title) VALUES (:id, :title)",
+    )
+    with pool.connect() as conn:
+        conn.execute(insert_stmt, id="book1", title="Book One")
+        conn.execute(insert_stmt, id="book2", title="Book Two")
+
+    select_stmt = sqlalchemy.text(f"SELECT title FROM {table_name} ORDER BY ID;")
+    with pool.connect() as conn:
+        rows = conn.execute(select_stmt).fetchall()
+        titles = [row[0] for row in rows]
+
+    assert titles == ["Book One", "Book Two"]

tests/unit/test_instance_connection_manager.py

@@ -54,10 +54,10 @@ def test_InstanceConnectionManager_init(async_loop: asyncio.AbstractEventLoop) -
 
 @pytest.mark.asyncio
 async def test_InstanceConnectionManager_perform_refresh(
-    icm: InstanceConnectionManager,
+    icm: InstanceConnectionManager, async_loop: asyncio.AbstractEventLoop
 ) -> None:
     """
-    Test to check whether _get_perform works as described given valid
+    Test to check whether _perform_refresh works as described given valid
     conditions.
     """
     task = await icm._perform_refresh()