water-management-system/tests/performance/db_query_stress.py

#!/usr/bin/env python3
"""
大数据量数据库查询压力测试

测试百万级记录下各类查询的性能：
- 简单查询（单表 WHERE）
- 聚合查询（GROUP BY / SUM / AVG）
- JOIN 查询（多表关联）
- GIS 空间查询（PostGIS ST_DWithin 等）
- 对比有无索引的性能差异

运行方式：
  python db_query_stress.py [--host HOST] [--port PORT] [--db DB] [--user USER] [--password PASS]

示例：
  python db_query_stress.py --host localhost --db water_management --generate-data 1000000
"""

import argparse
import json
import time
import random
import statistics
import sys
from datetime import datetime, timedelta
from contextlib import contextmanager

try:
    import psycopg2
    from psycopg2 import pool
except ImportError:
    print("请先安装 psycopg2: pip install psycopg2-binary")
    sys.exit(1)


# ==================== 配置 ====================

DEFAULT_HOST = "localhost"
DEFAULT_PORT = 5432
DEFAULT_DB = "water_management"
DEFAULT_USER = "postgres"
DEFAULT_PASSWORD = "postgres"

# 测试查询配置
QUERY_ROUNDS = 5  # 每个查询重复次数取平均值


# ==================== 数据库连接 ====================

class DBConnectionPool:
    """数据库连接池管理"""

    def __init__(self, host, port, db, user, password, min_conn=2, max_conn=10):
        self.pool = pool.ThreadedConnectionPool(
            min_conn, max_conn,
            host=host, port=port, dbname=db, user=user, password=password,
            connect_timeout=10,
        )

    @contextmanager
    def get_connection(self):
        conn = self.pool.getconn()
        try:
            yield conn
        finally:
            self.pool.putconn(conn)

    def close(self):
        self.pool.closeall()


# ==================== 测试数据生成 ====================

def generate_test_data(db_pool, num_records):
    """生成测试数据"""
    print(f"\n📦 生成 {num_records} 条测试数据...")

    with db_pool.get_connection() as conn:
        cur = conn.cursor()

        # 创建传感器数据表（如不存在）
        cur.execute("""
            CREATE TABLE IF NOT EXISTS perf_sensor_data (
                id BIGSERIAL PRIMARY KEY,
                sensor_id INTEGER NOT NULL,
                device_id VARCHAR(32) NOT NULL,
                timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                pressure REAL,
                flow REAL,
                temperature REAL,
                ph REAL,
                turbidity REAL,
                chlorine REAL,
                battery REAL,
                signal_strength INTEGER,
                lng DOUBLE PRECISION,
                lat DOUBLE PRECISION,
                area_id INTEGER,
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        """)

        # 创建告警表
        cur.execute("""
            CREATE TABLE IF NOT EXISTS perf_alarms (
                id BIGSERIAL PRIMARY KEY,
                sensor_id INTEGER NOT NULL,
                alarm_type VARCHAR(32) NOT NULL,
                severity VARCHAR(16) NOT NULL,
                value REAL,
                threshold REAL,
                description TEXT,
                acknowledged BOOLEAN DEFAULT FALSE,
                timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        """)

        # 创建设备表
        cur.execute("""
            CREATE TABLE IF NOT EXISTS perf_devices (
                id SERIAL PRIMARY KEY,
                device_id VARCHAR(32) UNIQUE NOT NULL,
                name VARCHAR(128),
                device_type VARCHAR(32),
                area_id INTEGER,
                lng DOUBLE PRECISION,
                lat DOUBLE PRECISION,
                status VARCHAR(16) DEFAULT 'active',
                install_date DATE,
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        """)

        conn.commit()

        # 清空旧测试数据
        print("  清空旧数据...")
        cur.execute("TRUNCATE TABLE perf_sensor_data, perf_alarms, perf_devices")
        conn.commit()

        # 生成设备数据
        print("  生成设备数据...")
        device_count = min(5000, num_records // 100)
        devices = []
        for i in range(1, device_count + 1):
            device_id = f"DEV{i:06d}"
            devices.append((
                device_id,
                f"传感器-{i}",
                random.choice(["pressure", "flow", "quality", "level", "valve"]),
                random.randint(1, 50),
                round(random.uniform(116.0, 117.0), 6),
                round(random.uniform(39.5, 40.5), 6),
                random.choice(["active", "inactive", "maintenance"]),
                (datetime.now() - timedelta(days=random.randint(0, 730))).date(),
            ))

        cur.executemany("""
            INSERT INTO perf_devices (device_id, name, device_type, area_id, lng, lat, status, install_date)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        """, devices)
        conn.commit()
        print(f"  ✅ 已生成 {device_count} 个设备")

        # 批量插入传感器数据
        print(f"  生成 {num_records} 条传感器数据...")
        batch_size = 10000
        total_inserted = 0
        base_time = datetime.now() - timedelta(days=365)

        alarm_types = ["PRESSURE_HIGH", "PRESSURE_LOW", "FLOW_ABNORMAL",
                       "QUALITY_WARNING", "LEAK_DETECTED", "BATTERY_LOW"]
        severities = ["INFO", "WARNING", "CRITICAL"]

        while total_inserted < num_records:
            batch = []
            for _ in range(min(batch_size, num_records - total_inserted)):
                sensor_id = random.randint(1, device_count)
                ts = base_time + timedelta(
                    seconds=random.randint(0, 365 * 24 * 3600)
                )
                batch.append((
                    sensor_id,
                    f"DEV{sensor_id:06d}",
                    ts,
                    round(random.uniform(0.1, 0.8), 3),  # pressure
                    round(random.uniform(10, 500), 2),    # flow
                    round(random.uniform(5, 35), 1),      # temperature
                    round(random.uniform(6.5, 8.5), 2),   # ph
                    round(random.uniform(0, 5), 2),       # turbidity
                    round(random.uniform(0.1, 0.8), 3),   # chlorine
                    round(random.uniform(20, 100), 1),    # battery
                    random.randint(-90, -30),             # signal
                    round(random.uniform(116.0, 117.0), 6),  # lng
                    round(random.uniform(39.5, 40.5), 6),    # lat
                    random.randint(1, 50),                # area_id
                ))

            cur.executemany("""
                INSERT INTO perf_sensor_data
                (sensor_id, device_id, timestamp, pressure, flow, temperature,
                 ph, turbidity, chlorine, battery, signal_strength, lng, lat, area_id)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
            """, batch)
            conn.commit()
            total_inserted += len(batch)
            if total_inserted % 100000 == 0:
                print(f"  进度: {total_inserted}/{num_records} ({total_inserted*100//num_records}%)")

        print(f"  ✅ 已生成 {total_inserted} 条传感器数据")

        # 生成告警数据（传感器数据的 5%）
        alarm_count = num_records // 20
        print(f"  生成 {alarm_count} 条告警数据...")
        alarms = []
        for _ in range(alarm_count):
            sensor_id = random.randint(1, device_count)
            alarms.append((
                sensor_id,
                random.choice(alarm_types),
                random.choice(severities),
                round(random.uniform(0, 100), 2),
                round(random.uniform(50, 100), 2),
                f"自动告警描述 - {random.choice(alarm_types)}",
                random.choice([True, False]),
                base_time + timedelta(seconds=random.randint(0, 365 * 24 * 3600)),
            ))

        cur.executemany("""
            INSERT INTO perf_alarms
            (sensor_id, alarm_type, severity, value, threshold, description, acknowledged, timestamp)
            VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
        """, alarms)
        conn.commit()
        print(f"  ✅ 已生成 {alarm_count} 条告警数据")

        cur.close()

    print("✅ 测试数据生成完成\n")


# ==================== 索引管理 ====================

def create_indexes(db_pool):
    """创建性能测试用索引"""
    with db_pool.get_connection() as conn:
        cur = conn.cursor()
        indexes = [
            "CREATE INDEX IF NOT EXISTS idx_sensor_data_sensor_id ON perf_sensor_data(sensor_id)",
            "CREATE INDEX IF NOT EXISTS idx_sensor_data_timestamp ON perf_sensor_data(timestamp)",
            "CREATE INDEX IF NOT EXISTS idx_sensor_data_area_id ON perf_sensor_data(area_id)",
            "CREATE INDEX IF NOT EXISTS idx_sensor_data_device_id ON perf_sensor_data(device_id)",
            "CREATE INDEX IF NOT EXISTS idx_sensor_data_composite ON perf_sensor_data(sensor_id, timestamp)",
            "CREATE INDEX IF NOT EXISTS idx_alarms_sensor_id ON perf_alarms(sensor_id)",
            "CREATE INDEX IF NOT EXISTS idx_alarms_timestamp ON perf_alarms(timestamp)",
            "CREATE INDEX IF NOT EXISTS idx_alarms_type ON perf_alarms(alarm_type)",
            "CREATE INDEX IF NOT EXISTS idx_devices_area_id ON perf_devices(area_id)",
            "CREATE INDEX IF NOT EXISTS idx_devices_type ON perf_devices(device_type)",
        ]
        for idx in indexes:
            cur.execute(idx)
        conn.commit()
        cur.close()
    print("✅ 索引已创建")


def drop_indexes(db_pool):
    """删除性能测试用索引"""
    with db_pool.get_connection() as conn:
        cur = conn.cursor()
        indexes = [
            "idx_sensor_data_sensor_id", "idx_sensor_data_timestamp",
            "idx_sensor_data_area_id", "idx_sensor_data_device_id",
            "idx_sensor_data_composite", "idx_alarms_sensor_id",
            "idx_alarms_timestamp", "idx_alarms_type",
            "idx_devices_area_id", "idx_devices_type",
        ]
        for idx in indexes:
            cur.execute(f"DROP INDEX IF EXISTS {idx}")
        conn.commit()
        cur.close()
    print("🗑️ 索引已删除")


# ==================== 查询测试 ====================

class QueryBenchmark:
    """查询性能基准测试"""

    def __init__(self, db_pool):
        self.db_pool = db_pool
        self.results = []

    def run_query(self, name, sql, params=None, rounds=QUERY_ROUNDS):
        """执行查询并统计耗时"""
        times = []
        row_count = 0

        for i in range(rounds):
            with self.db_pool.get_connection() as conn:
                cur = conn.cursor()
                start = time.time()
                try:
                    cur.execute(sql, params)
                    rows = cur.fetchall()
                    row_count = len(rows)
                except Exception as e:
                    print(f"    ⚠️ 查询错误: {e}")
                    conn.rollback()
                    row_count = 0
                    times.append(-1)
                    cur.close()
                    continue
                elapsed_ms = (time.time() - start) * 1000
                times.append(elapsed_ms)
                cur.close()

        valid_times = [t for t in times if t >= 0]
        if valid_times:
            result = {
                "name": name,
                "sql": sql[:100] + ("..." if len(sql) > 100 else ""),
                "rounds": rounds,
                "success_rounds": len(valid_times),
                "rows_returned": row_count,
                "min_ms": round(min(valid_times), 2),
                "max_ms": round(max(valid_times), 2),
                "avg_ms": round(statistics.mean(valid_times), 2),
                "median_ms": round(statistics.median(valid_times), 2),
            }
        else:
            result = {
                "name": name,
                "sql": sql[:100] + ("..." if len(sql) > 100 else ""),
                "rounds": rounds,
                "success_rounds": 0,
                "error": "all rounds failed",
            }

        self.results.append(result)
        status = "✅" if valid_times else "❌"
        avg = result.get("avg_ms", "N/A")
        print(f"  {status} {name}: avg={avg}ms, rows={row_count}")
        return result

    def run_all_benchmarks(self):
        """运行所有基准测试"""
        print("\n🔍 运行查询性能测试...\n")

        # 1. 简单查询
        print("  --- 简单查询 ---")
        self.run_query(
            "简单查询 - 按传感器ID查询最近24小时",
            """SELECT * FROM perf_sensor_data
               WHERE sensor_id = %s AND timestamp > NOW() - INTERVAL '24 hours'
               ORDER BY timestamp DESC LIMIT 100""",
            (random.randint(1, 5000),),
        )

        self.run_query(
            "简单查询 - 按区域ID查询",
            """SELECT * FROM perf_sensor_data
               WHERE area_id = %s
               ORDER BY timestamp DESC LIMIT 100""",
            (random.randint(1, 50),),
        )

        self.run_query(
            "简单查询 - 时间范围查询",
            """SELECT * FROM perf_sensor_data
               WHERE timestamp BETWEEN %s AND %s
               ORDER BY timestamp DESC LIMIT 1000""",
            (
                (datetime.now() - timedelta(days=7)).isoformat(),
                datetime.now().isoformat(),
            ),
        )

        # 2. 聚合查询
        print("\n  --- 聚合查询 ---")
        self.run_query(
            "聚合查询 - 每小时平均压力",
            """SELECT date_trunc('hour', timestamp) AS hour,
                      AVG(pressure) AS avg_pressure,
                      MAX(pressure) AS max_pressure,
                      MIN(pressure) AS min_pressure,
                      COUNT(*) AS readings
               FROM perf_sensor_data
               WHERE sensor_id = %s AND timestamp > NOW() - INTERVAL '7 days'
               GROUP BY hour ORDER BY hour""",
            (random.randint(1, 5000),),
        )

        self.run_query(
            "聚合查询 - 各区域统计",
            """SELECT area_id,
                      COUNT(*) AS total_readings,
                      AVG(pressure) AS avg_pressure,
                      AVG(flow) AS avg_flow,
                      AVG(temperature) AS avg_temp
               FROM perf_sensor_data
               WHERE timestamp > NOW() - INTERVAL '1 day'
               GROUP BY area_id ORDER BY total_readings DESC""",
        )

        self.run_query(
            "聚合查询 - 日统计报表",
            """SELECT date_trunc('day', timestamp) AS day,
                      COUNT(*) AS readings,
                      AVG(pressure) AS avg_pressure,
                      AVG(flow) AS avg_flow,
                      STDDEV(pressure) AS pressure_stddev
               FROM perf_sensor_data
               WHERE timestamp > NOW() - INTERVAL '30 days'
               GROUP BY day ORDER BY day""",
        )

        # 3. JOIN 查询
        print("\n  --- JOIN 查询 ---")
        self.run_query(
            "JOIN 查询 - 传感器数据 + 设备信息",
            """SELECT s.timestamp, s.pressure, s.flow, s.temperature,
                      d.name AS device_name, d.device_type, d.area_id, d.status
               FROM perf_sensor_data s
               JOIN perf_devices d ON s.device_id = d.device_id
               WHERE s.sensor_id = %s
               ORDER BY s.timestamp DESC LIMIT 100""",
            (random.randint(1, 5000),),
        )

        self.run_query(
            "JOIN 查询 - 告警 + 设备信息",
            """SELECT a.timestamp, a.alarm_type, a.severity, a.value, a.threshold,
                      d.name AS device_name, d.device_type, d.lng, d.lat
               FROM perf_alarms a
               JOIN perf_devices d ON a.sensor_id = d.id
               WHERE a.timestamp > NOW() - INTERVAL '7 days'
               ORDER BY a.timestamp DESC LIMIT 200""",
        )

        self.run_query(
            "JOIN 查询 - 三表关联统计",
            """SELECT d.area_id,
                      COUNT(DISTINCT d.id) AS device_count,
                      COUNT(s.id) AS reading_count,
                      COUNT(a.id) AS alarm_count,
                      AVG(s.pressure) AS avg_pressure
               FROM perf_devices d
               LEFT JOIN perf_sensor_data s ON s.device_id = d.device_id
                   AND s.timestamp > NOW() - INTERVAL '1 day'
               LEFT JOIN perf_alarms a ON a.sensor_id = d.id
                   AND a.timestamp > NOW() - INTERVAL '1 day'
               GROUP BY d.area_id ORDER BY reading_count DESC""",
        )

        # 4. GIS 空间查询
        print("\n  --- GIS 空间查询 ---")

        # 先检查 PostGIS 是否可用
        postgis_available = False
        try:
            with self.db_pool.get_connection() as conn:
                cur = conn.cursor()
                cur.execute("SELECT PostGIS_Version()")
                postgis_available = True
                cur.close()
        except Exception:
            pass

        if postgis_available:
            self.run_query(
                "GIS 查询 - 圆形范围内设备 (1km)",
                """SELECT d.id, d.device_id, d.name, d.device_type,
                          ST_Distance(
                            ST_SetSRID(ST_MakePoint(d.lng, d.lat), 4326)::geography,
                            ST_SetSRID(ST_MakePoint(%s, %s), 4326)::geography
                          ) AS distance_m
                   FROM perf_devices d
                   WHERE ST_DWithin(
                     ST_SetSRID(ST_MakePoint(d.lng, d.lat), 4326)::geography,
                     ST_SetSRID(ST_MakePoint(%s, %s), 4326)::geography,
                     1000
                   )
                   ORDER BY distance_m""",
                (116.4, 39.9, 116.4, 39.9),
            )
        else:
            # 无 PostGIS 时用距离公式近似
            self.run_query(
                "空间查询 - 距离近似计算 (无PostGIS)",
                """SELECT id, device_id, name, device_type,
                          SQRT(POW(lng - %s, 2) + POW(lat - %s, 2)) * 111000 AS approx_distance_m
                   FROM perf_devices
                   WHERE ABS(lng - %s) < 0.01 AND ABS(lat - %s) < 0.01
                   ORDER BY approx_distance_m
                   LIMIT 100""",
                (116.4, 39.9, 116.4, 39.9),
            )

        self.run_query(
            "空间查询 - 区域设备统计",
            """SELECT area_id,
                      COUNT(*) AS device_count,
                      AVG(lng) AS center_lng,
                      AVG(lat) AS center_lat,
                      COUNT(CASE WHEN status = 'active' THEN 1 END) AS active_count
               FROM perf_devices
               GROUP BY area_id ORDER BY device_count DESC""",
        )

        # 5. 复杂查询
        print("\n  --- 复杂查询 ---")
        self.run_query(
            "复杂查询 - 异常检测 (压力突变)",
            """WITH sensor_stats AS (
                 SELECT sensor_id,
                        AVG(pressure) AS avg_p,
                        STDDEV(pressure) AS std_p
                 FROM perf_sensor_data
                 WHERE timestamp > NOW() - INTERVAL '7 days'
                 GROUP BY sensor_id
                 HAVING COUNT(*) > 10
               )
               SELECT s.sensor_id, s.timestamp, s.pressure,
                      ss.avg_p, ss.std_p,
                      ABS(s.pressure - ss.avg_p) / NULLIF(ss.std_p, 0) AS z_score
               FROM perf_sensor_data s
               JOIN sensor_stats ss ON s.sensor_id = ss.sensor_id
               WHERE ABS(s.pressure - ss.avg_p) > 3 * NULLIF(ss.std_p, 0)
                 AND s.timestamp > NOW() - INTERVAL '1 day'
               ORDER BY z_score DESC LIMIT 50""",
        )

        self.run_query(
            "复杂查询 - 窗口函数 (滑动平均)",
            """SELECT sensor_id, timestamp, pressure,
                      AVG(pressure) OVER (
                        PARTITION BY sensor_id
                        ORDER BY timestamp
                        ROWS BETWEEN 5 PRECEDING AND 5 FOLLOWING
                      ) AS moving_avg
               FROM perf_sensor_data
               WHERE sensor_id = %s
                 AND timestamp > NOW() - INTERVAL '1 day'
               ORDER BY timestamp LIMIT 500""",
            (random.randint(1, 5000),),
        )

        return self.results


# ==================== 索引对比测试 ====================

def run_index_comparison(db_pool):
    """对比有无索引的查询性能"""
    print("\n" + "=" * 60)
    print("📊 索引性能对比测试")
    print("=" * 60)

    # 无索引测试
    print("\n--- 无索引状态 ---")
    drop_indexes(db_pool)
    bench_no_idx = QueryBenchmark(db_pool)
    results_no_idx = bench_no_idx.run_all_benchmarks()

    # 有索引测试
    print("\n--- 有索引状态 ---")
    create_indexes(db_pool)

    # 先 ANALYZE 更新统计
    with db_pool.get_connection() as conn:
        cur = conn.cursor()
        cur.execute("ANALYZE perf_sensor_data, perf_alarms, perf_devices")
        conn.commit()
        cur.close()

    bench_with_idx = QueryBenchmark(db_pool)
    results_with_idx = bench_with_idx.run_all_benchmarks()

    # 对比
    print("\n" + "=" * 60)
    print("📊 索引性能对比结果")
    print("=" * 60)
    print(f"{'查询名称':<40} {'无索引(ms)':>12} {'有索引(ms)':>12} {'提升':>10}")
    print("-" * 74)

    comparison = []
    for no_idx, with_idx in zip(results_no_idx, results_with_idx):
        no_avg = no_idx.get("avg_ms", 0)
        with_avg = with_idx.get("avg_ms", 0)
        if no_avg > 0:
            improvement = f"{(1 - with_avg / no_avg) * 100:.1f}%"
        else:
            improvement = "N/A"
        print(f"{no_idx['name']:<40} {no_avg:>12.2f} {with_avg:>12.2f} {improvement:>10}")
        comparison.append({
            "query": no_idx["name"],
            "no_index_ms": no_avg,
            "with_index_ms": with_avg,
        })

    return comparison


# ==================== 入口 ====================

def main():
    parser = argparse.ArgumentParser(description="数据库查询压力测试")
    parser.add_argument("--host", default=DEFAULT_HOST, help="数据库主机")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="数据库端口")
    parser.add_argument("--db", default=DEFAULT_DB, help="数据库名")
    parser.add_argument("--user", default=DEFAULT_USER, help="数据库用户")
    parser.add_argument("--password", default=DEFAULT_PASSWORD, help="数据库密码")
    parser.add_argument("--generate-data", type=int, default=0,
                        help="生成指定条数的测试数据 (如 1000000)")
    parser.add_argument("--skip-data-gen", action="store_true", help="跳过数据生成")
    parser.add_argument("--index-compare", action="store_true", help="运行索引对比测试")
    parser.add_argument("--output", default=None, help="结果输出 JSON 文件路径")

    args = parser.parse_args()

    print(f"\n{'=' * 60}")
    print(f"🗄️ 数据库查询压力测试")
    print(f"{'=' * 60}")
    print(f"数据库: {args.host}:{args.port}/{args.db}")
    print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'=' * 60}")

    # 连接数据库
    try:
        db_pool = DBConnectionPool(
            args.host, args.port, args.db, args.user, args.password
        )
    except Exception as e:
        print(f"❌ 数据库连接失败: {e}")
        sys.exit(1)

    all_results = {}

    try:
        # 生成测试数据
        if args.generate_data > 0 and not args.skip_data_gen:
            generate_test_data(db_pool, args.generate_data)
            create_indexes(db_pool)
            # ANALYZE
            with db_pool.get_connection() as conn:
                cur = conn.cursor()
                cur.execute("ANALYZE perf_sensor_data, perf_alarms, perf_devices")
                conn.commit()
                cur.close()

        # 运行基准查询测试
        bench = QueryBenchmark(db_pool)
        results = bench.run_all_benchmarks()
        all_results["query_benchmark"] = results

        # 索引对比测试
        if args.index_compare:
            comparison = run_index_comparison(db_pool)
            all_results["index_comparison"] = comparison

    finally:
        db_pool.close()

    # 输出汇总
    print(f"\n{'=' * 60}")
    print("📊 查询测试汇总")
    print(f"{'=' * 60}")
    for r in all_results.get("query_benchmark", []):
        avg = r.get("avg_ms", "ERR")
        print(f"  {r['name']}: {avg}ms ({r.get('rows_returned', 0)} rows)")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(all_results, f, indent=2, ensure_ascii=False, default=str)
        print(f"\n结果已保存到: {args.output}")

    print(f"\n{'=' * 60}")
    return all_results


if __name__ == "__main__":
    main()