feat: add tidb

.gitignore

@@ -9,6 +9,7 @@ data/
 test_deploy_flow.py
 
 .Python
+hammer_tidb.py
 build/
 develop-eggs/
 dist/

app/api/server.py

@@ -7,16 +7,17 @@ from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 
+from app.api.v1.router import api_v1_router
+from app.api.v1.system import is_maintenance
 from app.config import get_settings
 from app.core.auth.deps import init_auth_db
 from app.core.database import pool_manager
 from app.core.logger import get_logger
-from app.core.redis_client import create_redis_client, close_redis
+from app.core.redis_client import close_redis, create_redis_client
 from app.core.scripts import load_scripts
+from app.core.tidb_manager import TiDBWriteError
 from app.services.embeddings_service import EmbeddingService
 from app.services.vector_store_service import VectorStoreService
-from app.api.v1.router import api_v1_router
-from app.api.v1.system import is_maintenance
 
 _logger = get_logger(__name__)
 _settings = get_settings()
@@ -43,13 +44,27 @@ async def _self_ping():
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    _logger.info("Initializing authentication database...")
-    await init_auth_db()
-    _logger.info("Authentication database initialized")
-
-    _logger.info("Initializing vector store database...")
-    await _vector_store_service.init_db()
-    _logger.info("Vector store database initialized with %d stores", len(_vector_store_service.list_stores()))
+    if not _settings.tidb_instance_list:
+        _logger.error(
+            "No TiDB instances configured! Set TIDB_INSTANCES environment variable "
+            "with comma-separated MySQL connection strings."
+        )
+    else:
+        _logger.info(
+            "Initializing TiDB databases across %d instance(s)...",
+            len(_settings.tidb_instance_list),
+        )
+        await init_auth_db()
+        _logger.info("Authentication database initialized on all TiDB instances")
+
+        _logger.info("Initializing vector store database...")
+        from app.core.vector_store.deps import init_vector_store_db
+        await init_vector_store_db()
+        await _vector_store_service.init_db()
+        _logger.info(
+            "Vector store database initialized with %d stores",
+            len(_vector_store_service.list_stores()),
+        )
 
     _logger.info("Initializing embedding service (loading 384-dim model)...")
     loop = asyncio.get_running_loop()
@@ -72,6 +87,11 @@ async def lifespan(app: FastAPI):
     await close_redis(redis)
     await _vector_store_service.close_all()
     await pool_manager.close_all()
+    from app.core.tidb_manager import get_tidb_manager
+    mgr = get_tidb_manager()
+    if mgr:
+        await mgr.close_all_pools()
+        _logger.info("TiDB connection pools closed")
 
 
 def create_application() -> FastAPI:
@@ -101,6 +121,18 @@ def create_application() -> FastAPI:
         allow_headers=["*"],
     )
 
+    @app.exception_handler(TiDBWriteError)
+    async def tidb_write_error_handler(request: Request, exc: TiDBWriteError):
+        from starlette.responses import JSONResponse
+        return JSONResponse(
+            status_code=503,
+            content={
+                "success": False,
+                "detail": "Failed to process request due to database storage limit. "
+                          "Please try again later or contact support.",
+            },
+        )
+
     @app.middleware("http")
     async def maintenance_middleware(request: Request, call_next):
         if is_maintenance():
@@ -109,7 +141,7 @@ def create_application() -> FastAPI:
             if method not in ("GET", "HEAD", "OPTIONS"):
                 if not path.startswith("/api/v1/maintenance"):
                     if method == "POST" and path.startswith("/api/v1/backup"):
-                        pass  # allow backup/restore during maintenance
+                        pass
                     else:
                         from starlette.responses import JSONResponse
                         return JSONResponse(

app/api/v1/auth.py

@@ -3,10 +3,9 @@ from __future__ import annotations
 from typing import Annotated
 
 from fastapi import APIRouter, Depends, Request, status
-from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.config import get_settings
-from app.core.auth.deps import get_current_user, get_db, get_temp_db_warning, require_application_id
+from app.core.auth.deps import get_current_user, get_db, require_application_id
 from app.core.auth.models import User
 from app.core.auth.schemas import (
     ChangePasswordSchema,
@@ -25,142 +24,134 @@ from app.core.auth.schemas import (
     UserSchemaField,
     UserSchemaResponse,
 )
+from app.core.tidb_manager import TiDBManager
 from app.services.auth_service import AuthService
 
 router = APIRouter(prefix="/auth", tags=["Authentication"])
 _settings = get_settings()
 
 
-def _build(data, warning):
+def _build(data, application_id=None):
     result = {"success": True, "data": data}
-    if _settings.application_id:
-        result["application_id"] = _settings.application_id
-    if warning:
-        result["warning"] = warning
+    if application_id:
+        result["application_id"] = application_id
     return result
 
 
 @router.post("/register", response_model=ProfileResponse, status_code=status.HTTP_201_CREATED)
 async def register(
     schema: RegisterSchema,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     user = await AuthService.register(db, schema)
-    warning = get_temp_db_warning()
-    return _build(AuthService.user_to_profile(user).model_dump(), warning)
+    profile = await AuthService.user_to_profile(db, user)
+    return _build(profile.model_dump(), _settings.application_id or None)
 
 
 @router.post("/login", response_model=TokenResponse)
 async def login(
     schema: LoginSchema,
     request: Request,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     tokens = await AuthService.login(db, request, schema)
-    warning = get_temp_db_warning()
-    return _build(tokens.model_dump(), warning)
+    return _build(tokens.model_dump(), _settings.application_id or None)
 
 
 @router.post("/refresh", response_model=TokenResponse)
 async def refresh_token(
     schema: TokenRefreshSchema,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     tokens = await AuthService.refresh(db, schema.refresh_token)
-    warning = get_temp_db_warning()
-    return _build(tokens.model_dump(), warning)
+    return _build(tokens.model_dump(), _settings.application_id or None)
 
 
 @router.post("/logout", response_model=MessageDataResponse)
 async def logout(
     schema: TokenRefreshSchema,
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.logout(db, current_user, schema.refresh_token)
-    warning = get_temp_db_warning()
-    return _build({"message": "Logged out successfully"}, warning)
+    return _build({"message": "Logged out successfully"}, _settings.application_id or None)
 
 
 @router.post("/logout-all", response_model=MessageDataResponse)
 async def logout_all(
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.logout_all(db, current_user)
-    warning = get_temp_db_warning()
-    return _build({"message": "All sessions revoked"}, warning)
+    return _build({"message": "All sessions revoked"}, _settings.application_id or None)
 
 
 @router.post("/forgot-password", response_model=MessageDataResponse)
 async def forgot_password(
     schema: ForgotPasswordSchema,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.forgot_password(db, schema)
-    warning = get_temp_db_warning()
-    return _build({"message": "If that email exists, a password reset link has been sent."}, warning)
+    return _build({"message": "If that email exists, a password reset link has been sent."}, _settings.application_id or None)
 
 
 @router.post("/reset-password", response_model=MessageDataResponse)
 async def reset_password(
     schema: ResetPasswordSchema,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.reset_password(db, schema)
-    warning = get_temp_db_warning()
-    return _build({"message": "Password reset successfully"}, warning)
+    return _build({"message": "Password reset successfully"}, _settings.application_id or None)
 
 
 @router.post("/change-password", response_model=MessageDataResponse)
 async def change_password(
     schema: ChangePasswordSchema,
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.change_password(db, current_user, schema)
-    warning = get_temp_db_warning()
-    return _build({"message": "Password changed successfully"}, warning)
+    return _build({"message": "Password changed successfully"}, _settings.application_id or None)
 
 
 @router.get("/me", response_model=ProfileResponse)
 async def get_me(
     current_user: Annotated[User, Depends(get_current_user)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
-    warning = get_temp_db_warning()
-    return _build(AuthService.user_to_profile(current_user).model_dump(), warning)
+    profile = await AuthService.user_to_profile(db, current_user)
+    return _build(profile.model_dump(), _settings.application_id or None)
 
 
 @router.patch("/me", response_model=ProfileResponse)
 async def update_me(
     schema: UpdateProfileSchema,
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     user = await AuthService.update_profile(db, current_user, schema)
-    warning = get_temp_db_warning()
-    return _build(AuthService.user_to_profile(user).model_dump(), warning)
+    profile = await AuthService.user_to_profile(db, user)
+    return _build(profile.model_dump(), _settings.application_id or None)
 
 
 @router.delete("/me", response_model=MessageDataResponse)
 async def delete_me(
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.soft_delete(db, current_user)
-    warning = get_temp_db_warning()
-    return _build({"message": "Account deleted successfully"}, warning)
+    return _build({"message": "Account deleted successfully"}, _settings.application_id or None)
 
 
 @router.get("/schema", response_model=SchemaResponse)
@@ -185,28 +176,25 @@ async def get_user_schema(
         UserSchemaField(field="deleted_at", type="datetime (ISO 8601)", required=False, description="Soft delete timestamp", constraints="Nullable"),
         UserSchemaField(field="roles", type="array[string]", required=False, description="Assigned role names", constraints="Via user_roles association table"),
     ]
-    warning = get_temp_db_warning()
-    return _build(UserSchemaResponse(table_name="users", columns=columns).model_dump(), warning)
+    return _build(UserSchemaResponse(table_name="users", columns=columns).model_dump(), _settings.application_id or None)
 
 
 @router.get("/sessions", response_model=SessionListResponse)
 async def list_sessions(
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     sessions = await AuthService.list_sessions(db, current_user)
-    warning = get_temp_db_warning()
-    return _build([SessionOut.model_validate(s).model_dump() for s in sessions], warning)
+    return _build([SessionOut.model_validate(s).model_dump() for s in sessions], _settings.application_id or None)
 
 
 @router.delete("/sessions/{session_id}", response_model=MessageDataResponse)
 async def revoke_session(
     session_id: str,
     current_user: Annotated[User, Depends(get_current_user)],
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
     _: Annotated[bool, Depends(require_application_id)],
 ):
     await AuthService.revoke_session(db, current_user, session_id)
-    warning = get_temp_db_warning()
-    return _build({"message": "Session revoked successfully"}, warning)
+    return _build({"message": "Session revoked successfully"}, _settings.application_id or None)

app/api/v1/batch.py

@@ -9,7 +9,12 @@ from urllib.parse import urlparse
 import httpx
 from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile
 
-from app.api.deps import get_converter_service, get_extraction_service, get_text_cleaner_service, require_auth
+from app.api.deps import (
+    get_converter_service,
+    get_extraction_service,
+    get_text_cleaner_service,
+    require_auth,
+)
 from app.api.v1.convert import _build_metadata, _thread_pool
 from app.config import get_settings
 from app.core.logger import get_logger

app/api/v1/chat.py

@@ -9,7 +9,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.responses import StreamingResponse
 
 from app.api.deps import require_auth
-from app.services.chat_service import chat_completion, _stream_chat_completion
+from app.services.chat_service import _stream_chat_completion, chat_completion
 
 logger = logging.getLogger(__name__)
 

app/api/v1/convert.py

@@ -9,16 +9,25 @@ from typing import Annotated, Any, Dict, Optional
 from urllib.parse import urlparse
 
 import httpx
+from fastapi import (
+    APIRouter,
+    Depends,
+    File,
+    Form,
+    HTTPException,
+    Query,
+    UploadFile,
+    status,
+)
 from fastapi.responses import PlainTextResponse
-from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile, status
 
-from app.config import get_settings
 from app.api.deps import (
     get_converter_service,
     get_extraction_service,
     get_text_cleaner_service,
     require_auth,
 )
+from app.config import get_settings
 from app.core.logger import get_logger
 from app.models.domain import ConversionError, count_tokens
 from app.models.schemas import ConversionMetadata, ConversionResponse, UrlRequest

app/api/v1/csv_analysis.py

@@ -3,15 +3,13 @@ from __future__ import annotations
 import json
 from typing import Annotated, Any, Dict, List, Optional
 
-from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile
-from pydantic import BaseModel, Field, ValidationError
+from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
+from pydantic import BaseModel, ValidationError
 
 from app.api.deps import require_auth
 from app.config import get_settings
 from app.services.chat_service import chat_completion
 from app.services.csv_analysis_service import (
-    analyze_csv_dataset,
-    create_csv_chart,
     execute_csv_chat_blocks,
     get_dataset_info,
 )

app/api/v1/database.py

@@ -6,7 +6,12 @@ from fastapi import APIRouter, Depends, HTTPException, status
 
 from app.api.deps import require_auth
 from app.core.logger import get_logger
-from app.models.schemas import DatabaseQueryRequest, DatabaseQueryResponse, DatabaseValidateRequest, DatabaseValidateResponse
+from app.models.schemas import (
+    DatabaseQueryRequest,
+    DatabaseQueryResponse,
+    DatabaseValidateRequest,
+    DatabaseValidateResponse,
+)
 from app.services.database_service import DatabaseService
 
 router = APIRouter()

app/api/v1/embeddings.py

@@ -4,9 +4,10 @@ import asyncio
 import concurrent.futures
 import os
 import time
+
 from fastapi import APIRouter, Depends, HTTPException
 
-from app.api.deps import require_auth, get_embeddings_service
+from app.api.deps import get_embeddings_service, require_auth
 from app.config import get_settings
 from app.core.logger import get_logger
 from app.models.schemas import EmbeddingItem, EmbeddingRequest, EmbeddingResponse

app/api/v1/qr_generator.py

@@ -6,7 +6,7 @@ from typing import List, Optional
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, Field, field_validator
 
-from app.services.qr_generator_service import QRGeneratorService, QRGeneratorError
+from app.services.qr_generator_service import QRGeneratorError, QRGeneratorService
 
 router = APIRouter()
 

app/api/v1/router.py

@@ -2,7 +2,28 @@ from __future__ import annotations
 
 from fastapi import APIRouter
 
-from app.api.v1 import auth, batch, chat, code_executor, convert, csv_analysis, database, embeddings, qr_generator, reconcile, scraper, semantic_router, sql_validator, system, token_counter, token_generator, url_shortener, vector_stores, web_search, webhook_socket
+from app.api.v1 import (
+    auth,
+    batch,
+    chat,
+    code_executor,
+    convert,
+    csv_analysis,
+    database,
+    embeddings,
+    qr_generator,
+    reconcile,
+    scraper,
+    semantic_router,
+    sql_validator,
+    system,
+    token_counter,
+    token_generator,
+    url_shortener,
+    vector_stores,
+    web_search,
+    webhook_socket,
+)
 from app.api.verify import router as verify_router
 
 api_v1_router = APIRouter()

app/api/v1/url_shortener.py

@@ -9,7 +9,6 @@ from pydantic import BaseModel, Field
 
 from app.config import get_settings
 from app.services.url_shortener_service import (
-    URLShortenerService,
     AliasTakenError,
     AuthenticationError,
     AuthorizationError,
@@ -17,6 +16,7 @@ from app.services.url_shortener_service import (
     LinkNotFoundError,
     PlanLimitExceededError,
     URLShortenerError,
+    URLShortenerService,
     ValidationError,
 )
 

app/api/v1/vector_stores.py

@@ -4,10 +4,18 @@ import asyncio
 import time
 
 import httpx
-from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile, status
+from fastapi import (
+    APIRouter,
+    Depends,
+    File,
+    Form,
+    HTTPException,
+    Query,
+    UploadFile,
+    status,
+)
 
 from app.api.deps import get_vector_store_service, require_auth
-from app.core.auth.deps import get_temp_db_warning
 from app.core.logger import get_logger
 from app.models.domain import ConversionError
 from app.models.schemas import (
@@ -72,7 +80,6 @@ async def create_vector_store(
         document_count=stats["document_count"],
         created_at=stats["created_at"],
         metadata=stats["metadata"],
-        warning=get_temp_db_warning(),
     )
 
 

app/api/v1/webhook_socket.py

@@ -4,7 +4,14 @@ import asyncio
 import hmac as hmac_mod
 import json
 
-from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, WebSocketDisconnect
+from fastapi import (
+    APIRouter,
+    Depends,
+    HTTPException,
+    Request,
+    WebSocket,
+    WebSocketDisconnect,
+)
 
 from app.api.deps import require_auth
 from app.models.schemas import (

app/config.py

@@ -73,7 +73,8 @@ class Settings(BaseSettings):
     jwt_default_expiry_minutes: int = 30
     jwt_issuer: str = "all-api-collection"
 
-    database_url: str = "sqlite+aiosqlite:///data/auth.db"
+    tidb_instances: str = ""
+    tidb_ssl: bool = True
     access_token_expire_minutes: int = 15
     refresh_token_expire_days: int = 7
     max_login_attempts: int = 5
@@ -81,6 +82,10 @@ class Settings(BaseSettings):
     application_id: str = ""
     admin_password: str = ""
 
+    @property
+    def tidb_instance_list(self) -> list[str]:
+        return [u.strip() for u in self.tidb_instances.split(",") if u.strip()]
+
     @property
     def max_upload_mb(self) -> int:
         return self.max_upload_bytes // (1024 * 1024)

app/core/auth/__init__.py

@@ -1,27 +1,28 @@
 from __future__ import annotations
 
-from app.core.auth.models import (
-    Base, Permission, Role, User, RefreshSession,
-    user_roles, role_permissions,
-)
+from app.core.auth.deps import get_current_user, get_db, require_permissions
+from app.core.auth.models import Permission, RefreshSession, Role, User
 from app.core.auth.schemas import (
-    RegisterSchema, LoginSchema, TokenResponse, TokenRefreshSchema,
-    UserProfile, UpdateProfileSchema, ChangePasswordSchema,
-    ForgotPasswordSchema, ResetPasswordSchema, SessionOut,
-    MessageResponse, UserSchemaResponse, UserSchemaField,
-)
-from app.core.auth.deps import (
-    get_db, get_current_user, require_permissions,
-    get_temp_db_warning, TempDatabaseWarning,
+    ChangePasswordSchema,
+    ForgotPasswordSchema,
+    LoginSchema,
+    MessageResponse,
+    RegisterSchema,
+    ResetPasswordSchema,
+    SessionOut,
+    TokenRefreshSchema,
+    TokenResponse,
+    UpdateProfileSchema,
+    UserProfile,
+    UserSchemaField,
+    UserSchemaResponse,
 )
 
 __all__ = [
-    "Base", "Permission", "Role", "User", "RefreshSession",
-    "user_roles", "role_permissions",
+    "User", "RefreshSession", "Role", "Permission",
     "RegisterSchema", "LoginSchema", "TokenResponse", "TokenRefreshSchema",
     "UserProfile", "UpdateProfileSchema", "ChangePasswordSchema",
     "ForgotPasswordSchema", "ResetPasswordSchema", "SessionOut",
     "MessageResponse", "UserSchemaResponse", "UserSchemaField",
     "get_db", "get_current_user", "require_permissions",
-    "get_temp_db_warning", "TempDatabaseWarning",
 ]

app/core/auth/deps.py

@@ -1,117 +1,211 @@
 from __future__ import annotations
 
 import logging
-from typing import Annotated, AsyncGenerator, Optional
+from datetime import datetime
+from typing import Annotated, Optional
 
 import jwt
 from argon2 import PasswordHasher
 from fastapi import Depends, HTTPException, Request, status
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 
 from app.config import get_settings
-from app.core.auth.models import Base, Permission, Role, User
+from app.core.auth.models import RefreshSession, User, _uuid
+from app.core.auth.models import _utcnow as _now
+from app.core.tidb_manager import TiDBManager, get_tidb_manager, set_tidb_manager
 
 logger = logging.getLogger("auth")
 
 _settings = get_settings()
-
-_is_temp_db = (
-    "sqlite" in _settings.database_url
-    and "localhost" not in _settings.database_url
-    and "postgres" not in _settings.database_url.lower()
-    and "mysql" not in _settings.database_url.lower()
-)
-
-_raw_url: str = getattr(_settings, "database_url", None) or "sqlite+aiosqlite:///data/auth.db"
-
-if "sqlite" in _raw_url and "sqlite+aiosqlite" not in _raw_url:
-    _raw_url = _raw_url.replace("sqlite:///", "sqlite+aiosqlite:///")
-
-DATABASE_URL: str = _raw_url
-
-engine = create_async_engine(DATABASE_URL, echo=False)
-AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
-
 ph = PasswordHasher()
 
+AUTH_TABLES: dict[str, str] = {
+    "users": """
+        CREATE TABLE IF NOT EXISTS users (
+            id VARCHAR(36) PRIMARY KEY,
+            email VARCHAR(255) NOT NULL,
+            username VARCHAR(50) NULL,
+            full_name VARCHAR(255) NULL,
+            password_hash TEXT NOT NULL,
+            is_active TINYINT(1) NOT NULL DEFAULT 1,
+            is_verified TINYINT(1) NOT NULL DEFAULT 0,
+            failed_login_attempts INT NOT NULL DEFAULT 0,
+            locked_until DATETIME NULL,
+            last_login DATETIME NULL,
+            password_changed_at DATETIME NULL,
+            created_at DATETIME NOT NULL,
+            updated_at DATETIME NOT NULL,
+            deleted_at DATETIME NULL,
+            UNIQUE KEY uk_users_email (email),
+            UNIQUE KEY uk_users_username (username),
+            INDEX idx_users_email (email),
+            INDEX idx_users_username (username)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "roles": """
+        CREATE TABLE IF NOT EXISTS roles (
+            id VARCHAR(36) PRIMARY KEY,
+            name VARCHAR(50) NOT NULL,
+            description VARCHAR(255) NULL,
+            UNIQUE KEY uk_roles_name (name)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "permissions": """
+        CREATE TABLE IF NOT EXISTS permissions (
+            id VARCHAR(36) PRIMARY KEY,
+            code VARCHAR(100) NOT NULL,
+            description VARCHAR(255) NULL,
+            UNIQUE KEY uk_permissions_code (code),
+            INDEX idx_permissions_code (code)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "user_roles": """
+        CREATE TABLE IF NOT EXISTS user_roles (
+            user_id VARCHAR(36) NOT NULL,
+            role_id VARCHAR(36) NOT NULL,
+            PRIMARY KEY (user_id, role_id),
+            CONSTRAINT fk_ur_user FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
+            CONSTRAINT fk_ur_role FOREIGN KEY (role_id) REFERENCES roles(id) ON DELETE CASCADE
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "role_permissions": """
+        CREATE TABLE IF NOT EXISTS role_permissions (
+            role_id VARCHAR(36) NOT NULL,
+            permission_id VARCHAR(36) NOT NULL,
+            PRIMARY KEY (role_id, permission_id),
+            CONSTRAINT fk_rp_role FOREIGN KEY (role_id) REFERENCES roles(id) ON DELETE CASCADE,
+            CONSTRAINT fk_rp_perm FOREIGN KEY (permission_id) REFERENCES permissions(id) ON DELETE CASCADE
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "refresh_sessions": """
+        CREATE TABLE IF NOT EXISTS refresh_sessions (
+            id VARCHAR(36) PRIMARY KEY,
+            user_id VARCHAR(36) NOT NULL,
+            token_key VARCHAR(64) NOT NULL,
+            token_hash VARCHAR(255) NOT NULL,
+            device_info VARCHAR(255) NULL,
+            ip_address VARCHAR(45) NULL,
+            expires_at DATETIME NOT NULL,
+            revoked_at DATETIME NULL,
+            created_at DATETIME NOT NULL,
+            updated_at DATETIME NOT NULL,
+            UNIQUE KEY uk_rs_token_key (token_key),
+            INDEX idx_rs_user_id (user_id),
+            INDEX idx_rs_token_key (token_key),
+            CONSTRAINT fk_rs_user FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+}
+
+
+def _create_manager() -> Optional[TiDBManager]:
+    urls = _settings.tidb_instance_list
+    if not urls:
+        logger.warning("No TiDB instances configured (TIDB_INSTANCES is empty)")
+        return None
+    import ssl
+    if _settings.tidb_ssl:
+        ctx = ssl.create_default_context()
+        ctx.check_hostname = False
+        ctx.verify_mode = ssl.CERT_NONE
+    else:
+        ctx = None
+    mgr = TiDBManager(urls, ssl=ctx)
+    return mgr
 
-class TempDatabaseWarning:
-    code: str = "TEMP_DATABASE"
-    message: str = (
-        "No external database configuration was provided. "
-        "The service is currently using its built-in SQLite database intended for development and temporary use. "
-        "Data stored in this database should not be considered permanent."
-    )
 
+async def init_auth_db():
+    mgr = _create_manager()
+    if mgr is None:
+        logger.error("Cannot initialize auth DB: no TiDB instances configured")
+        return
+    await mgr.initialize_pools()
+    set_tidb_manager(mgr)
+
+    await mgr.ensure_tables(AUTH_TABLES)
+
+    roles_result = await mgr.fetchone("SELECT id FROM roles WHERE name = 'SuperAdmin'")
+    if not roles_result:
+        sa_role_id = _uuid()
+        admin_role_id = _uuid()
+        user_role_id = _uuid()
+        perm_ids = {}
+        for code, desc in [
+            ("users:read", "Read users"),
+            ("users:write", "Modify users"),
+            ("users:delete", "Delete users"),
+            ("admin:access", "Access admin panel"),
+        ]:
+            pid = _uuid()
+            await mgr.execute(
+                "INSERT INTO permissions (id, code, description) VALUES (%s, %s, %s)",
+                (pid, code, desc),
+            )
+            perm_ids[code] = pid
 
-def get_temp_db_warning() -> Optional[dict]:
-    if _is_temp_db:
-        return {"code": TempDatabaseWarning.code, "message": TempDatabaseWarning.message}
-    return None
+        await mgr.execute(
+            "INSERT INTO roles (id, name, description) VALUES (%s, %s, %s)",
+            (sa_role_id, "SuperAdmin", "Full system access"),
+        )
+        await mgr.execute(
+            "INSERT INTO roles (id, name, description) VALUES (%s, %s, %s)",
+            (admin_role_id, "Admin", "Administrative access"),
+        )
+        await mgr.execute(
+            "INSERT INTO roles (id, name, description) VALUES (%s, %s, %s)",
+            (user_role_id, "User", "Standard user access"),
+        )
 
+        perms_list = list(perm_ids.values())
+        for pid in perms_list:
+            await mgr.execute(
+                "INSERT INTO role_permissions (role_id, permission_id) VALUES (%s, %s)",
+                (sa_role_id, pid),
+            )
+        for pid in perms_list[:-1]:
+            await mgr.execute(
+                "INSERT INTO role_permissions (role_id, permission_id) VALUES (%s, %s)",
+                (admin_role_id, pid),
+            )
+        await mgr.execute(
+            "INSERT INTO role_permissions (role_id, permission_id) VALUES (%s, %s)",
+            (user_role_id, perms_list[0]),
+        )
 
-async def init_auth_db():
-    async with engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
-
-    async with AsyncSessionLocal() as session:
-        result = await session.execute(select(Role).where(Role.name == "SuperAdmin"))
-        if not result.scalars().first():
-            sa_role = Role(name="SuperAdmin", description="Full system access")
-            admin_role = Role(name="Admin", description="Administrative access")
-            user_role = Role(name="User", description="Standard user access")
-            session.add_all([sa_role, admin_role, user_role])
-
-            perms = []
-            for code, desc in [
-                ("users:read", "Read users"),
-                ("users:write", "Modify users"),
-                ("users:delete", "Delete users"),
-                ("admin:access", "Access admin panel"),
-            ]:
-                p = Permission(code=code, description=desc)
-                perms.append(p)
-                session.add(p)
-
-            sa_role.permissions.extend(perms)
-            admin_role.permissions.extend(perms[:-1])
-            user_role.permissions.append(perms[0])
-
-            await session.commit()
-
-        admin_email = "admin@example.com"
-        result = await session.execute(select(User).where(User.email == admin_email))
-        if not result.scalars().first():
-            admin_password = _settings.admin_password or "Admin123!"
-            admin_user = User(
-                email=admin_email,
-                full_name="System Administrator",
-                password_hash=ph.hash(admin_password),
-                is_verified=True,
+    admin_email = "admin@example.com"
+    admin_result = await mgr.fetchone("SELECT id FROM users WHERE email = %s", (admin_email,))
+    if not admin_result:
+        admin_password = _settings.admin_password or "Admin123!"
+        admin_id = _uuid()
+        now = _now()
+        await mgr.execute(
+            "INSERT INTO users (id, email, full_name, password_hash, is_verified, created_at, updated_at) "
+            "VALUES (%s, %s, %s, %s, 1, %s, %s)",
+            (admin_id, admin_email, "System Administrator", ph.hash(admin_password), now, now),
+        )
+        role_result = await mgr.fetchone("SELECT id FROM roles WHERE name = 'SuperAdmin'")
+        if role_result:
+            await mgr.execute(
+                "INSERT INTO user_roles (user_id, role_id) VALUES (%s, %s)",
+                (admin_id, role_result["id"]),
             )
-            role_result = await session.execute(select(Role).where(Role.name == "SuperAdmin"))
-            admin_role = role_result.scalars().first()
-            if admin_role:
-                admin_user.roles.append(admin_role)
-            session.add(admin_user)
-            await session.commit()
 
-        from app.services.auth_service import AuthService
-        await AuthService.cleanup_expired_sessions(session)
+    from app.services.auth_service import AuthService
+    await AuthService.cleanup_expired_sessions(mgr)
 
 
-async def get_db() -> AsyncGenerator[AsyncSession, None]:
-    async with AsyncSessionLocal() as session:
-        try:
-            yield session
-        finally:
-            await session.close()
+async def get_db() -> TiDBManager:
+    mgr = get_tidb_manager()
+    if mgr is None:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Database not initialized. Configure TIDB_INSTANCES.",
+        )
+    return mgr
 
 
 async def get_current_user(
     request: Request,
-    db: Annotated[AsyncSession, Depends(get_db)],
+    db: Annotated[TiDBManager, Depends(get_db)],
 ) -> User:
     credentials_exception = HTTPException(
         status_code=status.HTTP_401_UNAUTHORIZED,
@@ -134,9 +228,14 @@ async def get_current_user(
     except jwt.PyJWTError:
         raise credentials_exception
 
-    result = await db.execute(select(User).where(User.id == user_id, User.deleted_at.is_(None)))
-    user = result.scalars().first()
-    if user is None or not user.is_active:
+    row = await db.fetchone_all(
+        "SELECT * FROM users WHERE id = %s AND deleted_at IS NULL",
+        (user_id,),
+    )
+    if row is None:
+        raise credentials_exception
+    user = _row_to_user(row)
+    if not user.is_active:
         raise credentials_exception
     return user
 
@@ -164,10 +263,16 @@ def require_permissions(*required_perms: str):
         user: Annotated[User, Depends(get_current_user)],
     ) -> User:
         user_perms = set()
-        for role in user.roles:
-            for perm in role.permissions:
-                user_perms.add(perm.code)
-
+        mgr = get_tidb_manager()
+        if mgr:
+            rows = await mgr.fetchall_all(
+                "SELECT p.code FROM permissions p "
+                "JOIN role_permissions rp ON rp.permission_id = p.id "
+                "JOIN user_roles ur ON ur.role_id = rp.role_id "
+                "WHERE ur.user_id = %s",
+                (user.id,),
+            )
+            user_perms = {r["code"] for r in rows}
         missing = [p for p in required_perms if p not in user_perms]
         if missing:
             raise HTTPException(
@@ -177,3 +282,53 @@ def require_permissions(*required_perms: str):
         return user
 
     return permission_checker
+
+
+def _row_to_user(row: dict) -> User:
+    return User(
+        id=row["id"],
+        email=row["email"],
+        username=row.get("username"),
+        full_name=row.get("full_name"),
+        password_hash=row["password_hash"],
+        is_active=bool(row["is_active"]),
+        is_verified=bool(row["is_verified"]),
+        failed_login_attempts=row["failed_login_attempts"],
+        locked_until=_deserialize_dt(row.get("locked_until")),
+        last_login=_deserialize_dt(row.get("last_login")),
+        password_changed_at=_deserialize_dt(row.get("password_changed_at")),
+        created_at=_deserialize_dt(row["created_at"]),
+        updated_at=_deserialize_dt(row["updated_at"]),
+        deleted_at=_deserialize_dt(row.get("deleted_at")),
+    )
+
+
+def _row_to_refresh_session(row: dict) -> RefreshSession:
+    return RefreshSession(
+        id=row["id"],
+        user_id=row["user_id"],
+        token_key=row["token_key"],
+        token_hash=row["token_hash"],
+        device_info=row.get("device_info"),
+        ip_address=row.get("ip_address"),
+        expires_at=_deserialize_dt(row["expires_at"]),
+        revoked_at=_deserialize_dt(row.get("revoked_at")),
+        created_at=_deserialize_dt(row["created_at"]),
+        updated_at=_deserialize_dt(row["updated_at"]),
+    )
+
+
+def _serialize_dt(dt: Optional[datetime]) -> Optional[str]:
+    if dt is None:
+        return None
+    return dt.strftime("%Y-%m-%d %H:%M:%S.%f")
+
+
+def _deserialize_dt(val) -> Optional[datetime]:
+    if val is None:
+        return None
+    if isinstance(val, datetime):
+        return val
+    if isinstance(val, str):
+        return datetime.fromisoformat(val)
+    return val

app/core/auth/models.py

@@ -1,16 +1,9 @@
 from __future__ import annotations
 
 import uuid
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
-
-from sqlalchemy import (
-    Boolean, Column, DateTime, ForeignKey, Integer, String, Text, Table,
-)
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
-
-
-class Base(DeclarativeBase):
-    pass
+from typing import Optional
 
 
 def _utcnow() -> datetime:
@@ -21,80 +14,49 @@ def _uuid() -> str:
     return str(uuid.uuid4())
 
 
-user_roles = Table(
-    "user_roles",
-    Base.metadata,
-    Column("user_id", String(36), ForeignKey("users.id", ondelete="CASCADE"), primary_key=True),
-    Column("role_id", String(36), ForeignKey("roles.id", ondelete="CASCADE"), primary_key=True),
-)
-
-role_permissions = Table(
-    "role_permissions",
-    Base.metadata,
-    Column("role_id", String(36), ForeignKey("roles.id", ondelete="CASCADE"), primary_key=True),
-    Column("permission_id", String(36), ForeignKey("permissions.id", ondelete="CASCADE"), primary_key=True),
-)
-
-
-class Permission(Base):
-    __tablename__ = "permissions"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid)
-    code: Mapped[str] = mapped_column(String(100), unique=True, index=True, nullable=False)
-    description: Mapped[str | None] = mapped_column(String(255), nullable=True)
-
-    roles: Mapped[list[Role]] = relationship(secondary=role_permissions, back_populates="permissions", lazy="selectin")
-
-
-class Role(Base):
-    __tablename__ = "roles"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid)
-    name: Mapped[str] = mapped_column(String(50), unique=True, nullable=False)
-    description: Mapped[str | None] = mapped_column(String(255), nullable=True)
-
-    users: Mapped[list[User]] = relationship(secondary=user_roles, back_populates="roles", lazy="selectin")
-    permissions: Mapped[list[Permission]] = relationship(secondary=role_permissions, back_populates="roles", lazy="selectin")
-
-
-class User(Base):
-    __tablename__ = "users"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid)
-    email: Mapped[str] = mapped_column(String(255), unique=True, index=True, nullable=False)
-    username: Mapped[str | None] = mapped_column(String(50), unique=True, index=True, nullable=True)
-    full_name: Mapped[str | None] = mapped_column(String(255), nullable=True)
-    password_hash: Mapped[str] = mapped_column(Text, nullable=False)
-
-    is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
-    is_verified: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
-
-    failed_login_attempts: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
-    locked_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
-    last_login: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
-    password_changed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
-
-    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow, nullable=False)
-    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, nullable=False)
-    deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
-
-    roles: Mapped[list[Role]] = relationship(secondary=user_roles, back_populates="users", lazy="selectin")
-    sessions: Mapped[list[RefreshSession]] = relationship(back_populates="user", cascade="all, delete-orphan", lazy="selectin")
-
-
-class RefreshSession(Base):
-    __tablename__ = "refresh_sessions"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid)
-    user_id: Mapped[str] = mapped_column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True)
-    token_key: Mapped[str] = mapped_column(String(64), unique=True, index=True, nullable=False)
-    token_hash: Mapped[str] = mapped_column(String(255), nullable=False)
-    device_info: Mapped[str | None] = mapped_column(String(255), nullable=True)
-    ip_address: Mapped[str | None] = mapped_column(String(45), nullable=True)
-    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
-    revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
-
-    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow, nullable=False)
-    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, nullable=False)
-
-    user: Mapped[User] = relationship(back_populates="sessions")
+@dataclass
+class User:
+    id: str = field(default_factory=_uuid)
+    email: str = ""
+    username: Optional[str] = None
+    full_name: Optional[str] = None
+    password_hash: str = ""
+    is_active: bool = True
+    is_verified: bool = False
+    failed_login_attempts: int = 0
+    locked_until: Optional[datetime] = None
+    last_login: Optional[datetime] = None
+    password_changed_at: Optional[datetime] = None
+    created_at: datetime = field(default_factory=_utcnow)
+    updated_at: datetime = field(default_factory=_utcnow)
+    deleted_at: Optional[datetime] = None
+    roles: list[str] = field(default_factory=list)
+
+
+@dataclass
+class RefreshSession:
+    id: str = field(default_factory=_uuid)
+    user_id: str = ""
+    token_key: str = ""
+    token_hash: str = ""
+    device_info: Optional[str] = None
+    ip_address: Optional[str] = None
+    expires_at: Optional[datetime] = None
+    revoked_at: Optional[datetime] = None
+    created_at: datetime = field(default_factory=_utcnow)
+    updated_at: datetime = field(default_factory=_utcnow)
+
+
+@dataclass
+class Role:
+    id: str = field(default_factory=_uuid)
+    name: str = ""
+    description: Optional[str] = None
+    permissions: list[str] = field(default_factory=list)
+
+
+@dataclass
+class Permission:
+    id: str = field(default_factory=_uuid)
+    code: str = ""
+    description: Optional[str] = None

app/core/auth/schemas.py

@@ -117,8 +117,6 @@ class UserSchemaResponse(BaseModel):
     columns: list[UserSchemaField]
 
 
-# --- Wrapped API response schemas ---
-
 class ApiResponse(BaseModel):
     success: bool = True
 
@@ -126,34 +124,28 @@ class ApiResponse(BaseModel):
 class DataResponse(ApiResponse):
     data: dict
     application_id: Optional[str] = None
-    warning: Optional[dict] = None
 
 
 class TokenResponse(ApiResponse):
     data: TokenData
     application_id: Optional[str] = None
-    warning: Optional[dict] = None
 
 
 class ProfileResponse(ApiResponse):
     data: UserProfile
     application_id: Optional[str] = None
-    warning: Optional[dict] = None
 
 
 class MessageDataResponse(ApiResponse):
     data: MessageResponse
     application_id: Optional[str] = None
-    warning: Optional[dict] = None
 
 
 class SchemaResponse(ApiResponse):
     data: UserSchemaResponse
     application_id: Optional[str] = None
-    warning: Optional[dict] = None
 
 
 class SessionListResponse(ApiResponse):
     data: list[SessionOut]
     application_id: Optional[str] = None
-    warning: Optional[dict] = None

app/core/database/__init__.py

@@ -1,8 +1,8 @@
 from app.core.database.base import BaseExecutor, ConnectionConfig, StatementResult
-from app.core.database.mysql import MySQLExecutor
-from app.core.database.postgresql import PostgreSQLExecutor
 from app.core.database.mongodb import MongoDBExecutor
+from app.core.database.mysql import MySQLExecutor
 from app.core.database.pool import pool_manager
+from app.core.database.postgresql import PostgreSQLExecutor
 
 __all__ = [
     "BaseExecutor",

app/core/database/pool.py

@@ -3,9 +3,9 @@ from __future__ import annotations
 import asyncio
 
 from app.core.database.base import BaseExecutor, ConnectionConfig
+from app.core.database.mongodb import MongoDBExecutor
 from app.core.database.mysql import MySQLExecutor
 from app.core.database.postgresql import PostgreSQLExecutor
-from app.core.database.mongodb import MongoDBExecutor
 from app.core.logger import get_logger
 
 _logger = get_logger(__name__)

app/core/redis_client.py

@@ -5,7 +5,6 @@ from typing import Optional
 
 from redis.asyncio import Redis
 
-
 logger = logging.getLogger(__name__)
 
 

app/core/tidb_manager.py

@@ -0,0 +1,340 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import ssl
+from typing import Any, Dict, List, Optional, Set
+from urllib.parse import urlparse
+
+import aiomysql
+
+logger = logging.getLogger(__name__)
+
+# MySQL/TiDB error codes that indicate the instance cannot accept more writes:
+#   1114  (ER_RECORD_FILE_FULL)  – "The table '...' is full"
+#   1021  (ER_DISK_FULL)         – "Disk full (...)"
+#   126   (ER_INDEX_FILE_FULL)   – "Index file is full"
+#   1877  (ER_TABLE_FULL)        – "Table full"
+_STORAGE_FULL_CODES: Set[int] = {1114, 1021, 126, 1877}
+
+
+class TiDBWriteError(RuntimeError):
+    """All TiDB instances are full or unavailable — writes cannot be accepted."""
+    def __init__(self):
+        super().__init__("No TiDB instances available for write. All instances are full or unreachable.")
+
+
+def _is_storage_full_error(exc: Exception) -> bool:
+    if isinstance(exc, aiomysql.MySQLError):
+        code = exc.args[0] if exc.args else None
+        if isinstance(code, int) and code in _STORAGE_FULL_CODES:
+            return True
+        msg = str(exc).lower()
+        for keyword in ("table is full", "disk full", "storage full",
+                        "quota exceeded", "out of storage"):
+            if keyword in msg:
+                return True
+    return False
+
+
+class TiDBInstance:
+    def __init__(self, host: str, port: int, user: str, password: str, database: str,
+                 ssl: Optional[ssl.SSLContext] = None):
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        self.database = database
+        self.ssl = ssl
+        self.pool: Optional[aiomysql.Pool] = None
+        self.is_full = False
+        self.is_available = True
+
+    @classmethod
+    def from_url(cls, url: str,
+                 ssl: Optional[ssl.SSLContext] = None) -> TiDBInstance:
+        parsed = urlparse(url)
+        host = parsed.hostname or "localhost"
+        port = parsed.port or 4000
+        user = parsed.username or "root"
+        password = parsed.password or ""
+        database = parsed.path.lstrip("/") or "test"
+        return cls(host=host, port=port, user=user, password=password,
+                   database=database, ssl=ssl)
+
+    async def create_pool(self, minsize: int = 1, maxsize: int = 10):
+        self.pool = await aiomysql.create_pool(
+            host=self.host,
+            port=self.port,
+            user=self.user,
+            password=self.password,
+            db=self.database,
+            minsize=minsize,
+            maxsize=maxsize,
+            autocommit=False,
+            cursorclass=aiomysql.DictCursor,
+            ssl=self.ssl,
+        )
+        self.is_available = True
+        return self.pool
+
+    async def close_pool(self):
+        if self.pool:
+            await self.pool.close()
+            await self.pool.wait_closed()
+            self.pool = None
+
+    def __repr__(self) -> str:
+        return f"TiDBInstance({self.user}@{self.host}:{self.port}/{self.database})"
+
+
+class TiDBManager:
+    def __init__(self, instance_urls: List[str],
+                 ssl: Optional[ssl.SSLContext] = None):
+        self.instances = [TiDBInstance.from_url(url, ssl=ssl) for url in instance_urls]
+        self._current_instance_index = 0
+
+    @property
+    def current_instance(self) -> TiDBInstance:
+        return self.instances[self._current_instance_index]
+
+    @property
+    def current_index(self) -> int:
+        return self._current_instance_index
+
+    async def initialize_pools(self, minsize: int = 1, maxsize: int = 10):
+        for inst in self.instances:
+            try:
+                await inst.create_pool(minsize=minsize, maxsize=maxsize)
+                logger.info("Connected to TiDB: %s", inst)
+            except Exception as exc:
+                logger.error("Failed to connect to TiDB %s: %s", inst, exc)
+                inst.is_available = False
+
+    async def close_all_pools(self):
+        for inst in self.instances:
+            await inst.close_pool()
+
+    async def ensure_tables(self, table_ddls: Dict[str, str]):
+        for inst in self.instances:
+            if not inst.is_available or not inst.pool:
+                continue
+            async with inst.pool.acquire() as conn:
+                async with conn.cursor() as cur:
+                    for table_name, ddl in table_ddls.items():
+                        try:
+                            await cur.execute(
+                                "SELECT 1 FROM information_schema.tables "
+                                "WHERE table_schema = %s AND table_name = %s",
+                                (inst.database, table_name),
+                            )
+                            exists = await cur.fetchone()
+                            if not exists:
+                                for statement in ddl.split(";"):
+                                    stmt = statement.strip()
+                                    if stmt:
+                                        await cur.execute(stmt)
+                                logger.info("Created table '%s' on %s", table_name, inst)
+                        except Exception as exc:
+                            logger.error(
+                                "Failed to create table '%s' on %s: %s",
+                                table_name, inst, exc,
+                            )
+                await conn.commit()
+
+    async def has_table(self, table_name: str) -> bool:
+        inst = self.current_instance
+        if not inst.pool:
+            return False
+        async with inst.pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT 1 FROM information_schema.tables "
+                    "WHERE table_schema = %s AND table_name = %s",
+                    (inst.database, table_name),
+                )
+                return await cur.fetchone() is not None
+
+    async def check_and_rotate(self) -> bool:
+        if self.current_instance.is_available and not self.current_instance.is_full:
+            return True
+
+        for i in range(1, len(self.instances)):
+            idx = (self._current_instance_index + i) % len(self.instances)
+            inst = self.instances[idx]
+            if inst.is_available and not inst.is_full:
+                logger.info(
+                    "Rotating TiDB write target: %s -> %s",
+                    self.instances[self._current_instance_index], inst,
+                )
+                self._current_instance_index = idx
+                return True
+
+        logger.error("All TiDB instances are full or unavailable")
+        return False
+
+    async def fetchone(self, query: str, params=None) -> Optional[Dict[str, Any]]:
+        inst = self.current_instance
+        if not inst.pool:
+            return None
+        async with inst.pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(query, params or ())
+                return await cur.fetchone()
+
+    async def fetchone_all(self, query: str, params=None) -> Optional[Dict[str, Any]]:
+        async def _query(inst_: TiDBInstance):
+            try:
+                async with inst_.pool.acquire() as conn:
+                    async with conn.cursor() as cur:
+                        await cur.execute(query, params or ())
+                        return await cur.fetchone()
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                logger.error("Error reading from %s: %s", inst_, exc)
+                return None
+
+        tasks = [
+            asyncio.create_task(_query(inst))
+            for inst in self.instances
+            if inst.pool and inst.is_available
+        ]
+        if not tasks:
+            return None
+
+        for coro in asyncio.as_completed(tasks):
+            result = await coro
+            if result is not None:
+                for t in tasks:
+                    if not t.done():
+                        t.cancel()
+                return result
+        return None
+
+    async def fetchall(self, query: str, params=None) -> List[Dict[str, Any]]:
+        inst = self.current_instance
+        if not inst.pool:
+            return []
+        async with inst.pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(query, params or ())
+                return await cur.fetchall()
+
+    async def fetchall_all(self, query: str, params=None) -> List[Dict[str, Any]]:
+        async def _query(inst_: TiDBInstance):
+            try:
+                async with inst_.pool.acquire() as conn:
+                    async with conn.cursor() as cur:
+                        await cur.execute(query, params or ())
+                        return await cur.fetchall()
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                logger.error("Error reading from %s: %s", inst_, exc)
+                return []
+
+        tasks = [
+            asyncio.create_task(_query(inst))
+            for inst in self.instances
+            if inst.pool and inst.is_available
+        ]
+        if not tasks:
+            return []
+
+        all_rows = await asyncio.gather(*tasks)
+        seen = set()
+        merged = []
+        for rows in all_rows:
+            for row in rows:
+                key = tuple(row.items())
+                if key not in seen:
+                    seen.add(key)
+                    merged.append(row)
+        return merged
+
+    async def _try_write(self, query: str, params) -> int:
+        inst = self.current_instance
+        if not inst.pool:
+            raise RuntimeError(f"TiDB pool not initialized for {inst}")
+        async with inst.pool.acquire() as conn:
+            async with conn.cursor(aiomysql.DictCursor) as cur:
+                await cur.execute(query, params or ())
+                await conn.commit()
+                return cur.lastrowid if cur.lastrowid is not None else 0
+
+    async def execute(self, query: str, params=None) -> int:
+        ok = await self.check_and_rotate()
+        if not ok:
+            raise TiDBWriteError()
+        try:
+            return await self._try_write(query, params)
+        except Exception as exc:
+            if _is_storage_full_error(exc):
+                logger.warning("TiDB instance %s is full, rotating...", self.current_instance)
+                self.current_instance.is_full = True
+                ok = await self.check_and_rotate()
+                if not ok:
+                    raise TiDBWriteError() from exc
+                try:
+                    return await self._try_write(query, params)
+                except Exception:
+                    raise TiDBWriteError() from exc
+            raise
+
+    async def execute_on_all(self, query: str, params=None):
+        for inst in self.instances:
+            if not inst.is_available or not inst.pool:
+                continue
+            try:
+                async with inst.pool.acquire() as conn:
+                    async with conn.cursor(aiomysql.DictCursor) as cur:
+                        await cur.execute(query, params or ())
+                    await conn.commit()
+            except Exception as exc:
+                logger.error("Error executing on %s: %s", inst, exc)
+
+    async def insert_and_get_id(self, query: str, params=None) -> int:
+        ok = await self.check_and_rotate()
+        if not ok:
+            raise TiDBWriteError()
+        try:
+            return await self._try_write(query, params)
+        except Exception as exc:
+            if _is_storage_full_error(exc):
+                logger.warning("TiDB instance %s is full, rotating...", self.current_instance)
+                self.current_instance.is_full = True
+                ok = await self.check_and_rotate()
+                if not ok:
+                    raise TiDBWriteError() from exc
+                try:
+                    return await self._try_write(query, params)
+                except Exception:
+                    raise TiDBWriteError() from exc
+            raise
+
+    def get_instances_summary(self) -> List[Dict[str, Any]]:
+        return [
+            {
+                "host": inst.host,
+                "port": inst.port,
+                "database": inst.database,
+                "is_available": inst.is_available,
+                "is_full": inst.is_full,
+                "is_current": i == self._current_instance_index,
+            }
+            for i, inst in enumerate(self.instances)
+        ]
+
+
+_tidb_manager: Optional[TiDBManager] = None
+
+
+def get_tidb_manager() -> Optional[TiDBManager]:
+    global _tidb_manager
+    return _tidb_manager
+
+
+def set_tidb_manager(mgr: TiDBManager):
+    global _tidb_manager
+    _tidb_manager = mgr

app/core/vector_store/__init__.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
-from app.core.vector_store.models import VectorStoreIndex, Base
 from app.core.vector_store.deps import (
-    engine,
-    AsyncSessionLocal,
-    init_vector_store_db,
     get_vs_db,
+    get_vs_manager,
+    init_vector_store_db,
+    set_vs_manager,
 )
+from app.core.vector_store.models import VectorStoreIndex
 
 __all__ = [
-    "Base",
     "VectorStoreIndex",
-    "engine",
-    "AsyncSessionLocal",
     "init_vector_store_db",
     "get_vs_db",
+    "set_vs_manager",
+    "get_vs_manager",
 ]

app/core/vector_store/deps.py

@@ -1,32 +1,52 @@
 from __future__ import annotations
 
-import os
-from typing import AsyncGenerator
-
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+import logging
+from typing import Optional
 
 from app.config import get_settings
-from app.core.vector_store.models import Base
+from app.core.tidb_manager import TiDBManager, get_tidb_manager
+
+logger = logging.getLogger(__name__)
 
 _settings = get_settings()
 
-DATA_DIR = _settings.data_dir
-os.makedirs(DATA_DIR, exist_ok=True)
+VS_TABLES: dict[str, str] = {
+    "vector_store_index": """
+        CREATE TABLE IF NOT EXISTS vector_store_index (
+            store_id VARCHAR(36) PRIMARY KEY,
+            name VARCHAR(255) NOT NULL,
+            path VARCHAR(1024) NOT NULL,
+            description VARCHAR(1024) DEFAULT '',
+            metadata_json TEXT DEFAULT '{}',
+            created_at VARCHAR(64) NOT NULL
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+}
 
-VECTOR_STORE_DB_URL = f"sqlite+aiosqlite:///{os.path.join(DATA_DIR, 'vector_stores.db')}"
 
-engine = create_async_engine(VECTOR_STORE_DB_URL, echo=False)
-AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+async def init_vector_store_db():
+    mgr = get_tidb_manager()
+    if mgr is None:
+        logger.error("Cannot init vector store DB: TiDB manager not initialized")
+        return
+    await mgr.ensure_tables(VS_TABLES)
+    logger.info("Vector store tables ensured on all TiDB instances")
 
 
-async def init_vector_store_db():
-    async with engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
+async def get_vs_db() -> TiDBManager:
+    mgr = get_tidb_manager()
+    if mgr is None:
+        raise RuntimeError("TiDB manager not initialized")
+    return mgr
+
+
+_tidb_vs_manager: Optional[TiDBManager] = None
+
+
+def set_vs_manager(mgr: TiDBManager):
+    global _tidb_vs_manager
+    _tidb_vs_manager = mgr
 
 
-async def get_vs_db() -> AsyncGenerator[AsyncSession, None]:
-    async with AsyncSessionLocal() as session:
-        try:
-            yield session
-        finally:
-            await session.close()
+def get_vs_manager() -> Optional[TiDBManager]:
+    return _tidb_vs_manager

app/core/vector_store/models.py

@@ -1,30 +1,23 @@
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any, Dict
 
-from sqlalchemy import String, Text
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
-
-
-class Base(DeclarativeBase):
-    pass
-
 
 def _utcnow() -> str:
     return datetime.now(timezone.utc).isoformat()
 
 
-class VectorStoreIndex(Base):
-    __tablename__ = "vector_store_index"
-
-    store_id: Mapped[str] = mapped_column(String(36), primary_key=True)
-    name: Mapped[str] = mapped_column(String(255), nullable=False)
-    path: Mapped[str] = mapped_column(String(1024), nullable=False)
-    description: Mapped[str] = mapped_column(String(1024), default="")
-    metadata_json: Mapped[str] = mapped_column(Text, default="{}")
-    created_at: Mapped[str] = mapped_column(String(64), default=_utcnow)
+@dataclass
+class VectorStoreIndex:
+    store_id: str = ""
+    name: str = ""
+    path: str = ""
+    description: str = ""
+    metadata_json: str = "{}"
+    created_at: str = field(default_factory=_utcnow)
 
     def to_dict(self) -> Dict[str, Any]:
         return {

app/models/domain.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass, field
 
 import tiktoken
 
-
 _ENCODING_CACHE: dict[str, tiktoken.Encoding] = {}
 
 

app/models/schemas.py

@@ -573,7 +573,6 @@ class VectorStoreResponse(BaseModel):
     document_count: int
     created_at: str
     metadata: Dict[str, Any] = {}
-    warning: Optional[Dict[str, str]] = None
 
 
 class VectorStoreListResponse(BaseModel):

app/services/__init__.py

@@ -2,21 +2,21 @@ from __future__ import annotations
 
 from app.services.auth_service import AuthService
 from app.services.converter_service import ConverterService
-from app.services.embeddings_service import EmbeddingService
-from app.services.extraction_service import ExtractionService
-from app.services.ocr_service import OCRService
 from app.services.dataset_metadata_service import (
     ExtractionConfig,
-    extract_metadata,
-    FileType,
     FileSource,
+    FileType,
     MetadataExtractionError,
+    extract_metadata,
 )
+from app.services.embeddings_service import EmbeddingService
+from app.services.extraction_service import ExtractionService
+from app.services.ocr_service import OCRService
 from app.services.prompt_service import (
+    CSV_STRICT_OUTPUT_PROMPT,
+    CSV_SYSTEM_PROMPT,
     build_csv_context_prompt,
     build_csv_system_prompt_with_context,
-    CSV_SYSTEM_PROMPT,
-    CSV_STRICT_OUTPUT_PROMPT,
 )
 
 __all__ = [

app/services/auth_service.py

@@ -9,12 +9,11 @@ import jwt
 from argon2 import PasswordHasher
 from argon2.exceptions import VerifyMismatchError
 from fastapi import HTTPException, Request
-from sqlalchemy import select, update
-from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.config import get_settings
+from app.core.auth.deps import _row_to_refresh_session, _row_to_user, _serialize_dt
+from app.core.auth.models import RefreshSession, User, _uuid
 from app.core.auth.models import _utcnow as _now
-from app.core.auth.models import RefreshSession, Role, User
 from app.core.auth.schemas import (
     ChangePasswordSchema,
     ForgotPasswordSchema,
@@ -25,6 +24,7 @@ from app.core.auth.schemas import (
     UpdateProfileSchema,
     UserProfile,
 )
+from app.core.tidb_manager import TiDBManager
 
 logger = logging.getLogger("auth_service")
 _settings = get_settings()
@@ -38,43 +38,51 @@ def _token_key(raw: str) -> str:
 class AuthService:
 
     @staticmethod
-    async def register(db: AsyncSession, schema: RegisterSchema) -> User:
-        result = await db.execute(select(User).where(User.email == schema.email))
-        if result.scalars().first():
+    async def register(db: TiDBManager, schema: RegisterSchema) -> User:
+        existing = await db.fetchone_all(
+            "SELECT id FROM users WHERE email = %s AND deleted_at IS NULL",
+            (schema.email,),
+        )
+        if existing:
             raise HTTPException(status_code=409, detail="Email already registered")
 
         if schema.username:
-            result = await db.execute(select(User).where(User.username == schema.username))
-            if result.scalars().first():
+            existing = await db.fetchone_all(
+                "SELECT id FROM users WHERE username = %s AND deleted_at IS NULL",
+                (schema.username,),
+            )
+            if existing:
                 raise HTTPException(status_code=409, detail="Username already taken")
 
-        user = User(
-            email=schema.email,
-            username=schema.username,
-            full_name=schema.full_name,
-            password_hash=ph.hash(schema.password),
+        user_id = _uuid()
+        now = _now()
+        await db.insert_and_get_id(
+            "INSERT INTO users (id, email, username, full_name, password_hash, created_at, updated_at) "
+            "VALUES (%s, %s, %s, %s, %s, %s, %s)",
+            (user_id, schema.email, schema.username, schema.full_name, ph.hash(schema.password), now, now),
         )
 
-        role_result = await db.execute(select(Role).where(Role.name == "User"))
-        default_role = role_result.scalars().first()
-        if default_role:
-            user.roles.append(default_role)
+        role_row = await db.fetchone("SELECT id FROM roles WHERE name = 'User'")
+        if role_row:
+            await db.execute(
+                "INSERT INTO user_roles (user_id, role_id) VALUES (%s, %s)",
+                (user_id, role_row["id"]),
+            )
 
-        db.add(user)
-        await db.commit()
-        await db.refresh(user)
-        return user
+        user_row = await db.fetchone("SELECT * FROM users WHERE id = %s", (user_id,))
+        return _row_to_user(user_row)
 
     @staticmethod
-    async def login(db: AsyncSession, request: Request, schema: LoginSchema) -> TokenData:
-        result = await db.execute(
-            select(User).where(User.email == schema.email, User.deleted_at.is_(None))
+    async def login(db: TiDBManager, request: Request, schema: LoginSchema) -> TokenData:
+        row = await db.fetchone_all(
+            "SELECT * FROM users WHERE email = %s AND deleted_at IS NULL",
+            (schema.email,),
         )
-        user = result.scalars().first()
-
-        if not user:
+        if not row:
             raise HTTPException(status_code=401, detail="Incorrect email or password")
 
+        user = _row_to_user(row)
+
         if user.locked_until and user.locked_until > _now():
             raise HTTPException(
                 status_code=403,
@@ -85,149 +93,158 @@ class AuthService:
             user.failed_login_attempts += 1
             if user.failed_login_attempts >= _settings.max_login_attempts:
                 user.locked_until = _now() + timedelta(minutes=_settings.lockout_minutes)
-            await db.commit()
+            await db.execute(
+                "UPDATE users SET failed_login_attempts = %s, locked_until = %s WHERE id = %s",
+                (user.failed_login_attempts, _serialize_dt(user.locked_until), user.id),
+            )
             raise HTTPException(status_code=401, detail="Incorrect email or password")
 
-        user.failed_login_attempts = 0
-        user.locked_until = None
-        user.last_login = _now()
+        now = _now()
+        await db.execute(
+            "UPDATE users SET failed_login_attempts = 0, locked_until = NULL, last_login = %s "
+            "WHERE id = %s",
+            (now, user.id),
+        )
 
         access_token = AuthService._create_access_token(user.id)
         raw_refresh, refresh_hash, token_key, expires_at = AuthService._create_refresh_token()
 
-        session = RefreshSession(
-            user_id=user.id,
-            token_key=token_key,
-            token_hash=refresh_hash,
-            expires_at=expires_at,
-            device_info=request.headers.get("User-Agent", "Unknown"),
-            ip_address=request.client.host if request.client else "Unknown",
+        session_id = _uuid()
+        await db.execute(
+            "INSERT INTO refresh_sessions (id, user_id, token_key, token_hash, device_info, "
+            "ip_address, expires_at, created_at, updated_at) "
+            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
+            (
+                session_id, user.id, token_key, refresh_hash,
+                request.headers.get("User-Agent", "Unknown"),
+                request.client.host if request.client else "Unknown",
+                _serialize_dt(expires_at), now, now,
+            ),
         )
-        db.add(session)
-        await db.commit()
 
         return TokenData(access_token=access_token, refresh_token=raw_refresh)
 
     @staticmethod
-    async def refresh(db: AsyncSession, raw_refresh_token: str) -> TokenData:
+    async def refresh(db: TiDBManager, raw_refresh_token: str) -> TokenData:
         key = _token_key(raw_refresh_token)
 
-        result = await db.execute(
-            select(RefreshSession).where(
-                RefreshSession.token_key == key,
-                RefreshSession.revoked_at.is_(None),
-                RefreshSession.expires_at > _now(),
-            )
+        row = await db.fetchone_all(
+            "SELECT * FROM refresh_sessions WHERE token_key = %s AND revoked_at IS NULL "
+            "AND expires_at > %s",
+            (key, _serialize_dt(_now())),
         )
-        session = result.scalars().first()
-
-        if not session:
-            session = await db.execute(
-                select(RefreshSession).where(RefreshSession.token_key == key)
+        if not row:
+            comp = await db.fetchone_all(
+                "SELECT * FROM refresh_sessions WHERE token_key = %s", (key,),
             )
-            existing = session.scalars().first()
-            if existing and existing.revoked_at is not None:
-                await db.execute(
-                    update(RefreshSession)
-                    .where(RefreshSession.user_id == existing.user_id, RefreshSession.revoked_at.is_(None))
-                    .values(revoked_at=_now())
+            if comp and comp.get("revoked_at") is not None:
+                now = _serialize_dt(_now())
+                await db.execute_on_all(
+                    "UPDATE refresh_sessions SET revoked_at = %s "
+                    "WHERE user_id = %s AND revoked_at IS NULL",
+                    (now, comp["user_id"]),
                 )
-                await db.commit()
                 raise HTTPException(
                     status_code=401,
                     detail="Session compromised. All sessions revoked. Please login again.",
                 )
             raise HTTPException(status_code=401, detail="Invalid or expired refresh token")
 
+        session = _row_to_refresh_session(row)
+
         if not AuthService._verify_token(raw_refresh_token, session.token_hash):
             raise HTTPException(status_code=401, detail="Invalid refresh token")
 
-        user_result = await db.execute(
-            select(User).where(
-                User.id == session.user_id, User.deleted_at.is_(None), User.is_active.is_(True)
-            )
+        user_row = await db.fetchone_all(
+            "SELECT * FROM users WHERE id = %s AND deleted_at IS NULL AND is_active = 1",
+            (session.user_id,),
         )
-        user = user_result.scalars().first()
-        if not user:
+        if not user_row:
             raise HTTPException(status_code=401, detail="User not found or inactive")
+        user = _row_to_user(user_row)
 
-        session.revoked_at = _now()
+        now = _serialize_dt(_now())
+        await db.execute(
+            "UPDATE refresh_sessions SET revoked_at = %s WHERE id = %s",
+            (now, session.id),
+        )
 
         new_access = AuthService._create_access_token(user.id)
         new_raw_refresh, new_hash, new_key, new_expires = AuthService._create_refresh_token()
 
-        new_session = RefreshSession(
-            user_id=user.id,
-            token_key=new_key,
-            token_hash=new_hash,
-            expires_at=new_expires,
-            device_info=session.device_info,
-            ip_address=session.ip_address,
+        new_session_id = _uuid()
+        await db.execute(
+            "INSERT INTO refresh_sessions (id, user_id, token_key, token_hash, device_info, "
+            "ip_address, expires_at, created_at, updated_at) "
+            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
+            (
+                new_session_id, user.id, new_key, new_hash,
+                session.device_info, session.ip_address,
+                _serialize_dt(new_expires), _serialize_dt(_now()), _serialize_dt(_now()),
+            ),
         )
-        db.add(new_session)
 
-        user.last_login = _now()
-        await db.commit()
+        await db.execute(
+            "UPDATE users SET last_login = %s WHERE id = %s",
+            (_serialize_dt(_now()), user.id),
+        )
 
         return TokenData(access_token=new_access, refresh_token=new_raw_refresh)
 
     @staticmethod
-    async def logout(db: AsyncSession, user: User, raw_refresh_token: str):
+    async def logout(db: TiDBManager, user: User, raw_refresh_token: str):
         key = _token_key(raw_refresh_token)
-        result = await db.execute(
-            select(RefreshSession).where(
-                RefreshSession.token_key == key,
-                RefreshSession.user_id == user.id,
-                RefreshSession.revoked_at.is_(None),
-            )
+        row = await db.fetchone_all(
+            "SELECT id FROM refresh_sessions WHERE token_key = %s "
+            "AND user_id = %s AND revoked_at IS NULL",
+            (key, user.id),
         )
-        session = result.scalars().first()
-        if not session:
+        if not row:
             raise HTTPException(status_code=404, detail="Session not found")
-        session.revoked_at = _now()
-        await db.commit()
+        await db.execute(
+            "UPDATE refresh_sessions SET revoked_at = %s WHERE id = %s",
+            (_serialize_dt(_now()), row["id"]),
+        )
 
     @staticmethod
-    async def logout_all(db: AsyncSession, user: User):
-        now = _now()
+    async def logout_all(db: TiDBManager, user: User):
+        now = _serialize_dt(_now())
         await db.execute(
-            update(RefreshSession)
-            .where(
-                RefreshSession.user_id == user.id,
-                RefreshSession.revoked_at.is_(None),
-            )
-            .values(revoked_at=now)
+            "UPDATE refresh_sessions SET revoked_at = %s "
+            "WHERE user_id = %s AND revoked_at IS NULL",
+            (now, user.id),
         )
-        await db.commit()
 
     @staticmethod
-    async def change_password(db: AsyncSession, user: User, schema: ChangePasswordSchema):
+    async def change_password(db: TiDBManager, user: User, schema: ChangePasswordSchema):
         if not AuthService._verify_password(schema.current_password, user.password_hash):
             raise HTTPException(status_code=400, detail="Incorrect current password")
-        user.password_hash = ph.hash(schema.new_password)
-        user.password_changed_at = _now()
-        await db.commit()
+        now = _serialize_dt(_now())
+        await db.execute(
+            "UPDATE users SET password_hash = %s, password_changed_at = %s WHERE id = %s",
+            (ph.hash(schema.new_password), now, user.id),
+        )
 
     @staticmethod
-    async def forgot_password(db: AsyncSession, schema: ForgotPasswordSchema):
-        result = await db.execute(
-            select(User).where(User.email == schema.email, User.deleted_at.is_(None))
+    async def forgot_password(db: TiDBManager, schema: ForgotPasswordSchema):
+        row = await db.fetchone_all(
+            "SELECT id FROM users WHERE email = %s AND deleted_at IS NULL",
+            (schema.email,),
         )
-        user = result.scalars().first()
-        if user:
+        if row:
             reset_token = jwt.encode(
                 {
-                    "sub": user.id,
+                    "sub": row["id"],
                     "type": "reset_password",
                     "exp": _now() + timedelta(hours=1),
                 },
                 _settings.jwt_secret_key,
                 algorithm=_settings.jwt_algorithm,
             )
-            logger.info("Password reset token for %s: %s", user.email, reset_token)
+            logger.info("Password reset token for %s: %s", schema.email, reset_token)
 
     @staticmethod
-    async def reset_password(db: AsyncSession, schema: ResetPasswordSchema):
+    async def reset_password(db: TiDBManager, schema: ResetPasswordSchema):
         try:
             payload = jwt.decode(
                 schema.token,
@@ -240,86 +257,89 @@ class AuthService:
         except jwt.PyJWTError:
             raise HTTPException(status_code=400, detail="Invalid or expired reset token")
 
-        result = await db.execute(select(User).where(User.id == user_id))
-        user = result.scalars().first()
-        if not user:
+        row = await db.fetchone_all("SELECT id FROM users WHERE id = %s", (user_id,))
+        if not row:
             raise HTTPException(status_code=404, detail="User not found")
 
-        user.password_hash = ph.hash(schema.new_password)
-        user.password_changed_at = _now()
-
+        now = _serialize_dt(_now())
+        await db.execute(
+            "UPDATE users SET password_hash = %s, password_changed_at = %s WHERE id = %s",
+            (ph.hash(schema.new_password), now, user_id),
+        )
         await db.execute(
-            update(RefreshSession)
-            .where(RefreshSession.user_id == user.id, RefreshSession.revoked_at.is_(None))
-            .values(revoked_at=_now())
+            "UPDATE refresh_sessions SET revoked_at = %s "
+            "WHERE user_id = %s AND revoked_at IS NULL",
+            (now, user_id),
         )
-        await db.commit()
 
     @staticmethod
-    async def update_profile(db: AsyncSession, user: User, schema: UpdateProfileSchema) -> User:
+    async def update_profile(db: TiDBManager, user: User, schema: UpdateProfileSchema) -> User:
         if schema.username is not None:
-            result = await db.execute(
-                select(User).where(User.username == schema.username, User.id != user.id)
+            existing = await db.fetchone_all(
+                "SELECT id FROM users WHERE username = %s AND id != %s AND deleted_at IS NULL",
+                (schema.username, user.id),
             )
-            if result.scalars().first():
+            if existing:
                 raise HTTPException(status_code=409, detail="Username already taken")
-            user.username = schema.username
+            await db.execute("UPDATE users SET username = %s WHERE id = %s", (schema.username, user.id))
         if schema.full_name is not None:
-            user.full_name = schema.full_name
-        await db.commit()
-        await db.refresh(user)
-        return user
+            await db.execute(
+                "UPDATE users SET full_name = %s WHERE id = %s",
+                (schema.full_name, user.id),
+            )
+        row = await db.fetchone("SELECT * FROM users WHERE id = %s", (user.id,))
+        return _row_to_user(row)
 
     @staticmethod
-    async def soft_delete(db: AsyncSession, user: User):
-        user.deleted_at = _now()
-        user.is_active = False
+    async def soft_delete(db: TiDBManager, user: User):
+        now = _serialize_dt(_now())
+        await db.execute(
+            "UPDATE users SET deleted_at = %s, is_active = 0 WHERE id = %s",
+            (now, user.id),
+        )
         await db.execute(
-            update(RefreshSession)
-            .where(RefreshSession.user_id == user.id, RefreshSession.revoked_at.is_(None))
-            .values(revoked_at=_now())
+            "UPDATE refresh_sessions SET revoked_at = %s "
+            "WHERE user_id = %s AND revoked_at IS NULL",
+            (now, user.id),
         )
-        await db.commit()
 
     @staticmethod
-    async def list_sessions(db: AsyncSession, user: User) -> list[RefreshSession]:
-        result = await db.execute(
-            select(RefreshSession)
-            .where(
-                RefreshSession.user_id == user.id,
-                RefreshSession.revoked_at.is_(None),
-            )
-            .order_by(RefreshSession.created_at.desc())
+    async def list_sessions(db: TiDBManager, user: User) -> list[RefreshSession]:
+        rows = await db.fetchall_all(
+            "SELECT * FROM refresh_sessions WHERE user_id = %s AND revoked_at IS NULL "
+            "ORDER BY created_at DESC",
+            (user.id,),
         )
-        return list(result.scalars().all())
+        return [_row_to_refresh_session(r) for r in rows]
 
     @staticmethod
-    async def revoke_session(db: AsyncSession, user: User, session_id: str):
-        result = await db.execute(
-            select(RefreshSession).where(
-                RefreshSession.id == session_id,
-                RefreshSession.user_id == user.id,
-                RefreshSession.revoked_at.is_(None),
-            )
+    async def revoke_session(db: TiDBManager, user: User, session_id: str):
+        row = await db.fetchone_all(
+            "SELECT id FROM refresh_sessions WHERE id = %s AND user_id = %s AND revoked_at IS NULL",
+            (session_id, user.id),
         )
-        session = result.scalars().first()
-        if not session:
+        if not row:
             raise HTTPException(status_code=404, detail="Session not found")
-        session.revoked_at = _now()
-        await db.commit()
+        await db.execute(
+            "UPDATE refresh_sessions SET revoked_at = %s WHERE id = %s",
+            (_serialize_dt(_now()), session_id),
+        )
 
     @staticmethod
-    async def cleanup_expired_sessions(db: AsyncSession):
-        result = await db.execute(
-            select(RefreshSession).where(RefreshSession.expires_at < _now())
+    async def cleanup_expired_sessions(db: TiDBManager):
+        rows = await db.fetchall_all(
+            "SELECT * FROM refresh_sessions WHERE expires_at < %s",
+            (_serialize_dt(_now()),),
         )
         count = 0
-        for s in result.scalars().all():
-            if s.revoked_at is None:
-                s.revoked_at = _now()
+        for row in rows:
+            if row.get("revoked_at") is None:
+                await db.execute(
+                    "UPDATE refresh_sessions SET revoked_at = %s WHERE id = %s",
+                    (_serialize_dt(_now()), row["id"]),
+                )
                 count += 1
         if count:
-            await db.commit()
             logger.info("Cleaned up %s expired sessions", count)
 
     @staticmethod
@@ -359,7 +379,19 @@ class AuthService:
             return False
 
     @staticmethod
-    def user_to_profile(user: User) -> UserProfile:
+    async def user_to_profile(db: TiDBManager, user: User) -> UserProfile:
+        roles = []
+        if db:
+            try:
+                rows = await db.fetchall_all(
+                    "SELECT r.name FROM roles r "
+                    "JOIN user_roles ur ON ur.role_id = r.id "
+                    "WHERE ur.user_id = %s",
+                    (user.id,),
+                )
+                roles = [r["name"] for r in rows]
+            except Exception:
+                pass
         return UserProfile(
             id=user.id,
             email=user.email,
@@ -367,6 +399,6 @@ class AuthService:
             full_name=user.full_name,
             is_active=user.is_active,
             is_verified=user.is_verified,
-            roles=[r.name for r in user.roles],
+            roles=roles,
             created_at=user.created_at,
         )

app/services/csv_analysis_service.py

@@ -5,8 +5,8 @@ import json
 import logging
 import os
 import shutil
-import subprocess
 import signal
+import subprocess
 import sys
 import tempfile
 import time

app/services/dataset_metadata_service.py

@@ -31,9 +31,9 @@ from urllib.parse import unquote, urlparse
 import aiohttp
 import chardet
 import numpy as np
+import openpyxl  # noqa: F401 – needed as engine for .xlsx
 import pandas as pd
 import xlrd  # noqa: F401 – needed as engine for .xls
-import openpyxl  # noqa: F401 – needed as engine for .xlsx
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)

app/services/prompts/csv_system_prompt.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-import os
 from typing import Any, Dict
 
 _JSON_EXAMPLE = (

app/services/py_sandbox.py

@@ -1,31 +1,31 @@
 #!/usr/bin/env python3
-import sys
-import io
-import os
-import math as _math
-import json as _json
-import re as _re
-import random as _random
-import collections as _collections
-import itertools as _itertools
-import functools as _functools
-import operator as _operator
-import string as _string
-import decimal as _decimal
-import fractions as _fractions
-import statistics as _statistics
-import heapq as _heapq
+import base64 as _base64
 import bisect as _bisect
+import collections as _collections
 import copy as _copy
-import typing as _typing
-import textwrap as _textwrap
 import datetime as _datetime
+import decimal as _decimal
+import enum as _enum
+import fractions as _fractions
+import functools as _functools
 import hashlib as _hashlib
-import uuid as _uuid
+import heapq as _heapq
+import io
+import itertools as _itertools
+import json as _json
+import math as _math
+import operator as _operator
+import os
+import random as _random
+import re as _re
 import secrets as _secrets
-import enum as _enum
-import base64 as _base64
+import statistics as _statistics
+import string as _string
+import sys
+import textwrap as _textwrap
 import time as _time
+import typing as _typing
+import uuid as _uuid
 
 _SAFE_MODULES = {
     'math': _math, 'json': _json, 're': _re, 'random': _random,

app/services/text_cleaner_service.py

@@ -1,5 +1,6 @@
-import re
 import html
+import re
+
 from cleantext import clean
 from text_unidecode import unidecode
 

app/services/url_shortener_service.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import csv
 import hashlib
 import hmac
@@ -9,11 +10,9 @@ import json
 import os
 import secrets
 import socket
-import sqlite3
 import string
 import threading
 import unicodedata
-from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List, Optional
@@ -21,6 +20,7 @@ from urllib.parse import urlparse
 
 from app.config import get_settings
 from app.core.logger import get_logger
+from app.core.tidb_manager import TiDBManager, get_tidb_manager
 
 logger = get_logger(__name__)
 
@@ -278,15 +278,12 @@ def _apply_geo_targeting(rules: Dict[str, Any], default_url: str,
     ua_lower = (user_agent or "").lower()
     is_mobile = "mobile" in ua_lower or "android" in ua_lower or "iphone" in ua_lower
     is_tablet = "tablet" in ua_lower or "ipad" in ua_lower
-    # Device rule takes precedence
     if is_tablet and "tablet" in devices:
         return devices["tablet"]
     if is_mobile and "mobile" in devices:
         return devices["mobile"]
     if not (is_mobile or is_tablet) and "desktop" in devices:
         return devices["desktop"]
-    # Country-specific override (requires IP geolocation service)
-    # For now, limited to simplistic check
     return default_url
 
 
@@ -308,266 +305,290 @@ def _fire_click_webhook(webhook_url: str, short_code: str, long_url: str,
             client.post(webhook_url, json=payload)
     except ImportError:
         import urllib.request
-        import json as _json
-        data = _json.dumps(payload).encode()
+        data = json.dumps(payload).encode()
         req = urllib.request.Request(webhook_url, data=data,
                                      headers={"Content-Type": "application/json"},
                                      method="POST")
         urllib.request.urlopen(req, timeout=5)
 
 
+# ---------------------------------------------------------------------------
+# TiDB-backed Storage
+# ---------------------------------------------------------------------------
+
+US_TABLES: dict[str, str] = {
+    "url_shortener_owners": """
+        CREATE TABLE IF NOT EXISTS url_shortener_owners (
+            owner_id      VARCHAR(64) PRIMARY KEY,
+            name          VARCHAR(255) NOT NULL,
+            api_key_hash  VARCHAR(128) NOT NULL,
+            plan          VARCHAR(32) NOT NULL DEFAULT 'free',
+            created_at    VARCHAR(64) NOT NULL,
+            UNIQUE KEY uk_uso_api_key_hash (api_key_hash)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "url_shortener_links": """
+        CREATE TABLE IF NOT EXISTS url_shortener_links (
+            short_code       VARCHAR(32) PRIMARY KEY,
+            long_url         TEXT NOT NULL,
+            owner_id         VARCHAR(64) NOT NULL,
+            created_at       VARCHAR(64) NOT NULL,
+            expires_at       VARCHAR(64),
+            max_clicks       INT,
+            click_count      INT NOT NULL DEFAULT 0,
+            last_accessed_at VARCHAR(64),
+            is_active        TINYINT(1) NOT NULL DEFAULT 1,
+            tags             VARCHAR(512),
+            note             TEXT,
+            password_hash    VARCHAR(128),
+            utm_source       VARCHAR(512),
+            utm_medium       VARCHAR(512),
+            utm_campaign     VARCHAR(512),
+            campaign_id      VARCHAR(64),
+            custom_domain    VARCHAR(255),
+            fallback_url     TEXT,
+            webhook_url      TEXT,
+            geo_targeting    TEXT,
+            INDEX idx_usl_owner (owner_id),
+            INDEX idx_usl_campaign (campaign_id),
+            CONSTRAINT fk_usl_owner FOREIGN KEY (owner_id) REFERENCES url_shortener_owners(owner_id)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "url_shortener_clicks": """
+        CREATE TABLE IF NOT EXISTS url_shortener_clicks (
+            id          BIGINT AUTO_INCREMENT PRIMARY KEY,
+            short_code  VARCHAR(32) NOT NULL,
+            referrer    TEXT,
+            user_agent  TEXT,
+            ip_address  VARCHAR(45),
+            country     VARCHAR(64),
+            browser     VARCHAR(64),
+            device      VARCHAR(64),
+            os          VARCHAR(64),
+            clicked_at  VARCHAR(64) NOT NULL,
+            INDEX idx_usc_code (short_code),
+            INDEX idx_usc_clicked_at (clicked_at),
+            CONSTRAINT fk_usc_link FOREIGN KEY (short_code) REFERENCES url_shortener_links(short_code)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "url_shortener_audit_log": """
+        CREATE TABLE IF NOT EXISTS url_shortener_audit_log (
+            id         BIGINT AUTO_INCREMENT PRIMARY KEY,
+            timestamp  VARCHAR(64) NOT NULL,
+            owner_id   VARCHAR(64),
+            action     VARCHAR(64) NOT NULL,
+            short_code VARCHAR(32),
+            detail     TEXT
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+    "url_shortener_campaigns": """
+        CREATE TABLE IF NOT EXISTS url_shortener_campaigns (
+            campaign_id  VARCHAR(64) PRIMARY KEY,
+            owner_id     VARCHAR(64) NOT NULL,
+            name         VARCHAR(255) NOT NULL,
+            description  TEXT,
+            created_at   VARCHAR(64) NOT NULL,
+            is_active    TINYINT(1) NOT NULL DEFAULT 1,
+            INDEX idx_uscamp_owner (owner_id),
+            CONSTRAINT fk_uscamp_owner FOREIGN KEY (owner_id) REFERENCES url_shortener_owners(owner_id)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
+    """,
+}
+
+
+def _run_async(coro):
+    try:
+        loop = asyncio.get_running_loop()
+        if loop.is_running():
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1):
+                future = asyncio.run_coroutine_threadsafe(coro, loop)
+                return future.result()
+    except RuntimeError:
+        pass
+    return asyncio.run(coro)
+
+
 class Storage:
-    def __init__(self, db_path: str):
-        self._db_path = db_path
+    def __init__(self):
         self._lock = threading.RLock()
-        self._local = threading.local()
-        self._init_schema()
-
-    @property
-    def _conn(self) -> sqlite3.Connection:
-        if not hasattr(self._local, "conn"):
-            conn = sqlite3.connect(self._db_path, check_same_thread=False)
-            conn.row_factory = sqlite3.Row
-            conn.execute("PRAGMA foreign_keys = ON")
-            conn.execute("PRAGMA journal_mode = WAL")
-            self._local.conn = conn
-        return self._local.conn
-
-    @contextmanager
-    def writer_lock(self):
-        with self._lock:
-            yield
-
-    def _init_schema(self) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS url_shortener_owners (
-                    owner_id      TEXT PRIMARY KEY,
-                    name          TEXT NOT NULL,
-                    api_key_hash  TEXT NOT NULL UNIQUE,
-                    plan          TEXT NOT NULL DEFAULT 'free',
-                    created_at    TEXT NOT NULL
-                )
-            """)
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS url_shortener_links (
-                    short_code       TEXT PRIMARY KEY,
-                    long_url         TEXT NOT NULL,
-                    owner_id         TEXT NOT NULL,
-                    created_at       TEXT NOT NULL,
-                    expires_at       TEXT,
-                    max_clicks       INTEGER,
-                    click_count      INTEGER NOT NULL DEFAULT 0,
-                    last_accessed_at TEXT,
-                    is_active        INTEGER NOT NULL DEFAULT 1,
-                    tags             TEXT,
-                    note             TEXT,
-                    password_hash    TEXT,
-                    utm_source       TEXT,
-                    utm_medium       TEXT,
-                    utm_campaign     TEXT,
-                    campaign_id      TEXT,
-                    custom_domain    TEXT,
-                    fallback_url     TEXT,
-                    webhook_url      TEXT,
-                    geo_targeting    TEXT,
-                    FOREIGN KEY (owner_id) REFERENCES url_shortener_owners(owner_id)
-                )
-            """)
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS url_shortener_clicks (
-                    id          INTEGER PRIMARY KEY AUTOINCREMENT,
-                    short_code  TEXT NOT NULL,
-                    referrer    TEXT,
-                    user_agent  TEXT,
-                    ip_address  TEXT,
-                    country     TEXT,
-                    browser     TEXT,
-                    device      TEXT,
-                    os          TEXT,
-                    clicked_at  TEXT NOT NULL,
-                    FOREIGN KEY (short_code) REFERENCES url_shortener_links(short_code)
-                )
-            """)
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS url_shortener_audit_log (
-                    id         INTEGER PRIMARY KEY AUTOINCREMENT,
-                    timestamp  TEXT NOT NULL,
-                    owner_id   TEXT,
-                    action     TEXT NOT NULL,
-                    short_code TEXT,
-                    detail     TEXT
-                )
-            """)
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS url_shortener_campaigns (
-                    campaign_id  TEXT PRIMARY KEY,
-                    owner_id     TEXT NOT NULL,
-                    name         TEXT NOT NULL,
-                    description  TEXT,
-                    created_at   TEXT NOT NULL,
-                    is_active    INTEGER NOT NULL DEFAULT 1,
-                    FOREIGN KEY (owner_id) REFERENCES url_shortener_owners(owner_id)
-                )
-            """)
-            conn.execute("CREATE INDEX IF NOT EXISTS idx_us_links_owner ON url_shortener_links(owner_id)")
-            conn.execute("CREATE INDEX IF NOT EXISTS idx_us_clicks_code ON url_shortener_clicks(short_code)")
-            conn.execute("CREATE INDEX IF NOT EXISTS idx_us_clicks_at ON url_shortener_clicks(clicked_at)")
-            conn.execute("CREATE INDEX IF NOT EXISTS idx_us_links_campaign ON url_shortener_links(campaign_id)")
-            self._migrate_schema(conn)
-
-    def _migrate_schema(self, conn: sqlite3.Connection) -> None:
-        existing = {r["name"] for r in conn.execute("PRAGMA table_info(url_shortener_links)").fetchall()}
-        migrations = {
-            "campaign_id": "ALTER TABLE url_shortener_links ADD COLUMN campaign_id TEXT",
-            "custom_domain": "ALTER TABLE url_shortener_links ADD COLUMN custom_domain TEXT",
-            "fallback_url": "ALTER TABLE url_shortener_links ADD COLUMN fallback_url TEXT",
-            "webhook_url": "ALTER TABLE url_shortener_links ADD COLUMN webhook_url TEXT",
-            "geo_targeting": "ALTER TABLE url_shortener_links ADD COLUMN geo_targeting TEXT",
-        }
-        for col, ddl in migrations.items():
-            if col not in existing:
-                conn.execute(ddl)
 
-    def create_owner(self, owner_id: str, name: str, api_key_hash: str, plan: str = "free") -> None:
-        with self._lock, self._conn as conn:
-            conn.execute(
+    def _get_mgr(self) -> TiDBManager:
+        mgr = get_tidb_manager()
+        if mgr is None:
+            raise RuntimeError("TiDB manager not initialized")
+        return mgr
+
+    def _init_schema(self):
+        mgr = self._get_mgr()
+        _run_async(mgr.ensure_tables(US_TABLES))
+
+    def close(self):
+        pass
+
+    # ---------- Owners ----------
+
+    def create_owner(self, owner_id: str, name: str, api_key_hash: str, plan: str = "free"):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
                 "INSERT INTO url_shortener_owners (owner_id, name, api_key_hash, plan, created_at) "
-                "VALUES (?, ?, ?, ?, ?)",
+                "VALUES (%s, %s, %s, %s, %s)",
                 (owner_id, name, api_key_hash, plan, _now_iso()),
-            )
+            ))
 
-    def get_owner_by_key_hash(self, api_key_hash: str) -> Optional[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_owners WHERE api_key_hash = ?", (api_key_hash,)
-        ).fetchone()
+    def get_owner_by_key_hash(self, api_key_hash: str) -> Optional[dict]:
+        return _run_async(self._get_mgr().fetchone_all(
+            "SELECT * FROM url_shortener_owners WHERE api_key_hash = %s",
+            (api_key_hash,),
+        ))
 
-    def get_owner(self, owner_id: str) -> Optional[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_owners WHERE owner_id = ?", (owner_id,)
-        ).fetchone()
+    def get_owner(self, owner_id: str) -> Optional[dict]:
+        return _run_async(self._get_mgr().fetchone_all(
+            "SELECT * FROM url_shortener_owners WHERE owner_id = %s",
+            (owner_id,),
+        ))
 
-    def update_owner_plan(self, owner_id: str, plan: str) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute("UPDATE url_shortener_owners SET plan = ? WHERE owner_id = ?", (plan, owner_id))
+    def update_owner_plan(self, owner_id: str, plan: str):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
+                "UPDATE url_shortener_owners SET plan = %s WHERE owner_id = %s",
+                (plan, owner_id),
+            ))
+
+    # ---------- Links ----------
 
     def code_exists(self, short_code: str) -> bool:
-        return self._conn.execute(
-            "SELECT 1 FROM url_shortener_links WHERE short_code = ?", (short_code,)
-        ).fetchone() is not None
+        row = _run_async(self._get_mgr().fetchone_all(
+            "SELECT 1 AS chk FROM url_shortener_links WHERE short_code = %s",
+            (short_code,),
+        ))
+        return row is not None
 
-    def insert_link(self, **kwargs: Any) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute("""
+    def insert_link(self, **kwargs):
+        with self._lock:
+            _run_async(self._get_mgr().insert_and_get_id("""
                 INSERT INTO url_shortener_links
                     (short_code, long_url, owner_id, created_at, expires_at,
                      max_clicks, click_count, last_accessed_at, is_active, tags, note,
                      password_hash, utm_source, utm_medium, utm_campaign,
                      campaign_id, custom_domain, fallback_url, webhook_url, geo_targeting)
                 VALUES
-                    (:short_code, :long_url, :owner_id, :created_at, :expires_at,
-                     :max_clicks, 0, NULL, 1, :tags, :note,
-                     :password_hash, :utm_source, :utm_medium, :utm_campaign,
-                     :campaign_id, :custom_domain, :fallback_url, :webhook_url,
-                     :geo_targeting)
-            """, kwargs)
-
-    def update_link(self, short_code: str, **updates: Any) -> None:
-        with self._lock, self._conn as conn:
-            sets = ", ".join(f"{k} = ?" for k in updates)
+                    (%(short_code)s, %(long_url)s, %(owner_id)s, %(created_at)s, %(expires_at)s,
+                     %(max_clicks)s, 0, NULL, 1, %(tags)s, %(note)s,
+                     %(password_hash)s, %(utm_source)s, %(utm_medium)s, %(utm_campaign)s,
+                     %(campaign_id)s, %(custom_domain)s, %(fallback_url)s, %(webhook_url)s,
+                     %(geo_targeting)s)
+            """, kwargs))
+
+    def update_link(self, short_code: str, **updates):
+        with self._lock:
+            sets = ", ".join(f"{k} = %s" for k in updates)
             vals = list(updates.values()) + [short_code]
-            conn.execute(f"UPDATE url_shortener_links SET {sets} WHERE short_code = ?", vals)
-
-    def get_link(self, short_code: str) -> Optional[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_links WHERE short_code = ?", (short_code,)
-        ).fetchone()
-
-    def list_links_for_owner(self, owner_id: str) -> List[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_links WHERE owner_id = ? ORDER BY created_at DESC",
+            _run_async(self._get_mgr().execute(
+                f"UPDATE url_shortener_links SET {sets} WHERE short_code = %s",
+                tuple(vals),
+            ))
+
+    def get_link(self, short_code: str) -> Optional[dict]:
+        return _run_async(self._get_mgr().fetchone_all(
+            "SELECT * FROM url_shortener_links WHERE short_code = %s",
+            (short_code,),
+        ))
+
+    def list_links_for_owner(self, owner_id: str) -> List[dict]:
+        return _run_async(self._get_mgr().fetchall_all(
+            "SELECT * FROM url_shortener_links WHERE owner_id = %s ORDER BY created_at DESC",
             (owner_id,),
-        ).fetchall()
+        ))
 
-    def deactivate_link(self, short_code: str) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute("UPDATE url_shortener_links SET is_active = 0 WHERE short_code = ?", (short_code,))
+    def deactivate_link(self, short_code: str):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
+                "UPDATE url_shortener_links SET is_active = 0 WHERE short_code = %s",
+                (short_code,),
+            ))
 
     def record_click(self, short_code: str, referrer: Optional[str] = None,
-                     user_agent: Optional[str] = None, ip_address: Optional[str] = None) -> None:
+                     user_agent: Optional[str] = None, ip_address: Optional[str] = None):
         now = _now_iso()
         parsed = _parse_user_agent(user_agent or "")
-        with self._lock, self._conn as conn:
-            conn.execute(
+        with self._lock:
+            _run_async(self._get_mgr().execute(
                 "INSERT INTO url_shortener_clicks "
                 "(short_code, referrer, user_agent, ip_address, country, browser, device, os, clicked_at) "
-                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
                 (short_code, referrer, user_agent, ip_address, None,
                  parsed["browser"], parsed["device"], parsed["os"], now),
-            )
-            conn.execute(
-                "UPDATE url_shortener_links SET click_count = click_count + 1, last_accessed_at = ? "
-                "WHERE short_code = ?",
+            ))
+            _run_async(self._get_mgr().execute(
+                "UPDATE url_shortener_links SET click_count = click_count + 1, last_accessed_at = %s "
+                "WHERE short_code = %s",
                 (now, short_code),
-            )
+            ))
 
     def get_click_analytics(self, short_code: str) -> Dict[str, Any]:
-        total = self._conn.execute(
-            "SELECT COUNT(*) FROM url_shortener_clicks WHERE short_code = ?", (short_code,)
-        ).fetchone()[0]
-        browsers = dict(self._conn.execute(
-            "SELECT browser, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-            "AND browser IS NOT NULL GROUP BY browser ORDER BY COUNT(*) DESC",
-            (short_code,)
-        ).fetchall())
-        devices = dict(self._conn.execute(
-            "SELECT device, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-            "AND device IS NOT NULL GROUP BY device ORDER BY COUNT(*) DESC",
-            (short_code,)
-        ).fetchall())
-        referrers = dict(self._conn.execute(
+        mgr = self._get_mgr()
+        total = _run_async(mgr.fetchone_all(
+            "SELECT COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s",
+            (short_code,),
+        ))
+        total_val = total["cnt"] if total else 0
+        browsers_rows = _run_async(mgr.fetchall_all(
+            "SELECT browser, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+            "AND browser IS NOT NULL GROUP BY browser ORDER BY cnt DESC",
+            (short_code,),
+        ))
+        devices_rows = _run_async(mgr.fetchall_all(
+            "SELECT device, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+            "AND device IS NOT NULL GROUP BY device ORDER BY cnt DESC",
+            (short_code,),
+        ))
+        referrers_rows = _run_async(mgr.fetchall_all(
             "SELECT COALESCE(NULLIF(referrer, ''), '(direct)') AS ref, COUNT(*) AS n "
-            "FROM url_shortener_clicks WHERE short_code = ? "
+            "FROM url_shortener_clicks WHERE short_code = %s "
             "GROUP BY ref ORDER BY n DESC LIMIT 10",
-            (short_code,)
-        ).fetchall())
-        os_data = dict(self._conn.execute(
-            "SELECT os, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-            "AND os IS NOT NULL GROUP BY os ORDER BY COUNT(*) DESC",
-            (short_code,)
-        ).fetchall())
-        recent = self._conn.execute(
-            "SELECT clicked_at FROM url_shortener_clicks WHERE short_code = ? "
+            (short_code,),
+        ))
+        os_rows = _run_async(mgr.fetchall_all(
+            "SELECT os, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+            "AND os IS NOT NULL GROUP BY os ORDER BY cnt DESC",
+            (short_code,),
+        ))
+        recent = _run_async(mgr.fetchall_all(
+            "SELECT clicked_at FROM url_shortener_clicks WHERE short_code = %s "
             "ORDER BY clicked_at DESC LIMIT 50",
-            (short_code,)
-        ).fetchall()
+            (short_code,),
+        ))
         return {
-            "total_clicks": total,
-            "browsers": browsers,
-            "devices": devices,
-            "operating_systems": os_data,
-            "top_referrers": referrers,
+            "total_clicks": total_val,
+            "browsers": {r["browser"]: r["cnt"] for r in browsers_rows},
+            "devices": {r["device"]: r["cnt"] for r in devices_rows},
+            "operating_systems": {r["os"]: r["cnt"] for r in os_rows},
+            "top_referrers": {r["ref"]: r["n"] for r in referrers_rows},
             "recent_clicks": [r["clicked_at"] for r in recent],
         }
 
     def get_link_count_for_owner(self, owner_id: str) -> int:
-        return self._conn.execute(
-            "SELECT COUNT(*) FROM url_shortener_links WHERE owner_id = ?", (owner_id,)
-        ).fetchone()[0]
+        row = _run_async(self._get_mgr().fetchone_all(
+            "SELECT COUNT(*) AS cnt FROM url_shortener_links WHERE owner_id = %s",
+            (owner_id,),
+        ))
+        return row["cnt"] if row else 0
 
-    def audit(self, owner_id: Optional[str], action: str, short_code: Optional[str], detail: str = "") -> None:
-        with self._lock, self._conn as conn:
-            conn.execute(
+    def audit(self, owner_id: Optional[str], action: str, short_code: Optional[str], detail: str = ""):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
                 "INSERT INTO url_shortener_audit_log (timestamp, owner_id, action, short_code, detail) "
-                "VALUES (?, ?, ?, ?, ?)",
+                "VALUES (%s, %s, %s, %s, %s)",
                 (_now_iso(), owner_id, action, short_code, detail),
-            )
+            ))
 
     def export_links_csv(self, owner_id: str) -> str:
-        rows = self._conn.execute(
+        rows = _run_async(self._get_mgr().fetchall_all(
             "SELECT short_code, long_url, created_at, click_count, is_active, tags, note "
-            "FROM url_shortener_links WHERE owner_id = ? ORDER BY created_at DESC", (owner_id,)
-        ).fetchall()
+            "FROM url_shortener_links WHERE owner_id = %s ORDER BY created_at DESC",
+            (owner_id,),
+        ))
         buf = io.StringIO()
         w = csv.writer(buf)
         w.writerow(["short_code", "long_url", "created_at", "click_count", "is_active", "tags", "note"])
@@ -577,10 +598,11 @@ class Storage:
         return buf.getvalue()
 
     def export_clicks_csv(self, short_code: str) -> str:
-        rows = self._conn.execute(
+        rows = _run_async(self._get_mgr().fetchall_all(
             "SELECT id, referrer, user_agent, ip_address, browser, device, os, clicked_at "
-            "FROM url_shortener_clicks WHERE short_code = ? ORDER BY clicked_at DESC", (short_code,)
-        ).fetchall()
+            "FROM url_shortener_clicks WHERE short_code = %s ORDER BY clicked_at DESC",
+            (short_code,),
+        ))
         buf = io.StringIO()
         w = csv.writer(buf)
         w.writerow(["id", "referrer", "user_agent", "ip_address", "browser", "device", "os", "clicked_at"])
@@ -589,76 +611,85 @@ class Storage:
                        r["browser"], r["device"], r["os"], r["clicked_at"]])
         return buf.getvalue()
 
-    def list_links_by_campaign(self, campaign_id: str) -> List[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_links WHERE campaign_id = ? ORDER BY created_at DESC",
+    def list_links_by_campaign(self, campaign_id: str) -> List[dict]:
+        return _run_async(self._get_mgr().fetchall_all(
+            "SELECT * FROM url_shortener_links WHERE campaign_id = %s ORDER BY created_at DESC",
             (campaign_id,),
-        ).fetchall()
+        ))
 
-    def create_campaign(self, campaign_id: str, owner_id: str, name: str, description: Optional[str] = None) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute(
+    def create_campaign(self, campaign_id: str, owner_id: str, name: str, description: Optional[str] = None):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
                 "INSERT INTO url_shortener_campaigns (campaign_id, owner_id, name, description, created_at) "
-                "VALUES (?, ?, ?, ?, ?)",
+                "VALUES (%s, %s, %s, %s, %s)",
                 (campaign_id, owner_id, name, description, _now_iso()),
-            )
+            ))
 
-    def get_campaign(self, campaign_id: str) -> Optional[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_campaigns WHERE campaign_id = ?", (campaign_id,)
-        ).fetchone()
+    def get_campaign(self, campaign_id: str) -> Optional[dict]:
+        return _run_async(self._get_mgr().fetchone_all(
+            "SELECT * FROM url_shortener_campaigns WHERE campaign_id = %s",
+            (campaign_id,),
+        ))
 
-    def list_campaigns_for_owner(self, owner_id: str) -> List[sqlite3.Row]:
-        return self._conn.execute(
-            "SELECT * FROM url_shortener_campaigns WHERE owner_id = ? ORDER BY created_at DESC",
+    def list_campaigns_for_owner(self, owner_id: str) -> List[dict]:
+        return _run_async(self._get_mgr().fetchall_all(
+            "SELECT * FROM url_shortener_campaigns WHERE owner_id = %s ORDER BY created_at DESC",
             (owner_id,),
-        ).fetchall()
+        ))
 
-    def deactivate_campaign(self, campaign_id: str) -> None:
-        with self._lock, self._conn as conn:
-            conn.execute(
-                "UPDATE url_shortener_campaigns SET is_active = 0 WHERE campaign_id = ?",
+    def deactivate_campaign(self, campaign_id: str):
+        with self._lock:
+            _run_async(self._get_mgr().execute(
+                "UPDATE url_shortener_campaigns SET is_active = 0 WHERE campaign_id = %s",
                 (campaign_id,),
-            )
+            ))
 
     def get_campaign_link_count(self, campaign_id: str) -> int:
-        return self._conn.execute(
-            "SELECT COUNT(*) FROM url_shortener_links WHERE campaign_id = ?", (campaign_id,)
-        ).fetchone()[0]
+        row = _run_async(self._get_mgr().fetchone_all(
+            "SELECT COUNT(*) AS cnt FROM url_shortener_links WHERE campaign_id = %s",
+            (campaign_id,),
+        ))
+        return row["cnt"] if row else 0
 
     def get_campaign_total_clicks(self, campaign_id: str) -> int:
-        return self._conn.execute(
-            "SELECT COALESCE(SUM(click_count), 0) FROM url_shortener_links WHERE campaign_id = ?",
+        row = _run_async(self._get_mgr().fetchone_all(
+            "SELECT COALESCE(SUM(click_count), 0) AS cnt FROM url_shortener_links WHERE campaign_id = %s",
             (campaign_id,),
-        ).fetchone()[0]
+        ))
+        return row["cnt"] if row else 0
 
     def get_campaign_analytics(self, campaign_id: str) -> Dict[str, Any]:
-        codes = self._conn.execute(
-            "SELECT short_code FROM url_shortener_links WHERE campaign_id = ?", (campaign_id,)
-        ).fetchall()
+        codes = _run_async(self._get_mgr().fetchall_all(
+            "SELECT short_code FROM url_shortener_links WHERE campaign_id = %s",
+            (campaign_id,),
+        ))
         total = 0
         browsers: Dict[str, int] = {}
         devices: Dict[str, int] = {}
         os_data: Dict[str, int] = {}
-        for (code,) in codes:
-            total += self._conn.execute(
-                "SELECT COUNT(*) FROM url_shortener_clicks WHERE short_code = ?", (code,)
-            ).fetchone()[0]
-            for b, n in self._conn.execute(
-                "SELECT browser, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-                "AND browser IS NOT NULL GROUP BY browser", (code,)
-            ).fetchall():
-                browsers[b] = browsers.get(b, 0) + n
-            for d, n in self._conn.execute(
-                "SELECT device, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-                "AND device IS NOT NULL GROUP BY device", (code,)
-            ).fetchall():
-                devices[d] = devices.get(d, 0) + n
-            for o, n in self._conn.execute(
-                "SELECT os, COUNT(*) FROM url_shortener_clicks WHERE short_code = ? "
-                "AND os IS NOT NULL GROUP BY os", (code,)
-            ).fetchall():
-                os_data[o] = os_data.get(o, 0) + n
+        mgr = self._get_mgr()
+        for row in codes:
+            code = row["short_code"]
+            cnt_row = _run_async(mgr.fetchone_all(
+                "SELECT COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s",
+                (code,),
+            ))
+            total += cnt_row["cnt"] if cnt_row else 0
+            for b in _run_async(mgr.fetchall_all(
+                "SELECT browser, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+                "AND browser IS NOT NULL GROUP BY browser", (code,),
+            )):
+                browsers[b["browser"]] = browsers.get(b["browser"], 0) + b["cnt"]
+            for d in _run_async(mgr.fetchall_all(
+                "SELECT device, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+                "AND device IS NOT NULL GROUP BY device", (code,),
+            )):
+                devices[d["device"]] = devices.get(d["device"], 0) + d["cnt"]
+            for o in _run_async(mgr.fetchall_all(
+                "SELECT os, COUNT(*) AS cnt FROM url_shortener_clicks WHERE short_code = %s "
+                "AND os IS NOT NULL GROUP BY os", (code,),
+            )):
+                os_data[o["os"]] = os_data.get(o["os"], 0) + o["cnt"]
         return {"total_clicks": total, "browsers": browsers, "devices": devices, "operating_systems": os_data}
 
     def owner_summary(self, owner_id: str) -> Dict[str, Any]:
@@ -666,30 +697,31 @@ class Storage:
         if row is None:
             return {}
         link_count = self.get_link_count_for_owner(owner_id)
-        total_clicks = self._conn.execute(
-            "SELECT COALESCE(SUM(click_count), 0) FROM url_shortener_links WHERE owner_id = ?",
-            (owner_id,)
-        ).fetchone()[0]
-        active = self._conn.execute(
-            "SELECT COUNT(*) FROM url_shortener_links WHERE owner_id = ? AND is_active = 1",
-            (owner_id,)
-        ).fetchone()[0]
+        total_clicks_row = _run_async(self._get_mgr().fetchone_all(
+            "SELECT COALESCE(SUM(click_count), 0) AS cnt FROM url_shortener_links WHERE owner_id = %s",
+            (owner_id,),
+        ))
+        total_clicks = total_clicks_row["cnt"] if total_clicks_row else 0
+        active = _run_async(self._get_mgr().fetchone_all(
+            "SELECT COUNT(*) AS cnt FROM url_shortener_links WHERE owner_id = %s AND is_active = 1",
+            (owner_id,),
+        ))
         return {
             "owner_id": owner_id,
             "name": row["name"],
             "plan": row["plan"],
             "total_links": link_count,
-            "active_links": active,
+            "active_links": active["cnt"] if active else 0,
             "total_clicks": total_clicks,
             "created_at": row["created_at"],
         }
 
 
 class URLShortenerService:
-    def __init__(self, db_path: Optional[str] = None):
-        self.db_path = db_path or os.environ.get("URL_SHORTENER_DB", "data/url_shortener.db")
-        self.storage = Storage(self.db_path)
-        logger.info("URLShortenerService initialized (db=%s)", self.db_path)
+    def __init__(self):
+        self.storage = Storage()
+        self.storage._init_schema()
+        logger.info("URLShortenerService initialized (TiDB)")
 
     def close(self) -> None:
         self.storage.close()
@@ -727,26 +759,15 @@ class URLShortenerService:
         logger.info("Owner %s plan updated to %s", owner_id, plan)
         return {"owner_id": owner_id, "plan": plan}
 
-    def shorten(
-        self,
-        long_url: str,
-        owner_id: str,
-        custom_alias: Optional[str] = None,
-        expires_in_days: Optional[int] = None,
-        expires_at: Optional[datetime] = None,
-        max_clicks: Optional[int] = None,
-        tags: Optional[List[str]] = None,
-        note: Optional[str] = None,
-        password: Optional[str] = None,
-        utm_source: Optional[str] = None,
-        utm_medium: Optional[str] = None,
-        utm_campaign: Optional[str] = None,
-        campaign_id: Optional[str] = None,
-        custom_domain: Optional[str] = None,
-        fallback_url: Optional[str] = None,
-        webhook_url: Optional[str] = None,
-        geo_targeting: Optional[Dict[str, Any]] = None,
-    ) -> ShortLink:
+    def shorten(self, long_url: str, owner_id: str, custom_alias: Optional[str] = None,
+                expires_in_days: Optional[int] = None, expires_at: Optional[datetime] = None,
+                max_clicks: Optional[int] = None, tags: Optional[List[str]] = None,
+                note: Optional[str] = None, password: Optional[str] = None,
+                utm_source: Optional[str] = None, utm_medium: Optional[str] = None,
+                utm_campaign: Optional[str] = None, campaign_id: Optional[str] = None,
+                custom_domain: Optional[str] = None, fallback_url: Optional[str] = None,
+                webhook_url: Optional[str] = None,
+                geo_targeting: Optional[Dict[str, Any]] = None) -> ShortLink:
         plan_limits = self._get_owner_plan(owner_id)
 
         if custom_alias is not None:
@@ -811,7 +832,7 @@ class URLShortenerService:
         if fallback_url:
             validate_url(fallback_url)
 
-        with self.storage.writer_lock():
+        with self.storage._lock:
             code = custom_alias
             if code is not None:
                 if self.storage.code_exists(code):
@@ -861,7 +882,7 @@ class URLShortenerService:
 
     def resolve(self, short_code: str, referrer: Optional[str] = None,
                 user_agent: Optional[str] = None, ip_address: Optional[str] = None) -> str:
-        with self.storage.writer_lock():
+        with self.storage._lock:
             row = self.storage.get_link(short_code)
             if row is None:
                 raise LinkNotFoundError(f"No link found for code '{short_code}'")
@@ -889,7 +910,6 @@ class URLShortenerService:
             self.storage.record_click(short_code, referrer, user_agent, ip_address)
             long_url = row["long_url"]
 
-            # Geo-targeting: override destination based on geo/device rules
             geo_raw = row["geo_targeting"]
             if isinstance(geo_raw, str) and geo_raw:
                 try:
@@ -910,7 +930,6 @@ class URLShortenerService:
                 new_qs = "&".join(f"{k}={v}" for k, v in existing.items())
                 long_url = parsed._replace(query=new_qs).geturl()
 
-        # Fire webhook asynchronously outside the lock
         webhook = row["webhook_url"]
         if webhook:
             try:
@@ -923,7 +942,7 @@ class URLShortenerService:
 
     def resolve_with_password(self, short_code: str, password: str, referrer: Optional[str] = None,
                                user_agent: Optional[str] = None, ip_address: Optional[str] = None) -> str:
-        with self.storage.writer_lock():
+        with self.storage._lock:
             row = self.storage.get_link(short_code)
             if row is None:
                 raise LinkNotFoundError(f"No link found for code '{short_code}'")
@@ -944,7 +963,7 @@ class URLShortenerService:
     def list_links(self, owner_id: str) -> List[ShortLink]:
         return [self._row_to_link(r) for r in self.storage.list_links_for_owner(owner_id)]
 
-    def update_link(self, short_code: str, owner_id: str, **updates: Any) -> ShortLink:
+    def update_link(self, short_code: str, owner_id: str, **updates) -> ShortLink:
         row = self.storage.get_link(short_code)
         if row is None:
             raise LinkNotFoundError(f"No link found for code '{short_code}'")
@@ -1048,12 +1067,15 @@ class URLShortenerService:
             raise LinkNotFoundError(f"No link found for code '{short_code}'")
         if row["owner_id"] != owner_id:
             raise AuthorizationError("You do not own this link")
-        rows = self.storage._conn.execute(
+        mgr = get_tidb_manager()
+        if mgr is None:
+            return []
+        rows = _run_async(mgr.fetchall_all(
             "SELECT id, referrer, user_agent, ip_address, browser, device, os, clicked_at "
-            "FROM url_shortener_clicks WHERE short_code = ? ORDER BY clicked_at DESC LIMIT ?",
+            "FROM url_shortener_clicks WHERE short_code = %s ORDER BY clicked_at DESC LIMIT %s",
             (short_code, limit),
-        ).fetchall()
-        return [dict(r) for r in rows]
+        ))
+        return list(rows)
 
     def get_qr_code(self, short_code: str) -> Dict[str, str]:
         row = self.storage.get_link(short_code)
@@ -1170,8 +1192,8 @@ class URLShortenerService:
             result["error"] = str(exc)
         return result
 
-    def _row_to_link(self, row: sqlite3.Row) -> ShortLink:
-        gt_raw = row["geo_targeting"]
+    def _row_to_link(self, row: dict) -> ShortLink:
+        gt_raw = row.get("geo_targeting")
         geo = json.loads(gt_raw) if isinstance(gt_raw, str) and gt_raw else None
         return ShortLink(
             short_code=row["short_code"],
@@ -1182,15 +1204,15 @@ class URLShortenerService:
             max_clicks=row["max_clicks"],
             click_count=row["click_count"],
             is_active=bool(row["is_active"]),
-            tags=row["tags"].split(",") if row["tags"] else [],
-            note=row["note"],
-            password_hash=row["password_hash"],
-            utm_source=row["utm_source"],
-            utm_medium=row["utm_medium"],
-            utm_campaign=row["utm_campaign"],
-            campaign_id=row["campaign_id"],
-            custom_domain=row["custom_domain"],
-            fallback_url=row["fallback_url"],
-            webhook_url=row["webhook_url"],
+            tags=row["tags"].split(",") if row.get("tags") else [],
+            note=row.get("note"),
+            password_hash=row.get("password_hash"),
+            utm_source=row.get("utm_source"),
+            utm_medium=row.get("utm_medium"),
+            utm_campaign=row.get("utm_campaign"),
+            campaign_id=row.get("campaign_id"),
+            custom_domain=row.get("custom_domain"),
+            fallback_url=row.get("fallback_url"),
+            webhook_url=row.get("webhook_url"),
             geo_targeting=geo,
         )

app/services/vector_store_service.py

@@ -11,12 +11,10 @@ from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Tuple
 
 import zvec
-from sqlalchemy import select
 
 from app.config import get_settings
 from app.core.logger import get_logger
-from app.core.vector_store.deps import AsyncSessionLocal
-from app.core.vector_store.models import VectorStoreIndex
+from app.core.tidb_manager import get_tidb_manager
 from app.services.chunking_service import chunk_text_async
 from app.services.embeddings_service import EmbeddingService
 
@@ -64,58 +62,71 @@ class VectorStoreService:
     def _get_collection(self, store_id: str) -> Optional["zvec.Collection"]:
         return self._collections.get(store_id)
 
-    # --- SQLite persistence ---
+    # --- TiDB persistence ---
 
     async def init_db(self) -> None:
         from app.core.vector_store.deps import init_vector_store_db
         await init_vector_store_db()
-        async with AsyncSessionLocal() as session:
-            result = await session.execute(select(VectorStoreIndex))
-            rows = result.scalars().all()
-            for row in rows:
-                d = row.to_dict()
-                record = VectorStoreRecord(
-                    store_id=d["store_id"],
-                    name=d["name"],
-                    path=d["path"],
-                    description=d["description"],
-                    metadata=d["metadata"],
-                    created_at=d["created_at"],
-                )
-                self._stores[record.store_id] = record
-                store_path = d["path"]
-                if os.path.exists(os.path.join(store_path, "__zvec_meta")):
-                    try:
-                        col = zvec.open(store_path)
-                        if col is not None:
-                            self._collections[record.store_id] = col
-                    except Exception as exc:
-                        logger.warning("Could not open collection %s: %s", record.store_id, exc)
+        mgr = get_tidb_manager()
+        if mgr is None:
+            logger.warning("TiDB manager not available, vector store metadata not loaded")
+            return
+        rows = await mgr.fetchall_all(
+            "SELECT * FROM vector_store_index ORDER BY created_at ASC"
+        )
+        for row in rows:
+            record = VectorStoreRecord(
+                store_id=row["store_id"],
+                name=row["name"],
+                path=row["path"],
+                description=row.get("description", ""),
+                metadata=json.loads(row.get("metadata_json", "{}") or "{}"),
+                created_at=row["created_at"],
+            )
+            self._stores[record.store_id] = record
+            store_path = row["path"]
+            if os.path.exists(os.path.join(store_path, "__zvec_meta")):
+                try:
+                    col = zvec.open(store_path)
+                    if col is not None:
+                        self._collections[record.store_id] = col
+                except Exception as exc:
+                    logger.warning("Could not open collection %s: %s", record.store_id, exc)
 
     async def _persist_store(self, record: VectorStoreRecord) -> None:
-        async with AsyncSessionLocal() as session:
-            existing = await session.get(VectorStoreIndex, record.store_id)
-            if existing:
-                existing.name = record.name
-                existing.description = record.description
-                existing.metadata_json = json.dumps(record.metadata)
-            else:
-                session.add(VectorStoreIndex.from_dict({
-                    "store_id": record.store_id,
-                    "name": record.name,
-                    "path": record.path,
-                    "description": record.description,
-                    "metadata": record.metadata,
-                    "created_at": record.created_at,
-                }))
-            await session.commit()
+        mgr = get_tidb_manager()
+        if mgr is None:
+            logger.error("TiDB manager not available, cannot persist store")
+            return
+        existing = await mgr.fetchone(
+            "SELECT store_id FROM vector_store_index WHERE store_id = %s",
+            (record.store_id,),
+        )
+        if existing:
+            await mgr.execute(
+                "UPDATE vector_store_index SET name = %s, description = %s, "
+                "metadata_json = %s WHERE store_id = %s",
+                (record.name, record.description, json.dumps(record.metadata), record.store_id),
+            )
+        else:
+            await mgr.execute(
+                "INSERT INTO vector_store_index (store_id, name, path, description, metadata_json, created_at) "
+                "VALUES (%s, %s, %s, %s, %s, %s)",
+                (
+                    record.store_id, record.name, record.path,
+                    record.description, json.dumps(record.metadata),
+                    record.created_at,
+                ),
+            )
 
     async def _remove_persisted_store(self, store_id: str) -> None:
-        async with AsyncSessionLocal() as session:
-            row = await session.get(VectorStoreIndex, store_id)
-            if row:
-                await session.delete(row)
-                await session.commit()
+        mgr = get_tidb_manager()
+        if mgr is None:
+            return
+        await mgr.execute(
+            "DELETE FROM vector_store_index WHERE store_id = %s",
+            (store_id,),
+        )
 
     # --- Synchronous helpers (run in thread pool) ---
 

app/services/verify_service.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import phonenumbers
-from phonenumbers import carrier, geocoder, PhoneNumberType
+from phonenumbers import PhoneNumberType, carrier, geocoder
 
 
 class PhoneVerificationResult:

pyproject.toml

@@ -36,6 +36,10 @@ dependencies = [
 [project.optional-dependencies]
 dev = ["pytest>=8", "pytest-asyncio>=0.23"]
 
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["All API Collection*"]

requirements.txt

@@ -36,8 +36,6 @@ email-validator>=2.1.0
 slowapi>=0.1.9
 
 # Async database drivers
-aiosqlite>=0.20.0
-sqlalchemy[asyncio]>=2.0.0
 aiomysql>=0.3.2
 asyncpg>=0.31.0
 motor>=3.7.1