mirror of
https://github.com/KohakuBlueleaf/KohakuHub.git
synced 2026-03-11 17:34:08 -05:00
Allow user provide external token
This commit is contained in:
Kohaku-Blueleaf
133
scripts/db_migrations/014_user_external_tokens.py
Normal file
133
scripts/db_migrations/014_user_external_tokens.py
Normal file
@@ -0,0 +1,133 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Migration 014: Add UserExternalToken table for user-specific external fallback tokens.
|
||||
|
||||
Allows users to provide their own tokens for external fallback sources (HuggingFace, etc.).
|
||||
User tokens override admin-configured tokens for matching URLs.
|
||||
|
||||
Changes:
|
||||
- Add UserExternalToken table with encrypted token storage
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add src to path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
|
||||
# Add db_migrations to path (for _migration_utils)
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
from kohakuhub.config import cfg
|
||||
from kohakuhub.db import UserExternalToken, db
|
||||
from _migration_utils import check_table_exists, should_skip_due_to_future_migrations
|
||||
|
||||
MIGRATION_NUMBER = 14
|
||||
|
||||
|
||||
def is_applied(db, cfg):
|
||||
"""Check if THIS migration has been applied.
|
||||
|
||||
Returns True if UserExternalToken table exists.
|
||||
"""
|
||||
return check_table_exists(db, "userexternaltoken")
|
||||
|
||||
|
||||
def migrate_postgres():
|
||||
"""Create UserExternalToken table in PostgreSQL."""
|
||||
cursor = db.cursor()
|
||||
|
||||
print("Creating UserExternalToken table...")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS userexternaltoken (
|
||||
id SERIAL PRIMARY KEY,
|
||||
user_id INTEGER NOT NULL REFERENCES "user"(id) ON DELETE CASCADE,
|
||||
url VARCHAR(255) NOT NULL,
|
||||
encrypted_token TEXT NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE (user_id, url)
|
||||
)
|
||||
"""
|
||||
)
|
||||
print(" ✓ Created UserExternalToken table")
|
||||
|
||||
# Create index on user_id for faster lookups
|
||||
print("Creating index on user_id...")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS userexternaltoken_user_id
|
||||
ON userexternaltoken(user_id)
|
||||
"""
|
||||
)
|
||||
print(" ✓ Created index on user_id")
|
||||
|
||||
|
||||
def migrate_sqlite():
|
||||
"""Create UserExternalToken table in SQLite."""
|
||||
cursor = db.cursor()
|
||||
|
||||
print("Creating UserExternalToken table...")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS userexternaltoken (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL REFERENCES "user"(id) ON DELETE CASCADE,
|
||||
url VARCHAR(255) NOT NULL,
|
||||
encrypted_token TEXT NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE (user_id, url)
|
||||
)
|
||||
"""
|
||||
)
|
||||
print(" ✓ Created UserExternalToken table")
|
||||
|
||||
# Create index on user_id for faster lookups
|
||||
print("Creating index on user_id...")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS userexternaltoken_user_id
|
||||
ON userexternaltoken(user_id)
|
||||
"""
|
||||
)
|
||||
print(" ✓ Created index on user_id")
|
||||
|
||||
|
||||
def run():
|
||||
"""Run migration 014."""
|
||||
print("\n" + "=" * 60)
|
||||
print(f"Migration {MIGRATION_NUMBER}: Add UserExternalToken table")
|
||||
print("=" * 60)
|
||||
|
||||
# Connect to database
|
||||
db.connect(reuse_if_open=True)
|
||||
|
||||
# Check if should skip
|
||||
if should_skip_due_to_future_migrations(db, cfg, MIGRATION_NUMBER):
|
||||
print(
|
||||
f"\n⚠️ Skipping migration {MIGRATION_NUMBER} - already applied or superseded by future migrations\n"
|
||||
)
|
||||
return
|
||||
|
||||
# Check if already applied
|
||||
if is_applied(db, cfg):
|
||||
print(
|
||||
f"\n✓ Migration {MIGRATION_NUMBER} already applied (UserExternalToken table exists)\n"
|
||||
)
|
||||
return
|
||||
|
||||
# Run migration in transaction
|
||||
print(f"\nApplying migration {MIGRATION_NUMBER}...\n")
|
||||
|
||||
with db.atomic():
|
||||
if cfg.app.db_backend == "postgres":
|
||||
migrate_postgres()
|
||||
else:
|
||||
migrate_sqlite()
|
||||
|
||||
print(f"\n✓ Migration {MIGRATION_NUMBER} completed successfully!\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -273,6 +273,7 @@ def generate_hub_api_service(config: dict) -> str:
|
||||
## ===== CRITICAL: Security Configuration (MUST CHANGE) =====
|
||||
- KOHAKU_HUB_SESSION_SECRET={config['session_secret']}
|
||||
- KOHAKU_HUB_ADMIN_SECRET_TOKEN={config['admin_secret']}
|
||||
- KOHAKU_HUB_DATABASE_KEY={config['database_key']}
|
||||
|
||||
## ===== Performance Configuration =====
|
||||
- KOHAKU_HUB_WORKERS=4 # Number of worker processes (1-8, recommend: CPU cores)
|
||||
@@ -454,9 +455,13 @@ def load_config_file(config_path: Path) -> dict:
|
||||
config["admin_secret"] = sec.get(
|
||||
"admin_secret", fallback=generate_secret(48)
|
||||
) # 64 chars
|
||||
config["database_key"] = sec.get(
|
||||
"database_key", fallback=generate_secret(32)
|
||||
) # 43 chars
|
||||
else:
|
||||
config["session_secret"] = generate_secret(48) # 64 chars
|
||||
config["admin_secret"] = generate_secret(48) # 64 chars
|
||||
config["database_key"] = generate_secret(32) # 43 chars for encryption
|
||||
|
||||
# Network section
|
||||
if parser.has_section("network"):
|
||||
@@ -519,6 +524,7 @@ builtin = true
|
||||
# Session and admin secrets (auto-generated if not specified)
|
||||
# session_secret = your-session-secret-here
|
||||
# admin_secret = your-admin-secret-here
|
||||
# database_key = your-database-encryption-key-here # For encrypting external fallback tokens
|
||||
|
||||
[network]
|
||||
# External bridge network (optional)
|
||||
@@ -708,6 +714,17 @@ def interactive_config() -> dict:
|
||||
else:
|
||||
config["admin_secret"] = ask_string("Admin secret token")
|
||||
|
||||
# Database encryption key (for external tokens)
|
||||
print()
|
||||
default_database_key = generate_secret(32) # 43 chars for Fernet encryption
|
||||
print(f"Generated database encryption key: {default_database_key}")
|
||||
use_generated_db = ask_yes_no("Use generated database key?", default=True)
|
||||
|
||||
if use_generated_db:
|
||||
config["database_key"] = default_database_key
|
||||
else:
|
||||
config["database_key"] = ask_string("Database encryption key")
|
||||
|
||||
# LakeFS encryption key
|
||||
config["lakefs_encrypt_key"] = generate_secret(32) # 43 chars
|
||||
|
||||
@@ -796,6 +813,9 @@ token_expire_days = 365
|
||||
enabled = true
|
||||
secret_token = "{config['admin_secret']}"
|
||||
|
||||
[app]
|
||||
database_key = "{config['database_key']}" # For encrypting external fallback tokens
|
||||
|
||||
[quota]
|
||||
default_user_private_quota_bytes = 10_000_000 # 10MB
|
||||
default_user_public_quota_bytes = 100_000_000 # 100MB
|
||||
|
||||
Reference in New Issue
Block a user