如何在Alembic升级脚本中执行插入和更新?


94

我需要在Alembic升级期间更改数据。

我目前在第一个修订版中有一个“玩家”表:

def upgrade():
    op.create_table('player',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.Unicode(length=200), nullable=False),
        sa.Column('position', sa.Unicode(length=200), nullable=True),
        sa.Column('team', sa.Unicode(length=100), nullable=True)
        sa.PrimaryKeyConstraint('id')
    )

我想介绍一个“团队”表。我创建了第二个修订版:

def upgrade():
    op.create_table('teams',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=80), nullable=False)
    )
    op.add_column('players', sa.Column('team_id', sa.Integer(), nullable=False))

我希望第二次迁移也添加以下数据:

  1. 填充团队表:

    INSERT INTO teams (name) SELECT DISTINCT team FROM players;
  2. 根据players.team名称更新players.team_id:

    UPDATE players AS p JOIN teams AS t SET p.team_id = t.id WHERE p.team = t.name;

如何在升级脚本中执行插入和更新?

Answers:


145

您需要的是数据迁移,而不是Alembic文档中最普遍的模式迁移

该答案假设您使用声明式(而不是class-Mapper-Table或core)定义模型。使它适应其他形式应该相对简单。

请注意,Alembic提供了一些基本数据功能:op.bulk_insert()op.execute()。如果操作很少,请使用这些操作。如果迁移需要关系或其他复杂的交互作用,我更喜欢使用模型和会话的全部功能,如下所述。

以下是示例迁移脚本,该脚本设置了一些声明性模型,这些声明性模型将用于处理会话中的数据。关键点是:

  1. 使用所需的列定义所需的基本模型。您不需要每一列,只需要主键和将要使用的那些。

  2. 在升级功能中,用于op.get_bind()获取当前连接并与其建立会话。

    • 或者使用bind.execute()SQLAlchemy的较低级别直接编写SQL查询。这对于简单的迁移很有用。
  3. 像通常在应用程序中一样使用模型和会话。

"""create teams table

Revision ID: 169ad57156f0
Revises: 29b4c2bfce6d
Create Date: 2014-06-25 09:00:06.784170
"""

revision = '169ad57156f0'
down_revision = '29b4c2bfce6d'

from alembic import op
import sqlalchemy as sa
from sqlalchemy import orm
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()


class Player(Base):
    __tablename__ = 'players'

    id = sa.Column(sa.Integer, primary_key=True)
    name = sa.Column(sa.String, nullable=False)
    team_name = sa.Column('team', sa.String, nullable=False)
    team_id = sa.Column(sa.Integer, sa.ForeignKey('teams.id'), nullable=False)

    team = orm.relationship('Team', backref='players')


class Team(Base):
    __tablename__ = 'teams'

    id = sa.Column(sa.Integer, primary_key=True)
    name = sa.Column(sa.String, nullable=False, unique=True)


def upgrade():
    bind = op.get_bind()
    session = orm.Session(bind=bind)

    # create the teams table and the players.team_id column
    Team.__table__.create(bind)
    op.add_column('players', sa.Column('team_id', sa.ForeignKey('teams.id'), nullable=False)

    # create teams for each team name
    teams = {name: Team(name=name) for name in session.query(Player.team).distinct()}
    session.add_all(teams.values())

    # set player team based on team name
    for player in session.query(Player):
        player.team = teams[player.team_name]

    session.commit()

    # don't need team name now that team relationship is set
    op.drop_column('players', 'team')


def downgrade():
    bind = op.get_bind()
    session = orm.Session(bind=bind)

    # re-add the players.team column
    op.add_column('players', sa.Column('team', sa.String, nullable=False)

    # set players.team based on team relationship
    for player in session.query(Player):
        player.team_name = player.team.name

    session.commit()

    op.drop_column('players', 'team_id')
    op.drop_table('teams')

迁移定义了单独的模型,因为代码中的模型表示数据库的当前状态,而迁移表示过程中的步骤。您的数据库可能沿着该路径处于任何状态,因此模型可能尚未与数据库同步。除非您非常小心,否则直接使用真实模型会导致缺少列,无效数据等问题。更清晰地明确说明要在迁移中使用的列和模型。


11

您还可以使用直接SQL参见(Alembic操作参考),如以下示例所示:

from alembic import op

# revision identifiers, used by Alembic.
revision = '1ce7873ac4ced2'
down_revision = '1cea0ac4ced2'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands made by andrew ###
    op.execute('UPDATE STOCK SET IN_STOCK = -1 WHERE IN_STOCK IS NULL')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###

如果我一直想从外部文件中读取一条SQL语句,然后将其传递给op.executein upgrade(),是否可以提供一种默认的模板以供alembic revision命令使用(生成.py文件的默认主体)?
昆汀

1
我不知道@Quentin。这是一个有趣的想法。
Martlark

6

我建议使用临时表来使用SQLAlchemy核心语句(如官方文档中所述),因为它允许使用不可知论的SQL和pythonic编写,并且也是独立的。对于迁移脚本,SQLAlchemy Core是两全其美的。

这是概念的示例:

from sqlalchemy.sql import table, column
from sqlalchemy import String
from alembic import op

account = table('account',
    column('name', String)
)
op.execute(
    account.update().\\
    where(account.c.name==op.inline_literal('account 1')).\\
        values({'name':op.inline_literal('account 2')})
        )

# If insert is required
from sqlalchemy.sql import insert
from sqlalchemy import orm

session = orm.Session(bind=bind)
bind = op.get_bind()

data = {
    "name": "John",
}
ret = session.execute(insert(account).values(data))
# for use in other insert calls
account_id = ret.lastrowid
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.