feat: add caching to routing header calculation (#526)

daniel-sanche · parthea · web-flow · commit 6251eab3fca5 · 2023-10-09T13:00:15.000-07:00
Co-authored-by: Anthonios Partheniou &lt;partheniou@google.com&gt;
diff --git a/google/api_core/gapic_v1/routing_header.py b/google/api_core/gapic_v1/routing_header.py
@@ -20,45 +20,42 @@
 Generally, these headers are specified as gRPC metadata.
 """
 
+import functools
 from enum import Enum
 from urllib.parse import urlencode
 
 ROUTING_METADATA_KEY = "x-goog-request-params"
+# This is the value for the `maxsize` argument of @functools.lru_cache
+# https://docs.python.org/3/library/functools.html#functools.lru_cache
+# This represents the number of recent function calls to store.
+ROUTING_PARAM_CACHE_SIZE = 32
 
 
 def to_routing_header(params, qualified_enums=True):
     """Returns a routing header string for the given request parameters.
 
     Args:
-        params (Mapping[str, Any]): A dictionary containing the request
+        params (Mapping[str, str | bytes | Enum]): A dictionary containing the request
             parameters used for routing.
         qualified_enums (bool): Whether to represent enum values
             as their type-qualified symbol names instead of as their
             unqualified symbol names.
 
     Returns:
         str: The routing header string.
-
     """
+    tuples = params.items() if isinstance(params, dict) else params
     if not qualified_enums:
-        if isinstance(params, dict):
-            tuples = params.items()
-        else:
-            tuples = params
-        params = [(x[0], x[1].name) if isinstance(x[1], Enum) else x for x in tuples]
-    return urlencode(
-        params,
-        # Per Google API policy (go/api-url-encoding), / is not encoded.
-        safe="/",
-    )
+        tuples = [(x[0], x[1].name) if isinstance(x[1], Enum) else x for x in tuples]
+    return "&".join([_urlencode_param(*t) for t in tuples])
 
 
 def to_grpc_metadata(params, qualified_enums=True):
     """Returns the gRPC metadata containing the routing headers for the given
     request parameters.
 
     Args:
-        params (Mapping[str, Any]): A dictionary containing the request
+        params (Mapping[str, str | bytes | Enum]): A dictionary containing the request
             parameters used for routing.
         qualified_enums (bool): Whether to represent enum values
             as their type-qualified symbol names instead of as their
@@ -69,3 +66,22 @@ def to_grpc_metadata(params, qualified_enums=True):
             and value.
     """
     return (ROUTING_METADATA_KEY, to_routing_header(params, qualified_enums))
+
+
+# use caching to avoid repeated computation
+@functools.lru_cache(maxsize=ROUTING_PARAM_CACHE_SIZE)
+def _urlencode_param(key, value):
+    """Cacheable wrapper over urlencode
+
+    Args:
+        key (str): The key of the parameter to encode.
+        value (str | bytes | Enum): The value of the parameter to encode.
+
+    Returns:
+        str: The encoded parameter.
+    """
+    return urlencode(
+        {key: value},
+        # Per Google API policy (go/api-url-encoding), / is not encoded.
+        safe="/",
+    )
diff --git a/tests/unit/gapic/test_routing_header.py b/tests/unit/gapic/test_routing_header.py
@@ -70,3 +70,34 @@ def test_to_grpc_metadata():
     params = [("name", "meep"), ("book.read", "1")]
     metadata = routing_header.to_grpc_metadata(params)
     assert metadata == (routing_header.ROUTING_METADATA_KEY, "name=meep&book.read=1")
+
+
+@pytest.mark.parametrize(
+    "key,value,expected",
+    [
+        ("book.read", "1", "book.read=1"),
+        ("name", "me/ep", "name=me/ep"),
+        ("\\", "=", "%5C=%3D"),
+        (b"hello", "world", "hello=world"),
+        ("✔️", "✌️", "%E2%9C%94%EF%B8%8F=%E2%9C%8C%EF%B8%8F"),
+    ],
+)
+def test__urlencode_param(key, value, expected):
+    result = routing_header._urlencode_param(key, value)
+    assert result == expected
+
+
+def test__urlencode_param_caching_performance():
+    import time
+
+    key = "key" * 100
+    value = "value" * 100
+    # time with empty cache
+    start_time = time.perf_counter()
+    routing_header._urlencode_param(key, value)
+    duration = time.perf_counter() - start_time
+    second_start_time = time.perf_counter()
+    routing_header._urlencode_param(key, value)
+    second_duration = time.perf_counter() - second_start_time
+    # second call should be approximately 10 times faster
+    assert second_duration < duration / 10
