lmszx / main.py
Sunrusojsis's picture
Implemented caching
d471223 verified
raw
history blame
2.64 kB
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import Response
import aiohttp
from urllib.parse import urlparse
from aiohttp_client_cache import CachedSession,CacheBackend
app = FastAPI()
# Initialize the aiohttp ClientSession
#http_client = aiohttp.CachedSession()
@app.on_event("startup")
async def startup_event():
cache = CacheBackend(
allowed_methods=('GET', 'POST'),
allowed_codes=(200,503,502),
include_headers=True,
expire_after = 50
)
global http_client
http_client = CachedSession(cache=cache)
@app.on_event("shutdown")
async def shutdown_event():
await http_client.close()
@app.api_route("/{url:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH", "TRACE"])
async def proxy(request: Request, url: str):
# The base_url includes scheme (http/https), so ensure it is provided
target_url = request.query_params.get("url")
if not target_url:
raise HTTPException(status_code=400, detail="Missing 'url' query parameter.")
# Optional: Validate the URL schema
parsed_url = urlparse(target_url)
if parsed_url.scheme not in ['http', 'https']:
raise HTTPException(status_code=400, detail="Invalid URL scheme. Only HTTP and HTTPS are supported.")
# Prepare headers from the original request
headers = dict(request.headers)
headers.pop('host', None)
# Prepare the body for methods that allow a body
body = None
if request.method in ["POST", "PUT", "PATCH"]:
body = await request.body()
try:
async with http_client.request(
method=request.method,
url=target_url,
headers=headers,
data=body,
allow_redirects=False,
ssl = False
) as upstream_response:
# Prepare the response headers, excluding some that might cause issues
response_headers = {k: v for k, v in upstream_response.headers.items()
if k.lower() not in ['transfer-encoding', 'content-encoding']}
# Stream the response back to the client
content = await upstream_response.read()
return Response(content=content, status_code=upstream_response.status,
headers=response_headers)
except aiohttp.ClientError as e:
raise HTTPException(status_code=500, detail=str(e))
# Run the application
# Note: If running this file directly, uncomment the following lines:
# if __name__ == "__main__":
# import uvicorn
# uvicorn.run(app, host="0.0.0.0", port=8000)