import ast import json from datetime import date import geopandas as gpd import pandas as pd import psycopg from sqlalchemy import create_engine from app.core.config import get_pgconn_string # 2025/03/23 def create_user(name: str, username: str, password: str): """ 创建用户 :param name: 数据库名称 :param username: 用户名 :param password: 密码 :return: """ try: # 动态替换数据库名称 conn_string = get_pgconn_string(db_name=name) # 连接到 PostgreSQL 数据库(这里是数据库 "bb") with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( "INSERT INTO users (username, password) VALUES (%s, %s)", (username, password), ) # 提交事务 conn.commit() print("新用户创建成功!") except Exception as e: print(f"创建用户出错:{e}") # 2025/03/23 def delete_user(name: str, username: str): """ 删除用户 :param name: 数据库名称 :param username: 用户名 :return: """ try: # 动态替换数据库名称 conn_string = get_pgconn_string(db_name=name) # 连接到 PostgreSQL 数据库(这里是数据库 "bb") with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute("DELETE FROM users WHERE username = %s", (username,)) conn.commit() print(f"用户 {username} 删除成功!") except Exception as e: print(f"删除用户出错:{e}") # 2025/03/23 def scheme_name_exists(name: str, scheme_name: str) -> bool: """ 判断传入的 scheme_name 是否已存在于 scheme_list 表中,用于输入框判断 :param name: 数据库名称 :param scheme_name: 需要判断的方案名称 :return: 如果存在返回 True,否则返回 False """ try: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( "SELECT COUNT(*) FROM scheme_list WHERE scheme_name = %s", (scheme_name,), ) result = cur.fetchone() if result is not None and result[0] > 0: return True else: return False except Exception as e: print(f"查询 scheme_name 时出错:{e}") return False # 2025/03/23 def store_scheme_info( name: str, scheme_name: str, scheme_type: str, username: str, scheme_start_time: str, scheme_detail: dict, ): """ 将一条方案记录插入 scheme_list 表中 :param name: 数据库名称 :param scheme_name: 方案名称 :param scheme_type: 方案类型 :param username: 用户名(需在 users 表中已存在) :param scheme_start_time: 方案起始时间(字符串) :param scheme_detail: 方案详情(字典,会转换为 JSON) :return: """ try: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: sql = """ INSERT INTO scheme_list (scheme_name, scheme_type, username, scheme_start_time, scheme_detail) VALUES (%s, %s, %s, %s, %s) """ # 将字典转换为 JSON 字符串 scheme_detail_json = json.dumps(scheme_detail) cur.execute( sql, ( scheme_name, scheme_type, username, scheme_start_time, scheme_detail_json, ), ) conn.commit() print("方案信息存储成功!") except Exception as e: print(f"存储方案信息时出错:{e}") # 2025/03/23 def delete_scheme_info(name: str, scheme_name: str) -> None: """ 从 scheme_list 表中删除指定的方案 :param name: 数据库名称 :param scheme_name: 要删除的方案名称 """ try: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: # 使用参数化查询删除方案记录 cur.execute( "DELETE FROM scheme_list WHERE scheme_name = %s", (scheme_name,) ) conn.commit() print(f"方案 {scheme_name} 删除成功!") except Exception as e: print(f"删除方案时出错:{e}") # 2025/03/23 def query_scheme_list(name: str) -> list: """ 查询pg数据库中的scheme_list,按照 create_time 降序排列,离现在时间最近的记录排在最前面 :param name: 项目名称(数据库名称) :return: 返回查询结果的所有行 """ try: # 动态替换数据库名称 conn_string = get_pgconn_string(db_name=name) # 连接到 PostgreSQL 数据库(这里是数据库 "bb") with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: # 按 create_time 降序排列 cur.execute("SELECT * FROM scheme_list ORDER BY create_time DESC") rows = cur.fetchall() return rows except Exception as e: print(f"查询错误:{e}") def store_leakage_identify_result( name: str, scheme_name: str, network: str, sensor_nodes: list[str], result_rows: list[dict], node_area_map: dict[str, str], areas: list[dict], drawing_payload: dict | None = None, run_status: str = "completed", error_message: str | None = None, ) -> None: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( """ INSERT INTO public.leakage_identify_result ( scheme_name, network, run_status, error_message, sensor_nodes, result_rows, node_area_map, areas, drawing_payload ) VALUES (%s, %s, %s, %s, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb) ON CONFLICT (scheme_name) DO UPDATE SET network = EXCLUDED.network, run_status = EXCLUDED.run_status, error_message = EXCLUDED.error_message, sensor_nodes = EXCLUDED.sensor_nodes, result_rows = EXCLUDED.result_rows, node_area_map = EXCLUDED.node_area_map, areas = EXCLUDED.areas, drawing_payload = EXCLUDED.drawing_payload, created_at = NOW(); """, ( scheme_name, network, run_status, error_message, json.dumps(sensor_nodes), json.dumps(result_rows), json.dumps(node_area_map), json.dumps(areas), json.dumps(drawing_payload or {}), ), ) conn.commit() def query_leakage_identify_schemes( name: str, network: str, scheme_type: str = "dma_leak_identification", query_date: date | None = None, ) -> list[dict]: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: if query_date is None: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s ORDER BY create_time DESC """, (scheme_type,), ) else: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s AND DATE(create_time) = %s ORDER BY create_time DESC """, (scheme_type, query_date), ) rows = cur.fetchall() result = [] for row in rows: detail = row[6] if isinstance(row[6], dict) else {} if network and detail.get("network") not in (None, network): continue result.append( { "scheme_id": row[0], "scheme_name": row[1], "scheme_type": row[2], "username": row[3], "create_time": row[4], "scheme_start_time": row[5], "scheme_detail": detail, } ) return result def query_leakage_identify_scheme_detail(name: str, scheme_name: str) -> dict: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_name = %s LIMIT 1 """, (scheme_name,), ) base_row = cur.fetchone() if base_row is None: return {} cur.execute( """ SELECT network, created_at, run_status, error_message, sensor_nodes, result_rows, node_area_map, areas, drawing_payload FROM public.leakage_identify_result WHERE scheme_name = %s LIMIT 1 """, (scheme_name,), ) result_row = cur.fetchone() if result_row is None: return {} return { "scheme_id": base_row[0], "scheme_name": base_row[1], "scheme_type": base_row[2], "username": base_row[3], "create_time": base_row[4], "scheme_start_time": base_row[5], "scheme_detail": base_row[6] if isinstance(base_row[6], dict) else {}, "network": result_row[0], "result_created_at": result_row[1], "run_status": result_row[2], "error_message": result_row[3], "sensor_nodes": result_row[4] if isinstance(result_row[4], list) else [], "rows": result_row[5] if isinstance(result_row[5], list) else [], "node_area_map": result_row[6] if isinstance(result_row[6], dict) else {}, "areas": result_row[7] if isinstance(result_row[7], list) else [], "drawing_payload": ( result_row[8] if isinstance(result_row[8], dict) else {"type": "FeatureCollection", "features": []} ), } def query_burst_location_schemes( name: str, network: str, scheme_type: str = "burst_location", query_date: date | None = None, ) -> list[dict]: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: if query_date is None: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s ORDER BY create_time DESC """, (scheme_type,), ) else: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s AND DATE(create_time) = %s ORDER BY create_time DESC """, (scheme_type, query_date), ) rows = cur.fetchall() result = [] for row in rows: detail = row[6] if isinstance(row[6], dict) else {} if network and detail.get("network") not in (None, network): continue result.append( { "scheme_id": row[0], "scheme_name": row[1], "scheme_type": row[2], "username": row[3], "create_time": row[4], "scheme_start_time": row[5], "scheme_detail": detail, } ) return result def query_burst_location_scheme_detail(name: str, scheme_name: str) -> dict: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_name = %s LIMIT 1 """, (scheme_name,), ) base_row = cur.fetchone() if base_row is None: return {} detail = base_row[6] if isinstance(base_row[6], dict) else {} return { "scheme_id": base_row[0], "scheme_name": base_row[1], "scheme_type": base_row[2], "username": base_row[3], "create_time": base_row[4], "scheme_start_time": base_row[5], "scheme_detail": detail, "network": detail.get("network"), "result_payload": detail.get("result_payload", {}), } def query_burst_detection_schemes( name: str, network: str, scheme_type: str = "burst_detection", query_date: date | None = None, ) -> list[dict]: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: if query_date is None: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s ORDER BY create_time DESC """, (scheme_type,), ) else: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_type = %s AND DATE(create_time) = %s ORDER BY create_time DESC """, (scheme_type, query_date), ) rows = cur.fetchall() result = [] for row in rows: detail = row[6] if isinstance(row[6], dict) else {} if network and detail.get("network") not in (None, network): continue result.append( { "scheme_id": row[0], "scheme_name": row[1], "scheme_type": row[2], "username": row[3], "create_time": row[4], "scheme_start_time": row[5], "scheme_detail": detail, } ) return result def query_burst_detection_scheme_detail(name: str, scheme_name: str) -> dict: conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: cur.execute( """ SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail FROM public.scheme_list WHERE scheme_name = %s LIMIT 1 """, (scheme_name,), ) base_row = cur.fetchone() if base_row is None: return {} detail = base_row[6] if isinstance(base_row[6], dict) else {} return { "scheme_id": base_row[0], "scheme_name": base_row[1], "scheme_type": base_row[2], "username": base_row[3], "create_time": base_row[4], "scheme_start_time": base_row[5], "scheme_detail": detail, "network": detail.get("network"), "result_payload": detail.get("result_payload", {}), } # 2025/03/23 def upload_shp_to_pg(name: str, table_name: str, role: str, shp_file_path: str): """ 将 Shapefile 文件上传到 PostgreSQL 数据库 :param name: 项目名称(数据库名称) :param table_name: 创建表的名字 :param role: 数据库角色名,位于c盘user中查看 :param shp_file_path: shp文件的路径 :return: """ try: # 动态连接到指定的数据库 conn_string = get_pgconn_string(db_name=name) with psycopg.connect(conn_string) as conn: # 读取 Shapefile 文件 gdf = gpd.read_file(shp_file_path) # 检查投影坐标系(CRS),并确保是 EPSG:4326 if gdf.crs.to_string() != "EPSG:4490": gdf = gdf.to_crs(epsg=4490) # 使用 GeoDataFrame 的 .to_postgis 方法将数据写入 PostgreSQL # 需要在数据库中提前安装 PostGIS 扩展 engine = create_engine(f"postgresql+psycopg2://{role}:@127.0.0.1/{name}") gdf.to_postgis( table_name, engine, if_exists="replace", index=True, index_label="id" ) print( f"Shapefile 文件成功上传到 PostgreSQL 数据库 '{name}' 的表 '{table_name}'." ) except Exception as e: print(f"上传 Shapefile 到 PostgreSQL 时出错:{e}") def submit_risk_probability_result(name: str, result_file_path: str) -> None: """ 将管网风险评估结果导入pg数据库 :param name: 项目名称(数据库名称) :param result_file_path: 结果文件路径 :return: """ # 自动检测文件编码 # with open({result_file_path}, 'rb') as file: # raw_data = file.read() # detected = chardet.detect(raw_data) # file_encoding = detected['encoding'] # print(f"检测到的文件编码:{file_encoding}") try: # 动态替换数据库名称 conn_string = get_pgconn_string(db_name=name) # 连接到 PostgreSQL 数据库 with psycopg.connect(conn_string) as conn: with conn.cursor() as cur: # 检查 scada_info 表是否为空 cur.execute("SELECT COUNT(*) FROM pipe_risk_probability;") count = cur.fetchone()[0] if count > 0: print("pipe_risk_probability表中已有数据,正在清空记录...") cur.execute("DELETE FROM pipe_risk_probability;") print("表记录已清空。") # 读取Excel并转换x/y列为列表 df = pd.read_excel(result_file_path, sheet_name="Sheet1") df["x"] = df["x"].apply(ast.literal_eval) df["y"] = df["y"].apply(ast.literal_eval) # 批量插入数据 for index, row in df.iterrows(): insert_query = """ INSERT INTO pipe_risk_probability (pipeID, pipeage, risk_probability_now, x, y) VALUES (%s, %s, %s, %s, %s) """ cur.execute( insert_query, ( row["pipeID"], row["pipeage"], row["risk_probability_now"], row["x"], # 直接传递列表 row["y"], # 同上 ), ) conn.commit() print("风险评估结果导入成功") except Exception as e: print(f"导入时出错:{e}")