feat: Add scheduled collections refresh

main.py

@@ -23,6 +23,7 @@ from data_loader import refresh_data
 login(token=os.getenv("HF_TOKEN"))
 
 UPDATE_SCHEDULE = {"hour": os.getenv("UPDATE_INTERVAL_HOURS", "*/6")}
+COLLECTION_UPDATE_SCHEDULE = {"hour": "0"}  # Run at midnight every day
 
 cache.setup("mem://?check_interval=10&size=10000")
 logger = logging.getLogger(__name__)
@@ -95,11 +96,9 @@ def background_refresh_data():
 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
@@ -151,7 +150,6 @@ async def update_database():
     finally:
         conn.close()
 
-    # Upload the database file to Hugging Face Hub
     try:
         upload_file(
             path_or_fileobj="datasets.db",
@@ -187,14 +185,13 @@ async def lifespan(app: FastAPI):
     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))
-    # Schedule the update_collections function to run daily at midnight
-    scheduler.add_job(update_collections, CronTrigger(hour=0, minute=0))
+    scheduler.add_job(update_collections, CronTrigger(**COLLECTION_UPDATE_SCHEDULE))
     scheduler.start()
 
+    await update_collections()
+
     yield
 
     scheduler.shutdown()
@@ -256,7 +253,6 @@ async def search_datasets(
 
         results = [dict(row) for row in c.fetchall()]
 
-        # Get total count
         if match_all:
             count_query = """
             SELECT COUNT(*) as total FROM datasets