209 lines
7.7 KiB
Python
209 lines
7.7 KiB
Python
"""
|
||
舷号-型号映射数据库模块
|
||
用于管理舰船型号、舷号、别名、舰名的映射关系
|
||
数据存储在 PostgreSQL 的 ship_model_mapping 表中
|
||
"""
|
||
import json
|
||
from typing import Dict, List, Any, Optional
|
||
|
||
from config import POSTGRES_CONNECTION_STRING
|
||
|
||
|
||
async def init_ship_model_mapping_table():
|
||
"""
|
||
初始化 ship_model_mapping 表并插入初始数据
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[ship_model_db] psycopg_pool 未安装，无法初始化表")
|
||
return
|
||
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS ship_model_mapping (
|
||
id SERIAL PRIMARY KEY,
|
||
model_name TEXT NOT NULL UNIQUE,
|
||
numbers JSONB NOT NULL DEFAULT '[]',
|
||
aliases JSONB NOT NULL DEFAULT '[]',
|
||
names JSONB NOT NULL DEFAULT '[]'
|
||
)
|
||
""")
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS ship_model_mapping_model_name_idx ON ship_model_mapping(model_name)
|
||
""")
|
||
print("[ship_model_db] ship_model_mapping 表初始化完成")
|
||
|
||
# 检查是否已有数据，没有则插入初始数据
|
||
await cur.execute("SELECT COUNT(*) FROM ship_model_mapping")
|
||
count = (await cur.fetchone())[0]
|
||
if count == 0:
|
||
initial_data = [
|
||
("055型驱逐舰", ["102", "103", "104", "105", "106", "107", "108"],
|
||
["055", "055型", "055驱逐舰", "万吨大驱"],
|
||
["拉萨舰", "鞍山舰", "无锡舰", "大连舰", "延安舰", "遵义舰", "咸阳舰"]),
|
||
("052D型驱逐舰", ["172", "173", "174", "175", "163"],
|
||
["052D", "052D型", "052D驱逐舰"],
|
||
["昆明舰", "长沙舰", "合肥舰", "银川舰", "南昌舰"]),
|
||
("054A型护卫舰", ["529", "530", "547", "568", "570"],
|
||
["054A", "054A型", "054A护卫舰"],
|
||
["舟山舰", "徐州舰", "临沂舰", "衡阳舰", "黄山舰"]),
|
||
]
|
||
for model_name, numbers, aliases, names in initial_data:
|
||
await cur.execute(
|
||
"""
|
||
INSERT INTO ship_model_mapping (model_name, numbers, aliases, names)
|
||
VALUES (%s, %s, %s, %s)
|
||
""",
|
||
(model_name,
|
||
json.dumps(numbers, ensure_ascii=False),
|
||
json.dumps(aliases, ensure_ascii=False),
|
||
json.dumps(names, ensure_ascii=False))
|
||
)
|
||
print(f"[ship_model_db] 已插入 {len(initial_data)} 条初始数据")
|
||
|
||
|
||
async def load_ship_model_mapping() -> Dict[str, Dict[str, List[str]]]:
|
||
"""
|
||
从数据库加载全部舷号-型号映射数据
|
||
|
||
Returns:
|
||
与原 config.SHIP_MODEL_MAPPING 格式一致的字典
|
||
{
|
||
"055型驱逐舰": {
|
||
"numbers": ["102", ...],
|
||
"aliases": ["055", ...],
|
||
"names": ["拉萨舰", ...]
|
||
},
|
||
...
|
||
}
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[ship_model_db] psycopg_pool 未安装，返回空映射")
|
||
return {}
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute(
|
||
"SELECT model_name, numbers, aliases, names FROM ship_model_mapping"
|
||
)
|
||
rows = await cur.fetchall()
|
||
result = {}
|
||
for row in rows:
|
||
model_name, numbers, aliases, names = row
|
||
# JSONB 返回的可能是 list 或 str，统一处理
|
||
if isinstance(numbers, str):
|
||
numbers = json.loads(numbers)
|
||
if isinstance(aliases, str):
|
||
aliases = json.loads(aliases)
|
||
if isinstance(names, str):
|
||
names = json.loads(names)
|
||
result[model_name] = {
|
||
"numbers": numbers,
|
||
"aliases": aliases,
|
||
"names": names,
|
||
}
|
||
return result
|
||
except Exception as e:
|
||
print(f"[ship_model_db] 加载映射数据失败: {str(e)}")
|
||
return {}
|
||
|
||
|
||
async def add_ship_model(
|
||
model_name: str,
|
||
numbers: List[str],
|
||
aliases: List[str],
|
||
names: List[str]
|
||
) -> bool:
|
||
"""
|
||
添加舰船型号映射
|
||
|
||
Args:
|
||
model_name: 型号名称
|
||
numbers: 舷号列表
|
||
aliases: 别名列表
|
||
names: 舰名列表
|
||
|
||
Returns:
|
||
是否添加成功
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[ship_model_db] psycopg_pool 未安装，无法添加")
|
||
return False
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute(
|
||
"""
|
||
INSERT INTO ship_model_mapping (model_name, numbers, aliases, names)
|
||
VALUES (%s, %s, %s, %s)
|
||
ON CONFLICT (model_name) DO UPDATE SET
|
||
numbers = EXCLUDED.numbers,
|
||
aliases = EXCLUDED.aliases,
|
||
names = EXCLUDED.names
|
||
""",
|
||
(model_name,
|
||
json.dumps(numbers, ensure_ascii=False),
|
||
json.dumps(aliases, ensure_ascii=False),
|
||
json.dumps(names, ensure_ascii=False))
|
||
)
|
||
print(f"[ship_model_db] 舰船型号 '{model_name}' 已保存")
|
||
return True
|
||
except Exception as e:
|
||
print(f"[ship_model_db] 添加舰船型号失败: {str(e)}")
|
||
return False
|
||
|
||
|
||
async def delete_ship_model(model_name: str) -> bool:
|
||
"""
|
||
删除舰船型号映射
|
||
|
||
Args:
|
||
model_name: 型号名称
|
||
|
||
Returns:
|
||
是否删除成功
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
print("[ship_model_db] psycopg_pool 未安装，无法删除")
|
||
return False
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
await cur.execute(
|
||
"DELETE FROM ship_model_mapping WHERE model_name = %s",
|
||
(model_name,)
|
||
)
|
||
print(f"[ship_model_db] 舰船型号 '{model_name}' 已删除")
|
||
return True
|
||
except Exception as e:
|
||
print(f"[ship_model_db] 删除舰船型号失败: {str(e)}")
|
||
return False
|
||
|