psycopg2でスキーマを指定してCOPYする方法
要約
業務中に取り組んだ不具合解決のフローを記事にしてみました。
詳細は下記のとおりです。
- PostgreSQLにCSVファイルをCOPYする際、スキーマ指定ができない不具合が発生
- 原因はライブラリの更新で、psycopg2 2.9以降では
copy_from関数でスキーマを指定出来なくなった為 - 上記の事象に対して、スキーマ検索パスを設定して解決した
今回は、不具合を再現する所から始め、問題の解決フローをなるべく細かく記事にしてみます！
検証手順
- COPYの動作を検証するDB・入力データを用意
- 古いバージョンのpsycopg2で、
copy_fromメソッドが正常に動作することを確認 - 新しいバージョンのpsycopg2で、同様のコードが動かない不具合を確認
- スキーマ検索パスを指定して上記問題を解決する
検証環境
OS・ソフトウェア
- Windows10 x86-64
- PostgreSQL 16.1
- Python 3.8.10
- Python 3.10.5
- psycopg2 2.9.9
- psycopg2 2.8.6
- pyenv 3.1.1
データベース構成
検証目的で簡単なDBを用意しました。
スキーマを指定できるか確認するため、異なるスキーマ名で同じ名前のテーブルを作成しています。
データベース：copy_test
スキーマ：public, sub_schema
テーブル：users（下図のテーブルを2つのスキーマに作成）
入力データ
usersテーブルと同じ構造のCSVファイルを用意しました。
ユーザー名とメールアドレスのみのシンプルな入力ファイルです。
Azusa Nakano,azusa11@example.com
Megumi Kitamura,megumi@example.com
Naoe Riki,naoenaoe@example.com
Satoshi Akatsuka,s4toshi@example.com
古いバージョンでcopy_fromメソッドの動作確認
まずは、古いバージョンのライブラリでスキーマを指定したCOPYができるか動作検証を行います。
Pythonのバージョンは3.8.10、psycopg2のバージョンは2.8.6を利用しました。
下記コードでsub_schema.usersにデータをCOPYします。
import psycopg2
def main():
"""psycopg2 2.8.6でCOPYの動作を確認
"""
schema = "sub_schema"
table_name = "users"
with get_connection() as conn:
with conn.cursor() as cur:
with open(r"input_file\user_list.csv", mode="r", encoding="utf-8") as f:
cur.copy_from(f, table=f"{schema}.{table_name}", sep=",", columns=('name', 'email'))
conn.commit()
def get_connection():
dsn = {
"host":"localhost",
"port":"5432",
"database":"copy_test",
"user":"postgres",
"password":"postgres",
}
return psycopg2.connect(**dsn)
if __name__ == "__main__":
main()
ソースコードを実行し、無事にCOPYされた事を確認しました。
新しいバージョンでcopy_fromメソッドの動作確認
次に、新しいバージョンのライブラリでスキーマを指定したCOPYができるか動作検証を行います。
Pythonのバージョンは3.10.5、psycopg2のバージョンは2.9.9を利用しました。
同じコードを実行しましたが、下記のようなエラーが発生します。
PS C:\git\copy-test> & c:/git/copy-test/venv/Scripts/python.exe c:/git/copy-test/venv/copy_test.py
Traceback (most recent call last):
File "c:\git\copy-test\venv\copy_test.py", line 25, in <module>
main()
File "c:\git\copy-test\venv\copy_test.py", line 11, in main
cur.copy_from(f, table=f"{schema}.{table_name}", sep=",", columns=('name', 'email'))
psycopg2.errors.UndefinedTable: リレーション"sub_schema.users"は存在しません
psycopg2.errors.UndefinedTable: リレーション"sub_schema.users"は存在しません
...と記載されていますが、スキーマは実際に存在します。
初めてこのエラーに当たったときは混乱しました。。
解決策
DB側を調査しても解決せず、google検索したところ、あまり良い情報にたどり着けませんでした。
そこで、psycopg2のGithubでエラーを検索すると、下記Issueがヒットしました。
cursor.copy_fromメソッドがテーブル引数でスキーマを受け付けなくなった。とのことで、同じ事象についてのスレッドから解決策を探しました。
スレッド内では「スキーマ検索パスを指定する」という方法で回避する方法が紹介されていました。
具体的には、下記ソースコードを追加する方法です。
cursor.execute(f'SET search_path TO {schema}')
上記コードを追記して、テーブルからスキーマ指定を外すことで、無事にCOPYが実行できました。
まとめ
問題解決のフローから学んだことを簡単にまとめます。
- ライブラリのバージョンアップで不具合が発生しないかきちんと確認する
- マイナーなライブラリの不具合で詰まった時は、一次情報（GithubのIssues等）から調べる
今回検証に利用したコードは下記Githubにアップロードしています。詳細が見たい方は是非ご確認ください！
注意事項
psycopg2開発者から、スキーマ検索パスの指定がセッション全体に影響する場合もあるとコメントが来ています。
複数セッションで異なるスキーマ並行して扱うケースにはご注意ください。
dvarrazzo commented on Nov 6, 2021
I was able to adapt my existing function by adding this line just prior to calling "copy_from":
curs.execute(f'SET search_path TO {schema}')
This affects the whole session. It might be a good solution for you, as well as for a program, but it's not a good solution with a driver because it has global (session-wide) repercussions.
参考記事
COPY-related methods
cursor.copy_from table argument no longer accepting schema #1294
スキーマ検索パスを設定する
Discussion