270 lines
8.8 KiB
Bash
270 lines
8.8 KiB
Bash
|
||
#!/bin/bash
|
||
|
||
# 定义颜色输出
|
||
RED='\033[0;31m'
|
||
GREEN='\033[0;32m'
|
||
YELLOW='\033[1;33m'
|
||
BLUE='\033[0;34m'
|
||
NC='\033[0m' # No Color
|
||
|
||
# 设置脚本所在目录为工作目录
|
||
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||
SQL_DIR="${SCRIPT_DIR}/sql"
|
||
|
||
# 打印带时间戳的日志
|
||
log() {
|
||
local level=$1
|
||
local message=$2
|
||
local color=$NC
|
||
|
||
case $level in
|
||
"INFO") color=$BLUE;;
|
||
"SUCCESS") color=$GREEN;;
|
||
"WARN") color=$YELLOW;;
|
||
"ERROR") color=$RED;;
|
||
esac
|
||
|
||
echo -e "[$(date '+%Y-%m-%d %H:%M:%S')] ${color}${level}${NC}: ${message}"
|
||
}
|
||
|
||
# 数据库连接信息（可通过环境变量覆盖）
|
||
DB_HOST=${POSTGRES_HOST:-"localhost"}
|
||
DB_PORT=${POSTGRES_PORT:-"5432"}
|
||
DB_NAME=${POSTGRES_DB:-"urban-lifeline"}
|
||
DB_USER=${POSTGRES_USER:-"postgres"}
|
||
DB_PASSWORD=${POSTGRES_PASSWORD:-"postgres"}
|
||
|
||
# 设置 PSQL 环境变量以支持中文
|
||
export PGCLIENTENCODING=UTF8
|
||
|
||
# 检查psql命令是否可用
|
||
check_psql() {
|
||
if ! command -v psql &> /dev/null; then
|
||
echo -e "${RED}Error: psql command not found. Please install PostgreSQL client.${NC}"
|
||
exit 1
|
||
fi
|
||
}
|
||
|
||
# 检查并创建数据库用户
|
||
check_and_create_user() {
|
||
local new_user=$1
|
||
local new_password=$2
|
||
|
||
# 使用 postgres 用户执行
|
||
if sudo -u postgres psql -c "SELECT 1 FROM pg_roles WHERE rolname = '$new_user'" | grep -q 1; then
|
||
echo -e "${GREEN}User $new_user already exists${NC}"
|
||
else
|
||
echo -e "${YELLOW}Creating user $new_user...${NC}"
|
||
sudo -u postgres psql -c "CREATE USER $new_user WITH PASSWORD '$new_password' CREATEDB;"
|
||
if [ $? -eq 0 ]; then
|
||
echo -e "${GREEN}User $new_user created successfully${NC}"
|
||
else
|
||
echo -e "${RED}Failed to create user $new_user${NC}"
|
||
exit 1
|
||
fi
|
||
fi
|
||
}
|
||
|
||
# 检查数据库连接
|
||
check_db_connection() {
|
||
# 首先尝试以当前用户身份连接
|
||
if ! psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c '\q' &> /dev/null; then
|
||
echo -e "${YELLOW}Could not connect with current settings, attempting to create user...${NC}"
|
||
# 创建用户并设置权限
|
||
check_and_create_user "$DB_USER" "$DB_PASSWORD"
|
||
|
||
# 再次检查连接
|
||
if ! psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c '\q' &> /dev/null; then
|
||
echo -e "${RED}Error: Could not connect to PostgreSQL server.${NC}"
|
||
echo "Please check your connection settings:"
|
||
echo "Host: $DB_HOST"
|
||
echo "Port: $DB_PORT"
|
||
echo "User: $DB_USER"
|
||
exit 1
|
||
fi
|
||
fi
|
||
}
|
||
|
||
# 执行SQL文件
|
||
execute_sql_file() {
|
||
local sql_file=$1
|
||
if [ ! -f "$sql_file" ]; then
|
||
echo -e "${RED}Error: SQL file not found: $sql_file${NC}"
|
||
return 1
|
||
fi
|
||
|
||
echo -e "${YELLOW}Executing SQL file: $sql_file${NC}"
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$sql_file"
|
||
local status=$?
|
||
if [ $status -eq 0 ]; then
|
||
echo -e "${GREEN}Successfully executed: $sql_file${NC}"
|
||
else
|
||
echo -e "${RED}Failed to execute: $sql_file${NC}"
|
||
return $status
|
||
fi
|
||
}
|
||
|
||
# 初始化数据库
|
||
init() {
|
||
echo -e "${YELLOW}Initializing database...${NC}"
|
||
|
||
# 执行完整的初始化脚本
|
||
log "INFO" "Executing initialization script..."
|
||
# Run from inside the SQL_DIR so relative \i includes in initAll.sql (like createDB.sql)
|
||
# resolve relative to the SQL directory.
|
||
(
|
||
if [ ! -d "$SQL_DIR" ]; then
|
||
echo -e "${RED}Error: SQL directory not found: $SQL_DIR${NC}"
|
||
exit 1
|
||
fi
|
||
cd "$SQL_DIR" || exit 1
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -v ON_ERROR_STOP=1 -f "initAll.sql"
|
||
)
|
||
|
||
if [ $? -eq 0 ]; then
|
||
log "SUCCESS" "Database initialization completed successfully"
|
||
else
|
||
log "ERROR" "Database initialization failed"
|
||
return 1
|
||
fi
|
||
|
||
if [ $? -ne 0 ]; then
|
||
echo -e "${RED}Failed to create database.${NC}"
|
||
return 1
|
||
fi
|
||
|
||
# 2. 创建扩展和设置搜索路径
|
||
echo -e "${YELLOW}Creating extensions...${NC}"
|
||
check_extensions_availability() {
|
||
# 检查服务器上是否存在需创建的扩展
|
||
local missing=()
|
||
local exts=("uuid-ossp" "pg_trgm" "btree_gist")
|
||
for ext in "${exts[@]}"; do
|
||
# 查询 pg_available_extensions 来判断扩展是否已安装到服务器目录
|
||
if ! PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -tAc "SELECT 1 FROM pg_available_extensions WHERE name = '$ext';" | grep -q 1; then
|
||
missing+=("$ext")
|
||
fi
|
||
done
|
||
|
||
if [ ${#missing[@]} -ne 0 ]; then
|
||
echo -e "${RED}Error: The following server-side extensions are not available: ${missing[*]}${NC}"
|
||
echo "If you compiled PostgreSQL from source, you need to build and install the contrib modules into the server's installation prefix. Example steps:"
|
||
echo " # 在 PostgreSQL 源码目录下运行："
|
||
echo " cd /path/to/postgresql-source/contrib"
|
||
echo " make"
|
||
echo " sudo make install"
|
||
echo "或者只安装缺失的模块（例如 uuid-ossp）："
|
||
echo " cd /path/to/postgresql-source/contrib/uuid-ossp"
|
||
echo " make"
|
||
echo " sudo make install"
|
||
echo "安装完成后，重启 PostgreSQL 服务并重新运行此脚本："
|
||
echo " sudo systemctl restart postgresql"
|
||
echo "如果你使用的是容器或自定义路径，请确保将编译安装的扩展安装到 PostgreSQL 的 \$(pg_config --sharedir)/extension 目录下。"
|
||
return 1
|
||
fi
|
||
return 0
|
||
}
|
||
|
||
# 检查扩展可用性，若缺失则给出建议并退出
|
||
if ! check_extensions_availability; then
|
||
return 1
|
||
fi
|
||
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
|
||
CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";
|
||
CREATE EXTENSION IF NOT EXISTS \"pg_trgm\";
|
||
CREATE EXTENSION IF NOT EXISTS \"btree_gist\";"
|
||
|
||
if [ $? -ne 0 ]; then
|
||
echo -e "${RED}Failed to create extensions.${NC}"
|
||
return 1
|
||
fi
|
||
|
||
# 3. 逐个执行初始化SQL文件
|
||
echo -e "${YELLOW}Initializing tables...${NC}"
|
||
while IFS= read -r line || [[ -n "$line" ]]; do
|
||
# 跳过注释和空行
|
||
[[ $line =~ ^--.*$ ]] && continue
|
||
[[ -z "${line// }" ]] && continue
|
||
|
||
# 从 \i 命令中提取文件名
|
||
if [[ $line =~ \\i[[:space:]]+([^[:space:]]+) ]]; then
|
||
sql_file="${SQL_DIR}/${BASH_REMATCH[1]}"
|
||
if [[ $sql_file != *"createDB.sql"* ]]; then # 跳过createDB.sql
|
||
echo -e "${YELLOW}Executing: $sql_file${NC}"
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$sql_file"
|
||
if [ $? -ne 0 ]; then
|
||
echo -e "${RED}Failed to execute: $sql_file${NC}"
|
||
return 1
|
||
fi
|
||
fi
|
||
fi
|
||
done < "${SQL_DIR}/initAll.sql"
|
||
|
||
# 4. 设置搜索路径
|
||
echo -e "${YELLOW}Setting search path...${NC}"
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
|
||
ALTER DATABASE $DB_NAME SET search_path TO sys, public;"
|
||
|
||
echo -e "${GREEN}Database initialization completed successfully.${NC}"
|
||
return 0
|
||
}
|
||
|
||
# 重新初始化数据库
|
||
reinit() {
|
||
echo -e "${YELLOW}Reinitializing database...${NC}"
|
||
delete
|
||
init
|
||
}
|
||
|
||
# 删除数据库
|
||
delete() {
|
||
echo -e "${YELLOW}Deleting database...${NC}"
|
||
|
||
# 确保没有活动连接
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c "
|
||
SELECT pg_terminate_backend(pg_stat_activity.pid)
|
||
FROM pg_stat_activity
|
||
WHERE pg_stat_activity.datname = '$DB_NAME'
|
||
AND pid <> pg_backend_pid();"
|
||
|
||
# 删除数据库
|
||
PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c "DROP DATABASE IF EXISTS $DB_NAME;"
|
||
|
||
echo -e "${GREEN}Database deleted.${NC}"
|
||
}
|
||
|
||
# 显示帮助信息
|
||
show_help() {
|
||
echo "Usage: $0 {init|reinit|delete}"
|
||
echo "Commands:"
|
||
echo " init Initialize the database"
|
||
echo " reinit Reinitialize the database (delete and create)"
|
||
echo " delete Delete the database"
|
||
}
|
||
|
||
# 主函数
|
||
main() {
|
||
check_psql
|
||
check_db_connection
|
||
|
||
case "$1" in
|
||
"init")
|
||
init
|
||
;;
|
||
"reinit")
|
||
reinit
|
||
;;
|
||
"delete")
|
||
delete
|
||
;;
|
||
*)
|
||
show_help
|
||
exit 1
|
||
;;
|
||
esac
|
||
}
|
||
|
||
# Call main with all passed arguments so the script runs when invoked
|
||
main "$@" |