feat: add use_auth_w_custom_endpoint support (#941)

* feat: add support for use_auth_w_custom_endpoint * improve readability w bool _is_emulator_set * update tests Co-authored-by: Andrew Gorcester <[email protected]>
googleapis · Dec 7, 2022 · 5291c08 · 5291c08
1 parent 64406ca
commit 5291c08
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 44 deletions.
diff --git a/google/cloud/storage/_helpers.py b/google/cloud/storage/_helpers.py
@@ -33,17 +33,20 @@
 STORAGE_EMULATOR_ENV_VAR = "STORAGE_EMULATOR_HOST"
 """Environment variable defining host for Storage emulator."""
 
+_API_ENDPOINT_OVERRIDE_ENV_VAR = "API_ENDPOINT_OVERRIDE"
+"""This is an experimental configuration variable. Use api_endpoint instead."""
+
+_API_VERSION_OVERRIDE_ENV_VAR = "API_VERSION_OVERRIDE"
+"""This is an experimental configuration variable used for internal testing."""
+
 _DEFAULT_STORAGE_HOST = os.getenv(
-    "API_ENDPOINT_OVERRIDE", "https://storage.googleapis.com"
+    _API_ENDPOINT_OVERRIDE_ENV_VAR, "https://storage.googleapis.com"
 )
 """Default storage host for JSON API."""
 
-_API_VERSION = os.getenv("API_VERSION_OVERRIDE", "v1")
+_API_VERSION = os.getenv(_API_VERSION_OVERRIDE_ENV_VAR, "v1")
 """API version of the default storage host"""
 
-_BASE_STORAGE_URI = "storage.googleapis.com"
-"""Base request endpoint URI for JSON API."""
-
 # etag match parameters in snake case and equivalent header
 _ETAG_MATCH_PARAMETERS = (
     ("if_etag_match", "If-Match"),

diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py
@@ -34,7 +34,6 @@
 from google.cloud.storage._helpers import _get_default_headers
 from google.cloud.storage._helpers import _get_environ_project
 from google.cloud.storage._helpers import _get_storage_host
-from google.cloud.storage._helpers import _BASE_STORAGE_URI
 from google.cloud.storage._helpers import _DEFAULT_STORAGE_HOST
 from google.cloud.storage._helpers import _bucket_bound_hostname_url
 from google.cloud.storage._helpers import _add_etag_match_headers
@@ -96,6 +95,12 @@ class Client(ClientWithProject):
     :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.
         API Endpoint should be set through client_options.
+
+    :type use_auth_w_custom_endpoint: bool
+    :param use_auth_w_custom_endpoint:
+        (Optional) Whether authentication is required under custom endpoints.
+        If false, uses AnonymousCredentials and bypasses authentication.
+        Defaults to True. Note this is only used when a custom endpoint is set in conjunction.
     """
 
     SCOPE = (
@@ -112,6 +117,7 @@ def __init__(
         _http=None,
         client_info=None,
         client_options=None,
+        use_auth_w_custom_endpoint=True,
     ):
         self._base_connection = None
 
@@ -127,13 +133,12 @@ def __init__(
         kw_args = {"client_info": client_info}
 
         # `api_endpoint` should be only set by the user via `client_options`,
-        # or if the _get_storage_host() returns a non-default value.
+        # or if the _get_storage_host() returns a non-default value (_is_emulator_set).
         # `api_endpoint` plays an important role for mTLS, if it is not set,
         # then mTLS logic will be applied to decide which endpoint will be used.
         storage_host = _get_storage_host()
-        kw_args["api_endpoint"] = (
-            storage_host if storage_host != _DEFAULT_STORAGE_HOST else None
-        )
+        _is_emulator_set = storage_host != _DEFAULT_STORAGE_HOST
+        kw_args["api_endpoint"] = storage_host if _is_emulator_set else None
 
         if client_options:
             if type(client_options) == dict:
@@ -144,19 +149,20 @@ def __init__(
                 api_endpoint = client_options.api_endpoint
                 kw_args["api_endpoint"] = api_endpoint
 
-        # Use anonymous credentials and no project when
-        # STORAGE_EMULATOR_HOST or a non-default api_endpoint is set.
-        if (
-            kw_args["api_endpoint"] is not None
-            and _BASE_STORAGE_URI not in kw_args["api_endpoint"]
-        ):
-            if credentials is None:
-                credentials = AnonymousCredentials()
-            if project is None:
-                project = _get_environ_project()
-            if project is None:
-                no_project = True
-                project = "<none>"
+        # If a custom endpoint is set, the client checks for credentials
+        # or finds the default credentials based on the current environment.
+        # Authentication may be bypassed under certain conditions:
+        # (1) STORAGE_EMULATOR_HOST is set (for backwards compatibility), OR
+        # (2) use_auth_w_custom_endpoint is set to False.
+        if kw_args["api_endpoint"] is not None:
+            if _is_emulator_set or not use_auth_w_custom_endpoint:
+                if credentials is None:
+                    credentials = AnonymousCredentials()
+                if project is None:
+                    project = _get_environ_project()
+                if project is None:
+                    no_project = True
+                    project = "<none>"
 
         super(Client, self).__init__(
             project=project,
@@ -897,7 +903,8 @@ def create_bucket(
             project = self.project
 
         # Use no project if STORAGE_EMULATOR_HOST is set
-        if _BASE_STORAGE_URI not in _get_storage_host():
+        _is_emulator_set = _get_storage_host() != _DEFAULT_STORAGE_HOST
+        if _is_emulator_set:
             if project is None:
                 project = _get_environ_project()
             if project is None:
@@ -1327,7 +1334,8 @@ def list_buckets(
             project = self.project
 
         # Use no project if STORAGE_EMULATOR_HOST is set
-        if _BASE_STORAGE_URI not in _get_storage_host():
+        _is_emulator_set = _get_storage_host() != _DEFAULT_STORAGE_HOST
+        if _is_emulator_set:
             if project is None:
                 project = _get_environ_project()
             if project is None:

diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -28,9 +28,10 @@
 from google.auth.credentials import AnonymousCredentials
 from google.oauth2.service_account import Credentials
 
+from google.cloud.storage import _helpers
 from google.cloud.storage._helpers import STORAGE_EMULATOR_ENV_VAR
 from google.cloud.storage._helpers import _get_default_headers
-from google.cloud.storage import _helpers
+from google.cloud.storage._http import Connection
 from google.cloud.storage.retry import DEFAULT_RETRY
 from google.cloud.storage.retry import DEFAULT_RETRY_IF_GENERATION_SPECIFIED
 from tests.unit.test__helpers import GCCL_INVOCATION_TEST_CONST
@@ -119,7 +120,6 @@ def _make_one(self, *args, **kw):
 
     def test_ctor_connection_type(self):
         from google.cloud._http import ClientInfo
-        from google.cloud.storage._http import Connection
 
         PROJECT = "PROJECT"
         credentials = _make_credentials()
@@ -179,8 +179,6 @@ def test_ctor_w_client_options_object(self):
         )
 
     def test_ctor_wo_project(self):
-        from google.cloud.storage._http import Connection
-
         PROJECT = "PROJECT"
         credentials = _make_credentials(project=PROJECT)
 
@@ -193,8 +191,6 @@ def test_ctor_wo_project(self):
         self.assertEqual(list(client._batch_stack), [])
 
     def test_ctor_w_project_explicit_none(self):
-        from google.cloud.storage._http import Connection
-
         credentials = _make_credentials()
 
         client = self._make_one(project=None, credentials=credentials)
@@ -207,7 +203,6 @@ def test_ctor_w_project_explicit_none(self):
 
     def test_ctor_w_client_info(self):
         from google.cloud._http import ClientInfo
-        from google.cloud.storage._http import Connection
 
         credentials = _make_credentials()
         client_info = ClientInfo()
@@ -239,8 +234,40 @@ def test_ctor_mtls(self):
         self.assertEqual(client._connection.ALLOW_AUTO_SWITCH_TO_MTLS_URL, False)
         self.assertEqual(client._connection.API_BASE_URL, "http://foo")
 
+    def test_ctor_w_custom_endpoint_use_auth(self):
+        custom_endpoint = "storage-example.p.googleapis.com"
+        client = self._make_one(client_options={"api_endpoint": custom_endpoint})
+        self.assertEqual(client._connection.API_BASE_URL, custom_endpoint)
+        self.assertIsNotNone(client.project)
+        self.assertIsInstance(client._connection, Connection)
+        self.assertIsNotNone(client._connection.credentials)
+        self.assertNotIsInstance(client._connection.credentials, AnonymousCredentials)
+
+    def test_ctor_w_custom_endpoint_bypass_auth(self):
+        custom_endpoint = "storage-example.p.googleapis.com"
+        client = self._make_one(
+            client_options={"api_endpoint": custom_endpoint},
+            use_auth_w_custom_endpoint=False,
+        )
+        self.assertEqual(client._connection.API_BASE_URL, custom_endpoint)
+        self.assertEqual(client.project, None)
+        self.assertIsInstance(client._connection, Connection)
+        self.assertIsInstance(client._connection.credentials, AnonymousCredentials)
+
+    def test_ctor_w_custom_endpoint_w_credentials(self):
+        PROJECT = "PROJECT"
+        custom_endpoint = "storage-example.p.googleapis.com"
+        credentials = _make_credentials(project=PROJECT)
+        client = self._make_one(
+            credentials=credentials, client_options={"api_endpoint": custom_endpoint}
+        )
+        self.assertEqual(client._connection.API_BASE_URL, custom_endpoint)
+        self.assertEqual(client.project, PROJECT)
+        self.assertIsInstance(client._connection, Connection)
+        self.assertIs(client._connection.credentials, credentials)
+
     def test_ctor_w_emulator_wo_project(self):
-        # avoids authentication if STORAGE_EMULATOR_ENV_VAR is set
+        # bypasses authentication if STORAGE_EMULATOR_ENV_VAR is set
         host = "http://localhost:8080"
         environ = {STORAGE_EMULATOR_ENV_VAR: host}
         with mock.patch("os.environ", environ):
@@ -250,16 +277,8 @@ def test_ctor_w_emulator_wo_project(self):
         self.assertEqual(client._connection.API_BASE_URL, host)
         self.assertIsInstance(client._connection.credentials, AnonymousCredentials)
 
-        # avoids authentication if storage emulator is set through api_endpoint
-        client = self._make_one(
-            client_options={"api_endpoint": "http://localhost:8080"}
-        )
-        self.assertIsNone(client.project)
-        self.assertEqual(client._connection.API_BASE_URL, host)
-        self.assertIsInstance(client._connection.credentials, AnonymousCredentials)
-
     def test_ctor_w_emulator_w_environ_project(self):
-        # avoids authentication and infers the project from the environment
+        # bypasses authentication and infers the project from the environment
         host = "http://localhost:8080"
         environ_project = "environ-project"
         environ = {
@@ -289,9 +308,17 @@ def test_ctor_w_emulator_w_project_arg(self):
         self.assertEqual(client._connection.API_BASE_URL, host)
         self.assertIsInstance(client._connection.credentials, AnonymousCredentials)
 
-    def test_create_anonymous_client(self):
-        from google.cloud.storage._http import Connection
+    def test_ctor_w_emulator_w_credentials(self):
+        host = "http://localhost:8080"
+        environ = {STORAGE_EMULATOR_ENV_VAR: host}
+        credentials = _make_credentials()
+        with mock.patch("os.environ", environ):
+            client = self._make_one(credentials=credentials)
 
+        self.assertEqual(client._connection.API_BASE_URL, host)
+        self.assertIs(client._connection.credentials, credentials)
+
+    def test_create_anonymous_client(self):
         klass = self._get_target_class()
         client = klass.create_anonymous_client()
 
@@ -1269,6 +1296,28 @@ def test_create_bucket_w_environ_project_w_emulator(self):
             _target_object=bucket,
         )
 
+    def test_create_bucket_w_custom_endpoint(self):
+        custom_endpoint = "storage-example.p.googleapis.com"
+        client = self._make_one(client_options={"api_endpoint": custom_endpoint})
+        bucket_name = "bucket-name"
+        api_response = {"name": bucket_name}
+        client._post_resource = mock.Mock()
+        client._post_resource.return_value = api_response
+
+        bucket = client.create_bucket(bucket_name)
+
+        expected_path = "/b"
+        expected_data = api_response
+        expected_query_params = {"project": client.project}
+        client._post_resource.assert_called_once_with(
+            expected_path,
+            expected_data,
+            query_params=expected_query_params,
+            timeout=self._get_default_timeout(),
+            retry=DEFAULT_RETRY,
+            _target_object=bucket,
+        )
+
     def test_create_bucket_w_conflict_w_user_project(self):
         from google.cloud.exceptions import Conflict
 
@@ -2055,6 +2104,37 @@ def test_list_buckets_w_environ_project_w_emulator(self):
             retry=DEFAULT_RETRY,
         )
 
+    def test_list_buckets_w_custom_endpoint(self):
+        from google.cloud.storage.client import _item_to_bucket
+
+        custom_endpoint = "storage-example.p.googleapis.com"
+        client = self._make_one(client_options={"api_endpoint": custom_endpoint})
+        client._list_resource = mock.Mock(spec=[])
+
+        iterator = client.list_buckets()
+
+        self.assertIs(iterator, client._list_resource.return_value)
+
+        expected_path = "/b"
+        expected_item_to_value = _item_to_bucket
+        expected_page_token = None
+        expected_max_results = None
+        expected_page_size = None
+        expected_extra_params = {
+            "project": client.project,
+            "projection": "noAcl",
+        }
+        client._list_resource.assert_called_once_with(
+            expected_path,
+            expected_item_to_value,
+            page_token=expected_page_token,
+            max_results=expected_max_results,
+            extra_params=expected_extra_params,
+            page_size=expected_page_size,
+            timeout=self._get_default_timeout(),
+            retry=DEFAULT_RETRY,
+        )
+
     def test_list_buckets_w_defaults(self):
         from google.cloud.storage.client import _item_to_bucket