PostgreSQLから並列COPYするスクリプト

先にPythonで書き殴ったやつをブログに書いたけれど、もうちょっと使いやすいようにならんかと、コマンドっぽく書き直したので、githubに置いた。

要はpsql -c ‘\COPY’を並列に実行するだけなんだけど、毎回のようにシェルであれこれやったりするのがいい加減面倒なので。numericなカラム（例：bigserialのid）があってインデックスが張られてる場合に、それを利用して並列にダウンロードするためのツール、と言えば良いのかな。

pg_dumpすると5時間ぐらいかかるところを18分で終わらせられたので、まあ良し。

$ ./pg_parcopy.py --count 32 -t users --host riosourcesingle.postgres.database.azure.com -U rifujita@riosourcesingle --dbname postgres
['id']
[?] Which columns do you use to split dump files? Indexed columns are ['id']: id
 > id

2021-11-25 09:03:13.777056: started.
2021-11-25 09:21:21.252234: finished.

使い方はhelpで分かるかと。

$ ./pg_parcopy.py --help
usage: pg_parcopy.py [-h] [--dbname DBNAME] --table TABLE [--count COUNT] [--directory DIRECTORY]
                     [--size SIZE] [--format {CSV,TEXT,BINARY}] [--host HOST] [--port PORT]
                     [--username USERNAME] [--password PASSWORD]

COPY table in parallel using psql command

optional arguments:
  -h, --help            show this help message and exit

General options:
  --dbname DBNAME, -d DBNAME
                        Database to be dumped (default: root)
  --table TABLE, -t TABLE
                        Table to be dumped
  --count COUNT, -c COUNT
                        Number of parallelized processes (should set to number of CPUs that
                        PostgreSQL server is using. default: 16)

Output options:
  --directory DIRECTORY
                        path of directory to save dump files
  --size SIZE, -s SIZE  target size of each dump files in MB
  --format {CSV,TEXT,BINARY}, -f {CSV,TEXT,BINARY}
                        table output format

Connection options:
  --host HOST           database server host (default: "localhost")
  --port PORT, -p PORT  database server port (default: "5432")
  --username USERNAME, -U USERNAME
                        database user name (default: "root")
  --password PASSWORD, -W PASSWORD
                        force password prompt (should happen automatically, default: $PGPASSWORD)

追記：macOS 12 / RHEL8 / Ubuntu 20.04で動くように書いてあるけど、他はPythonのバージョンが古かったりするので怪しい。

関連