Merge pull request #118 from andrewm4894/sqlite-add-retry

Enhance SQLite save function with retry logic for locked database errors

Author: andrewm4894

anomstack/df/utils.py

@@ -1,6 +1,7 @@
 from io import StringIO
-from dagster import get_dagster_logger
+
 import pandas as pd
+from dagster import get_dagster_logger
 
 
 def log_df_info(df: pd.DataFrame, logger=None):

anomstack/external/sqlite/sqlite.py

@@ -4,10 +4,14 @@
 
 import os
 import sqlite3
+import time
 
 import pandas as pd
 from dagster import get_dagster_logger
 
+MAX_RETRIES = 5
+RETRY_DELAY = 1
+
 
 def read_sql_sqlite(sql: str) -> pd.DataFrame:
     """
@@ -19,25 +23,19 @@ def read_sql_sqlite(sql: str) -> pd.DataFrame:
     Returns:
         pd.DataFrame: The result of the SQL query as a pandas DataFrame.
     """
-
     logger = get_dagster_logger()
-
     sqlite_path = os.environ.get("ANOMSTACK_SQLITE_PATH", "tmpdata/anomstack.db")
-    logger.info(f"sqlite_path:{sqlite_path}")
-
+    logger.info(f"sqlite_path: {sqlite_path}")
     os.makedirs(os.path.dirname(sqlite_path), exist_ok=True)
 
     conn = sqlite3.connect(sqlite_path)
-
     df = pd.read_sql_query(sql, conn)
-
     conn.close()
     return df
 
-
 def save_df_sqlite(df: pd.DataFrame, table_key: str) -> pd.DataFrame:
     """
-    Save df to db.
+    Save df to db with retry logic.
 
     Args:
         df (pd.DataFrame): The DataFrame to save.
@@ -46,21 +44,27 @@ def save_df_sqlite(df: pd.DataFrame, table_key: str) -> pd.DataFrame:
     Returns:
         pd.DataFrame: The input DataFrame.
     """
-
     logger = get_dagster_logger()
-
     sqlite_path = os.environ.get("ANOMSTACK_SQLITE_PATH", "tmpdata/anomstack.db")
-    logger.info(f"sqlite_path:{sqlite_path}")
-
+    logger.info(f"sqlite_path: {sqlite_path}")
     os.makedirs(os.path.dirname(sqlite_path), exist_ok=True)
 
-    conn = sqlite3.connect(sqlite_path)
-
-    try:
-        df.to_sql(table_key, conn, if_exists='append', index=False)
-    except Exception as e:
-        logger.error(f"Error saving DataFrame to SQLite: {e}")
-        raise
-
-    conn.close()
-    return df
+    attempt = 0
+    while attempt < MAX_RETRIES:
+        try:
+            conn = sqlite3.connect(sqlite_path)
+            df.to_sql(table_key, conn, if_exists='append', index=False)
+            conn.close()
+            return df
+        except sqlite3.OperationalError as e:
+            if "database is locked" in str(e):
+                attempt += 1
+                logger.warning(
+                    f"Database is locked; attempt {attempt} of {MAX_RETRIES}. Retrying in {RETRY_DELAY} seconds..."
+                )
+                time.sleep(RETRY_DELAY)
+            else:
+                logger.error(f"Error saving DataFrame to SQLite: {e}")
+                raise
+    # If all retries fail, raise an error
+    raise sqlite3.OperationalError("Database is locked after multiple attempts.")