feat: centralize additional client helpers to api_core

Author: hebaalazzeh

packages/google-api-core/google/api_core/gapic_v1/client_helpers.py

@@ -18,13 +18,15 @@
 
 import os
 import re
-from typing import Any, Optional
+import uuid
+from typing import Any, Optional, Tuple
 
 from google.auth.exceptions import MutualTLSChannelError  # type: ignore
 from google.auth.transport import mtls  # type: ignore
 
 _MTLS_ENDPOINT_RE = re.compile(
-    r"(?P<name>[^.]+)(?P<mtls>\.mtls)?(?P<sandbox>\.sandbox)?(?P<googledomain>\.googleapis\.com)?"
+    r"(?P<name>[^.]+)(?P<mtls>\.mtls)?(?P<sandbox>\.sandbox)?"
+    r"(?P<googledomain>\.googleapis\.com)?"
 )
 
 
@@ -60,17 +62,20 @@ def _get_default_mtls_endpoint(api_endpoint: Optional[str]) -> Optional[str]:
 
 def _use_client_cert_effective() -> bool:
     """Returns whether client certificate should be used for mTLS if the
-    google-auth version supports should_use_client_cert automatic mTLS enablement.
+    google-auth version supports should_use_client_cert automatic mTLS
+    enablement.
 
     Alternatively, read from the GOOGLE_API_USE_CLIENT_CERTIFICATE env var.
 
     Returns:
         bool: whether client certificate should be used for mTLS
     Raises:
-        ValueError: (If using a version of google-auth without should_use_client_cert and
-        GOOGLE_API_USE_CLIENT_CERTIFICATE is set to an unexpected value.)
+        ValueError: (If using a version of google-auth without
+        should_use_client_cert and GOOGLE_API_USE_CLIENT_CERTIFICATE is
+        set to an unexpected value.)
     """
-    # check if google-auth version supports should_use_client_cert for automatic mTLS enablement
+    # check if google-auth version supports should_use_client_cert for
+    # automatic mTLS enablement
     if hasattr(mtls, "should_use_client_cert"):  # pragma: NO COVER
         return mtls.should_use_client_cert()
     else:  # pragma: NO COVER
@@ -80,8 +85,8 @@ def _use_client_cert_effective() -> bool:
         ).lower()
         if use_client_cert_str not in ("true", "false"):
             raise ValueError(
-                "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be"
-                " either `true` or `false`"
+                "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` "
+                "must be either `true` or `false`"
             )
         return use_client_cert_str == "true"
 
@@ -116,8 +121,127 @@ def _get_api_endpoint(
     ):
         if universe_domain != default_universe:
             raise MutualTLSChannelError(
-                f"mTLS is not supported in any universe other than {default_universe}."
+                f"mTLS is not supported in any universe other than "
+                f"{default_universe}."
             )
         return default_mtls_endpoint
     else:
-        return default_endpoint_template.format(UNIVERSE_DOMAIN=universe_domain)
+        return default_endpoint_template.format(
+            UNIVERSE_DOMAIN=universe_domain
+        )
+
+
+def _read_environment_variables() -> Tuple[bool, str, Optional[str]]:
+    """Returns the environment variables used by the client.
+
+    Returns:
+        Tuple[bool, str, Optional[str]]: returns the
+        GOOGLE_API_USE_CLIENT_CERTIFICATE, GOOGLE_API_USE_MTLS_ENDPOINT,
+        and GOOGLE_CLOUD_UNIVERSE_DOMAIN environment variables.
+
+    Raises:
+        ValueError: If GOOGLE_API_USE_CLIENT_CERTIFICATE is not
+            any of ["true", "false"].
+        google.auth.exceptions.MutualTLSChannelError: If
+            GOOGLE_API_USE_MTLS_ENDPOINT is not any of
+            ["auto", "never", "always"].
+    """
+    use_client_cert = _use_client_cert_effective()
+    use_mtls_endpoint = os.getenv(
+        "GOOGLE_API_USE_MTLS_ENDPOINT", "auto"
+    ).lower()
+    universe_domain_env = os.getenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN")
+    if use_mtls_endpoint not in ("auto", "never", "always"):
+        raise MutualTLSChannelError(
+            "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` "
+            "must be `never`, `auto` or `always`"
+        )
+    return use_client_cert, use_mtls_endpoint, universe_domain_env
+
+
+def _get_client_cert_source(
+    provided_cert_source: Optional[Any], use_cert_flag: bool
+) -> Optional[Any]:
+    """Return the client cert source to be used by the client.
+
+    Args:
+        provided_cert_source (bytes): The client certificate source provided.
+        use_cert_flag (bool): A flag indicating whether to use the
+            client certificate.
+
+    Returns:
+        bytes or None: The client cert source to be used by the client.
+    """
+    client_cert_source = None
+    if use_cert_flag:
+        if provided_cert_source:
+            client_cert_source = provided_cert_source
+        elif (
+            hasattr(mtls, "has_default_client_cert_source")
+            and mtls.has_default_client_cert_source()
+        ):
+            client_cert_source = mtls.default_client_cert_source()
+    return client_cert_source
+
+
+def _get_universe_domain(
+    client_universe_domain: Optional[str],
+    universe_domain_env: Optional[str],
+    default_universe: str,
+) -> str:
+    """Return the universe domain used by the client.
+
+    Args:
+        client_universe_domain (Optional[str]): The universe domain
+            configured via the client options.
+        universe_domain_env (Optional[str]): The universe domain
+            configured via the "GOOGLE_CLOUD_UNIVERSE_DOMAIN" env var.
+        default_universe (str): The default universe domain.
+
+    Returns:
+        str: The universe domain to be used by the client.
+
+    Raises:
+        ValueError: If the universe domain is an empty string.
+    """
+    universe_domain = default_universe
+    if client_universe_domain is not None:
+        universe_domain = client_universe_domain
+    elif universe_domain_env is not None:
+        universe_domain = universe_domain_env
+    if len(universe_domain.strip()) == 0:
+        raise ValueError("Universe Domain cannot be an empty string.")
+    return universe_domain
+
+
+def _setup_request_id(
+    request: Any, field_name: str, is_proto3_optional: bool
+) -> None:
+    """Populate a UUID4 field in the request if it is not already set.
+
+    Args:
+        request (Union[google.protobuf.message.Message, dict]): The
+            request object.
+        field_name (str): The name of the field to populate.
+        is_proto3_optional (bool): Whether the field is proto3 optional.
+    """
+    if isinstance(request, dict):
+        if is_proto3_optional:
+            if field_name not in request:
+                request[field_name] = str(uuid.uuid4())
+        elif not request.get(field_name):
+            request[field_name] = str(uuid.uuid4())
+        return
+
+    if is_proto3_optional:
+        try:
+            # Pure protobuf messages
+            if not request.HasField(field_name):
+                setattr(request, field_name, str(uuid.uuid4()))
+        except (AttributeError, ValueError):
+            # Proto-plus messages or other objects
+            if field_name not in request:
+                setattr(request, field_name, str(uuid.uuid4()))
+    else:
+        if not getattr(request, field_name):
+            setattr(request, field_name, str(uuid.uuid4()))

packages/google-api-core/tests/unit/gapic/test_client_helpers.py

@@ -149,3 +149,112 @@ def test__get_api_endpoint_mtls_universe_mismatch():
             default_mtls_endpoint="foo.mtls.googleapis.com",
             default_endpoint_template="foo.{UNIVERSE_DOMAIN}",
         )
+
+
+@mock.patch(
+    "google.api_core.gapic_v1.client_helpers._use_client_cert_effective"
+)
+@mock.patch.dict(os.environ, clear=True)
+def test__read_environment_variables(mock_effective):
+    mock_effective.return_value = True
+    os.environ["GOOGLE_API_USE_MTLS_ENDPOINT"] = "always"
+    os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"] = "custom.com"
+
+    cert, mtls, domain = client_helpers._read_environment_variables()
+    assert cert is True
+    assert mtls == "always"
+    assert domain == "custom.com"
+
+
+@mock.patch.dict(os.environ, clear=True)
+def test__read_environment_variables_invalid_mtls():
+    os.environ["GOOGLE_API_USE_MTLS_ENDPOINT"] = "invalid"
+    with pytest.raises(
+        MutualTLSChannelError, match="must be `never`, `auto` or `always`"
+    ):
+        client_helpers._read_environment_variables()
+
+
+@mock.patch("google.auth.transport.mtls.has_default_client_cert_source")
+@mock.patch("google.auth.transport.mtls.default_client_cert_source")
+def test__get_client_cert_source(mock_default, mock_has_default):
+    mock_default.return_value = b"default_cert"
+    mock_has_default.return_value = True
+
+    # When use_cert_flag is False, return None
+    assert client_helpers._get_client_cert_source(b"provided", False) is None
+
+    # When provided_cert_source is given, return provided
+    assert (
+        client_helpers._get_client_cert_source(b"provided", True)
+        == b"provided"
+    )
+
+    # When no provided cert but default is available
+    assert (
+        client_helpers._get_client_cert_source(None, True) == b"default_cert"
+    )
+
+
+def test__get_universe_domain():
+    # client_universe_domain takes precedence
+    assert (
+        client_helpers._get_universe_domain(
+            "client.com", "env.com", "default.com"
+        )
+        == "client.com"
+    )
+
+    # env takes precedence over default
+    assert (
+        client_helpers._get_universe_domain(None, "env.com", "default.com")
+        == "env.com"
+    )
+
+    # fallback to default
+    assert (
+        client_helpers._get_universe_domain(None, None, "default.com")
+        == "default.com"
+    )
+
+
+def test__get_universe_domain_empty():
+    with pytest.raises(ValueError, match="cannot be an empty string"):
+        client_helpers._get_universe_domain("", None, "default.com")
+
+
+def test__setup_request_id():
+    import uuid
+
+    # test dict request
+    req = {}
+    client_helpers._setup_request_id(req, "request_id", True)
+    assert "request_id" in req
+    uuid_str = req["request_id"]
+    uuid.UUID(uuid_str)  # verify it is a valid UUID
+
+    # test dict request when already set
+    req = {"request_id": "existing"}
+    client_helpers._setup_request_id(req, "request_id", True)
+    assert req["request_id"] == "existing"
+
+    class DummyRequest:
+        def __init__(self):
+            self.request_id = ""
+
+        def HasField(self, field_name):
+            if not hasattr(self, field_name):
+                raise ValueError()
+            return bool(getattr(self, field_name))
+
+    # test object request proto3 optional true
+    req_obj = DummyRequest()
+    client_helpers._setup_request_id(req_obj, "request_id", True)
+    assert req_obj.request_id != ""
+    uuid.UUID(req_obj.request_id)
+
+    # test object request proto3 optional false
+    req_obj2 = DummyRequest()
+    client_helpers._setup_request_id(req_obj2, "request_id", False)
+    assert req_obj2.request_id != ""
+    uuid.UUID(req_obj2.request_id)