SmartEDT/backend/database/check_db.py

"""数据库连通性检查脚本。

用途：
- 快速验证 PostgreSQL 是否可连接
- 尝试加载 TimescaleDB 扩展，并输出版本信息（若可用）

注意：脚本只做连通性验证，不会打印密码。
"""

import os
import sys

def check_database():
    """检查数据库连接与 TimescaleDB 扩展可用性。"""
    print("正在检查数据库连接...")
    
    # 连接参数：通过环境变量覆盖（不在输出中打印密码）
    user = os.getenv("PG_USER", "smartedt")
    password = os.getenv("PG_PASSWORD", "postgres")
    host = os.getenv("PG_HOST", "127.0.0.1")
    port = os.getenv("PG_PORT", "5432")
    dbname = "smartedt"

    try:
        import psycopg
        conn = psycopg.connect(
            dbname=dbname,
            user=user,
            password=password,
            host=host,
            port=port,
            autocommit=True,
        )
        server_version = getattr(conn.info, "server_version", None)
        print(f"✅ 成功连接到 PostgreSQL 数据库 '{dbname}' (用户: {user}, v{server_version})")
        
        cur = conn.cursor()
        
        # 检查 TimescaleDB 扩展是否已安装
        print("正在检查 TimescaleDB 扩展...")
        try:
            # 尝试创建扩展（如果不存在）
            cur.execute("CREATE EXTENSION IF NOT EXISTS timescaledb;")
            print("✅ TimescaleDB 扩展加载成功")
            
            # 检查版本
            cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'timescaledb';")
            version = cur.fetchone()
            if version:
                print(f"✅ TimescaleDB 版本: {version[0]}")
            else:
                print("❌ 未找到 TimescaleDB 版本信息")
                
        except Exception as e:
            print(f"❌ TimescaleDB 检查失败: {e}")
            
        cur.close()
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        print("\n可能有以下原因：")
        print("1. 数据库服务未启动 (请运行 'net start postgresql-x64-17')")
        print("2. 密码错误 (请设置环境变量 PG_PASSWORD)")
        print("3. 端口被占用或配置不同")
        return False

if __name__ == "__main__":
    """作为脚本运行时的入口。"""
    try:
        import psycopg  # noqa: F401
    except ImportError:
        print("缺少依赖 psycopg，请先安装后再运行该脚本。")
        raise

    check_database()