From bb93a5181772c798bf5c982faeba8f2315938ebf Mon Sep 17 00:00:00 2001
From: wqy <wqyuyu@outlook.com>
Date: Sat, 17 Sep 2022 09:25:21 +0800
Subject: [PATCH] Refactor reservoir

---
 api/s3_reservoirs.py | 108 +++++++++++--------------------------------
 1 file changed, 27 insertions(+), 81 deletions(-)

diff --git a/api/s3_reservoirs.py b/api/s3_reservoirs.py
index fd3f4b9..7dc649c 100644
--- a/api/s3_reservoirs.py
+++ b/api/s3_reservoirs.py
@@ -4,74 +4,42 @@ from .s0_base import *
 from .operation import *
 from .change_set import ChangeSet
 from .s24_coordinates import *
+import utility
 
 
 def add_reservoir(name: str, id: str, x: float, y: float, head: float) -> ChangeSet:
-    if is_node(name, id):
-        return
+    sql = f"insert into reservoirs (id, head) values ('{id}', {head});"
+    undo_sql = f'delete from reservoirs where id = "{id}";'
+    return add_node(name, RESERVOIR, id, x, y, sql, undo_sql)
 
-    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 _get_reservoir(name: str, id: str) -> Row | None:
+    return utility.query(f"select head, pattern from reservoirs where id = '{id}'")
 
 
 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
+    row = _get_reservoir(name, id)
+    if row == None:
+        return
 
-        head = row['head']
-        pattern = 'NULL' if row['pattern'] == None else row['pattern']
-        pattern = f'"{pattern}"' if pattern != 'NULL' else pattern
+    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
+    sql = f"delete from reservoirs where id = '{id}';"
+    undo_sql = f'insert into reservoirs (id, head, pattern) values ("{id}", {head}, {pattern});'
 
-        coord = row['coord']
+    return delete_node(name, RESERVOIR, id, sql, undo_sql)
 
-        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:
@@ -79,50 +47,28 @@ def get_reservoir_pattern(name: str, id: str) -> str | None:
     else:
         return None
 
+
 def get_reservoir_coord(name: str, id: str) -> dict[str, float] | None:
     return get_node_coord(name, id)
 
-def set_reservoir_head(name: str, id: str, head: float) -> ChangeSet:
+
+def _set_reservoir(name: str, id: str, key: str, key_type: str, value: str, optional: bool = False) -> ChangeSet:
     if not is_reservoir(name, id):
         return
 
-    old = get_reservoir_head(name, id)
-    if old == None:
+    row = _get_reservoir(name, id)
+    if row == 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)
+    return utility.update(name, RESERVOIR, 'reservoirs', 'id', id, key, key_type, row[key], value, optional)
+
+
+def set_reservoir_head(name: str, id: str, head: float) -> ChangeSet:
+    return _set_reservoir(name, id, 'head', 'float', str(head))
 
-    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
+    return _set_reservoir(name, id, 'pattern', 'str', pattern, True)
 
 
 def set_reservoir_coord(name: str, id: str, x: float, y: float) -> ChangeSet: