refactor: combine query results into one endpoint

server/app.py

@@ -1,4 +1,5 @@
 import os
+import datetime
 
 from dotenv import load_dotenv
 from flask import Flask, jsonify, request
@@ -42,7 +43,6 @@ auth_manager = AuthManager(db, bcrypt)
 
 stat_gen = StatGen()
 
-
 @app.route("/register", methods=["POST"])
 def register_user():
     data = request.get_json()
@@ -112,7 +112,7 @@ def upload_data():
     post_file = request.files["posts"]
     topic_file = request.files["topics"]
 
-    if post_file.filename == "" or topic_file == "":
+    if post_file.filename == "" or topic_file.filename == "":
         return jsonify({"error": "Empty filename"}), 400
 
     if not post_file.filename.endswith(".jsonl") or not topic_file.filename.endswith(
@@ -280,75 +280,34 @@ def get_interaction_analysis(dataset_id):
         return jsonify({"error": f"An unexpected error occurred: {str(e)}"}), 500
 
 
-# @app.route("/filter/query", methods=["POST"])
-# def filter_query():
-#     if stat_obj is None:
-#         return jsonify({"error": "No data uploaded"}), 400
+@app.route("/dataset/query", methods=["POST"])
+@jwt_required()
+def filter_query():
+    data = request.get_json()
 
-#     data = request.get_json(silent=True) or {}
+    if "query" in data:
+        stat_gen.set_search_query(data["query"])
 
-#     if "query" not in data:
-#         return jsonify(stat_obj.df.to_dict(orient="records")), 200
+    if "start" in data:
+        start_timestamp = datetime.datetime.fromisoformat(data["start"])
+        stat_gen.set_start_date(start_timestamp)
 
-#     query = data["query"]
-#     filtered_df = stat_obj.filter_by_query(query)
+    if "end" in data:
+        end_timestamp = datetime.datetime.fromisoformat(data["end"])
+        stat_gen.set_end_date(end_timestamp)
 
-#     return jsonify(filtered_df), 200
+    if "sources" in data:
+        data_sources = set(data["sources"])
+        stat_gen.set_data_sources(data_sources)
+
+    return jsonify({"message": "Filters set successfully"}), 200
 
 
-# @app.route("/filter/time", methods=["POST"])
-# def filter_time():
-#     if stat_obj is None:
-#         return jsonify({"error": "No data uploaded"}), 400
-
-#     data = request.get_json(silent=True)
-#     if not data:
-#         return jsonify({"error": "Invalid or missing JSON body"}), 400
-
-#     if "start" not in data or "end" not in data:
-#         return jsonify({"error": "Please include both start and end dates"}), 400
-
-#     try:
-#         start = pd.to_datetime(data["start"], utc=True)
-#         end = pd.to_datetime(data["end"], utc=True)
-#         filtered_df = stat_obj.set_time_range(start, end)
-#         return jsonify(filtered_df), 200
-#     except Exception:
-#         return jsonify({"error": "Invalid datetime format"}), 400
-
-
-# @app.route("/filter/sources", methods=["POST"])
-# def filter_sources():
-#     if stat_obj is None:
-#         return jsonify({"error": "No data uploaded"}), 400
-
-#     data = request.get_json(silent=True)
-#     if not data:
-#         return jsonify({"error": "Invalid or missing JSON body"}), 400
-
-#     if "sources" not in data:
-#         return jsonify({"error": "Ensure sources hash map is in 'sources' key"}), 400
-
-#     try:
-#         filtered_df = stat_obj.filter_data_sources(data["sources"])
-#         return jsonify(filtered_df), 200
-#     except ValueError:
-#         return jsonify({"error": "Please enable at least one data source"}), 400
-#     except Exception as e:
-#         return jsonify({"error": "An unexpected server error occured: " + str(e)}), 500
-
-
-# @app.route("/filter/reset", methods=["GET"])
-# def reset_dataset():
-#     if stat_obj is None:
-#         return jsonify({"error": "No data uploaded"}), 400
-
-#     try:
-#         stat_obj.reset_dataset()
-#         return jsonify({"success": "Dataset successfully reset"})
-#     except Exception as e:
-#         print(traceback.format_exc())
-#         return jsonify({"error": f"An unexpected error occurred: {str(e)}"}), 500
+@app.route("/database/query/reset", methods=["GET"])
+@jwt_required()
+def reset_dataset():
+    stat_gen.reset_filters()
+    return jsonify({"message": "Filters reset successfully"}), 200
 
 
 if __name__ == "__main__":

server/stat_gen.py

@@ -39,6 +39,37 @@ class StatGen:
         self.linguistic_analysis = LinguisticAnalysis(EXCLUDE_WORDS)
         self.cultural_analysis = CulturalAnalysis()
 
+        self.search_query = ""
+        self.start_date_filter = None
+        self.end_date_filter = None
+        self.data_source_filters = set()
+
+    ## Private Methods
+    def _prepare_filtered_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        filtered_df = df.copy()
+
+        if self.search_query:
+            mask = (
+                filtered_df["content"].str.contains(self.search_query, case=False, na=False)
+                | filtered_df["author"].str.contains(self.search_query, case=False, na=False).fillna(False)
+                | filtered_df["title"].str.contains(self.search_query, case=False, na=False, regex=False).fillna(False)
+            )
+            filtered_df = filtered_df[mask]
+
+        if self.start_date_filter and self.end_date_filter:
+            filtered_df = filtered_df[
+                (filtered_df["dt"] >= self.start_date_filter) & (filtered_df["dt"] <= self.end_date_filter)
+            ]
+
+        if self.data_source_filters:
+            enabled_sources = [src for src, enabled in self.data_source_filters.items() if enabled]
+            if enabled_sources:
+                filtered_df = filtered_df[filtered_df["source"].isin(enabled_sources)]
+
+        return filtered_df
+
+    ## Public Methods
+
     def get_time_analysis(self, df: pd.DataFrame) -> dict:
         return {
             "events_per_day": self.temporal_analysis.posts_per_day(df),
@@ -93,43 +124,20 @@ class StatGen:
             "sources": df["source"].dropna().unique().tolist(),
         }
 
-    # def filter_by_query(self, df: pd.DataFrame, search_query: str) -> dict:
-    #     filtered_df = df[df["content"].str.contains(search_query, na=False)]
+    def set_search_query(self, search_query: str) -> None:
+        self.search_query = search_query
 
-    #     return {
-    #         "rows": len(filtered_df),
-    #         "data": filtered_df.to_dict(orient="records"),
-    #     }
+    def set_start_date(self, start: datetime.datetime) -> None:
+        self.start_date_filter = start
 
-    # def set_time_range(
-    #     self,
-    #     original_df: pd.DataFrame,
-    #     start: datetime.datetime,
-    #     end: datetime.datetime,
-    # ) -> dict:
-    #     df = self._prepare_df(original_df)
-    #     filtered_df = df[(df["dt"] >= start) & (df["dt"] <= end)]
+    def set_end_date(self, end: datetime.datetime) -> None:
+        self.end_date_filter = end
 
-    #     return {
-    #         "rows": len(filtered_df),
-    #         "data": filtered_df.to_dict(orient="records"),
-    #     }
-
-    # def filter_data_sources(
-    #     self, original_df: pd.DataFrame, data_sources: dict
-    # ) -> dict:
-    #     df = self._prepare_df(original_df)
-    #     enabled_sources = [src for src, enabled in data_sources.items() if enabled]
-
-    #     if not enabled_sources:
-    #         raise ValueError("Please choose at least one data source")
-
-    #     filtered_df = df[df["source"].isin(enabled_sources)]
-
-    #     return {
-    #         "rows": len(filtered_df),
-    #         "data": filtered_df.to_dict(orient="records"),
-    #     }
-
-    # def reset_dataset(self, original_df: pd.DataFrame) -> pd.DataFrame:
-    #     return self._prepare_df(original_df)
+    def set_data_sources(self, data_sources: set) -> None:
+        self.data_source_filters = data_sources
+        
+    def reset_filters(self) -> None:
+        self.search_query = ""
+        self.start_date_filter = None
+        self.end_date_filter = None
+        self.data_source_filters = set()