feat: 🎸 缓存ArkStore的id_token (#1382)

api/v1/views/auth.py

@@ -6,6 +6,7 @@
 from arkid.core.error import ErrorCode, ErrorDict
 from api.v1.schema.auth import *
 from django.http import HttpResponse, JsonResponse
+from arkid.core.tasks.celery import dispatch_task
 
 @api.post("/tenant/{tenant_id}/auth/", response=AuthOut, tags=['登录与注册'], auth=None)
 @operation(AuthOut, use_id=True)
@@ -22,6 +23,7 @@ def auth(request, tenant_id: str, event_tag: str, data: AuthIn):
 
     # 生成 token
     token = refresh_token(user)
+    dispatch_task.delay('async_get_arkstore_access_token', tenant.id.hex, token)
     netloc = request.get_host().split(':')[0]
 
     is_ip_addr = False

arkid/common/arkstore.py

@@ -16,18 +16,19 @@
 from arkid.extension.utils import import_extension, unload_extension, load_extension_apps, restart_celery
 from pathlib import Path
 from arkid.common.logger import logger
+from django.core.cache import cache
 
 
-arkid_saas_token_cache = {}
-
 def get_saas_token(tenant, token, use_cache=True):
     """
     获取saas平台token
     """
     # 缓存 saas_token
     key = (str(tenant.id) + '_' + str(token)).replace('-', '')
-    if use_cache and key in arkid_saas_token_cache:
-        return arkid_saas_token_cache[key]
+    saas_cache = cache.get(key)
+    if use_cache and saas_cache:
+        return saas_cache
+
     app = Application.objects.filter(name='arkid_saas', uuid = tenant.id).first()
     nonce = uuid.uuid4().hex
     params = {
@@ -55,11 +56,12 @@ def get_saas_token(tenant, token, use_cache=True):
 
     resp = requests.get(resp.url)
     if resp.status_code != 200:
-        arkid_saas_token_cache.pop(key, None)
+        cache.delete(key)
         raise Exception(f'Error get_saas_token: {resp.status_code}')
     resp = resp.json()
-    arkid_saas_token_cache[key] = (resp['token'], resp['tenant_id'], resp['tenant_slug'])
-    return arkid_saas_token_cache[key]
+    result = (resp['token'], resp['tenant_id'], resp['tenant_slug'])
+    cache.set(key, result, timeout=60*60)
+    return result
 
 
 def get_arkstore_access_token(tenant, token, use_cache=True):
@@ -77,7 +79,6 @@ def get_arkstore_access_token(tenant, token, use_cache=True):
             use_cache=False, local_tenant=tenant, local_token=token)
 
 
-arkstore_access_token_saas_cache = {}
 
 def get_arkstore_access_token_with_saas_token(saas_tenant_slug, saas_tenant_id, saas_token, 
         use_cache=True, local_tenant=None, local_token=None):
@@ -86,19 +87,20 @@ def get_arkstore_access_token_with_saas_token(saas_tenant_slug, saas_tenant_id,
     """
     # 缓存 idtoken
     key = (str(saas_tenant_id) + '_' + str(saas_token)).replace('-', '')
-    if use_cache and key in arkstore_access_token_saas_cache:
+    id_token = cache.get(key)
+    if use_cache and id_token:
         try:
-            payload = jwt.decode(arkstore_access_token_saas_cache[key], options={"verify_signature": False})
+            payload = jwt.decode(id_token, options={"verify_signature": False})
         except Exception:
-            arkstore_access_token_saas_cache.pop(key, None)
+            cache.delete(key)
             raise Exception("Unable to parse id_token")
         exp_dt = datetime.fromtimestamp(payload["exp"])
         expire_time = timezone.make_aware(exp_dt, timezone.get_default_timezone())
         now = timezone.localtime()
         if now <= expire_time:
-            return arkstore_access_token_saas_cache[key]
+            return id_token
         else:
-            arkstore_access_token_saas_cache.pop(key, None)
+            cache.delete(key)
             # id_token 过期后，重新获取saas_token和id_token
             if local_tenant and local_token:
                 saas_token, saas_tenant_id, saas_tenant_slug = get_saas_token(local_tenant, local_token, use_cache=False)
@@ -107,15 +109,16 @@ def get_arkstore_access_token_with_saas_token(saas_tenant_slug, saas_tenant_id,
     app_login_url = settings.ARKSTOER_URL + '/api/v1/login'
     resp = requests.get(app_login_url, params=params)
     if resp.status_code != 200:
-        arkstore_access_token_saas_cache.pop(key, None)
+        cache.delete(key)
         raise Exception(f'Error get_arkstore_access_token_with_saas_token: {resp.status_code}, url: {resp.url}')
     try:
         resp = resp.json()
     except:
         from urllib.parse import urlencode, unquote
         raise Exception(f'Error get_arkstore_access_token_with_saas_token: {resp.status_code}, url: {unquote(resp.url)}')
-    arkstore_access_token_saas_cache[key] = resp['access_token']
-    return arkstore_access_token_saas_cache[key] 
+    id_token = resp['access_token']
+    cache.set(key, id_token, timeout=60*60)
+    return id_token
 
 
 def get_arkstore_extensions(access_token, purchased=None, rented=False, type=None, offset=0, limit=10, extra_params={}):

arkid/extension/utils.py

@@ -152,7 +152,12 @@ def unload_extension(ext_dir: str) -> any:
     if not Path(ext_dir).exists():
         return
     ext_name = f'{Path(ext_dir).parent}.{Path(ext_dir).name}'
-    ext = importlib.import_module(ext_name)
+    try:
+        ext = importlib.import_module(ext_name)
+    except Exception as e:
+        logger.error(f"import_module {ext_name} failed: {str(e)}")
+        return
+
     if ext:
         ext.extension.stop()
     sys.modules.pop(ext_name, None)