Refactored DB connection sections

* Connect for each post instead of once per session
* Restart the SSH tunnel if it was closed
* Avoid issue where the connection(s) were broken by aborting Prime95

.bin/Scripts/functions/hw_diags.py

@@ -113,31 +113,33 @@ def connect_to_db():
         ]
 
     # Establish SSH tunnel unless one already exists
-    if not ost_db['Tunnel']:
+    if not ost_db['Tunnel'] or ost_db['Tunnel'].poll() is not None:
         ost_db['Tunnel'] = popen_program(cmd)
 
     # Establish SQL connection (try a few times in case SSH is slow)
     for x in range(5):
-        sleep(1)
+        sleep(2)
         try:
             ost_db['Connection'] = mariadb.connect(
                 user=DB_USER, password=DB_PASS, database=DB_NAME)
+            ost_db['Cursor'] = ost_db['Connection'].cursor()
         except:
             # Just try again
             pass
         else:
             break
-    ost_db['Cursor'] = ost_db['Connection'].cursor()
-    ost_db['Errors'] = False
 
-def disconnect_from_db():
+def disconnect_from_db(reset_errors=False):
-    """Disconnect from SQL DB and close SSH tunnel."""
+    """Disconnect from SQL DB."""
-    if ost_db['Cursor']:
+    for c in ['Cursor', 'Connection']:
-        ost_db['Cursor'].close()
+        try:
-    if ost_db['Connection']:
+            ost_db[c].close()
-        ost_db['Connection'].close()
+        except:
-    if ost_db['Tunnel']:
+            # Ignore
-        ost_db['Tunnel'].kill()
+            pass
+        ost_db[c] = None
+    if reset_errors:
+        ost_db['Errors'] = False
 
 def export_png_graph(name, dev):
     """Exports PNG graph using gnuplot, returns file path as str."""
@@ -358,6 +360,8 @@ def menu_diags(*args):
                 if not result['CS']:
                     print_warning('osTicket integration disabled for this run.')
                     pause()
+                ticket_number = get_osticket_number()
+                disconnect_from_db()
                 # Save log for non-quick tests
                 global_vars['Date-Time'] = time.strftime("%Y-%m-%d_%H%M_%z")
                 global_vars['LogDir'] = '{}/Logs/{}_{}'.format(
@@ -367,7 +371,6 @@ def menu_diags(*args):
                 os.makedirs(global_vars['LogDir'], exist_ok=True)
                 global_vars['LogFile'] = '{}/Hardware Diagnostics.log'.format(
                     global_vars['LogDir'])
-                ticket_number = get_osticket_number()
             run_tests(diag_modes[int(selection)-1]['Tests'], ticket_number)
         elif selection == 'A':
             run_program(['hw-diags-audio'], check=False, pipe=False)
@@ -391,7 +394,7 @@ def menu_diags(*args):
             break
 
     # Done
-    disconnect_from_db()
+    disconnect_from_db(reset_errors=True)
 
 def osticket_get_ticket_name(ticket_id):
     """Lookup ticket and return name as str."""
@@ -420,6 +423,9 @@ def osticket_needs_attention(ticket_id):
     return # This function has been DISABLED due to a repurposing of that flag
     if not ticket_id:
         raise GenericError
+
+    # Connect to DB
+    connect_to_db()
     if not ost_db['Cursor']:
         # Skip section
         return
@@ -435,11 +441,15 @@ def osticket_needs_attention(ticket_id):
         ost_db['Cursor'].execute(sql_cmd)
     except:
         ost_db['Errors'] = True
+    disconnect_from_db()
 
 def osticket_post_reply(ticket_id, response):
     """Post a reply to a ticket in osTicket."""
     if not ticket_id:
         raise GenericError
+
+    # Connect to DB
+    connect_to_db()
     if not ost_db['Cursor']:
         # Skip section
         return
@@ -458,11 +468,15 @@ def osticket_post_reply(ticket_id, response):
         ost_db['Cursor'].execute(sql_cmd)
     except:
         ost_db['Errors'] = True
+    disconnect_from_db()
 
 def osticket_set_drive_result(ticket_id, passed):
     """Marks the pass/fail box for the drive(s) in osTicket."""
     if not ticket_id:
         raise GenericError
+
+    # Connect to DB
+    connect_to_db()
     if not ost_db['Cursor']:
         # Skip section
         return
@@ -479,6 +493,7 @@ def osticket_set_drive_result(ticket_id, passed):
         ost_db['Cursor'].execute(sql_cmd)
     except:
         ost_db['Errors'] = True
+    disconnect_from_db()
 
 def pad_with_dots(s, left_pad=True):
     """Replace ' ' padding with '..' for osTicket posts."""