Add reservoir api

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