fix(sqldata): restrict fallback to FAILED_PRECONDITION, defer metadata, and support psycopg

google/cloud/sql/connector/connector.py

@@ -123,7 +123,7 @@ def __init__(
         refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
         resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
         failover_period: int = 30,
-        sql_data_endpoint: str = "sqladmin.googleapis.com",
+        sql_data_endpoint: str | None = None,
         sql_data_stream_timeout: int = 7200,
         resource_exhausted_cooldown_period: float = 5.0,
     ) -> None:
@@ -180,7 +180,8 @@ def __init__(
                 Must be used with `resolver=DnsResolver` to have any effect.
                 Default: 30
 
-            sql_data_endpoint (str): Endpoint host for SQL Data Service calls.
+            sql_data_endpoint (str | None): Endpoint host for SQL Data Service
+                calls. If not specified, derived from `sqladmin_api_endpoint`.
                 Default: "sqladmin.googleapis.com".
 
             sql_data_stream_timeout (int): Timeout in seconds for the SQL Data
@@ -276,14 +277,20 @@ def __init__(
                 "configured the universe domain explicitly, `googleapis.com` "
                 "is the default."
             )
-        self._sql_data_endpoint = sql_data_endpoint
+        if sql_data_endpoint:
+            self._sql_data_endpoint = sql_data_endpoint
+        else:
+            self._sql_data_endpoint = (
+                self._sqladmin_api_endpoint.split("://", 1)[-1].split("/", 1)[0]
+            )
         self._sql_data_stream_timeout = sql_data_stream_timeout
         self._resource_exhausted_cooldown_period = (
             resource_exhausted_cooldown_period
         )
         self._sql_data_fallback_cache: set[str] = set()
         self._sql_data_conn_state: dict[str, SqlDataConnState] = {}
         self._sqldata_clients: set[Any] = set()
+        self._sqldata_client: SqlDataClient | None = None
 
 
 
@@ -499,16 +506,20 @@ async def connect_async(
                         )
                         kwargs["user"] = formatted_user
 
-                sqldata_client = SqlDataClient(
-                    endpoint=self._sql_data_endpoint,
-                    credentials=self._credentials,
-                    quota_project=self._quota_project,
-                    timeout=self._sql_data_stream_timeout,
-                )
-                self._sqldata_clients.add(sqldata_client)
-                sqldata_client._on_close_callbacks.append(
-                    lambda: self._sqldata_clients.discard(sqldata_client)
-                )
+                if self._sqldata_client is None:
+                    sqldata_client = SqlDataClient(
+                        endpoint=self._sql_data_endpoint,
+                        credentials=self._credentials,
+                        quota_project=self._quota_project,
+                        timeout=self._sql_data_stream_timeout,
+                    )
+                    self._sqldata_client = sqldata_client
+                    self._sqldata_clients.add(sqldata_client)
+                    sqldata_client._on_close_callbacks.append(
+                        lambda: self._sqldata_clients.discard(sqldata_client)
+                    )
+                else:
+                    sqldata_client = self._sqldata_client
 
                 def on_resource_exhausted(err: Exception) -> None:
                     backoff = state.record_exhausted(
@@ -736,6 +747,7 @@ async def close_async(self) -> None:
         """Helper function to cancel the cache's tasks
         and close aiohttp.ClientSession."""
         self._closed = True
+        self._sqldata_client = None
         if self._client:
             await self._client.close()
         await asyncio.gather(
@@ -758,7 +770,7 @@ async def create_async_connector(
     refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
     resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
     failover_period: int = 30,
-    sql_data_endpoint: str = "sqladmin.googleapis.com",
+    sql_data_endpoint: str | None = None,
     sql_data_stream_timeout: int = 7200,
     resource_exhausted_cooldown_period: float = 5.0,
 ) -> Connector:
@@ -818,7 +830,8 @@ async def create_async_connector(
             Must be used with `resolver=DnsResolver` to have any effect.
             Default: 30
 
-        sql_data_endpoint (str): Endpoint host for SQL Data Service calls.
+        sql_data_endpoint (str | None): Endpoint host for SQL Data Service
+            calls. If not specified, derived from `sqladmin_api_endpoint`.
             Default: "sqladmin.googleapis.com".
 
         sql_data_stream_timeout (int): Timeout in seconds for the SQL Data

google/cloud/sql/connector/psycopg.py

@@ -28,8 +28,46 @@
 logger = logging.getLogger(name=__name__)
 
 
+def _proxy_sqldata(local: socket.socket, remote: socket.socket) -> None:
+    """Bidirectional thread-based proxy for queue-backed SqlDataSocket."""
+
+    def forward_local_to_remote() -> None:
+        try:
+            while data := local.recv(8192):
+                remote.sendall(data)
+        except OSError:
+            pass
+        finally:
+            remote.close()
+
+    threading.Thread(target=forward_local_to_remote, daemon=True).start()
+    try:
+        while True:
+            try:
+                if not (data := remote.recv(8192)):
+                    break
+                local.sendall(data)
+            except socket.timeout:
+                if not getattr(remote, "_first_read_done", False) or getattr(
+                    remote, "_closed", False
+                ):
+                    break
+    except OSError:
+        pass
+    finally:
+        try:
+            local.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+        local.close()
+        remote.close()
+
+
 def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None:
     """Single-threaded selectors-based proxy to avoid SSLSocket thread-safety issues."""
+    if hasattr(remote, "_read_queue"):
+        _proxy_sqldata(local, remote)
+        return
     sel = selectors.DefaultSelector()
     sel.register(local, selectors.EVENT_READ, data="local")
     sel.register(remote, selectors.EVENT_READ, data="remote")