こちらの記事を参考にしてインサートを試みましたが躓いたので書き記します。大方同じなのですが、文字列挿入で二つ困ったところがありました。
日本語文字列を挿入する時のエンコーディングエラーとバックスラッシュを含んだ文字列のエンコーディングエラーです。
実装
from io import StringIO
import pandas as pd
import psycopg2
class Client:
def __init__(self, dsn: str) -> None:
"""
Arguments:
dsn(str): 'postgresql://{username}:{password}@{hostname}:{port}/{dbname}'
"""
self._cur = None
self._conn = psycopg2.connect(dsn)
self._conn.set_client_encoding('UTF8')
self._cur = self._conn.cursor()
def __del__(self) -> None:
if self._cur is not None:
self._cur.close()
self._conn.close()
def insert(self, table: str, values: pd.DataFrame) -> None:
buf = StringIO()
df.to_csv(buf, sep='\t', na_rep='\\N', index=False, header=False)
buf.seek(0)
columns = df.columns.values.tolist()
self._cur.copy_from(buf, table, columns=columns)
self._conn.commit()
今回は発生しませんでしたが、 pd.DataFrame
で NULLABLE
な整数型を扱う場合もエラーになる可能性があるようです。
参照: https://qiita.com/hoto17296/items/b6c90db4b9bcdb7b6d78
日本語のエンコーディングの問題
set_client_encoding('UTF8')
で解決
参考: https://www.psycopg.org/docs/connection.html#connection.set_client_encoding
バックスラッシュの問題
これは外から実装してしまいましたが、次のようにして回避しました。
ここでは a
というカラムが文字列だとします
import os
df = get_dataframe() # 任意の方法でデータフレームを取得
df.a = df.a.str.replace('\\', '\\\\')
client = Client(os.environ.get('POSTGRES_DSN')
client.insert(table, df)
他の実装
上では copy_from
を用いましたが、 copy_expert
を使う場合は次のように書けます。
cur.copy_expert(
f"""
COPY {table}(
{','.join(columns)}
)
FROM STDIN
WITH
DELIMITER AS ' '
NULL AS '\\N'
ENCODING 'utf8'
""",
buf,
)
こちらではクエリの中でエンコーディングを指定しているため、 set_client_encoding('UTF8')
は不要です。
参考: https://www.psycopg.org/docs/cursor.html#cursor.copy_expert