From 369298a83e563d8bc77adeb83a5046ead28bc3a0 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Fri, 28 Nov 2025 06:29:41 -0500 Subject: [PATCH] refac: user table db migration --- .../b10670c03dd5_update_user_table.py | 264 ++++++++++++++++++ backend/open_webui/models/users.py | 114 +++++--- 2 files changed, 345 insertions(+), 33 deletions(-) create mode 100644 backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py diff --git a/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py b/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py new file mode 100644 index 0000000000..ce7d3d9870 --- /dev/null +++ b/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py @@ -0,0 +1,264 @@ +"""Update user table + +Revision ID: b10670c03dd5 +Revises: 2f1211949ecc +Create Date: 2025-11-28 04:55:31.737538 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +import open_webui.internal.db +import json +import time + +# revision identifiers, used by Alembic. +revision: str = "b10670c03dd5" +down_revision: Union[str, None] = "2f1211949ecc" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def _drop_sqlite_indexes_for_column(table_name, column_name, conn): + """ + SQLite requires manual removal of any indexes referencing a column + before ALTER TABLE ... DROP COLUMN can succeed. + """ + indexes = conn.execute(sa.text(f"PRAGMA index_list('{table_name}')")).fetchall() + + for idx in indexes: + index_name = idx[1] # index name + # Get indexed columns + idx_info = conn.execute( + sa.text(f"PRAGMA index_info('{index_name}')") + ).fetchall() + + indexed_cols = [row[2] for row in idx_info] # col names + if column_name in indexed_cols: + conn.execute(sa.text(f"DROP INDEX IF EXISTS {index_name}")) + + +def _convert_column_to_json(table: str, column: str): + conn = op.get_bind() + dialect = conn.dialect.name + + # SQLite cannot ALTER COLUMN → must recreate column + if dialect == "sqlite": + # 1. Add temporary column + op.add_column(table, sa.Column(f"{column}_json", sa.JSON(), nullable=True)) + + # 2. Load old data + rows = conn.execute(sa.text(f'SELECT id, {column} FROM "{table}"')).fetchall() + + for row in rows: + uid, raw = row + if raw is None: + parsed = None + else: + try: + parsed = json.loads(raw) + except Exception: + parsed = None # fallback safe behavior + + conn.execute( + sa.text(f'UPDATE "{table}" SET {column}_json = :val WHERE id = :id'), + {"val": json.dumps(parsed) if parsed else None, "id": uid}, + ) + + # 3. Drop old TEXT column + op.drop_column(table, column) + + # 4. Rename new JSON column → original name + op.alter_column(table, f"{column}_json", new_column_name=column) + + else: + # PostgreSQL supports direct CAST + op.alter_column( + table, + column, + type_=sa.JSON(), + postgresql_using=f"{column}::json", + ) + + +def _convert_column_to_text(table: str, column: str): + conn = op.get_bind() + dialect = conn.dialect.name + + if dialect == "sqlite": + op.add_column(table, sa.Column(f"{column}_text", sa.Text(), nullable=True)) + + rows = conn.execute(sa.text(f'SELECT id, {column} FROM "{table}"')).fetchall() + + for uid, raw in rows: + conn.execute( + sa.text(f'UPDATE "{table}" SET {column}_text = :val WHERE id = :id'), + {"val": json.dumps(raw) if raw else None, "id": uid}, + ) + + op.drop_column(table, column) + op.alter_column(table, f"{column}_text", new_column_name=column) + + else: + op.alter_column( + table, + column, + type_=sa.Text(), + postgresql_using=f"to_json({column})::text", + ) + + +def upgrade() -> None: + + op.add_column( + "user", + sa.Column( + "is_active", + sa.Boolean(), + nullable=False, + default=False, + server_default=sa.sql.expression.false(), + ), + ) + + op.add_column( + "user", sa.Column("profile_banner_image_url", sa.Text(), nullable=True) + ) + + op.add_column("user", sa.Column("timezone", sa.String(), nullable=True)) + + op.add_column("user", sa.Column("status_emoji", sa.String(), nullable=True)) + op.add_column("user", sa.Column("status_message", sa.Text(), nullable=True)) + op.add_column( + "user", sa.Column("status_expires_at", sa.BigInteger(), nullable=True) + ) + + op.add_column("user", sa.Column("oauth", sa.JSON(), nullable=True)) + + # Convert info (TEXT/JSONField) → JSON + _convert_column_to_json("user", "info") + # Convert settings (TEXT/JSONField) → JSON + _convert_column_to_json("user", "settings") + + op.create_table( + "api_key", + sa.Column("id", sa.Text(), primary_key=True, unique=True), + sa.Column("user_id", sa.Text(), sa.ForeignKey("user.id", ondelete="CASCADE")), + sa.Column("key", sa.Text(), unique=True, nullable=False), + sa.Column("data", sa.JSON(), nullable=True), + sa.Column("expires_at", sa.BigInteger(), nullable=True), + sa.Column("last_used_at", sa.BigInteger(), nullable=True), + sa.Column("created_at", sa.BigInteger(), nullable=False), + sa.Column("updated_at", sa.BigInteger(), nullable=False), + ) + + conn = op.get_bind() + users = conn.execute( + sa.text('SELECT id, oauth_sub FROM "user" WHERE oauth_sub IS NOT NULL') + ).fetchall() + + for uid, oauth_sub in users: + if oauth_sub: + # Example formats supported: + # provider@sub + # plain sub (stored as {"oidc": {"sub": sub}}) + if "@" in oauth_sub: + provider, sub = oauth_sub.split("@", 1) + else: + provider, sub = "oidc", oauth_sub + + oauth_json = json.dumps({provider: {"sub": sub}}) + conn.execute( + sa.text('UPDATE "user" SET oauth = :oauth WHERE id = :id'), + {"oauth": oauth_json, "id": uid}, + ) + + users_with_keys = conn.execute( + sa.text('SELECT id, api_key FROM "user" WHERE api_key IS NOT NULL') + ).fetchall() + now = int(time.time()) + + for uid, api_key in users_with_keys: + if api_key: + conn.execute( + sa.text( + """ + INSERT INTO api_key (id, user_id, key, created_at, updated_at) + VALUES (:id, :user_id, :key, :created_at, :updated_at) + """ + ), + { + "id": f"key_{uid}", + "user_id": uid, + "key": api_key, + "created_at": now, + "updated_at": now, + }, + ) + + if conn.dialect.name == "sqlite": + _drop_sqlite_indexes_for_column("user", "api_key", conn) + _drop_sqlite_indexes_for_column("user", "oauth_sub", conn) + + with op.batch_alter_table("user") as batch_op: + batch_op.drop_column("api_key") + batch_op.drop_column("oauth_sub") + + +def downgrade() -> None: + # --- 1. Restore old oauth_sub column --- + op.add_column("user", sa.Column("oauth_sub", sa.Text(), nullable=True)) + + conn = op.get_bind() + users = conn.execute( + sa.text('SELECT id, oauth FROM "user" WHERE oauth IS NOT NULL') + ).fetchall() + + for uid, oauth in users: + try: + data = json.loads(oauth) + provider = list(data.keys())[0] + sub = data[provider].get("sub") + oauth_sub = f"{provider}@{sub}" + except Exception: + oauth_sub = None + + conn.execute( + sa.text('UPDATE "user" SET oauth_sub = :oauth_sub WHERE id = :id'), + {"oauth_sub": oauth_sub, "id": uid}, + ) + + op.drop_column("user", "oauth") + + # --- 2. Restore api_key field --- + op.add_column("user", sa.Column("api_key", sa.String(), nullable=True)) + + # Restore values from api_key + keys = conn.execute(sa.text("SELECT user_id, key FROM api_key")).fetchall() + for uid, key in keys: + conn.execute( + sa.text('UPDATE "user" SET api_key = :key WHERE id = :id'), + {"key": key, "id": uid}, + ) + + # Drop new table + op.drop_table("api_key") + + with op.batch_alter_table("user") as batch_op: + batch_op.drop_column("is_active") + + batch_op.drop_column("profile_banner_image_url") + batch_op.drop_column("timezone") + + batch_op.drop_column("status_emoji") + batch_op.drop_column("status_message") + batch_op.drop_column("status_expires_at") + + # Convert info (JSON) → TEXT + _convert_column_to_text("user", "info") + # Convert settings (JSON) → TEXT + _convert_column_to_text("user", "settings") diff --git a/backend/open_webui/models/users.py b/backend/open_webui/models/users.py index e7beeee1bf..b68b195cd2 100644 --- a/backend/open_webui/models/users.py +++ b/backend/open_webui/models/users.py @@ -11,7 +11,17 @@ from open_webui.utils.misc import throttle from pydantic import BaseModel, ConfigDict -from sqlalchemy import BigInteger, Column, String, Text, Date, exists, select +from sqlalchemy import ( + BigInteger, + JSON, + Column, + String, + Boolean, + Text, + Date, + exists, + select, +) from sqlalchemy import or_, case import datetime @@ -21,59 +31,71 @@ import datetime #################### -class User(Base): - __tablename__ = "user" - - id = Column(String, primary_key=True, unique=True) - name = Column(String) - - email = Column(String) - username = Column(String(50), nullable=True) - - role = Column(String) - profile_image_url = Column(Text) - - bio = Column(Text, nullable=True) - gender = Column(Text, nullable=True) - date_of_birth = Column(Date, nullable=True) - - info = Column(JSONField, nullable=True) - settings = Column(JSONField, nullable=True) - - api_key = Column(String, nullable=True, unique=True) - oauth_sub = Column(Text, unique=True) - - last_active_at = Column(BigInteger) - - updated_at = Column(BigInteger) - created_at = Column(BigInteger) - - class UserSettings(BaseModel): ui: Optional[dict] = {} model_config = ConfigDict(extra="allow") pass +class User(Base): + __tablename__ = "user" + + id = Column(String, primary_key=True, unique=True) + email = Column(String) + username = Column(String(50), nullable=True) + role = Column(String) + + name = Column(String) + is_active = Column(Boolean, nullable=False, default=False) + + profile_image_url = Column(Text) + profile_banner_image_url = Column(Text, nullable=True) + + bio = Column(Text, nullable=True) + gender = Column(Text, nullable=True) + date_of_birth = Column(Date, nullable=True) + timezone = Column(String, nullable=True) + + status_emoji = Column(String, nullable=True) + status_message = Column(Text, nullable=True) + status_expires_at = Column(BigInteger, nullable=True) + + info = Column(JSON, nullable=True) + settings = Column(JSON, nullable=True) + + oauth = Column(JSON, nullable=True) + + last_active_at = Column(BigInteger) + updated_at = Column(BigInteger) + created_at = Column(BigInteger) + + class UserModel(BaseModel): id: str - name: str email: str username: Optional[str] = None - role: str = "pending" + + name: str + is_active: bool = False + profile_image_url: str + profile_banner_image_url: Optional[str] = None bio: Optional[str] = None gender: Optional[str] = None date_of_birth: Optional[datetime.date] = None + timezone: Optional[str] = None + + status_emoji: Optional[str] = None + status_message: Optional[str] = None + status_expires_at: Optional[int] = None info: Optional[dict] = None settings: Optional[UserSettings] = None - api_key: Optional[str] = None - oauth_sub: Optional[str] = None + oauth: Optional[dict] = None last_active_at: int # timestamp in epoch updated_at: int # timestamp in epoch @@ -82,6 +104,32 @@ class UserModel(BaseModel): model_config = ConfigDict(from_attributes=True) +class ApiKey(Base): + __tablename__ = "api_key" + + id = Column(Text, primary_key=True, unique=True) + user_id = Column(Text, nullable=False) + key = Column(Text, unique=True, nullable=False) + data = Column(JSON, nullable=True) + expires_at = Column(BigInteger, nullable=True) + last_used_at = Column(BigInteger, nullable=True) + created_at = Column(BigInteger, nullable=False) + updated_at = Column(BigInteger, nullable=False) + + +class ApiKeyModel(BaseModel): + id: str + user_id: str + key: str + data: Optional[dict] = None + expires_at: Optional[int] = None + last_used_at: Optional[int] = None + created_at: int # timestamp in epoch + updated_at: int # timestamp in epoch + + model_config = ConfigDict(from_attributes=True) + + #################### # Forms ####################