Alembicを使ったPostgreSQLマイグレーションのセットアップと実行方法


Poetryを使用してプロジェクトを初期化し、Docker ComposeでPostgreSQLデータベースをセットアップして、pgvectorエクステンションを利用したデータベースマイグレーションを実行する方法をまとめました。



# Githubリポジトリ

DBへマイグレーションに必要となる最低限のコードをリポジトリにまとめました。

https://github.com/phlin0424/pgvetor_base


以下は、リポジトリ内格納されているコードを例として説明します。


# セットアップ

## パッケージのインストール

まずは、Poetryを使用してプロジェクトを初期化します。
```shell
poetry init
```

下記のパッケージを記入し、`poetry install`にてインストールします。

```

[tool.poetry.dependencies]
python = "^3.11"
sqlalchemy = "^2.0.31"
alembic = "^1.13.2"
pgvector = "^0.3.0"
psycopg2-binary = "^2.9.9"

```

## PostgreSQL DBのセットアップ

次に、以下のような`docker-compose.yaml`ファイルを作成して、PostgreSQLデータベースコンテナをセットアップします。


```
version: '3.8'
services:
  db:
    image: ankane/pgvector
    environment:
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
      POSTGRES_DB: pgvector_db
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data
    deploy:
      resources:
        limits:
          cpus: '0.50'
          memory: '512M'
        reservations:
          cpus: '0.25'
          memory: '256M'
volumes:
  pgdata:

```




## データベースモデルの作成

次に、データベースにマイグレーションするテーブルを定義します。以下のコードは、`db_model.py`ファイル内にテーブルを定義する例です。


```python
from pgvector.sqlalchemy import Vector
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import declarative_base

Base = declarative_base()

# Baseクラスを継承してテーブルモデルを定義
class AnkiNoteModel(Base):
    """このモデルに基づいてテーブルを作成します。"""

    __tablename__ = "vector_table_1"

    product_id = Column(Integer, primary_key=True)
    product_name = Column(String, nullable=False)
    description = Column(String, nullable=False)
    vector = Column(Vector(dim=1536))

```

このリポジトリでは、上記のファイルは`src/models`ディレクトリに配置されていますが、独自に作成しても構いませんし、列を必要に応じて変更することも可能です。


## Alembicの初期化

下記のコマンドでAlembiプロジェクトを初期化します。

```shell
alembic init
```

これにより、下記のようにプロジェクトルートに`alembic.ini`ファイルと、`migrations/`ディレクトリが生成されます。

```
$ tree
.
├── alembic.ini
└── migrations
    ├── README
    ├── env.py
    ├── script.py.mako
    └── versions

```


## `alembic.ini`の編集

次は、`alembic.ini`ファイルの`sqlalchemy.url`行を以下のように編集します。

```ini
sqlalchemy.url = postgresql+psycopg2://postgres:postgres@localhost/pgvector_db
```

## `env.py`の編集

`db_model.py`内で定義した`Base`を`env.py`にインポートします。

```python
from models.db_model import Base
```

また、次の部分を編集します。
```python
target_metadata = Base.metadata
```

さらに、`Alembic`はデフォルトで`pgvector`に対応していないため、`env.py`に以下の関数を追加します。

```python
def create_vector_extension(connection) -> None:
    try:
        with Session(connection) as session:  # type: ignore[arg-type]
            statement = sqlalchemy.text(
                "BEGIN;" "CREATE EXTENSION IF NOT EXISTS vector;" "COMMIT;"
            )
            session.execute(statement)
            session.commit()
    except Exception as e:
        raise Exception(f"Failed to create vector extension: {e}") from e

def do_run_migrations(connection) -> None:
    connection.dialect.ischema_names["vector"] = pgvector.sqlalchemy.Vector

    context.configure(
        connection=connection,
        target_metadata=target_metadata,
    )

    with context.begin_transaction():
        context.run_migrations()

```


:::note info 
`env.py`はマイグレーションコードを生成するテンプレートとして機能しますが、`Alembic`はデフォルトで`pgvector`カラムを認識しないため、追加の関数を組み込む必要があります。これにより、マイグレーション中のエラーを防止することができます。
:::




# マイグレーション


## PostgreSQLの起動

作成した`docker-compose.yaml`を使用して、PostgreSQLコンテナをビルドします。

```
docker-compose up --build
```


## マイグレーション準備

:::note warn
PostgreSQLコンテナが実行している状態でマイグレーションを実施してください。
:::

すべての設定が完了したら、以下のコマンドを実行してマイグレーションスクリプトを生成します。


```shell
alembic revision --autogenerate -m "Create initial tables"
```

:::note info
-m 後のメッセジーは自分が分かる文言を入れ替えてください
:::



このコマンドにより、`migrations/versions/`ディレクトリ内にマイグレーションスクリプト（例：`{unique_id}_create_initial_tables.py`）が生成されます。中身は下記のように：

```python
def upgrade() -> None:
    op.create_table(
        "vector_table_1",
        sa.Column("product_id", sa.Integer(), nullable=False),
        sa.Column("product_name", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=False),
        sa.Column("vector", pgvector.sqlalchemy.vector.VECTOR(dim=1536), nullable=True),
        sa.PrimaryKeyConstraint("product_id"),
    )

```
生成されたコードが正しいことを確認してください。


:::note info
`db_model.py`内定義されているデータモデルを確認しながら生成されたコードを編集してください（必要があれば）
:::


また、デフォルトでは、`Alembic`は`pgvector`パッケージをインポートしないため、生成されたマイグレーションスクリプトに以下の行を手動で追加する必要があります。


```python
import pgvector
```

## マイグレーションの実施

以下のコマンドを実行して、マイグレーションをデータベースに適用します。

```shell
alembic upgrade head
```


## データベースの確認

マイグレーションが正常に適用されたかどうかを確認するために、データベース内のテーブルを確認します。

```
docker-compose exec db psql -U postgres -d pgvector_db
```

データベース内に`vector_table_1`というテーブルが存在するはずです。


# 他

データベースモデルの定義を変えるたびに再びマイグレーションを実施する必要があります。この場合は：

`alembic revision --autogenerate -m "<message>"`

→生成されたコードを確認

→`alembic upgrade head`

を繰り返します。