fix(main): fix feedback PR

Changelog:
- Make local_socket_path configurable
- Set right file permissions
- Handle exceptions properly
- Use asyncio and the main loop to stop the local proxy and clear the file when the connector is stopped

Author: Uziel Silva

google/cloud/sql/connector/connector.py

@@ -44,10 +44,12 @@
 from google.cloud.sql.connector.resolver import DnsResolver
 from google.cloud.sql.connector.utils import format_database_user
 from google.cloud.sql.connector.utils import generate_keys
+from google.cloud.sql.connector.proxy import start_local_proxy
 
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
+LOCAL_PROXY_DRIVERS = ["psycopg"]
 SERVER_PROXY_PORT = 3307
 _DEFAULT_SCHEME = "https://"
 _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
@@ -383,7 +385,7 @@ async def connect_async(
             # async drivers are unblocking and can be awaited directly
             if driver in ASYNC_DRIVERS:
                 return await connector(
-                    ip_address,
+                    host,
                     await conn_info.create_ssl_context(enable_iam_auth),
                     **kwargs,
                 )
@@ -393,14 +395,26 @@ async def connect_async(
                 socket.create_connection((ip_address, SERVER_PROXY_PORT)),
                 server_hostname=ip_address,
             )
+
+            host = ip_address
+            # start local proxy if driver needs it
+            if driver in LOCAL_PROXY_DRIVERS:
+                local_socket_path = kwargs.pop("local_socket_path", "/tmp/connector-socket")
+                host = local_socket_path
+                start_local_proxy(
+                    sock,
+                    socket_path=f"{local_socket_path}/.s.PGSQL.{SERVER_PROXY_PORT}",
+                    loop=self._loop
+                )
+
             # If this connection was opened using a domain name, then store it
             # for later in case we need to forcibly close it on failover.
             if conn_info.conn_name.domain_name:
                 monitored_cache.sockets.append(sock)
             # Synchronous drivers are blocking and run using executor
             connect_partial = partial(
                 connector,
-                ip_address,
+                host,
                 sock,
                 **kwargs,
             )

google/cloud/sql/connector/exceptions.py

@@ -84,3 +84,9 @@ class CacheClosedError(Exception):
     Exception to be raised when a ConnectionInfoCache can not be accessed after
     it is closed.
     """
+
+
+class LocalProxyStartupError(Exception):
+    """
+    Exception to be raised when a the local UNIX-socket based proxy can not be started.
+    """

google/cloud/sql/connector/proxy.py

@@ -16,48 +16,75 @@
 
 import socket
 import os
-import threading
+import ssl
+import asyncio
 from pathlib import Path
+from typing import Optional
+
+from google.cloud.sql.connector.exceptions import LocalProxyStartupError
 
 SERVER_PROXY_PORT = 3307
+LOCAL_PROXY_MAX_MESSAGE_SIZE = 10485760
 
 def start_local_proxy(
-  ssl_sock,
-  socket_path,
+    ssl_sock: ssl.SSLSocket,
+    socket_path: Optional[str] = "/tmp/connector-socket",
+    loop: Optional[asyncio.AbstractEventLoop] = None,
 ):
-  path_parts = socket_path.rsplit('/', 1)
-  parent_directory = '/'.join(path_parts[:-1])
+    """Helper function to start a UNIX based local proxy for
+    transport messages through the SSL Socket.
+
+    Args:
+        ssl_sock (ssl.SSLSocket): An SSLSocket object created from the Cloud SQL
+            server CA cert and ephemeral cert.
+        socket_path: A system path that is going to be used to store the socket.
+        loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks.
+
+    Raises:
+        LocalProxyStartupError: Local UNIX socket based proxy was not able to
+        get started.
+    """
+    unix_socket = None
 
-  desired_path = Path(parent_directory)
-  desired_path.mkdir(parents=True, exist_ok=True)
+    try:
+        path_parts = socket_path.rsplit('/', 1)
+        parent_directory = '/'.join(path_parts[:-1])
 
-  if os.path.exists(socket_path):
-      os.remove(socket_path)
-  conn_unix = None
-  unix_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        desired_path = Path(parent_directory)
+        desired_path.mkdir(parents=True, exist_ok=True)
 
-  unix_socket.bind(socket_path)
-  unix_socket.listen(1)
+        if os.path.exists(socket_path):
+            os.remove(socket_path)
+        unix_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 
-  threading.Thread(target=local_communication, args=(unix_socket, ssl_sock, socket_path)).start()
+        unix_socket.bind(socket_path)
+        unix_socket.listen(1)
+        unix_socket.setblocking(False)
+        os.chmod(socket_path, 0o600)
+    except Exception:
+        raise LocalProxyStartupError(
+            'Local UNIX socket based proxy was not able to get started.'
+        )
 
+    loop.create_task(local_communication(unix_socket, ssl_sock, socket_path, loop))
 
-def local_communication(
-  unix_socket, ssl_sock, socket_path
+
+async def local_communication(
+  unix_socket, ssl_sock, socket_path, loop
 ):
-  try:
-    conn_unix, addr_unix = unix_socket.accept()
-
-    while True:
-      data = conn_unix.recv(10485760)
-      if not data:
-        break
-      ssl_sock.sendall(data)
-      response = ssl_sock.recv(10485760)
-      conn_unix.sendall(response)
-
-  finally:
-    if conn_unix is not None:
-      conn_unix.close()
-    unix_socket.close()
-    os.remove(socket_path) # Clean up the socket file
+    try:
+        client, _ = await loop.sock_accept(unix_socket)
+
+        while True:
+            data = await loop.sock_recv(client, LOCAL_PROXY_MAX_MESSAGE_SIZE)
+            if not data:
+              client.close()
+              break
+            ssl_sock.sendall(data)
+            response = ssl_sock.recv(LOCAL_PROXY_MAX_MESSAGE_SIZE)
+            await loop.sock_sendall(client, response)
+    except Exception:
+        pass
+    finally:
+        client.close()
+        os.remove(socket_path) # Clean up the socket file

google/cloud/sql/connector/psycopg.py

@@ -25,13 +25,12 @@
 
 
 def connect(
-    ip_address: str, sock: ssl.SSLSocket, **kwargs: Any
+    host: str, sock: ssl.SSLSocket, **kwargs: Any
 ) -> "psycopg.Connection":
     """Helper function to create a psycopg DB-API connection object.
 
     Args:
-        ip_address (str): A string containing an IP address for the Cloud SQL
-            instance.
+        host (str): A string containing the socket path used by the local proxy.
         sock (ssl.SSLSocket): An SSLSocket object created from the Cloud SQL
             server CA cert and ephemeral cert.
         kwargs: Additional arguments to pass to the psycopg connect method.
@@ -44,10 +43,7 @@ def connect(
         ImportError: The psycopg module cannot be imported.
     """
     try:
-        from psycopg.rows import dict_row
         from psycopg import Connection
-        import threading
-        from google.cloud.sql.connector.proxy import start_local_proxy
     except ImportError:
         raise ImportError(
             'Unable to import module "psycopg." Please install and try again.'
@@ -59,14 +55,9 @@ def connect(
 
     kwargs.pop("timeout", None)
 
-    start_local_proxy(sock, f"/tmp/connector-socket/.s.PGSQL.3307")
-
     conn = Connection.connect(
-        f"host=/tmp/connector-socket port={SERVER_PROXY_PORT} dbname={db} user={user} password={passwd} sslmode=require",
-        autocommit=True,
-        row_factory=dict_row,
+        f"host={host} port={SERVER_PROXY_PORT} dbname={db} user={user} password={passwd} sslmode=require",
         **kwargs
     )
 
-    conn.autocommit = True
     return conn

tests/system/test_psycopg_connection.py

@@ -18,12 +18,84 @@
 import os
 
 # [START cloud_sql_connector_postgres_psycopg]
+from typing import Union
+
+import sqlalchemy
 
 from google.cloud.sql.connector import Connector
 from google.cloud.sql.connector import DefaultResolver
+from google.cloud.sql.connector import DnsResolver
+
+
+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)
+
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+psycopg://",
+        creator=lambda: connector.connect(
+            instance_connection_name,
+            "psycopg",
+            user=user,
+            password=password,
+            db=db,
+            local_socket_path="/tmp/conn",
+            ip_type=ip_type,  # can be "public", "private" or "psc"
+            autocommit=True,
+        ),
+    )
+    return engine, connector
 
-from sqlalchemy.dialects.postgresql.base import PGDialect
-PGDialect._get_server_version_info = lambda *args: (9, 2)
 
 # [END cloud_sql_connector_postgres_psycopg]
 
@@ -36,25 +108,12 @@ def test_psycopg_connection() -> None:
     db = os.environ["POSTGRES_DB"]
     ip_type = os.environ.get("IP_TYPE", "public")
 
-    connector = Connector(refresh_strategy="background", resolver=DefaultResolver)
-
-    pool = connector.connect(
-        inst_conn_name,
-        "psycopg",
-        user=user,
-        password=password,
-        db=db,
-        ip_type=ip_type,  # can be "public", "private" or "psc"
+    engine, connector = create_sqlalchemy_engine(
+        inst_conn_name, user, password, db, ip_type
     )
-
-    with pool as conn:
-
-        # Open a cursor to perform database operations
-        with conn.cursor() as cur:
-
-            # Query the database and obtain data as Python objects.
-            cur.execute("SELECT NOW()")
-            curr_time = cur.fetchone()["now"]
-            assert type(curr_time) is datetime
-
-
+    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()