Fixed bug in @retried when exception subtypes were not respected

databricks/sdk/retries.py

@@ -39,8 +39,11 @@ def wrapper(*args, **kwargs):
                         retry_reason = 'throttled by platform'
                     elif is_retryable is not None:
                         retry_reason = is_retryable(err)
-                    elif type(err) in on:
-                        retry_reason = f'{type(err).__name__} is allowed to retry'
+                    elif on is not None:
+                        for err_type in on:
+                            if not isinstance(err, err_type):
+                                continue
+                            retry_reason = f'{type(err).__name__} is allowed to retry'
 
                     if retry_reason is None:
                         # raise if exception is not retryable

tests/test_retries.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from databricks.sdk.errors import NotFound, ResourceDoesNotExist
 from databricks.sdk.retries import retried
 
 
@@ -41,6 +42,16 @@ def foo():
         foo()
 
 
+def test_match_on_subclass():
+    with pytest.raises(TimeoutError):
+
+        @retried(on=[NotFound], timeout=timedelta(seconds=0.5))
+        def foo():
+            raise ResourceDoesNotExist(...)
+
+        foo()
+
+
 def test_propagates_outside_exception():
     with pytest.raises(KeyError):