Как выполнить вставки и обновления в скрипте обновления Alembic?

Мне нужно изменить данные во время обновления 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. Обновите player.team_id на основе названия player.team:

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

Как выполнять вставки и обновления внутри сценария обновления?

То, что вы просите, - это миграция данных , в отличие от миграции схемы, которая наиболее распространена в документации 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')

The migration defines separate models because the models in your code represent the current state of the database, while the migrations represent steps along the way. Your database might be in any state along that path, so the models might not sync up with the database yet. Unless you're very careful, using the real models directly will cause problems with missing columns, invalid data, etc. It's clearer to explicitly state exactly what columns and models you will use in the migration.

You can also use direct SQL see (Alembic Operation Reference) as in the following example:

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 ###

I recommend using SQLAlchemy core statements using an ad-hoc table, as detailed in the official documentation, because it allows the use of agnostic SQL and pythonic writing and is also self-contained. SQLAlchemy Core is the best of both worlds for migration scripts.

Here is an example of the concept:

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