Files
TJWaterServer/api/s0_base.py
2022-09-16 23:15:16 +08:00

76 lines
2.1 KiB
Python

from psycopg.rows import dict_row, Row
from .connection import g_conn_dict as conn
_NODE = "_node"
_LINK = "_link"
_CURVE = "_curve"
_PATTERN = "_pattern"
JUNCTION = "junction"
RESERVOIR = "reservoir"
TANK = "tank"
PIPE = "pipe"
PUMP = "pump"
VALVE = "valve"
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) != {}
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 _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)