From da5d7597fe491738e3c4329f2425f5840ba9cde7 Mon Sep 17 00:00:00 2001
From: Xevion <xevioni@yandex.com>
Date: Thu, 21 Jan 2021 15:52:00 -0600
Subject: [PATCH] improve message recording, use try/finally cur cursor
 closing, on-demand cursor creation

---
 server/db.py | 38 ++++++++++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/server/db.py b/server/db.py
index 98c17e8..ef333b0 100644
--- a/server/db.py
+++ b/server/db.py
@@ -1,6 +1,6 @@
+import logging
 import sqlite3
 from typing import List
-import logging
 
 import constants
 
@@ -10,7 +10,6 @@ conn = sqlite3.connect(constants.DATABASE)
 logger.debug(f"Connected to '{constants.DATABASE}'")
 
 logger.debug("Constructing 'message' table.")
-cur = conn.cursor()
 conn.execute('''CREATE TABLE IF NOT EXISTS message
             (id INTEGER PRIMARY KEY,
             nickname TEXT NOT NULL,
@@ -21,14 +20,33 @@ conn.execute('''CREATE TABLE IF NOT EXISTS message
 conn.commit()
 
 
-def add_message(nickname: str, hash: str, color: str, message: str, timestamp: int):
-    cur.execute('''INSERT INTO message (nickname, connection_hash, color, message, timestamp)
-                VALUES (?, ?, ?, ?, ?)''', [nickname, hash, color, message, timestamp])
-    conn.commit()
+def add_message(nickname: str, user_hash: str, color: str, message: str, timestamp: int) -> int:
+    """
+    Insert a message into the database. Returns the message ID.
+
+    :param nickname: A non-unique identifier for the user.
+    :param user_hash: A unique hash (usually) denoting the sender's identity.
+    :param color: The color of the user who sent the message.
+    :param message: The string content of the message echoed to all clients.
+    :param timestamp: The epoch time of the sent message.
+    :return: The unique integer primary key chosen for the message, i.e. it's ID.
+    """
+    cur = conn.cursor()
+    try:
+        cur.execute('''INSERT INTO message (nickname, connection_hash, color, message, timestamp)
+                    VALUES (?, ?, ?, ?, ?)''', [nickname, user_hash, color, message, timestamp])
+        conn.commit()
+        logger.debug(f'Message {cur.lastrowid} recorded.')
+        return cur.lastrowid
+    finally:
+        cur.close()
 
 
 def get_messages(columns: List[str] = None):
-    if columns is None:
-        cur.execute('''SELECT * FROM message''')
-        return cur.fetchall()
-
+    cur = conn.cursor()
+    try:
+        if columns is None:
+            cur.execute('''SELECT * FROM message''')
+            return cur.fetchall()
+    finally:
+        cur.close()