perf(api): optimize token read/write during api call

api/controllers/console/apikey.py

@@ -12,6 +12,7 @@
 from . import api
 from .setup import setup_required
 from .wraps import account_initialization_required
+from ..service_api.wraps import get_api_token_from_db
 
 api_key_fields = {
     "id": fields.String,
@@ -137,6 +138,25 @@ def after_request(self, resp):
         resp.headers["Access-Control-Allow-Credentials"] = "true"
         return resp
 
+    def delete(self, resource_id, api_key_id):
+        key = (
+            db.session.query(ApiToken)
+            .filter(
+                getattr(ApiToken, self.resource_id_field) == str(resource_id),
+                ApiToken.type == self.resource_type,
+                ApiToken.id == str(api_key_id),
+            )
+            .first()
+        )
+
+        response = super().delete(resource_id, api_key_id)
+
+        if response[1] == 204 and key:
+            from extensions.ext_redis import cache
+            cache.delete_memoized(get_api_token_from_db, key.token, "app")
+
+        return response
+
     resource_type = "app"
     resource_model = App
     resource_id_field = "app_id"
@@ -160,6 +180,25 @@ def after_request(self, resp):
         resp.headers["Access-Control-Allow-Credentials"] = "true"
         return resp
 
+    def delete(self, resource_id, api_key_id):
+        key = (
+            db.session.query(ApiToken)
+            .filter(
+                getattr(ApiToken, self.resource_id_field) == str(resource_id),
+                ApiToken.type == self.resource_type,
+                ApiToken.id == str(api_key_id),
+            )
+            .first()
+        )
+
+        response = super().delete(resource_id, api_key_id)
+
+        if response[1] == 204 and key:
+            from extensions.ext_redis import cache
+            cache.delete_memoized(get_api_token_from_db, key.token, "dataset")
+
+        return response
+
     resource_type = "dataset"
     resource_model = Dataset
     resource_id_field = "dataset_id"

api/controllers/service_api/wraps.py

@@ -4,13 +4,15 @@
 from functools import wraps
 from typing import Optional
 
+from celery import shared_task
 from flask import current_app, request
 from flask_login import user_logged_in
 from flask_restful import Resource
 from pydantic import BaseModel
 from werkzeug.exceptions import Forbidden, Unauthorized
 
 from extensions.ext_database import db
+from extensions.ext_redis import cache
 from libs.login import _get_user
 from models.account import Account, Tenant, TenantAccountJoin, TenantStatus
 from models.model import ApiToken, App, EndUser
@@ -172,6 +174,26 @@ def decorated(*args, **kwargs):
     return decorator
 
 
+@cache.memoize(timeout=3600)
+def get_api_token_from_db(auth_token, scope):
+    return db.session.query(ApiToken).filter(
+        ApiToken.token == auth_token,
+        ApiToken.type == scope
+    ).first()
+
+
+@shared_task(queue="ops_trace", bind=True, max_retries=3)
+def update_token_last_used_at(self, token_id, last_used_time):
+    try:
+        with db.session.begin():
+            token = db.session.query(ApiToken).get(token_id)
+            if token:
+                token.last_used_at = last_used_time
+    except Exception as e:
+        # 如果失败，最多重试3次
+        self.retry(exc=e, countdown=60)  # 1分钟后重试
+
+
 def validate_and_get_api_token(scope=None):
     """
     Validate and get API token.
@@ -186,20 +208,12 @@ def validate_and_get_api_token(scope=None):
     if auth_scheme != "bearer":
         raise Unauthorized("Authorization scheme must be 'Bearer'")
 
-    api_token = (
-        db.session.query(ApiToken)
-        .filter(
-            ApiToken.token == auth_token,
-            ApiToken.type == scope,
-        )
-        .first()
-    )
+    api_token = get_api_token_from_db(auth_token, scope)
 
     if not api_token:
         raise Unauthorized("Access token is invalid")
 
-    api_token.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)
-    db.session.commit()
+    update_token_last_used_at.delay(api_token.id, datetime.now(timezone.utc).replace(tzinfo=None))
 
     return api_token
 

api/extensions/ext_redis.py

@@ -1,6 +1,7 @@
 import redis
 from redis.connection import Connection, SSLConnection
 from redis.sentinel import Sentinel
+from flask_caching import Cache
 
 from configs import dify_config
 
@@ -40,6 +41,7 @@ def __getattr__(self, item):
 
 
 redis_client = RedisClientWrapper()
+cache = Cache()
 
 
 def init_app(app):
@@ -71,6 +73,14 @@ def init_app(app):
         )
         master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
         redis_client.initialize(master)
+
+        cache_config = {
+            'CACHE_TYPE': 'redis',
+            'CACHE_REDIS_SENTINELS': app.config.get('REDIS_SENTINELS'),
+            'CACHE_REDIS_SENTINEL_MASTER': app.config.get('REDIS_SENTINEL_SERVICE_NAME'),
+            'CACHE_REDIS_SENTINEL_PASSWORD': app.config.get('REDIS_SENTINEL_PASSWORD'),
+            'CACHE_DEFAULT_TIMEOUT': 3600
+        }
     else:
         redis_params.update(
             {
@@ -82,4 +92,15 @@ def init_app(app):
         pool = redis.ConnectionPool(**redis_params)
         redis_client.initialize(redis.Redis(connection_pool=pool))
 
+        cache_config = {
+            'CACHE_TYPE': 'redis',
+            'CACHE_REDIS_HOST': app.config.get('REDIS_HOST'),
+            'CACHE_REDIS_PORT': app.config.get('REDIS_PORT'),
+            'CACHE_REDIS_PASSWORD': app.config.get('REDIS_PASSWORD'),
+            'CACHE_REDIS_DB': app.config.get('REDIS_DB'),
+            'CACHE_DEFAULT_TIMEOUT': 3600
+        }
+
+    cache.init_app(app, config=cache_config)
+
     app.extensions["redis"] = redis_client

api/pyproject.toml

@@ -192,6 +192,7 @@ zhipuai = "~2.1.5"
 # Related transparent dependencies with pinned version
 # required by main implementations
 ############################################################
+flask-caching = "^2.3.0"
 [tool.poetry.group.indirect.dependencies]
 kaleido = "0.2.1"
 rank-bm25 = "~0.2.2"