Merge pull request #9 from atlarge-research/mariadb

Migration to MariaDB

1 files changed, 14 insertions, 8 deletions

diff --git a/opendc/util/database.py b/opendc/util/database.py
index 2ef8b982..337f8fc7 100644
--- a/opendc/util/database.py
+++ b/opendc/util/database.py
@@ -3,40 +3,46 @@ import json
 import sqlite3
 import sys
 
+from mysql.connector.pooling import MySQLConnectionPool
+
 # Get keys from config file
 with open(sys.argv[1]) as file:
     KEYS = json.load(file)
 
 DATETIME_STRING_FORMAT = '%Y-%m-%dT%H:%M:%S'
+CONNECTION_POOL = None
+
+def init_connection_pool(user, password, database, host, port):
+    global CONNECTION_POOL
+    CONNECTION_POOL = MySQLConnectionPool(pool_name = "opendcpool", pool_size = 5, \
+                        user=user, password=password, database=database, host=host, port=port)
 
 def execute(statement, t):
     """Open a database connection and execute the statement."""
 
     # Connect to the database
-    connection = sqlite3.connect(KEYS['DATABASE_LOCATION'])
+    connection = CONNECTION_POOL.get_connection()
     cursor = connection.cursor()
-
-    # Turn on foreign key checks
-    cursor.execute('pragma foreign_keys=ON')
     
     # Execute the statement
     cursor.execute(statement, t)
 
     # Get the id
-    database_id = cursor.execute('SELECT last_insert_rowid()').fetchone()[0]
+    cursor.execute('SELECT last_insert_id();')
+    row_id = cursor.fetchone()[0]
 
     # Disconnect from the database
     connection.commit()
     connection.close()
 
     # Return the id
-    return database_id
+    return row_id
 
 def fetchone(statement, t=None):
     """Open a database connection and return the first row matched by the SELECT statement."""
 
     # Connect to the database
-    connection = sqlite3.connect(KEYS['DATABASE_LOCATION'])
+    connection = CONNECTION_POOL.get_connection()
     cursor = connection.cursor()
 
     # Execute the SELECT statement
@@ -56,7 +62,7 @@ def fetchall(statement, t=None):
     """Open a database connection and return all rows matched by the SELECT statement."""
 
     # Connect to the database
-    connection = sqlite3.connect(KEYS['DATABASE_LOCATION'])
+    connection = CONNECTION_POOL.get_connection()
     cursor = connection.cursor()
 
     # Execute the SELECT statement