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

Githubリポジトリ

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

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

セットアップ

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

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

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ファイル内にテーブルを定義する例です。

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プロジェクトを初期化します。

alembic init

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

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

`alembic.ini`の編集

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

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

`env.py`の編集

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

from models.db_model import Base

また、次の部分を編集します。

target_metadata = Base.metadata

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

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()

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

マイグレーション

PostgreSQLの起動

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

docker-compose up --build

マイグレーション準備

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

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

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

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

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

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"),
    )

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

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

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

import pgvector

マイグレーションの実施

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

alembic upgrade head

データベースの確認

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

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

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

他

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

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

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

→alembic upgrade head

を繰り返します。