fix startup crash if database not found

2 files changed, 24 insertions, 4 deletions

diff --git a/api.py b/api.py
index fd5dfe9..db7a881 100644
--- a/api.py
+++ b/api.py
@@ -16,9 +16,16 @@ def gen_word(N, min_N_dig, min_N_low):
     shuffle(chars)
     return ''.join(chars)
 
+dbFailing = False
+
 import mysql.connector as con
-mycon = con.connect(host=flask_app.config['DB_HOST'],user=flask_app.config['DB_USER'],password=flask_app.config['DB_PASSWORD'],port=flask_app.config['DB_PORT'],database=flask_app.config['DB_NAME'])
-mycursor = mycon.cursor()
+from mysql.connector.errors import InterfaceError
+try:
+    mycon = con.connect(host=flask_app.config['DB_HOST'],user=flask_app.config['DB_USER'],password=flask_app.config['DB_PASSWORD'],port=flask_app.config['DB_PORT'],database=flask_app.config['DB_NAME'])
+    mycursor = mycon.cursor()
+except InterfaceError:
+    print("Could not connect to the database!")
+    dbFailing = True
 
 """
 @flask_app.route("/")
@@ -38,6 +45,8 @@ async def API_Version():
 
 @app.get("/v1/status/{job_id}")
 async def get_status(job_id: str):
+    if dbFailing:
+        return {"message":"Could not connect to the database"}
     sqlQuery = 'select id, protein_name, ligand_name, date, description, done from curieweb where id="%s"' % (job_id)
     mycursor.execute(sqlQuery)
     records = mycursor.fetchall()
@@ -48,6 +57,8 @@ async def get_status(job_id: str):
 
 @app.get("/v1/3DModels/{job_id}")
 async def get_models(job_id: str):
+    if dbFailing:
+        return {"message":"Could not connect to the database"}
     sqlQuery = 'select done from curieweb where id="%s"' % (job_id)
     mycursor.execute(sqlQuery)
     records = mycursor.fetchall()
@@ -59,6 +70,8 @@ async def get_models(job_id: str):
 
 @app.get("/v1/Report/{job_id}")
 async def get_report(job_id:str):
+    if dbFailing:
+        return {"message":"Could not connect to the database"}
     sqlQuery = 'select done from curieweb where id="%s"' % (job_id)
     mycursor.execute(sqlQuery)
     records = mycursor.fetchall()
@@ -71,6 +84,8 @@ async def get_report(job_id:str):
 
 @app.post("/v1/docking/automatic")
 async def docking_automatic(pdb: str, smiles:str,compound_name:str,email:str,description:str):
+    if dbFailing:
+        return {"message":"Could not connect to the database"}
     if len(pdb) != 0:
         return {"message": "Invalid PDB ID"}
     
diff --git a/tests/backendTest.py b/tests/backendTest.py
index 0965e0a..12cefad 100644
--- a/tests/backendTest.py
+++ b/tests/backendTest.py
@@ -1,5 +1,6 @@
 import mysql.connector as con
-
+from mysql.connector.errors import InterfaceError
+import sys
 import configparser
 config = configparser.ConfigParser()
 config.read('config.ini')
@@ -9,7 +10,11 @@ try:
 except KeyError:
     config.read("../config.ini")
 
-mycon = con.connect(host=config['DATABASE']['HOST'],user=config['DATABASE']['USER'],password=config['DATABASE']['PASSWORD'],port=config['DATABASE']['PORT'],database=config['DATABASE']['NAME'])
+try:
+    mycon = con.connect(host=config['DATABASE']['HOST'],user=config['DATABASE']['USER'],password=config['DATABASE']['PASSWORD'],port=config['DATABASE']['PORT'],database=config['DATABASE']['NAME'])
+except InterfaceError:
+    print("Could not connect to the database")
+    sys.exit(1)
 mycursor = mycon.cursor()
 
 # If we are running the CI on an actual server, try using the 6LU7 Mpro and Geraniin Job ID because Eucalyptol fails