102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
import aiohttp
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.responses import FileResponse
|
|
import asyncmy
|
|
import asyncio
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from contextlib import asynccontextmanager
|
|
import logging
|
|
|
|
load_dotenv()
|
|
|
|
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
|
|
log_formatter = logging.Formatter(
|
|
fmt="%(asctime)s [%(levelname)s] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
file_handler = logging.FileHandler("data/api.log", encoding="utf-8")
|
|
file_handler.setFormatter(log_formatter)
|
|
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setFormatter(log_formatter)
|
|
|
|
logging.basicConfig(
|
|
level=getattr(logging, log_level, logging.INFO),
|
|
handlers=[file_handler, console_handler],
|
|
)
|
|
|
|
@asynccontextmanager
|
|
async def connect_db(app: FastAPI):
|
|
app.state.pool = await asyncmy.create_pool(
|
|
host=os.getenv("DB_HOST", "localhost"),
|
|
port=int(os.getenv("DB_PORT", 3306)),
|
|
user=os.getenv("DB_USER"),
|
|
password=os.getenv("DB_PASSWORD"),
|
|
db=os.getenv("DB_NAME"),
|
|
minsize=5,
|
|
maxsize=20
|
|
)
|
|
await create_tables(app.state.pool)
|
|
task = asyncio.create_task(fetch_images())
|
|
try:
|
|
yield
|
|
finally:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
app.state.pool.close()
|
|
await app.state.pool.wait_closed()
|
|
|
|
app = FastAPI(lifespan=connect_db)
|
|
|
|
async def create_tables(pool):
|
|
async with pool.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS images (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
url VARCHAR(255) NOT NULL,
|
|
filename VARCHAR(255) NOT NULL,
|
|
subreddit VARCHAR(255) NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
await conn.commit()
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return {"message": "yes the api works"}
|
|
|
|
@app.get("/random")
|
|
async def get_random_bnuy():
|
|
return {"message": "here could be a bnuy, if I would've implemented it"}
|
|
|
|
@app.get("/images/{filename}")
|
|
async def get_image(filename: str):
|
|
async with app.state.pool.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("SELECT filename FROM images WHERE filename = %s", (filename,))
|
|
result = await cursor.fetchone()
|
|
if result:
|
|
filepath = os.path.join("data/images", result[0])
|
|
if os.path.exists(filepath):
|
|
return FileResponse(filepath)
|
|
else:
|
|
raise HTTPException(status_code=404, detail="Image file not found")
|
|
else:
|
|
raise HTTPException(status_code=404, detail="Image not found")
|
|
|
|
async def fetch_images():
|
|
from collector import save_picture
|
|
while True:
|
|
try:
|
|
logging.info("Starting image collection...")
|
|
await save_picture(app.state.pool)
|
|
logging.info("Image collection completed. Sleeping for 1 hour...")
|
|
except Exception as e:
|
|
logging.error(f"Error during image collection: {e}")
|
|
await asyncio.sleep(86400) # Sleep for 24 hours |