Создаем начальную миграцию с alembic для существующей базы

Не так давно я пришел на проект, где пользовали SQLAlchermy и Alembic. Но по воле рока так случилось, что alembic подключили после того, как создали в базе кучу объектов. Для тех, кто не в курсе, SQLAlchemy - это библиотека и ORM для питона, а Alembic - это инстумент для работы с миграциями для SQLAlchemy.

Мы захотели научиться делать 2 вещи: создавать базу с нуля прогоняя миграции и генерить миграции автоматом. Для начала создадим начальную миграцию, от которой и будем уже потом генерить последующие миграции автоматом.

Начнем с моделей. В статье я публиковать модели из проекта не буду, а определю довольно простенькие модели пользователей, публикаций и комментов в каком-нибудь блоге. Создадим файл db.py в корне проекта со следующим содержанием.

from sqlalchemy import create_engine, Column, Integer, String, Text, ForeignKey
from sqlalchemy.orm import declarative_base

engine = create_engine('postgresql://romblin@localhost/db')

Base = declarative_base()


class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True)
    username = Column(String(255), unique=True, nullable=False)
    password = Column(String(255), nullable=False)
    email = Column(String(255), unique=True, nullable=False)


class Comment(Base):
    __tablename__ = 'comments'

    id = Column(Integer, primary_key=True)
    body = Column(Text, nullable=False)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    post_id = Column(Integer, ForeignKey('posts.id'), nullable=False)


class Post(Base):
    __tablename__ = 'posts'

    id = Column(Integer, primary_key=True)
    title = Column(String(255), nullable=True)
    text = Column(Text, nullable=False)
    author_id = Column(Integer, ForeignKey('users.id'), nullable=False)

Для нашего случая можно предполагать, что все эти объекты уже есть в базе.

db=# \dt
          List of relations
 Schema |   Name   | Type  |  Owner  
--------+----------+-------+---------
 public | comments | table | romblin
 public | posts    | table | romblin
 public | users    | table | romblin
(3 rows)

Подключаем и настраиваем алембик:

$ alembic init migration

В файле alembic.ini указываем адрес базы:

[alembic]
...
sqlalchemy.url = postgresql://romblin@localhost/db

В файле migration/env.py импортируем все модели и указываем target_metadata:

from db import *
target_metadata = Base.metadata

В самом простом случае у нас есть строгое соответствие между объектами, которые есть базе и моделями. В этом случае можно просто создать пустую базу и попросить алебмик создать миграцию для этой базы, на основе моделей и он как раз создаст нужную нам началную миграцию, и уже от этой миграции и плясать, переключившись на основную базу.

Создаем базу:

db=# CREATE DATABASE tmp;
CREATE DATABASE

Генерим миграцию, не забыв перед этим указать правильное имя базы в настройках.

$ alembic revision --autogenerate -m 'initial'

На выходе получаем migration/versions/8cb9616d5ef0_initial.py:

"""initial

Revision ID: 8cb9616d5ef0
Revises:
Create Date: 2021-10-24 19:27:31.205081

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '8cb9616d5ef0'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
   # ### commands auto generated by Alembic - please adjust! ###
   op.create_table('users',
   sa.Column('id', sa.Integer(), nullable=False),
   sa.Column('username', sa.String(length=255), nullable=False),
   sa.Column('password', sa.String(length=255), nullable=False),
   sa.Column('email', sa.String(length=255), nullable=False),
   sa.PrimaryKeyConstraint('id'),
   sa.UniqueConstraint('email'),
   sa.UniqueConstraint('username')
   )
   op.create_table('posts',
   sa.Column('id', sa.Integer(), nullable=False),
   sa.Column('title', sa.String(length=255), nullable=True),
   sa.Column('text', sa.Text(), nullable=False),
   sa.Column('author_id', sa.Integer(), nullable=False),
   sa.ForeignKeyConstraint(['author_id'], ['users.id'], ),
   sa.PrimaryKeyConstraint('id')
   )
   op.create_table('comments',
   sa.Column('id', sa.Integer(), nullable=False),
   sa.Column('body', sa.Text(), nullable=False),
   sa.Column('user_id', sa.Integer(), nullable=False),
   sa.Column('post_id', sa.Integer(), nullable=False),
   sa.ForeignKeyConstraint(['post_id'], ['posts.id'], ),
   sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
   sa.PrimaryKeyConstraint('id')
   )
   # ### end Alembic commands ###


def downgrade():
   # ### commands auto generated by Alembic - please adjust! ###
   op.drop_table('comments')
   op.drop_table('posts')
   op.drop_table('users')
   # ### end Alembic commands ###

Теперь возвращаем настройки и вуа-ля, у нас есть initial миграция, от которой алембик будет генерить миграции автоматом и прогоняя этот набор миграций можно создать базу с нуля. Для автотетов тестов, к примеру.

Это был самый простой случай. Теперь рассмотрим ситуцию, когда в базе есть объекты, которые не определены на уровне моделей, к примеру хранимые процедуры, представления или типы. Это был как раз наш случай, и таких объектов в базе было довольно много. Что делать? Я создал дамп схемы и импортировал его в начальной миграции.

Дамп создавал так:

$ pg_dump -h localhost -U romblin --dbname=db --no-owner --schema-only --no-privileges > migration/schema.dump

Далее создаем пустую миграцию и импортируем в ней нам дамп.

$ alembic revision -m 'initial'

migration/versions/3f81b707dbe4_initial.py

"""init

Revision ID: 3f81b707dbe4
Revises:
Create Date: 2021-10-06 12:33:17.827125

"""
from pathlib import Path

from alembic import op
# revision identifiers, used by Alembic.
from sqlalchemy import text

revision = '3f81b707dbe4'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
   dump_path = Path(__file__).parent.parent.absolute() / 'schema.dump'

   with open(dump_path, 'r') as sql_reader:
       op.execute(text(sql_reader.read()))

   op.execute(text('SET search_path = public'))


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

Теперь как эту миграцию накатывать на уже существующую базу. Здесь нам поможет команда stamp алембика. На просто сделает запись в служебную таблицу алебмика о миграции без реального ее применения.

$ alembic stamp 3f81b707dbe4

Готово, теперь мы можем автоматически генерировать миграции и воссоздавать базу в нуля с помощью alembic-а.