psycopg2でパラメータを渡す

背景

psycopg2を使ってherokuのPostgreSQLを使おうとしたところ、パラメータの渡し方が思ったよりも面倒だった。自分用のメモを兼ねて作成。

環境

OS : macOS Big Sur 11.6.5
python : Python 3.9.12
psycopg2-binary : 2.9.3

使うパッケージ群

psycopg2を先にダウンロードしておく。

ssh

pip install psycopg2-binary

最低限のパッケージをインポート。
（sqlはパラメータを安全にクエリに渡すために必要。）

python

from psycopg2 import (
    connect, sql
)

おまじない

データベースのURLは、herokuでは "Settings > Config Vars"の中に入っている。
ここではauto commitを有効にしておく。
（デフォルトは無効。その場合は都度コミットする必要がある。）

python

DATABASE_URL = "hogehoge"

# データベースに接続
conn = connect(DATABASE_URL)

# auto commit を有効にする
conn.autocommit = True

# カーソルの設定
curs = conn.cursor()

クエリのベタ書き

テーブルの作成

今回は"hoge"という名前のテーブルを作り、ユーザー名（string型）とタイムスタンプ（datetime型）を格納する。
まずはテーブルを作る。

python

# hoge という名前のテーブルがすでにある場合は削除する
curs.execute("DROP TABLE IF EXISTS hoge")

# 新たに hoge というテーブルを作成
curs.execute("CREATE TABLE hoge (user_name TEXT, timestamp TIMESTAMP)")

# 作成できているか確認
curs.execute("""SELECT table_name 
    FROM information_schema.tables
    WHERE table_schema = 'public';"""
)
curs.fetchall()

outputs
 [('hoge',)]

行の追加

INSERT文を使って行を追加する。

python

# タイムスタンプを作る
import datetime
ts = datetime.datetime.now()

# データの挿入
curs.execute("INSERT INTO hoge (user_name, timestamp) VALUES ('John Smith', ts)")

これだとうまくいかない。正しい方法でパラメータをクエリに渡す必要がある。

outputs
UndefinedColumn: column "ts" does not exist
LINE 1: INSERT INTO hoge (user_name, timestamp) VALUES ('John Smith', ts)
                                                                      ^

個数のわかっているパラメータを渡す

psycopg2のドキュメントに記載がある。
パラメータを渡したい場所に%sを入れて、executeの第二引数にパラメータのタプルを指定する。

python

que = "INSERT INTO hoge (user_name, timestamp) VALUES (%s, %s)"
curs.execute(que,("John Smith",ts))

このようにすると、データを挿入することができた。

python

curs.execute("SELECT * FROM hoge")
curs.fethall()

outputs
[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))]

%sでパラメータが渡せない！

"SELECT * FROM hoge"の"hoge"をパラメータとして渡したい。%sを使って書くと

python

curs.execute("SELECT FROM * %s",("hoge",))

となりそうだが、これはうまくいかない。おそらく、SQL injectionを防ぐためのもの。

outputs
SyntaxError: syntax error at or near "'hoge'"
LINE 1: SELECT * FROM 'hoge'
                      ^

この問題の解決方法はいくつかある。

1. String.formatを使う

一旦クエリをstringとして作って、executeに渡す方法。SQL injectionに対して脆弱なので推奨されない。

python

que = "SELECT * FROM {}".format("hoge")
print(que)

curs.execute(que)
curs.fetchall()

outputs
SELECT * FROM hoge
[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))]

同様にf-stringや+ %演算子を使った方法があるが、どれもSQL injectionに脆弱であり、使用しないことが強く推奨される(参考："The problem with the query parameters" Psycopg 2.9.3 documentation)。

2. sql.SQL, sql.Identifierを使う

安全にパラメータを渡すためにpsycopg2.sqlを用いる方法がある。クエリをsql.SQLで、パラメータをsql.Identifier変換してからformatする。

python

SQL = sql.SQL("SELECT * FROM {}").format(sql.Identifier("hoge"))

print(SQL)
print(type(SQL))
print("\n")

curs.execute(SQL)
curs.fetchall()

outputs
Composed([SQL('SELECT * FROM '), Identifier('hoge')])
<class 'psycopg2.sql.Composed'>

[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))]

この方法では、元のクエリとパラメータが別々にサーバーに送られて、サーバーサイドで処理が行われる。
クエリをstringとして確認したい場合は次のようにする。

python

print(SQL.as_string(conn))

outputs
SELECT * FROM "hoge"

個数のわからないパラメータを渡す

パラメータが全てstring型の場合は非常に簡単。
例えばhogeというテーブルからuser_nameとtimestampを取り出したいときは次のように
、sql.Identifier(parameter)とsql.SQL(',')を結合させれば良い。
結合させるためには、お馴染みの区切り文字.join(list)を使う。

python

SQL = sql.SQL("SELECT {field} FROM {name}").format(
    field = sql.SQL(',').join([
        sql.Identifier("user_name"),
        sql.Identifier("timestamp")
    ]),
    name = sql.Identifier("hoge")
)

print(SQL.as_string(conn))
print("\n")

curs.execute(SQL)
curs.fetchall()

outputs
SELECT "user_name","timestamp" FROM "hoge"

[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))]

しかし、渡すパラメータがstring型以外の場合はうまくいかない。sql.Identifierはstring型の引数しか取らないからだ。

python

sql.Identifier(ts)

outputs
TypeError: SQL identifier parts must be strings

この場合、sql.Literalを用いることで解決できる。

python

sql.Literal(ts)

outputs
Literal(datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))

sql.Literalを用いてhogeにデータを挿入する。

python

SQL = sql.SQL("INSERT INTO {name} ({field}) VALUES ({values})").format(
    name = sql.Identifier("hoge"),
    field = sql.SQL(",").join([
        sql.Identifier("user_name"),
        sql.Identifier("timestamp")
    ]),
    values = sql.SQL(",").join([
        sql.Literal("Mary Smith"),
        sql.Literal(ts)
    ])
)

print(SQL.as_string(conn))
print("\n")

curs.execute(SQL)

curs.execute("SELECT * FROM hoge")
curs.fetchall()

outputs
INSERT INTO "hoge" ("user_name","timestamp") VALUES ('Mary Smith','2022-04-03T08:51:40.312203'::timestamp)

[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203)),
('Mary Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))]

なお、valuesをjoinする際にsql.Identifierを使って

python

values = sql.SQL(",").join([
    sql.Identifier('Mary Smith'),
    sql.Literal(ts)
])

とするとうまくいかない。

outputs
UndefinedColumn: column "Mary Smith" does not exist
LINE 1: INSERT INTO "hoge" ("user_name","timestamp") VALUES ("Mary Smith",'2022...

これはクエリの中の"Mary Smith"がダブルクオーテーションで囲まれていることが原因のようだ（参考："cannot get simple PostgreSQL insert to work" Stack Overflow）。

classを作ってパラメータを渡す

これまではパラメータを明示的にクエリに渡していた。パラメータを渡すもう一つの方法は、自作のclassを使う方法である。
まず、次のようなclassを定義する。

python

class User():
    def __init__(self, name, ts):
        self.user_name = name
        self.timestamp = ts

このclassで生成されたインスタンスをパラメータに渡すのである。

ts = datetime.datetime.now()

# User classのインスタンス user1を作る
user1 = User("Tom Brown", ts)

そのためにいくつかのモジュールと関数が必要になる。

python

from psycopg2 import (
    adapt, register_adapter, AsIs
)

# User classをクエリに渡せるようにする関数
def user_adapter(user):
    return AsIs("{}, {}".format(
        adapt(user.user_name),
        adapt(user.timestamp)
    ))

# user_adapter関数を登録する
register_adapter(User,user_adapter)

Userclassがパラメータとして渡されると、user_adapterが呼び出されるようになった。
これでデータを挿入することができる。

python

SQL = sql.SQL("INSERT INTO {name} ({fields}) VALUES (%s)").format(
    name = sql.Identifier("hoge"),
    fields = sql.SQL(',').join([
        sql.Identifier("user_name"),
        sql.Identifier("timestamp")
    ])
)
curs.execute(SQL,(user1,))

curs.execute("SELECT * FROM hoge")
curs.fetchall()

outputs
[('John Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203)),
('Mary Smith', datetime.datetime(2022, 4, 3, 8, 51, 40, 312203)),
('Tom Brown', datetime.datetime(2022, 4, 3, 14, 36, 58, 333255))]

まとめ

• パラメータの個数がわかっている場合は%sを用いれば良い。
  （ただし、クオーテーションの関係でErrorが出ることもある）
• パラメータの個数が可変な場合は、psycopg2.sqlを用いる。
• 何度もデータを挿入するならclassを作成して渡すのが良い。