128 lines
3.4 KiB
Python
128 lines
3.4 KiB
Python
from psycopg.rows import dict_row, Row
|
|
from .connection import g_conn_dict as conn
|
|
from .database import read
|
|
|
|
|
|
_NODE = '_node'
|
|
_LINK = '_link'
|
|
_CURVE = '_curve'
|
|
_PATTERN = '_pattern'
|
|
_REGION = '_region'
|
|
|
|
JUNCTION = 'junction'
|
|
RESERVOIR = 'reservoir'
|
|
TANK = 'tank'
|
|
PIPE = 'pipe'
|
|
PUMP = 'pump'
|
|
VALVE = 'valve'
|
|
|
|
PATTERN = 'pattern'
|
|
CURVE = 'curve'
|
|
|
|
REGION = 'region'
|
|
|
|
|
|
def _get_from(name: str, id: str, base_type: str) -> Row | None:
|
|
with conn[name].cursor(row_factory=dict_row) as cur:
|
|
cur.execute(f"select * from {base_type} where id = '{id}'")
|
|
return cur.fetchone()
|
|
|
|
|
|
def is_node(name: str, id: str) -> bool:
|
|
return _get_from(name, id, _NODE) != None
|
|
|
|
|
|
def is_junction(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _NODE)
|
|
return row != None and row['type'] == JUNCTION
|
|
|
|
|
|
def is_reservoir(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _NODE)
|
|
return row != None and row['type'] == RESERVOIR
|
|
|
|
|
|
def is_tank(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _NODE)
|
|
return row != None and row['type'] == TANK
|
|
|
|
|
|
def is_link(name: str, id: str) -> bool:
|
|
return _get_from(name, id, _LINK) != None
|
|
|
|
|
|
def is_pipe(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _LINK)
|
|
return row != None and row['type'] == PIPE
|
|
|
|
|
|
def is_pump(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _LINK)
|
|
return row != None and row['type'] == PUMP
|
|
|
|
|
|
def is_valve(name: str, id: str) -> bool:
|
|
row = _get_from(name, id, _LINK)
|
|
return row != None and row['type'] == VALVE
|
|
|
|
|
|
def is_curve(name: str, id: str) -> bool:
|
|
return _get_from(name, id, _CURVE) != None
|
|
|
|
|
|
def is_pattern(name: str, id: str) -> bool:
|
|
return _get_from(name, id, _PATTERN) != None
|
|
|
|
def is_region(name: str, id: str) -> bool:
|
|
return _get_from(name, id, _REGION) != None
|
|
|
|
|
|
def _get_all(name: str, base_type: str) -> list[str]:
|
|
ids : list[str] = []
|
|
with conn[name].cursor(row_factory=dict_row) as cur:
|
|
cur.execute(f"select id from {base_type} order by id")
|
|
for record in cur:
|
|
ids.append(record['id'])
|
|
return ids
|
|
|
|
|
|
def get_nodes(name: str) -> list[str]:
|
|
return _get_all(name, _NODE)
|
|
|
|
|
|
def get_links(name: str) -> list[str]:
|
|
return _get_all(name, _LINK)
|
|
|
|
|
|
def get_curves(name: str) -> list[str]:
|
|
return _get_all(name, _CURVE)
|
|
|
|
|
|
def get_patterns(name: str) -> list[str]:
|
|
return _get_all(name, _PATTERN)
|
|
|
|
def get_regions(name: str) -> list[str]:
|
|
return _get_all(name, _REGION)
|
|
|
|
def get_node_links(name: str, id: str) -> list[str]:
|
|
with conn[name].cursor(row_factory=dict_row) as cur:
|
|
links: list[str] = []
|
|
for p in cur.execute(f"select id from pipes where node1 = '{id}' or node2 = '{id}'").fetchall():
|
|
links.append(p['id'])
|
|
for p in cur.execute(f"select id from pumps where node1 = '{id}' or node2 = '{id}'").fetchall():
|
|
links.append(p['id'])
|
|
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'])]
|