Update app.py

app.py

@@ -10,6 +10,7 @@ import sys
 from datetime import datetime
 import time
 import logging
+import threading
 
 # Set up logging
 logging.basicConfig(
@@ -34,6 +35,7 @@ DB_CONFIG = {
 # Global variables for model and tokenizer
 GLOBAL_MODEL = None
 GLOBAL_TOKENIZER = None
+db_connection_status = False
 
 def initialize_model():
     """Initialize model and tokenizer globally"""
@@ -56,6 +58,7 @@ def initialize_model():
 
 def test_db_connection():
     """Test database connection with timeout"""
+    global db_connection_status
     try:
         logging.info("Testing database connection...")
         connection = mysql.connector.connect(
@@ -69,9 +72,11 @@ def test_db_connection():
             db_name = cursor.fetchone()[0]
             cursor.close()
             connection.close()
+            db_connection_status = True
             logging.info(f"Successfully connected to MySQL Server version {db_info} - Database: {db_name}")
             return True, f"Successfully connected to MySQL Server version {db_info}\nDatabase: {db_name}"
     except Error as e:
+        db_connection_status = False
         logging.error(f"Error connecting to MySQL database: {e}")
         return False, f"Error connecting to MySQL database: {e}"
     return False, "Unable to establish database connection"
@@ -187,19 +192,25 @@ def format_result(query_result):
     
     return "\n".join(results)
 
+def check_live_connection():
+    """Check the database connection status periodically."""
+    while True:
+        test_db_connection()
+        time.sleep(10)  # Check every 10 seconds
+
 def main():
     """Main function with Streamlit UI components"""
     st.title("Natural Language to SQL Query")
     st.write("Ask questions about pizza sales data in plain English.")
     
-    # Test and display database connection status
-    db_success, db_message = test_db_connection()
-    st.write(db_message)
+    # Start the live connection check in a separate thread
+    threading.Thread(target=check_live_connection, daemon=True).start()
     
-    if not db_success:
-        logging.error("Could not connect to the database. Exiting.")
-        st.write("Could not connect to the database. Exiting.")
-        return
+    # Test and display database connection status
+    if db_connection_status:
+        st.success("Database connection is live.")
+    else:
+        st.error("Database connection is not live.")
     
     # Initialize model
     initialize_model()