Commit cd235ec
Changed files (1)
src
src/database.py
@@ -12,6 +12,7 @@ import json
import os
import shutil
import warnings
+from datetime import timedelta
from pathlib import Path
from urllib.parse import quote_plus, unquote_plus
@@ -25,7 +26,7 @@ from loguru import logger
from config import DB, DOWNLOAD_DIR, cache
from networking import hx_req
-from utils import bare_url, stringfy
+from utils import bare_url, nowdt, stringfy
# hot fix: https://developers.cloudflare.com/r2/examples/aws/boto3/
os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"] = "when_required"
@@ -133,21 +134,23 @@ async def get_cf_r2(key: str) -> dict:
return {}
-async def set_db(key: str, data: dict | list, ttl: int = 600, metadata: dict | None = None) -> bool:
+async def set_db(key: str, data: dict | list, ttl: int | None = None, metadata: dict | None = None) -> bool:
"""Set KV."""
key = quote_plus(key)
success = False
if DB.ENGINE == "Cloudflare-KV":
- success = await set_cf_kv(key, data, ttl=None) # no expiration for CF-KV
+ success = await set_cf_kv(key, data, ttl=ttl)
if DB.ENGINE == "Cloudflare-R2":
- success = await set_cf_r2(key, data, metadata)
+ success = await set_cf_r2(key, data, metadata, ttl=ttl)
if success:
set_memory_kv(key, data, ttl)
return success
-def set_memory_kv(key: str, data: dict | list | str, ttl: int = 600) -> None:
+def set_memory_kv(key: str, data: dict | list | str, ttl: int | None = None) -> None:
"""Set to memory cache."""
+ if ttl is None:
+ ttl = 600
cache.set(key, data, ttl=ttl)
logger.trace(f"SET KV to memory cache for {key}: {data}")
@@ -182,7 +185,7 @@ async def set_cf_kv(key: str, data: dict | list | str, ttl: int | None = None, *
return False
-async def set_cf_r2(key: str, data: dict | list | str | None = None, metadata: dict | None = None, *, skip_in_memory: bool = True) -> bool:
+async def set_cf_r2(key: str, data: dict | list | str | None = None, metadata: dict | None = None, ttl: int | None = None, *, skip_in_memory: bool = True) -> bool:
"""Set to Cloudflare R2 via boto3.
We do not put data to R2, just use metadata to store data.
@@ -206,6 +209,8 @@ async def set_cf_r2(key: str, data: dict | list | str | None = None, metadata: d
payload |= {"Body": json.dumps(data).encode("utf-8")}
if metadata:
payload |= {"Metadata": stringfy(metadata)}
+ if ttl is not None:
+ payload |= {"Expires": nowdt() + timedelta(seconds=ttl)}
async with Session().client(
service_name="s3",
endpoint_url=f"https://{DB.CF_ACCOUNT_ID}.r2.cloudflarestorage.com",
@@ -434,8 +439,8 @@ async def delete_pastbin(url: str):
if __name__ == "__main__":
import asyncio
- url, manage_url = asyncio.run(upload_pastbin("测试.mp3", ttl=10))
- asyncio.run(delete_pastbin(manage_url))
+ # url, manage_url = asyncio.run(upload_pastbin("测试.mp3", ttl=10))
+ # asyncio.run(delete_pastbin(manage_url))
# asyncio.run(upload_uguu("测试.mp3"))
# asyncio.run(list_alist())
# asyncio.run(upload_alist("测试.mp3"))
@@ -443,4 +448,5 @@ if __name__ == "__main__":
# asyncio.run(delete_alist("测试.mp3"))
# asyncio.run(download_alist("test.py"))
# asyncio.run(set_cf_r2("test2", metadata={"finished": "1"}))
- # asyncio.run(set_cf_r2("test2", data={"finished": "1"}))
+ asyncio.run(set_cf_r2("test2", data={"finished": "1"}, ttl=60))
+ asyncio.run(get_cf_r2("test2"))