diff --git a/docker/startup.py b/docker/startup.py index a8156a3..4adae78 100644 --- a/docker/startup.py +++ b/docker/startup.py @@ -107,7 +107,16 @@ def main(): print("[startup] Starting API server...") subprocess.run( - ["uvicorn", "kohakuhub.main:app", "--host", "0.0.0.0", "--port", "48888"], + [ + "uvicorn", + "kohakuhub.main:app", + "--workers", + "4", + "--host", + "0.0.0.0", + "--port", + "48888", + ], check=True, env=os.environ, ) diff --git a/src/kohakuhub/api/admin.py b/src/kohakuhub/api/admin.py index a076bde..85eef95 100644 --- a/src/kohakuhub/api/admin.py +++ b/src/kohakuhub/api/admin.py @@ -10,22 +10,29 @@ from fastapi import APIRouter, Depends, HTTPException, Header, Query from peewee import fn from pydantic import BaseModel +from kohakuhub.async_utils import run_in_s3_executor from kohakuhub.config import cfg -from kohakuhub.db import Commit, File, LFSObjectHistory, Organization, Repository, User -from kohakuhub.db_async import ( +from kohakuhub.db import ( + db, + Commit, + File, + LFSObjectHistory, + Organization, + Repository, + User, +) +from kohakuhub.db_operations import ( create_user, delete_repository, delete_user, - execute_db_query, ) from kohakuhub.logger import get_logger +from kohakuhub.utils.s3 import get_s3_client from kohakuhub.api.quota.util import ( get_storage_info, set_quota, update_namespace_storage, ) -from kohakuhub.utils.s3 import get_s3_client -from kohakuhub.async_utils import run_in_s3_executor logger = get_logger("ADMIN") router = APIRouter() @@ -132,10 +139,7 @@ async def get_user_info( HTTPException: If user not found """ - def _get_user(): - return User.get_or_none(User.username == username) - - user = await execute_db_query(_get_user) + user = User.get_or_none(User.username == username) if not user: raise HTTPException(404, detail={"error": f"User not found: {username}"}) @@ -171,25 +175,22 @@ async def list_users( List of users with quota information """ - def _list_users(): - users = User.select().limit(limit).offset(offset) - return [ - { - "id": u.id, - "username": u.username, - "email": u.email, - "email_verified": u.email_verified, - "is_active": u.is_active, - "private_quota_bytes": u.private_quota_bytes, - "public_quota_bytes": u.public_quota_bytes, - "private_used_bytes": u.private_used_bytes, - "public_used_bytes": u.public_used_bytes, - "created_at": u.created_at.isoformat(), - } - for u in users - ] - - users = await execute_db_query(_list_users) + users_query = User.select().limit(limit).offset(offset) + users = [ + { + "id": u.id, + "username": u.username, + "email": u.email, + "email_verified": u.email_verified, + "is_active": u.is_active, + "private_quota_bytes": u.private_quota_bytes, + "public_quota_bytes": u.public_quota_bytes, + "private_used_bytes": u.private_used_bytes, + "public_used_bytes": u.public_used_bytes, + "created_at": u.created_at.isoformat(), + } + for u in users_query + ] return {"users": users, "limit": limit, "offset": offset} @@ -212,38 +213,33 @@ async def create_user_admin( HTTPException: If username or email already exists """ - # Check if user already exists - def _check_exists(): - return ( - User.get_or_none(User.username == request.username), - User.get_or_none(User.email == request.email), - ) - - existing_username, existing_email = await execute_db_query(_check_exists) - - if existing_username: - raise HTTPException( - 400, detail={"error": f"Username already exists: {request.username}"} - ) - - if existing_email: - raise HTTPException( - 400, detail={"error": f"Email already exists: {request.email}"} - ) - # Hash password password_hash = bcrypt.hashpw(request.password.encode(), bcrypt.gensalt()).decode() - # Create user with quotas (defaults applied if not specified) - user = await create_user( - username=request.username, - email=request.email, - password_hash=password_hash, - email_verified=request.email_verified, - is_active=request.is_active, - private_quota_bytes=request.private_quota_bytes, - public_quota_bytes=request.public_quota_bytes, - ) + # Check if user already exists and create atomically + with db.atomic(): + existing_username = User.get_or_none(User.username == request.username) + if existing_username: + raise HTTPException( + 400, detail={"error": f"Username already exists: {request.username}"} + ) + + existing_email = User.get_or_none(User.email == request.email) + if existing_email: + raise HTTPException( + 400, detail={"error": f"Email already exists: {request.email}"} + ) + + # Create user with quotas (defaults applied if not specified) + user = create_user( + username=request.username, + email=request.email, + password_hash=password_hash, + email_verified=request.email_verified, + is_active=request.is_active, + private_quota_bytes=request.private_quota_bytes, + public_quota_bytes=request.public_quota_bytes, + ) logger.info(f"Admin created user: {user.username}") @@ -285,20 +281,14 @@ async def delete_user_admin( HTTPException: If user not found or has repositories (without force) """ - def _get_user_and_repos(): - user = User.get_or_none(User.username == username) - if not user: - return None, [] - - # Get all repositories owned by this user - repos = list(Repository.select().where(Repository.owner_id == user.id)) - return user, repos - - user, owned_repos = await execute_db_query(_get_user_and_repos) + user = User.get_or_none(User.username == username) if not user: raise HTTPException(404, detail={"error": f"User not found: {username}"}) + # Get all repositories owned by this user + owned_repos = list(Repository.select().where(Repository.owner_id == user.id)) + # Check if user owns repositories if owned_repos and not force: raise HTTPException( @@ -315,18 +305,14 @@ async def delete_user_admin( # Delete user's repositories if force=true deleted_repos = [] if owned_repos and force: - # Delete repositories concurrently - async def delete_repo_and_log(repo): - await delete_repository(repo) + # Delete repositories sequentially (sync DB operations) + for repo in owned_repos: + delete_repository(repo) logger.warning(f"Admin deleted repository: {repo.full_id}") - return f"{repo.repo_type}:{repo.full_id}" + deleted_repos.append(f"{repo.repo_type}:{repo.full_id}") - deleted_repos = await asyncio.gather( - *[delete_repo_and_log(repo) for repo in owned_repos] - ) - - # Delete user - await delete_user(user) + # Delete user (already has db.atomic() inside) + delete_user(user) logger.warning( f"Admin deleted user: {username} (deleted {len(deleted_repos)} repositories)" @@ -358,19 +344,13 @@ async def set_email_verification( HTTPException: If user not found """ - def _update_user(): - user = User.get_or_none(User.username == username) - if not user: - return None - user.email_verified = verified - user.save() - return user - - user = await execute_db_query(_update_user) - + user = User.get_or_none(User.username == username) if not user: raise HTTPException(404, detail={"error": f"User not found: {username}"}) + user.email_verified = verified + user.save() + logger.info(f"Admin set email_verified={verified} for user: {username}") return { @@ -404,13 +384,10 @@ async def get_quota_admin( """ # Check if namespace exists - def _check_exists(): - if is_org: - return Organization.get_or_none(Organization.name == namespace) - else: - return User.get_or_none(User.username == namespace) - - entity = await execute_db_query(_check_exists) + if is_org: + entity = Organization.get_or_none(Organization.name == namespace) + else: + entity = User.get_or_none(User.username == namespace) if not entity: raise HTTPException( @@ -420,7 +397,7 @@ async def get_quota_admin( }, ) - info = await get_storage_info(namespace, is_org) + info = get_storage_info(namespace, is_org) return { "namespace": namespace, @@ -452,13 +429,10 @@ async def set_quota_admin( """ # Check if namespace exists - def _check_exists(): - if is_org: - return Organization.get_or_none(Organization.name == namespace) - else: - return User.get_or_none(User.username == namespace) - - entity = await execute_db_query(_check_exists) + if is_org: + entity = Organization.get_or_none(Organization.name == namespace) + else: + entity = User.get_or_none(User.username == namespace) if not entity: raise HTTPException( @@ -468,7 +442,7 @@ async def set_quota_admin( }, ) - info = await set_quota( + info = set_quota( namespace, private_quota_bytes=request.private_quota_bytes, public_quota_bytes=request.public_quota_bytes, @@ -508,13 +482,10 @@ async def recalculate_quota_admin( """ # Check if namespace exists - def _check_exists(): - if is_org: - return Organization.get_or_none(Organization.name == namespace) - else: - return User.get_or_none(User.username == namespace) - - entity = await execute_db_query(_check_exists) + if is_org: + entity = Organization.get_or_none(Organization.name == namespace) + else: + entity = User.get_or_none(User.username == namespace) if not entity: raise HTTPException( @@ -528,8 +499,8 @@ async def recalculate_quota_admin( f"Admin recalculating storage for {'org' if is_org else 'user'} {namespace}" ) - storage = await update_namespace_storage(namespace, is_org) - info = await get_storage_info(namespace, is_org) + storage = update_namespace_storage(namespace, is_org) + info = get_storage_info(namespace, is_org) return { "namespace": namespace, @@ -555,30 +526,21 @@ async def get_system_stats( System statistics """ - def _get_stats(): - user_count = User.select().count() - org_count = Organization.select().count() - repo_count = Repository.select().count() - private_repo_count = ( - Repository.select().where(Repository.private == True).count() - ) - public_repo_count = ( - Repository.select().where(Repository.private == False).count() - ) + user_count = User.select().count() + org_count = Organization.select().count() + repo_count = Repository.select().count() + private_repo_count = Repository.select().where(Repository.private == True).count() + public_repo_count = Repository.select().where(Repository.private == False).count() - return { - "users": user_count, - "organizations": org_count, - "repositories": { - "total": repo_count, - "private": private_repo_count, - "public": public_repo_count, - }, - } - - stats = await execute_db_query(_get_stats) - - return stats + return { + "users": user_count, + "organizations": org_count, + "repositories": { + "total": repo_count, + "private": private_repo_count, + "public": public_repo_count, + }, + } # ===== Repository Management Endpoints ===== @@ -605,38 +567,33 @@ async def list_repositories_admin( List of repositories with metadata """ - def _list_repos(): - query = Repository.select() - if repo_type: - query = query.where(Repository.repo_type == repo_type) - if namespace: - query = query.where(Repository.namespace == namespace) + query = Repository.select() + if repo_type: + query = query.where(Repository.repo_type == repo_type) + if namespace: + query = query.where(Repository.namespace == namespace) - query = query.order_by(Repository.created_at.desc()).limit(limit).offset(offset) + query = query.order_by(Repository.created_at.desc()).limit(limit).offset(offset) - repos = [] - for repo in query: - # Get owner username - owner = User.get_or_none(User.id == repo.owner_id) - owner_username = owner.username if owner else "unknown" + repos = [] + for repo in query: + # Get owner username + owner = User.get_or_none(User.id == repo.owner_id) + owner_username = owner.username if owner else "unknown" - repos.append( - { - "id": repo.id, - "repo_type": repo.repo_type, - "namespace": repo.namespace, - "name": repo.name, - "full_id": repo.full_id, - "private": repo.private, - "owner_id": repo.owner_id, - "owner_username": owner_username, - "created_at": repo.created_at.isoformat(), - } - ) - - return repos - - repos = await execute_db_query(_list_repos) + repos.append( + { + "id": repo.id, + "repo_type": repo.repo_type, + "namespace": repo.namespace, + "name": repo.name, + "full_id": repo.full_id, + "private": repo.private, + "owner_id": repo.owner_id, + "owner_username": owner_username, + "created_at": repo.created_at.isoformat(), + } + ) return {"repositories": repos, "limit": limit, "offset": offset} @@ -663,65 +620,55 @@ async def get_repository_admin( HTTPException: If repository not found """ - def _get_repo_details(): - repo = Repository.get_or_none( - Repository.repo_type == repo_type, - Repository.namespace == namespace, - Repository.name == name, - ) + repo = Repository.get_or_none( + Repository.repo_type == repo_type, + Repository.namespace == namespace, + Repository.name == name, + ) - if not repo: - return None - - # Get owner - owner = User.get_or_none(User.id == repo.owner_id) - - # Count files - file_count = File.select().where(File.repo_full_id == repo.full_id).count() - - # Count commits - commit_count = ( - Commit.select() - .where( - Commit.repo_full_id == repo.full_id, Commit.repo_type == repo.repo_type - ) - .count() - ) - - # Get total file size - total_size = ( - File.select() - .where(File.repo_full_id == repo.full_id) - .select(File.size) - .scalar(as_tuple=False) - ) - if total_size is None: - total_size = 0 - - return { - "id": repo.id, - "repo_type": repo.repo_type, - "namespace": repo.namespace, - "name": repo.name, - "full_id": repo.full_id, - "private": repo.private, - "owner_id": repo.owner_id, - "owner_username": owner.username if owner else "unknown", - "created_at": repo.created_at.isoformat(), - "file_count": file_count, - "commit_count": commit_count, - "total_size": total_size, - } - - repo_info = await execute_db_query(_get_repo_details) - - if not repo_info: + if not repo: raise HTTPException( 404, detail={"error": f"Repository not found: {repo_type}/{namespace}/{name}"}, ) - return repo_info + # Get owner + owner = User.get_or_none(User.id == repo.owner_id) + + # Count files + file_count = File.select().where(File.repo_full_id == repo.full_id).count() + + # Count commits + commit_count = ( + Commit.select() + .where(Commit.repo_full_id == repo.full_id, Commit.repo_type == repo.repo_type) + .count() + ) + + # Get total file size + total_size = ( + File.select() + .where(File.repo_full_id == repo.full_id) + .select(File.size) + .scalar(as_tuple=False) + ) + if total_size is None: + total_size = 0 + + return { + "id": repo.id, + "repo_type": repo.repo_type, + "namespace": repo.namespace, + "name": repo.name, + "full_id": repo.full_id, + "private": repo.private, + "owner_id": repo.owner_id, + "owner_username": owner.username if owner else "unknown", + "created_at": repo.created_at.isoformat(), + "file_count": file_count, + "commit_count": commit_count, + "total_size": total_size, + } # ===== Commit History Endpoints ===== @@ -748,36 +695,31 @@ async def list_commits_admin( List of commits """ - def _list_commits(): - query = Commit.select() + query = Commit.select() - if repo_full_id: - query = query.where(Commit.repo_full_id == repo_full_id) - if username: - query = query.where(Commit.username == username) + if repo_full_id: + query = query.where(Commit.repo_full_id == repo_full_id) + if username: + query = query.where(Commit.username == username) - query = query.order_by(Commit.created_at.desc()).limit(limit).offset(offset) + query = query.order_by(Commit.created_at.desc()).limit(limit).offset(offset) - commits = [] - for commit in query: - commits.append( - { - "id": commit.id, - "commit_id": commit.commit_id, - "repo_full_id": commit.repo_full_id, - "repo_type": commit.repo_type, - "branch": commit.branch, - "user_id": commit.user_id, - "username": commit.username, - "message": commit.message, - "description": commit.description, - "created_at": commit.created_at.isoformat(), - } - ) - - return commits - - commits = await execute_db_query(_list_commits) + commits = [] + for commit in query: + commits.append( + { + "id": commit.id, + "commit_id": commit.commit_id, + "repo_full_id": commit.repo_full_id, + "repo_type": commit.repo_type, + "branch": commit.branch, + "user_id": commit.user_id, + "username": commit.username, + "message": commit.message, + "description": commit.description, + "created_at": commit.created_at.isoformat(), + } + ) return {"commits": commits, "limit": limit, "offset": offset} @@ -912,98 +854,89 @@ async def get_detailed_stats( Detailed statistics from database """ - def _get_detailed_stats(): - # User stats - total_users = User.select().count() - active_users = User.select().where(User.is_active == True).count() - verified_users = User.select().where(User.email_verified == True).count() + # User stats + total_users = User.select().count() + active_users = User.select().where(User.is_active == True).count() + verified_users = User.select().where(User.email_verified == True).count() - # Organization stats - total_orgs = Organization.select().count() + # Organization stats + total_orgs = Organization.select().count() - # Repository stats - total_repos = Repository.select().count() - private_repos = Repository.select().where(Repository.private == True).count() - public_repos = Repository.select().where(Repository.private == False).count() + # Repository stats + total_repos = Repository.select().count() + private_repos = Repository.select().where(Repository.private == True).count() + public_repos = Repository.select().where(Repository.private == False).count() - # Repos by type - model_repos = Repository.select().where(Repository.repo_type == "model").count() - dataset_repos = ( - Repository.select().where(Repository.repo_type == "dataset").count() - ) - space_repos = Repository.select().where(Repository.repo_type == "space").count() + # Repos by type + model_repos = Repository.select().where(Repository.repo_type == "model").count() + dataset_repos = Repository.select().where(Repository.repo_type == "dataset").count() + space_repos = Repository.select().where(Repository.repo_type == "space").count() - # Commit stats - total_commits = Commit.select().count() + # Commit stats + total_commits = Commit.select().count() - # Top contributors - top_contributors = ( - Commit.select(Commit.username, fn.COUNT(Commit.id).alias("commit_count")) - .group_by(Commit.username) - .order_by(fn.COUNT(Commit.id).desc()) - .limit(10) - ) + # Top contributors + top_contributors = ( + Commit.select(Commit.username, fn.COUNT(Commit.id).alias("commit_count")) + .group_by(Commit.username) + .order_by(fn.COUNT(Commit.id).desc()) + .limit(10) + ) - contributors = [ - {"username": c.username, "commit_count": c.commit_count} - for c in top_contributors - ] + contributors = [ + {"username": c.username, "commit_count": c.commit_count} + for c in top_contributors + ] - # LFS object stats - total_lfs_objects = LFSObjectHistory.select().count() - total_lfs_size = ( - LFSObjectHistory.select( - fn.SUM(LFSObjectHistory.size).alias("total") - ).scalar() - or 0 - ) + # LFS object stats + total_lfs_objects = LFSObjectHistory.select().count() + total_lfs_size = ( + LFSObjectHistory.select(fn.SUM(LFSObjectHistory.size).alias("total")).scalar() + or 0 + ) - # Storage stats - total_private_used = ( - User.select(fn.SUM(User.private_used_bytes).alias("total")).scalar() or 0 - ) - total_public_used = ( - User.select(fn.SUM(User.public_used_bytes).alias("total")).scalar() or 0 - ) + # Storage stats + total_private_used = ( + User.select(fn.SUM(User.private_used_bytes).alias("total")).scalar() or 0 + ) + total_public_used = ( + User.select(fn.SUM(User.public_used_bytes).alias("total")).scalar() or 0 + ) - return { - "users": { - "total": total_users, - "active": active_users, - "verified": verified_users, - "inactive": total_users - active_users, + return { + "users": { + "total": total_users, + "active": active_users, + "verified": verified_users, + "inactive": total_users - active_users, + }, + "organizations": { + "total": total_orgs, + }, + "repositories": { + "total": total_repos, + "private": private_repos, + "public": public_repos, + "by_type": { + "model": model_repos, + "dataset": dataset_repos, + "space": space_repos, }, - "organizations": { - "total": total_orgs, - }, - "repositories": { - "total": total_repos, - "private": private_repos, - "public": public_repos, - "by_type": { - "model": model_repos, - "dataset": dataset_repos, - "space": space_repos, - }, - }, - "commits": { - "total": total_commits, - "top_contributors": contributors, - }, - "lfs": { - "total_objects": total_lfs_objects, - "total_size": total_lfs_size, - }, - "storage": { - "private_used": total_private_used, - "public_used": total_public_used, - "total_used": total_private_used + total_public_used, - }, - } - - stats = await execute_db_query(_get_detailed_stats) - - return stats + }, + "commits": { + "total": total_commits, + "top_contributors": contributors, + }, + "lfs": { + "total_objects": total_lfs_objects, + "total_size": total_lfs_size, + }, + "storage": { + "private_used": total_private_used, + "public_used": total_public_used, + "total_used": total_private_used + total_public_used, + }, + } @router.get("/stats/timeseries") @@ -1021,60 +954,55 @@ async def get_timeseries_stats( Time-series data for various metrics """ - def _get_timeseries(): - cutoff_date = datetime.now(timezone.utc) - timedelta(days=days) + cutoff_date = datetime.now(timezone.utc) - timedelta(days=days) - # Repositories created per day - repos_by_day = ( - Repository.select( - Repository.created_at, - Repository.repo_type, - ) - .where(Repository.created_at >= cutoff_date) - .order_by(Repository.created_at.asc()) + # Repositories created per day + repos_by_day = ( + Repository.select( + Repository.created_at, + Repository.repo_type, ) + .where(Repository.created_at >= cutoff_date) + .order_by(Repository.created_at.asc()) + ) - # Group by date - daily_repos = {} - for repo in repos_by_day: - date_key = repo.created_at.date().isoformat() - if date_key not in daily_repos: - daily_repos[date_key] = {"model": 0, "dataset": 0, "space": 0} - daily_repos[date_key][repo.repo_type] += 1 + # Group by date + daily_repos = {} + for repo in repos_by_day: + date_key = repo.created_at.date().isoformat() + if date_key not in daily_repos: + daily_repos[date_key] = {"model": 0, "dataset": 0, "space": 0} + daily_repos[date_key][repo.repo_type] += 1 - # Commits created per day - commits_by_day = ( - Commit.select(Commit.created_at) - .where(Commit.created_at >= cutoff_date) - .order_by(Commit.created_at.asc()) - ) + # Commits created per day + commits_by_day = ( + Commit.select(Commit.created_at) + .where(Commit.created_at >= cutoff_date) + .order_by(Commit.created_at.asc()) + ) - daily_commits = {} - for commit in commits_by_day: - date_key = commit.created_at.date().isoformat() - daily_commits[date_key] = daily_commits.get(date_key, 0) + 1 + daily_commits = {} + for commit in commits_by_day: + date_key = commit.created_at.date().isoformat() + daily_commits[date_key] = daily_commits.get(date_key, 0) + 1 - # Users created per day - users_by_day = ( - User.select(User.created_at) - .where(User.created_at >= cutoff_date) - .order_by(User.created_at.asc()) - ) + # Users created per day + users_by_day = ( + User.select(User.created_at) + .where(User.created_at >= cutoff_date) + .order_by(User.created_at.asc()) + ) - daily_users = {} - for user in users_by_day: - date_key = user.created_at.date().isoformat() - daily_users[date_key] = daily_users.get(date_key, 0) + 1 + daily_users = {} + for user in users_by_day: + date_key = user.created_at.date().isoformat() + daily_users[date_key] = daily_users.get(date_key, 0) + 1 - return { - "repositories_by_day": daily_repos, - "commits_by_day": daily_commits, - "users_by_day": daily_users, - } - - stats = await execute_db_query(_get_timeseries) - - return stats + return { + "repositories_by_day": daily_repos, + "commits_by_day": daily_commits, + "users_by_day": daily_users, + } @router.get("/stats/top-repos") @@ -1094,68 +1022,63 @@ async def get_top_repositories( List of top repositories """ - def _get_top_repos(): - if by == "commits": - # Top repos by commit count - top_repos = ( - Commit.select( - Commit.repo_full_id, - Commit.repo_type, - fn.COUNT(Commit.id).alias("count"), - ) - .group_by(Commit.repo_full_id, Commit.repo_type) - .order_by(fn.COUNT(Commit.id).desc()) - .limit(limit) + if by == "commits": + # Top repos by commit count + top_repos = ( + Commit.select( + Commit.repo_full_id, + Commit.repo_type, + fn.COUNT(Commit.id).alias("count"), + ) + .group_by(Commit.repo_full_id, Commit.repo_type) + .order_by(fn.COUNT(Commit.id).desc()) + .limit(limit) + ) + + result = [] + for item in top_repos: + repo = Repository.get_or_none( + Repository.full_id == item.repo_full_id, + Repository.repo_type == item.repo_type, + ) + result.append( + { + "repo_full_id": item.repo_full_id, + "repo_type": item.repo_type, + "commit_count": item.count, + "private": repo.private if repo else False, + } ) - result = [] - for item in top_repos: - repo = Repository.get_or_none( + else: # by size + # Top repos by total file size + top_repos = ( + File.select(File.repo_full_id, fn.SUM(File.size).alias("total_size")) + .group_by(File.repo_full_id) + .order_by(fn.SUM(File.size).desc()) + .limit(limit) + ) + + result = [] + for item in top_repos: + # Find repo in any type + repo = None + for repo_type in ["model", "dataset", "space"]: + r = Repository.get_or_none( Repository.full_id == item.repo_full_id, - Repository.repo_type == item.repo_type, - ) - result.append( - { - "repo_full_id": item.repo_full_id, - "repo_type": item.repo_type, - "commit_count": item.count, - "private": repo.private if repo else False, - } + Repository.repo_type == repo_type, ) + if r: + repo = r + break - else: # by size - # Top repos by total file size - top_repos = ( - File.select(File.repo_full_id, fn.SUM(File.size).alias("total_size")) - .group_by(File.repo_full_id) - .order_by(fn.SUM(File.size).desc()) - .limit(limit) + result.append( + { + "repo_full_id": item.repo_full_id, + "repo_type": repo.repo_type if repo else "unknown", + "total_size": item.total_size, + "private": repo.private if repo else False, + } ) - result = [] - for item in top_repos: - # Find repo in any type - repo = None - for repo_type in ["model", "dataset", "space"]: - r = Repository.get_or_none( - Repository.full_id == item.repo_full_id, - Repository.repo_type == repo_type, - ) - if r: - repo = r - break - - result.append( - { - "repo_full_id": item.repo_full_id, - "repo_type": repo.repo_type if repo else "unknown", - "total_size": item.total_size, - "private": repo.private if repo else False, - } - ) - - return result - - repos = await execute_db_query(_get_top_repos) - - return {"top_repositories": repos, "sorted_by": by} + return {"top_repositories": result, "sorted_by": by} diff --git a/src/kohakuhub/api/branches.py b/src/kohakuhub/api/branches.py index a63aab6..28b5d6a 100644 --- a/src/kohakuhub/api/branches.py +++ b/src/kohakuhub/api/branches.py @@ -6,8 +6,14 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from kohakuhub.db import Repository, User +from kohakuhub.db_operations import create_commit, get_repository from kohakuhub.logger import get_logger -from kohakuhub.db_async import create_commit, get_repository +from kohakuhub.auth.dependencies import get_current_user +from kohakuhub.auth.permissions import ( + check_repo_delete_permission, + check_repo_write_permission, +) +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.api.repo.utils.gc import ( check_commit_range_recoverability, check_lfs_recoverability, @@ -20,12 +26,6 @@ from kohakuhub.api.repo.utils.hf import ( hf_repo_not_found, hf_server_error, ) -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.auth.dependencies import get_current_user -from kohakuhub.auth.permissions import ( - check_repo_delete_permission, - check_repo_write_permission, -) logger = get_logger("BRANCHES") @@ -62,9 +62,7 @@ async def create_branch( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) - ) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -131,7 +129,7 @@ async def delete_branch( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -219,7 +217,7 @@ async def create_tag( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -275,7 +273,7 @@ async def delete_tag( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -326,7 +324,7 @@ async def revert_branch( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -409,7 +407,7 @@ async def revert_branch( # Record commit in database commit_msg = payload.message or f"Revert commit {commit_id[:8]}" try: - await create_commit( + create_commit( commit_id=new_commit_id, repo_full_id=repo_id, repo_type=repo_type, @@ -468,7 +466,7 @@ async def merge_branches( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -538,7 +536,7 @@ async def merge_branches( payload.message or f"Merge {source_ref} into {destination_branch}" ) try: - await create_commit( + create_commit( commit_id=merge_commit_id, repo_full_id=repo_id, repo_type=repo_type, @@ -598,7 +596,7 @@ async def reset_branch( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -787,7 +785,7 @@ async def reset_branch( # Record reset commit in database try: - await create_commit( + create_commit( commit_id=commit_result["id"], repo_full_id=repo_id, repo_type=repo_type, diff --git a/src/kohakuhub/api/commit/__init__.py b/src/kohakuhub/api/commit/__init__.py index a604f22..b24d6b8 100644 --- a/src/kohakuhub/api/commit/__init__.py +++ b/src/kohakuhub/api/commit/__init__.py @@ -1,6 +1,6 @@ """Commit management module.""" -from kohakuhub.api.commit.routers.operations import router from kohakuhub.api.commit.routers import history +from kohakuhub.api.commit.routers.operations import router __all__ = ["router", "history"] diff --git a/src/kohakuhub/api/commit/routers/history.py b/src/kohakuhub/api/commit/routers/history.py index 0782a0d..00d4181 100644 --- a/src/kohakuhub/api/commit/routers/history.py +++ b/src/kohakuhub/api/commit/routers/history.py @@ -1,21 +1,21 @@ """Commit history API endpoints.""" +from typing import Optional import asyncio import difflib -from typing import Optional -import httpx from fastapi import APIRouter, Depends, Query +import httpx -from kohakuhub.api.repo.utils.hf import hf_repo_not_found, hf_server_error -from kohakuhub.utils.lakefs import lakefs_repo_name -from kohakuhub.auth.dependencies import get_optional_user -from kohakuhub.auth.permissions import check_repo_read_permission from kohakuhub.config import cfg -from kohakuhub.db_async import execute_db_query, get_commit, get_repository from kohakuhub.db import Commit, Repository, User +from kohakuhub.db_operations import get_commit, get_repository from kohakuhub.lakefs_rest_client import get_lakefs_rest_client from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_optional_user +from kohakuhub.auth.permissions import check_repo_read_permission +from kohakuhub.utils.lakefs import lakefs_repo_name +from kohakuhub.api.repo.utils.hf import hf_repo_not_found, hf_server_error logger = get_logger("COMMITS") router = APIRouter() @@ -48,7 +48,7 @@ async def list_commits( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -92,16 +92,13 @@ async def list_commits( commit_ids = [c["id"] for c in log_result["results"]] # Fetch our commit records (with actual user info) - def _get_our_commits(): - return { - c.commit_id: c - for c in Commit.select().where( - Commit.commit_id.in_(commit_ids), - Commit.repo_full_id == repo_id, - ) - } - - our_commits = await execute_db_query(_get_our_commits) + our_commits = { + c.commit_id: c + for c in Commit.select().where( + Commit.commit_id.in_(commit_ids), + Commit.repo_full_id == repo_id, + ) + } for commit in log_result["results"]: try: @@ -171,7 +168,7 @@ async def get_commit_detail( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -189,7 +186,7 @@ async def get_commit_detail( ) # Get our commit record (with actual user info) - our_commit = await get_commit(commit_id, repo_id) + our_commit = get_commit(commit_id, repo_id) # Build response response = { @@ -242,7 +239,7 @@ async def get_commit_diff( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -283,17 +280,17 @@ async def get_commit_diff( file_paths = [d["path"] for d in diff_results] # Fetch File records for LFS status - def _get_file_records(): + if file_paths: from kohakuhub.db import File - return { + file_records = { f.path_in_repo: f for f in File.select().where( File.repo_full_id == repo_id, File.path_in_repo.in_(file_paths) ) } - - file_records = await execute_db_query(_get_file_records) if file_paths else {} + else: + file_records = {} # Process diff entries concurrently async def process_diff_entry(diff_entry): @@ -454,7 +451,7 @@ async def get_commit_diff( ) # Get our commit record for user info - our_commit = await get_commit(commit_id, repo_id) + our_commit = get_commit(commit_id, repo_id) response = { "commit_id": commit_id, diff --git a/src/kohakuhub/api/commit/routers/operations.py b/src/kohakuhub/api/commit/routers/operations.py index 583659b..6c51a90 100644 --- a/src/kohakuhub/api/commit/routers/operations.py +++ b/src/kohakuhub/api/commit/routers/operations.py @@ -1,31 +1,30 @@ """Commit creation endpoint - Refactored version with smaller functions.""" +from datetime import datetime, timezone +from enum import Enum import asyncio import base64 import hashlib import json -from datetime import datetime, timezone -from enum import Enum from fastapi import APIRouter, Depends, HTTPException, Request -from kohakuhub.api.repo.utils.gc import run_gc_for_file, track_lfs_object -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.api.quota.util import update_namespace_storage -from kohakuhub.utils.s3 import get_object_metadata, object_exists -from kohakuhub.auth.dependencies import get_current_user -from kohakuhub.auth.permissions import check_repo_write_permission from kohakuhub.config import cfg -from kohakuhub.db_async import ( +from kohakuhub.db import File, Organization, Repository, User +from kohakuhub.db_operations import ( create_commit, create_file, delete_file, - execute_db_query, get_file, update_file, ) -from kohakuhub.db import File, Organization, Repository, User from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user +from kohakuhub.auth.permissions import check_repo_write_permission +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name +from kohakuhub.utils.s3 import get_object_metadata, object_exists +from kohakuhub.api.quota.util import update_namespace_storage +from kohakuhub.api.repo.utils.gc import run_gc_for_file, track_lfs_object logger = get_logger("FILE") router = APIRouter() @@ -94,7 +93,7 @@ async def process_regular_file( git_blob_sha1 = calculate_git_blob_sha1(data) # Check if file unchanged (deduplication) - existing = await get_file(repo_id, path) + existing = get_file(repo_id, path) if existing and existing.sha256 == git_blob_sha1 and existing.size == len(data): logger.info(f"Skipping unchanged file: {path}") return False @@ -115,23 +114,20 @@ async def process_regular_file( raise HTTPException(500, detail={"error": f"Failed to upload {path}: {e}"}) # Update database - store git blob SHA1 in sha256 column for non-LFS files - def _update_db(): - File.insert( - repo_full_id=repo_id, - path_in_repo=path, - size=len(data), - sha256=git_blob_sha1, - lfs=False, - ).on_conflict( - conflict_target=(File.repo_full_id, File.path_in_repo), - update={ - File.sha256: git_blob_sha1, - File.size: len(data), - File.updated_at: datetime.now(timezone.utc), - }, - ).execute() - - await execute_db_query(_update_db) + File.insert( + repo_full_id=repo_id, + path_in_repo=path, + size=len(data), + sha256=git_blob_sha1, + lfs=False, + ).on_conflict( + conflict_target=(File.repo_full_id, File.path_in_repo), + update={ + File.sha256: git_blob_sha1, + File.size: len(data), + File.updated_at: datetime.now(timezone.utc), + }, + ).execute() return True @@ -166,7 +162,7 @@ async def process_lfs_file( raise HTTPException(400, detail={"error": f"Missing OID for LFS file {path}"}) # Check for existing file - existing = await get_file(repo_id, path) + existing = get_file(repo_id, path) # Track old LFS object for potential deletion old_lfs_oid = None @@ -238,24 +234,21 @@ async def process_lfs_file( ) # Update database - def _update_db(): - File.insert( - repo_full_id=repo_id, - path_in_repo=path, - size=size, - sha256=oid, - lfs=True, - ).on_conflict( - conflict_target=(File.repo_full_id, File.path_in_repo), - update={ - File.sha256: oid, - File.size: size, - File.lfs: True, - File.updated_at: datetime.now(timezone.utc), - }, - ).execute() - - await execute_db_query(_update_db) + File.insert( + repo_full_id=repo_id, + path_in_repo=path, + size=size, + sha256=oid, + lfs=True, + ).on_conflict( + conflict_target=(File.repo_full_id, File.path_in_repo), + update={ + File.sha256: oid, + File.size: size, + File.lfs: True, + File.updated_at: datetime.now(timezone.utc), + }, + ).execute() logger.success(f"Updated database record for LFS file: {path}") @@ -293,14 +286,11 @@ async def process_deleted_file( logger.warning(f"Failed to delete {path} from LakeFS: {e}") # Remove from database - def _delete_from_db(): - return ( - File.delete() - .where((File.repo_full_id == repo_id) & (File.path_in_repo == path)) - .execute() - ) - - deleted_count = await execute_db_query(_delete_from_db) + deleted_count = ( + File.delete() + .where((File.repo_full_id == repo_id) & (File.path_in_repo == path)) + .execute() + ) if deleted_count > 0: logger.success(f"Removed {path} from database") @@ -377,18 +367,14 @@ async def process_deleted_folder( # Remove from database if deleted_files: - - def _delete_from_db(): - return ( - File.delete() - .where( - (File.repo_full_id == repo_id) - & (File.path_in_repo.startswith(folder_path)) - ) - .execute() + deleted_count = ( + File.delete() + .where( + (File.repo_full_id == repo_id) + & (File.path_in_repo.startswith(folder_path)) ) - - deleted_count = await execute_db_query(_delete_from_db) + .execute() + ) logger.success(f"Removed {deleted_count} records from database") except Exception as e: @@ -458,45 +444,42 @@ async def process_copy_file( ) # Update database - copy file metadata - src_file = await get_file(repo_id, src_path) + src_file = get_file(repo_id, src_path) - def _update_db(): - if src_file: - File.insert( - repo_full_id=repo_id, - path_in_repo=dest_path, - size=src_file.size, - sha256=src_file.sha256, - lfs=src_file.lfs, - ).on_conflict( - conflict_target=(File.repo_full_id, File.path_in_repo), - update={ - File.sha256: src_file.sha256, - File.size: src_file.size, - File.lfs: src_file.lfs, - File.updated_at: datetime.now(timezone.utc), - }, - ).execute() - else: - # If not in database, create entry based on LakeFS info - is_lfs = src_obj["size_bytes"] >= cfg.app.lfs_threshold_bytes - File.insert( - repo_full_id=repo_id, - path_in_repo=dest_path, - size=src_obj["size_bytes"], - sha256=src_obj["checksum"], - lfs=is_lfs, - ).on_conflict( - conflict_target=(File.repo_full_id, File.path_in_repo), - update={ - File.sha256: src_obj["checksum"], - File.size: src_obj["size_bytes"], - File.lfs: is_lfs, - File.updated_at: datetime.now(timezone.utc), - }, - ).execute() - - await execute_db_query(_update_db) + if src_file: + File.insert( + repo_full_id=repo_id, + path_in_repo=dest_path, + size=src_file.size, + sha256=src_file.sha256, + lfs=src_file.lfs, + ).on_conflict( + conflict_target=(File.repo_full_id, File.path_in_repo), + update={ + File.sha256: src_file.sha256, + File.size: src_file.size, + File.lfs: src_file.lfs, + File.updated_at: datetime.now(timezone.utc), + }, + ).execute() + else: + # If not in database, create entry based on LakeFS info + is_lfs = src_obj["size_bytes"] >= cfg.app.lfs_threshold_bytes + File.insert( + repo_full_id=repo_id, + path_in_repo=dest_path, + size=src_obj["size_bytes"], + sha256=src_obj["checksum"], + lfs=is_lfs, + ).on_conflict( + conflict_target=(File.repo_full_id, File.path_in_repo), + update={ + File.sha256: src_obj["checksum"], + File.size: src_obj["size_bytes"], + File.lfs: is_lfs, + File.updated_at: datetime.now(timezone.utc), + }, + ).execute() logger.success(f"Successfully copied {src_path} to {dest_path}") @@ -542,12 +525,9 @@ async def commit( repo_id = f"{namespace}/{name}" # Check repository exists and write permission - def _get_repo(): - return Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type.value) - ) - - repo_row = await execute_db_query(_get_repo) + repo_row = Repository.get_or_none( + (Repository.full_id == repo_id) & (Repository.repo_type == repo_type.value) + ) if not repo_row: raise HTTPException(404, detail={"error": "Repository not found"}) @@ -692,7 +672,7 @@ async def commit( # Record commit in our database (track the actual user) try: - await create_commit( + create_commit( commit_id=commit_result["id"], repo_full_id=repo_id, repo_type=repo_type, @@ -736,14 +716,11 @@ async def commit( # Update storage usage for namespace after successful commit try: # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Recalculate storage usage - await update_namespace_storage(namespace, is_org) + update_namespace_storage(namespace, is_org) logger.debug( f"Updated storage usage for {'org' if is_org else 'user'} {namespace}" ) diff --git a/src/kohakuhub/api/files.py b/src/kohakuhub/api/files.py index 604c041..3b41c99 100644 --- a/src/kohakuhub/api/files.py +++ b/src/kohakuhub/api/files.py @@ -12,24 +12,24 @@ from urllib.parse import quote from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse, Response +from kohakuhub.config import cfg +from kohakuhub.db import File, Organization, Repository, User +from kohakuhub.db_operations import get_file, get_repository +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user, get_optional_user +from kohakuhub.auth.permissions import ( + check_repo_read_permission, + check_repo_write_permission, +) +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name +from kohakuhub.utils.s3 import generate_download_presigned_url, parse_s3_uri +from kohakuhub.api.quota.util import check_quota from kohakuhub.api.repo.utils.hf import ( hf_repo_not_found, hf_revision_not_found, hf_server_error, is_lakefs_not_found_error, ) -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.api.quota.util import check_quota -from kohakuhub.utils.s3 import generate_download_presigned_url, parse_s3_uri -from kohakuhub.auth.dependencies import get_current_user, get_optional_user -from kohakuhub.auth.permissions import ( - check_repo_read_permission, - check_repo_write_permission, -) -from kohakuhub.config import cfg -from kohakuhub.db_async import execute_db_query, get_file, get_repository -from kohakuhub.db import File, Organization, Repository, User -from kohakuhub.logger import get_logger logger = get_logger("FILE") router = APIRouter() @@ -58,7 +58,7 @@ async def check_file_by_sha256(repo_id: str, path: str, sha256: str, size: int) Returns: True if file should be ignored (already exists), False otherwise """ - existing = await get_file(repo_id, path) + existing = get_file(repo_id, path) if existing and existing.sha256 == sha256 and existing.size == size: return True return False @@ -189,7 +189,7 @@ async def preupload( """ repo_id = f"{namespace}/{name}" # Verify repository exists - repo_row = await get_repository(repo_type.value, namespace, name) + repo_row = get_repository(repo_type.value, namespace, name) if not repo_row: raise HTTPException(404, detail={"error": "Repository not found"}) @@ -214,17 +214,12 @@ async def preupload( total_upload_bytes = sum(int(f.get("size", 0)) for f in files) # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Check storage quota before upload (based on repo privacy) is_private = repo_row.private - allowed, error_msg = await check_quota( - namespace, total_upload_bytes, is_private, is_org - ) + allowed, error_msg = check_quota(namespace, total_upload_bytes, is_private, is_org) if not allowed: raise HTTPException( status_code=413, # Payload Too Large @@ -275,7 +270,7 @@ async def get_revision( """ repo_id = f"{namespace}/{name}" # Check if repository exists in database first - repo_row = await get_repository(repo_type.value, namespace, name) + repo_row = get_repository(repo_type.value, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type.value) @@ -353,12 +348,9 @@ async def _get_file_metadata( repo_id = f"{namespace}/{name}" # Check repository exists and read permission - def _get_repo(): - return Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) - ) - - repo_row = await execute_db_query(_get_repo) + repo_row = Repository.get_or_none( + (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) + ) if repo_row: check_repo_read_permission(repo_row, user) @@ -402,7 +394,7 @@ async def _get_file_metadata( # Get correct checksum from database # sha256 column stores: git blob SHA1 for non-LFS, SHA256 for LFS - file_record = await get_file(repo_id, path) + file_record = get_file(repo_id, path) # HuggingFace expects plain SHA256 hex (64 characters, unquoted) # For non-LFS: use git blob SHA1, for LFS: use SHA256 diff --git a/src/kohakuhub/api/git/routers/http.py b/src/kohakuhub/api/git/routers/http.py index 655b800..31d545e 100644 --- a/src/kohakuhub/api/git/routers/http.py +++ b/src/kohakuhub/api/git/routers/http.py @@ -8,20 +8,19 @@ from datetime import datetime, timezone from fastapi import APIRouter, Header, HTTPException, Request, Response from kohakuhub.db import Repository, Token, User -from kohakuhub.db_async import execute_db_query from kohakuhub.logger import get_logger -from kohakuhub.api.git.utils.lakefs_bridge import GitLakeFSBridge -from kohakuhub.api.git.utils.server import ( - GitReceivePackHandler, - GitUploadPackHandler, - parse_git_credentials, -) from kohakuhub.auth.dependencies import get_optional_user from kohakuhub.auth.permissions import ( check_repo_read_permission, check_repo_write_permission, ) from kohakuhub.auth.utils import hash_token +from kohakuhub.api.git.utils.lakefs_bridge import GitLakeFSBridge +from kohakuhub.api.git.utils.server import ( + GitReceivePackHandler, + GitUploadPackHandler, + parse_git_credentials, +) logger = get_logger("GIT_HTTP") router = APIRouter() @@ -46,30 +45,21 @@ async def get_user_from_git_auth(authorization: str | None) -> User | None: # Hash token and lookup token_hash = hash_token(token_str) - def _get_token(): - return Token.get_or_none(Token.token_hash == token_hash) - - token = await execute_db_query(_get_token) + token = Token.get_or_none(Token.token_hash == token_hash) if not token: logger.debug(f"Invalid token for user {username}") return None # Get user - def _get_user(): - return User.get_or_none(User.id == token.user_id) - - user = await execute_db_query(_get_user) + user = User.get_or_none(User.id == token.user_id) if not user or not user.is_active: logger.warning(f"User {username} not found or inactive") return None # Update token last used - def _update_token(): - Token.update(last_used=datetime.now(timezone.utc)).where( - Token.id == token.id - ).execute() - - await execute_db_query(_update_token) + Token.update(last_used=datetime.now(timezone.utc)).where( + Token.id == token.id + ).execute() logger.info(f"Authenticated {username} via Git HTTP") return user @@ -99,23 +89,20 @@ async def git_info_refs( logger.info(f"Git info/refs: {service} for {repo_id}") # Get repository - try all repo types since we don't know from URL - def _get_repo(): - for repo_type in ["model", "dataset", "space"]: - repo = Repository.get_or_none( - Repository.namespace == namespace, - Repository.name == name, - Repository.repo_type == repo_type, - ) - if repo: - return repo - return None - - repo = await execute_db_query(_get_repo) + repo = None + for repo_type in ["model", "dataset", "space"]: + repo = Repository.get_or_none( + Repository.namespace == namespace, + Repository.name == name, + Repository.repo_type == repo_type, + ) + if repo: + break if not repo: raise HTTPException(404, detail="Repository not found") # Authenticate user - user = await get_user_from_git_auth(authorization) + user = get_user_from_git_auth(authorization) # Check permissions based on service if service == "git-upload-pack": @@ -174,23 +161,20 @@ async def git_upload_pack( logger.info(f"Git upload-pack for {repo_id}") # Get repository - try all repo types since we don't know from URL - def _get_repo(): - for repo_type in ["model", "dataset", "space"]: - repo = Repository.get_or_none( - Repository.namespace == namespace, - Repository.name == name, - Repository.repo_type == repo_type, - ) - if repo: - return repo - return None - - repo = await execute_db_query(_get_repo) + repo = None + for repo_type in ["model", "dataset", "space"]: + repo = Repository.get_or_none( + Repository.namespace == namespace, + Repository.name == name, + Repository.repo_type == repo_type, + ) + if repo: + break if not repo: raise HTTPException(404, detail="Repository not found") # Authenticate and check read permission - user = await get_user_from_git_auth(authorization) + user = get_user_from_git_auth(authorization) check_repo_read_permission(repo, user) # Read request body @@ -234,23 +218,20 @@ async def git_receive_pack( logger.info(f"Git receive-pack for {repo_id}") # Get repository - try all repo types since we don't know from URL - def _get_repo(): - for repo_type in ["model", "dataset", "space"]: - repo = Repository.get_or_none( - Repository.namespace == namespace, - Repository.name == name, - Repository.repo_type == repo_type, - ) - if repo: - return repo - return None - - repo = await execute_db_query(_get_repo) + repo = None + for repo_type in ["model", "dataset", "space"]: + repo = Repository.get_or_none( + Repository.namespace == namespace, + Repository.name == name, + Repository.repo_type == repo_type, + ) + if repo: + break if not repo: raise HTTPException(404, detail="Repository not found") # Authenticate and check write permission - user = await get_user_from_git_auth(authorization) + user = get_user_from_git_auth(authorization) if not user: raise HTTPException(401, detail="Authentication required for push") @@ -291,23 +272,20 @@ async def git_head( repo_id = f"{namespace}/{name}" # Get repository - try all repo types since we don't know from URL - def _get_repo(): - for repo_type in ["model", "dataset", "space"]: - repo = Repository.get_or_none( - Repository.namespace == namespace, - Repository.name == name, - Repository.repo_type == repo_type, - ) - if repo: - return repo - return None - - repo = await execute_db_query(_get_repo) + repo = None + for repo_type in ["model", "dataset", "space"]: + repo = Repository.get_or_none( + Repository.namespace == namespace, + Repository.name == name, + Repository.repo_type == repo_type, + ) + if repo: + break if not repo: raise HTTPException(404, detail="Repository not found") # Authenticate and check read permission - user = await get_user_from_git_auth(authorization) + user = get_user_from_git_auth(authorization) check_repo_read_permission(repo, user) # Return HEAD reference diff --git a/src/kohakuhub/api/git/routers/lfs.py b/src/kohakuhub/api/git/routers/lfs.py index f41497b..c7f7d39 100644 --- a/src/kohakuhub/api/git/routers/lfs.py +++ b/src/kohakuhub/api/git/routers/lfs.py @@ -13,22 +13,22 @@ from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import JSONResponse from pydantic import BaseModel -from kohakuhub.api.quota.util import check_quota +from kohakuhub.config import cfg +from kohakuhub.db import Organization, Repository, User +from kohakuhub.db_operations import get_file_by_sha256 +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_optional_user +from kohakuhub.auth.permissions import ( + check_repo_read_permission, + check_repo_write_permission, +) from kohakuhub.utils.s3 import ( generate_download_presigned_url, generate_upload_presigned_url, get_object_metadata, object_exists, ) -from kohakuhub.auth.dependencies import get_optional_user -from kohakuhub.auth.permissions import ( - check_repo_read_permission, - check_repo_write_permission, -) -from kohakuhub.config import cfg -from kohakuhub.db_async import execute_db_query, get_file_by_sha256 -from kohakuhub.db import Organization, Repository, User -from kohakuhub.logger import get_logger +from kohakuhub.api.quota.util import check_quota logger = get_logger("LFS") router = APIRouter() @@ -101,7 +101,7 @@ async def process_upload_object(oid: str, size: int, repo_id: str) -> LFSObjectR lfs_key = get_lfs_key(oid) # Check if object already exists (deduplication) - existing = await get_file_by_sha256(oid) + existing = get_file_by_sha256(oid) if existing and existing.size == size: # Object exists, no upload needed @@ -178,7 +178,7 @@ async def process_download_object(oid: str, size: int) -> LFSObjectResponse: lfs_key = get_lfs_key(oid) # Check if object exists - existing = await get_file_by_sha256(oid) + existing = get_file_by_sha256(oid) if not existing: return LFSObjectResponse( @@ -258,12 +258,9 @@ async def lfs_batch( raise HTTPException(400, detail={"error": f"Invalid LFS batch request: {e}"}) # Check repository exists and permissions - def _get_repo(): - return Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) - ) - - repo_row = await execute_db_query(_get_repo) + repo_row = Repository.get_or_none( + (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) + ) if repo_row: operation = batch_req.operation @@ -281,15 +278,12 @@ async def lfs_batch( total_upload_bytes = sum(obj.size for obj in batch_req.objects) # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Check quota (based on repo privacy) is_private = repo_row.private - allowed, error_msg = await check_quota( + allowed, error_msg = check_quota( namespace, total_upload_bytes, is_private, is_org ) if not allowed: diff --git a/src/kohakuhub/api/git/routers/ssh_keys.py b/src/kohakuhub/api/git/routers/ssh_keys.py index 00645ba..6ae84f9 100644 --- a/src/kohakuhub/api/git/routers/ssh_keys.py +++ b/src/kohakuhub/api/git/routers/ssh_keys.py @@ -3,13 +3,13 @@ import base64 import hashlib -from cryptography.hazmat.backends import default_backend -from cryptography.hazmat.primitives.serialization import load_ssh_public_key from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives.serialization import load_ssh_public_key from kohakuhub.db import SSHKey, User -from kohakuhub.db_async import ( +from kohakuhub.db_operations import ( create_ssh_key, delete_ssh_key, get_ssh_key_by_fingerprint, @@ -153,7 +153,7 @@ async def list_ssh_keys(user: User = Depends(get_current_user)): Returns: List of SSH keys """ - keys = await list_user_ssh_keys(user.id) + keys = list_user_ssh_keys(user.id) return [ SSHKeyResponse( @@ -193,7 +193,7 @@ async def add_ssh_key( key_type, fingerprint = validate_ssh_key(key_request.key) # Check if key already exists - existing_key = await get_ssh_key_by_fingerprint(fingerprint) + existing_key = get_ssh_key_by_fingerprint(fingerprint) if existing_key: raise HTTPException( 409, @@ -201,7 +201,7 @@ async def add_ssh_key( ) # Create key - new_key = await create_ssh_key( + new_key = create_ssh_key( user_id=user.id, key_type=key_type, public_key=key_request.key.strip(), @@ -239,7 +239,7 @@ async def remove_ssh_key(key_id: int, user: User = Depends(get_current_user)): HTTPException: If key not found or not owned by user """ # Get key - key = await get_ssh_key_by_id(key_id) + key = get_ssh_key_by_id(key_id) if not key: raise HTTPException(404, detail="SSH key not found") @@ -248,7 +248,7 @@ async def remove_ssh_key(key_id: int, user: User = Depends(get_current_user)): raise HTTPException(403, detail="Not authorized to delete this key") # Delete key - await delete_ssh_key(key) + delete_ssh_key(key) logger.info( f"User {user.username} removed SSH key: {key.key_type} {key.fingerprint}" @@ -272,7 +272,7 @@ async def get_ssh_key(key_id: int, user: User = Depends(get_current_user)): HTTPException: If key not found or not owned by user """ # Get key - key = await get_ssh_key_by_id(key_id) + key = get_ssh_key_by_id(key_id) if not key: raise HTTPException(404, detail="SSH key not found") diff --git a/src/kohakuhub/api/git/utils/lakefs_bridge.py b/src/kohakuhub/api/git/utils/lakefs_bridge.py index ad62012..e795bd8 100644 --- a/src/kohakuhub/api/git/utils/lakefs_bridge.py +++ b/src/kohakuhub/api/git/utils/lakefs_bridge.py @@ -1,14 +1,14 @@ """Git-LakeFS bridge - Pure Python implementation (no pygit2, no file I/O).""" +from datetime import datetime import asyncio import fnmatch import hashlib -from datetime import datetime from kohakuhub.config import cfg from kohakuhub.db import File -from kohakuhub.db_async import execute_db_query from kohakuhub.logger import get_logger +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.api.git.utils.objects import ( build_nested_trees, create_blob_object, @@ -16,7 +16,6 @@ from kohakuhub.api.git.utils.objects import ( create_pack_file, ) from kohakuhub.api.git.utils.server import create_empty_pack -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name logger = get_logger("GIT_LAKEFS") @@ -172,13 +171,10 @@ class GitLakeFSBridge: """ # Get File table records for LFS tracking - def _get_all_files(): - return { - f.path_in_repo: f - for f in File.select().where(File.repo_full_id == self.repo_id) - } - - file_records = await execute_db_query(_get_all_files) + file_records = { + f.path_in_repo: f + for f in File.select().where(File.repo_full_id == self.repo_id) + } # Check for .gitattributes and parse LFS patterns existing_lfs_patterns = set() diff --git a/src/kohakuhub/api/git/utils/objects.py b/src/kohakuhub/api/git/utils/objects.py index 396fd0e..066a0e7 100644 --- a/src/kohakuhub/api/git/utils/objects.py +++ b/src/kohakuhub/api/git/utils/objects.py @@ -7,7 +7,6 @@ import hashlib import struct import zlib - def compute_git_object_sha1(obj_type: str, content: bytes) -> str: """Compute Git object SHA-1. diff --git a/src/kohakuhub/api/misc.py b/src/kohakuhub/api/misc.py index 66156ab..38b8554 100644 --- a/src/kohakuhub/api/misc.py +++ b/src/kohakuhub/api/misc.py @@ -1,14 +1,13 @@ """Utility API endpoints for Kohaku Hub.""" -import yaml - from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel +import yaml -from kohakuhub.auth.dependencies import get_optional_user from kohakuhub.config import cfg from kohakuhub.db import Organization, User, UserOrganization from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_optional_user logger = get_logger("UTILS") diff --git a/src/kohakuhub/api/org/router.py b/src/kohakuhub/api/org/router.py index 33c696c..536696d 100644 --- a/src/kohakuhub/api/org/router.py +++ b/src/kohakuhub/api/org/router.py @@ -4,16 +4,16 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from kohakuhub.db import Organization, User, UserOrganization -from kohakuhub.db_async import ( +from kohakuhub.db_operations import ( create_organization as create_org_async, create_user_organization, - execute_db_query, get_organization, get_user_by_username, get_user_organization, list_organization_members as list_org_members_async, ) from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user from kohakuhub.api.org.util import ( add_member_to_organization as add_member_util, get_organization_details as get_org_details_util, @@ -21,7 +21,6 @@ from kohakuhub.api.org.util import ( remove_member_from_organization, update_member_role as update_member_role_util, ) -from kohakuhub.auth.dependencies import get_current_user # Error messages _ERR_ORG_NOT_FOUND = "Organization not found" @@ -43,18 +42,15 @@ async def create_organization( """Create a new organization with default quotas.""" # Check if organization already exists - def _check_exists(): - return Organization.get_or_none(Organization.name == payload.name) - - existing_org = await execute_db_query(_check_exists) + existing_org = Organization.get_or_none(Organization.name == payload.name) if existing_org: raise HTTPException(400, detail="Organization name already exists") # Create organization with default quotas - org = await create_org_async(payload.name, payload.description) + org = create_org_async(payload.name, payload.description) # Add creator as super-admin - await create_user_organization(user.id, org.id, "super-admin") + create_user_organization(user.id, org.id, "super-admin") logger.info(f"User {user.username} created organization: {org.name}") @@ -91,7 +87,7 @@ async def add_member( raise HTTPException(404, detail=_ERR_ORG_NOT_FOUND) # Check if the current user is an admin of the organization - user_org = await get_user_organization(current_user.id, org.id) + user_org = get_user_organization(current_user.id, org.id) if not user_org or user_org.role not in ["admin", "super-admin"]: raise HTTPException(403, detail="Not authorized to add members") @@ -111,7 +107,7 @@ async def remove_member( raise HTTPException(404, detail=_ERR_ORG_NOT_FOUND) # Check if the current user is an admin of the organization - user_org = await get_user_organization(current_user.id, org.id) + user_org = get_user_organization(current_user.id, org.id) if not user_org or user_org.role not in ["admin", "super-admin"]: raise HTTPException(403, detail="Not authorized to remove members") @@ -122,7 +118,7 @@ async def remove_member( @router.get("/users/{username}/orgs") async def list_user_organizations(username: str): """List organizations a user belongs to.""" - user = await get_user_by_username(username) + user = get_user_by_username(username) if not user: raise HTTPException(404, detail="User not found") @@ -156,7 +152,7 @@ async def update_member_role( raise HTTPException(404, detail=_ERR_ORG_NOT_FOUND) # Check if the current user is an admin of the organization - user_org = await get_user_organization(current_user.id, org.id) + user_org = get_user_organization(current_user.id, org.id) if not user_org or user_org.role not in ["admin", "super-admin"]: raise HTTPException(403, detail="Not authorized to update member roles") @@ -167,12 +163,12 @@ async def update_member_role( @router.get("/{org_name}/members") async def list_organization_members(org_name: str): """List organization members.""" - org = await get_organization(org_name) + org = get_organization(org_name) if not org: raise HTTPException(404, detail=_ERR_ORG_NOT_FOUND) # Get all members - members = await list_org_members_async(org.id) + members = list_org_members_async(org.id) return { "members": [ diff --git a/src/kohakuhub/api/org/util.py b/src/kohakuhub/api/org/util.py index 4da79b9..98fd225 100644 --- a/src/kohakuhub/api/org/util.py +++ b/src/kohakuhub/api/org/util.py @@ -2,6 +2,7 @@ from fastapi import HTTPException +from kohakuhub.config import cfg from kohakuhub.db import Organization, User, UserOrganization # Error messages @@ -13,8 +14,6 @@ def create_organization(name: str, description: str | None, user: User): Note: This function is deprecated. Use async version from db_async instead. """ - from kohakuhub.config import cfg - if Organization.get_or_none(Organization.name == name): raise HTTPException(400, detail="Organization name already exists") diff --git a/src/kohakuhub/api/quota/router.py b/src/kohakuhub/api/quota/router.py index 52acb09..d759c72 100644 --- a/src/kohakuhub/api/quota/router.py +++ b/src/kohakuhub/api/quota/router.py @@ -3,15 +3,14 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel +from kohakuhub.db import Organization, User, UserOrganization +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user, get_optional_user from kohakuhub.api.quota.util import ( get_storage_info, set_quota, update_namespace_storage, ) -from kohakuhub.auth.dependencies import get_current_user, get_optional_user -from kohakuhub.db import Organization, User, UserOrganization -from kohakuhub.db_async import execute_db_query -from kohakuhub.logger import get_logger logger = get_logger("QUOTA") router = APIRouter() @@ -74,25 +73,19 @@ async def get_quota( """ # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Check if namespace exists - def _check_user(): - return User.get_or_none(User.username == namespace) - if not is_org: - target_user = await execute_db_query(_check_user) + target_user = User.get_or_none(User.username == namespace) if not target_user: raise HTTPException( 404, detail={"error": f"User or organization not found: {namespace}"} ) # Get storage info - info = await get_storage_info(namespace, is_org) + info = get_storage_info(namespace, is_org) return QuotaInfo( namespace=namespace, @@ -127,25 +120,17 @@ async def update_quota( """ # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Authorization check if is_org: # For orgs, must be admin member - def _check_org_admin(): - from kohakuhub.db import UserOrganization - - member = UserOrganization.get_or_none( - (UserOrganization.user == user.id) - & (UserOrganization.organization == org.id) - ) - return member and member.role in ("admin", "super-admin") - - is_admin = await execute_db_query(_check_org_admin) + member = UserOrganization.get_or_none( + (UserOrganization.user == user.id) + & (UserOrganization.organization == org.id) + ) + is_admin = member and member.role in ("admin", "super-admin") if not is_admin: raise HTTPException( 403, @@ -161,7 +146,7 @@ async def update_quota( ) # Set quota - info = await set_quota(namespace, request.quota_bytes, is_org) + info = set_quota(namespace, request.quota_bytes, is_org) return QuotaInfo( namespace=namespace, @@ -194,25 +179,17 @@ async def recalculate_storage( """ # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Authorization check (same as update_quota) if is_org: # For orgs, must be admin member - def _check_org_admin(): - from kohakuhub.db import UserOrganization - - member = UserOrganization.get_or_none( - (UserOrganization.user == user.id) - & (UserOrganization.organization == org.id) - ) - return member and member.role in ("admin", "super-admin") - - is_admin = await execute_db_query(_check_org_admin) + member = UserOrganization.get_or_none( + (UserOrganization.user == user.id) + & (UserOrganization.organization == org.id) + ) + is_admin = member and member.role in ("admin", "super-admin") if not is_admin: raise HTTPException( 403, @@ -229,10 +206,10 @@ async def recalculate_storage( # Recalculate storage logger.info(f"Recalculating storage for {'org' if is_org else 'user'} {namespace}") - await update_namespace_storage(namespace, is_org) + update_namespace_storage(namespace, is_org) # Get updated info - info = await get_storage_info(namespace, is_org) + info = get_storage_info(namespace, is_org) return QuotaInfo( namespace=namespace, @@ -271,19 +248,13 @@ async def get_public_quota( """ # Check if namespace is organization - def _check_org(): - return Organization.get_or_none(Organization.name == namespace) - - org = await execute_db_query(_check_org) + org = Organization.get_or_none(Organization.name == namespace) is_org = org is not None # Check if namespace exists - def _check_user(): - return User.get_or_none(User.username == namespace) - target_user = None if not is_org: - target_user = await execute_db_query(_check_user) + target_user = User.get_or_none(User.username == namespace) if not target_user: raise HTTPException( 404, detail={"error": f"User or organization not found: {namespace}"} @@ -295,20 +266,17 @@ async def get_public_quota( if user: if is_org: # For organizations: user must be a member - def _check_membership(): - return UserOrganization.get_or_none( - (UserOrganization.user == user.id) - & (UserOrganization.organization == org.id) - ) - - membership = await execute_db_query(_check_membership) + membership = UserOrganization.get_or_none( + (UserOrganization.user == user.id) + & (UserOrganization.organization == org.id) + ) can_see_private = membership is not None else: # For users: user must be viewing their own profile can_see_private = user.username == namespace # Get storage info - info = await get_storage_info(namespace, is_org) + info = get_storage_info(namespace, is_org) # Build response based on permissions response_data = { diff --git a/src/kohakuhub/api/quota/util.py b/src/kohakuhub/api/quota/util.py index d506aa8..d116742 100644 --- a/src/kohakuhub/api/quota/util.py +++ b/src/kohakuhub/api/quota/util.py @@ -6,11 +6,10 @@ for users and organizations with separate tracking for private and public reposi import asyncio -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.config import cfg from kohakuhub.db import LFSObjectHistory, Organization, Repository, User -from kohakuhub.db_async import execute_db_query from kohakuhub.logger import get_logger +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name logger = get_logger("QUOTA") @@ -63,36 +62,31 @@ async def calculate_repository_storage(repo: Repository) -> dict[str, int]: ) # Calculate LFS storage from history - def _get_lfs_stats(): - # Total LFS storage (all versions) - lfs_total = ( - LFSObjectHistory.select() - .where(LFSObjectHistory.repo_full_id == repo.full_id) - .count() - ) + # Total LFS storage (all versions) + lfs_total = ( + LFSObjectHistory.select() + .where(LFSObjectHistory.repo_full_id == repo.full_id) + .count() + ) - if lfs_total > 0: - lfs_total_bytes = sum( - obj.size - for obj in LFSObjectHistory.select().where( - LFSObjectHistory.repo_full_id == repo.full_id - ) + if lfs_total > 0: + lfs_total_bytes = sum( + obj.size + for obj in LFSObjectHistory.select().where( + LFSObjectHistory.repo_full_id == repo.full_id ) - else: - lfs_total_bytes = 0 - - # Unique LFS storage (deduplicated by SHA256) - unique_lfs = ( - LFSObjectHistory.select(LFSObjectHistory.sha256, LFSObjectHistory.size) - .where(LFSObjectHistory.repo_full_id == repo.full_id) - .distinct() ) + else: + lfs_total_bytes = 0 - lfs_unique_bytes = sum(obj.size for obj in unique_lfs) + # Unique LFS storage (deduplicated by SHA256) + unique_lfs = ( + LFSObjectHistory.select(LFSObjectHistory.sha256, LFSObjectHistory.size) + .where(LFSObjectHistory.repo_full_id == repo.full_id) + .distinct() + ) - return lfs_total_bytes, lfs_unique_bytes - - lfs_total_bytes, lfs_unique_bytes = await execute_db_query(_get_lfs_stats) + lfs_unique_bytes = sum(obj.size for obj in unique_lfs) # Total storage = current branch + all LFS versions total_bytes = current_branch_bytes + lfs_total_bytes @@ -122,10 +116,7 @@ async def calculate_namespace_storage( """ # Get all repositories in this namespace - def _get_repos(): - return list(Repository.select().where(Repository.namespace == namespace)) - - repos = await execute_db_query(_get_repos) + repos = list(Repository.select().where(Repository.namespace == namespace)) # Separate repos by privacy private_repos = [r for r in repos if r.private] @@ -176,27 +167,24 @@ async def update_namespace_storage( storage = await calculate_namespace_storage(namespace, is_org) # Update database - def _update(): - if is_org: - org = Organization.get(Organization.name == namespace) - org.private_used_bytes = storage["private_bytes"] - org.public_used_bytes = storage["public_bytes"] - org.save() - else: - user = User.get(User.username == namespace) - user.private_used_bytes = storage["private_bytes"] - user.public_used_bytes = storage["public_bytes"] - user.save() - - await execute_db_query(_update) + if is_org: + org = Organization.get(Organization.name == namespace) + org.private_used_bytes = storage["private_bytes"] + org.public_used_bytes = storage["public_bytes"] + org.save() + else: + user = User.get(User.username == namespace) + user.private_used_bytes = storage["private_bytes"] + user.public_used_bytes = storage["public_bytes"] + user.save() return storage -async def check_quota( +def check_quota( namespace: str, additional_bytes: int, is_private: bool, is_org: bool = False ) -> tuple[bool, str | None]: - """Check if adding additional storage would exceed quota. + """Check if adding additional storage would exceed quota (SYNCHRONOUS). Args: namespace: Username or organization name @@ -209,33 +197,38 @@ async def check_quota( """ # Get current quota and usage - def _get_quota(): - if is_org: - org = Organization.get_or_none(Organization.name == namespace) - if not org: - return None, 0, 0, 0, 0 - return ( - org, - org.private_quota_bytes, - org.public_quota_bytes, - org.private_used_bytes, - org.public_used_bytes, + if is_org: + org = Organization.get_or_none(Organization.name == namespace) + if not org: + entity, private_quota, public_quota, private_used, public_used = ( + None, + 0, + 0, + 0, + 0, ) else: - user = User.get_or_none(User.username == namespace) - if not user: - return None, 0, 0, 0, 0 - return ( - user, - user.private_quota_bytes, - user.public_quota_bytes, - user.private_used_bytes, - user.public_used_bytes, + entity = org + private_quota = org.private_quota_bytes + public_quota = org.public_quota_bytes + private_used = org.private_used_bytes + public_used = org.public_used_bytes + else: + user = User.get_or_none(User.username == namespace) + if not user: + entity, private_quota, public_quota, private_used, public_used = ( + None, + 0, + 0, + 0, + 0, ) - - entity, private_quota, public_quota, private_used, public_used = ( - await execute_db_query(_get_quota) - ) + else: + entity = user + private_quota = user.private_quota_bytes + public_quota = user.public_quota_bytes + private_used = user.private_used_bytes + public_used = user.public_used_bytes if not entity: return False, f"{'Organization' if is_org else 'User'} not found: {namespace}" @@ -268,10 +261,10 @@ async def check_quota( return True, None -async def increment_storage( +def increment_storage( namespace: str, bytes_delta: int, is_private: bool, is_org: bool = False ) -> tuple[int, int]: - """Increment storage usage for a user or organization. + """Increment storage usage for a user or organization (SYNCHRONOUS). Args: namespace: Username or organization name @@ -283,25 +276,22 @@ async def increment_storage( Tuple of (private_used_bytes, public_used_bytes) """ - def _increment(): - if is_org: - org = Organization.get(Organization.name == namespace) - if is_private: - org.private_used_bytes = max(0, org.private_used_bytes + bytes_delta) - else: - org.public_used_bytes = max(0, org.public_used_bytes + bytes_delta) - org.save() - return org.private_used_bytes, org.public_used_bytes + if is_org: + org = Organization.get(Organization.name == namespace) + if is_private: + org.private_used_bytes = max(0, org.private_used_bytes + bytes_delta) else: - user = User.get(User.username == namespace) - if is_private: - user.private_used_bytes = max(0, user.private_used_bytes + bytes_delta) - else: - user.public_used_bytes = max(0, user.public_used_bytes + bytes_delta) - user.save() - return user.private_used_bytes, user.public_used_bytes - - private_used, public_used = await execute_db_query(_increment) + org.public_used_bytes = max(0, org.public_used_bytes + bytes_delta) + org.save() + private_used, public_used = org.private_used_bytes, org.public_used_bytes + else: + user = User.get(User.username == namespace) + if is_private: + user.private_used_bytes = max(0, user.private_used_bytes + bytes_delta) + else: + user.public_used_bytes = max(0, user.public_used_bytes + bytes_delta) + user.save() + private_used, public_used = user.private_used_bytes, user.public_used_bytes logger.debug( f"Updated storage for {'org' if is_org else 'user'} {namespace}: " @@ -312,10 +302,10 @@ async def increment_storage( return private_used, public_used -async def get_storage_info( +def get_storage_info( namespace: str, is_org: bool = False ) -> dict[str, int | float | None]: - """Get storage quota and usage information. + """Get storage quota and usage information (SYNCHRONOUS). Args: namespace: Username or organization name @@ -325,31 +315,24 @@ async def get_storage_info( Dict with quota information for both private and public storage """ - def _get_info(): - if is_org: - org = Organization.get_or_none(Organization.name == namespace) - if not org: - return None, None, 0, 0 - return ( - org.private_quota_bytes, - org.public_quota_bytes, - org.private_used_bytes, - org.public_used_bytes, - ) + if is_org: + org = Organization.get_or_none(Organization.name == namespace) + if not org: + private_quota, public_quota, private_used, public_used = None, None, 0, 0 else: - user = User.get_or_none(User.username == namespace) - if not user: - return None, None, 0, 0 - return ( - user.private_quota_bytes, - user.public_quota_bytes, - user.private_used_bytes, - user.public_used_bytes, - ) - - private_quota, public_quota, private_used, public_used = await execute_db_query( - _get_info - ) + private_quota = org.private_quota_bytes + public_quota = org.public_quota_bytes + private_used = org.private_used_bytes + public_used = org.public_used_bytes + else: + user = User.get_or_none(User.username == namespace) + if not user: + private_quota, public_quota, private_used, public_used = None, None, 0, 0 + else: + private_quota = user.private_quota_bytes + public_quota = user.public_quota_bytes + private_used = user.private_used_bytes + public_used = user.public_used_bytes # Calculate availability and percentages private_available = ( @@ -385,13 +368,13 @@ async def get_storage_info( } -async def set_quota( +def set_quota( namespace: str, private_quota_bytes: int | None = None, public_quota_bytes: int | None = None, is_org: bool = False, ) -> dict[str, int | float | None]: - """Set storage quota for a user or organization. + """Set storage quota for a user or organization (SYNCHRONOUS). Args: namespace: Username or organization name @@ -403,27 +386,24 @@ async def set_quota( Updated storage info dict """ - def _set_quota(): - if is_org: - org = Organization.get(Organization.name == namespace) - if private_quota_bytes is not None: - org.private_quota_bytes = private_quota_bytes - if public_quota_bytes is not None: - org.public_quota_bytes = public_quota_bytes - org.save() - else: - user = User.get(User.username == namespace) - if private_quota_bytes is not None: - user.private_quota_bytes = private_quota_bytes - if public_quota_bytes is not None: - user.public_quota_bytes = public_quota_bytes - user.save() - - await execute_db_query(_set_quota) + if is_org: + org = Organization.get(Organization.name == namespace) + if private_quota_bytes is not None: + org.private_quota_bytes = private_quota_bytes + if public_quota_bytes is not None: + org.public_quota_bytes = public_quota_bytes + org.save() + else: + user = User.get(User.username == namespace) + if private_quota_bytes is not None: + user.private_quota_bytes = private_quota_bytes + if public_quota_bytes is not None: + user.public_quota_bytes = public_quota_bytes + user.save() logger.info( f"Set quota for {'org' if is_org else 'user'} {namespace}: " f"private={private_quota_bytes}, public={public_quota_bytes}" ) - return await get_storage_info(namespace, is_org) + return get_storage_info(namespace, is_org) # No await - it's sync now diff --git a/src/kohakuhub/api/repo/routers/crud.py b/src/kohakuhub/api/repo/routers/crud.py index a7cf1d0..46c4ea1 100644 --- a/src/kohakuhub/api/repo/routers/crud.py +++ b/src/kohakuhub/api/repo/routers/crud.py @@ -5,6 +5,16 @@ from typing import Literal, Optional from fastapi import APIRouter, Depends from pydantic import BaseModel +from kohakuhub.config import cfg +from kohakuhub.db import File, Repository, StagingUpload, User, init_db +from kohakuhub.db_operations import get_repository +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user +from kohakuhub.auth.permissions import ( + check_namespace_permission, + check_repo_delete_permission, +) +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.api.repo.utils.hf import ( HFErrorCode, hf_error_response, @@ -12,16 +22,6 @@ from kohakuhub.api.repo.utils.hf import ( hf_server_error, is_lakefs_not_found_error, ) -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.auth.dependencies import get_current_user -from kohakuhub.auth.permissions import ( - check_namespace_permission, - check_repo_delete_permission, -) -from kohakuhub.config import cfg -from kohakuhub.db_async import execute_db_query, get_repository -from kohakuhub.db import File, Repository, StagingUpload, User, init_db -from kohakuhub.logger import get_logger logger = get_logger("REPO") router = APIRouter() @@ -64,7 +64,7 @@ async def create_repo( full_id = f"{namespace}/{payload.name}" lakefs_repo = lakefs_repo_name(payload.type, full_id) - existing_repo = await get_repository(payload.type, namespace, payload.name) + existing_repo = get_repository(payload.type, namespace, payload.name) if existing_repo: return hf_error_response( 400, @@ -87,16 +87,13 @@ async def create_repo( return hf_server_error(f"LakeFS repository creation failed: {str(e)}") # Store in database for listing/metadata - def _create_repo(): - Repository.get_or_create( - repo_type=payload.type, - namespace=namespace, - name=payload.name, - full_id=full_id, - defaults={"private": payload.private, "owner_id": user.id}, - ) - - await execute_db_query(_create_repo) + Repository.get_or_create( + repo_type=payload.type, + namespace=namespace, + name=payload.name, + full_id=full_id, + defaults={"private": payload.private, "owner_id": user.id}, + ) return { "url": f"{cfg.app.base_url}/{payload.type}s/{full_id}", @@ -135,7 +132,7 @@ async def delete_repo( lakefs_repo = lakefs_repo_name(repo_type, full_id) # 1. Check if repository exists in database - repo_row = await get_repository(repo_type, namespace, payload.name) + repo_row = get_repository(repo_type, namespace, payload.name) if not repo_row: # NOTE: HuggingFace client expects 400 for delete repo not found @@ -160,22 +157,14 @@ async def delete_repo( logger.info(f"LakeFS repository {lakefs_repo} not found/already deleted (OK)") # 3. Delete related metadata from database (manual cascade) - def _delete_db_records(): - try: - # Delete related file records first - File.delete().where(File.repo_full_id == full_id).execute() - StagingUpload.delete().where( - StagingUpload.repo_full_id == full_id - ).execute() - repo_row.delete_instance() - logger.success(f"Successfully deleted database records for: {full_id}") - except Exception as e: - logger.exception(f"Database deletion failed for {full_id}", e) - raise - try: - await execute_db_query(_delete_db_records) + # Delete related file records first + File.delete().where(File.repo_full_id == full_id).execute() + StagingUpload.delete().where(StagingUpload.repo_full_id == full_id).execute() + repo_row.delete_instance() + logger.success(f"Successfully deleted database records for: {full_id}") except Exception as e: + logger.exception(f"Database deletion failed for {full_id}", e) return hf_server_error(f"Database deletion failed for {full_id}: {str(e)}") # 4. Return success response (200 OK with a simple message) @@ -219,7 +208,7 @@ async def move_repo( ) from_namespace, from_name = from_parts - repo_row = await get_repository(repo_type, from_namespace, from_name) + repo_row = get_repository(repo_type, from_namespace, from_name) if not repo_row: return hf_repo_not_found(from_id, repo_type) @@ -235,7 +224,7 @@ async def move_repo( ) to_namespace, to_name = to_parts - existing = await get_repository(repo_type, to_namespace, to_name) + existing = get_repository(repo_type, to_namespace, to_name) if existing: return hf_error_response( 400, @@ -247,23 +236,20 @@ async def move_repo( check_namespace_permission(to_namespace, user) # Update database records - def _update_db_records(): - # Update repository record - Repository.update( - namespace=to_namespace, - name=to_name, - full_id=to_id, - ).where(Repository.id == repo_row.id).execute() + # Update repository record + Repository.update( + namespace=to_namespace, + name=to_name, + full_id=to_id, + ).where(Repository.id == repo_row.id).execute() - # Update related file records - File.update(repo_full_id=to_id).where(File.repo_full_id == from_id).execute() + # Update related file records + File.update(repo_full_id=to_id).where(File.repo_full_id == from_id).execute() - # Update staging uploads - StagingUpload.update(repo_full_id=to_id).where( - StagingUpload.repo_full_id == from_id - ).execute() - - await execute_db_query(_update_db_records) + # Update staging uploads + StagingUpload.update(repo_full_id=to_id).where( + StagingUpload.repo_full_id == from_id + ).execute() # Note: LakeFS repository rename not implemented yet # Would require creating new LakeFS repo and migrating data diff --git a/src/kohakuhub/api/repo/routers/info.py b/src/kohakuhub/api/repo/routers/info.py index caffc64..7fdce64 100644 --- a/src/kohakuhub/api/repo/routers/info.py +++ b/src/kohakuhub/api/repo/routers/info.py @@ -1,30 +1,29 @@ """Repository information and listing endpoints.""" -import asyncio from datetime import datetime from typing import Literal, Optional +import asyncio from fastapi import APIRouter, Depends, Query, Request +from kohakuhub.config import cfg +from kohakuhub.db import Organization, Repository, User, UserOrganization +from kohakuhub.db_operations import ( + get_file, + get_organization, + get_repository, + get_user_by_username, +) +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_optional_user +from kohakuhub.auth.permissions import check_repo_read_permission +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.api.repo.utils.hf import ( HFErrorCode, format_hf_datetime, hf_error_response, hf_repo_not_found, ) -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.auth.dependencies import get_optional_user -from kohakuhub.auth.permissions import check_repo_read_permission -from kohakuhub.config import cfg -from kohakuhub.db_async import ( - execute_db_query, - get_file, - get_organization, - get_repository, - get_user_by_username, -) -from kohakuhub.db import Organization, Repository, User, UserOrganization -from kohakuhub.logger import get_logger logger = get_logger("REPO") router = APIRouter() @@ -79,7 +78,7 @@ async def get_repo_info( ) # Check if repository exists in database - repo_row = await get_repository(repo_type, namespace, repo_name) + repo_row = get_repository(repo_type, namespace, repo_name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -152,18 +151,15 @@ async def get_repo_info( file_records = {} if lfs_files: - # Fetch all file records in parallel - file_record_tasks = [ - get_file(repo_id, obj["path"]) for obj in lfs_files - ] - file_record_results = await asyncio.gather( - *file_record_tasks, return_exceptions=True - ) - - # Map path to file record - for obj, record in zip(lfs_files, file_record_results): - if not isinstance(record, Exception): - file_records[obj["path"]] = record + # Fetch all file records sequentially (sync DB operations) + for obj in lfs_files: + try: + record = get_file(repo_id, obj["path"]) + if record: + file_records[obj["path"]] = record + except Exception: + # Skip files that fail to fetch + pass # Convert to siblings format for obj in file_objects: @@ -304,19 +300,16 @@ async def list_repos( ) # Query database - def _query_repos(): - q = Repository.select().where(Repository.repo_type == rt) + q = Repository.select().where(Repository.repo_type == rt) - # Filter by author if specified - if author: - q = q.where(Repository.namespace == author) + # Filter by author if specified + if author: + q = q.where(Repository.namespace == author) - # Apply privacy filtering - q = _filter_repos_by_privacy(q, user, author) + # Apply privacy filtering + q = _filter_repos_by_privacy(q, user, author) - return list(q.limit(limit)) - - rows = await execute_db_query(_query_repos) + rows = list(q.limit(limit)) # Format response with lastModified from LakeFS client = get_lakefs_client() @@ -390,10 +383,10 @@ async def list_user_repos( Dict with models, datasets, and spaces lists """ # Check if the username exists - target_user = await get_user_by_username(username) + target_user = get_user_by_username(username) if not target_user: # Could also be an organization - target_org = await get_organization(username) + target_org = get_organization(username) if not target_org: return hf_error_response( 404, @@ -408,18 +401,14 @@ async def list_user_repos( } for repo_type in ["model", "dataset", "space"]: + q = Repository.select().where( + (Repository.repo_type == repo_type) & (Repository.namespace == username) + ) - def _query_repos(): - q = Repository.select().where( - (Repository.repo_type == repo_type) & (Repository.namespace == username) - ) + # Apply privacy filtering + q = _filter_repos_by_privacy(q, user, username) - # Apply privacy filtering - q = _filter_repos_by_privacy(q, user, username) - - return list(q.limit(limit)) - - rows = await execute_db_query(_query_repos) + rows = list(q.limit(limit)) key = repo_type + "s" repos_list = [] diff --git a/src/kohakuhub/api/repo/routers/tree.py b/src/kohakuhub/api/repo/routers/tree.py index 3d037bf..1022762 100644 --- a/src/kohakuhub/api/repo/routers/tree.py +++ b/src/kohakuhub/api/repo/routers/tree.py @@ -1,11 +1,18 @@ """Repository tree listing and path information endpoints - Refactored version.""" -import asyncio from datetime import datetime from typing import Literal +import asyncio from fastapi import APIRouter, Depends, Form +from kohakuhub.config import cfg +from kohakuhub.db import File, Repository, User +from kohakuhub.db_operations import get_file +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_optional_user +from kohakuhub.auth.permissions import check_repo_read_permission +from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name from kohakuhub.api.repo.utils.hf import ( hf_repo_not_found, hf_revision_not_found, @@ -13,13 +20,6 @@ from kohakuhub.api.repo.utils.hf import ( is_lakefs_not_found_error, is_lakefs_revision_error, ) -from kohakuhub.utils.lakefs import get_lakefs_client, lakefs_repo_name -from kohakuhub.auth.dependencies import get_optional_user -from kohakuhub.auth.permissions import check_repo_read_permission -from kohakuhub.config import cfg -from kohakuhub.db_async import execute_db_query, get_file -from kohakuhub.db import File, Repository, User -from kohakuhub.logger import get_logger logger = get_logger("REPO") router = APIRouter() @@ -132,7 +132,7 @@ async def convert_file_object(obj, repo_id: str, prefix_len: int) -> dict: relative_path = obj["path"][prefix_len:] if prefix_len else obj["path"] # Get correct checksum from database - file_record = await get_file(repo_id, obj["path"]) + file_record = get_file(repo_id, obj["path"]) checksum = ( file_record.sha256 if file_record and file_record.sha256 else obj["checksum"] @@ -236,12 +236,9 @@ async def list_repo_tree( repo_id = f"{namespace}/{repo_name}" # Check if repository exists - def _get_repo(): - return Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) - ) - - repo_row = await execute_db_query(_get_repo) + repo_row = Repository.get_or_none( + (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) + ) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -325,12 +322,9 @@ async def get_paths_info( repo_id = f"{namespace}/{repo_name}" # Check if repository exists - def _get_repo(): - return Repository.get_or_none( - (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) - ) - - repo_row = await execute_db_query(_get_repo) + repo_row = Repository.get_or_none( + (Repository.full_id == repo_id) & (Repository.repo_type == repo_type) + ) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -358,7 +352,7 @@ async def get_paths_info( is_lfs = obj_stats["size_bytes"] > cfg.app.lfs_threshold_bytes # Get correct checksum from database - file_record = await get_file(repo_id, clean_path) + file_record = get_file(repo_id, clean_path) checksum = ( file_record.sha256 diff --git a/src/kohakuhub/api/repo/utils/gc.py b/src/kohakuhub/api/repo/utils/gc.py index e2a79dd..4a824b4 100644 --- a/src/kohakuhub/api/repo/utils/gc.py +++ b/src/kohakuhub/api/repo/utils/gc.py @@ -1,12 +1,11 @@ """Garbage collection utilities for LFS objects.""" -import asyncio from datetime import datetime, timezone from typing import List, Optional +import asyncio from kohakuhub.config import cfg from kohakuhub.db import File, LFSObjectHistory -from kohakuhub.db_async import execute_db_query from kohakuhub.logger import get_logger from kohakuhub.utils.lakefs import get_lakefs_client from kohakuhub.utils.s3 import get_s3_client, object_exists @@ -446,19 +445,17 @@ async def sync_file_table_with_commit( synced_count += 1 # Remove files from File table that no longer exist in commit - def _cleanup_removed_files(): - if file_paths: - return ( - File.delete() - .where( - (File.repo_full_id == repo_full_id) - & (File.path_in_repo.not_in(file_paths)) - ) - .execute() + if file_paths: + removed_count = ( + File.delete() + .where( + (File.repo_full_id == repo_full_id) + & (File.path_in_repo.not_in(file_paths)) ) - return 0 - - removed_count = await execute_db_query(_cleanup_removed_files) + .execute() + ) + else: + removed_count = 0 if removed_count > 0: logger.info(f"Removed {removed_count} stale file(s) from File table") @@ -586,18 +583,14 @@ async def track_commit_lfs_objects( # Remove deleted files from File table if files_to_remove: - - def _cleanup_removed_files(): - return ( - File.delete() - .where( - (File.repo_full_id == repo_full_id) - & (File.path_in_repo.in_(files_to_remove)) - ) - .execute() + removed_count = ( + File.delete() + .where( + (File.repo_full_id == repo_full_id) + & (File.path_in_repo.in_(files_to_remove)) ) - - removed_count = await execute_db_query(_cleanup_removed_files) + .execute() + ) logger.info(f"Removed {removed_count} deleted file(s) from File table") if tracked_count > 0: diff --git a/src/kohakuhub/api/repo/utils/hf.py b/src/kohakuhub/api/repo/utils/hf.py index 0501b3a..f8d5b6b 100644 --- a/src/kohakuhub/api/repo/utils/hf.py +++ b/src/kohakuhub/api/repo/utils/hf.py @@ -7,7 +7,6 @@ from typing import Optional from fastapi.responses import Response - class HFErrorCode: """HuggingFace error codes for X-Error-Code header. diff --git a/src/kohakuhub/api/settings.py b/src/kohakuhub/api/settings.py index d575164..be7da5d 100644 --- a/src/kohakuhub/api/settings.py +++ b/src/kohakuhub/api/settings.py @@ -5,17 +5,16 @@ from typing import Optional from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, EmailStr -from kohakuhub.api.repo.utils.hf import hf_repo_not_found -from kohakuhub.auth.dependencies import get_current_user -from kohakuhub.auth.permissions import check_repo_delete_permission -from kohakuhub.db_async import ( - execute_db_query, +from kohakuhub.db import Organization, Repository, User, UserOrganization +from kohakuhub.db_operations import ( get_organization, get_repository, get_user_organization, ) -from kohakuhub.db import Organization, Repository, User, UserOrganization from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user +from kohakuhub.auth.permissions import check_repo_delete_permission +from kohakuhub.api.repo.utils.hf import hf_repo_not_found logger = get_logger("SETTINGS") @@ -56,19 +55,13 @@ async def update_user_settings( # Update fields if provided if req.email is not None: # Check if email is already taken by another user - def _check_email(): - return User.get_or_none((User.email == req.email) & (User.id != user.id)) - - existing = await execute_db_query(_check_email) + existing = User.get_or_none((User.email == req.email) & (User.id != user.id)) if existing: raise HTTPException(400, detail="Email already in use") - def _update_email(): - User.update(email=req.email, email_verified=False).where( - User.id == user.id - ).execute() - - await execute_db_query(_update_email) + User.update(email=req.email, email_verified=False).where( + User.id == user.id + ).execute() # TODO: Send new verification email return {"success": True, "message": "User settings updated successfully"} @@ -99,12 +92,12 @@ async def update_organization_settings( Returns: Success message """ - org = await get_organization(org_name) + org = get_organization(org_name) if not org: raise HTTPException(404, detail="Organization not found") # Check if user is admin of the organization - user_org = await get_user_organization(user.id, org.id) + user_org = get_user_organization(user.id, org.id) if not user_org or user_org.role not in ["admin", "super-admin"]: raise HTTPException( 403, detail="Not authorized to update organization settings" @@ -112,13 +105,9 @@ async def update_organization_settings( # Update fields if provided if req.description is not None: - - def _update_org(): - Organization.update(description=req.description).where( - Organization.id == org.id - ).execute() - - await execute_db_query(_update_org) + Organization.update(description=req.description).where( + Organization.id == org.id + ).execute() return {"success": True, "message": "Organization settings updated successfully"} @@ -160,7 +149,7 @@ async def update_repo_settings( repo_id = f"{namespace}/{name}" # Check if repository exists - repo_row = await get_repository(repo_type, namespace, name) + repo_row = get_repository(repo_type, namespace, name) if not repo_row: return hf_repo_not_found(repo_id, repo_type) @@ -170,13 +159,9 @@ async def update_repo_settings( # Update fields if provided if payload.private is not None: - - def _update_private(): - Repository.update(private=payload.private).where( - Repository.id == repo_row.id - ).execute() - - await execute_db_query(_update_private) + Repository.update(private=payload.private).where( + Repository.id == repo_row.id + ).execute() # Note: gated functionality not yet implemented in database schema # Would require adding a 'gated' field to Repository model diff --git a/src/kohakuhub/async_utils.py b/src/kohakuhub/async_utils.py index 7928796..4e7a018 100644 --- a/src/kohakuhub/async_utils.py +++ b/src/kohakuhub/async_utils.py @@ -2,7 +2,7 @@ import asyncio from concurrent.futures import ThreadPoolExecutor -from functools import wraps, partial +from functools import partial, wraps from typing import Callable, TypeVar # Create separate thread pool executors for different types of operations @@ -14,8 +14,9 @@ _lakefs_executor = ThreadPoolExecutor( max_workers=32, thread_name_prefix="kohakuhub_lakefs" ) -# Peewee operations should use single worker to avoid database concurrency issues -# (unless we implement proper transaction handling for multi-instance deployment) +# DEPRECATED: DB operations are now synchronous and don't need thread pool +# Database operations use db.atomic() for transactions instead +# Keeping this for backward compatibility only _db_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="kohakuhub_db") T = TypeVar("T") @@ -56,7 +57,10 @@ async def run_in_lakefs_executor(func: Callable[..., T], *args, **kwargs) -> T: async def run_in_db_executor(func: Callable[..., T], *args, **kwargs) -> T: - """Run database operation in dedicated single-worker DB thread pool. + """DEPRECATED: Run database operation in dedicated single-worker DB thread pool. + + Database operations are now synchronous and use db.atomic() for transactions. + This function is kept for backward compatibility only. Args: func: Database operation to run @@ -268,6 +272,7 @@ def get_async_lakefs_client(): Returns: LakeFSRestClient instance (already async, no wrapping needed) """ + # Import here to avoid circular dependency from kohakuhub.lakefs_rest_client import get_lakefs_rest_client return get_lakefs_rest_client() diff --git a/src/kohakuhub/auth/__init__.py b/src/kohakuhub/auth/__init__.py index e5bd38b..c06462a 100644 --- a/src/kohakuhub/auth/__init__.py +++ b/src/kohakuhub/auth/__init__.py @@ -1,6 +1,6 @@ """Authentication module.""" -from .routes import router from .dependencies import get_current_user, get_optional_user +from .routes import router __all__ = ["router", "get_current_user", "get_optional_user"] diff --git a/src/kohakuhub/auth/dependencies.py b/src/kohakuhub/auth/dependencies.py index 16e5605..adac511 100644 --- a/src/kohakuhub/auth/dependencies.py +++ b/src/kohakuhub/auth/dependencies.py @@ -2,11 +2,12 @@ from datetime import datetime, timezone from typing import Optional + from fastapi import Cookie, Header, HTTPException -from ..db import User, Session, Token -from .utils import hash_token -from ..logger import get_logger +from kohakuhub.db import Session, Token, User +from kohakuhub.logger import get_logger +from kohakuhub.auth.utils import hash_token logger = get_logger("AUTH") diff --git a/src/kohakuhub/auth/email.py b/src/kohakuhub/auth/email.py index 4989f3d..3a17f5e 100644 --- a/src/kohakuhub/auth/email.py +++ b/src/kohakuhub/auth/email.py @@ -1,11 +1,11 @@ """Email utilities for authentication.""" import smtplib -from email.mime.text import MIMEText from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText -from ..config import cfg -from ..logger import get_logger +from kohakuhub.config import cfg +from kohakuhub.logger import get_logger logger = get_logger("EMAIL") diff --git a/src/kohakuhub/auth/permissions.py b/src/kohakuhub/auth/permissions.py index 9870bbf..d1c84b8 100644 --- a/src/kohakuhub/auth/permissions.py +++ b/src/kohakuhub/auth/permissions.py @@ -1,10 +1,10 @@ """Authorization and permission checking utilities.""" from typing import Optional + from fastapi import HTTPException -from ..db import User, Repository, Organization, UserOrganization - +from kohakuhub.db import Organization, Repository, User, UserOrganization def check_namespace_permission( namespace: str, user: User, require_admin: bool = False diff --git a/src/kohakuhub/auth/routes.py b/src/kohakuhub/auth/routes.py index 343f4d7..dfed78c 100644 --- a/src/kohakuhub/auth/routes.py +++ b/src/kohakuhub/auth/routes.py @@ -3,12 +3,13 @@ import asyncio from datetime import datetime, timezone -from fastapi import APIRouter, HTTPException, Response, Depends +from fastapi import APIRouter, Depends, HTTPException, Response from fastapi.responses import RedirectResponse from pydantic import BaseModel, EmailStr -from ..config import cfg -from ..db_async import ( +from kohakuhub.config import cfg +from kohakuhub.db import EmailVerification, Session, Token, User, db +from kohakuhub.db_operations import ( create_email_verification, create_session, create_token, @@ -16,7 +17,6 @@ from ..db_async import ( delete_email_verification, delete_session, delete_token, - execute_db_query, get_email_verification, get_session, get_token_by_hash, @@ -26,20 +26,19 @@ from ..db_async import ( list_user_tokens, update_user, ) -from ..db import User, EmailVerification, Session, Token -from ..logger import get_logger +from kohakuhub.logger import get_logger +from kohakuhub.auth.dependencies import get_current_user, get_optional_user +from kohakuhub.auth.email import send_verification_email +from kohakuhub.auth.utils import ( + generate_session_secret, + generate_token, + get_expiry_time, + hash_password, + hash_token, + verify_password, +) logger = get_logger("AUTH") -from .utils import ( - hash_password, - verify_password, - generate_token, - hash_token, - generate_session_secret, - get_expiry_time, -) -from .email import send_verification_email -from .dependencies import get_current_user, get_optional_user router = APIRouter(prefix="/auth", tags=["auth"]) @@ -66,29 +65,32 @@ async def register(req: RegisterRequest): logger.info(f"Registration attempt for username: {req.username}") - # Check if username or email already exists - if await get_user_by_username(req.username): - logger.warning(f"Registration failed: username '{req.username}' already exists") - raise HTTPException(400, detail="Username already exists") + # Check if username or email already exists and create user atomically + with db.atomic(): + if get_user_by_username(req.username): + logger.warning( + f"Registration failed: username '{req.username}' already exists" + ) + raise HTTPException(400, detail="Username already exists") - if await get_user_by_email(req.email): - logger.warning(f"Registration failed: email '{req.email}' already exists") - raise HTTPException(400, detail="Email already exists") + if get_user_by_email(req.email): + logger.warning(f"Registration failed: email '{req.email}' already exists") + raise HTTPException(400, detail="Email already exists") - # Create user - user = await create_user( - username=req.username, - email=req.email, - password_hash=hash_password(req.password), - email_verified=not cfg.auth.require_email_verification, - ) + # Create user + user = create_user( + username=req.username, + email=req.email, + password_hash=hash_password(req.password), + email_verified=not cfg.auth.require_email_verification, + ) logger.success(f"User registered: {user.username} (id={user.id})") # Send verification email if required if cfg.auth.require_email_verification: token = generate_token() - await create_email_verification( + create_email_verification( user_id=user.id, token=token, expires_at=get_expiry_time(24) ) @@ -120,7 +122,7 @@ async def verify_email(token: str, response: Response): """Verify email with token and automatically log in user.""" logger.info(f"Email verification attempt with token: {token[:8]}...") - verification = await get_email_verification(token) + verification = get_email_verification(token) if not verification: logger.warning(f"Invalid verification token: {token[:8]}...") @@ -143,29 +145,28 @@ async def verify_email(token: str, response: Response): ) # Get user - user = await get_user_by_id(verification.user) + user = get_user_by_id(verification.user) if not user: logger.error(f"User not found for verification token: {token[:8]}...") return RedirectResponse(url="/?error=user_not_found", status_code=302) - # Update user email verification status - await update_user(user, email_verified=True) + # Update user, delete verification, and create session atomically + with db.atomic(): + update_user(user, email_verified=True) + delete_email_verification(verification) - # Delete verification token - await delete_email_verification(verification) + logger.success(f"Email verified for user: {user.username}") - logger.success(f"Email verified for user: {user.username}") + # Create session for auto-login + session_id = generate_token() + session_secret = generate_session_secret() - # Create session for auto-login - session_id = generate_token() - session_secret = generate_session_secret() - - await create_session( - session_id=session_id, - user_id=user.id, - secret=session_secret, - expires_at=get_expiry_time(cfg.auth.session_expire_hours), - ) + create_session( + session_id=session_id, + user_id=user.id, + secret=session_secret, + expires_at=get_expiry_time(cfg.auth.session_expire_hours), + ) # Set session cookie redirect_response = RedirectResponse(url=f"/{user.username}", status_code=302) @@ -190,7 +191,7 @@ async def login(req: LoginRequest, response: Response): logger.info(f"Login attempt for user: {req.username}") - user = await get_user_by_username(req.username) + user = get_user_by_username(req.username) if not user or not verify_password(req.password, user.password_hash): logger.warning(f"Failed login attempt for: {req.username}") @@ -208,7 +209,7 @@ async def login(req: LoginRequest, response: Response): session_id = generate_token() session_secret = generate_session_secret() - await create_session( + create_session( session_id=session_id, user_id=user.id, secret=session_secret, @@ -241,10 +242,7 @@ async def logout(response: Response, user: User = Depends(get_current_user)): logger.info(f"Logout request for user: {user.username}") # Delete all user sessions - def _delete_sessions(): - return Session.delete().where(Session.user_id == user.id).execute() - - deleted_count = await execute_db_query(_delete_sessions) + deleted_count = Session.delete().where(Session.user_id == user.id).execute() # Clear cookie response.delete_cookie(key="session_id") @@ -273,7 +271,7 @@ def get_me(user: User = Depends(get_current_user)): async def list_tokens(user: User = Depends(get_current_user)): """List user's API tokens.""" - tokens = await list_user_tokens(user.id) + tokens = list_user_tokens(user.id) return { "tokens": [ @@ -299,15 +297,10 @@ async def create_token_endpoint( token_hash_val = hash_token(token_str) # Save to database - token = await create_token( - user_id=user.id, token_hash=token_hash_val, name=req.name - ) + token = create_token(user_id=user.id, token_hash=token_hash_val, name=req.name) # Get session secret for encryption (if in web session) - def _get_session(): - return Session.get_or_none(Session.user_id == user.id) - - session = await execute_db_query(_get_session) + session = Session.get_or_none(Session.user_id == user.id) session_secret = session.secret if session else None return { @@ -323,14 +316,11 @@ async def create_token_endpoint( async def revoke_token(token_id: int, user: User = Depends(get_current_user)): """Revoke an API token.""" - def _get_token(): - return Token.get_or_none((Token.id == token_id) & (Token.user_id == user.id)) - - token = await execute_db_query(_get_token) + token = Token.get_or_none((Token.id == token_id) & (Token.user_id == user.id)) if not token: raise HTTPException(404, detail="Token not found") - await delete_token(token.id) + delete_token(token.id) return {"success": True, "message": "Token revoked successfully"} diff --git a/src/kohakuhub/auth/utils.py b/src/kohakuhub/auth/utils.py index cb06c33..7ce6cc7 100644 --- a/src/kohakuhub/auth/utils.py +++ b/src/kohakuhub/auth/utils.py @@ -1,10 +1,10 @@ """Authentication utilities.""" -import secrets -import hashlib -import bcrypt from datetime import datetime, timedelta, timezone +import hashlib +import secrets +import bcrypt def hash_password(password: str) -> str: """Hash password with bcrypt.""" diff --git a/src/kohakuhub/config.py b/src/kohakuhub/config.py index a58d24d..0a57ec8 100644 --- a/src/kohakuhub/config.py +++ b/src/kohakuhub/config.py @@ -3,6 +3,7 @@ import os import tomllib from functools import lru_cache + from pydantic import BaseModel # Default configuration values diff --git a/src/kohakuhub/db.py b/src/kohakuhub/db.py index f462e60..7a28fa6 100644 --- a/src/kohakuhub/db.py +++ b/src/kohakuhub/db.py @@ -1,22 +1,24 @@ """Database models for Kohaku Hub.""" -from functools import partial from datetime import datetime, timezone +from functools import partial + from peewee import ( AutoField, BigIntegerField, BooleanField, CharField, DateTimeField, + ForeignKeyField, IntegerField, Model, - SqliteDatabase, PostgresqlDatabase, + SqliteDatabase, TextField, - ForeignKeyField, ) -from .config import cfg -from .logger import get_logger + +from kohakuhub.config import cfg +from kohakuhub.logger import get_logger logger = get_logger("DB") diff --git a/src/kohakuhub/db_async.py b/src/kohakuhub/db_async.py deleted file mode 100644 index 2c4d7dc..0000000 --- a/src/kohakuhub/db_async.py +++ /dev/null @@ -1,660 +0,0 @@ -"""Async wrappers for Peewee database operations. - -All database operations should go through this module to ensure they run -in the dedicated single-worker DB thread pool. -""" - -from typing import TypeVar - -from kohakuhub.async_utils import run_in_db_executor -from kohakuhub.db import ( - Commit, - EmailVerification, - File, - LFSObjectHistory, - Organization, - Repository, - Session, - SSHKey, - StagingUpload, - Token, - User, - UserOrganization, - db, -) - -T = TypeVar("T") - - -# User operations -async def get_user_by_id(user_id: int) -> User | None: - """Get user by ID.""" - return await run_in_db_executor(lambda: User.get_or_none(User.id == user_id)) - - -async def get_user_by_username(username: str) -> User | None: - """Get user by username.""" - return await run_in_db_executor(lambda: User.get_or_none(User.username == username)) - - -async def get_user_by_email(email: str) -> User | None: - """Get user by email.""" - return await run_in_db_executor(lambda: User.get_or_none(User.email == email)) - - -async def create_user( - username: str, - email: str, - password_hash: str, - email_verified: bool = False, - is_active: bool = True, - private_quota_bytes: int | None = None, - public_quota_bytes: int | None = None, -) -> User: - """Create a new user with default quota values if not specified.""" - from kohakuhub.config import cfg - - # Apply default quotas if not explicitly provided - if private_quota_bytes is None: - private_quota_bytes = cfg.quota.default_user_private_quota_bytes - if public_quota_bytes is None: - public_quota_bytes = cfg.quota.default_user_public_quota_bytes - - def _create(): - return User.create( - username=username, - email=email, - password_hash=password_hash, - email_verified=email_verified, - is_active=is_active, - private_quota_bytes=private_quota_bytes, - public_quota_bytes=public_quota_bytes, - ) - - return await run_in_db_executor(_create) - - -async def update_user(user: User, **fields) -> None: - """Update user fields.""" - - def _update(): - for key, value in fields.items(): - setattr(user, key, value) - user.save() - - await run_in_db_executor(_update) - - -async def delete_user(user: User) -> None: - """Delete a user and all related data. - - This will delete: - - All sessions - - All tokens - - All email verifications - - All user-organization memberships - - The user record itself - - Note: Repositories owned by the user are NOT deleted automatically. - Delete repositories separately before calling this function. - """ - - def _delete(): - user_id = user.id - - # Delete all sessions - Session.delete().where(Session.user_id == user_id).execute() - - # Delete all tokens - Token.delete().where(Token.user_id == user_id).execute() - - # Delete all email verifications - EmailVerification.delete().where(EmailVerification.user == user_id).execute() - - # Delete all user-organization memberships - UserOrganization.delete().where(UserOrganization.user == user_id).execute() - - # Finally delete the user - user.delete_instance() - - await run_in_db_executor(_delete) - - -# Session operations -async def get_session(session_id: str) -> Session | None: - """Get session by session ID.""" - return await run_in_db_executor( - lambda: Session.get_or_none(Session.session_id == session_id) - ) - - -async def create_session( - session_id: str, user_id: int, secret: str, expires_at -) -> Session: - """Create a new session.""" - - def _create(): - return Session.create( - session_id=session_id, user_id=user_id, secret=secret, expires_at=expires_at - ) - - return await run_in_db_executor(_create) - - -async def delete_session(session_id: str) -> None: - """Delete a session.""" - - def _delete(): - Session.delete().where(Session.session_id == session_id).execute() - - await run_in_db_executor(_delete) - - -# Token operations -async def get_token_by_hash(token_hash: str) -> Token | None: - """Get token by hash.""" - return await run_in_db_executor( - lambda: Token.get_or_none(Token.token_hash == token_hash) - ) - - -async def create_token(user_id: int, token_hash: str, name: str) -> Token: - """Create a new API token.""" - - def _create(): - return Token.create(user_id=user_id, token_hash=token_hash, name=name) - - return await run_in_db_executor(_create) - - -async def list_user_tokens(user_id: int) -> list[Token]: - """List all tokens for a user.""" - - def _list(): - return list(Token.select().where(Token.user_id == user_id)) - - return await run_in_db_executor(_list) - - -async def delete_token(token_id: int) -> None: - """Delete a token.""" - - def _delete(): - Token.delete().where(Token.id == token_id).execute() - - await run_in_db_executor(_delete) - - -async def update_token_last_used(token: Token, last_used) -> None: - """Update token last used timestamp.""" - - def _update(): - token.last_used = last_used - token.save() - - await run_in_db_executor(_update) - - -# Repository operations -async def get_repository( - repo_type: str, namespace: str, name: str -) -> Repository | None: - """Get repository by type, namespace, and name.""" - return await run_in_db_executor( - lambda: Repository.get_or_none( - Repository.repo_type == repo_type, - Repository.namespace == namespace, - Repository.name == name, - ) - ) - - -async def get_repository_by_full_id(full_id: str, repo_type: str) -> Repository | None: - """Get repository by full ID and type.""" - return await run_in_db_executor( - lambda: Repository.get_or_none( - Repository.full_id == full_id, Repository.repo_type == repo_type - ) - ) - - -async def create_repository( - repo_type: str, - namespace: str, - name: str, - full_id: str, - private: bool, - owner_id: int, -) -> Repository: - """Create a new repository.""" - - def _create(): - return Repository.create( - repo_type=repo_type, - namespace=namespace, - name=name, - full_id=full_id, - private=private, - owner_id=owner_id, - ) - - return await run_in_db_executor(_create) - - -async def delete_repository(repo: Repository) -> None: - """Delete a repository.""" - await run_in_db_executor(repo.delete_instance) - - -async def update_repository(repo: Repository, **fields) -> None: - """Update repository fields.""" - - def _update(): - for key, value in fields.items(): - setattr(repo, key, value) - repo.save() - - await run_in_db_executor(_update) - - -async def list_repositories( - repo_type: str | None = None, namespace: str | None = None, limit: int | None = None -) -> list[Repository]: - """List repositories with optional filters.""" - - def _list(): - query = Repository.select() - if repo_type: - query = query.where(Repository.repo_type == repo_type) - if namespace: - query = query.where(Repository.namespace == namespace) - if limit: - query = query.limit(limit) - return list(query) - - return await run_in_db_executor(_list) - - -# File operations -async def get_file(repo_full_id: str, path_in_repo: str) -> File | None: - """Get file by repository and path.""" - return await run_in_db_executor( - lambda: File.get_or_none( - File.repo_full_id == repo_full_id, File.path_in_repo == path_in_repo - ) - ) - - -async def get_file_by_sha256(sha256: str) -> File | None: - """Get file by SHA256 hash.""" - return await run_in_db_executor(lambda: File.get_or_none(File.sha256 == sha256)) - - -async def create_file( - repo_full_id: str, path_in_repo: str, size: int, sha256: str, lfs: bool -) -> File: - """Create a new file record.""" - - def _create(): - return File.create( - repo_full_id=repo_full_id, - path_in_repo=path_in_repo, - size=size, - sha256=sha256, - lfs=lfs, - ) - - return await run_in_db_executor(_create) - - -async def update_file(file: File, **fields) -> None: - """Update file fields.""" - - def _update(): - for key, value in fields.items(): - setattr(file, key, value) - file.save() - - await run_in_db_executor(_update) - - -async def delete_file(file: File) -> None: - """Delete a file record.""" - await run_in_db_executor(file.delete_instance) - - -# Organization operations -async def get_organization(name: str) -> Organization | None: - """Get organization by name.""" - return await run_in_db_executor( - lambda: Organization.get_or_none(Organization.name == name) - ) - - -async def create_organization( - name: str, - description: str | None = None, - private_quota_bytes: int | None = None, - public_quota_bytes: int | None = None, -) -> Organization: - """Create a new organization with default quota values if not specified.""" - from kohakuhub.config import cfg - - # Apply default quotas if not explicitly provided - if private_quota_bytes is None: - private_quota_bytes = cfg.quota.default_org_private_quota_bytes - if public_quota_bytes is None: - public_quota_bytes = cfg.quota.default_org_public_quota_bytes - - def _create(): - return Organization.create( - name=name, - description=description, - private_quota_bytes=private_quota_bytes, - public_quota_bytes=public_quota_bytes, - ) - - return await run_in_db_executor(_create) - - -async def update_organization(org: Organization, **fields) -> None: - """Update organization fields.""" - - def _update(): - for key, value in fields.items(): - setattr(org, key, value) - org.save() - - await run_in_db_executor(_update) - - -async def delete_organization(org: Organization) -> None: - """Delete an organization.""" - await run_in_db_executor(org.delete_instance) - - -# UserOrganization operations -async def get_user_organization(user_id: int, org_id: int) -> UserOrganization | None: - """Get user-organization relationship.""" - return await run_in_db_executor( - lambda: UserOrganization.get_or_none( - UserOrganization.user == user_id, UserOrganization.organization == org_id - ) - ) - - -async def create_user_organization( - user_id: int, org_id: int, role: str -) -> UserOrganization: - """Create user-organization relationship.""" - - def _create(): - return UserOrganization.create(user=user_id, organization=org_id, role=role) - - return await run_in_db_executor(_create) - - -async def update_user_organization(user_org: UserOrganization, **fields) -> None: - """Update user-organization fields.""" - - def _update(): - for key, value in fields.items(): - setattr(user_org, key, value) - user_org.save() - - await run_in_db_executor(_update) - - -async def delete_user_organization(user_org: UserOrganization) -> None: - """Delete user-organization relationship.""" - await run_in_db_executor(user_org.delete_instance) - - -async def list_user_organizations(user_id: int) -> list[UserOrganization]: - """List all organizations for a user.""" - - def _list(): - return list(UserOrganization.select().where(UserOrganization.user == user_id)) - - return await run_in_db_executor(_list) - - -async def list_organization_members(org_id: int) -> list[UserOrganization]: - """List all members of an organization.""" - - def _list(): - return list( - UserOrganization.select().where(UserOrganization.organization == org_id) - ) - - return await run_in_db_executor(_list) - - -# StagingUpload operations -async def create_staging_upload( - repo_full_id: str, - repo_type: str, - revision: str, - path_in_repo: str, - sha256: str, - size: int, - storage_key: str, - lfs: bool, - upload_id: str | None = None, -) -> StagingUpload: - """Create a staging upload record.""" - - def _create(): - return StagingUpload.create( - repo_full_id=repo_full_id, - repo_type=repo_type, - revision=revision, - path_in_repo=path_in_repo, - sha256=sha256, - size=size, - storage_key=storage_key, - lfs=lfs, - upload_id=upload_id, - ) - - return await run_in_db_executor(_create) - - -async def delete_staging_upload(staging: StagingUpload) -> None: - """Delete a staging upload record.""" - await run_in_db_executor(staging.delete_instance) - - -# LFSObjectHistory operations -async def create_lfs_history( - repo_full_id: str, path_in_repo: str, sha256: str, size: int, commit_id: str -) -> LFSObjectHistory: - """Create LFS object history record.""" - - def _create(): - return LFSObjectHistory.create( - repo_full_id=repo_full_id, - path_in_repo=path_in_repo, - sha256=sha256, - size=size, - commit_id=commit_id, - ) - - return await run_in_db_executor(_create) - - -async def list_lfs_history( - repo_full_id: str, path_in_repo: str, limit: int | None = None -) -> list[LFSObjectHistory]: - """List LFS history for a file.""" - - def _list(): - query = ( - LFSObjectHistory.select() - .where( - LFSObjectHistory.repo_full_id == repo_full_id, - LFSObjectHistory.path_in_repo == path_in_repo, - ) - .order_by(LFSObjectHistory.created_at.desc()) - ) - if limit: - query = query.limit(limit) - return list(query) - - return await run_in_db_executor(_list) - - -# Commit operations -async def create_commit( - commit_id: str, - repo_full_id: str, - repo_type: str, - branch: str, - user_id: int, - username: str, - message: str, - description: str = "", -) -> Commit: - """Create commit record.""" - from kohakuhub.db import Commit - - def _create(): - return Commit.create( - commit_id=commit_id, - repo_full_id=repo_full_id, - repo_type=repo_type, - branch=branch, - user_id=user_id, - username=username, - message=message, - description=description, - ) - - return await run_in_db_executor(_create) - - -async def get_commit(commit_id: str, repo_full_id: str) -> Commit | None: - """Get commit by ID and repo.""" - from kohakuhub.db import Commit - - return await run_in_db_executor( - lambda: Commit.get_or_none( - Commit.commit_id == commit_id, Commit.repo_full_id == repo_full_id - ) - ) - - -async def list_commits_by_repo( - repo_full_id: str, branch: str | None = None, limit: int | None = None -) -> list[Commit]: - """List commits for a repository.""" - from kohakuhub.db import Commit - - def _list(): - query = Commit.select().where(Commit.repo_full_id == repo_full_id) - if branch: - query = query.where(Commit.branch == branch) - query = query.order_by(Commit.created_at.desc()) - if limit: - query = query.limit(limit) - return list(query) - - return await run_in_db_executor(_list) - - -# EmailVerification operations -async def create_email_verification( - user_id: int, token: str, expires_at -) -> EmailVerification: - """Create email verification token.""" - - def _create(): - return EmailVerification.create( - user=user_id, token=token, expires_at=expires_at - ) - - return await run_in_db_executor(_create) - - -async def get_email_verification(token: str) -> EmailVerification | None: - """Get email verification by token.""" - return await run_in_db_executor( - lambda: EmailVerification.get_or_none(EmailVerification.token == token) - ) - - -async def delete_email_verification(verification: EmailVerification) -> None: - """Delete email verification.""" - await run_in_db_executor(verification.delete_instance) - - -# SSH Key operations -async def get_ssh_key_by_id(key_id: int) -> SSHKey | None: - """Get SSH key by ID.""" - return await run_in_db_executor(lambda: SSHKey.get_or_none(SSHKey.id == key_id)) - - -async def get_ssh_key_by_fingerprint(fingerprint: str) -> SSHKey | None: - """Get SSH key by fingerprint.""" - return await run_in_db_executor( - lambda: SSHKey.get_or_none(SSHKey.fingerprint == fingerprint) - ) - - -async def list_user_ssh_keys(user_id: int) -> list[SSHKey]: - """List all SSH keys for a user.""" - - def _list(): - return list( - SSHKey.select() - .where(SSHKey.user_id == user_id) - .order_by(SSHKey.created_at.desc()) - ) - - return await run_in_db_executor(_list) - - -async def create_ssh_key( - user_id: int, key_type: str, public_key: str, fingerprint: str, title: str -) -> SSHKey: - """Create a new SSH key.""" - - def _create(): - return SSHKey.create( - user_id=user_id, - key_type=key_type, - public_key=public_key, - fingerprint=fingerprint, - title=title, - ) - - return await run_in_db_executor(_create) - - -async def update_ssh_key(key: SSHKey, **fields) -> None: - """Update SSH key fields.""" - - def _update(): - for field_name, value in fields.items(): - setattr(key, field_name, value) - key.save() - - await run_in_db_executor(_update) - - -async def delete_ssh_key(key: SSHKey) -> None: - """Delete an SSH key.""" - await run_in_db_executor(key.delete_instance) - - -# Generic query helper -async def execute_db_query(query_func): - """Execute a custom database query in the DB executor. - - Args: - query_func: A callable that performs the database operation - - Returns: - Result of the query function - """ - return await run_in_db_executor(query_func) diff --git a/src/kohakuhub/db_operations.py b/src/kohakuhub/db_operations.py new file mode 100644 index 0000000..79b3eec --- /dev/null +++ b/src/kohakuhub/db_operations.py @@ -0,0 +1,501 @@ +"""Synchronous database operations with transaction support. + +All database operations are now synchronous. Use db.atomic() for transactions. +Multi-worker deployment is safe with proper transaction usage. +""" + +from kohakuhub.config import cfg +from kohakuhub.db import ( + Commit, + EmailVerification, + File, + LFSObjectHistory, + Organization, + Repository, + Session, + SSHKey, + StagingUpload, + Token, + User, + UserOrganization, + db, +) + +# User operations +def get_user_by_id(user_id: int) -> User | None: + """Get user by ID.""" + return User.get_or_none(User.id == user_id) + + +def get_user_by_username(username: str) -> User | None: + """Get user by username.""" + return User.get_or_none(User.username == username) + + +def get_user_by_email(email: str) -> User | None: + """Get user by email.""" + return User.get_or_none(User.email == email) + + +def create_user( + username: str, + email: str, + password_hash: str, + email_verified: bool = False, + is_active: bool = True, + private_quota_bytes: int | None = None, + public_quota_bytes: int | None = None, +) -> User: + """Create a new user with default quota values if not specified. + + NOTE: Wrap in db.atomic() if checking existence first to prevent race conditions. + """ + # Apply default quotas if not explicitly provided + if private_quota_bytes is None: + private_quota_bytes = cfg.quota.default_user_private_quota_bytes + if public_quota_bytes is None: + public_quota_bytes = cfg.quota.default_user_public_quota_bytes + + return User.create( + username=username, + email=email, + password_hash=password_hash, + email_verified=email_verified, + is_active=is_active, + private_quota_bytes=private_quota_bytes, + public_quota_bytes=public_quota_bytes, + ) + + +def update_user(user: User, **fields) -> None: + """Update user fields.""" + for key, value in fields.items(): + setattr(user, key, value) + user.save() + + +def delete_user(user: User) -> None: + """Delete a user and all related data. + + This will delete: + - All sessions + - All tokens + - All SSH keys + - All email verifications + - All user-organization memberships + - The user record itself + + NOTE: This is automatically transactional via db.atomic() wrapper. + """ + with db.atomic(): + user_id = user.id + + # Delete all sessions + Session.delete().where(Session.user_id == user_id).execute() + + # Delete all tokens + Token.delete().where(Token.user_id == user_id).execute() + + # Delete all SSH keys + SSHKey.delete().where(SSHKey.user_id == user_id).execute() + + # Delete all email verifications + EmailVerification.delete().where(EmailVerification.user == user_id).execute() + + # Delete all user-organization memberships + UserOrganization.delete().where(UserOrganization.user == user_id).execute() + + # Finally delete the user + user.delete_instance() + + +# Repository operations +def get_repository(repo_type: str, namespace: str, name: str) -> Repository | None: + """Get repository by type, namespace, and name.""" + return Repository.get_or_none( + Repository.repo_type == repo_type, + Repository.namespace == namespace, + Repository.name == name, + ) + + +def get_repository_by_full_id(full_id: str, repo_type: str) -> Repository | None: + """Get repository by full ID and type.""" + return Repository.get_or_none( + Repository.full_id == full_id, Repository.repo_type == repo_type + ) + + +def create_repository( + repo_type: str, + namespace: str, + name: str, + full_id: str, + private: bool, + owner_id: int, +) -> Repository: + """Create a new repository. + + NOTE: Wrap in db.atomic() if checking existence first. + """ + return Repository.create( + repo_type=repo_type, + namespace=namespace, + name=name, + full_id=full_id, + private=private, + owner_id=owner_id, + ) + + +def delete_repository(repo: Repository) -> None: + """Delete a repository.""" + repo.delete_instance() + + +def update_repository(repo: Repository, **fields) -> None: + """Update repository fields.""" + for key, value in fields.items(): + setattr(repo, key, value) + repo.save() + + +def list_repositories( + repo_type: str | None = None, namespace: str | None = None, limit: int | None = None +) -> list[Repository]: + """List repositories with optional filters.""" + query = Repository.select() + if repo_type: + query = query.where(Repository.repo_type == repo_type) + if namespace: + query = query.where(Repository.namespace == namespace) + if limit: + query = query.limit(limit) + return list(query) + + +# Session operations +def get_session(session_id: str) -> Session | None: + """Get session by session ID.""" + return Session.get_or_none(Session.session_id == session_id) + + +def create_session(session_id: str, user_id: int, secret: str, expires_at) -> Session: + """Create a new session.""" + return Session.create( + session_id=session_id, user_id=user_id, secret=secret, expires_at=expires_at + ) + + +def delete_session(session_id: str) -> None: + """Delete a session.""" + Session.delete().where(Session.session_id == session_id).execute() + + +# Token operations +def get_token_by_hash(token_hash: str) -> Token | None: + """Get token by hash.""" + return Token.get_or_none(Token.token_hash == token_hash) + + +def create_token(user_id: int, token_hash: str, name: str) -> Token: + """Create a new API token.""" + return Token.create(user_id=user_id, token_hash=token_hash, name=name) + + +def list_user_tokens(user_id: int) -> list[Token]: + """List all tokens for a user.""" + return list(Token.select().where(Token.user_id == user_id)) + + +def delete_token(token_id: int) -> None: + """Delete a token.""" + Token.delete().where(Token.id == token_id).execute() + + +def update_token_last_used(token: Token, last_used) -> None: + """Update token last used timestamp.""" + token.last_used = last_used + token.save() + + +# Organization operations +def get_organization(name: str) -> Organization | None: + """Get organization by name.""" + return Organization.get_or_none(Organization.name == name) + + +def create_organization( + name: str, + description: str | None = None, + private_quota_bytes: int | None = None, + public_quota_bytes: int | None = None, +) -> Organization: + """Create a new organization with default quota values if not specified.""" + if private_quota_bytes is None: + private_quota_bytes = cfg.quota.default_org_private_quota_bytes + if public_quota_bytes is None: + public_quota_bytes = cfg.quota.default_org_public_quota_bytes + + return Organization.create( + name=name, + description=description, + private_quota_bytes=private_quota_bytes, + public_quota_bytes=public_quota_bytes, + ) + + +def update_organization(org: Organization, **fields) -> None: + """Update organization fields.""" + for key, value in fields.items(): + setattr(org, key, value) + org.save() + + +def delete_organization(org: Organization) -> None: + """Delete an organization.""" + org.delete_instance() + + +# UserOrganization operations +def get_user_organization(user_id: int, org_id: int) -> UserOrganization | None: + """Get user-organization relationship.""" + return UserOrganization.get_or_none( + UserOrganization.user == user_id, UserOrganization.organization == org_id + ) + + +def create_user_organization(user_id: int, org_id: int, role: str) -> UserOrganization: + """Create user-organization relationship.""" + return UserOrganization.create(user=user_id, organization=org_id, role=role) + + +def update_user_organization(user_org: UserOrganization, **fields) -> None: + """Update user-organization fields.""" + for key, value in fields.items(): + setattr(user_org, key, value) + user_org.save() + + +def delete_user_organization(user_org: UserOrganization) -> None: + """Delete user-organization relationship.""" + user_org.delete_instance() + + +def list_user_organizations(user_id: int) -> list[UserOrganization]: + """List all organizations for a user.""" + return list(UserOrganization.select().where(UserOrganization.user == user_id)) + + +def list_organization_members(org_id: int) -> list[UserOrganization]: + """List all members of an organization.""" + return list( + UserOrganization.select().where(UserOrganization.organization == org_id) + ) + + +# File operations +def get_file(repo_full_id: str, path_in_repo: str) -> File | None: + """Get file by repository and path.""" + return File.get_or_none( + File.repo_full_id == repo_full_id, File.path_in_repo == path_in_repo + ) + + +def get_file_by_sha256(sha256: str) -> File | None: + """Get file by SHA256 hash.""" + return File.get_or_none(File.sha256 == sha256) + + +def create_file( + repo_full_id: str, path_in_repo: str, size: int, sha256: str, lfs: bool +) -> File: + """Create a new file record.""" + return File.create( + repo_full_id=repo_full_id, + path_in_repo=path_in_repo, + size=size, + sha256=sha256, + lfs=lfs, + ) + + +def update_file(file: File, **fields) -> None: + """Update file fields.""" + for key, value in fields.items(): + setattr(file, key, value) + file.save() + + +def delete_file(file: File) -> None: + """Delete a file record.""" + file.delete_instance() + + +# Commit operations +def create_commit( + commit_id: str, + repo_full_id: str, + repo_type: str, + branch: str, + user_id: int, + username: str, + message: str, + description: str = "", +) -> Commit: + """Create commit record.""" + return Commit.create( + commit_id=commit_id, + repo_full_id=repo_full_id, + repo_type=repo_type, + branch=branch, + user_id=user_id, + username=username, + message=message, + description=description, + ) + + +def get_commit(commit_id: str, repo_full_id: str) -> Commit | None: + """Get commit by ID and repo.""" + return Commit.get_or_none( + Commit.commit_id == commit_id, Commit.repo_full_id == repo_full_id + ) + + +def list_commits_by_repo( + repo_full_id: str, branch: str | None = None, limit: int | None = None +) -> list[Commit]: + """List commits for a repository.""" + query = Commit.select().where(Commit.repo_full_id == repo_full_id) + if branch: + query = query.where(Commit.branch == branch) + query = query.order_by(Commit.created_at.desc()) + if limit: + query = query.limit(limit) + return list(query) + + +# SSH Key operations +def get_ssh_key_by_id(key_id: int) -> SSHKey | None: + """Get SSH key by ID.""" + return SSHKey.get_or_none(SSHKey.id == key_id) + + +def get_ssh_key_by_fingerprint(fingerprint: str) -> SSHKey | None: + """Get SSH key by fingerprint.""" + return SSHKey.get_or_none(SSHKey.fingerprint == fingerprint) + + +def list_user_ssh_keys(user_id: int) -> list[SSHKey]: + """List all SSH keys for a user.""" + return list( + SSHKey.select() + .where(SSHKey.user_id == user_id) + .order_by(SSHKey.created_at.desc()) + ) + + +def create_ssh_key( + user_id: int, key_type: str, public_key: str, fingerprint: str, title: str +) -> SSHKey: + """Create a new SSH key.""" + return SSHKey.create( + user_id=user_id, + key_type=key_type, + public_key=public_key, + fingerprint=fingerprint, + title=title, + ) + + +def update_ssh_key(key: SSHKey, **fields) -> None: + """Update SSH key fields.""" + for field_name, value in fields.items(): + setattr(key, field_name, value) + key.save() + + +def delete_ssh_key(key: SSHKey) -> None: + """Delete an SSH key.""" + key.delete_instance() + + +# LFS operations +def create_lfs_history( + repo_full_id: str, path_in_repo: str, sha256: str, size: int, commit_id: str +) -> LFSObjectHistory: + """Create LFS object history record.""" + return LFSObjectHistory.create( + repo_full_id=repo_full_id, + path_in_repo=path_in_repo, + sha256=sha256, + size=size, + commit_id=commit_id, + ) + + +def list_lfs_history( + repo_full_id: str, path_in_repo: str, limit: int | None = None +) -> list[LFSObjectHistory]: + """List LFS history for a file.""" + query = ( + LFSObjectHistory.select() + .where( + LFSObjectHistory.repo_full_id == repo_full_id, + LFSObjectHistory.path_in_repo == path_in_repo, + ) + .order_by(LFSObjectHistory.created_at.desc()) + ) + if limit: + query = query.limit(limit) + return list(query) + + +# Email verification operations +def create_email_verification( + user_id: int, token: str, expires_at +) -> EmailVerification: + """Create email verification token.""" + return EmailVerification.create(user=user_id, token=token, expires_at=expires_at) + + +def get_email_verification(token: str) -> EmailVerification | None: + """Get email verification by token.""" + return EmailVerification.get_or_none(EmailVerification.token == token) + + +def delete_email_verification(verification: EmailVerification) -> None: + """Delete email verification.""" + verification.delete_instance() + + +# Staging upload operations +def create_staging_upload( + repo_full_id: str, + repo_type: str, + revision: str, + path_in_repo: str, + sha256: str, + size: int, + storage_key: str, + lfs: bool, + upload_id: str | None = None, +) -> StagingUpload: + """Create a staging upload record.""" + return StagingUpload.create( + repo_full_id=repo_full_id, + repo_type=repo_type, + revision=revision, + path_in_repo=path_in_repo, + sha256=sha256, + size=size, + storage_key=storage_key, + lfs=lfs, + upload_id=upload_id, + ) + + +def delete_staging_upload(staging: StagingUpload) -> None: + """Delete a staging upload record.""" + staging.delete_instance() diff --git a/src/kohakuhub/logger.py b/src/kohakuhub/logger.py index 2cb48c3..25b44c6 100644 --- a/src/kohakuhub/logger.py +++ b/src/kohakuhub/logger.py @@ -1,11 +1,11 @@ """Custom logging module for KohakuHub with colored output and formatted tracebacks.""" +import os import sys import traceback as tb from datetime import datetime -from typing import Optional from enum import Enum - +from typing import Optional class Color: """ANSI color codes for terminal output.""" @@ -73,7 +73,7 @@ class Logger: def _format_message(self, level: LogLevel, message: str) -> str: """Format log message with colors and structure. - Format: [LEVEL][API-NAME][HH:MM:SS] message + Format: [LEVEL][API-NAME][Worker:PID][HH:MM:SS] message Args: level: Log level @@ -84,11 +84,13 @@ class Logger: """ level_name, level_color = level.value timestamp = self._get_timestamp() + worker_id = os.getpid() # Process ID identifies the worker # Build formatted message parts = [ f"{level_color}[{level_name}]{Color.RESET}", f"{Color.BRIGHT_MAGENTA}[{self.api_name}]{Color.RESET}", + f"{Color.BLUE}[W:{worker_id}]{Color.RESET}", f"{Color.BRIGHT_BLACK}[{timestamp}]{Color.RESET}", message, ] diff --git a/src/kohakuhub/main.py b/src/kohakuhub/main.py index 43be876..7ed23f3 100644 --- a/src/kohakuhub/main.py +++ b/src/kohakuhub/main.py @@ -5,36 +5,23 @@ from contextlib import asynccontextmanager from fastapi import Depends, FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware +from kohakuhub.api import admin, branches, files, misc, settings +from kohakuhub.auth import router as auth_router from kohakuhub.config import cfg from kohakuhub.db import Repository, User from kohakuhub.logger import get_logger -from kohakuhub.utils.s3 import init_storage -from kohakuhub.api import ( - admin, - branches, - files, - misc, - settings, -) +from kohakuhub.api.commit import history as commit_history +from kohakuhub.api.commit import router as commits from kohakuhub.api.files import resolve_file_get, resolve_file_head from kohakuhub.api.org import router as org from kohakuhub.api.quota import router as quota -from kohakuhub.api.commit import ( - history as commit_history, - router as commits, -) -from kohakuhub.api.git.routers import ( - http as git_http, - lfs, - ssh_keys, -) -from kohakuhub.api.repo.routers import ( - crud as repo_crud, - info as repo_info, - tree as repo_tree, -) -from kohakuhub.auth import router as auth_router from kohakuhub.auth.dependencies import get_optional_user +from kohakuhub.utils.s3 import init_storage +from kohakuhub.api.git.routers import http as git_http +from kohakuhub.api.git.routers import lfs, ssh_keys +from kohakuhub.api.repo.routers import crud as repo_crud +from kohakuhub.api.repo.routers import info as repo_info +from kohakuhub.api.repo.routers import tree as repo_tree logger = get_logger("MAIN") diff --git a/src/kohakuhub/utils/lakefs.py b/src/kohakuhub/utils/lakefs.py index 631ed89..eb96a9a 100644 --- a/src/kohakuhub/utils/lakefs.py +++ b/src/kohakuhub/utils/lakefs.py @@ -3,7 +3,6 @@ from kohakuhub.config import cfg from kohakuhub.lakefs_rest_client import LakeFSRestClient, get_lakefs_rest_client - def get_lakefs_client() -> LakeFSRestClient: """Get configured LakeFS REST client. diff --git a/tests/conftest.py b/tests/conftest.py index c61f4dc..a3ba31c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -190,7 +190,7 @@ def test_org(authenticated_http_client): This fixture creates an organization for testing org-related features. """ - from kohakuhub.org.routes import CreateOrganizationPayload + from kohakuhub.api.org.router import CreateOrganizationPayload client = authenticated_http_client diff --git a/tests/test_auth.py b/tests/test_auth.py index a407ccd..6a93ba0 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -151,7 +151,7 @@ class TestAuthentication: assert resp.status_code == 401, "Should require authentication" # Try to create repo without auth - from kohakuhub.api.routers.repo_crud import CreateRepoPayload + from kohakuhub.api.repo.routers.crud import CreateRepoPayload payload = CreateRepoPayload(type="model", name="test-repo") diff --git a/tests/test_org.py b/tests/test_org.py index 0dadb08..c2321c7 100644 --- a/tests/test_org.py +++ b/tests/test_org.py @@ -16,7 +16,7 @@ class TestOrganization: def test_create_organization(self, authenticated_http_client): """Test organization creation.""" - from kohakuhub.org.routes import CreateOrganizationPayload + from kohakuhub.api.org.router import CreateOrganizationPayload unique_id = uuid.uuid4().hex[:6] org_name = f"org-{unique_id}" @@ -51,7 +51,7 @@ class TestOrganization: def test_add_remove_member(self, authenticated_http_client, test_org): """Test adding and removing organization members.""" from kohakuhub.auth.routes import RegisterRequest - from kohakuhub.org.routes import AddMemberPayload + from kohakuhub.api.org.router import AddMemberPayload # Create a new user to add as member unique_id = uuid.uuid4().hex[:6] @@ -99,7 +99,7 @@ class TestOrganization: def test_update_member_role(self, authenticated_http_client, test_org): """Test updating organization member role.""" from kohakuhub.auth.routes import RegisterRequest - from kohakuhub.org.routes import AddMemberPayload, UpdateMemberRolePayload + from kohakuhub.api.org.router import AddMemberPayload, UpdateMemberRolePayload # Create a new user unique_id = uuid.uuid4().hex[:6] @@ -157,7 +157,7 @@ class TestOrganization: def test_duplicate_organization(self, authenticated_http_client, test_org): """Test that creating duplicate organization fails.""" - from kohakuhub.org.routes import CreateOrganizationPayload + from kohakuhub.api.org.router import CreateOrganizationPayload # Try to create organization with same name payload = CreateOrganizationPayload(name=test_org, description="Duplicate org") diff --git a/tests/test_repo_crud.py b/tests/test_repo_crud.py index 5a674a1..ca6a7f7 100644 --- a/tests/test_repo_crud.py +++ b/tests/test_repo_crud.py @@ -38,7 +38,7 @@ class TestRepositoryCRUD: def test_create_repo_http_client(self, authenticated_http_client): """Test repository creation using custom HTTP client with Pydantic model.""" - from kohakuhub.api.routers.repo_crud import CreateRepoPayload + from kohakuhub.api.repo.routers.crud import CreateRepoPayload unique_id = uuid.uuid4().hex[:6] repo_name = f"htc-{unique_id}" # htc = http-create @@ -65,7 +65,7 @@ class TestRepositoryCRUD: assert resp.status_code == 200, f"Get repo info failed: {resp.text}" # Cleanup using actual delete model - from kohakuhub.api.routers.repo_crud import DeleteRepoPayload + from kohakuhub.api.repo.routers.crud import DeleteRepoPayload delete_payload = DeleteRepoPayload( type="model", name=repo_name, organization=None @@ -189,7 +189,7 @@ class TestRepositoryCRUD: def test_move_repo(self, authenticated_http_client, hf_client): """Test moving/renaming repository.""" - from kohakuhub.api.routers.repo_crud import MoveRepoPayload + from kohakuhub.api.repo.routers.crud import MoveRepoPayload unique_id = uuid.uuid4().hex[:6] old_name = f"old-{unique_id}" @@ -244,7 +244,7 @@ class TestRepositoryCRUD: def test_create_org_repo(self, authenticated_http_client, hf_client, test_org): """Test creating repository under organization.""" - from kohakuhub.api.routers.repo_crud import CreateRepoPayload + from kohakuhub.api.repo.routers.crud import CreateRepoPayload unique_id = uuid.uuid4().hex[:6] repo_name = f"org-{unique_id}"