feat: add support for auto IAM authentication to Connector

google/api/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 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
+#
+#     http://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.

google/api/field_behavior_pb2.py

@@ -0,0 +1,43 @@
+# Copyright 2023 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
+#
+#     http://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.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: google/api/field_behavior.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1fgoogle/api/field_behavior.proto\x12\ngoogle.api\x1a google/protobuf/descriptor.proto*\xb6\x01\n\rFieldBehavior\x12\x1e\n\x1a\x46IELD_BEHAVIOR_UNSPECIFIED\x10\x00\x12\x0c\n\x08OPTIONAL\x10\x01\x12\x0c\n\x08REQUIRED\x10\x02\x12\x0f\n\x0bOUTPUT_ONLY\x10\x03\x12\x0e\n\nINPUT_ONLY\x10\x04\x12\r\n\tIMMUTABLE\x10\x05\x12\x12\n\x0eUNORDERED_LIST\x10\x06\x12\x15\n\x11NON_EMPTY_DEFAULT\x10\x07\x12\x0e\n\nIDENTIFIER\x10\x08:Q\n\x0e\x66ield_behavior\x12\x1d.google.protobuf.FieldOptions\x18\x9c\x08 \x03(\x0e\x32\x19.google.api.FieldBehaviorBp\n\x0e\x63om.google.apiB\x12\x46ieldBehaviorProtoP\x01ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\xa2\x02\x04GAPIb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.api.field_behavior_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension(field_behavior)
+
+  DESCRIPTOR._options = None
+  DESCRIPTOR._serialized_options = b'\n\016com.google.apiB\022FieldBehaviorProtoP\001ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\242\002\004GAPI'
+  _globals['_FIELDBEHAVIOR']._serialized_start=82
+  _globals['_FIELDBEHAVIOR']._serialized_end=264
+# @@protoc_insertion_point(module_scope)

google/cloud/alloydb/connector/client.py

@@ -67,6 +67,7 @@ def __init__(
         self._client = client if client else aiohttp.ClientSession(headers=headers)
         self._credentials = credentials
         self._alloydb_api_endpoint = alloydb_api_endpoint
+        self._user_agent = USER_AGENT
 
     async def _get_metadata(
         self,

google/cloud/alloydb/connector/connector.py

@@ -16,6 +16,8 @@
 
 import asyncio
 from functools import partial
+import socket
+import struct
 from threading import Thread
 from types import TracebackType
 from typing import Any, Dict, Optional, Type, TYPE_CHECKING
@@ -26,10 +28,14 @@
 from google.cloud.alloydb.connector.instance import Instance
 import google.cloud.alloydb.connector.pg8000 as pg8000
 from google.cloud.alloydb.connector.utils import generate_keys
+import google.cloud.alloydb_connectors_v1.proto.resources_pb2 as connectorspb
 
 if TYPE_CHECKING:
+    import ssl
     from google.auth.credentials import Credentials
 
+SERVER_PROXY_PORT = 5433
+
 
 class Connector:
     """A class to configure and create connections to Cloud SQL instances.
@@ -44,13 +50,15 @@ class Connector:
             Defaults to None, picking up project from environment.
         alloydb_api_endpoint (str): Base URL to use when calling
             the AlloyDB API endpoint. Defaults to "https://alloydb.googleapis.com".
+        enable_iam_auth (bool): Enables automatic IAM database authentication.
     """
 
     def __init__(
         self,
         credentials: Optional[Credentials] = None,
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
+        enable_iam_auth: bool = False,
     ) -> None:
         # create event loop and start it in background thread
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
@@ -60,6 +68,7 @@ def __init__(
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
+        self._enable_iam_auth = enable_iam_auth
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
@@ -122,6 +131,7 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
             self._client = AlloyDBClient(
                 self._alloydb_api_endpoint, self._quota_project, self._credentials
             )
+        enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
         # use existing connection info if possible
         if instance_uri in self._instances:
             instance = self._instances[instance_uri]
@@ -149,13 +159,80 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
 
         # synchronous drivers are blocking and run using executor
         try:
-            connect_partial = partial(connector, ip_address, context, **kwargs)
+            metadata_partial = partial(
+                self.metadata_exchange, ip_address, context, enable_iam_auth, driver
+            )
+            sock = await self._loop.run_in_executor(None, metadata_partial)
+            connect_partial = partial(connector, sock, **kwargs)
             return await self._loop.run_in_executor(None, connect_partial)
         except Exception:
             # we attempt a force refresh, then throw the error
             await instance.force_refresh()
             raise
 
+    def metadata_exchange(
+        self, ip_address: str, ctx: ssl.SSLContext, enable_iam_auth: bool, driver: str
+    ):
+        # Create socket and wrap with SSL/TLS context
+        sock = ctx.wrap_socket(
+            socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+            server_hostname=ip_address,
+        )
+        # set auth type for metadata exchange
+        if enable_iam_auth:
+            auth_type = connectorspb.MetadataExchangeRequest.AUTO_IAM
+        else:
+            auth_type = connectorspb.MetadataExchangeRequest.DB_NATIVE
+
+        # form metadata exchange request
+        req = connectorspb.MetadataExchangeRequest(
+            user_agent=self._client._user_agent + f"+{driver}",
+            auth_type=auth_type,
+            oauth2_token=self._credentials.token,
+        )
+
+        # set I/O timeout
+        sock.settimeout(30)
+
+        # pack big-endian unsigned integer
+        packed_len = struct.pack(">I", req.ByteSize())
+
+        # send message length
+        sock.sendall(packed_len)
+        # send message
+        sock.sendall(req.SerializeToString())
+
+        # form metadata exchange response
+        resp = connectorspb.MetadataExchangeResponse()
+
+        # read message length
+        message_len_buffer_size = struct.Struct("I").size
+        message_len_buffer = b""
+        while message_len_buffer_size > 0:
+            chunk = sock.recv(message_len_buffer_size)
+            if not chunk:
+                raise RuntimeError("connection closed before chunk was read")
+            message_len_buffer += chunk
+            message_len_buffer_size -= len(chunk)
+
+        (message_len,) = struct.unpack(">I", message_len_buffer)
+
+        # read message
+        buffer = b""
+        while message_len > 0:
+            chunk = sock.recv(message_len)
+            if not chunk:
+                raise RuntimeError("connection closed before chunk was read")
+            buffer += chunk
+            message_len -= len(chunk)
+        # parse mdx resp from buffer
+        resp.ParseFromString(buffer)
+
+        if resp.response_code != connectorspb.MetadataExchangeResponse.OK:
+            raise ValueError("Metadata Exchange request has failed")
+
+        return sock
+
     def __enter__(self) -> "Connector":
         """Enter context manager by returning Connector object"""
         return self

google/cloud/alloydb/connector/pg8000.py

@@ -11,45 +11,32 @@
 # 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 socket
-import ssl
 from typing import Any, TYPE_CHECKING
 
-SERVER_PROXY_PORT = 5433
-
 if TYPE_CHECKING:
     import pg8000
+    import ssl
 
 
 def connect(
-    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+    sock: "ssl.SSLSocket", **kwargs: Any
 ) -> "pg8000.dbapi.Connection":
     """Create a pg8000 DBAPI connection object.
 
     Args:
-        ip_address (str): IP address of AlloyDB instance to connect to.
-        ctx (ssl.SSLContext): Context used to create a TLS connection
-            with AlloyDB instance ssl certificates.
+        sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the
+            AlloyDB proxy server.
 
     Returns:
         pg8000.dbapi.Connection: A pg8000 Connection object for
         the AlloyDB instance.
     """
-    # Connecting through pg8000 is done by passing in an SSL Context and setting the
-    # "request_ssl" attr to false. This works because when "request_ssl" is false,
-    # the driver skips the database level SSL/TLS exchange, but still uses the
-    # ssl_context (if it is not None) to create the connection.
     try:
         import pg8000
     except ImportError:
         raise ImportError(
             'Unable to import module "pg8000." Please install and try again.'
         )
-    # Create socket and wrap with context.
-    sock = ctx.wrap_socket(
-        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
-        server_hostname=ip_address,
-    )
 
     user = kwargs.pop("user")
     db = kwargs.pop("db")

google/cloud/alloydb_connectors_v1/proto/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 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
+#
+#     http://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.

google/cloud/alloydb_connectors_v1/proto/resources_pb2.py

@@ -0,0 +1,52 @@
+# Copyright 2023 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
+#
+#     http://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.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: google/cloud/alloydb_connectors_v1/proto/resources.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.api import field_behavior_pb2 as google_dot_api_dot_field__behavior__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n8google/cloud/alloydb_connectors_v1/proto/resources.proto\x12\"google.cloud.alloydb.connectors.v1\x1a\x1fgoogle/api/field_behavior.proto\"\xe6\x01\n\x17MetadataExchangeRequest\x12\x18\n\nuser_agent\x18\x01 \x01(\tB\x04\xe2\x41\x01\x01\x12W\n\tauth_type\x18\x02 \x01(\x0e\x32\x44.google.cloud.alloydb.connectors.v1.MetadataExchangeRequest.AuthType\x12\x14\n\x0coauth2_token\x18\x03 \x01(\t\"B\n\x08\x41uthType\x12\x19\n\x15\x41UTH_TYPE_UNSPECIFIED\x10\x00\x12\r\n\tDB_NATIVE\x10\x01\x12\x0c\n\x08\x41UTO_IAM\x10\x02\"\xd3\x01\n\x18MetadataExchangeResponse\x12`\n\rresponse_code\x18\x01 \x01(\x0e\x32I.google.cloud.alloydb.connectors.v1.MetadataExchangeResponse.ResponseCode\x12\x13\n\x05\x65rror\x18\x02 \x01(\tB\x04\xe2\x41\x01\x01\"@\n\x0cResponseCode\x12\x1d\n\x19RESPONSE_CODE_UNSPECIFIED\x10\x00\x12\x06\n\x02OK\x10\x01\x12\t\n\x05\x45RROR\x10\x02\x42\xf5\x01\n&com.google.cloud.alloydb.connectors.v1B\x0eResourcesProtoP\x01ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\xaa\x02\"Google.Cloud.AlloyDb.Connectors.V1\xca\x02\"Google\\Cloud\\AlloyDb\\Connectors\\V1\xea\x02&Google::Cloud::AlloyDb::Connectors::V1b\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.cloud.alloydb_connectors_v1.proto.resources_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  DESCRIPTOR._serialized_options = b'\n&com.google.cloud.alloydb.connectors.v1B\016ResourcesProtoP\001ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\252\002\"Google.Cloud.AlloyDb.Connectors.V1\312\002\"Google\\Cloud\\AlloyDb\\Connectors\\V1\352\002&Google::Cloud::AlloyDb::Connectors::V1'
+  _METADATAEXCHANGEREQUEST.fields_by_name['user_agent']._options = None
+  _METADATAEXCHANGEREQUEST.fields_by_name['user_agent']._serialized_options = b'\342A\001\001'
+  _METADATAEXCHANGERESPONSE.fields_by_name['error']._options = None
+  _METADATAEXCHANGERESPONSE.fields_by_name['error']._serialized_options = b'\342A\001\001'
+  _globals['_METADATAEXCHANGEREQUEST']._serialized_start=130
+  _globals['_METADATAEXCHANGEREQUEST']._serialized_end=360
+  _globals['_METADATAEXCHANGEREQUEST_AUTHTYPE']._serialized_start=294
+  _globals['_METADATAEXCHANGEREQUEST_AUTHTYPE']._serialized_end=360
+  _globals['_METADATAEXCHANGERESPONSE']._serialized_start=363
+  _globals['_METADATAEXCHANGERESPONSE']._serialized_end=574
+  _globals['_METADATAEXCHANGERESPONSE_RESPONSECODE']._serialized_start=510
+  _globals['_METADATAEXCHANGERESPONSE_RESPONSECODE']._serialized_end=574
+# @@protoc_insertion_point(module_scope)

requirements.txt

@@ -2,3 +2,4 @@ aiohttp==3.9.1
 cryptography==41.0.7
 google-auth==2.25.2
 requests==2.31.0
+protobuf==4.25.1

setup.py

@@ -25,6 +25,7 @@
     "cryptography>=38.0.3",
     "requests",
     "google-auth",
+    "protobuf",
 ]
 
 package_root = os.path.abspath(os.path.dirname(__file__))