【PostgreSQL】リストアの進捗状況を監視・表示するシェルスクリプト

概要

コマンドを用いて大きいdumpファイルをリストアする際、ちゃんと動いてるのか分からなくてつらいので作りました。

実行結果例：

[2024-11-06 20:33:14] Database dump_2: IN PROGRESS (現在のサイズ: 48 MB)
[2024-11-06 20:33:14] Database dump_1: IN PROGRESS (現在のサイズ: 48 MB)
[2024-11-06 20:33:18] Restore of dump_2 completed successfully
[2024-11-06 20:33:18] Restore of dump_1 completed successfully
[2024-11-06 20:33:18] Database dump_2 exists with size: 56 MB
[2024-11-06 20:33:18] Database dump_1 exists with size: 55 MB

項目	内容
取り扱う内容	• Linux環境でのDBリストアにおいて、サイズ表示により進捗を出力するシェルスクリプト
想定読者	• PostgreSQLでリストアを行う人
ゴール	• 数秒ごとにリストア済サイズが表示されるスクリプトにより、心の安寧を得る

構築環境

ディレクトリ構造

今回はDockerコンテナ上で動作確認を行いました。Linux上のPostgreSQLが動作する環境であれば、ディレクトリの差分さえ変えれば、ほぼ使いまわしが効くはず。

dbtest/
├── Dockerfile
├── docker-compose.yml
├── dump/ # ダンプファイル置き場
│   ├── dump_1.backup
│   └── dump_2.backup
└── init/
    └── init.sh # 今回扱うスクリプト

設定ファイル

コンテナ内に送り込んだinit.shがコンテナビルド時に実行されるイメージで作成しています。
docker-compose.yml

services:
  db:
    build:
      context: .
      dockerfile: Dockerfile
    environment:
      POSTGRES_USER: 'postgres'
      POSTGRES_PASSWORD: 'pgpass'
    ports:
      - 5440:5432
    volumes:
      - db_data:/var/lib/postgresql/data
      # DBリストアを行う初期化スクリプト
      - ./init/init.sh:/docker-entrypoint-initdb.d/init.sh

volumes:
  db_data:

Dockerfile

FROM postgres:10

# dumpファイルをまとめてコンテナ内にコピー　コピー先ディレクトリはもうちょい、ちゃんとしたパスにしても
COPY ./dump/* /var/lib/postgresql/

スクリプト本体

ユーザーが編集する部分は以下2点。

• リストアを並列で実行するかどうかを定めたフラグ（PARALLEL_RESTORE）
• 作成するDATABASE名と、対応するdumpファイル名（DATABASE_BACKUPS）
• リストア対象のDBdumpに拡張機能の利用が含まれている場合、拡張機能の追加（コメントでCREATE EXTENSIONしている箇所）

#!/bin/bash

# エラー発生時も無視して続ける
# リストアするdumpにpublicスキーマがあると、CREATE DATABASEした際に自動作成されるpublicスキーマと衝突してエラーになる
# 無視して続けてほしいので解除した
# set -euo pipefail

# =====================================
# 定数・変数定義
# =====================================
readonly POSTGRES_USER
readonly POSTGRES_PASSWORD

readonly POSTGRES_PORT="5432" # docker-compose.ymlで指定した右辺と合わせる
readonly MONITOR_USER="restore_monitor_user" # リストア状況を監視する使い捨てユーザー
readonly PARALLEL_RESTORE="true" # リストアを並列で行うかどうか

# パス設定
readonly BASE_DIR="/var/lib/postgresql"
readonly LOGFILE="${BASE_DIR}/restore.log"
readonly STATUS_FILE="${BASE_DIR}/restore_status.log"
readonly MONITOR_FILE="${BASE_DIR}/monitor.pid"

# DB設定
# データベース名やdumpファイル名に追加・変更がある場合はここを変更
declare -A DATABASE_BACKUPS=(
    # ["データベース名"]="バックアップファイルのパス"
    ["dump_1"]="${BASE_DIR}/dump_1.backup"
    ["dump_2"]="${BASE_DIR}/dump_2.backup"
)

# =====================================
# ユーティリティ関数
# =====================================

# タイムスタンプ付きのログメッセージを出力
# 引数:
#   $1 - ログメッセージ
log_message() {
    local timestamp
    timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[${timestamp}] $1" | tee -a "${STATUS_FILE}"
}

# エラーメッセージを出力して終了
# 引数:
#   $1 - エラーメッセージ
#   $2 - 終了コード（デフォルト: 1）
error_exit() {
    log_message "ERROR: $1"
    exit "${2:-1}"
}

# =====================================
# データベース操作関数
# =====================================
# データベースの存在確認と状態検証
# 引数:
#   $1 - データベース名
verify_database() {
    local db_name=$1
    local exists
    local size

    # 引数チェック
    if [ -z "$db_name" ]; then
        error_exit "Database name not provided"
    fi

    # PostgreSQLシステムテーブル（pg_database）から対象DBの存在確認
    # -t: タプルのみ出力（ヘッダーなし）, -A: 位置揃えなし（整形なし）, -c: SQLコマンド実行
    # 存在する場合は1、存在しない場合は0が代入される
    exists=$(psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -tAc \
        "SELECT 1 FROM pg_database WHERE datname='$db_name'" || echo "0")
    
    # データベースが存在する場合、接続可能か確認
    if [ "$exists" = "1" ]; then
        # データベースのサイズを取得・表示
        size=$(psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -tAc \
            "SELECT pg_size_pretty(pg_database_size('$db_name'))")
        log_message "Database ${db_name} exists with size: ${size}"
        
        # データベースに接続できるか確認
        # /dev/null: select 1の結果をデバイスの墓場へ破棄して画面出力させない
        # 2>&1: 標準エラー出力を標準出力にリダイレクト　要するにエラーが起きても画面に表示させない
        if psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -d "$db_name" \
            -c "SELECT 1" >/dev/null 2>&1; then
            log_message "Successfully connected to ${db_name}"
            return 0
        fi

        # 接続できない場合
        log_message "WARNING: Could not connect to database ${db_name}"
    fi
    
    # データベースが存在しないか接続できない場合
    log_message "WARNING: Database ${db_name} does not exist or is inaccessible"
    return 1
}

# データベースのリストア処理
# 引数:
#   $1 - データベース名
#   $2 - バックアップファイルのパス
restore_database() {
    local db_name=$1
    local backup_file=$2
    local progress_file="${BASE_DIR}/${db_name}_in_progress"
    local restored_file="${BASE_DIR}/${db_name}_restored"
    local log_file="${LOGFILE}.${db_name}"

    # バックアップファイルの事前チェック
    if [ ! -f "${backup_file}" ] || [ ! -r "${backup_file}" ]; then
        log_message "ERROR: Backup file ${backup_file} not found or not readable"
        return 1
    fi
    
    log_message "Starting restore of ${db_name} from ${backup_file}..."
    touch "${progress_file}"
    
    # リストア処理
    # --clean: データベースをリストアする前にDROP DATABASEを実行
    # --if-exists: データベースが存在しない場合は何もしない
    # -v: 詳細ログを出力
    if pg_restore \
        --clean \
        --if-exists \
        -v \
        -p "${POSTGRES_PORT}" \
        -d "${db_name}" \
        -U "${POSTGRES_USER}" \
        "${backup_file}" > "${log_file}" 2>&1; then
        
        log_message "Restore of ${db_name} completed successfully"
        rm -f "${progress_file}"
        touch "${restored_file}"
        
        # リストア後のデータベース状態を確認
        if verify_database "${db_name}"; then
            log_message "${db_name} restored and verified successfully"
            return 0
        fi
    else
        log_message "Restore of ${db_name} failed. Check ${log_file} for details"
        rm -f "${progress_file}"
    fi
    
    return 1
}

# =====================================
# モニタリング関数
# =====================================
# リストア進捗を監視
# 引数: なし
monitor_progress() {
    # 現在のプロセスIDをモニタリングPIDとしてファイルに保存
    echo $$ > "${MONITOR_FILE}"
    
    # モニタリングファイルが存在する間、リストア状況を監視
    while [ -f "${MONITOR_FILE}" ]; do
        sleep 5
        # モニタリングファイルが削除された場合はループを抜ける
        [ ! -f "${MONITOR_FILE}" ] && break
        
        log_message "Current restore status:"
        local all_completed=true
        
        # データベースごとのリストア状況を表示
        # 表示順は順不同 制御するためにはDATABASE_ORDERを別途定義しないといけないので省略
        for db in "${!DATABASE_BACKUPS[@]}"; do
            local status
            local restored_file="${BASE_DIR}/${db}_restored"
            local progress_file="${BASE_DIR}/${db}_in_progress"
            local log_file="${LOGFILE}.${db}"
            
            # 表示例：Database dump_1: IN PROGRESS (現在のサイズ: 77 MB)
            if [ -f "${restored_file}" ]; then
                status="COMPLETED"
            elif [ -f "${progress_file}" ]; then
                all_completed=false
                # 現在のDBサイズを人が読める形式で取得
                local db_size=$(psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -tAc \
                    "SELECT pg_size_pretty(pg_database_size('${db}'))" 2>/dev/null || echo "calculating...")
                status="IN PROGRESS (現在のサイズ: ${db_size})"
                
                # リストアログにエラーが含まれているかチェック
                if [ -f "${log_file}" ] && grep -q "ERROR:" "${log_file}"; then
                    status="${status} - WARNING: Errors detected, check ${log_file}"
                fi
            else
                status="PENDING"
                all_completed=false
            fi
            log_message "Database ${db}: ${status}"
        done

        # すべてのDBが完了していたら監視を終了
        if [ "${all_completed}" = "true" ]; then
            log_message "All database restores have completed or failed. Stopping monitoring."
            rm -f "${MONITOR_FILE}"
            break
        fi
    done
}

# =====================================
# クリーンアップ関数
# =====================================
cleanup() {
    rm -f "${MONITOR_FILE}"
    rm -f "${BASE_DIR}"/*_in_progress
    rm -f "${BASE_DIR}"/*_restored
    rm -f "${STATUS_FILE}"
}

# =====================================
# メイン処理
# =====================================
main() {
    local failed_restores=0
    
    # PostgreSQLの起動待ち
    until pg_isready -p "${POSTGRES_PORT}" -U "${POSTGRES_USER}"; do
        log_message "Waiting for PostgreSQL to start..."
        sleep 2
    done
    
    log_message "PostgreSQL is ready. Starting database creation and restore process..."
    
    # データベースの再作成（既存の接続を強制切断してからデータベースを削除・作成）
    for db in "${!DATABASE_BACKUPS[@]}"; do
        log_message "Recreating database: ${db}"
        # 既存の接続を強制終了
        psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -d postgres -c \
            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '${db}' AND pid <> pg_backend_pid();" > /dev/null 2>&1 || true
        # データベースを削除（存在する場合）
        psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -d postgres -c \
            "DROP DATABASE IF EXISTS ${db};" || true
        # データベースを作成
        psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -d postgres -c \
            "CREATE DATABASE ${db};"

        # dblink拡張機能のインストール（追加）
        psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -d "${db}" -c \
            "CREATE EXTENSION IF NOT EXISTS dblink;"
    done
    
    # 進捗監視用ユーザー作成
    if ! psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -tAc \
        "SELECT 1 FROM pg_roles WHERE rolname='${MONITOR_USER}'" | grep -q 1; then
        log_message "Creating monitor user: ${MONITOR_USER}"
        createuser -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" "${MONITOR_USER}" -s || \
            error_exit "Failed to create monitor user"
    else
        log_message "Monitor user ${MONITOR_USER} already exists"
    fi
    
    # モニタリング開始
    monitor_progress &
    monitor_pid=$!  # モニタリングプロセスのPIDを保存
    
    # データベースのリストア実行
    if [ "${PARALLEL_RESTORE}" = "true" ]; then
        log_message "Starting parallel database restore..."
        for db in "${!DATABASE_BACKUPS[@]}"; do
            restore_database "$db" "${DATABASE_BACKUPS[$db]}" &
        done
        
        # すべてのリストア処理の完了を待機
        wait
        
        # リストア結果の確認
        for db in "${!DATABASE_BACKUPS[@]}"; do
            if [ ! -f "${BASE_DIR}/${db}_restored" ]; then
                ((failed_restores++))
            fi
        done
    else
        log_message "Starting sequential database restore..."
        for db in "${!DATABASE_BACKUPS[@]}"; do
            if ! restore_database "$db" "${DATABASE_BACKUPS[$db]}"; then
                ((failed_restores++))
            fi
        done
    fi
    
    # モニタリング終了
    rm -f "${MONITOR_FILE}"    
    
    # 最終確認
    if [ $failed_restores -eq 0 ]; then
        log_message "All database restores completed successfully"
        psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" \
            -c "SELECT pg_reload_conf();" > /dev/null

        # 監視用ユーザーの削除
        if psql -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" -tAc \
            "SELECT 1 FROM pg_roles WHERE rolname='${MONITOR_USER}'" | grep -q 1; then
            log_message "Dropping monitor user: ${MONITOR_USER}"
            dropuser -U "${POSTGRES_USER}" -p "${POSTGRES_PORT}" "${MONITOR_USER}" || \
                log_message "WARNING: Failed to drop monitor user ${MONITOR_USER}"
        fi
    else
        log_message "WARNING: Some database restores failed (failed_restores=${failed_restores})"
        return 1
    fi
}

# =====================================
# スクリプトの実行
# =====================================
# トラップ設定（スクリプト終了時にcleanup関数を実行）
trap cleanup EXIT
main "$@"A

実行結果

[2024-11-06 20:33:14] Database dump_2: IN PROGRESS (現在のサイズ: 48 MB)
[2024-11-06 20:33:14] Database dump_1: IN PROGRESS (現在のサイズ: 48 MB)
[2024-11-06 20:33:18] Restore of dump_2 completed successfully
[2024-11-06 20:33:18] Restore of dump_1 completed successfully
[2024-11-06 20:33:18] Database dump_2 exists with size: 56 MB
[2024-11-06 20:33:18] Database dump_1 exists with size: 55 MB

処理イメージ