switch to xds-proto

switch to xds-proto

Author: tarun-google

sdks/python/apache_beam/io/components/rate_limiter.py

@@ -28,11 +28,9 @@
 from typing import List
 
 import grpc
-from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptor
-from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptorEntry
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitRequest
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode
+from envoy.extensions.common.ratelimit.v3 import ratelimit_pb2
+from envoy.service.ratelimit.v3 import rls_pb2
+from envoy.service.ratelimit.v3 import rls_pb2_grpc
 
 from apache_beam.io.components import adaptive_throttler
 from apache_beam.metrics import Metrics
@@ -114,30 +112,14 @@ def __init__(
     self._stub = None
     self._lock = threading.Lock()
 
-  class RateLimitServiceStub(object):
-    """ 
-    Wrapper for gRPC stub to be compatible with envoy_data_plane messages.
-    
-    The envoy-data-plane package uses 'betterproto' which generates async stubs
-    for 'grpclib'. As Beam uses standard synchronous 'grpcio', RateLimitServiceStub
-    is a bridge class to use the betterproto Message types (RateLimitRequest) 
-    with a standard grpcio Channel.
-    """
-    def __init__(self, channel):
-      self.ShouldRateLimit = channel.unary_unary(
-          '/envoy.service.ratelimit.v3.RateLimitService/ShouldRateLimit',
-          request_serializer=RateLimitRequest.SerializeToString,
-          response_deserializer=RateLimitResponse.FromString,
-      )
-
   def init_connection(self):
     if self._stub is None:
       # Acquire lock to safegaurd againest multiple DoFn threads sharing the same
       # RateLimiter instance, which is the case when using Shared().
       with self._lock:
         if self._stub is None:
           channel = grpc.insecure_channel(self.service_address)
-          self._stub = EnvoyRateLimiter.RateLimitServiceStub(channel)
+          self._stub = rls_pb2_grpc.RateLimitServiceStub(channel)
 
   def throttle(self, hits_added: int = 1) -> bool:
     """Calls the Envoy RLS to check for rate limits.
@@ -156,10 +138,10 @@ def throttle(self, hits_added: int = 1) -> bool:
     for d in self.descriptors:
       entries = []
       for k, v in d.items():
-        entries.append(RateLimitDescriptorEntry(key=k, value=v))
-      proto_descriptors.append(RateLimitDescriptor(entries=entries))
+        entries.append(ratelimit_pb2.RateLimitDescriptor.Entry(key=k, value=v))
+      proto_descriptors.append(ratelimit_pb2.RateLimitDescriptor(entries=entries))
 
-    request = RateLimitRequest(
+    request = rls_pb2.RateLimitRequest(
         domain=self.domain,
         descriptors=proto_descriptors,
         hits_addend=hits_added)
@@ -188,22 +170,21 @@ def throttle(self, hits_added: int = 1) -> bool:
               "[EnvoyRateLimiter] Connection Failed, retrying: %s", e)
           time.sleep(_RETRY_DELAY_SECONDS)
 
-      if response.overall_code == RateLimitResponseCode.OK:
+      if response.overall_code == rls_pb2.RateLimitResponse.OK:
         self.requests_allowed.inc()
         throttled = True
         break
-      elif response.overall_code == RateLimitResponseCode.OVER_LIMIT:
+      elif response.overall_code == rls_pb2.RateLimitResponse.OVER_LIMIT:
         self.requests_throttled.inc()
         # Ratelimit exceeded, sleep for duration until reset and retry
         # multiple rules can be set in the RLS config, so we need to find the max duration
         sleep_s = 0.0
         if response.statuses:
           for status in response.statuses:
-            if status.code == RateLimitResponseCode.OVER_LIMIT:
+            if status.code == rls_pb2.RateLimitResponse.OVER_LIMIT:
               dur = status.duration_until_reset
-              # duration_until_reset is converted to timedelta by betterproto
-              # duration_until_reset has microsecond precision
-              val = dur.total_seconds()
+              # duration_until_reset is google.protobuf.Duration
+              val = dur.seconds + dur.nanos / 1e9
               if val > sleep_s:
                 sleep_s = val
 

sdks/python/apache_beam/io/components/rate_limiter_test.py

@@ -21,9 +21,8 @@
 from datetime import timedelta
 
 from apache_beam.io.components import rate_limiter
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode
-from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseDescriptorStatus
+from envoy.service.ratelimit.v3 import rls_pb2
+from google.protobuf.duration_pb2 import Duration
 
 
 class EnvoyRateLimiterTest(unittest.TestCase):
@@ -44,7 +43,8 @@ def setUp(self):
   def test_throttle_allowed(self, mock_channel):
     # Mock successful OK response
     mock_stub = mock.Mock()
-    mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK)
+    mock_response = rls_pb2.RateLimitResponse(
+        overall_code=rls_pb2.RateLimitResponse.OK)
     mock_stub.ShouldRateLimit.return_value = mock_response
 
     # Inject mock stub
@@ -59,8 +59,8 @@ def test_throttle_allowed(self, mock_channel):
   def test_throttle_over_limit_retries_exceeded(self, mock_channel):
     # Mock OVER_LIMIT response
     mock_stub = mock.Mock()
-    mock_response = RateLimitResponse(
-        overall_code=RateLimitResponseCode.OVER_LIMIT)
+    mock_response = rls_pb2.RateLimitResponse(
+        overall_code=rls_pb2.RateLimitResponse.OVER_LIMIT)
     mock_stub.ShouldRateLimit.return_value = mock_response
 
     self.limiter._stub = mock_stub
@@ -84,7 +84,8 @@ def test_throttle_over_limit_retries_exceeded(self, mock_channel):
   def test_throttle_rpc_error_retry(self, mock_channel):
     # Mock RpcError then Success
     mock_stub = mock.Mock()
-    mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK)
+    mock_response = rls_pb2.RateLimitResponse(
+        overall_code=rls_pb2.RateLimitResponse.OK)
 
     # Side effect: Error, Error, Success
     error = grpc.RpcError()
@@ -121,11 +122,12 @@ def test_extract_duration_from_response(self, mock_random, mock_channel):
     mock_stub = mock.Mock()
 
     # Valid until 5 seconds
-    status = RateLimitResponseDescriptorStatus(
-        code=RateLimitResponseCode.OVER_LIMIT,
-        duration_until_reset=timedelta(seconds=5))
-    mock_response = RateLimitResponse(
-        overall_code=RateLimitResponseCode.OVER_LIMIT, statuses=[status])
+    # Valid until 5 seconds
+    status = rls_pb2.RateLimitResponse.DescriptorStatus(
+        code=rls_pb2.RateLimitResponse.OVER_LIMIT,
+        duration_until_reset=Duration(seconds=5))
+    mock_response = rls_pb2.RateLimitResponse(
+        overall_code=rls_pb2.RateLimitResponse.OVER_LIMIT, statuses=[status])
 
     mock_stub.ShouldRateLimit.return_value = mock_response
     self.limiter._stub = mock_stub

sdks/python/setup.py

@@ -376,8 +376,7 @@ def get_portability_package_data():
           'cryptography>=39.0.0,<48.0.0',
           'fastavro>=0.23.6,<2',
           'fasteners>=0.3,<1.0',
-          'envoy-data-plane>=1.0.3,<2; python_version >= "3.11"',
-          'envoy-data-plane==0.8.1; python_version < "3.11"',
+          'xds-protos<=1.75.0',
           # TODO(https://github.com/grpc/grpc/issues/37710): Unpin grpc
           'grpcio>=1.33.1,<2,!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,<1.66.0; python_version <= "3.12"',  # pylint: disable=line-too-long
           'grpcio>=1.67.0; python_version >= "3.13"',