convert from connection to pools in mysql

backend/src/db_client.py

@@ -1,7 +1,7 @@
 from enum import Enum
-import sys
 import mysql.connector
 import os
+from mysql.connector import pooling
 
 STARTUP_TABLE_CREATION_QUERIES = {
     "users": """CREATE TABLE IF NOT EXISTS users (
@@ -21,7 +21,8 @@ class Severity(Enum):
     ERROR = "ERROR"
 
 class DBClient:
-    def __init__(self):
+    def __init__(self, app):
+        self.app = app
         self.db_server = os.environ.get('DB_SERVER')
         self.db_port = os.environ.get('DB_PORT')
         self.user = 'root'
@@ -29,57 +30,54 @@ class DBClient:
         self.database = os.environ.get('DB_NAME')
 
         self.validate_env_variables()  # Check for required environment variables
-        self.connection = self.open()
-        self.cursor = self.connection.cursor()
+        self.pool = self.create_pool()  # Create a connection pool
         
         self.initialize_database()
 
     def validate_env_variables(self):
         if not self.db_server or not self.db_port or not self.password or not self.database:
-            self.error("Missing one or more environment variables.")
-    
-    def open(self):
-        return mysql.connector.connect(
+            self.app.logger.error("Missing one or more environment variables.")
+
+    def create_pool(self):
+        return pooling.MySQLConnectionPool(
+            pool_name="mypool",
+            pool_size=5,  # Adjust size as needed
             host=self.db_server,
             port=self.db_port,
             user=self.user,
             password=self.password,
-            database=self.database
+            database=self.database,
+            connection_timeout=10  # Timeout in seconds
         )
 
-    def close(self):
-        self.cursor.close()
-        self.connection.close()
-
     def initialize_database(self):
         self.query(STARTUP_TABLE_CREATION_QUERIES['users'])
         self.query(STARTUP_TABLE_CREATION_QUERIES['sessions'])
 
-    def query(self, query_str, params=None):
-        try:
-            self.info(f'Executing query: {query_str}')
-            self.cursor.execute(query_str, params)
-
-            if 'SELECT' in query_str:
-                return self.cursor.fetchall()  # Return results for SELECT queries
-            else:
-                self.commit()  # Commit if it's a non-SELECT query
-        except Exception as e:
-            self.error(f"Query failed: {str(e)}")
-
-    def commit(self):
-        self.info('Committing actions to DB')
-        self.connection.commit()
-
-    def info(self, message):
-        self.message(severity=Severity.INFO, message=message)
-
-    def warning(self, message):
-        self.message(severity=Severity.WARNING, message=message)
-
-    def error(self, message):
-        self.message(severity=Severity.ERROR, message=message)
-
-    def message(self, severity, message):
-        print(f'DBClient [{severity.value}]: {message}')
-       
+    def query(self, query_str, params=None):    
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                self.app.logger.info(f'Executing query: {query_str}')
+                # Get a connection from the pool
+                connection = self.pool.get_connection()
+                with connection.cursor() as cursor:
+                    cursor.execute(query_str, params)
+                    if 'SELECT' in query_str:
+                        results = cursor.fetchall()
+                        self.app.logger.info(f'Query results: {results}')
+                        return results
+                    else:
+                        connection.commit()
+                        self.app.logger.info('Query executed successfully, changes committed.')
+                        break
+            except mysql.connector.Error as e:
+                if e.errno in (2013, 2006):  # Lost connection or connection no longer available
+                    self.app.logger.warning(f"Lost connection to MySQL, retrying... {attempt + 1}/{max_retries}")
+                    continue  # Retry the query
+                else:
+                    self.app.logger.error(f"Query failed: {str(e)}")
+                    return None  # Handle the error accordingly
+            finally:
+                if connection.is_connected():
+                    connection.close()  # Close the connection to return it to the pool

backend/src/server.py

@@ -10,16 +10,18 @@ from flask_cors import CORS
 from dotenv import load_dotenv
 from db_client import DBClient
 from user import registerUserEndpoints
+import logging
 
 load_dotenv()
 app = Flask(__name__)
 app.config['JSON_AS_ASCII'] = False  # Ensures non-ASCII characters are preserved
+logging.basicConfig(level=logging.INFO)
 allowed_origins = [
     "https://nyipyatki.davydovcloud.com",
     "https://nyipyatki-dev.davydovcloud.com",
 ]
 CORS(app, resources={r"*": {"origins": allowed_origins}})  # Only allow example.com
-database = DBClient()
+database = DBClient(app)
 registerUserEndpoints(app=app, database=database)
 
 if __name__ == "__main__":

backend/src/user.py

@@ -14,8 +14,10 @@ def registerUserEndpoints(app, database):
     def user_is_set():
         user_name = request.args.get('userName')
         try:
+            app.logger.info(f'Searching for user {user_name}')
             query = "SELECT * FROM users WHERE Name=%s"
             result = database.query(query, params=(user_name,))
+            app.logger.info(f'Got: {result}')
             return jsonify(bool(result)), 200
         except mysql.connector.Error as err:
             # Log the error or handle it as necessary
@@ -79,6 +81,7 @@ def registerUserEndpoints(app, database):
         query = "SELECT * FROM sessions WHERE Token=%s AND Name=%s"
         try:
             result = database.query(query, params=(token, user_name))
+            app.logger.info(f'Got result: {result}')
             return jsonify(tokenValid=bool(result)), 200
         except Exception as e:
             return jsonify(success=False, message=str(e)), 500