581 lines
20 KiB
Python
581 lines
20 KiB
Python
import ast
|
||
import json
|
||
from datetime import date, datetime
|
||
|
||
import geopandas as gpd
|
||
import pandas as pd
|
||
import psycopg
|
||
from sqlalchemy import create_engine
|
||
|
||
from app.core.config import get_pgconn_string
|
||
from app.services.time_api import parse_utc_time
|
||
|
||
|
||
# 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: datetime | 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)
|
||
normalized_scheme_start_time = parse_utc_time(
|
||
scheme_start_time, field_name="scheme_start_time"
|
||
)
|
||
cur.execute(
|
||
sql,
|
||
(
|
||
scheme_name,
|
||
scheme_type,
|
||
username,
|
||
normalized_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}")
|