pyhaya’s diary

機械学習系の記事をメインで書きます

Python から PostgreSQL を操作する

Python はデータの解析で使われることが多く、様々な記事で解析の方法や結果の可視化手法が紹介されています。多くの記事では構造化データを扱う例として CSV を pandas を使って解析していますよね。しかし、扱うデータが非常に多い場合にはデータベースを叩いてほしいデータを取り出して解析するということも往々にして行われています。この記事ではデータベースとして有名な PostgreSQL を、Python から操作する方法について紹介したいと思います。

動作環境

データベースの準備

PostgreSQL の Ubuntu へのインストールは以下の記事を参考にしました。

wiki.debian.org

Debian用の記事ですが apt install を使うだけなので問題はありませんでした。無事インストール出来たか確認します。

psql -V
# psql (PostgreSQL) 12.4 (Ubuntu 12.4-0ubuntu0.20.04.1)

上のようにバージョンがちゃんと表示されていればOKです。次にデータベースを１つ作ります。

createdb mydb
psql mydb  # 接続確認
mydb=>\l  # 存在するデータベースのリストを表示

                                  List of databases
   Name    |  Owner   | Encoding |   Collate   |    Ctype    |   Access privileges
-----------+----------+----------+-------------+-------------+-----------------------
 mydb      | pyhaya   | UTF8     | ja_JP.UTF-8 | ja_JP.UTF-8 |
 postgres  | postgres | UTF8     | ja_JP.UTF-8 | ja_JP.UTF-8 |
 template0 | postgres | UTF8     | ja_JP.UTF-8 | ja_JP.UTF-8 | =c/postgres          +
           |          |          |             |             | postgres=CTc/postgres
 template1 | postgres | UTF8     | ja_JP.UTF-8 | ja_JP.UTF-8 | =c/postgres          +
           |          |          |             |             | postgres=CTc/postgres

テーブルは Python で作ることもできますが、1回しか実行しないものなので SQL でやっておきます。

CREATE TABLE repo (
  id INTEGER,
  node_id VARCHAR(100),
  name VARCHAR(100),
  full_name VARCHAR(100),
  owner VARCHAR(100),
  private BOOLEAN,
  PRIMARY KEY(id)
);

このテーブルは後に GitHub API を使ってリポジトリの情報を取ることを想定しています。APIの使い方は下のリンクを参考にしてください。
docs.github.com

この SQL を実行するには

mydb=>\i create_table.sql

とします。

Python の準備

データベースと通信するためにパッケージを入れます。

pip install psycopg2

Python から PostgreSQL を操作する

import psycopg2
import dotenv
import os
import requests
from requests.auth import HTTPBasicAuth


def get_repos():
    auth = HTTPBasicAuth(
        os.environ.get("GITHUB_USERNAME"), os.environ.get("GITHUB_TOKEN")
    )
    data = requests.get("https://api.github.com/users/(USER_NAME)/repos", auth=auth)
    data = data.json()

    return data[0]


def connect(user, dbname, password):
    return psycopg2.connect(
        " user=" + user + " dbname=" + dbname + " password=" + password
    )


if __name__ == "__main__":
    dotenv.load_dotenv()

    user = "pyhaya"
    dbname = "mydb"
    password = os.environ.get("PG_PASSWORD")

    data = get_repos()

    ID = data["id"]
    node_id = data["node_id"]
    full_name = data["full_name"]
    owner = data["owner"]["login"]
    private = data["private"]

    query = (
        "INSERT INTO repo (id, node_id, full_name, owner, private) VALUES"
        + f" ({ID}, '{node_id}', '{full_name}', '{owner}', {private});"
    )

    with connect(user, dbname, password) as conn:
        with conn.cursor() as cur:
            cur.execute(query)
            conn.commit()

データベースへの接続のために with を使っています。これを使わずに操作の終了後にコネクションを connect.close()のようにすることもできますが、withを使ったほうが安全なのでこちらを使うようにします。

ちゃんとデータベースに入ったか確認してみます。

mydb=> SELECT id, node_id, full_name, owner, private FROM repo;

    id     |             node_id              |          full_name          |     owner     | private
-----------+----------------------------------+-----------------------------+---------------+---------
 251618774 | MDEwOlJlcG9zaXRvcnkyNTE2MTg3NzQ= |    pyhaya/analysis_memo     |    pyhaya     | f
(1 row)

ちゃんと入っていることが確認できました。

SQL 第2版 ゼロからはじめるデータベース操作

SQL 第2版 ゼロからはじめるデータベース操作