PostgreSQLをPythonからpsycopg2を使っていじる

PostgreSQLをPythonから扱うには、psygopg2を使うのが簡単です。

インストール
% pip install psycopg2
で簡単です。ただ、C関数を作るので、Python.hなどのdevライブラリやgccといったツールを事前にインストールしておく必要があります。

簡単な使い方
>>> import psycopg2
# コネクション作成
>>> conn = psycopg2.connect("dbname=test host=localhost user=postgres")
# カーソル作成
>>> cur = conn.cursor()
# SQLコマンド実行 (今回はテーブル作成)
>>> cur.execute("CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);")
# SQLコマンド実行 (プレースホルダー使用。エスケープも勝手にされる)
>>> cur.execute("INSERT INTO test (num, data) VALUES (%s, %s)", (100, "abc'def"))
# SQL結果を受け取る
>>> cur.execute("SELECT * FROM test;")
>>> cur.fetchone()
(1, 100, "abc'def")
# コミット
>>> conn.commit()
# クローズ
>>> cur.close()
>>> conn.close()
はい。connect、cursor、execute、fetchone、commit、close、だいたいどのDBIでも変わりません。簡単ですね。

なお、execute()への渡し方ですが、ちょっと癖がありますので気をつけて下さい。

>>> cur.execute("INSERT INTO foo VALUES (%s)", "bar") # だめ
>>> cur.execute("INSERT INTO foo VALUES (%s)", ("bar")) # だめ
>>> cur.execute("INSERT INTO foo VALUES (%s)", ("bar",)) # OK
>>> cur.execute("INSERT INTO foo VALUES (%s)", ["bar"]) # OK

# intとStringの型について
>>> cur.execute("INSERT INTO numbers VALUES (%d)", (42,)) # だめ
>>> cur.execute("INSERT INTO numbers VALUES (%s)", (42,)) # OK
もちろん、 + でつなげたり %やformat()を使ってはだめ、ぜったい！ですよ。

Pythonオブジェクトのやりとり
さて本題です。PostgreSQLは配列型をはじめとして豊富な型を最初から持っています。これをうまくPythonと連携させたいですね。

配列
>>> cur.execute("DROP TABLE IF EXISTS test")
# integer[] のテーブルを作る
>>> cur.execute("CREATE TABLE test (id serial PRIMARY KEY, num_array integer[]);")
>>> a = [20,3040]
# データ挿入
>>> cur.execute("INSERT INTO test (num_array) VALUES (%s)", [a])
>>> cur.execute("SELECT * FROM test;")
>>> b = cur.fetchone()
>>> b
(1, [20, 3040])
# リスト型で取得できた
>>> type(b[1])
<type 'list'>
配列の挿入と取得ができました。

時刻型
>>> import datetime

>>> cur.execute("DROP TABLE IF EXISTS test")
# timestamp型 のテーブルを作る
>>> cur.execute("CREATE TABLE test (id serial PRIMARY KEY, t timestamp);")
>>> dt = datetime.datetime.now()
>>> cur.execute("INSERT INTO test (t) VALUES (%s)", [dt])
>>> b = cur.fetchone()
>>> b
(1, datetime.datetime(2012, 12, 01, 01, 47, 28, 711776))
>>> type(b[1])
<type 'datetime.datetime'>
はい。できました。あんまり面白みがないですね。

辞書型
PostgreSQLといえば配列型の他にhstoreという辞書型があります。これを使ってみましょう。

事前にこちらなどを参考にして、hstoreをPostgreSQLに入れておく必要があります。

>>> cur.execute("DROP TABLE IF EXISTS test")
# hstore型 のテーブルを作る
>>> cur.execute("CREATE TABLE test (id serial PRIMARY KEY, data hstore);")
>>> a = {'you': 'me', 'spam': 'ham', 'hoge': None}
>>> cur.execute("INSERT INTO test (num_array) VALUES (%s)", [a])
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
psycopg2.ProgrammingError: can't adapt type 'dict'
おっと。さすがにそのままではだめなようです。実はcursorかconnectionにregister_hstoreという関数を使ってhstoreへの変換を事前に登録しておく必要があります。

# register_hstore用
>>> import psycopg2.extras
# cursorにhstoreを登録
>>> psycopg2.extras.register_hstore(cur)
# dictを作る
>>> a = {'you': 'me', 'spam': 'ham', 'hoge': None}
# INSERT
>>> cur.execute("INSERT INTO test (id serial PRIMARY KEY, data) VALUES (%s)", [a])
>>> cur.execute("SELECT * FROM test;")
>>> b = cur.fetchone()
>>> b
(1, {'you': 'me', 'spam': 'ham', 'hoge': None})
>>> type(b[1])
<type 'dict'>
はい。というわけで、辞書型を入れてさらに取り出すことができました。

ただ、残念なことにdictに数字は使えないようです。Noneはvalueとしては使えますが、Keyとしては使えません。

>>> a = {1: 122, 2: 22}
>>> cur.execute("INSERT INTO test (id serial PRIMARY KEY, data) VALUES (%s)", [a])
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
psycopg2.ProgrammingError: function hstore(integer[], integer[]) does not exist
LINE 1: INSERT INTO test (data) VALUES (hstore(ARRAY[1, 2], ARRAY[12...
^
HINT: No function matches the given name and argument types. You might need to add explicit type casts.
その他
inet型とか、9.2から入ったrange型やJSON型とかも扱えたりしますが、詳しくは割愛します。

独自のクラス
さて、ここまでは通常使える型に関して述べてきました。ここから、独自のクラスの場合について説明します。

Python -> PostgreSQL
PythonのオブジェクトからPostgreSQLに入れるためにはadaptなどの関数を使います。

まずクラスを作りましょう。(このBookクラスはあとでも使いますから覚えておいてください)

class Book(object):
def __init__(self, author, title, publish_date):
self.author = author
self.title = title
self.publish_date = publish_date
その後に、変換関数を定義します。

from psycopg2.extensions import adapt, register_adapter, AsIs

def adapt_book(book):
return AsIs("%s, %s, %s" % (adapt(book.author),
adapt(book.title),
adapt(book.publish_date)))
date型とかを突っ込んでもちゃんとやってくれます。

定義したところで、実際に使ってみましょう。この例では分かりやすいようにINSERTとかexecuteとかを直接実行していますが、Book型でラッパーを作るのは簡単なはずです。

# 変換関数を登録
>>> register_adapter(Book, adapt_book)
# テーブル作成
>>> cur.execute("DROP TABLE IF EXISTS books")
>>> cur.execute("CREATE TABLE books (id serial PRIMARY KEY, author text, title text, publish_date timestamp);")
# データ作成
>>> b = Book('shirou', 'TestBook', datetime.date.today())
# データINSERT
>>> cur.execute("INSERT INTO books (author, title, publish_date) VALUES (%s)", [b])
# コミット
>>> conn.commit()
>>> cur.close()
>>> conn.close()
さて、コミットしたのでpsqlで覗いてみましょう。

test=# select * from books;
id | author | title | publish_date
----+--------+----------+---------------------
1 | shirou | TestBook | 2012-12-01 00:00:00
(1 row)
ちゃんと入っていますね。これで、どんなクラスを作ってもおーけー！

PostgreSQL -> Python
PostgreSQLでは比較的簡単に独自の型を作ることができます。

-- book型を定義
test=# CREATE TYPE book AS (author text, title text, publish_date timestamp);
-- book型を使うテーブルを定義
test=# DROP TABLE IF EXISTS books;
test=# CREATE TABLE books (id serial PRIMARY KEY, a_book book);
-- book型を挿入
test=# INSERT INTO books (a_book) VALUES (('shirou', 'TestBook', '2012-12-01'));
INSERT 0 1
test=# SELECT * FROM books;
id | a_book
----+-----------------------------------------
1 | (shirou,TestBook,"2012-12-01 00:00:00")
(1 row)
では、このPostgreSQLのbook型をBookクラスに変換してみます。

まず、変換関数を定義します。

実はpython側の変換関数に渡されるときには単にtextになっています。具体的にはpsqlで見た値と同じです。

(shirou,TestBook,"2012-12-01 00:00:00")
変換関数はこの文字列を解釈して、オブジェクトを作って返せば良いです。

def cast_book(value, cur):
if value is None:
return None

import datetime
v = value[1:-1].split(",") # 両はじの()を除く
fmt = "%Y-%m-%d %H:%M:%S"
d = datetime.datetime.strptime(v[2][1:-1], fmt) # 両はじの"を除く
return Book(v[0], v[1], d)
変換に関しては、PostgreSQL側で入出力関数を定義するとちょっと変わってくるのかもしれません(未確認ですいません)。

さて、この関数をpsycopg2に登録します。登録にはOIDが必要なので取っておきます。

>>> cur.execute("SELECT NULL::book")
>>> book_oid = cur.description[0][1]
>>> book_oid
23396640
このoidを元に extensions.register_type() で登録します。

>>> BOOK = psycopg2.extensions.new_type((book_oid,), "BOOK", cast_book)
>>> psycopg2.extensions.register_type(BOOK)
>>> cur.execute("SELECT * FROM books;")
>>> b = cur.fetchone()
>>> b
(1, <__main__.Book object at 0x7f84c16a7a90>)
>>> type(b[1])
<class '__main__.Book'>
>>> b[1].title
'TestBook'
>>> b[1].author
'shirou'
>>> b[1].publish_date
datetime.datetime(2012, 12, 1, 0, 0)
>>> type(b[1].publish_date)
<type 'datetime.datetime'>
これでPostgreSQLとPythonのオブジェクトとが相互に変換できましたね。やったね！