perf: Optimizing API Key Authentication

Author: fit2cloud-chenyw

backend/alembic/versions/056_api_key_ddl.py

@@ -0,0 +1,39 @@
+"""056_api_key_ddl
+
+Revision ID: d9a5589fc00b
+Revises: 3d4bd2d673dc
+Create Date: 2025-12-23 13:41:26.705947
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = 'd9a5589fc00b'
+down_revision = '3d4bd2d673dc'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('sys_apikey',
+    sa.Column('access_key', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False),
+    sa.Column('secret_key', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False),
+    sa.Column('create_time', sa.BigInteger(), nullable=False),
+    sa.Column('uid', sa.BigInteger(), nullable=False),
+    sa.Column('status', sa.Boolean(), nullable=False),
+    sa.Column('id', sa.BigInteger(), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_sys_apikey_id'), 'sys_apikey', ['id'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_sys_apikey_id'), table_name='sys_apikey')
+    op.drop_table('sys_apikey')
+    # ### end Alembic commands ###

backend/apps/api.py

@@ -5,7 +5,7 @@
 from apps.data_training.api import data_training
 from apps.datasource.api import datasource, table_relation, recommended_problem
 from apps.mcp import mcp
-from apps.system.api import login, user, aimodel, workspace, assistant, parameter
+from apps.system.api import login, user, aimodel, workspace, assistant, parameter, apikey
 from apps.terminology.api import terminology
 from apps.settings.api import base
 
@@ -24,5 +24,6 @@
 api_router.include_router(mcp.router)
 api_router.include_router(table_relation.router)
 api_router.include_router(parameter.router)
+api_router.include_router(apikey.router)
 
 api_router.include_router(recommended_problem.router)

backend/apps/system/api/apikey.py

@@ -0,0 +1,60 @@
+
+from fastapi import APIRouter
+from sqlmodel import func, select
+from apps.system.crud.apikey_manage import clear_api_key_cache
+from apps.system.models.system_model import ApiKeyModel
+from apps.system.schemas.system_schema import ApikeyGridItem, ApikeyStatus
+from common.core.deps import CurrentUser, SessionDep
+from common.utils.time import get_timestamp
+import secrets
+
+router = APIRouter(tags=["system_apikey"], prefix="/system/apikey")
+
+@router.get("")
+async def grid(session: SessionDep, current_user: CurrentUser) -> list[ApikeyGridItem]:
+    query = select(ApiKeyModel).where(ApiKeyModel.uid == current_user.id).order_by(ApiKeyModel.create_time.desc())
+    return session.exec(query).all()
+
+@router.post("")
+async def create(session: SessionDep, current_user: CurrentUser):
+    count = session.exec(select(func.count()).select_from(ApiKeyModel)).one()
+    if count >= 5:
+        raise ValueError("Maximum of 5 API keys allowed")
+    access_key = secrets.token_urlsafe(16)
+    secret_key = secrets.token_urlsafe(32)
+    api_key = ApiKeyModel(
+        access_key=access_key,
+        secret_key=secret_key,
+        create_time=get_timestamp(),
+        uid=current_user.id,
+        status=True
+    )
+    session.add(api_key)
+    session.commit()
+    return api_key.id
+
+@router.put("/status")
+async def status(session: SessionDep, current_user: CurrentUser, dto: ApikeyStatus):
+    api_key = session.get(ApiKeyModel, dto.id)
+    if not api_key:
+        raise ValueError("API Key not found")
+    if api_key.uid != current_user.id:
+        raise PermissionError("No permission to modify this API Key")
+    if dto.status == api_key.status:
+        return
+    api_key.status = dto.status
+    await clear_api_key_cache(api_key.access_key)
+    session.add(api_key)
+    session.commit()
+
+@router.delete("/{id}")
+async def delete(session: SessionDep, current_user: CurrentUser, id: int):
+    api_key = session.get(ApiKeyModel, id)
+    if not api_key:
+        raise ValueError("API Key not found")
+    if api_key.uid != current_user.id:
+        raise PermissionError("No permission to delete this API Key")
+    await clear_api_key_cache(api_key.access_key)
+    session.delete(api_key)
+    session.commit()
+    

backend/apps/system/crud/apikey_manage.py

@@ -0,0 +1,17 @@
+
+from sqlmodel import select
+
+from apps.system.models.system_model import ApiKeyModel
+from apps.system.schemas.auth import CacheName, CacheNamespace
+from common.core.deps import SessionDep
+from common.core.sqlbot_cache import cache, clear_cache
+from common.utils.utils import SQLBotLogUtil
+
+@cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.ASK_INFO, keyExpression="access_key")
+async def get_api_key(session: SessionDep, access_key: str) -> ApiKeyModel | None:
+    query = select(ApiKeyModel).where(ApiKeyModel.access_key == access_key)
+    return session.exec(query).first()
+
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.ASK_INFO, keyExpression="access_key")
+async def clear_api_key_cache(access_key: str):
+     SQLBotLogUtil.info(f"Api key cache for [{access_key}] has been cleaned")

backend/apps/system/middleware/auth.py

@@ -7,7 +7,8 @@
 import jwt
 from sqlmodel import Session
 from starlette.middleware.base import BaseHTTPMiddleware
-from apps.system.models.system_model import AssistantModel
+from apps.system.crud.apikey_manage import get_api_key
+from apps.system.models.system_model import ApiKeyModel, AssistantModel
 from common.core.db import engine 
 from apps.system.crud.assistant import get_assistant_info, get_assistant_user
 from apps.system.crud.user import get_user_by_account, get_user_info
@@ -33,7 +34,15 @@ async def dispatch(self, request, call_next):
             return await call_next(request)
         assistantTokenKey = settings.ASSISTANT_TOKEN_KEY
         assistantToken = request.headers.get(assistantTokenKey)
+        askToken = request.headers.get("X-SQLBOT-ASK-TOKEN")
         trans = await get_i18n(request)
+        if askToken:
+            validate_pass, data = await self.validateAskToken(askToken, trans)
+            if validate_pass:
+                request.state.current_user = data
+                return await call_next(request)
+            message = trans('i18n_permission.authenticate_invalid', msg = data)
+            return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
         #if assistantToken and assistantToken.lower().startswith("assistant "):
         if assistantToken:
             validator: tuple[any] = await self.validateAssistant(assistantToken, trans)
@@ -62,6 +71,50 @@ async def dispatch(self, request, call_next):
     def is_options(self, request: Request):
         return request.method == "OPTIONS"
 
+    async def validateAskToken(self, askToken: Optional[str], trans: I18n):
+        if not askToken:
+            return False, f"Miss Token[X-SQLBOT-ASK-TOKEN]!"
+        schema, param = get_authorization_scheme_param(askToken)
+        if schema.lower() != "sk":
+            return False, f"Token schema error!"
+        try: 
+            payload = jwt.decode(
+                param, options={"verify_signature": False, "verify_exp": False}, algorithms=[security.ALGORITHM]
+            )
+            access_key = payload.get('access_key', None)
+            
+            if not access_key:
+                return False, f"Miss access_key payload error!"
+            with Session(engine) as session:
+                api_key_model = await get_api_key(session, access_key)
+                api_key_model = ApiKeyModel.model_validate(api_key_model) if api_key_model else None
+                if not api_key_model:
+                    return False, f"Invalid access_key!"
+                if not api_key_model.status:
+                    return False, f"Disabled access_key!"
+                payload = jwt.decode(
+                    param, api_key_model.secret_key, algorithms=[security.ALGORITHM]
+                )
+                uid = api_key_model.uid
+                session_user = await get_user_info(session = session, user_id = uid)
+                if not session_user:
+                    message = trans('i18n_not_exist', msg = trans('i18n_user.account'))
+                    raise Exception(message)
+                session_user = UserInfoDTO.model_validate(session_user)
+                if session_user.status != 1:
+                    message = trans('i18n_login.user_disable', msg = trans('i18n_concat_admin'))
+                    raise Exception(message)
+                if not session_user.oid or session_user.oid == 0:
+                    message = trans('i18n_login.no_associated_ws', msg = trans('i18n_concat_admin'))
+                    raise Exception(message)
+                return True, session_user
+        except Exception as e:
+            msg = str(e)
+            SQLBotLogUtil.exception(f"Token validation error: {msg}")
+            if 'expired' in msg:
+                return False, jwt.ExpiredSignatureError(trans('i18n_permission.token_expired')) 
+            return False, e
+    
     async def validateToken(self, token: Optional[str], trans: I18n):
         if not token:
             return False, f"Miss Token[{settings.TOKEN_KEY}]!"

backend/apps/system/models/system_model.py

@@ -67,4 +67,15 @@ class AuthenticationModel(SnowflakeBase, AuthenticationBaseModel, table=True):
     __tablename__ = "sys_authentication"
     create_time: Optional[int] = Field(default=0, sa_type=BigInteger())
     enable: bool = Field(default=False, nullable=False)
-    valid: bool = Field(default=False, nullable=False)
+    valid: bool = Field(default=False, nullable=False)
+    
+
+class ApiKeyBaseModel(SQLModel):
+    access_key: str = Field(max_length=255, nullable=False)
+    secret_key: str = Field(max_length=255, nullable=False)
+    create_time: int = Field(default=0, sa_type=BigInteger())
+    uid: int = Field(default=0,nullable=False, sa_type=BigInteger())
+    status: bool = Field(default=True, nullable=False)
+    
+class ApiKeyModel(SnowflakeBase, ApiKeyBaseModel, table=True):
+    __tablename__ = "sys_apikey"

backend/apps/system/schemas/auth.py

@@ -16,6 +16,7 @@ class CacheName(Enum):
     USER_INFO = "user:info"
     ASSISTANT_INFO = "assistant:info"
     ASSISTANT_DS = "assistant:ds"
+    ASK_INFO = "ask:info"
     def __str__(self):
         return self.value
 

backend/apps/system/schemas/system_schema.py

@@ -207,3 +207,13 @@ class AssistantUiSchema(BaseCreatorDTO):
     name: Optional[str] = None
     welcome: Optional[str] = None
     welcome_desc: Optional[str] = None
+
+class ApikeyStatus(BaseModel):
+    id: int = Field(description=f"{PLACEHOLDER_PREFIX}id")
+    status: bool = Field(description=f"{PLACEHOLDER_PREFIX}status")
+
+class ApikeyGridItem(BaseCreatorDTO):
+    access_key: str = Field(description=f"Access Key")
+    secret_key: str = Field(description=f"Secret Key")
+    status: bool = Field(description=f"{PLACEHOLDER_PREFIX}status")
+    create_time: int = Field(description=f"{PLACEHOLDER_PREFIX}create_time")

frontend/src/components/layout/Apikey.vue

@@ -1,5 +1,5 @@
 <script lang="ts" setup>
-import { computed, reactive, ref } from 'vue'
+import { computed, onMounted, reactive, ref } from 'vue'
 import { useI18n } from 'vue-i18n'
 import icon_warning_filled from '@/assets/svg/icon_info_colorful.svg'
 import icon_add_outlined from '@/assets/svg/icon_add_outlined.svg'
@@ -11,6 +11,8 @@ import icon_visible_outlined from '@/assets/embedded/icon_visible_outlined.svg'
 import { formatTimestamp } from '@/utils/date'
 import { useClipboard } from '@vueuse/core'
 import EmptyBackground from '@/views/dashboard/common/EmptyBackground.vue'
+import { request } from '@/utils/request'
+
 const { t } = useI18n()
 
 const limitCount = ref(5)
@@ -23,37 +25,13 @@ const state = reactive({
   tableData: [] as any,
 })
 
-state.tableData = [
-  {
-    access_key: 'fwafwafwafwaf',
-    secret_key: '1234567',
-    status: false,
-    create_time: 1766455902237,
-  },
-  {
-    access_key: 'asdasdasdasd',
-    secret_key: '987654321',
-    status: true,
-    create_time: 1766455902237,
-  },
-  {
-    access_key: 'asdasdasdasd',
-    secret_key: '987654321',
-    status: true,
-    create_time: 1766455902237,
-  },
-  {
-    access_key: 'asdasdasdasd',
-    secret_key: '987654321',
-    status: true,
-    create_time: 1766455902237,
-  },
-]
 const handleAdd = () => {
   if (triggerLimit.value) {
     return
   }
-  console.log('Add API Key')
+  request.post('/system/apikey', {}).then(() => {
+    loadGridData()
+  })
 }
 const pwd = ref('**********')
 const toApiDoc = () => {
@@ -63,14 +41,13 @@ const toApiDoc = () => {
 }
 
 const statusHandler = (row: any) => {
-  /* state.form = { ...row }
-  editTerm() */
   const param = {
     id: row.id,
     status: row.status,
   }
-  console.log(row, param)
-  // userApi.status(param)
+  request.put('/system/apikey/status', param).then(() => {
+    loadGridData()
+  })
 }
 const { copy } = useClipboard({ legacy: true })
 
@@ -90,18 +67,27 @@ const deleteHandler = (row: any) => {
     cancelButtonText: t('common.cancel'),
     customClass: 'confirm-no_icon',
     autofocus: false,
-  }).then(() => {
-    /* userApi.delete(row.id).then(() => {
-      multipleSelectionAll.value = multipleSelectionAll.value.filter((ele) => ele.id !== row.id)
-      ElMessage({
-        type: 'success',
-        message: t('dashboard.delete_success'),
-      })
-      search()
-    }) */
-    console.log('execute delete')
+    callback: (action: any) => {
+      if (action === 'confirm') {
+        request.delete(`/system/apikey/${row.id}`).then(() => {
+          loadGridData()
+          ElMessage({
+            type: 'success',
+            message: t('dashboard.delete_success'),
+          })
+        })
+      }
+    },
+  })
+}
+const loadGridData = () => {
+  request.get('/system/apikey').then((res: any) => {
+    state.tableData = res || []
   })
 }
+onMounted(() => {
+  loadGridData()
+})
 </script>
 
 <template>