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

この記事は PostgreSQL Advent Calendar 2012 の5日目です。

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のオブジェクトとが相互に変換できましたね。やったね!

Comments

comments powered by Disqus