Retry deploy schema on InternalError (#283)

Retry deploy schema on `InternalError` to improve resilience against sporadic Databricks issues Co-authored-by: Serge Smertin <[email protected]>
databrickslabs · Sep 16, 2024 · 2ad66b6 · 2ad66b6
1 parent 2216a12
commit 2ad66b6
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/databricks/labs/lsql/deployment.py b/src/databricks/labs/lsql/deployment.py
@@ -1,8 +1,11 @@
+import datetime as dt
 import logging
 import pkgutil
 from typing import Any
 
 from databricks.labs.lsql.backends import Dataclass, SqlBackend
+from databricks.sdk.errors import InternalError
+from databricks.sdk.retries import retried
 
 logger = logging.getLogger(__name__)
 
@@ -13,6 +16,8 @@ def __init__(self, sql_backend: SqlBackend, inventory_schema: str, mod: Any):
         self._inventory_schema = inventory_schema
         self._module = mod
 
+    # InternalError are retried for resilience on sporadic Databricks issues
+    @retried(on=[InternalError], timeout=dt.timedelta(minutes=1))
     def deploy_schema(self):
         logger.info(f"Ensuring {self._inventory_schema} database exists")
         self._sql_backend.execute(f"CREATE SCHEMA IF NOT EXISTS hive_metastore.{self._inventory_schema}")