feat: replace ssl_object with sslcontext and build a socket manually

Author: Uziel Silva

google/cloud/sql/connector/connector.py

@@ -20,6 +20,7 @@
 from functools import partial
 import logging
 import os
+import socket
 from threading import Thread
 from types import TracebackType
 from typing import Any, Callable, Optional, Union
@@ -514,7 +515,11 @@ async def connect_async(
                     instance_connection_string, asyncio.Protocol, **kwargs
                 )
                 # See https://docs.python.org/3/library/asyncio-protocol.html#asyncio.BaseTransport.get_extra_info
-                sock = tx.get_extra_info("ssl_object")
+                ctx = tx.get_extra_info("sslcontext")
+                sock = ctx.wrap_socket(
+                    socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+                    server_hostname=ip_address,
+                )
                 connect_partial = partial(connector, ip_address, sock, **kwargs)
                 return await self._loop.run_in_executor(None, connect_partial)
 

tests/system/test_psycopg_connection.py

@@ -14,13 +14,16 @@
 limitations under the License.
 """
 
+import asyncio
 from datetime import datetime
 import os
 
 # [START cloud_sql_connector_postgres_psycopg]
 from typing import Union
 
 from psycopg import Connection
+import pytest
+import logging
 import sqlalchemy
 
 from google.cloud.sql.connector import Connector
@@ -29,98 +32,56 @@
 
 SERVER_PROXY_PORT = 3307
 
-async def create_sqlalchemy_engine(
-    instance_connection_name: str,
-    user: str,
-    password: str,
-    db: str,
-    ip_type: str = "public",
-    refresh_strategy: str = "background",
-    resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
-) -> tuple[sqlalchemy.engine.Engine, Connector]:
-    """Creates a connection pool for a Cloud SQL instance and returns the pool
-    and the connector. Callers are responsible for closing the pool and the
-    connector.
-
-    A sample invocation looks like:
-
-        engine, connector = create_sqlalchemy_engine(
-            inst_conn_name,
-            user,
-            password,
-            db,
-        )
-        with engine.connect() as conn:
-            time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
-            conn.commit()
-            curr_time = time[0]
-            # do something with query result
-            connector.close()
-
-    Args:
-        instance_connection_name (str):
-            The instance connection name specifies the instance relative to the
-            project and region. For example: "my-project:my-region:my-instance"
-        user (str):
-            The database user name, e.g., root
-        password (str):
-            The database user's password, e.g., secret-password
-        db (str):
-            The name of the database, e.g., mydb
-        ip_type (str):
-            The IP type of the Cloud SQL instance to connect to. Can be one
-            of "public", "private", or "psc".
-        refresh_strategy (Optional[str]):
-            Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
-            or "background". For serverless environments use "lazy" to avoid
-            errors resulting from CPU being throttled.
-        resolver (Optional[google.cloud.sql.connector.DefaultResolver]):
-            Resolver class for resolving instance connection name. Use
-            google.cloud.sql.connector.DnsResolver when resolving DNS domain
-            names or google.cloud.sql.connector.DefaultResolver for regular
-            instance connection names ("my-project:my-region:my-instance").
-    """
-    connector = Connector(refresh_strategy=refresh_strategy, resolver=resolver)
-    unix_socket_folder = "/tmp/conn"
-    unix_socket_path = f"{unix_socket_folder}/.s.PGSQL.3307"
-    await connector.start_unix_socket_proxy_async(
-        instance_connection_name,
-        unix_socket_path,
-        ip_type=ip_type,  # can be "public", "private" or "psc"
-    )
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.create_engine(
-        "postgresql+psycopg://",
-        creator=lambda: Connection.connect(
-            f"host={unix_socket_folder} port={SERVER_PROXY_PORT} dbname={db} user={user} password={password} sslmode=require",
-            user=user,
-            password=password,
-            dbname=db,
-            autocommit=True,
-        )
-    )
-    
-    return engine, connector
-
+logger = logging.getLogger(name=__name__)
 
 # [END cloud_sql_connector_postgres_psycopg]
 
 
+@pytest.mark.asyncio
 async def test_psycopg_connection() -> None:
     """Basic test to get time from database."""
-    inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
+    instance_connection_name = os.environ["POSTGRES_CONNECTION_NAME"]
     user = os.environ["POSTGRES_USER"]
     password = os.environ["POSTGRES_PASS"]
     db = os.environ["POSTGRES_DB"]
     ip_type = os.environ.get("IP_TYPE", "public")
 
-    engine, connector = await create_sqlalchemy_engine(
-        inst_conn_name, user, password, db, ip_type
-    )
-    with engine.connect() as conn:
-        time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
-        conn.commit()
-        curr_time = time[0]
-        assert type(curr_time) is datetime
-    connector.close()
+    unix_socket_folder = "/tmp/conn"
+    unix_socket_path = f"{unix_socket_folder}/.s.PGSQL.3307"
+
+    async with Connector(
+        refresh_strategy='lazy', resolver=DefaultResolver
+    ) as connector:
+        # Open proxy connection
+        # start the proxy server
+        
+        await connector.start_unix_socket_proxy_async(
+            instance_connection_name,
+            unix_socket_path,
+            driver="psycopg",
+            user=user,
+            password=password,
+            db=db,
+            ip_type=ip_type,  # can be "public", "private" or "psc"
+        )
+        
+        # Wait for server to start
+        await asyncio.sleep(0.5)
+
+        engine = sqlalchemy.create_engine(
+            "postgresql+psycopg://",
+            creator=lambda: Connection.connect(
+                f"host={unix_socket_folder} port={SERVER_PROXY_PORT} dbname={db} user={user} password={password} sslmode=require",
+                user=user,
+                password=password,
+                dbname=db,
+                autocommit=True,
+            )
+        )
+
+        with engine.connect() as conn:
+            time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
+            conn.commit()
+            curr_time = time[0]
+            assert type(curr_time) is datetime
+        connector.close()