feat(api): add status returns for NLP processing

server/core/datasets.py

@@ -7,13 +7,16 @@ class DatasetManager:
     def __init__(self, db: PostgresConnector):
         self.db = db
 
-    def get_dataset_and_validate(self, dataset_id: int, user_id: int) -> pd.DataFrame:
+    def authorize_user_dataset(self, dataset_id: int, user_id: int) -> bool:
         dataset_info = self.get_dataset_info(dataset_id)
 
+        if dataset_info.get("user_id", None) == None:
+            return False
+
         if dataset_info.get("user_id") != user_id:
-            raise NotAuthorisedException("This user is not authorised to access this dataset")
+            return False
         
-        return self.get_dataset_content(dataset_id)
+        return True
 
     def get_dataset_content(self, dataset_id: int) -> pd.DataFrame:
         query = "SELECT * FROM events WHERE dataset_id = %s"
@@ -102,4 +105,36 @@ class DatasetManager:
             for _, row in event_data.iterrows()
         ]
 
-        self.db.execute_batch(query, values)
+        self.db.execute_batch(query, values)
+
+    def set_dataset_status(self, dataset_id: int, status: str, status_message: str | None = None):
+        if status not in ["processing", "complete", "error"]:
+            raise ValueError("Invalid status")
+
+        query = """
+            UPDATE datasets
+            SET status = %s,
+                status_message = %s,
+                completed_at = CASE
+                    WHEN %s = 'complete' THEN NOW()
+                    ELSE NULL
+                END
+            WHERE id = %s
+        """
+
+        self.db.execute(query, (status, status_message, status, dataset_id))
+
+    def get_dataset_status(self, dataset_id: int):
+        query = """
+            SELECT status, status_message, completed_at
+            FROM datasets
+            WHERE id = %s
+        """
+
+        result = self.db.execute(query, (dataset_id, ), fetch=True)
+        
+        if not result:
+            print(result)
+            raise NonExistentDatasetException(f"Dataset {dataset_id} does not exist")
+        
+        return result[0]