diff --git a/ai2/config.yaml b/ai2/config.yaml new file mode 100644 index 0000000..4ea0d31 --- /dev/null +++ b/ai2/config.yaml @@ -0,0 +1,13 @@ +minio: + endpoint: "222.212.85.86:9000" + access_key: "adminjdskfj" + secret_key: "123456ksldjfal@Y" + secure: false + web: "http://222.212.85.86" + +sql: + host: '222.212.85.86' + port: 5432 + dbname: 'postgres' + user: 'postgres' + password: 'root' \ No newline at end of file diff --git a/ai2/sqlhelp.py b/ai2/sqlhelp.py new file mode 100644 index 0000000..acc1d03 --- /dev/null +++ b/ai2/sqlhelp.py @@ -0,0 +1,44 @@ +import yaml +import psycopg2 + +def read_sql_config(yaml_name): + """ + 读取 SQL 配置 + """ + yaml_path = f"{yaml_name}.yaml" + with open(yaml_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) + sql_config = config.get('sql') + if not sql_config: + raise ValueError("未找到 'sql' 配置块") + return sql_config + +def get_user_power(user_id: str, yaml_name: str) -> int | None: + """ + 根据 user_id 查询数据库中对应的 power,找不到返回 None + """ + conn = None + try: + sql_config = read_sql_config(yaml_name) + conn = psycopg2.connect( + dbname=sql_config['dbname'], + user=sql_config['user'], + password=sql_config['password'], + host=sql_config['host'], + port=sql_config['port'] + ) + + conn.autocommit = True # 只读查询推荐设置 + with conn.cursor() as cur: + cur.execute("SELECT power FROM outusers WHERE user_id = %s", (user_id,)) + row = cur.fetchone() + if row: + return row[0] + else: + return None + except Exception as e: + print(f"数据库操作异常: {e}") + return None + finally: + if conn: + conn.close() \ No newline at end of file diff --git a/ai2/yolo_api.py b/ai2/yolo_api.py index 9dd7fd7..6963da2 100644 --- a/ai2/yolo_api.py +++ b/ai2/yolo_api.py @@ -10,7 +10,7 @@ import traceback from datetime import datetime from cv_video import startAIVideo,stopAIVideo,getIfAI from sanic_cors import CORS - +from sqlhelp import get_user_power # 配置日志 logging.basicConfig( level=logging.INFO, @@ -166,13 +166,18 @@ def verify_token(request) -> None: if not token or token != Config.VALID_TOKEN: logger.warning("Invalid token attempt") raise Unauthorized("Invalid token") - -@app.post("/ai/stream/detect") -async def start_detection(request): - try: - verify_token(request) - - # 检查服务健康状态 + +def verify_userid(request) -> None: + """验证请求userid""" + # 解析并验证请求数据 + stream_request = StreamRequest.from_dict(request.json) + userid = stream_request.user_id + if not userid or get_user_power(userid,"config") < 1: + logger.warning("userid not define or user have not power") + raise Unauthorized("Invalid userid or user have not power") + +async def detection(request): + # 检查服务健康状态 if not service_status["is_healthy"]: logger.warning(f"服务处于不健康状态,上次错误: {service_status['last_error']} 于 {service_status['error_time']}") # 尝试恢复服务 @@ -230,7 +235,29 @@ async def start_detection(request): "task_id": task_id, "message": "Detection started successfully" }) - + +@app.post("/ai/stream/detect1") +async def start_detection1(request): + try: + verify_userid(request) + try: + verify_token(request) + detection(request) + except Exception as e: + logger.error(f"Unexpected error: {str(e)}", exc_info=True) + return json_response({"status": "error", "message": f"Internal server error: {str(e)}"}, status=500) + except ValueError as e: + logger.error(f"Validation error: {str(e)}") + return json_response({"status": "error", "message": str(e)}, status=400) + except Exception as e: + logger.error(f"Unexpected error: {str(e)}", exc_info=True) + return json_response({"status": "error", "message": f"Internal server error: {str(e)}"}, status=500) + +@app.post("/ai/stream/detect") +async def start_detection(request): + try: + verify_token(request) + detection(request) except ValueError as e: logger.error(f"Validation error: {str(e)}") return json_response({"status": "error", "message": str(e)}, status=400)