J’hopper’s Heaven (ジョッパーズヘブン)

ジョブホッパー・ボブがジョブホッパーでなくなる日を夢見るブログ

リストファイルから値を読み込んでその値をWHERE句で指定してPostgreSQLをSELECTしちゃうPythonコード

thought

おはよう!そして会えないときのためにこんにちは、こんばんは、そして、おやすみなさい。ジョブホッパー・ボブです。

あるカラムで特定の値をもつエントリーの別のカラムをSELECTしたい、ってときありますよね。えぇ、ボブにはそれはもうよくあります。(๑˃́ꇴ˂̀๑)

ひとつやふたつだったら、

$ psql -d <db_name>

で接続しちゃいなYO!(*δωδ*) そんで、

# SELECT column_a FROM table WHERE column_b = 'xxx'

でいーじゃない!ってなるんですけど、それがもし、5000個とかだったら、もはや思考停止。誰だそんな数検索するやつ、アタマオカシーンジャネーノ。寝よー寝よー、もう寝よぅ _(ˇωˇ」∠)_ スヤァ… ってなってしまいまつね。そうなる前に!!!今夜寝ちゃう前に!ナウい感じ(←死語)のPythonで、5000個でも何万個でもSELECTしちゃうコードをご紹介したいと思いますよ!(←もはや需要があるのは自分だけだけどな!)

リストファイルをつくる

スクリプトファイルの中にそんな大量の値を埋め込むのはよろしくないのでね、別のファイルで用意しましょうかね。。まぁたいてい、csvとか何かでどっかから依頼されてきた結果の5000個とかだったりするので、そこらへんは、適当にそのファイルを修正してだな。改行区切りで検索条件となる値を羅列していく感じで。

$ cat searchList.txt
59c07c1853b5dc000838e693
59c891378c51df00061261a9
59c8948e8c51df000a126396
:

PostgreSQL用ドライバをインストールする

ここ、大事。

$ pip install psycopg2

スクリプトを書く

おまじないとインポート
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import psycopg2
リストファイルを読み込む

うーん、このsplitのところは、改行コードに合わせる必要がありますけれども。Macはこれでいけた。

listFile = 'searchList.txt'
searchList = open(listFile)
searchData = searchList.read()
searchList.close()
searchLines = searchData.split('\r')
PostgreSQLに接続する
conn = psycopg2.connect("dbname=<db_name> host=localhost")
PostgreSQLにQuery

rowはリストなのだよ、うむ。( ・⊝・∞)

query = "SELECT column_a FROM table WHERE column_b = '%(searchId)s';"
cursor = conn.cursor()
cursor.execute(query % {'searchId':searchId})
row = cursor.fetchone()
row[0]

ミソは、%(searchId)s のところですな!!

結果をファイルに出力する

検索数がパないですからね、printだともう後が大変。そんなわけで、SELECTした結果は別ファイルに出力する感じにしちゃおぅ。

outFile = '<OUT FILE PATH>'
if os.path.exists(outFile):
  os.remove(outFile)
f = csv.writer(open(outFile, "wb+"))
f.writerow([
  'column_a',
  'column_b'
])

できあがったスクリプト

そんなわけで、こんなスクリプトができあがりましたょv(*´∀`*)v

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import psycopg2
import os
import csv

# CREATE OUT FILE AND SET HEADERS
outFile = '<OUT FILE PATH>'
if os.path.exists(outFile):
  os.remove(outFile)
f = csv.writer(open(outFile, "wb+"))
f.writerow([
  'column_a',
  'column_b'
])

# SET SEARCH FILE
listFile = 'searchList.txt'
searchList = open(listFile)
searchData = searchList.read()
searchList.close()
searchLines = searchData.split('\r')

# EXEC SEARCH
try:
  conn = psycopg2.connect("dbname=<db_name> host=localhost")
  for searchLine in searchLines:
    searchId = searchLine.strip()
    query = "SELECT column_a FROM table WHERE column_b = '%(searchId)s';"
    cursor = conn.cursor()
    cursor.execute(query % {'searchId':searchId})
    row = cursor.fetchone()
    cursor.close()
    f.writerow([
      searchId,
      row[0]
    ])

except Exception as e:
  print(e)

finally:
  conn.close()

まとめ

文字コードと改行コードには気をつけろ。(° ꈊ °)✧˖°オホッ!