chore: refactor setting of auto-populated field in request (#17004)

This PR addresses the feedback in
#16944 (comment)

Author: parthea

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2

@@ -14,7 +14,7 @@
  # limitations under the License.
 #}
 
-{% macro auto_populate_uuid4_fields(api, method) %}
+{% macro auto_populate_uuid4_fields(api, method, is_async=False) %}
 {#
     Automatically populate UUID4 fields according to
     https://google.aip.dev/client-libraries/4235 when the
@@ -27,12 +27,12 @@
 {% with method_settings = api.all_method_settings.get(method.meta.address.proto) %}
 {% if method_settings is not none %}
 {% for auto_populated_field in method_settings.auto_populated_fields %}
-    {% if method.input.fields[auto_populated_field].proto3_optional %}
-        if '{{ auto_populated_field }}' not in request:
+    {% set is_proto3_optional = method.input.fields[auto_populated_field].proto3_optional %}
+    {% if is_async %}
+        self._client._setup_request_id(request, '{{ auto_populated_field }}', {{ is_proto3_optional }})
     {% else %}
-        if not request.{{ auto_populated_field }}:
+        self._setup_request_id(request, '{{ auto_populated_field }}', {{ is_proto3_optional }})
     {% endif %}
-            request.{{ auto_populated_field }} = str(uuid.uuid4())
 {% endfor %}
 {% endif %}{# if method_settings is not none #}
 {% endwith %}{# method_settings #}

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2

@@ -395,7 +395,7 @@ class {{ service.async_client_name }}:
 
 {{ shared_macros.create_metadata(method) }}
 {{ shared_macros.add_api_version_header_to_metadata(service.version) }}
-{{ shared_macros.auto_populate_uuid4_fields(api, method) }}
+{{ shared_macros.auto_populate_uuid4_fields(api, method, is_async=True) }}
 
         # Validate the universe domain.
         self._client._validate_universe_domain()

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2

@@ -453,6 +453,38 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
         # NOTE (b/349488459): universe validation is disabled until further notice.
         return True
 
+    {% if api.all_method_settings.values()|map(attribute="auto_populated_fields", default=[])|select|list %}
+    @staticmethod
+    def _setup_request_id(request, field_name: str, is_proto3_optional: bool):
+        """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()))
+    {% endif %}
+
     def _add_cred_info_for_auth_errors(
         self,
         error: core_exceptions.GoogleAPICallError

packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -415,6 +415,84 @@ def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code):
     client._add_cred_info_for_auth_errors(error)
     assert error.details == []
 
+{% if api.all_method_settings.values()|map(attribute="auto_populated_fields", default=[])|select|list %}
+def test__setup_request_id():
+    class MockRequest:
+        def __init__(self, **kwargs):
+            for k, v in kwargs.items():
+                setattr(self, k, v)
+        def __contains__(self, key):
+            return hasattr(self, key)
+
+    class MockProtoRequest:
+        def __init__(self, **kwargs):
+            for k, v in kwargs.items():
+                setattr(self, k, v)
+        def HasField(self, key):
+            return hasattr(self, key)
+
+    # Test with proto3 optional field not in request
+    request = MockRequest()
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request.request_id)
+
+    # Test with proto3 optional field already in request
+    request = MockRequest(request_id="already_set")
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert request.request_id == "already_set"
+
+    # Test with non-proto3 optional field empty
+    request = MockRequest(request_id="")
+    {{ service.client_name }}._setup_request_id(request, "request_id", False)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request.request_id)
+
+    # Test with non-proto3 optional field already set
+    request = MockRequest(request_id="already_set")
+    {{ service.client_name }}._setup_request_id(request, "request_id", False)
+    assert request.request_id == "already_set"
+
+    # Test with proto3 optional field not in request (MockProtoRequest)
+    request = MockProtoRequest()
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request.request_id)
+
+    # Test with proto3 optional field already in request (MockProtoRequest)
+    request = MockProtoRequest(request_id="already_set")
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert request.request_id == "already_set"
+
+    # Test with ValueError
+    class MockValueErrorRequest:
+        def HasField(self, key):
+            raise ValueError("Mismatched field")
+        def __contains__(self, key):
+            return hasattr(self, key)
+
+    request = MockValueErrorRequest()
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request.request_id)
+
+    # Test with dict and proto3 optional field not in request
+    request = {}
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request["request_id"])
+
+    # Test with dict and proto3 optional field already in request
+    request = {"request_id": "already_set"}
+    {{ service.client_name }}._setup_request_id(request, "request_id", True)
+    assert request["request_id"] == "already_set"
+
+    # Test with dict and non-proto3 optional field empty
+    request = {"request_id": ""}
+    {{ service.client_name }}._setup_request_id(request, "request_id", False)
+    assert re.match(r"{{ test_macros.get_uuid4_re() }}", request["request_id"])
+
+    # Test with dict and non-proto3 optional field already set
+    request = {"request_id": "already_set"}
+    {{ service.client_name }}._setup_request_id(request, "request_id", False)
+    assert request["request_id"] == "already_set"
+
+{% endif %}
 @pytest.mark.parametrize("client_class,transport_name", [
     {% if 'grpc' in opts.transport %}
     ({{ service.client_name }}, "grpc"),

packages/gapic-generator/tests/integration/goldens/storagebatchoperations/google/cloud/storagebatchoperations_v1/services/storage_batch_operations/async_client.py

@@ -621,8 +621,7 @@ async def sample_create_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._client._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._client._validate_universe_domain()
@@ -728,8 +727,7 @@ async def sample_delete_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._client._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._client._validate_universe_domain()
@@ -831,8 +829,7 @@ async def sample_cancel_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._client._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._client._validate_universe_domain()

packages/gapic-generator/tests/integration/goldens/storagebatchoperations/google/cloud/storagebatchoperations_v1/services/storage_batch_operations/client.py

@@ -471,6 +471,36 @@ def _validate_universe_domain(self):
         # NOTE (b/349488459): universe validation is disabled until further notice.
         return True
 
+    @staticmethod
+    def _setup_request_id(request, field_name: str, is_proto3_optional: bool):
+        """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()))
+
     def _add_cred_info_for_auth_errors(
         self,
         error: core_exceptions.GoogleAPICallError
@@ -1005,8 +1035,7 @@ def sample_create_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._validate_universe_domain()
@@ -1111,8 +1140,7 @@ def sample_delete_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._validate_universe_domain()
@@ -1213,8 +1241,7 @@ def sample_cancel_job():
             )),
         )
 
-        if not request.request_id:
-            request.request_id = str(uuid.uuid4())
+        self._setup_request_id(request, 'request_id', False)
 
         # Validate the universe domain.
         self._validate_universe_domain()

packages/gapic-generator/tests/integration/goldens/storagebatchoperations/tests/unit/gapic/storagebatchoperations_v1/test_storage_batch_operations.py

@@ -334,6 +334,82 @@ def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code):
     client._add_cred_info_for_auth_errors(error)
     assert error.details == []
 
+def test__setup_request_id():
+    class MockRequest:
+        def __init__(self, **kwargs):
+            for k, v in kwargs.items():
+                setattr(self, k, v)
+        def __contains__(self, key):
+            return hasattr(self, key)
+
+    class MockProtoRequest:
+        def __init__(self, **kwargs):
+            for k, v in kwargs.items():
+                setattr(self, k, v)
+        def HasField(self, key):
+            return hasattr(self, key)
+
+    # Test with proto3 optional field not in request
+    request = MockRequest()
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request.request_id)
+
+    # Test with proto3 optional field already in request
+    request = MockRequest(request_id="already_set")
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert request.request_id == "already_set"
+
+    # Test with non-proto3 optional field empty
+    request = MockRequest(request_id="")
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", False)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request.request_id)
+
+    # Test with non-proto3 optional field already set
+    request = MockRequest(request_id="already_set")
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", False)
+    assert request.request_id == "already_set"
+
+    # Test with proto3 optional field not in request (MockProtoRequest)
+    request = MockProtoRequest()
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request.request_id)
+
+    # Test with proto3 optional field already in request (MockProtoRequest)
+    request = MockProtoRequest(request_id="already_set")
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert request.request_id == "already_set"
+
+    # Test with ValueError
+    class MockValueErrorRequest:
+        def HasField(self, key):
+            raise ValueError("Mismatched field")
+        def __contains__(self, key):
+            return hasattr(self, key)
+
+    request = MockValueErrorRequest()
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request.request_id)
+
+    # Test with dict and proto3 optional field not in request
+    request = {}
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request["request_id"])
+
+    # Test with dict and proto3 optional field already in request
+    request = {"request_id": "already_set"}
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", True)
+    assert request["request_id"] == "already_set"
+
+    # Test with dict and non-proto3 optional field empty
+    request = {"request_id": ""}
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", False)
+    assert re.match(r"[a-f0-9]{8}-?[a-f0-9]{4}-?4[a-f0-9]{3}-?[89ab][a-f0-9]{3}-?[a-f0-9]{12}", request["request_id"])
+
+    # Test with dict and non-proto3 optional field already set
+    request = {"request_id": "already_set"}
+    StorageBatchOperationsClient._setup_request_id(request, "request_id", False)
+    assert request["request_id"] == "already_set"
+
 @pytest.mark.parametrize("client_class,transport_name", [
     (StorageBatchOperationsClient, "grpc"),
     (StorageBatchOperationsAsyncClient, "grpc_asyncio"),