fix format issues

Author: tarun-google

sdks/python/apache_beam/examples/rate_limiter_simple.py

@@ -29,7 +29,6 @@
 
 class SampleApiDoFn(beam.DoFn):
   """A DoFn that simulates calling an external API with rate limiting."""
-
   def __init__(self, rls_address, domain, descriptors):
     self.rls_address = rls_address
     self.domain = domain
@@ -42,23 +41,24 @@ def setup(self):
     # We use shared.Shared() to ensure only one RateLimiter instance is created
     # per worker and shared across threads.
     def init_limiter():
-        logging.info(f"Connecting to Envoy RLS at {self.rls_address}")
-        return EnvoyRateLimiter(
-            service_address=self.rls_address,
-            domain=self.domain,
-            descriptors=self.descriptors,
-            namespace='example_pipeline'
-        )
+      logging.info(f"Connecting to Envoy RLS at {self.rls_address}")
+      return EnvoyRateLimiter(
+          service_address=self.rls_address,
+          domain=self.domain,
+          descriptors=self.descriptors,
+          namespace='example_pipeline')
+
     self.rate_limiter = self._shared.acquire(init_limiter)
 
   def process(self, element):
     self.rate_limiter.throttle()
 
     # Process the element mock API call
     logging.info(f"Processing element: {element}")
-    time.sleep(0.1) 
+    time.sleep(0.1)
     yield element
 
+
 def parse_known_args(argv):
   """Parses args for the workflow."""
   parser = argparse.ArgumentParser()
@@ -77,11 +77,13 @@ def run(argv=None):
     (
         p
         | 'Create' >> beam.Create(range(100))
-        | 'RateLimit' >> beam.ParDo(SampleApiDoFn(
-            rls_address=known_args.rls_address, 
-            domain="mongo_cps", 
-            descriptors=[{"database": "users"}]))
-    )
+        | 'RateLimit' >> beam.ParDo(
+            SampleApiDoFn(
+                rls_address=known_args.rls_address,
+                domain="mongo_cps",
+                descriptors=[{
+                    "database": "users"
+                }])))
 
 
 if __name__ == '__main__':

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

@@ -20,19 +20,19 @@
 """
 
 import abc
-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_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptor
-from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptorEntry
 import logging
-import time
 import threading
-import random
+import time
 from typing import Dict
 from typing import List
+
 import grpc
-import time
+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 apache_beam.io.components import adaptive_throttler
 from apache_beam.metrics import Metrics
 
@@ -41,19 +41,24 @@
 _MAX_CONNECTION_RETRIES = 5
 _RETRY_DELAY_SECONDS = 10
 
+
 class RateLimiter(abc.ABC):
   """Abstract base class for RateLimiters."""
   def __init__(self, namespace: str = ""):
     # Metrics collected from the RateLimiter
     # Metric updates are thread safe
     self.throttling_signaler = adaptive_throttler.ThrottlingSignaler(
         namespace=namespace)
-    self.requests_counter = Metrics.counter(namespace, 'envoyRatelimitRequestsTotal')
-    self.requests_allowed = Metrics.counter(namespace, 'envoyRatelimitRequestsAllowed')
-    self.requests_throttled = Metrics.counter(namespace, 'envoyRatelimitRequestsThrottled')
+    self.requests_counter = Metrics.counter(
+        namespace, 'envoyRatelimitRequestsTotal')
+    self.requests_allowed = Metrics.counter(
+        namespace, 'envoyRatelimitRequestsAllowed')
+    self.requests_throttled = Metrics.counter(
+        namespace, 'envoyRatelimitRequestsThrottled')
     self.rpc_errors = Metrics.counter(namespace, 'envoyRatelimitRpcErrors')
     self.rpc_retries = Metrics.counter(namespace, 'envoyRatelimitRpcRetries')
-    self.rpc_latency = Metrics.distribution(namespace, 'envoyRatelimitRpcLatencyMs')
+    self.rpc_latency = Metrics.distribution(
+        namespace, 'envoyRatelimitRpcLatencyMs')
 
   @abc.abstractmethod
   def throttle(self, **kwargs) -> bool:
@@ -69,7 +74,6 @@ def throttle(self, **kwargs) -> bool:
       Exception: If an underlying infrastructure error occurs (e.g. RPC failure).
     """
     pass
-    
 
 
 class EnvoyRateLimiter(RateLimiter):
@@ -99,7 +103,7 @@ def __init__(
         throttling is occurring.
     """
     super().__init__(namespace=namespace)
-    
+
     self.service_address = service_address
     self.domain = domain
     self.descriptors = descriptors
@@ -154,9 +158,10 @@ def throttle(self, hits_added: int = 1) -> bool:
         entries.append(RateLimitDescriptorEntry(key=k, value=v))
       proto_descriptors.append(RateLimitDescriptor(entries=entries))
 
-    
     request = RateLimitRequest(
-        domain=self.domain, descriptors=proto_descriptors, hits_addend=hits_added)
+        domain=self.domain,
+        descriptors=proto_descriptors,
+        hits_addend=hits_added)
 
     self.requests_counter.inc()
     attempt = 0
@@ -199,7 +204,7 @@ def throttle(self, hits_added: int = 1) -> bool:
               val = dur.total_seconds()
               if val > sleep_s:
                 sleep_s = val
-        
+
         _LOGGER.warning("Throttled for %s seconds", sleep_s)
         # signal throttled time to backend
         self.throttling_signaler.signal_throttled(int(sleep_s))