add retry for db

1 files changed, 17 insertions, 1 deletions

diff --git a/server/services/get_data.py b/server/services/get_data.py
index 3014ea1..fbbf12f 100644
--- a/server/services/get_data.py
+++ b/server/services/get_data.py
@@ -4,6 +4,8 @@ from db import SessionLocal
 from models.market import MarketDataDB
 from gridstatus import ISONE, MISO, NYISO
 from typing import Dict, Type
+from sqlalchemy.exc import OperationalError
+import time
 
 MARKET_CLASSES: Dict[str, Type] = {
     "ISONE": ISONE,
@@ -108,7 +110,21 @@ def update_market_data():
                 )
                 db.add(entry)
 
-    db.commit()
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            db.commit()
+            break
+        except OperationalError as e:
+            if "database is locked" in str(e).lower():
+                if attempt < max_retries - 1:
+                    time.sleep(2)
+                    continue
+                else:
+                    raise Exception("Database is locked, and retries exhausted.")
+            else:
+                raise e
+
     db.close()