背景
psycopg2を使ってherokuのPostgreSQLを使おうとしたところ、パラメータの渡し方が思ったよりも面倒だった。自分用のメモを兼ねて作成。
環境
OS : macOS Big Sur 11.6.5
python : Python 3.9.12
psycopg2-binary : 2.9.3
使うパッケージ群
psycopg2を先にダウンロードしておく。
pip install psycopg2-binary
最低限のパッケージをインポート。
(sqlはパラメータを安全にクエリに渡すために必要。)
from psycopg2 import (
connect, sql
)
おまじない
データベースのURLは、herokuでは "Settings > Config Vars"の中に入っている。
ここではauto commitを有効にしておく。
(デフォルトは無効。その場合は都度コミットする必要がある。)
DATABASE_URL = "hogehoge"
# データベースに接続
conn = connect(DATABASE_URL)
# auto commit を有効にする
conn.autocommit = True
# カーソルの設定
curs = conn.cursor()
クエリのベタ書き
テーブルの作成
今回は"hoge"という名前のテーブルを作り、ユーザー名(string型)とタイムスタンプ(datetime型)を格納する。
まずはテーブルを作る。
# 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文を使って行を追加する。
# タイムスタンプを作る
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の第二引数にパラメータのタプルを指定する。
que = "INSERT INTO hoge (user_name, timestamp) VALUES (%s, %s)"
curs.execute(que,("John Smith",ts))
このようにすると、データを挿入することができた。
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
を使って書くと
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に対して脆弱なので推奨されない。
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する。
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として確認したい場合は次のようにする。
print(SQL.as_string(conn))
outputs
SELECT * FROM "hoge"
個数のわからないパラメータを渡す
パラメータが全てstring型の場合は非常に簡単。
例えばhoge
というテーブルからuser_name
とtimestamp
を取り出したいときは次のように
、sql.Identifier(parameter)
とsql.SQL(',')
を結合させれば良い。
結合させるためには、お馴染みの区切り文字.join(list)
を使う。
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型の引数しか取らないからだ。
sql.Identifier(ts)
outputs
TypeError: SQL identifier parts must be strings
この場合、sql.Literal
を用いることで解決できる。
sql.Literal(ts)
outputs
Literal(datetime.datetime(2022, 4, 3, 8, 51, 40, 312203))
sql.Literal
を用いてhoge
にデータを挿入する。
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
を使って
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を定義する。
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)
そのためにいくつかのモジュールと関数が必要になる。
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)
User
classがパラメータとして渡されると、user_adapter
が呼び出されるようになった。
これでデータを挿入することができる。
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を作成して渡すのが良い。