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
を繰り返します。