Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac: async db

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-12 14:22:11 -05:00
parent b618d84065
commit 27169124f2
74 changed files with 4831 additions and 4479 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/access_grants.py
+101 -89
View File
@@ -3,8 +3,9 @@ import time
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db_context
from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from open_webui.internal.db import Base, get_async_db_context
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, Text, UniqueConstraint, or_, and_
@@ -281,20 +282,20 @@ def grants_to_access_control(grants: list) -> Optional[dict]:
class AccessGrantsTable:
def grant_access(
async def grant_access(
self,
resource_type: str,
resource_id: str,
principal_type: str,
principal_id: str,
permission: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[AccessGrantModel]:
"""Add a single access grant. Idempotent (ignores duplicates)."""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Check for existing grant
existing = (
db.query(AccessGrant)
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
@@ -302,8 +303,8 @@ class AccessGrantsTable:
principal_id=principal_id,
permission=permission,
)
.first()
)
existing = result.scalars().first()
if existing:
return AccessGrantModel.model_validate(existing)
@@ -317,23 +318,23 @@ class AccessGrantsTable:
created_at=int(time.time()),
)
db.add(grant)
db.commit()
db.refresh(grant)
await db.commit()
await db.refresh(grant)
return AccessGrantModel.model_validate(grant)
def revoke_access(
async def revoke_access(
self,
resource_type: str,
resource_id: str,
principal_type: str,
principal_id: str,
permission: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""Remove a single access grant."""
with get_db_context(db) as db:
deleted = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
delete(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
@@ -341,47 +342,47 @@ class AccessGrantsTable:
principal_id=principal_id,
permission=permission,
)
.delete()
)
db.commit()
return deleted > 0
await db.commit()
return result.rowcount > 0
def revoke_all_access(
async def revoke_all_access(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> int:
"""Remove all access grants for a resource."""
with get_db_context(db) as db:
deleted = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
delete(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.delete()
)
db.commit()
return deleted
await db.commit()
return result.rowcount
def set_access_control(
async def set_access_control(
self,
resource_type: str,
resource_id: str,
access_control: Optional[dict],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""
Replace all grants for a resource from an access_control JSON dict.
This is the primary bridge for backward compat with the frontend.
"""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Delete all existing grants for this resource
db.query(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
).delete()
await db.execute(
delete(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
)
# Convert JSON to grant dicts
grant_dicts = access_control_to_grants(resource_type, resource_id, access_control)
@@ -397,25 +398,27 @@ class AccessGrantsTable:
db.add(grant)
results.append(grant)
db.commit()
await db.commit()
return [AccessGrantModel.model_validate(g) for g in results]
def set_access_grants(
async def set_access_grants(
self,
resource_type: str,
resource_id: str,
access_grants: Optional[list],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""
Replace all grants for a resource from a direct access_grants list.
"""
with get_db_context(db) as db:
db.query(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
).delete()
async with get_async_db_context(db) as db:
await db.execute(
delete(AccessGrant).filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
)
normalized_grants = normalize_access_grants(access_grants)
@@ -433,80 +436,80 @@ class AccessGrantsTable:
db.add(grant)
results.append(grant)
db.commit()
await db.commit()
return [AccessGrantModel.model_validate(g) for g in results]
def get_access_control(
async def get_access_control(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> Optional[dict]:
"""
Reconstruct the old-style access_control JSON dict from grants.
For backward compat with the frontend.
"""
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.all()
)
grants = result.scalars().all()
grant_models = [AccessGrantModel.model_validate(g) for g in grants]
return grants_to_access_control(grant_models)
def get_grants_by_resource(
async def get_grants_by_resource(
self,
resource_type: str,
resource_id: str,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list[AccessGrantModel]:
"""Get all grants for a specific resource."""
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
)
.all()
)
grants = result.scalars().all()
return [AccessGrantModel.model_validate(g) for g in grants]
def get_grants_by_resources(
async def get_grants_by_resources(
self,
resource_type: str,
resource_ids: list[str],
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> dict[str, list[AccessGrantModel]]:
"""Batch-fetch grants for multiple resources. Returns {resource_id: [grants]}."""
if not resource_ids:
return {}
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id.in_(resource_ids),
)
.all()
)
result: dict[str, list[AccessGrantModel]] = {rid: [] for rid in resource_ids}
grants = result.scalars().all()
result_dict: dict[str, list[AccessGrantModel]] = {rid: [] for rid in resource_ids}
for g in grants:
result[g.resource_id].append(AccessGrantModel.model_validate(g))
return result
result_dict[g.resource_id].append(AccessGrantModel.model_validate(g))
return result_dict
def has_access(
async def has_access(
self,
user_id: str,
resource_type: str,
resource_id: str,
permission: str = 'read',
user_group_ids: Optional[set[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> bool:
"""
Check if a user has the specified permission on a resource.
@@ -516,7 +519,7 @@ class AccessGrantsTable:
- There's a grant for the specific user with the requested permission
- There's a grant for any of the user's groups with the requested permission
"""
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
# Build conditions for matching grants
conditions = [
# Public access
@@ -535,7 +538,7 @@ class AccessGrantsTable:
if user_group_ids is None:
from open_webui.models.groups import Groups
user_groups = Groups.get_groups_by_member_id(user_id, db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
if user_group_ids:
@@ -546,26 +549,27 @@ class AccessGrantsTable:
)
)
exists = (
db.query(AccessGrant)
result = await db.execute(
select(AccessGrant)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id == resource_id,
AccessGrant.permission == permission,
or_(*conditions),
)
.first()
.limit(1)
)
return exists is not None
grant = result.scalars().first()
return grant is not None
def get_accessible_resource_ids(
async def get_accessible_resource_ids(
self,
user_id: str,
resource_type: str,
resource_ids: list[str],
permission: str = 'read',
user_group_ids: Optional[set[str]] = None,
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> set[str]:
"""
Batch check: return the subset of resource_ids that the user can access.
@@ -575,7 +579,7 @@ class AccessGrantsTable:
if not resource_ids:
return set()
with get_db_context(db) as db:
async with get_async_db_context(db) as db:
conditions = [
and_(
AccessGrant.principal_type == 'user',
@@ -590,7 +594,7 @@ class AccessGrantsTable:
if user_group_ids is None:
from open_webui.models.groups import Groups
user_groups = Groups.get_groups_by_member_id(user_id, db=db)
user_groups = await Groups.get_groups_by_member_id(user_id, db=db)
user_group_ids = {group.id for group in user_groups}
if user_group_ids:
@@ -601,8 +605,8 @@ class AccessGrantsTable:
)
)
rows = (
db.query(AccessGrant.resource_id)
result = await db.execute(
select(AccessGrant.resource_id)
.filter(
AccessGrant.resource_type == resource_type,
AccessGrant.resource_id.in_(resource_ids),
@@ -610,16 +614,16 @@ class AccessGrantsTable:
or_(*conditions),
)
.distinct()
.all()
)
rows = result.all()
return {row[0] for row in rows}
def get_users_with_access(
async def get_users_with_access(
self,
resource_type: str,
resource_id: str,
permission: str = 'read',
db: Optional[Session] = None,
db: Optional[AsyncSession] = None,
) -> list:
"""
Get all users who have the specified permission on a resource.
@@ -628,21 +632,21 @@ class AccessGrantsTable:
from open_webui.models.users import Users, UserModel
from open_webui.models.groups import Groups
with get_db_context(db) as db:
grants = (
db.query(AccessGrant)
async with get_async_db_context(db) as db:
result = await db.execute(
select(AccessGrant)
.filter_by(
resource_type=resource_type,
resource_id=resource_id,
permission=permission,
)
.all()
)
grants = result.scalars().all()
# Check for public access
for grant in grants:
if grant.principal_type == 'user' and grant.principal_id == '*':
result = Users.get_users(filter={'roles': ['!pending']}, db=db)
result = await Users.get_users(filter={'roles': ['!pending']}, db=db)
return result.get('users', [])
user_ids_with_access = set()
@@ -651,14 +655,14 @@ class AccessGrantsTable:
if grant.principal_type == 'user':
user_ids_with_access.add(grant.principal_id)
elif grant.principal_type == 'group':
group_user_ids = Groups.get_group_user_ids_by_id(grant.principal_id, db=db)
group_user_ids = await Groups.get_group_user_ids_by_id(grant.principal_id, db=db)
if group_user_ids:
user_ids_with_access.update(group_user_ids)
if not user_ids_with_access:
return []
return Users.get_users_by_user_ids(list(user_ids_with_access), db=db)
return await Users.get_users_by_user_ids(list(user_ids_with_access), db=db)
def has_permission_filter(
self,
@@ -673,6 +677,10 @@ class AccessGrantsTable:
Apply access control filtering to a SQLAlchemy query by JOINing with access_grant.
This replaces the old JSON-column-based filtering with a proper relational JOIN.
Note: This method builds SQLAlchemy expressions and does NOT perform I/O itself,
so it remains synchronous. The caller is responsible for executing the query
asynchronously with `await db.execute(...)`.
"""
group_ids = filter.get('group_ids', [])
user_id = filter.get('user_id')
@@ -718,7 +726,7 @@ class AccessGrantsTable:
# LEFT JOIN access_grant and filter
# We use a subquery approach to avoid duplicates from multiple matching grants
from sqlalchemy import exists as sa_exists, select
from sqlalchemy import exists as sa_exists
grant_exists = (
select(AccessGrant.id)
@@ -776,11 +784,15 @@ class AccessGrantsTable:
"""
Filter for items where user has read BUT NOT write access.
Public items are NOT considered read_only.
Note: This method builds SQLAlchemy expressions and does NOT perform I/O itself,
so it remains synchronous. The caller is responsible for executing the query
asynchronously with `await db.execute(...)`.
"""
group_ids = filter.get('group_ids', [])
user_id = filter.get('user_id')
from sqlalchemy import exists as sa_exists, select
from sqlalchemy import exists as sa_exists
# Has read grant (not public)
read_grant_exists = (