Spaces:
Build error
Build error
| """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("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("presence_state", 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("profile_banner_image_url") | |
| batch_op.drop_column("timezone") | |
| batch_op.drop_column("presence_state") | |
| 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") | |