工具介紹 - Alembic - migration tool

前言

在使用FastAPI時,通常都會搭配sqlalchemy來連接database,角色介於python與db之間,示意圖像下方這樣

python <-> sqlachemy <-> database

而除了連接資料庫外,sqlalchemy最主要的功能還有提供ORM供開發人員使用,所謂的ORM(Object–relational mapping),能將python的類別(Class)和屬性(attributes)直接映射成關聯式資料庫,也提供了部份的function能夠輕鬆調用資料庫的數據。

而database頻繁調整的情況下,版次紀錄格外重要,但FastAPI並不像Django有內建的migration功能,FastAPI如同framework名稱一樣很單純,就是為了api而生,Django則是功能更為完整的Framework,可架設出包含完整前後端的網頁。

因此Fastapi需要另外安裝migration tool,也就是本篇要介紹的Alembic

(內容多出自下方的tutorial,內容很扎實,推薦觀看)

Python API Development - Comprehensive Course for Beginners

什麼是Alembic?

Alembic是一個migration tool,能夠用於產出migration file,並對資料庫執行調整的工具。FastAPI的官方也是介紹使用Alembic。

SQL (Relational) Databases - FastAPI

使用範例

首先安裝套件

pip install alembic

看一下有哪些指令可以用

alembic -h

認識指令

執行後顯示出許多可以下的指令,可以想像成database使用的git就OK了,裡頭最常使用到的會是init, current, revision, upgrade, downgrade知道到這即可,等等會用到

usage: alembic [-h] [--version] [-c CONFIG] [-n NAME] [-x X] [--raiseerr]
               {branches,current,downgrade,edit,ensure_version,heads,history,init,list_templates,merge,revision,show,stamp,upgrade}
               ...

positional arguments:
  {branches,current,downgrade,edit,ensure_version,heads,history,init,list_templates,merge,revision,show,stamp,upgrade}
    branches            Show current branch points.
    current             Display the current revision for a database.
    downgrade           Revert to a previous version.
    edit                Edit revision script(s) using $EDITOR.
    ensure_version      Create the alembic version table if it doesn't exist
                        already .
    heads               Show current available heads in the script directory.
    history             List changeset scripts in chronological order.
    init                Initialize a new scripts directory.
    list_templates      List available templates.
    merge               Merge two revisions together. Creates a new migration
                        file.
    revision            Create a new revision file.
    show                Show the revision(s) denoted by the given symbol.
    stamp               'stamp' the revision table with the given revision;
                        don't run any migrations.
    upgrade             Upgrade to a later version.

optional arguments:
  -h, --help            show this help message and exit
  --version             show program's version number and exit
  -c CONFIG, --config CONFIG
                        Alternate config file; defaults to value of
                        ALEMBIC_CONFIG environment variable, or "alembic.ini"
  -n NAME, --name NAME  Name of section in .ini file to use for Alembic config
  -x X                  Additional arguments consumed by custom env.py
                        scripts, e.g. -x setting1=somesetting -x
                        setting2=somesetting
  --raiseerr            Raise a full stack trace on error

建立環境和專案

先建立環境

mkdir alembic-demo && cd alembic-demo
virtualenv env
env\Scripts\activate

安裝需要的套件

pip install fastapi
pip install sqlalchemy
pip install alembic
pip install uvicorn
pip install psycopg2

先建立名為app的資料夾,再新增 main.py 並寫入

# app/main.py

from fastapi import FastAPI, status

app = FastAPI()

@app.get('/', status_code=status.HTTP_200_OK)
def root():
    return {'message': 'Hello FastAPI'}

接著運行並訪問 http://127.0.0.1:8000/,確定專案能成功運行

uvicorn app.main:app --reload

建立和database的connection

同樣在 app 資料夾中,新增 database.py,並建立engine,此處建立的 SessionLocal 是作為database session 使用,而 Base 是等等讓各個 model 繼承使用的。

# app/database.py

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

DB_USERNAME = 'your DB_USERNAME'
DB_PASSWORD = 'your DB_PASSWORD'
DB_HOST = 'your DB_HOST'
DB_NAME = 'your DB_NAME'

SQLALCHEMY_DATABASE_URL = f"postgresql://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}/{DB_NAME}"

engine = create_engine(SQLALCHEMY_DATABASE_URL)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

# Dependency
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

建立Model

同樣在 app 資料夾中,新增 models.py,並寫入兩個 Post 和 User 兩個類別,此處的類別繼承了在database.py 中所建立的 Base

# app/models.py

from .database import Base
from sqlalchemy import TIMESTAMP, Column, ForeignKey, Integer, String, Boolean
from sqlalchemy.sql.expression import text
from sqlalchemy.orm import relationship

class Post(Base):

    __tablename__ = 'posts'

    id = Column(Integer, primary_key=True, nullable=False)
    title = Column(String, nullable=False)
    content = Column(String, nullable=False)
    created_at = Column(
        TIMESTAMP(timezone=True),
        nullable=False,
        server_default=text('now()')
        )
    user = Column(Integer, ForeignKey(
        'users.id',
        ondelete="CASCADE"),
        nullable=False
        )
    user_relation = relationship('User')

class User(Base):

    __tablename__ = 'users'

    id = Column(Integer, primary_key=True, nullable=False)
    email = Column(String, nullable=False, unique=True)
    password = Column(String, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True),
                        nullable=False, server_default=text('now()'))

初始化Alembic

在和 app 資料夾同一層的地方,初始化一個名為 alembic 的 alembic 專案(聽起來有點饒口)

alembic init alembic

這時會多出一個名為 alembic 的資料夾和名為 alembic.ini 的檔案,此時先打開 alembic.ini ,並找到 sqlalchemy.url ,將值刪除並留空,等等這個變數的值會在別的地方填入。

# alembic.ini

...

sqlalchemy.url = 

...

接著到剛建立好的 alembic 資料夾中找到 env.py,從剛建立的models中import Base,用set_main_option設定剛刪除的值,以及更新 target_metadata 的值

# alembic/env.py

... 
from app.models import Base # <-- add this

DB_USERNAME = 'your DB_USERNAME'
DB_PASSWORD = 'your DB_PASSWORD'
DB_HOST = 'your DB_HOST'
DB_NAME = 'your DB_NAME'

config.set_main_option(
    "sqlalchemy.url",
    f"postgresql://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}/{DB_NAME}"
)

...

target_metadata = Base.metadata

...

編寫revision file

Alembic的運作流程為,先提出一個revision,接著編寫revision的內容,然後apply到資料庫中。此處先執行一個名為新增add post table的revision。

當執行後,在alembic資料夾中的version資料夾下會產生一個名稱由版本號加上該次revision名稱的.py檔案。

(env) C:\Users\User\alembic-demo>alembic revision -m "add post table"
Generating C:\Users\User\alembic-demo\alembic\versions\e4071dbdd1b8_add_post_table.py ...  done

以此例子來說,revision file如下,其中的upgrade即為本次revision所要做的變動,而downgrade即為本次revision如果rollback後要做的動作。

# alembic/versions/e4071dbdd1b8_add_post_table.py

"""add post table

Revision ID: e4071dbdd1b8
Revises: 
Create Date: 2022-08-03 15:38:58.692780

"""
from alembic import op
import sqlalchemy as sa

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

def upgrade() -> None:
    pass

def downgrade() -> None:
    pass

upgrade

加入要做的指令,這次的revision新增一個名為posts的table,並包含了id和title兩個欄位,如果要rollback這次的revision,則是直接drop table。

# alembic/versions/e4071dbdd1b8_add_post_table.py

"""add post table

Revision ID: e4071dbdd1b8
Revises:
Create Date: 2022-08-03 15:38:58.692780

"""
from alembic import op
import sqlalchemy as sa

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

def upgrade() -> None:
    op.create_table(
        'posts',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('title', sa.String(), nullable=False,)
    )
    pass

def downgrade() -> None:
    op.drop_table('posts')
    pass

接著用upgrade執行此次revision

(env) C:\Users\User\alembic-demo>alembic upgrade e4071dbdd1b8
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> e4071dbdd1b8, add post table

用current查看revision是否寫入,此時會看到剛剛的版本號出現在下方

(env) C:\Users\User\alembic-demo>alembic current
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
e4071dbdd1b8 (head)

此時的db中會出現兩張table,一個是剛剛建立的posts,另一個是alembic自動生成的用於紀錄當前版次的table

alembic-demo=# \dt+
                                  List of relations
 Schema |      Name       | Type  |  Owner   | Persistence |    Size    | Description
--------+-----------------+-------+----------+-------------+------------+-------------
 public | alembic_version | table | postgres | permanent   | 8192 bytes |
 public | posts           | table | postgres | permanent   | 8192 bytes |
(2 rows)

來看看版次table中寫了什麼? 就是剛剛revision的版本號

alembic-demo=# select * from alembic_version;
 version_num
--------------
 e4071dbdd1b8
(1 row)

而posts表格中也確實有了剛建立的兩個欄位

alembic-demo=# select * from posts;
 id | title
----+-------
(0 rows)

此時再用一個revision來新增content的欄位

(env) C:\Users\User\alembic-demo>alembic revision -m "add content column" 
Generating C:\Users\User\alembic-demo\alembic\versions\af2c1658d76e_add_content_column.py ...  done

此時的revision檔案中,就的down_revision就有值了,也就是前個版次。

# alembic/versions/af2c1658d76e_add_content_column.py

"""add content column

Revision ID: af2c1658d76e
Revises: e4071dbdd1b8
Create Date: 2022-08-03 16:14:59.225649

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'af2c1658d76e'
down_revision = 'e4071dbdd1b8'
branch_labels = None
depends_on = None

def upgrade() -> None:
    op.add_column(
        'posts',
        sa.Column('content', sa.String(), nullable=False)
    )
    pass

def downgrade() -> None:
    op.drop_column('posts', 'content')
    pass

接著一樣upgrade然後確認當前版次

(env) C:\Users\User\alembic-demo>alembic upgrade af2c1658d76e
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade e4071dbdd1b8 -> af2c1658d76e, add content column

(env) C:\Users\User\alembic-demo>alembic current
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
af2c1658d76e (head)

確認欄位是否新增

alembic-demo=# select * from posts;
 id | title | content
----+-------+---------
(0 rows)

downgrade

而如果想要回到前一個版次,不想要content這個欄位,可以用downgrade來完成,用downgrade指令再加上版本號,即可回到該版本,這邊用downgrade回到第一個版本號。

(env) C:\Users\User\alembic-demo>alembic downgrade e4071dbdd1b8
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running downgrade af2c1658d76e -> e4071dbdd1b8, 
add content column

再次確認,欄位已刪除

alembic-demo=# select * from posts;
 id | title
----+-------
(0 rows)

經過上面冗長的操作,可以瞭解migration是如何運作的,而總不可能在revision中慢慢敲每個欄位,這樣不是辦法,此時可以--autogenerate參數。

這個參數能自動偵測models中的各個Class,並生成revision檔案,並檢視當前資料庫與models中的差別,作為upgrade的內容。

此時需要先upgrade head,所謂的head即為執行current後出現的版本號

(env) C:\Users\User\alembic-demo>alembic upgrade head
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.INFO  [alembic.runtime.migration] Running upgrade e4071dbdd1b8 -> af2c1658d76e, add content column

執行後再建立一個新的revision,不過是加上了--autogenerate 加上參數後,將會自動偵測當前資料庫和models檔案中的差異,例如少了那些欄位,少了哪些table

(env) C:\Users\User\alembic-demo>alembic revision --autogenerate -m "add tables"  
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.INFO  [alembic.autogenerate.compare] Detected added table 'users'
INFO  [alembic.ddl.postgresql] Detected sequence named 'posts_id_seq' as owned by 
integer column 'posts(id)', assuming SERIAL and omitting
INFO  [alembic.autogenerate.compare] Detected added column 'posts.created_at'     
INFO  [alembic.autogenerate.compare] Detected added column 'posts.user'
INFO  [alembic.autogenerate.compare] Detected added foreign key (user)(id) on table posts
Generating C:\Users\User\alembic-demo\alembic\versions\63e1afe3fd86_add_tables.py ...  done

來看看revision檔案的內容,已經全部寫好了!

# alembic/versions/63e1afe3fd86_add_tables.py

"""add tables

Revision ID: 63e1afe3fd86
Revises: af2c1658d76e
Create Date: 2022-08-03 17:02:29.577274

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '63e1afe3fd86'
down_revision = 'af2c1658d76e'
branch_labels = None
depends_on = None

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('users',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('email', sa.String(), nullable=False),
    sa.Column('password', sa.String(), nullable=False),
    sa.Column('created_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('email')
    )
    op.add_column('posts', sa.Column('created_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False))
    op.add_column('posts', sa.Column('user', sa.Integer(), nullable=False))
    op.create_foreign_key(None, 'posts', 'users', ['user'], ['id'], ondelete='CASCADE')
    # ### end Alembic commands ###

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'posts', type_='foreignkey')
    op.drop_column('posts', 'user')
    op.drop_column('posts', 'created_at')
    op.drop_table('users')
    # ### end Alembic commands ###

這時再將內容apply到資料庫中

(env) C:\Users\User\alembic-demo>alembic upgrade 63e1afe3fd86
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.INFO  [alembic.runtime.migration] Running upgrade af2c1658d76e -> 63e1afe3fd86, add tables

再查看一下資料庫的內容,就變得和models.py中所定義的相同了,完成migration。

alembic-demo=# \dt+
                                  List of relations
 Schema |      Name       | Type  |  Owner   | Persistence |    Size    | Description
--------+-----------------+-------+----------+-------------+------------+-------------
 public | alembic_version | table | postgres | permanent   | 8192 bytes |
 public | posts           | table | postgres | permanent   | 8192 bytes |
 public | users           | table | postgres | permanent   | 8192 bytes |
(3 rows)

alembic-demo=# select * from alembic_version;
 version_num
--------------
 63e1afe3fd86
(1 row)

alembic-demo=# select * from posts;
 id | title | content | created_at | user
----+-------+---------+------------+------
(0 rows)

alembic-demo=# select * from users;
 id | email | password | created_at
----+-------+----------+------------
(0 rows)

參考資料

Python API Development - Comprehensive Course for Beginners

SQL (Relational) Databases - FastAPI

SQLAlchemy - The Database Toolkit for Python

Tags:
# python
# fastapi