From 2493c6d4657d5a0235e540fc7a622b403e2c971a Mon Sep 17 00:00:00 2001
From: Dylan De Faoite <dylanseandefaoite@gmail.com>
Date: Sun, 1 Mar 2026 15:01:07 +0000
Subject: [PATCH] feat: add save_dataset method to db

---
 db/database.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/db/database.py b/db/database.py
index 4f2209e..c9e60ed 100644
--- a/db/database.py
+++ b/db/database.py
@@ -1,5 +1,6 @@
 import os
 import psycopg2
+import pandas as pd
 from psycopg2.extras import RealDictCursor
 
 
@@ -49,6 +50,12 @@ class PostgresConnector:
         return result[0] if result else None
     
     # Dataset Management Methods
+    def save_dataset(self, user_id: int, dataset_name: str, dataset_content: pd.DataFrame, topics: dict):
+        query = """
+            INSERT INTO datasets (user_id, name, content, topics)
+            VALUES (%s, %s, %s, %s)
+        """
+        self.execute(query, (user_id, dataset_name, dataset_content.to_json(orient="records"), topics))
 
     def close(self):
         if self.connection: