Enrichment Transform with BigTable handler

sdks/python/apache_beam/io/requestresponseio_it_test.py → sdks/python/apache_beam/io/requestresponse_it_test.py

@@ -16,6 +16,7 @@
 #
 import base64
 import sys
+import typing
 import unittest
 from dataclasses import dataclass
 from typing import Tuple
@@ -24,13 +25,18 @@
 import urllib3
 
 import apache_beam as beam
-from apache_beam.io.requestresponseio import Caller
-from apache_beam.io.requestresponseio import RequestResponseIO
-from apache_beam.io.requestresponseio import UserCodeExecutionException
-from apache_beam.io.requestresponseio import UserCodeQuotaException
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.testing.test_pipeline import TestPipeline
 
+# pylint: disable=ungrouped-imports
+try:
+  from apache_beam.io.requestresponse import Caller
+  from apache_beam.io.requestresponse import RequestResponseIO
+  from apache_beam.io.requestresponse import UserCodeExecutionException
+  from apache_beam.io.requestresponse import UserCodeQuotaException
+except ImportError:
+  raise unittest.SkipTest('RequestResponseIO dependencies are not installed.')
+
 _HTTP_PATH = '/v1/echo'
 _PAYLOAD = base64.b64encode(bytes('payload', 'utf-8'))
 _HTTP_ENDPOINT_ADDRESS_FLAG = '--httpEndpointAddress'
@@ -61,28 +67,27 @@ def _add_argparse_args(cls, parser) -> None:
         help='The ID for an allocated quota that should exceed.')
 
 
-# TODO(riteshghorse,damondouglas) replace Echo(Request|Response) with proto
-#   generated classes from .test-infra/mock-apis:
 @dataclass
-class EchoRequest:
+class EchoResponse:
   id: str
   payload: bytes
 
 
-@dataclass
-class EchoResponse:
+# TODO(riteshghorse,damondouglas) replace Echo(Request|Response) with proto
+#   generated classes from .test-infra/mock-apis:
+class Request(typing.NamedTuple):
   id: str
   payload: bytes
 
 
-class EchoHTTPCaller(Caller):
+class EchoHTTPCaller(Caller[Request, EchoResponse]):
   """Implements ``Caller`` to call the ``EchoServiceGrpc``'s HTTP handler.
     The purpose of ``EchoHTTPCaller`` is to support integration tests.
     """
   def __init__(self, url: str):
     self.url = url + _HTTP_PATH
 
-  def __call__(self, request: EchoRequest, *args, **kwargs) -> EchoResponse:
+  def __call__(self, request: Request, *args, **kwargs) -> EchoResponse:
     """Overrides ``Caller``'s call method invoking the
         ``EchoServiceGrpc``'s HTTP handler with an ``EchoRequest``, returning
         either a successful ``EchoResponse`` or throwing either a
@@ -129,7 +134,7 @@ def setUpClass(cls) -> None:
   def setUp(self) -> None:
     client, options = EchoHTTPCallerTestIT._get_client_and_options()
 
-    req = EchoRequest(id=options.should_exceed_quota_id, payload=_PAYLOAD)
+    req = Request(id=options.should_exceed_quota_id, payload=_PAYLOAD)
     try:
       # The following is needed to exceed the API
       client(req)
@@ -148,7 +153,7 @@ def _get_client_and_options(cls) -> Tuple[EchoHTTPCaller, EchoITOptions]:
   def test_given_valid_request_receives_response(self):
     client, options = EchoHTTPCallerTestIT._get_client_and_options()
 
-    req = EchoRequest(id=options.never_exceed_quota_id, payload=_PAYLOAD)
+    req = Request(id=options.never_exceed_quota_id, payload=_PAYLOAD)
 
     response: EchoResponse = client(req)
 
@@ -158,20 +163,20 @@ def test_given_valid_request_receives_response(self):
   def test_given_exceeded_quota_should_raise(self):
     client, options = EchoHTTPCallerTestIT._get_client_and_options()
 
-    req = EchoRequest(id=options.should_exceed_quota_id, payload=_PAYLOAD)
+    req = Request(id=options.should_exceed_quota_id, payload=_PAYLOAD)
 
     self.assertRaises(UserCodeQuotaException, lambda: client(req))
 
   def test_not_found_should_raise(self):
     client, _ = EchoHTTPCallerTestIT._get_client_and_options()
 
-    req = EchoRequest(id='i-dont-exist-quota-id', payload=_PAYLOAD)
+    req = Request(id='i-dont-exist-quota-id', payload=_PAYLOAD)
     self.assertRaisesRegex(
         UserCodeExecutionException, "Not Found", lambda: client(req))
 
   def test_request_response_io(self):
     client, options = EchoHTTPCallerTestIT._get_client_and_options()
-    req = EchoRequest(id=options.never_exceed_quota_id, payload=_PAYLOAD)
+    req = Request(id=options.never_exceed_quota_id, payload=_PAYLOAD)
     with TestPipeline(is_integration_test=True) as test_pipeline:
       output = (
           test_pipeline

sdks/python/apache_beam/io/requestresponseio_test.py → sdks/python/apache_beam/io/requestresponse_test.py

@@ -18,14 +18,21 @@
 import unittest
 
 import apache_beam as beam
-from apache_beam.io.requestresponseio import Caller
-from apache_beam.io.requestresponseio import RequestResponseIO
-from apache_beam.io.requestresponseio import UserCodeExecutionException
-from apache_beam.io.requestresponseio import UserCodeTimeoutException
 from apache_beam.testing.test_pipeline import TestPipeline
 
+# pylint: disable=ungrouped-imports
+try:
+  from google.api_core.exceptions import TooManyRequests
+  from apache_beam.io.requestresponse import Caller
+  from apache_beam.io.requestresponse import RequestResponseIO
+  from apache_beam.io.requestresponse import UserCodeExecutionException
+  from apache_beam.io.requestresponse import UserCodeTimeoutException
+  from apache_beam.io.requestresponse import retry_on_exception
+except ImportError:
+  raise unittest.SkipTest('RequestResponseIO dependencies are not installed.')
 
-class AckCaller(Caller):
+
+class AckCaller(Caller[str, str]):
   """AckCaller acknowledges the incoming request by returning a
   request with ACK."""
   def __enter__(self):
@@ -54,6 +61,20 @@ def __call__(self, request: str, *args, **kwargs):
       raise RuntimeError("Exception expected, not an error.")
 
 
+class CallerThatRetries(AckCaller):
+  def __init__(self):
+    self.count = -1
+
+  def __call__(self, request: str, *args, **kwargs):
+    try:
+      pass
+    except Exception as e:
+      raise e
+    finally:
+      self.count += 1
+      raise TooManyRequests('retries = %d' % self.count)
+
+
 class TestCaller(unittest.TestCase):
   def test_valid_call(self):
     caller = AckCaller()
@@ -83,6 +104,30 @@ def test_call_runtime_error(self):
             | beam.Create([""])
             | RequestResponseIO(caller=caller))
 
+  def test_retry_on_exception(self):
+    self.assertFalse(retry_on_exception(RuntimeError()))
+    self.assertTrue(retry_on_exception(TooManyRequests("HTTP 429")))
+
+  def test_caller_backoff_retry_strategy(self):
+    caller = CallerThatRetries()
+    with self.assertRaises(TooManyRequests) as cm:
+      with TestPipeline() as test_pipeline:
+        _ = (
+            test_pipeline
+            | beam.Create(["sample_request"])
+            | RequestResponseIO(caller=caller))
+    self.assertRegex(cm.exception.message, 'retries = 2')
+
+  def test_caller_no_retry_strategy(self):
+    caller = CallerThatRetries()
+    with self.assertRaises(TooManyRequests) as cm:
+      with TestPipeline() as test_pipeline:
+        _ = (
+            test_pipeline
+            | beam.Create(["sample_request"])
+            | RequestResponseIO(caller=caller, repeater=None))
+    self.assertRegex(cm.exception.message, 'retries = 0')
+
 
 if __name__ == '__main__':
   unittest.main()