267 lines
9.2 KiB
Python
267 lines
9.2 KiB
Python
import ast
|
||
import json
|
||
|
||
import geopandas as gpd
|
||
import pandas as pd
|
||
import psycopg
|
||
from sqlalchemy import create_engine
|
||
|
||
from app.native.api.postgresql_info 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}")
|
||
|
||
|
||
# 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}")
|