You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
139 lines
4.4 KiB
139 lines
4.4 KiB
import sqlite3
|
|
import json
|
|
import time
|
|
import requests
|
|
from typing import Any, Optional, Dict, List
|
|
import pickle
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from threading import Lock
|
|
|
|
class CacheManager:
|
|
def __init__(self, db_name: str = 'pokemon_cache.db', max_connections: int = 10):
|
|
self.db_name = db_name
|
|
self.memory_cache: Dict[str, Any] = {}
|
|
self.expiry_disabled = True
|
|
self.lock = Lock()
|
|
self.connection_pool = ThreadPoolExecutor(max_workers=max_connections)
|
|
self._create_cache_table()
|
|
self.fetch_url_lock = Lock()
|
|
|
|
def _create_connection(self):
|
|
return sqlite3.connect(self.db_name)
|
|
|
|
def _create_cache_table(self):
|
|
with self._create_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS cache (
|
|
key TEXT PRIMARY KEY,
|
|
value BLOB,
|
|
timestamp FLOAT
|
|
)
|
|
''')
|
|
conn.commit()
|
|
|
|
def get(self, key: str) -> Optional[Any]:
|
|
# Check memory cache first
|
|
with self.lock:
|
|
if key in self.memory_cache:
|
|
return self.memory_cache[key]
|
|
|
|
# If not in memory, check database
|
|
def db_get():
|
|
with self._create_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT value, timestamp FROM cache WHERE key = ?', (key,))
|
|
result = cursor.fetchone()
|
|
if result:
|
|
value, timestamp = result
|
|
return pickle.loads(value)
|
|
return None
|
|
|
|
result = self.connection_pool.submit(db_get).result()
|
|
if result:
|
|
with self.lock:
|
|
self.memory_cache[key] = result
|
|
return result
|
|
|
|
def set(self, key: str, value: Any):
|
|
serialized_value = pickle.dumps(value)
|
|
timestamp = time.time()
|
|
|
|
with self.lock:
|
|
self.memory_cache[key] = value
|
|
|
|
def db_set():
|
|
with self._create_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO cache (key, value, timestamp)
|
|
VALUES (?, ?, ?)
|
|
''', (key, serialized_value, timestamp))
|
|
conn.commit()
|
|
|
|
self.connection_pool.submit(db_set)
|
|
|
|
def bulk_get(self, keys: List[str]) -> Dict[str, Any]:
|
|
results = {}
|
|
db_keys = []
|
|
|
|
with self.lock:
|
|
for key in keys:
|
|
if key in self.memory_cache:
|
|
results[key] = self.memory_cache[key]
|
|
else:
|
|
db_keys.append(key)
|
|
|
|
if db_keys:
|
|
def db_bulk_get():
|
|
with self._create_connection() as conn:
|
|
cursor = conn.cursor()
|
|
placeholders = ','.join('?' for _ in db_keys)
|
|
cursor.execute(f'SELECT key, value FROM cache WHERE key IN ({placeholders})', db_keys)
|
|
return {key: pickle.loads(value) for key, value in cursor.fetchall()}
|
|
|
|
db_results = self.connection_pool.submit(db_bulk_get).result()
|
|
with self.lock:
|
|
self.memory_cache.update(db_results)
|
|
results.update(db_results)
|
|
|
|
return results
|
|
|
|
def fetch_url(self, url: str, force_refresh: bool = False, expiry: int = 86400) -> Optional[str]:
|
|
cache_key = f"url_{url}"
|
|
if not force_refresh:
|
|
cached_data = self.get(cache_key)
|
|
if cached_data:
|
|
cached_time = cached_data['timestamp']
|
|
if time.time() - cached_time < expiry or self.expiry_disabled:
|
|
return cached_data['content']
|
|
|
|
#with self.fetch_url_lock:
|
|
print(f"Fetching URL: {url}")
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
content = response.text
|
|
self.set(cache_key, {
|
|
'content': content,
|
|
'timestamp': time.time()
|
|
})
|
|
time.sleep(1)
|
|
return content
|
|
return None
|
|
|
|
def close(self):
|
|
self.connection_pool.shutdown()
|
|
|
|
# Usage example
|
|
if __name__ == "__main__":
|
|
cache = CacheManager()
|
|
|
|
# Example usage
|
|
url = "https://example.com"
|
|
data = cache.fetch_url(url)
|
|
if data:
|
|
print("Data fetched successfully")
|
|
else:
|
|
print("Failed to fetch data")
|
|
|
|
cache.close()
|
|
|