postgres数据库中如统计N个数据库中users表数据总量

#!/bin/bash
#===============================================================================
# 脚本名称: count_db_users.sh
# 功能描述: 统计多个PostgreSQL数据库中用户表的用户总数
# 版本: 2.0
# 更新: 修复登录notice信息导致的解析错误
#===============================================================================

# 默认参数配置
DB_HOST="localhost"
DB_PORT="5432"
DB_USER="postgres"
DB_PASSWORD="xxx"
DATABASES=(test )
TABLE_NAME="users"
VERBOSE=0
OUTPUT_FILE=""
CONNECTION_TIMEOUT=10

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

#===============================================================================
# 函数: 显示帮助信息
#===============================================================================
show_help() {
    cat << EOF
${BLUE}PostgreSQL 多数据库用户统计工具 v2.0${NC}

用法: $0 [选项]

选项:
    -h HOST        PostgreSQL主机地址 (默认: localhost)
    -p PORT        PostgreSQL端口 (默认: 5432)
    -U USER        数据库用户名 (默认: postgres)
    -d "DB1 DB2"   要统计的数据库列表 (空格分隔，必需)
    -t TABLE       表名 (默认: users)
    -P PASSWORD    数据库密码 (不推荐，建议使用.pgpass)
    -v             详细模式
    -o FILE        输出结果到文件
    -T TIMEOUT     连接超时时间(秒) (默认: 10)
    --help         显示此帮助信息

示例:
    $0 -h localhost -p 5432 -U postgres -d "db1 db2 db3"
    $0 -d "prod_db test_db" -t user_accounts -v
    $0 -d "db1 db2" -o result.txt

EOF
}

#===============================================================================
# 函数: 打印消息
#===============================================================================
print_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
    if [[ -n "$2" ]]; then
        echo -e "${RED}详情: $2${NC}"
    fi
}

#===============================================================================
# 函数: 构建psql命令
#===============================================================================
build_psql_command() {
    local db="$1"
    local query="$2"
    local options="-h $DB_HOST -p $DB_PORT -U $DB_USER -d $db"
    
    # 添加静默选项，抑制通知和警告信息
    options="$options -q --no-psqlrc -v ON_ERROR_STOP=1"
    
    # 设置客户端消息级别（不显示NOTICE信息）
    options="$options -c 'SET client_min_messages = warning;'"
    
    # 超时设置
    options="$options -c statement_timeout = ${CONNECTION_TIMEOUT}s"
    
    # 添加查询
    options="$options -t -A -c \"$query\""
    
    echo "$options"
}

#===============================================================================
# 函数: 执行psql查询并返回纯数据
#===============================================================================
execute_query() {
    local db="$1"
    local query="$2"
    local result
    local exit_code
    
    # 使用临时文件避免子shell问题
    local tmp_file=$(mktemp 2>/dev/null || mktemp -t tmp 2>/dev/null)
    local tmp_error=$(mktemp 2>/dev/null || mktemp -t tmp 2>/dev/null)
    
    # 构建命令 - 重定向stderr到临时文件
    if [[ -z "$DB_PASSWORD" ]]; then
        psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$db" \
            -q --no-psqlrc \
            -v ON_ERROR_STOP=1 \
            -c "SET client_min_messages = warning;" \
            -c "SET statement_timeout = '${CONNECTION_TIMEOUT}s';" \
            -t -A -c "$query" \
            2>"$tmp_error" >"$tmp_file"
    else
        PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$db" \
            -q --no-psqlrc \
            -v ON_ERROR_STOP=1 \
            -c "SET client_min_messages = warning;" \
            -c "SET statement_timeout = '${CONNECTION_TIMEOUT}s';" \
            -t -A -c "$query" \
            2>"$tmp_error" >"$tmp_file"
    fi
    
    exit_code=$?
    
    # 读取输出
    result=$(cat "$tmp_file" 2>/dev/null | tr -d '\r' | sed '/^$/d')
    local error_msg=$(cat "$tmp_error" 2>/dev/null)
    
    # 清理临时文件
    rm -f "$tmp_file" "$tmp_error" 2>/dev/null
    
    # 检查是否执行成功
    if [[ $exit_code -ne 0 ]]; then
        # 过滤掉NOTICE信息（这些是警告，不是致命错误）
        if [[ -n "$error_msg" ]] && [[ ! "$error_msg" =~ ^NOTICE: ]]; then
            echo "ERROR|$error_msg"
            return 1
        fi
    fi
    
    # 如果没有结果，返回空值
    if [[ -z "$result" ]] || [[ "$result" =~ ^[[:space:]]*$ ]]; then
        echo "0"
        return 0
    fi
    
    echo "$result"
    return 0
}

#===============================================================================
# 函数: 检查数据库连接
#===============================================================================
check_connection() {
    local db="$1"
    local result
    result=$(execute_query "$db" "SELECT 1" 2>/dev/null)
    
    if [[ $? -eq 0 ]] && [[ "$result" == "1" ]]; then
        return 0
    else
        return 1
    fi
}

#===============================================================================
# 函数: 检查表是否存在
#===============================================================================
check_table_exists() {
    local db="$1"
    local table="$2"
    local result
    
    # 转义表名中的特殊字符
    local safe_table=$(echo "$table" | sed "s/'/''/g")
    
    result=$(execute_query "$db" \
        "SELECT COUNT(*) FROM information_schema.tables \
         WHERE table_schema NOT IN ('information_schema', 'pg_catalog') \
         AND table_name = '$safe_table';" 2>/dev/null)
    
    # 检查结果是否为数字且大于0
    if [[ "$result" =~ ^[0-9]+$ ]] && [[ "$result" -gt 0 ]]; then
        return 0
    else
        return 1
    fi
}

#===============================================================================
# 函数: 统计单个数据库的用户数
#===============================================================================
count_users_in_db() {
    local db="$1"
    local table="$2"
    local count=0
    local status="成功"
    local error_msg=""
    
    # 转义表名
    local safe_table=$(echo "$table" | sed 's/"/""/g')
    
    # 检查连接
    if ! check_connection "$db"; then
        status="失败"
        error_msg="无法连接到数据库"
        echo "0|$status|$error_msg"
        return 1
    fi
    
    # 检查表是否存在
    if ! check_table_exists "$db" "$table"; then
        status="失败"
        error_msg="表 '$table' 不存在"
        echo "0|$status|$error_msg"
        return 1
    fi
    
    # 统计用户数 - 使用双引号包裹表名以支持大小写敏感
    local query="SELECT COUNT(*) FROM \"$safe_table\";"
    local result
    result=$(execute_query "$db" "$query")
    
    # 检查是否返回错误
    if [[ "$result" =~ ^ERROR\| ]]; then
        status="失败"
        error_msg="${result#ERROR|}"
        echo "0|$status|$error_msg"
        return 1
    fi
    
    # 提取数字
    if [[ "$result" =~ ^[0-9]+$ ]]; then
        count="$result"
    else
        status="失败"
        error_msg="返回非数字结果: $result"
        echo "0|$status|$error_msg"
        return 1
    fi
    
    echo "$count|$status|"
    return 0
}

#===============================================================================
# 函数: 解析命令行参数
#===============================================================================
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h)
                DB_HOST="$2"
                shift 2
                ;;
            -p)
                DB_PORT="$2"
                shift 2
                ;;
            -U)
                DB_USER="$2"
                shift 2
                ;;
            -d)
                shift
                while [[ $# -gt 0 && ! "$1" =~ ^- ]]; do
                    DATABASES+=("$1")
                    shift
                done
                ;;
            -t)
                TABLE_NAME="$2"
                shift 2
                ;;
            -P)
                DB_PASSWORD="$2"
                shift 2
                ;;
            -T)
                CONNECTION_TIMEOUT="$2"
                shift 2
                ;;
            -v)
                VERBOSE=1
                shift
                ;;
            -o)
                OUTPUT_FILE="$2"
                shift 2
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                print_error "未知参数: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    if [[ ${#DATABASES[@]} -eq 0 ]]; then
        print_error "未指定数据库列表！请使用 -d 参数指定至少一个数据库。"
        show_help
        exit 1
    fi
}

#===============================================================================
# 函数: 输出格式化结果
#===============================================================================
output_result() {
    local total=0
    local db_count=${#DATABASES[@]}
    local success_count=0
    local result_lines=()
    
    # 构建表头
    local header="| 序号 | 数据库名 | 用户数量 | 状态 | 备注 |"
    local separator="|------|----------|----------|------|------|"
    
    if [[ -n "$OUTPUT_FILE" ]]; then
        {
            echo "PostgreSQL 多数据库用户统计报告"
            echo "=================================="
            echo "统计时间: $(date '+%Y-%m-%d %H:%M:%S')"
            echo "主机: $DB_HOST:$DB_PORT"
            echo "数据库用户: $DB_USER"
            echo "表名: $TABLE_NAME"
            echo "数据库列表: ${DATABASES[*]}"
            echo ""
            echo "$header"
            echo "$separator"
        } >> "$OUTPUT_FILE"
    else
        print_info "开始统计 ${#DATABASES[@]} 个数据库..."
        echo ""
    fi
    
    # 遍历统计每个数据库
    local index=1
    for db in "${DATABASES[@]}"; do
        if [[ "$VERBOSE" -eq 1 ]]; then
            echo -ne "${BLUE}正在统计数据库 '$db'...${NC} "
        fi
        
        local result
        result=$(count_users_in_db "$db" "$TABLE_NAME")
        local count=$(echo "$result" | cut -d'|' -f1)
        local status=$(echo "$result" | cut -d'|' -f2)
        local error_msg=$(echo "$result" | cut -d'|' -f3-)
        
        # 清理可能的多余字符
        count=$(echo "$count" | tr -d '[:space:]')
        status=$(echo "$status" | tr -d '[:space:]')
        error_msg=$(echo "$error_msg" | head -c 50)  # 截断过长的错误信息
        
        # 构建输出行
        local line="| $index | $db | $count | $status | $error_msg |"
        result_lines+=("$line")
        
        # 累积总数
        if [[ "$status" == "成功" ]]; then
            total=$((total + count))
            success_count=$((success_count + 1))
            if [[ "$VERBOSE" -eq 1 ]]; then
                echo -e "${GREEN}✓ 用户数: $count${NC}"
            fi
        else
            if [[ "$VERBOSE" -eq 1 ]]; then
                echo -e "${RED}✗ 失败: $error_msg${NC}"
            fi
        fi
        
        if [[ -n "$OUTPUT_FILE" ]]; then
            echo "$line" >> "$OUTPUT_FILE"
        fi
        
        index=$((index + 1))
    done
    
    # 显示汇总结果
    echo ""
    if [[ -n "$OUTPUT_FILE" ]]; then
        {
            echo "$separator"
            echo ""
            echo "统计汇总"
            echo "--------"
            echo "总数据库数: $db_count"
            echo "成功统计: $success_count"
            echo "失败: $((db_count - success_count))"
            echo "用户总数: $total"
            echo ""
            echo "报告已保存至: $OUTPUT_FILE"
        } >> "$OUTPUT_FILE"
        print_info "报告已保存至: $OUTPUT_FILE"
    fi
    
    # 控制台输出
    echo "$header"
    echo "$separator"
    for line in "${result_lines[@]}"; do
        echo "$line"
    done
    echo "$separator"
    echo ""
    echo "${BLUE}统计汇总${NC}"
    echo "总数据库数: $db_count"
    echo "成功统计: $success_count"
    echo "失败: $((db_count - success_count))"
    echo -e "${GREEN}用户总数: $total${NC}"
}

#===============================================================================
# 主程序
#===============================================================================
main() {
    # 解析参数
    parse_arguments "$@"
    
    # 显示配置信息
    if [[ "$VERBOSE" -eq 1 ]]; then
        print_info "配置信息"
        echo "  主机: $DB_HOST"
        echo "  端口: $DB_PORT"
        echo "  用户: $DB_USER"
        echo "  超时: ${CONNECTION_TIMEOUT}s"
        echo "  数据库: ${DATABASES[*]}"
        echo "  表名: $TABLE_NAME"
        echo ""
    fi
    
    # 执行统计
    output_result
}

# 调用主程序
main "$@"