Commit cd235ec

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-15 07:13:07
feat(database): add `TTL` for CF-R2
1 parent b6476db
Changed files (1)
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"))