This repository has been archived on 2024-10-27. You can view files and clone it, but cannot push or open issues or pull requests.
local-llm-server/llm_server/database/conn.py

41 lines
1.3 KiB
Python

from mysql.connector import pooling
class Database:
__connection_pool = None
@classmethod
def initialise(cls, maxconn, **kwargs):
if cls.__connection_pool is not None:
raise Exception('Database connection pool is already initialised')
cls.__connection_pool = pooling.MySQLConnectionPool(pool_size=maxconn,
pool_reset_session=True,
**kwargs)
@classmethod
def get_connection(cls):
return cls.__connection_pool.get_connection()
@classmethod
def return_connection(cls, connection):
connection.close()
class CursorFromConnectionFromPool:
def __init__(self):
self.conn = None
self.cursor = None
def __enter__(self):
self.conn = Database.get_connection()
self.cursor = self.conn.cursor()
return self.cursor
def __exit__(self, exception_type, exception_value, exception_traceback):
if exception_value is not None: # This is equivalent of saying if there is an exception
self.conn.rollback()
else:
self.cursor.close()
self.conn.commit()
Database.return_connection(self.conn)