Upload conversationDB.py

Updated the methods of the DataBase class

conversationDB.py

@@ -35,7 +35,7 @@ class DataBase():
         db.mysql_password=mysql_password
         db.mysql_db=mysql_db
     def getmessages(db,username = ""):
-        # start the SSH tunnel using sshtunnel
+        # get the entire conversation with the specific user
         with SSHTunnelForwarder(
                 (db.ssh_host, db.ssh_port),
                 ssh_username=db.ssh_user,
@@ -56,7 +56,7 @@ class DataBase():
                 cursor.execute(query)
                 rows = cursor.fetchall()
 
-                messages = [message[2] for message in rows]
+                messages = [(message[2],message[3]) for message in rows]
 
             # close the MySQL connection
             mysql_conn.close()
@@ -64,9 +64,8 @@ class DataBase():
         # close the SSH client
         db.ssh_client.close()
         return messages
-
-    def setmessages(db,username, message_text):
-        # start the SSH tunnel using sshtunnel
+    def setmessages(db,username, message_text="", bot_reply=""):
+        # Adding the record into the database
         with SSHTunnelForwarder(
                 (db.ssh_host, db.ssh_port),
                 ssh_username=db.ssh_user,
@@ -83,7 +82,8 @@ class DataBase():
 
             # send a query to the MySQL database and print the response table
             with mysql_conn.cursor() as cursor:
-                query = f"INSERT INTO {dbTableName} (username,message_text) VALUES ('{username}','{message_text}');"
+                query = f"INSERT INTO {dbTableName} (username,message_text,bot_reply) " \
+                        f"VALUES ('{username}','{message_text}','{bot_reply}');"
                 cursor.execute(query)
                 mysql_conn.commit()
 
@@ -92,11 +92,45 @@ class DataBase():
 
         # close the SSH client
         db.ssh_client.close()
+    def cleanup(db, username = "", remaining_messages = 3):
+        # Cleanup the records, except the last N rows
+        with SSHTunnelForwarder(
+                (db.ssh_host, db.ssh_port),
+                ssh_username=db.ssh_user,
+                ssh_pkey=db.private_key,
+                remote_bind_address=(db.mysql_host, db.mysql_port)) as tunnel:
+            # connect to the MySQL server through the SSH tunnel
+            mysql_conn = pymysql.connect(
+                host='localhost',  # will be due to the SSH tunneling issue
+                port=tunnel.local_bind_port,
+                user=db.mysql_user,
+                password=db.mysql_password,
+                db=db.mysql_db
+            )
+
+            # send a query to the MySQL database to delete the records except the last ones
+            with mysql_conn.cursor() as cursor:
+
+                query = f"DELETE FROM {dbTableName} WHERE username = '{username}' AND id NOT IN " \
+                        f"(SELECT id FROM (SELECT id FROM (SELECT id FROM {dbTableName} " \
+                        f"WHERE username = '{username}' ORDER BY id DESC LIMIT {remaining_messages}) " \
+                        f"subquery) subsubquery)"
+                cursor.execute(query)
+                mysql_conn.commit()
+
+            # close the MySQL connection
+            mysql_conn.close()
+
+        # close the SSH client
+        db.ssh_client.close()
+
 
 if __name__ == '__main__':
-    username = 'user2'
+    # This is for testing purpose only:
+    username = 'user1'
 
     db = DataBase()
-    db.setmessages(username='user2',message_text='some message')
+    db.setmessages(username='user2',message_text='some message',bot_reply='some reply')
+    #db.cleanup(username='user2',remaining_messages=1)
     messages = db.getmessages(username)
     print(messages)