refactor: Update main.py with scheduled data refresh and background task

main.py

@@ -1,10 +1,15 @@
+import asyncio
+import concurrent.futures
 import json
 import logging
+import os
 import sqlite3
 from contextlib import asynccontextmanager
 from typing import List
 
 import numpy as np
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.triggers.cron import CronTrigger
 from cashews import NOT_NONE, cache
 from fastapi import FastAPI, HTTPException, Query
 from pandas import Timestamp
@@ -13,6 +18,8 @@ from starlette.responses import RedirectResponse
 
 from data_loader import refresh_data
 
+UPDATE_SCHEDULE = {"hour": os.getenv("UPDATE_INTERVAL_HOURS", "*/6")}
+
 cache.setup("mem://?check_interval=10&size=10000")
 logger = logging.getLogger(__name__)
 
@@ -72,47 +79,91 @@ def serialize_numpy(obj):
     raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
 
 
+def background_refresh_data():
+    logger.info("Starting background data refresh")
+    try:
+        return refresh_data()
+    except Exception as e:
+        logger.error(f"Error in background data refresh: {str(e)}")
+        return None
+
+
+async def update_database():
+    logger.info("Starting scheduled data refresh")
+
+    # Run refresh_data in a background thread
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future = executor.submit(background_refresh_data)
+
+        # Wait for the background task to complete, but allow for cancellation
+        try:
+            datasets = await asyncio.get_event_loop().run_in_executor(
+                None, future.result
+            )
+        except asyncio.CancelledError:
+            future.cancel()
+            logger.info("Data refresh cancelled")
+            return
+
+    if datasets is None:
+        logger.error("Data refresh failed, skipping database update")
+        return
+
+    conn = get_db_connection()
+    try:
+        c = conn.cursor()
+        c.executemany(
+            """
+            INSERT OR REPLACE INTO datasets 
+            (hub_id, likes, downloads, tags, created_at, last_modified, license, language, config_name, column_names, features) 
+            VALUES (?, ?, ?, json(?), ?, ?, json(?), json(?), ?, json(?), json(?))
+            """,
+            [
+                (
+                    data["hub_id"],
+                    data.get("likes", 0),
+                    data.get("downloads", 0),
+                    json.dumps(data.get("tags", []), default=serialize_numpy),
+                    int(data["created_at"].timestamp())
+                    if isinstance(data["created_at"], Timestamp)
+                    else data.get("created_at", 0),
+                    int(data["last_modified"].timestamp())
+                    if isinstance(data["last_modified"], Timestamp)
+                    else data.get("last_modified", 0),
+                    json.dumps(data.get("license", []), default=serialize_numpy),
+                    json.dumps(data.get("language", []), default=serialize_numpy),
+                    data.get("config_name", ""),
+                    json.dumps(data.get("column_names", []), default=serialize_numpy),
+                    json.dumps(data.get("features", []), default=serialize_numpy),
+                )
+                for data in datasets
+            ],
+        )
+        conn.commit()
+        logger.info("Scheduled data refresh completed")
+    except Exception as e:
+        logger.error(f"Error during database update: {str(e)}")
+        conn.rollback()
+    finally:
+        conn.close()
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     setup_database()
-    logger.info("Creating database connection")
-    conn = get_db_connection()
-    logger.info("Refreshing data")
-    datasets = refresh_data()
+    logger.info("Performing initial data refresh")
+    await update_database()
+
+    # Set up the scheduler
+    scheduler = AsyncIOScheduler()
+    # Schedule the update_database function using the UPDATE_SCHEDULE configuration
+    scheduler.add_job(update_database, CronTrigger(**UPDATE_SCHEDULE))
+    scheduler.start()
 
-    c = conn.cursor()
-    c.executemany(
-        """
-        INSERT OR REPLACE INTO datasets 
-        (hub_id, likes, downloads, tags, created_at, last_modified, license, language, config_name, column_names, features) 
-        VALUES (?, ?, ?, json(?), ?, ?, json(?), json(?), ?, json(?), json(?))
-        """,
-        [
-            (
-                data["hub_id"],
-                data.get("likes", 0),
-                data.get("downloads", 0),
-                json.dumps(data.get("tags", []), default=serialize_numpy),
-                int(data["created_at"].timestamp())
-                if isinstance(data["created_at"], Timestamp)
-                else data.get("created_at", 0),
-                int(data["last_modified"].timestamp())
-                if isinstance(data["last_modified"], Timestamp)
-                else data.get("last_modified", 0),
-                json.dumps(data.get("license", []), default=serialize_numpy),
-                json.dumps(data.get("language", []), default=serialize_numpy),
-                data.get("config_name", ""),
-                json.dumps(data.get("column_names", []), default=serialize_numpy),
-                json.dumps(data.get("features", []), default=serialize_numpy),
-            )
-            for data in datasets
-        ],
-    )
-    conn.commit()
-    conn.close()
-    logger.info("Data refreshed")
     yield
 
+    scheduler.shutdown()
+
 
 app = FastAPI(lifespan=lifespan)