diff --git a/website/db/enterprise/__init__.py b/website/db/enterprise/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/website/service/enterprise.py b/website/db/enterprise/enterprise.py similarity index 100% rename from website/service/enterprise.py rename to website/db/enterprise/enterprise.py diff --git a/website/db/enterprise_busi_model/enterprise_busi_model_node_device.py b/website/db/enterprise_busi_model/enterprise_busi_model_node_device.py index 731fbc8..6bbf8bb 100644 --- a/website/db/enterprise_busi_model/enterprise_busi_model_node_device.py +++ b/website/db/enterprise_busi_model/enterprise_busi_model_node_device.py @@ -86,7 +86,7 @@ class EnterpriseBusiModelNodeDeviceRepository(object): ) return 1 if count > 0 else 0 - def get_device_ids(self, node_id: int, busi_model_id: int) -> int | None: + def get_device_ids(self, node_id: int, busi_model_id: int) -> list: with get_session() as session: sql = text( "select device_id from enterprise_busi_model_node_device where node_id=:node_id and busi_model_id=:busi_model_id" diff --git a/website/db/enterprise_device/enterprise_device.py b/website/db/enterprise_device/enterprise_device.py index 510c665..4a1097f 100644 --- a/website/db/enterprise_device/enterprise_device.py +++ b/website/db/enterprise_device/enterprise_device.py @@ -221,7 +221,7 @@ class EnterpriseDeviceRepository(object): "classification_name": device.classification_name, } device_dicts.append(device_dict) - return {"devices": device_dicts} + return device_dicts def get_device(self, device_id: int) -> dict: """ diff --git a/website/db/enterprise_node/enterprise_node.py b/website/db/enterprise_node/enterprise_node.py index 1573735..a748d46 100644 --- a/website/db/enterprise_node/enterprise_node.py +++ b/website/db/enterprise_node/enterprise_node.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -from typing import List, Dict, Any +from typing import List, Dict, Any, Union from sqlalchemy import text @@ -33,16 +33,20 @@ class EnterpriseNodeRepository(object): "insert into " "enterprise_node(suid, entity_id, entity_suid, name, parent, addr, lola, contact, phone, comment) " "values (:suid, :entity_id, :entity_suid, :name, :parent, :addr, :lola, :contact, :phone, :comment)" - ), node) + ), + node, + ) # last_insert_id = session.execute(text("SELECT LAST_INSERT_ID()")).scalar() # logging.info(f"last_insert_id: {last_insert_id}") return 0 - + def update_node(self, node: dict) -> int: with get_session() as session: - sql = ("update enterprise_node " - "set name=:name, parent=:parent, addr=:addr, lola=:lola, contact=:contact, phone=:phone, comment=:comment where id=:id") + sql = ( + "update enterprise_node " + "set name=:name, parent=:parent, addr=:addr, lola=:lola, contact=:contact, phone=:phone, comment=:comment where id=:id" + ) param = { "id": node["id"], "name": node["name"], @@ -51,8 +55,8 @@ class EnterpriseNodeRepository(object): "lola": node["lola"], "contact": node["contact"], "phone": node["phone"], - "comment": node["comment"] - } + "comment": node["comment"], + } session.execute(text(sql), param) return 0 @@ -74,14 +78,18 @@ class EnterpriseNodeRepository(object): "id": node["id"], "name": node["name"], "suid": node["suid"], - "children": self.build_tree(session, node, name) + "children": self.build_tree(session, node, name), } roots.append(root) # return node_list return roots - def build_tree(self, session: Any, node: Dict[str, Any], name: str = "") -> List[Any]: - sql = "select id, name, suid from enterprise_node where del=0 and parent=:parent" + def build_tree( + self, session: Any, node: Dict[str, Any], name: str = "" + ) -> List[Any]: + sql = ( + "select id, name, suid from enterprise_node where del=0 and parent=:parent" + ) param = {"parent": node["id"]} if name: @@ -97,17 +105,18 @@ class EnterpriseNodeRepository(object): "id": node["id"], "name": node["name"], "suid": node["suid"], - "children": self.build_tree(session, node, name) + "children": self.build_tree(session, node, name), } children.append(child) return children - def select_node(self, node_id: int) -> Dict[str, Any]: with get_session() as session: - sql = ("select id, name, parent, addr, lola, contact, phone, comment, suid from enterprise_node " - "where id=:id and del=0") + sql = ( + "select id, name, parent, addr, lola, contact, phone, comment, suid from enterprise_node " + "where id=:id and del=0" + ) param = {"id": node_id} res = session.execute(text(sql), param) node_list = to_json_list(res) @@ -116,23 +125,25 @@ class EnterpriseNodeRepository(object): def delete_node(self, node_id: int) -> int: with get_session() as session: - sql = ("update enterprise_node set del=1 where id=:id") + sql = "update enterprise_node set del=1 where id=:id" param = {"id": node_id} session.execute(text(sql), param) return 0 - + def get_node_by_id(self, node_id: int) -> dict: with get_session() as session: - sql = ("select suid, name from enterprise_node where id=:id") + sql = "select suid, name from enterprise_node where id=:id" param = {"id": node_id} res = session.execute(text(sql), param) node = to_json(res) return node - - def get_entity_suid_by_node_id(self, node_id: int) -> dict | None: + + def get_entity_suid_by_node_id(self, node_id: int) -> Union[dict, None]: with get_session() as session: - res = session.execute(text("select entity_suid from enterprise_node where id=:id"), - {"id": node_id}) + res = session.execute( + text("select entity_suid from enterprise_node where id=:id"), + {"id": node_id}, + ) entity = to_json(res) - - return entity if entity else None \ No newline at end of file + + return entity if entity else None diff --git a/website/db/enterprise_node/enterprise_node_alert.py b/website/db/enterprise_node/enterprise_node_alert.py index 7dfdb3c..c25cec6 100644 --- a/website/db/enterprise_node/enterprise_node_alert.py +++ b/website/db/enterprise_node/enterprise_node_alert.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -from typing import List, Dict, Any +from typing import List, Dict, Any, Union from sqlalchemy import text @@ -18,7 +18,9 @@ class EnterpriseNodeAlertRepository(object): sql, {"enterprise_node_id": self.enterprise_node_id} ) - def get_one(self, enterprise_suid: str, enterprise_node_id: int) -> Dict | None: + def get_one( + self, enterprise_suid: str, enterprise_node_id: int + ) -> Union[Dict, None]: with get_session() as session: sql = text( """SELECT * FROM `enterprise_alert` WHERE `enterprise_suid`=:enterprise_suid and `node_id`=:node_id;""" diff --git a/website/handlers/enterprise_entity/handler.py b/website/handlers/enterprise_entity/handler.py index e2ff87f..924c1d9 100644 --- a/website/handlers/enterprise_entity/handler.py +++ b/website/handlers/enterprise_entity/handler.py @@ -9,7 +9,7 @@ from website import db_mysql from website import errors from website import settings from website.handler import APIHandler, authenticated -from website.service import enterprise +from website.db.enterprise import enterprise from website.util import shortuuid, aes from concurrent.futures import ThreadPoolExecutor from functools import partial @@ -52,7 +52,6 @@ class EntityIndexHandler(APIHandler): count = db_mysql.to_json(cur_count) count = count["c"] if count else 0 - # data_index = [item["id"] for item in result] data = [] # for item in result: # modelCount = enterprise.get_enterprise_model_count(item["id"])