使用 alembic 升级表时设置列的值答案

【问题标题】：Set the value of column when upgrading table with alembic使用 alembic 升级表时设置列的值
【发布时间】：2014-10-13 17:48:28
【问题描述】：

我正在使用 PostgreSQL 和 Alembic 进行迁移。当我向我的用户表添加新列时，Alembic 使用以下脚本生成了迁移：

revision = '4824acf75bf3'
down_revision = '2f0fbdd56de1'

from alembic import op
import sqlalchemy as sa

def upgrade():
    op.add_column(
        'user', 
        sa.Column(
            'username', 
            sa.Unicode(length=255), 
            nullable=False
        )
    )

def downgrade():
    op.drop_column('user', 'username')

我真正想做的是在升级生产版本时自动生成用户名的值。换句话说，我的生产版本有很多用户，如果我在上面运行上面的升级，会出现一个错误，指出用户名不能为 NULL，所以我必须删除所有用户，升级用户表和再次添加用户后，这很痛苦。因此，我想将上述脚本更改为：

revision = '4824acf75bf3'
down_revision = '2f0fbdd56de1'

from alembic import op
import sqlalchemy as sa

def upgrade():
    op.add_column(
        'user', 
        sa.Column(
            'username', 
            sa.Unicode(length=255)
        )
    )
    op.execute(
        'UPDATE "user" set username = <email address with no '@' 
         and everything comes after '@' sign should be removed> 
         WHERE email is not null'
    )
    <only after the above code is executed 'nullable=False' must be set up>

def downgrade():
    op.drop_column('user', 'username')

正如上面代码中所述，我想执行一个 SQL 代码来检查电子邮件地址，如 test@example.com，并在“@”符号（在本例中为“@example.com”）之后抛出所有内容，并且在使 nullable=false 之后设置用户名的值（在本例中为“test”）。

我该怎么做？必须是什么脚本而不是username = <email address with no '@' and everything comes after '@' sign should be removed> 和设置nullable=false

或者如果有任何其他方式将username 默认值设置为不带@sing 的电子邮件地址以及它之后的所有内容？

【问题讨论】：

感谢您提出这个问题；希望 alembic 指令更清楚地表明 op.execute 是解决此问题的方法。
或许能帮上忙：medium.com/the-andela-way/…

标签： database postgresql upgrade auto-generate alembic

【解决方案1】：

这是解决问题的方法。

def upgrade():
    op.add_column(
        'user',
        sa.Column(
            'username',
            sa.Unicode(length=255)
        )
    )
    op.create_index('ix_user_username', 'user', ['username'], unique=True)
    op.execute(
        '''
        DO
        $do$
        DECLARE uid INTEGER;
        DECLARE username_candidate TEXT;
        BEGIN
        FOR uid, username_candidate IN (
            SELECT
                id,
                lower(
                    substring(email for position('@' in email) - 1)
                )
            FROM "user" WHERE username is null
        ) LOOP
            UPDATE "user"
            SET username = username_candidate
            WHERE
                id = uid AND
                NOT EXISTS (
                SELECT id FROM "user" WHERE username = username_candidate
            );
        END LOOP;
        END
        $do$
        '''
    )
    # Fix name colissions
    op.execute(
        '''
        DO
        $do$
        DECLARE uniqufier INTEGER := 0;
        DECLARE uid INTEGER;
        DECLARE username_candidate TEXT;
        BEGIN
        WHILE EXISTS (SELECT id FROM "user" WHERE username is null) LOOP
            uniqufier := uniqufier + 1;
            FOR uid, username_candidate IN (
                SELECT
                    id,
                    lower(
                        substring(email for position('@' in email) - 1)
                        || uniqufier
                    )
                FROM "user" WHERE username is null
            ) LOOP
                UPDATE "user"
                SET username = username_candidate
                WHERE
                    id = uid AND
                    NOT EXISTS (
                        SELECT id FROM "user" WHERE username = username_candidate
                    );
            END LOOP;
        END LOOP;
        END;
        $do$
        '''
    )
    op.alter_column(
        'user',
        'username',
        nullable=False,
    )


def downgrade():
    op.drop_index('ix_user_username', table_name='user')
    op.drop_column('user', 'username')

【讨论】：

【解决方案2】：

可以使用子查询编写根据同一表的另一列的值更新新列的脚本。唯一的技巧是，因为您从同一个表中查询，您需要为表名提供别名以确保您从相应的行中进行选择：

update 'user' as target set username = (
  select substring(email from '.+?(?=@)')
  from 'user' as source where source.id = target.id
);

【讨论】：