【Python】PostgreSQLを操作する② -SELECT結果を取得する

2022年4月27日
2023年3月4日
DB, Python
DB, postgresql, Python

前回はpsycopg2を使ってPythonからPostgreSQLに接続する方法を紹介しました。

今回はさらにSELECTのSQLを実行し、取得した結果の扱い方について触れていきたいと思います。本記事ではSQLでお望みのデータを取得する方法についてはスコープに入れていません。SELECT文で取得した結果をPythonでどう扱うかにフォーカスを当てていこうと思います。

1 接続
2 SELECT文実行
3 実行結果取得
4 データフレームに変換
- 4.1 列名を取得する
5 DBについてより学ぶ
6 まとめ

接続

前回詳しく紹介しましたが、まずはPostgreSQLに接続します。

import psycopg2

#接続
connector =  psycopg2.connect('postgresql://{user}:{password}@{host}:{port}/{dbname}'.format( 
            user="postgres",        #ユーザ
            password="postgresql",  #パスワード
            host="localhost",       #ホスト名
            port="5432",            #ポート
            dbname="dvdrental"))    #データベース名

#カーソル取得
cursor = connector.cursor()

SELECT文実行

まずはクエリでデータを取得します。今回はdvdrentalというサンプルDBのcategoryテーブルから全件取得するSQLを実行します。

DBへの接続はpsycopg2.connect()の時点で行っています。

cursor.execute('SELECT * FROM category')

実行結果取得

ここからはSQLを実行した結果を取得していきます。

1件だけ取得

SQLではSELECT ＊で全件取得していますが、その後cursor.fetchone()とすることで、その中の1件だけをプログラム的に取得することができます。

このとき結果を保持するrowは取得した列数分のタプルです（今回のケースでは３）。

row = cursor.fetchone() 
print(row)

>>　(1, 'Action', datetime.datetime(2006, 2, 15, 9, 46, 27))

全件まとめて取得

cursor.fetchall()とすることで、全件まとめて取得できます。

この時rowには、配列として1行分のタプルが件数分保持されます。

row = cursor.fetchall() 
print(row)

>> [(1, 'Action', datetime.datetime(2006, 2, 15, 9, 46, 27)),
   ...,
   (16, 'Travel', datetime.datetime(2006, 2, 15, 9, 46, 27))]

cursor.fetchall()した後1件ずつ取得するにはfor文で処理します。

for row in cursor:
    print(row)

辞書型で取得

psycopg2.extra.DictCursorをカーソル取得時に渡すことで、実行結果を辞書で取得することができます。

from psycopg2.extras import DictCursor

cursor = connector.cursor(cursor_factory=DictCursor)
row = cursor.fetchone() 
print(type(row))

>> psycopg2.extras.DictRow

データフレームに変換

Pythonでデータを扱う以上、Pandasを使いたいケースも多いと思います。cursor.fetchall()の結果はそのままデータフレームに変換可能です。

import pandas as pd

cursor.execute('SELECT * FROM category')
df = pd.DataFrame(cursor.fetchall())

列名を取得する

テーブルの列名を取得して、データフレームに反映していきます。

列名はcursorのdescriptionから取得できます。

import pandas as pd

cursor.execute('SELECT * FROM category')

cols = [col.name for col in cursor.description]
df = pd.DataFrame(cursor.fetchall(),columns = cols)