[Tables] wrapper for 429s (Azure#18996)

This will allow us to remove all sleeps in the code

Author: seankane-msft

sdk/tables/azure-data-tables/tests/_shared/asynctestcase.py

@@ -67,8 +67,6 @@ async def _tear_down(self):
         if is_live():
             async for table in self.ts.list_tables():
                 await self.ts.delete_table(table.name)
-            if self.ts._cosmos_endpoint:
-                self.sleep(SLEEP_DELAY)
             self.test_tables = []
             await self.ts.close()
 

sdk/tables/azure-data-tables/tests/_shared/testcase.py

@@ -99,8 +99,10 @@ def _create_table(self, ts, prefix=TEST_TABLE_PREFIX, table_list=None):
         return table
 
     def _delete_all_tables(self, ts):
-        for table in ts.list_tables():
-            ts.delete_table(table.name)
+        if self.is_live:
+            for table in ts.list_tables():
+                ts.delete_table(table.name)
+            self.sleep(10)
 
     def _create_pk_rk(self, pk, rk):
         try:
@@ -354,8 +356,6 @@ def _tear_down(self):
         if is_live():
             self._delete_all_tables(self.ts)
             self.test_tables = []
-            if self.ts._cosmos_endpoint:
-                self.sleep(SLEEP_DELAY)
             self.ts.close()
 
     def _create_query_table(self, entity_count):

sdk/tables/azure-data-tables/tests/async_preparers.py

@@ -1,10 +1,14 @@
-import functools
+import time
 
 from azure.core.credentials import AzureNamedKeyCredential
+from azure.core.exceptions import HttpResponseError
+
+from devtools_testutils import is_live
+
 from preparers import CosmosPreparer, TablesPreparer, trim_kwargs_from_test_function
 
-def cosmos_decorator_async(func, **kwargs):
 
+def cosmos_decorator_async(func, **kwargs):
     @CosmosPreparer()
     async def wrapper(*args, **kwargs):
         key = kwargs.pop("tables_primary_cosmos_account_key")
@@ -14,7 +18,7 @@ async def wrapper(*args, **kwargs):
         kwargs["tables_primary_cosmos_account_key"] = key
         kwargs["tables_cosmos_account_name"] = name
 
-        trimmed_kwargs = {k:v for k, v in kwargs.items()}
+        trimmed_kwargs = {k: v for k, v in kwargs.items()}
         trim_kwargs_from_test_function(func, trimmed_kwargs)
 
         return await func(*args, **trimmed_kwargs)
@@ -23,7 +27,6 @@ async def wrapper(*args, **kwargs):
 
 
 def tables_decorator_async(func, **kwargs):
-
     @TablesPreparer()
     async def wrapper(*args, **kwargs):
         key = kwargs.pop("tables_primary_storage_account_key")
@@ -33,9 +36,28 @@ async def wrapper(*args, **kwargs):
         kwargs["tables_primary_storage_account_key"] = key
         kwargs["tables_storage_account_name"] = name
 
-        trimmed_kwargs = {k:v for k, v in kwargs.items()}
+        trimmed_kwargs = {k: v for k, v in kwargs.items()}
         trim_kwargs_from_test_function(func, trimmed_kwargs)
 
-        return await func(*args, **trimmed_kwargs)
+        EXPONENTIAL_BACKOFF = 1.5
+        RETRY_COUNT = 0
+
+        try:
+            return await func(*args, **trimmed_kwargs)
+        except HttpResponseError as exc:
+            if exc.status_code != 429:
+                raise
+            print("Retrying: {} {}".format(RETRY_COUNT, EXPONENTIAL_BACKOFF))
+            while RETRY_COUNT < 6:
+                if is_live():
+                    time.sleep(EXPONENTIAL_BACKOFF)
+                try:
+                    return await func(*args, **trimmed_kwargs)
+                except HttpResponseError as exc:
+                    print("Retrying: {} {}".format(RETRY_COUNT, EXPONENTIAL_BACKOFF))
+                    EXPONENTIAL_BACKOFF **= 2
+                    RETRY_COUNT += 1
+                    if exc.status_code != 429 or RETRY_COUNT >= 6:
+                        raise
 
     return wrapper

sdk/tables/azure-data-tables/tests/preparers.py

@@ -1,38 +1,41 @@
 import functools
 import inspect
+import time
 
 from azure.core.credentials import AzureNamedKeyCredential
-from devtools_testutils import PowerShellPreparer
+from azure.core.exceptions import HttpResponseError
+from devtools_testutils import PowerShellPreparer, is_live
 
 CosmosPreparer = functools.partial(
-    PowerShellPreparer, "tables",
+    PowerShellPreparer,
+    "tables",
     tables_cosmos_account_name="fake_cosmos_account",
-    tables_primary_cosmos_account_key="fakecosmosaccountkey"
+    tables_primary_cosmos_account_key="fakecosmosaccountkey",
 )
 
 TablesPreparer = functools.partial(
-    PowerShellPreparer, "tables",
+    PowerShellPreparer,
+    "tables",
     tables_storage_account_name="fake_table_account",
-    tables_primary_storage_account_key="faketablesaccountkey"
+    tables_primary_storage_account_key="faketablesaccountkey",
 )
 
 
 def trim_kwargs_from_test_function(fn, kwargs):
     # the next function is the actual test function. the kwargs need to be trimmed so
     # that parameters which are not required will not be passed to it.
-    if not getattr(fn, '__is_preparer', False):
+    if not getattr(fn, "__is_preparer", False):
         try:
             args, _, kw, _, _, _, _ = inspect.getfullargspec(fn)
         except AttributeError:
-            args, _, kw, _ = inspect.getargspec(fn) # pylint: disable=deprecated-method
+            args, _, kw, _ = inspect.getargspec(fn)  # pylint: disable=deprecated-method
         if kw is None:
             args = set(args)
             for key in [k for k in kwargs if k not in args]:
                 del kwargs[key]
 
 
 def tables_decorator(func, **kwargs):
-
     @TablesPreparer()
     def wrapper(*args, **kwargs):
         key = kwargs.pop("tables_primary_storage_account_key")
@@ -42,7 +45,7 @@ def wrapper(*args, **kwargs):
         kwargs["tables_primary_storage_account_key"] = key
         kwargs["tables_storage_account_name"] = name
 
-        trimmed_kwargs = {k:v for k, v in kwargs.items()}
+        trimmed_kwargs = {k: v for k, v in kwargs.items()}
         trim_kwargs_from_test_function(func, trimmed_kwargs)
 
         func(*args, **trimmed_kwargs)
@@ -51,7 +54,6 @@ def wrapper(*args, **kwargs):
 
 
 def cosmos_decorator(func, **kwargs):
-
     @CosmosPreparer()
     def wrapper(*args, **kwargs):
         key = kwargs.pop("tables_primary_cosmos_account_key")
@@ -61,10 +63,28 @@ def wrapper(*args, **kwargs):
         kwargs["tables_primary_cosmos_account_key"] = key
         kwargs["tables_cosmos_account_name"] = name
 
-        trimmed_kwargs = {k:v for k, v in kwargs.items()}
+        trimmed_kwargs = {k: v for k, v in kwargs.items()}
         trim_kwargs_from_test_function(func, trimmed_kwargs)
 
-        func(*args, **trimmed_kwargs)
+        EXPONENTIAL_BACKOFF = 1
+        RETRY_COUNT = 0
 
-    return wrapper
+        try:
+            return func(*args, **trimmed_kwargs)
+        except HttpResponseError as exc:
+            if exc.status_code != 429:
+                raise
+            print("Retrying: {} {}".format(RETRY_COUNT, EXPONENTIAL_BACKOFF))
+            while RETRY_COUNT < 6:
+                if is_live():
+                    time.sleep(EXPONENTIAL_BACKOFF)
+                try:
+                    return func(*args, **trimmed_kwargs)
+                except HttpResponseError as exc:
+                    print("Retrying: {} {}".format(RETRY_COUNT, EXPONENTIAL_BACKOFF))
+                    EXPONENTIAL_BACKOFF **= 2
+                    RETRY_COUNT += 1
+                    if exc.status_code != 429 or RETRY_COUNT >= 6:
+                        raise
 
+    return wrapper