bot_dev1
5eae031679
- 实现REST API服务器，支持IoT数据、手动录入和批量导入接口 - 实现WebSocket服务器，支持实时数据推送和连接管理 - 实现批量导入模块，支持CSV、Excel、JSON多种格式 - 实现数据处理工具，包含字段映射和单位转换功能 - 实现数据模型定义和数据验证机制 - 创建主程序入口和配置文件 - 添加详细的使用文档和API说明
305 lines
9.4 KiB
Python
305 lines
9.4 KiB
Python
"""
|
|
水务管理系统主程序
|
|
集成REST API、WebSocket和批量导入功能
|
|
"""
|
|
import asyncio
|
|
import logging
|
|
import signal
|
|
import sys
|
|
from pathlib import Path
|
|
import argparse
|
|
from datetime import datetime
|
|
|
|
# 导入各个模块
|
|
from src.api.rest_api import app as rest_api_app
|
|
from src.websocket.websocket_server import websocket_server
|
|
from src.batch.batch_import import batch_manager
|
|
from src.utils.data_utils import data_converter, data_formatter, quality_checker
|
|
from src.models.models import validator
|
|
|
|
# 配置日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.FileHandler('water_management.log'),
|
|
logging.StreamHandler(sys.stdout)
|
|
]
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class WaterManagementSystem:
|
|
"""水务管理系统主类"""
|
|
|
|
def __init__(self, config_file: str = "config.json"):
|
|
self.config_file = config_file
|
|
self.running = False
|
|
self.tasks = []
|
|
self.config = self.load_config()
|
|
|
|
def load_config(self) -> dict:
|
|
"""加载配置文件"""
|
|
config_path = Path(self.config_file)
|
|
if config_path.exists():
|
|
import json
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
else:
|
|
# 默认配置
|
|
return {
|
|
"api": {
|
|
"host": "0.0.0.0",
|
|
"port": 8000
|
|
},
|
|
"websocket": {
|
|
"host": "0.0.0.0",
|
|
"port": 8765
|
|
},
|
|
"batch": {
|
|
"max_file_size_mb": 100,
|
|
"supported_formats": ["csv", "excel", "json"]
|
|
},
|
|
"logging": {
|
|
"level": "INFO",
|
|
"file": "water_management.log"
|
|
}
|
|
}
|
|
|
|
async def start_api_server(self):
|
|
"""启动API服务器"""
|
|
import uvicorn
|
|
logger.info("启动REST API服务器...")
|
|
|
|
# 在新的事件循环中运行uvicorn
|
|
api_config = uvicorn.Config(
|
|
app=rest_api_app,
|
|
host=self.config["api"]["host"],
|
|
port=self.config["api"]["port"],
|
|
log_level="info"
|
|
)
|
|
api_server = uvicorn.Server(api_config)
|
|
|
|
# 在后台任务中运行
|
|
await api_server.serve()
|
|
|
|
async def start_websocket_server(self):
|
|
"""启动WebSocket服务器"""
|
|
logger.info("启动WebSocket服务器...")
|
|
|
|
# 启动WebSocket服务器
|
|
server = await websocket_server.start_server()
|
|
|
|
# 添加服务器关闭处理
|
|
def cleanup():
|
|
logger.info("关闭WebSocket服务器...")
|
|
server.close()
|
|
asyncio.create_task(server.wait_closed())
|
|
|
|
return server, cleanup
|
|
|
|
async def start_data_generator(self):
|
|
"""启动数据生成器（用于演示）"""
|
|
logger.info("启动数据生成器...")
|
|
|
|
while self.running:
|
|
# 生成模拟数据
|
|
import random
|
|
sensor_types = ["LL", "YL", "SW", "ZD"]
|
|
sensor_type = random.choice(sensor_types)
|
|
|
|
# 根据传感器类型生成合理的数值范围
|
|
if sensor_type == "LL": # 流量
|
|
value = random.uniform(10, 100)
|
|
elif sensor_type == "YL": # 压力
|
|
value = random.uniform(0.1, 1.0)
|
|
elif sensor_type == "SW": # 水位
|
|
value = random.uniform(0, 10)
|
|
else: # ZD 浊度
|
|
value = random.uniform(0, 50)
|
|
|
|
sensor_data = {
|
|
"data_type": sensor_type,
|
|
"device_id": f"device_{random.randint(1, 10)}",
|
|
"value": round(value, 2),
|
|
"location": random.choice(["A区", "B区", "C区", "D区"])
|
|
}
|
|
|
|
# 通过WebSocket发送数据
|
|
await websocket_server.send_sensor_data(sensor_data)
|
|
|
|
# 等待5秒
|
|
await asyncio.sleep(5)
|
|
|
|
async def handle_batch_import(self, file_path: str, batch_id: str, data_source: str):
|
|
"""处理批量导入请求"""
|
|
try:
|
|
logger.info(f"开始批量导入: {file_path}")
|
|
|
|
# 验证文件
|
|
validation_result = await batch_manager.validate_file(file_path)
|
|
if not validation_result["valid"]:
|
|
raise Exception(f"文件验证失败: {validation_result['error']}")
|
|
|
|
# 导入文件
|
|
result = await batch_manager.import_file(
|
|
file_path=file_path,
|
|
batch_id=batch_id,
|
|
data_source=data_source,
|
|
file_type="auto"
|
|
)
|
|
|
|
logger.info(f"批量导入完成: {result}")
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"批量导入失败: {str(e)}")
|
|
raise
|
|
|
|
async def start_system(self):
|
|
"""启动系统"""
|
|
logger.info("启动水务管理系统...")
|
|
|
|
# 标记系统为运行状态
|
|
self.running = True
|
|
|
|
try:
|
|
# 启动WebSocket服务器
|
|
ws_server, ws_cleanup = await self.start_websocket_server()
|
|
self.tasks.append(ws_server)
|
|
|
|
# 启动数据生成器（如果启用）
|
|
if self.config.get("demo_mode", False):
|
|
generator_task = asyncio.create_task(self.start_data_generator())
|
|
self.tasks.append(generator_task)
|
|
|
|
# 启动API服务器
|
|
await self.start_api_server()
|
|
|
|
except Exception as e:
|
|
logger.error(f"系统启动失败: {str(e)}")
|
|
await self.stop_system()
|
|
raise
|
|
|
|
async def stop_system(self):
|
|
"""停止系统"""
|
|
logger.info("停止水务管理系统...")
|
|
|
|
self.running = False
|
|
|
|
# 取消所有任务
|
|
for task in self.tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
# 清理WebSocket服务器
|
|
if hasattr(websocket_server, 'server') and websocket_server.server:
|
|
websocket_server.server.close()
|
|
await websocket_server.server.wait_closed()
|
|
|
|
logger.info("水务管理系统已停止")
|
|
|
|
async def run(self):
|
|
"""运行系统"""
|
|
# 设置信号处理
|
|
def signal_handler():
|
|
logger.info("收到停止信号...")
|
|
asyncio.create_task(self.stop_system())
|
|
|
|
for sig in [signal.SIGINT, signal.SIGTERM]:
|
|
signal.signal(sig, signal_handler)
|
|
|
|
try:
|
|
await self.start_system()
|
|
|
|
# 保持运行直到收到停止信号
|
|
while self.running:
|
|
await asyncio.sleep(1)
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("收到键盘中断信号")
|
|
except Exception as e:
|
|
logger.error(f"系统运行时出错: {str(e)}")
|
|
finally:
|
|
await self.stop_system()
|
|
|
|
def create_sample_config():
|
|
"""创建示例配置文件"""
|
|
sample_config = {
|
|
"api": {
|
|
"host": "0.0.0.0",
|
|
"port": 8000
|
|
},
|
|
"websocket": {
|
|
"host": "0.0.0.0",
|
|
"port": 8765
|
|
},
|
|
"batch": {
|
|
"max_file_size_mb": 100,
|
|
"supported_formats": ["csv", "excel", "json"]
|
|
},
|
|
"demo_mode": True,
|
|
"logging": {
|
|
"level": "INFO",
|
|
"file": "water_management.log"
|
|
}
|
|
}
|
|
|
|
import json
|
|
with open("config.json", 'w', encoding='utf-8') as f:
|
|
json.dump(sample_config, f, indent=2, ensure_ascii=False)
|
|
|
|
logger.info("示例配置文件已创建: config.json")
|
|
|
|
def create_requirements():
|
|
"""创建requirements.txt文件"""
|
|
requirements = [
|
|
"fastapi==0.104.1",
|
|
"uvicorn[standard]==0.24.0",
|
|
"websockets==12.0",
|
|
"pandas==2.1.3",
|
|
"openpyxl==3.1.2",
|
|
"aiofiles==23.2.1",
|
|
"python-multipart==0.0.6",
|
|
"jinja2==3.1.2"
|
|
]
|
|
|
|
with open("requirements.txt", 'w') as f:
|
|
f.write('\n'.join(requirements))
|
|
|
|
logger.info("依赖文件已创建: requirements.txt")
|
|
|
|
def main():
|
|
"""主函数"""
|
|
parser = argparse.ArgumentParser(description="水务管理系统")
|
|
parser.add_argument("--config", "-c", default="config.json", help="配置文件路径")
|
|
parser.add_argument("--init", action="store_true", help="初始化项目（创建配置文件和依赖）")
|
|
parser.add_argument("--demo", action="store_true", help="启动演示模式")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.init:
|
|
create_sample_config()
|
|
create_requirements()
|
|
logger.info("项目初始化完成")
|
|
return
|
|
|
|
# 创建系统实例
|
|
system = WaterManagementSystem(args.config)
|
|
|
|
# 如果启用演示模式
|
|
if args.demo:
|
|
system.config["demo_mode"] = True
|
|
logger.info("启用演示模式")
|
|
|
|
# 运行系统
|
|
try:
|
|
asyncio.run(system.run())
|
|
except KeyboardInterrupt:
|
|
logger.info("程序已退出")
|
|
|
|
if __name__ == "__main__":
|
|
main() |