feat(experimental): Add grpc client (#1533)

Pulkit0110 · web-flow · commit 5674587f2aa3 · 2025-09-15T12:58:48.000+05:30
* feat: Add grpc client

* add an option to configure direct path

* fix lint errors

* move the files to _experimental folder

* move the test file under tests/unit

* fix patch errors

* update the imports for _storage_v2

* change the default directpath value to True

* add getter for grpc_client

* resolve comments

* add more tests

* minor change

* adding helper function for credentials
diff --git a/google/cloud/storage/_experimental/grpc_client.py b/google/cloud/storage/_experimental/grpc_client.py
@@ -0,0 +1,122 @@
+# Copyright 2025 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.
+
+"""A client for interacting with Google Cloud Storage using the gRPC API."""
+
+from google.cloud.client import ClientWithProject
+from google.cloud import _storage_v2 as storage_v2
+
+_marker = object()
+
+
+class GrpcClient(ClientWithProject):
+    """A client for interacting with Google Cloud Storage using the gRPC API.
+
+    :type project: str or None
+    :param project: The project which the client acts on behalf of. If not
+                    passed, falls back to the default inferred from the
+                    environment.
+
+    :type credentials: :class:`~google.auth.credentials.Credentials`
+    :param credentials: (Optional) The OAuth2 Credentials to use for this
+                        client. If not passed, falls back to the default
+                        inferred from the environment.
+
+    :type client_info: :class:`~google.api_core.client_info.ClientInfo`
+    :param client_info:
+        The client info used to send a user-agent string along with API
+        requests. If ``None``, then default info will be used. Generally,
+        you only need to set this if you're developing your own library
+        or partner tool.
+
+    :type client_options: :class:`~google.api_core.client_options.ClientOptions` or :class:`dict`
+    :param client_options: (Optional) Client options used to set user options
+        on the client. A non-default universe domain or API endpoint should be
+        set through client_options.
+
+    :type api_key: string
+    :param api_key:
+        (Optional) An API key. Mutually exclusive with any other credentials.
+        This parameter is an alias for setting `client_options.api_key` and
+        will supersede any API key set in the `client_options` parameter.
+
+    :type attempt_direct_path: bool
+    :param attempt_direct_path:
+        (Optional) Whether to attempt to use DirectPath for gRPC connections.
+        This provides a direct, unproxied connection to GCS for lower latency
+        and higher throughput, and is highly recommended when running on Google
+        Cloud infrastructure. Defaults to ``True``.
+    """
+
+    def __init__(
+        self,
+        project=_marker,
+        credentials=None,
+        client_info=None,
+        client_options=None,
+        *,
+        api_key=None,
+        attempt_direct_path=True,
+    ):
+        super(GrpcClient, self).__init__(project=project, credentials=credentials)
+
+        if isinstance(client_options, dict):
+            if api_key:
+                client_options["api_key"] = api_key
+        elif client_options is None:
+            client_options = {} if not api_key else {"api_key": api_key}
+        elif api_key:
+            client_options.api_key = api_key
+
+        self._grpc_client = self._create_gapic_client(
+            credentials=credentials,
+            client_info=client_info,
+            client_options=client_options,
+            attempt_direct_path=attempt_direct_path,
+        )
+
+    def _create_gapic_client(
+        self,
+        credentials=None,
+        client_info=None,
+        client_options=None,
+        attempt_direct_path=True,
+    ):
+        """Creates and configures the low-level GAPIC `storage_v2` client."""
+        transport_cls = storage_v2.StorageClient.get_transport_class("grpc")
+
+        channel = transport_cls.create_channel(attempt_direct_path=attempt_direct_path)
+
+        transport = transport_cls(credentials=credentials, channel=channel)
+
+        return storage_v2.StorageClient(
+            credentials=credentials,
+            transport=transport,
+            client_info=client_info,
+            client_options=client_options,
+        )
+
+    @property
+    def grpc_client(self):
+        """The underlying gRPC client.
+
+        This property gives users direct access to the `storage_v2.StorageClient`
+         instance. This can be useful for accessing
+        newly added or experimental RPCs that are not yet exposed through
+        the high-level GrpcClient.
+
+        Returns:
+            google.cloud.storage_v2.StorageClient: The configured GAPIC client.
+        """
+        return self._grpc_client
diff --git a/tests/unit/test_grpc_client.py b/tests/unit/test_grpc_client.py
@@ -0,0 +1,208 @@
+# Copyright 2025 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.
+
+import unittest
+from unittest import mock
+from google.auth import credentials as auth_credentials
+from google.api_core import client_options as client_options_lib
+
+
+def _make_credentials(spec=None):
+    if spec is None:
+        return mock.Mock(spec=auth_credentials.Credentials)
+    return mock.Mock(spec=spec)
+
+
+class TestGrpcClient(unittest.TestCase):
+    @mock.patch("google.cloud.client.ClientWithProject.__init__")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_constructor_defaults_and_options(
+        self, mock_storage_client, mock_base_client
+    ):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_transport_cls = mock.MagicMock()
+        mock_storage_client.get_transport_class.return_value = mock_transport_cls
+        mock_creds = _make_credentials(spec=["_base", "_get_project_id"])
+        mock_client_info = mock.Mock()
+        client_options_dict = {"api_endpoint": "test.endpoint"}
+
+        mock_base_instance = mock_base_client.return_value
+        mock_base_instance._credentials = mock_creds
+
+        client = grpc_client.GrpcClient(
+            project="test-project",
+            credentials=mock_creds,
+            client_info=mock_client_info,
+            client_options=client_options_dict,
+        )
+
+        # 1. Assert that the base class was initialized correctly.
+        mock_base_client.assert_called_once_with(
+            project="test-project", credentials=mock_creds
+        )
+
+        # 2. Assert DirectPath is ON by default.
+        mock_storage_client.get_transport_class.assert_called_once_with("grpc")
+        mock_transport_cls.create_channel.assert_called_once_with(
+            attempt_direct_path=True
+        )
+
+        # 3. Assert the GAPIC client was created with the correct options.
+        mock_transport = mock_transport_cls.return_value
+        mock_storage_client.assert_called_once_with(
+            credentials=mock_creds,
+            transport=mock_transport,
+            client_info=mock_client_info,
+            client_options=client_options_dict,
+        )
+
+        # 4. Assert the client instance holds the mocked GAPIC client.
+        self.assertIs(client.grpc_client, mock_storage_client.return_value)
+
+    @mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_constructor_disables_direct_path(
+        self, mock_storage_client, mock_base_client
+    ):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_transport_cls = mock.MagicMock()
+        mock_storage_client.get_transport_class.return_value = mock_transport_cls
+        mock_creds = _make_credentials()
+        mock_base_instance = mock_base_client.return_value
+        mock_base_instance._credentials = mock_creds
+
+        grpc_client.GrpcClient(
+            project="test-project",
+            credentials=mock_creds,
+            attempt_direct_path=False,
+        )
+
+        mock_transport_cls.create_channel.assert_called_once_with(
+            attempt_direct_path=False
+        )
+
+    @mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_constructor_initialize_with_api_key(
+        self, mock_storage_client, mock_base_client
+    ):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_transport_cls = mock.MagicMock()
+        mock_storage_client.get_transport_class.return_value = mock_transport_cls
+        mock_creds = _make_credentials()
+        mock_creds.project_id = None
+
+        mock_base_instance = mock_base_client.return_value
+        mock_base_instance._credentials = mock_creds
+
+        # Instantiate with just the api_key.
+        grpc_client.GrpcClient(
+            project="test-project", credentials=mock_creds, api_key="test-api-key"
+        )
+
+        # Assert that the GAPIC client was called with client_options
+        # that contains the api_key.
+        mock_transport = mock_transport_cls.return_value
+        mock_storage_client.assert_called_once_with(
+            credentials=mock_creds,
+            transport=mock_transport,
+            client_info=None,
+            client_options={"api_key": "test-api-key"},
+        )
+
+    @mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_grpc_client_property(self, mock_storage_client, mock_base_client):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_creds = _make_credentials()
+        mock_base_client.return_value._credentials = mock_creds
+
+        client = grpc_client.GrpcClient(project="test-project", credentials=mock_creds)
+
+        retrieved_client = client.grpc_client
+
+        self.assertIs(retrieved_client, mock_storage_client.return_value)
+
+    @mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_constructor_with_api_key_and_client_options(
+        self, mock_storage_client, mock_base_client
+    ):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_transport_cls = mock.MagicMock()
+        mock_storage_client.get_transport_class.return_value = mock_transport_cls
+        mock_transport = mock_transport_cls.return_value
+
+        mock_creds = _make_credentials()
+        mock_base_instance = mock_base_client.return_value
+        mock_base_instance._credentials = mock_creds
+
+        client_options_obj = client_options_lib.ClientOptions(
+            api_endpoint="test.endpoint"
+        )
+        self.assertIsNone(client_options_obj.api_key)
+
+        grpc_client.GrpcClient(
+            project="test-project",
+            credentials=mock_creds,
+            client_options=client_options_obj,
+            api_key="new-test-key",
+        )
+
+        mock_storage_client.assert_called_once_with(
+            credentials=mock_creds,
+            transport=mock_transport,
+            client_info=None,
+            client_options=client_options_obj,
+        )
+        self.assertEqual(client_options_obj.api_key, "new-test-key")
+
+    @mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
+    @mock.patch("google.cloud._storage_v2.StorageClient")
+    def test_constructor_with_api_key_and_dict_options(
+        self, mock_storage_client, mock_base_client
+    ):
+        from google.cloud.storage._experimental import grpc_client
+
+        mock_creds = _make_credentials()
+        mock_base_instance = mock_base_client.return_value
+        mock_base_instance._credentials = mock_creds
+        mock_transport_cls = mock.MagicMock()
+        mock_storage_client.get_transport_class.return_value = mock_transport_cls
+        mock_transport = mock_transport_cls.return_value
+
+        client_options_dict = {"api_endpoint": "test.endpoint"}
+
+        grpc_client.GrpcClient(
+            project="test-project",
+            credentials=mock_creds,
+            client_options=client_options_dict,
+            api_key="new-test-key",
+        )
+
+        expected_options = {
+            "api_endpoint": "test.endpoint",
+            "api_key": "new-test-key",
+        }
+        mock_storage_client.assert_called_once_with(
+            credentials=mock_creds,
+            transport=mock_transport,
+            client_info=None,
+            client_options=expected_options,
+        )
