更新

README.md

@@ -74,25 +74,12 @@ https://setu.yuban10703.xyz/docs
 | large    | string($uri) | 链接(画质:large)   |
 | medium   | string($uri) | 链接(画质:medium)  |
 
-### docker
-
-~~~
-docker build -t setuapi:v1.7 .
-~~~
-
-~~~
-docker run -d \
--p 9001:80 \
--e mongodb="mongodb+srv://username:password@cludn.mongodb.net/setu?retryWrites=true&w=majority" \
--e db="setu" \
--e col="setu_v5" \
--e LOG_LEVEL="debug" \
-setuapi:v1.7
-~~~
-
 ### mongodb
 
-要给r18,tags字段分别建索引
+要给r18,delete,tags建索引
+
+db.your_collection.createIndex({ r18: 1, delete: 1, tags: "text" })
+
 
 ### 感谢
 

app/database.py

@@ -0,0 +1,50 @@
+import asyncio
+import re
+from typing import List
+import os
+from fastapi.logger import logger
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from schema import Setu
+
+
+class Database:
+    def __init__(self):
+        self.client = AsyncIOMotorClient(os.getenv('mongodb'))
+        self.client.get_io_loop = asyncio.get_event_loop
+        self.database = self.client[os.getenv('database')]
+        self.collection = self.database[os.getenv('collection')]
+        self.Flag = None
+
+
+DB = Database()
+
+
+def get_setu_cursor(r18: int, num: int, tags: set):
+    condition = {
+        "$and": [{"r18": {"$nin": [None]}}, {"delete": False}] if r18 == 2 else [{"r18": bool(r18)}, {"delete": False}]
+    }
+    for tag in tags:
+        condition["$and"].append({"tags": {"$regex": re.compile(tag, re.I)}})
+    return DB.collection.aggregate(
+        [{"$match": condition}, {"$sample": {"size": num}}, {"$project": {"_id": 0, "delete": 0}}]
+    )
+
+
+
+async def get_setu_data(r18: int, num: int, tags: set, replace_url) -> List[Setu]:
+    logger.info(f"r18:[{r18}] num:[{num}] tags:[{tags}] replace_url:[{replace_url}]")
+    cursor = get_setu_cursor(r18, num, tags)
+    if replace_url:
+        setus = []
+        async for setu in cursor:
+
+            urls_dict: dict = setu["urls"]
+            for url_quality, url in urls_dict.items():
+                setu["urls"][url_quality] = url.replace(
+                    "https://i.pximg.net", replace_url
+                )
+            setus.append(setu)
+        return setus
+    else:
+        return await cursor.to_list(length=None)

app/main.py

@@ -1,63 +1,59 @@
-import logging
-from typing import Set, Optional
-
-from fastapi import FastAPI, Query, HTTPException
-from fastapi.logger import logger
-from fastapi.responses import ORJSONResponse
-from pydantic import HttpUrl
-
-from app.database import database, get_setu_data
-from app.model import Item, Setu_out
-
-gunicorn_logger = logging.getLogger('gunicorn.error')
-logger.handlers = gunicorn_logger.handlers
-
-if __name__ != "main":
-    logger.setLevel(gunicorn_logger.level)
-else:
-    logger.setLevel(logging.DEBUG)
-
-app = FastAPI(
-    title="setu",
-    description="https://github.com/yuban10703/SetuAPI",
-    version="0.1.4"
-)
-
-
-@app.on_event("startup")
-async def on_app_start():
-    import os
-    await database.connect(os.getenv('mongodb'), os.getenv('db'), os.getenv('col'))
-
-
-@app.on_event("shutdown")
-async def on_app_shutdown():
-    await database.close()
-
-
-@app.get('/setu', response_model=Setu_out, response_class=ORJSONResponse)
-async def setu_get(
-        r18: Optional[int] = Query(0, ge=0, le=2),
-        num: Optional[int] = Query(1, ge=1, le=50),
-        tags: Set[str] = Query(set()),
-        replace_url: Optional[HttpUrl] = None):
-    setus = await get_setu_data(r18, num, tags, replace_url)
-    if not setus:
-        raise HTTPException(status_code=404, detail="色图库中没找到色图~")
-    return {'detail': '', 'count': len(setus), 'tags': [tag for tag in tags if not tag.isspace() and len(tag) != 0],
-            'data': setus}
-
-
-@app.post("/setu", response_model=Setu_out, response_class=ORJSONResponse)
-async def setu_post(item: Item):
-    setus = await get_setu_data(item.r18, item.num, item.tags, item.replace_url)
-    if not setus:
-        raise HTTPException(status_code=404, detail="色图库中没找到色图~")
-    return {'detail': '', 'count': len(setus),
-            'tags': [tag for tag in item.tags if not tag.isspace() and len(tag) != 0], 'data': setus}
-
-
-if __name__ == '__main__':
-    import uvicorn
-
-    uvicorn.run('main:app', host="0.0.0.0", port=8080, reload=True)
+import logging
+from typing import Optional, Set
+
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.logger import logger
+from fastapi.responses import ORJSONResponse
+from pydantic import HttpUrl
+
+from database import get_setu_data
+from schema import Item, Setu_out
+
+gunicorn_logger = logging.getLogger("gunicorn.error")
+logger.handlers = gunicorn_logger.handlers
+
+if __name__ != "main":
+    logger.setLevel(gunicorn_logger.level)
+else:
+    logger.setLevel(logging.DEBUG)
+
+app = FastAPI(
+    title="setu", description="https://github.com/yuban10703/SetuAPI", version="0.1.5"
+)
+
+
+@app.get("/setu", response_model=Setu_out, response_class=ORJSONResponse)
+async def setu_get(
+        r18: Optional[int] = Query(0, ge=0, le=2),
+        num: Optional[int] = Query(1, ge=1, le=50),
+        tags: Set[str] = Query(set()),
+        replace_url: Optional[HttpUrl] = None,
+):
+    filtered_tags = {tag.strip() for tag in tags if tag.strip()}
+    setus = await get_setu_data(r18, num, filtered_tags, replace_url)
+    if not setus:
+        raise HTTPException(status_code=404, detail="色图库中没找到色图~")
+    return {
+        "detail": "",
+        "count": len(setus),
+        "tags": filtered_tags,
+        "data": setus,
+    }
+
+@app.post("/setu", response_class=ORJSONResponse)
+async def setu_post(item: Item):
+    setus = await get_setu_data(item.r18, item.num, item.tags, item.replace_url)
+    if not setus:
+        raise HTTPException(status_code=404, detail="色图库中没找到色图~")
+    return {
+        "detail": "",
+        "count": len(setus),
+        "tags": item.tags,
+        "data": setus,
+    }
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run("main:app", host="0.0.0.0", port=9999, reload=True)