diff --git a/api/s3_reservoirs.py b/api/s3_reservoirs.py new file mode 100644 index 0000000..d28622f --- /dev/null +++ b/api/s3_reservoirs.py @@ -0,0 +1,157 @@ +from psycopg.rows import dict_row, Row +from .connection import g_conn_dict as conn +from .s0_base import * +from .operation import * +from .change_set import ChangeSet + +def add_reservoir(name: str, id: str, x: float, y: float, head: float) -> ChangeSet: + if is_node(name, id): + return + + with conn[name].cursor() as cur: + sql = f"insert into _node (id, type) values ('{id}', 'RESERVOIR');" + sql += f" insert into reservoirs (id, head) values ('{id}', {head});" + sql += f" insert into coordinates (node, coord) values ('{id}', '({x}, {y})');" + cur.execute(sql) + + redo = sql.replace("'", '"') + undo = f'delete from coordinates where node = "{id}";' + undo += f' delete from reservoirs where id = "{id}";' + undo += f' delete from _node where id = "{id}";' + add_operation(name, redo, undo) + + change = ChangeSet() + change.add('reservoir', id) + return change + +def delete_reservoir(name: str, id: str) -> ChangeSet: + if not is_reservoir(name, id): + return + + with conn[name].cursor(row_factory=dict_row) as cur: + cur.execute(f"select * from reservoirs where id = '{id}'") + row = cur.fetchone() + if row == None: + return + + head = row['head'] + pattern = 'NULL' if row['pattern'] == None else row['pattern'] + pattern = f'"{pattern}"' if pattern != 'NULL' else pattern + + cur.execute(f"select * from coordinates where node = '{id}'") + row = cur.fetchone() + if row == None: + return + + coord = row['coord'] + + sql = f"delete from coordinates where node = '{id}';" + sql += f" delete from reservoirs where id = '{id}';" + sql += f" delete from _node where id = '{id}';" + cur.execute(sql) + + redo = sql.replace("'", '"') + undo = f'insert into _node (id, type) values ("{id}", "RESERVOIR");' + undo += f' insert into reservoirs (id, head, pattern) values ("{id}", {head}, {pattern});' + undo += f' insert into coordinates (node, coord) values ("{id}", "{coord}");' + add_operation(name, redo, undo) + + change = ChangeSet() + change.delete('reservoir', id) + return change + +def _get_reservoir(name: str, id: str) -> Row | None: + with conn[name].cursor(row_factory=dict_row) as cur: + cur.execute(f"select head, pattern from reservoirs where id = '{id}'") + return cur.fetchone() + +def get_reservoir_head(name: str, id: str) -> float | None: + row = _get_reservoir(name, id) + return float(row['head']) if row != None else None + +def get_reservoir_pattern(name: str, id: str) -> str | None: + row = _get_reservoir(name, id) + if row != None: + return row['pattern'] if row['pattern'] != None else 'NULL' + else: + return None + +def _to_point(coord: str) -> dict[str, float]: + coord = coord.removeprefix('(') + coord = coord.removesuffix(')') + coord = coord.split(',') + return { 'x': float(coord[0]), 'y': float(coord[1]) } + +def get_reservoir_coord(name: str, id: str) -> dict[str, float] | None: + with conn[name].cursor(row_factory=dict_row) as cur: + cur.execute(f"select * from coordinates where node = '{id}'") + row = cur.fetchone() + if row == None: + return None + + coord = str(row['coord']) + return _to_point(coord) + +def set_reservoir_head(name: str, id: str, head: float) -> ChangeSet: + if not is_reservoir(name, id): + return + + old = get_reservoir_head(name, id) + if old == None: + return + + with conn[name].cursor() as cur: + sql = f"update reservoirs set head = {head} where id = '{id}'" + cur.execute(sql) + redo = sql.replace("'", '"') + undo = f'update reservoirs set head = {old} where id = "{id}"' + add_operation(name, redo, undo) + + change = ChangeSet() + change.update('reservoir', id, 'head', 'float', str(head)) + return change + +def set_reservoir_pattern(name: str, id: str, pattern: str) -> ChangeSet: + if not is_reservoir(name, id): + return + if not is_pattern(name, id): + return + + old = get_reservoir_pattern(name, id) + if old == None: + return + + old = f'"{old}"' if old != 'NULL' else old + + with conn[name].cursor() as cur: + sql = f"update reservoirs set pattern = '{pattern}' where id = '{id}'" + cur.execute(sql) + redo = sql.replace("'", '"') + undo = f'update reservoirs set pattern = {old} where id = "{id}"' + add_operation(name, redo, undo) + + change = ChangeSet() + change.update('reservoir', id, 'pattern', 'str', str(pattern)) + return change + + +def set_reservoir_coord(name: str, id: str, x: float, y: float) -> ChangeSet: + if not is_reservoir(name, id): + return + + old = get_reservoir_coord(name, id) + if old == None: + return + old_x, old_y = old['x'], old['y'] + + with conn[name].cursor() as cur: + sql = f"update coordinates set coord = '({x},{y})' where node = '{id}'" + cur.execute(sql) + + redo = sql.replace("'", '"') + undo = f'update coordinates set coord = "({old_x},{old_y})" where node = "{id}"' + add_operation(name, redo, undo) + + change = ChangeSet() + change.update('reservoir', id, 'coord', 'point', str({'x': x, 'y': y})) + return change \ No newline at end of file