feat: Add attempt_direct_path argument to create_channel

google/api_core/grpc_helpers.py

@@ -271,11 +271,24 @@ def _create_composite_credentials(
     # Create a set of grpc.CallCredentials using the metadata plugin.
     google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)
 
-    if ssl_credentials is None:
-        ssl_credentials = grpc.ssl_channel_credentials()
+    # if `ssl_credentials` is set, use `grpc.composite_channel_credentials` instead of
+    # `grpc.compute_engine_channel_credentials` as the former supports passing
+    # `ssl_credentials` via `channel_credentials` which is needed for mTLS.
+    if ssl_credentials:
+        # Combine the ssl credentials and the authorization credentials.
+        # See https://grpc.github.io/grpc/python/grpc.html#grpc.composite_channel_credentials
+        return grpc.composite_channel_credentials(
+            ssl_credentials, google_auth_credentials
+        )
+    else:
+        # Use grpc.compute_engine_channel_credentials in order to support Direct Path.
+        # See https://grpc.github.io/grpc/python/grpc.html#grpc.compute_engine_channel_credentials
 
-    # Combine the ssl credentials and the authorization credentials.
-    return grpc.composite_channel_credentials(ssl_credentials, google_auth_credentials)
+        # TODO(<insert bug to github issue>): Although `grpc.compute_engine_channel_credentials`
+        # returns channel credentials outside of GCE, we should determine if there is a way to
+        # reliably detect when the client is in a GCE environment so that
+        # `grpc.compute_engine_channel_credentials` is not called outside of GCE.
+        return grpc.compute_engine_channel_credentials(google_auth_credentials)
 
 
 def create_channel(
@@ -288,6 +301,7 @@ def create_channel(
     default_scopes=None,
     default_host=None,
     compression=None,
+    attempt_direct_path: bool = None,
     **kwargs,
 ):
     """Create a secure channel with credentials.
@@ -311,6 +325,16 @@ def create_channel(
         default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
         compression (grpc.Compression): An optional value indicating the
             compression method to be used over the lifetime of the channel.
+        attempt_direct_path (bool): If set, Direct Path will be attempted when
+            the request is made. Direct Path provides a proxyless connection which
+            increases the available throughput, reduces latency, and increases
+            reliability. Outside of GCE, the direct path request may fallback
+            to DNS if this is configured by the Service. This argument should only
+            be set in a GCE environment and for Services that are known to support Direct Path.
+            If a `ServiceUnavailable` response is received when the request is sent, it is
+            recommended that the client repeat the request with `attempt_direct_path` set to `False`
+            as the Service may not support Direct Path. Using `ssl_credentials` with `attempt_direct_path`
+            set to `True` will result in `ValueError` as it is not yet supported.
         kwargs: Additional key-word args passed to
             :func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`.
             Note: `grpc_gcp` is only supported in environments with protobuf < 4.0.0.
@@ -320,8 +344,15 @@ def create_channel(
 
     Raises:
         google.api_core.DuplicateCredentialArgs: If both a credentials object and credentials_file are passed.
+        ValueError: If `ssl_credentials` is set and `attempt_direct_path` is set to `True`.
     """
 
+    # If `ssl_credentials` is set and `attempt_direct_path` is set to `True`,
+    # raise ValueError as this is not yet supported.
+    # TODO(<insert bug to github issue>): Add link to Github Issue
+    if ssl_credentials is not None and attempt_direct_path:
+        raise ValueError("Using ssl_credentials with Direct Path is not supported")
+
     composite_credentials = _create_composite_credentials(
         credentials=credentials,
         credentials_file=credentials_file,
@@ -332,12 +363,31 @@ def create_channel(
         default_host=default_host,
     )
 
+    # Note that grpcio-gcp is deprecated
     if HAS_GRPC_GCP:  # pragma: NO COVER
         if compression is not None and compression != grpc.Compression.NoCompression:
             _LOGGER.debug(
                 "Compression argument is being ignored for grpc_gcp.secure_channel creation."
             )
+        if attempt_direct_path:
+            warnings.warn(
+                """The `attempt_direct_path` argument is ignored for grpc_gcp.secure_channel creation.""",
+                DeprecationWarning,
+            )
         return grpc_gcp.secure_channel(target, composite_credentials, **kwargs)
+
+    if attempt_direct_path:
+        dns_prefix = "dns:///"
+        # Remove "dns:///" if `attempt_direct_path` is set to True as
+        # the Direct Path prefix `google-c2p:///` will be used instead.
+        target = target.replace(dns_prefix, "")
+
+        direct_path_prefix = ":///"
+        if direct_path_prefix not in target:
+            target_without_port = target.split(":")[0]
+            # Modify the target to use Direct Path by adding the `google-c2p:///` prefix
+            target = f"google-c2p{direct_path_prefix}{target_without_port}"
+
     return grpc.secure_channel(
         target, composite_credentials, compression=compression, **kwargs
     )