refac

backend/open_webui/migrations/versions/38d63c18f30f_add_oauth_session_table.py

@@ -19,50 +19,63 @@ depends_on: Union[str, Sequence[str], None] = None
 
 
 def upgrade() -> None:
-    # Ensure 'id' column in 'user' table is unique and primary key (ForeignKey constraint)
     inspector = sa.inspect(op.get_bind())
-    columns = inspector.get_columns('user')
+    existing_tables = set(inspector.get_table_names())
 
+    # ── Fix user.id PK only when it is genuinely wrong ────────────────
+    # On databases migrated from the old Peewee layer the PK may have
+    # drifted.  Fresh installs (init migration 7e5b5dc7342b) already
+    # have a correct PK on 'id', so we must NOT attempt to re-create it
+    # — PostgreSQL will reject a duplicate PRIMARY KEY constraint.
     pk_columns = inspector.get_pk_constraint('user')['constrained_columns']
-    id_column = next((col for col in columns if col['name'] == 'id'), None)
 
-    if id_column and not id_column.get('unique', False):
+    if pk_columns != ['id']:
+        columns = inspector.get_columns('user')
+        id_column = next((col for col in columns if col['name'] == 'id'), None)
         unique_constraints = inspector.get_unique_constraints('user')
         unique_columns = {tuple(u['column_names']) for u in unique_constraints}
 
-        with op.batch_alter_table('user') as batch_op:
-            # If primary key is wrong, drop it
-            if pk_columns and pk_columns != ['id']:
-                batch_op.drop_constraint(inspector.get_pk_constraint('user')['name'], type_='primary')
+        if id_column:
+            with op.batch_alter_table('user') as batch_op:
+                if pk_columns:
+                    batch_op.drop_constraint(
+                        inspector.get_pk_constraint('user')['name'],
+                        type_='primary',
+                    )
+                if ('id',) not in unique_columns:
+                    batch_op.create_unique_constraint('uq_user_id', ['id'])
+                batch_op.create_primary_key('pk_user_id', ['id'])
 
-            # Add unique constraint if missing
-            if ('id',) not in unique_columns:
-                batch_op.create_unique_constraint('uq_user_id', ['id'])
+    # ── Create oauth_session table (idempotent) ───────────────────────
+    if 'oauth_session' not in existing_tables:
+        op.create_table(
+            'oauth_session',
+            sa.Column('id', sa.Text(), primary_key=True, nullable=False, unique=True),
+            sa.Column(
+                'user_id',
+                sa.Text(),
+                sa.ForeignKey('user.id', ondelete='CASCADE'),
+                nullable=False,
+            ),
+            sa.Column('provider', sa.Text(), nullable=False),
+            sa.Column('token', sa.Text(), nullable=False),
+            sa.Column('expires_at', sa.BigInteger(), nullable=False),
+            sa.Column('created_at', sa.BigInteger(), nullable=False),
+            sa.Column('updated_at', sa.BigInteger(), nullable=False),
+        )
 
-            # Re-create correct primary key
-            batch_op.create_primary_key('pk_user_id', ['id'])
+    # Create indexes (idempotent — no-ops when table was just created
+    # with the columns above, and safe to call if indexes already exist).
+    existing_indexes = {
+        idx['name'] for idx in inspector.get_indexes('oauth_session')
+    } if 'oauth_session' in existing_tables else set()
 
-    # Create oauth_session table
-    op.create_table(
-        'oauth_session',
-        sa.Column('id', sa.Text(), primary_key=True, nullable=False, unique=True),
-        sa.Column(
-            'user_id',
-            sa.Text(),
-            sa.ForeignKey('user.id', ondelete='CASCADE'),
-            nullable=False,
-        ),
-        sa.Column('provider', sa.Text(), nullable=False),
-        sa.Column('token', sa.Text(), nullable=False),
-        sa.Column('expires_at', sa.BigInteger(), nullable=False),
-        sa.Column('created_at', sa.BigInteger(), nullable=False),
-        sa.Column('updated_at', sa.BigInteger(), nullable=False),
-    )
-
-    # Create indexes for better performance
-    op.create_index('idx_oauth_session_user_id', 'oauth_session', ['user_id'])
-    op.create_index('idx_oauth_session_expires_at', 'oauth_session', ['expires_at'])
-    op.create_index('idx_oauth_session_user_provider', 'oauth_session', ['user_id', 'provider'])
+    if 'idx_oauth_session_user_id' not in existing_indexes:
+        op.create_index('idx_oauth_session_user_id', 'oauth_session', ['user_id'])
+    if 'idx_oauth_session_expires_at' not in existing_indexes:
+        op.create_index('idx_oauth_session_expires_at', 'oauth_session', ['expires_at'])
+    if 'idx_oauth_session_user_provider' not in existing_indexes:
+        op.create_index('idx_oauth_session_user_provider', 'oauth_session', ['user_id', 'provider'])
 
 
 def downgrade() -> None:

backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py

@@ -20,20 +20,49 @@ down_revision: Union[str, None] = '2f1211949ecc'
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
+# ── Ad-hoc table references for Core DML ─────────────────────────────────
+# These are lightweight table() / column() references used only inside this
+# migration for SELECT / UPDATE / INSERT — they do NOT create or alter
+# anything on disk.
+
+_user = sa.table(
+    'user',
+    sa.column('id', sa.Text),
+    sa.column('oauth_sub', sa.Text),
+    sa.column('oauth', sa.JSON),
+    sa.column('api_key', sa.Text),
+    sa.column('info', sa.Text),
+    sa.column('settings', sa.Text),
+)
+
+_api_key = sa.table(
+    'api_key',
+    sa.column('id', sa.Text),
+    sa.column('user_id', sa.Text),
+    sa.column('key', sa.Text),
+    sa.column('created_at', sa.BigInteger),
+    sa.column('updated_at', sa.BigInteger),
+)
+
 
 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.
+    SQLite requires manual removal of any user-created indexes referencing
+    a column before ALTER TABLE ... DROP COLUMN can succeed.
+
+    NOTE: PRAGMAs have no Core equivalent — raw text is unavoidable here.
     """
     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
+        # Skip system-managed autoindexes (PK / UNIQUE constraints) — they
+        # cannot be dropped directly and will disappear when the column is
+        # removed via batch_alter_table.
+        if index_name.startswith('sqlite_autoindex_'):
+            continue
         idx_info = conn.execute(sa.text(f"PRAGMA index_info('{index_name}')")).fetchall()
-
-        indexed_cols = [row[2] for row in idx_info]  # col names
+        indexed_cols = [row[2] for row in idx_info]
         if column_name in indexed_cols:
             conn.execute(sa.text(f'DROP INDEX IF EXISTS {index_name}'))
 
@@ -42,33 +71,31 @@ def _convert_column_to_json(table: str, column: str):
     conn = op.get_bind()
     dialect = conn.dialect.name
 
+    t = sa.table(table, sa.column('id', sa.Text), sa.column(column, sa.Text))
+    t_json = sa.column(f'{column}_json', sa.JSON)
+
     # 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()
+        rows = conn.execute(sa.select(t.c.id, t.c[column])).fetchall()
 
-        for row in rows:
-            uid, raw = row
+        for uid, raw in rows:
             if raw is None:
                 parsed = None
             else:
                 try:
                     parsed = json.loads(raw)
                 except Exception:
-                    parsed = None  # fallback safe behavior
+                    parsed = None
 
             conn.execute(
-                sa.text(f'UPDATE "{table}" SET {column}_json = :val WHERE id = :id'),
-                {'val': json.dumps(parsed) if parsed else None, 'id': uid},
+                sa.update(sa.table(table, sa.column('id'), t_json))
+                .where(sa.column('id') == uid)
+                .values({f'{column}_json': json.dumps(parsed) if parsed else None})
             )
 
-        # 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:
@@ -85,15 +112,19 @@ def _convert_column_to_text(table: str, column: str):
     conn = op.get_bind()
     dialect = conn.dialect.name
 
+    t = sa.table(table, sa.column('id', sa.Text), sa.column(column))
+    t_text = sa.column(f'{column}_text', sa.Text)
+
     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()
+        rows = conn.execute(sa.select(t.c.id, t.c[column])).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},
+                sa.update(sa.table(table, sa.column('id'), t_text))
+                .where(sa.column('id') == uid)
+                .values({f'{column}_text': json.dumps(raw) if raw else None})
             )
 
         op.drop_column(table, column)
@@ -109,88 +140,102 @@ def _convert_column_to_text(table: str, column: str):
 
 
 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))
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
+    existing_tables = set(inspector.get_table_names())
+    user_columns = {c['name'] for c in inspector.get_columns('user')}
 
-    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))
+    # ── Add new columns (idempotent) ──────────────────────────────────
+    for col_name, col_type in [
+        ('profile_banner_image_url', sa.Text()),
+        ('timezone', sa.String()),
+        ('presence_state', sa.String()),
+        ('status_emoji', sa.String()),
+        ('status_message', sa.Text()),
+        ('status_expires_at', sa.BigInteger()),
+        ('oauth', sa.JSON()),
+    ]:
+        if col_name not in user_columns:
+            op.add_column('user', sa.Column(col_name, col_type, 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),
-    )
+    # ── Create api_key table (idempotent) ─────────────────────────────
+    if 'api_key' not in existing_tables:
+        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()
+    # ── Migrate oauth_sub → oauth JSON (only if old column still exists)
+    if 'oauth_sub' in user_columns:
+        rows = conn.execute(
+            sa.select(_user.c.id, _user.c.oauth_sub)
+            .where(_user.c.oauth_sub.is_not(None))
+        ).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
+        for uid, oauth_sub in rows:
+            if oauth_sub:
+                provider, sub = (oauth_sub.split('@', 1) if '@' in oauth_sub
+                                 else ('oidc', oauth_sub))
+                conn.execute(
+                    sa.update(_user)
+                    .where(_user.c.id == uid)
+                    .values(oauth=json.dumps({provider: {'sub': 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},
-            )
+    # ── Migrate api_key column → api_key table (only if old column still exists)
+    if 'api_key' in user_columns:
+        rows = conn.execute(
+            sa.select(_user.c.id, _user.c.api_key)
+            .where(_user.c.api_key.is_not(None))
+        ).fetchall()
+        now = int(time.time())
 
-    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, key_val in rows:
+            if key_val:
+                conn.execute(
+                    sa.insert(_api_key).values(
+                        id=f'key_{uid}',
+                        user_id=uid,
+                        key=key_val,
+                        created_at=now,
+                        updated_at=now,
+                    )
+                )
 
-    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,
-                },
-            )
+    # ── Drop legacy columns (idempotent) ──────────────────────────────
+    cols_to_drop = {'api_key', 'oauth_sub'} & user_columns
+    if cols_to_drop:
+        if conn.dialect.name == 'sqlite':
+            for col in cols_to_drop:
+                _drop_sqlite_indexes_for_column('user', col, conn)
 
-    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')
+        with op.batch_alter_table('user') as batch_op:
+            for col in cols_to_drop:
+                batch_op.drop_column(col)
 
 
 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()
+    rows = conn.execute(
+        sa.select(_user.c.id, _user.c.oauth)
+        .where(_user.c.oauth.is_not(None))
+    ).fetchall()
 
-    for uid, oauth in users:
+    for uid, oauth in rows:
         try:
             data = json.loads(oauth)
             provider = list(data.keys())[0]
@@ -200,24 +245,24 @@ def downgrade() -> None:
             oauth_sub = None
 
         conn.execute(
-            sa.text('UPDATE "user" SET oauth_sub = :oauth_sub WHERE id = :id'),
-            {'oauth_sub': oauth_sub, 'id': uid},
+            sa.update(_user)
+            .where(_user.c.id == uid)
+            .values(oauth_sub=oauth_sub)
         )
 
     op.drop_column('user', 'oauth')
 
-    # --- 2. Restore api_key field ---
+    # --- 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()
+    keys = conn.execute(sa.select(_api_key.c.user_id, _api_key.c.key)).fetchall()
     for uid, key in keys:
         conn.execute(
-            sa.text('UPDATE "user" SET api_key = :key WHERE id = :id'),
-            {'key': key, 'id': uid},
+            sa.update(_user)
+            .where(_user.c.id == uid)
+            .values(api_key=key)
         )
 
-    # Drop new table
     op.drop_table('api_key')
 
     with op.batch_alter_table('user') as batch_op: