diff --git a/api/s0_base.py b/api/s0_base.py index 5609f8f..bace826 100644 --- a/api/s0_base.py +++ b/api/s0_base.py @@ -1,5 +1,6 @@ from psycopg.rows import dict_row, Row from .connection import g_conn_dict as conn +from .database import read _NODE = '_node' @@ -105,3 +106,14 @@ def get_node_links(name: str, id: str) -> list[str]: for p in cur.execute(f"select id from valves where node1 = '{id}' or node2 = '{id}'").fetchall(): links.append(p['id']) return links + + +def get_link_nodes(name: str, id: str) -> list[str]: + row = {} + if is_pipe(name, id): + row = read(name, f"select node1, node2 from pipes where id = '{id}'") + elif is_pump(name, id): + row = read(name, f"select node1, node2 from pumps where id = '{id}'") + elif is_valve(name, id): + row = read(name, f"select node1, node2 from valves where id = '{id}'") + return [str(row['node1']), str(row['node2'])]