1 files changed, 23 insertions, 22 deletions
diff --git a/src/jwebmail/read_mails.py b/src/jwebmail/read_mails.py
index 44eba8d..404a242 100644
--- a/src/jwebmail/read_mails.py
+++ b/src/jwebmail/read_mails.py
@@ -1,3 +1,4 @@
+from contextlib import closing
 from datetime import datetime, timedelta
 
 from flask import current_app, g
@@ -36,7 +37,7 @@ class RedisTimeoutSession:
 
 
 class MysqlTimeoutSession:
-    def __init__(self, username, passwd, timeout, port=3306):
+    def __init__(self, username, passwd, timeout, database="jwebmaildb1", port=3306):
         import mysql.connector
 
         self.timeout = timeout
@@ -45,35 +46,35 @@ class MysqlTimeoutSession:
             port=port,
             username=username,
             password=passwd,
-            database="jwebmaildb1",
+            database=database,
         )
 
     def set(self, key, value):
         timeout = datetime.now() + timedelta(seconds=self.timeout)
-        cur = self.conn.cursor()
-        cur.execute("DELETE FROM session WHERE user = %s", [key])
-        cur.execute("INSERT INTO session VALUES (%s, %s, %s)", [key, value, timeout])
-        self.conn.commit()
-        cur.close()
 
-    def get(self, key):
-        cur = self.conn.cursor()
-        cur.execute("DELETE FROM session WHERE timeout < NOW()")
-        cur.execute("SELECT password FROM session WHERE user = %s", [key])
-        row = cur.fetchone()
-
-        if row is None:
-            self.conn.commit()
-            cur.close()
-            return None
-        else:
-            timeout = datetime.now() + timedelta(seconds=self.timeout)
+        with closing(self.conn.cursor()) as cur:
+            cur.execute("DELETE FROM session WHERE user = %s", [key])
             cur.execute(
-                "UPDATE session SET timeout = %s WHERE user = %s", [timeout, key]
+                "INSERT INTO session VALUES (%s, %s, %s)", [key, value, timeout]
             )
             self.conn.commit()
-            cur.close()
-            return row[0]
+
+    def get(self, key):
+        with closing(self.conn.cursor()) as cur:
+            cur.execute("DELETE FROM session WHERE timeout < NOW()")
+            cur.execute("SELECT password FROM session WHERE user = %s", [key])
+            row = cur.fetchone()
+
+            if row is None:
+                self.conn.commit()
+                return None
+            else:
+                timeout = datetime.now() + timedelta(seconds=self.timeout)
+                cur.execute(
+                    "UPDATE session SET timeout = %s WHERE user = %s", [timeout, key]
+                )
+                self.conn.commit()
+                return row[0]
 
 
 def select_timeout_session():