[BEAM-13599] Fix overflow error in Python Datastore RampupThrottlingFn

Author: danthev

sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py

@@ -71,7 +71,11 @@ def _calc_max_ops_budget(
     growth = max(
         0.0, (timedelta_since_first - self._RAMP_UP_INTERVAL) /
         self._RAMP_UP_INTERVAL)
-    max_ops_budget = int(self._BASE_BUDGET / self._num_workers * (1.5**growth))
+    try:
+      max_ops_budget = int(
+          self._BASE_BUDGET / self._num_workers * (1.5**growth))
+    except OverflowError:
+      max_ops_budget = float('inf')
     return max(1, max_ops_budget)
 
   def process(self, element, **kwargs):

sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py

@@ -16,6 +16,8 @@
 #
 
 import datetime
+import math
+import sys
 import unittest
 
 from mock import patch
@@ -57,6 +59,20 @@ def test_rampup_throttling(self, mock_sleep, mock_datetime):
       with self.assertRaises(_RampupDelayException):
         next(throttling_fn.process(None))
 
+  def test_budget_overflow(self):
+    throttling_fn = RampupThrottlingFn(num_workers=1)
+
+    normal_date = DATE_ZERO + datetime.timedelta(minutes=2000)
+    normal_budget = throttling_fn._calc_max_ops_budget(DATE_ZERO, normal_date)
+    self.assertNotEqual(normal_budget, float('inf'))
+
+    # This tests that a previously thrown OverflowError is caught.
+    overflow_minutes = math.log(sys.float_info.max) / math.log(1.5) * 5
+    overflow_date = DATE_ZERO + datetime.timedelta(minutes=overflow_minutes)
+    overflow_budget = throttling_fn._calc_max_ops_budget(
+        DATE_ZERO, overflow_date)
+    self.assertEqual(overflow_budget, float('inf'))
+
 
 if __name__ == '__main__':
   unittest.main()