refactor: clean up wiki edit implementation

* Add better error messages
* Create more modular functions

1 files changed, 48 insertions, 27 deletions

diff --git a/gn3/api/wiki.py b/gn3/api/wiki.py
index f96a3dc..2b8369f 100644
--- a/gn3/api/wiki.py
+++ b/gn3/api/wiki.py
@@ -1,14 +1,20 @@
 import datetime
+from typing import Any, Dict, List
 from flask import Blueprint, request, jsonify, current_app
 from gn3 import db_utils
 
+
 wiki = Blueprint("wiki", __name__)
 
 
-@wiki.route("/comments/<int:comment_id>/edit", methods=["POST"])
+class MissingDBDataException(Exception):
+    pass
+
+
+@wiki.route("/<int:comment_id>/edit", methods=["POST"])
 def edit_wiki(comment_id: int):
     # FIXME: attempt to check and fix for types here with relevant errors
-    payload = request.json
+    payload: Dict[str, Any] = request.json
     pubmed_ids = [str(x) for x in payload.get("pubmed_ids", [])]
 
     insert_dict = {
@@ -16,10 +22,7 @@ def edit_wiki(comment_id: int):
         "symbol": payload["symbol"],
         "PubMed_ID": " ".join(pubmed_ids),
         "comment": payload["comment"],
-        # does this need to be part of the payload or can we get this from session information
-        # e.g. https://github.com/genenetwork/genenetwork2/blob/0998033d0a7ea26ed96b00a360a334bae6de8c55/gn2/wqflask/oauth2/session.py#L22-L23
         "email": payload["email"],
-        # DB doesn't default to now
         "createtime": datetime.datetime.now(datetime.timezone.utc).strftime(
             "%Y-%m-%d %H:%M"
         ),
@@ -36,37 +39,55 @@ def edit_wiki(comment_id: int):
     """
     with db_utils.database_connection(current_app.config["SQL_URI"]) as conn:
         cursor = conn.cursor()
-        categories = get_categories(cursor)
-        category_ids = []
-        for category in payload["categories"]:
-            cat_id = categories.get(category.strip())
-            if cat_id is None:
-                return jsonify(error=f"Error editting wiki entry, category with Name={category} not found"), 500
-            category_ids.append(cat_id)
-        cursor.execute("SELECT SpeciesID from Species  WHERE Name = %s", (payload["species"],))
-        species_ids = cursor.fetchall()
-        if len(species_ids) != 1:
-            return jsonify(error=f"Error editting wiki entry, expected 1 species with Name={payload['species']} but found {len(species_ids)}!"), 500
-        insert_dict["SpeciesID"] = species_ids[0][0]
-
-        cursor.execute("SELECT MAX(versionId) as version_id from GeneRIF WHERE Id = %s", (comment_id,))
-        latest_version = cursor.fetchone()[0]
-        if latest_version is None:
-            return jsonify(error=f"Error editting wiki entry, No comments found with comment_id={comment_id}"), 500
-        insert_dict["versionId"] = latest_version + 1
+        try:
+            category_ids = get_categories_ids(cursor, payload["categories"])
+            species_id = get_species_id(cursor, payload["species"])
+            next_version = get_next_comment_version(cursor, comment_id)
+        except MissingDBDataException as missing_exc:
+            return jsonify(error=f"Error editting wiki entry, {missing_exc}"), 500
+        insert_dict["SpeciesID"] = species_id
+        insert_dict["versionId"] = next_version
         current_app.logger.debug(f"Running query: {insert_query}")
         cursor.execute(insert_query, insert_dict)
-
         category_addition_query = "INSERT INTO GeneRIFXRef (GeneRIFId, versionId, GeneCategoryId) VALUES (%s, %s, %s)"
 
         for cat_id in category_ids:
             current_app.logger.debug(f"Running query: {category_addition_query}")
-            cursor.execute(category_addition_query, (comment_id, insert_dict["versionId"], cat_id))
+            cursor.execute(
+                category_addition_query, (comment_id, insert_dict["versionId"], cat_id)
+            )
         return jsonify({"success": "ok"})
     return jsonify(error="Error editting wiki entry, most likely due to DB error!"), 500
 
 
-def get_categories(cursor) -> dict:
+def get_species_id(cursor, species_name: str) -> int:
+    cursor.execute("SELECT SpeciesID from Species  WHERE Name = %s", (species_name,))
+    species_ids = cursor.fetchall()
+    if len(species_ids) != 1:
+        raise MissingDBDataException(
+            f"expected 1 species with Name={species_name} but found {len(species_ids)}!"
+        )
+    return species_ids[0][0]
+
+
+def get_next_comment_version(cursor, comment_id: int) -> int:
+    cursor.execute(
+        "SELECT MAX(versionId) as version_id from GeneRIF WHERE Id = %s", (comment_id,)
+    )
+    latest_version = cursor.fetchone()[0]
+    if latest_version is None:
+        raise MissingDBDataException(f"No comment found with comment_id={comment_id}")
+    return latest_version + 1
+
+
+def get_categories_ids(cursor, categories: List[str]) -> List[int]:
     cursor.execute("SELECT Name, Id from GeneCategory")
     raw_categories = cursor.fetchall()
-    return dict(raw_categories)
+    dict_cats = dict(raw_categories)
+    category_ids = []
+    for category in set(categories):
+        cat_id = dict_cats.get(category.strip())
+        if cat_id is None:
+            raise MissingDBDataException(f"Category with Name={category} not found")
+        category_ids.append(cat_id)
+    return category_ids