add caching to puppeteer

Dockers/puppeteer-api/Dockerfile

@@ -5,6 +5,7 @@ RUN apt-get update && apt-get install -y \
     wget \
     gnupg2 \
     procps \
+    sqlite3 \
     && wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /usr/share/keyrings/google-chrome-keyring.gpg \
     && echo "deb [arch=amd64 signed-by=/usr/share/keyrings/google-chrome-keyring.gpg] http://dl.google.com/linux/chrome/deb/ stable main" | tee /etc/apt/sources.list.d/google-chrome.list \
     && apt-get update \

Dockers/puppeteer-api/main.py

@@ -4,6 +4,9 @@ import os
 import asyncio
 import json
 import re
+import sqlite3
+import time
+from datetime import datetime, timedelta
 from typing import Optional, Dict, List, Any
 from urllib.parse import unquote
 
@@ -17,6 +20,59 @@ if not API_KEY:
 # Define custom user agent
 CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36'
 
+# Initialize SQLite database
+def init_db():
+    conn = sqlite3.connect('cache.db')
+    cursor = conn.cursor()
+    cursor.execute('''
+    CREATE TABLE IF NOT EXISTS cache (
+        url TEXT,
+        route TEXT,
+        data TEXT,
+        timestamp INTEGER,
+        PRIMARY KEY (url, route)
+    )
+    ''')
+    conn.commit()
+    conn.close()
+
+# Get cached data if it exists and is not older than 36 hours
+def get_cached_data(url, route):
+    conn = sqlite3.connect('cache.db')
+    cursor = conn.cursor()
+    cache_expiry = int(time.time()) - (36 * 60 * 60)  # 36 hours in seconds
+    cursor.execute(
+        "SELECT data FROM cache WHERE url = ? AND route = ? AND timestamp > ?",
+        (url, route, cache_expiry)
+    )
+    result = cursor.fetchone()
+    conn.close()
+
+    if result:
+        print(f"Cache hit for {url} on route {route}")
+        return json.loads(result[0])
+    return None
+
+# Save data to cache
+def save_to_cache(url, route, data):
+    conn = sqlite3.connect('cache.db')
+    cursor = conn.cursor()
+    timestamp = int(time.time())
+
+    # Convert data to JSON string
+    data_json = json.dumps(data)
+
+    cursor.execute(
+        "INSERT OR REPLACE INTO cache (url, route, data, timestamp) VALUES (?, ?, ?, ?)",
+        (url, route, data_json, timestamp)
+    )
+    conn.commit()
+    conn.close()
+    print(f"Saved to cache: {url} on route {route}")
+
+# Initialize database on startup
+init_db()
+
 async def wait_for_network_idle(page):
     """Wait until no network requests are in flight"""
     await page.waitForNetworkIdle(idleTime=500, timeout=30000)
@@ -31,10 +87,16 @@ async def visit_url(url: str, x_api_key: Optional[str] = Header(None)):
     if not x_api_key or x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API key")
 
+    # Decode URL if it's encoded
+    decoded_url = unquote(url)
+
+    # Check cache first
+    cached_result = get_cached_data(decoded_url, "visit")
+    if cached_result:
+        return cached_result
+
     browser = None
     try:
-        # Decode URL if it's encoded
-        decoded_url = unquote(url)
         print(decoded_url)
 
         # Launch browser using installed Chrome
@@ -59,7 +121,12 @@ async def visit_url(url: str, x_api_key: Optional[str] = Header(None)):
         # Close page
         await page.close()
 
-        return {"status": "success", "content": content}
+        result = {"status": "success", "content": content}
+
+        # Save to cache
+        save_to_cache(decoded_url, "visit", result)
+
+        return result
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -76,10 +143,16 @@ async def extract_seo(url: str, x_api_key: Optional[str] = Header(None)):
     if not x_api_key or x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API key")
 
+    # Decode URL if it's encoded
+    decoded_url = unquote(url)
+
+    # Check cache first
+    cached_result = get_cached_data(decoded_url, "seo")
+    if cached_result:
+        return cached_result
+
     browser = None
     try:
-        # Decode URL if it's encoded
-        decoded_url = unquote(url)
         print(f"Extracting SEO from: {decoded_url}")
 
         # Launch browser using installed Chrome
@@ -169,7 +242,12 @@ async def extract_seo(url: str, x_api_key: Optional[str] = Header(None)):
             "twitter_card": twitter_data
         }
 
-        return {"status": "success", "seo_data": seo_data}
+        result = {"status": "success", "seo_data": seo_data}
+
+        # Save to cache
+        save_to_cache(decoded_url, "seo", result)
+
+        return result
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -186,10 +264,16 @@ async def extract_meta_tags(url: str, x_api_key: Optional[str] = Header(None)):
     if not x_api_key or x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API key")
 
+    # Decode URL if it's encoded
+    decoded_url = unquote(url)
+
+    # Check cache first
+    cached_result = get_cached_data(decoded_url, "meta")
+    if cached_result:
+        return cached_result
+
     browser = None
     try:
-        # Decode URL if it's encoded
-        decoded_url = unquote(url)
         print(f"Extracting meta tags from: {decoded_url}")
 
         # Launch browser using installed Chrome
@@ -227,13 +311,18 @@ async def extract_meta_tags(url: str, x_api_key: Optional[str] = Header(None)):
         # Close page
         await page.close()
 
-        return {
+        result = {
             "status": "success",
             "url": decoded_url,
             "title": title,
             "meta_tags": meta_tags
         }
 
+        # Save to cache
+        save_to_cache(decoded_url, "meta", result)
+
+        return result
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -242,6 +331,55 @@ async def extract_meta_tags(url: str, x_api_key: Optional[str] = Header(None)):
         if browser:
             await browser.close()
 
+@app.get("/cache/clear")
+async def clear_cache(x_api_key: Optional[str] = Header(None)):
+    """Clear the entire cache database"""
+    # Validate API key
+    if not x_api_key or x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    conn = sqlite3.connect('cache.db')
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM cache")
+    conn.commit()
+    conn.close()
+
+    return {"status": "success", "message": "Cache cleared successfully"}
+
+@app.get("/cache/stats")
+async def cache_stats(x_api_key: Optional[str] = Header(None)):
+    """Get cache statistics"""
+    # Validate API key
+    if not x_api_key or x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    conn = sqlite3.connect('cache.db')
+    cursor = conn.cursor()
+
+    # Get total entries
+    cursor.execute("SELECT COUNT(*) FROM cache")
+    total_entries = cursor.fetchone()[0]
+
+    # Get entries by route
+    cursor.execute("SELECT route, COUNT(*) FROM cache GROUP BY route")
+    routes = {route: count for route, count in cursor.fetchall()}
+
+    # Get recent entries (last 24 hours)
+    recent_timestamp = int(time.time()) - (24 * 60 * 60)
+    cursor.execute("SELECT COUNT(*) FROM cache WHERE timestamp > ?", (recent_timestamp,))
+    recent_entries = cursor.fetchone()[0]
+
+    conn.close()
+
+    return {
+        "status": "success",
+        "stats": {
+            "total_entries": total_entries,
+            "entries_by_route": routes,
+            "recent_entries": recent_entries
+        }
+    }
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)