代码之家  ›  专栏  ›  技术社区  ›  Max

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

  •  17
  • Max  · 技术社区  · 10 年前

    我正在使用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,因此我必须删除所有用户,升级User表,然后再次添加用户,这很痛苦。因此,我想通过以下方式更改上述脚本:

    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并在“@”sign(在本例中为“@example.com”)之后抛出所有内容,然后设置username的值(在本示例中为“test”),使nullable=false。

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

    或者如果有其他设置方式 username 默认值为电子邮件地址,不包含@ssing及其后的所有内容?

    2 回复  |  直到 10 年前
        1
  •  7
  •   Max    10 年前

    这是如何解决这个问题的。

    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
  •  0
  •   Tim    6 年前

    可以使用子查询编写基于同一表的另一列的值更新新列的脚本。唯一的诀窍是,因为您从同一个表中查询,所以需要为表名赋予别名,以确保您从相应的行中进行选择:

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