Add multi-threading for problem fetching and refactor code

Introduced multithreading to improve the efficiency of fetching and storing problems for study plans. Adapted the database schema to handle additional study plan details and provided functions for fetching problem counts and categories. Enhanced code readability by encapsulating details and resolving incorrect problem counts for study plans.

Author: nicholasadamou

.github/README.md

@@ -4,6 +4,7 @@
 ![Linux](https://img.shields.io/badge/-Linux-FCC624?style=flat-square&logo=linux&logoColor=black)
 ![Python](https://img.shields.io/badge/-Python-3776AB?style=flat-square&logo=python&logoColor=white)
 ![PostgreSQL](https://img.shields.io/badge/-PostgreSQL-336791?style=flat-square&logo=postgresql&logoColor=white)
+![Supabase](https://img.shields.io/badge/-Supabase-3ECF8E?style=flat-square&logo=supabase&logoColor=white)
 
 LeetCode Scraper is a Python-based tool designed to fetch and store details from LeetCode study plans into a PostgreSQL database. This tool leverages Docker for easy setup and environment management.
 
@@ -19,7 +20,7 @@ LeetCode Scraper is a Python-based tool designed to fetch and store details from
 ## Features
 
 - Fetches LeetCode problems and study plans
-- Stores data in a PostgreSQL database
+- Stores data in a PostgreSQL/Supabase database
 - Provides caching to reduce redundant requests
 - Handles rate limiting with retry mechanisms
 

database/database.py

@@ -8,6 +8,54 @@
 from leetcode.study_plan import StudyPlan
 
 
+def _create_study_plan_from_result(result) -> StudyPlan:
+    """
+    Create a StudyPlan object from a database result.
+    :param result: A tuple representing a row from the database query.
+    :return: A StudyPlan object.
+    """
+    slug = result[0]
+    name = result[1]
+    description = result[2]
+    expected_number_of_problems = result[3]
+    number_of_problems = result[4] if len(result) > 4 else None
+    number_of_categories = result[5] if len(result) > 5 else None
+
+    return StudyPlan(
+        slug=slug,
+        name=name,
+        description=description,
+        expected_number_of_problems=expected_number_of_problems,
+        number_of_problems=number_of_problems,
+        number_of_categories=number_of_categories,
+    )
+
+
+def _create_problem_from_result(result) -> Problem:
+    """
+    Create a Problem object from a database result.
+    :param result: A tuple representing a row from the database query.
+    :return: A Problem object.
+    """
+    problem_id = int(result[0])
+    title = result[1]
+    content = result[2]
+    difficulty = result[3]
+    topics = result[4] if result[4] else []
+    companies = result[5] if result[5] else []
+    hints = result[6] if result[6] else []
+
+    return Problem(
+        id=problem_id,
+        title=title,
+        content=content,
+        difficulty=difficulty,
+        topics=topics,
+        companies=companies,
+        hints=hints,
+    )
+
+
 def execute_insert(cursor, connection, sql: str, params: Dict[str, Any]) -> Any:
     """
     Execute an insert query and return the ID of the inserted row.
@@ -91,11 +139,12 @@ def insert_problem(self, problem: Problem) -> Any | None:
 
     def insert_study_plan(self, study_plan: StudyPlan) -> Any | None:
         sql = """
-        INSERT INTO leetcode.study_plans (slug, name, description)
-        VALUES (%(slug)s, %(name)s, %(description)s)
+        INSERT INTO leetcode.study_plans (slug, name, description, expected_number_of_problems)
+        VALUES (%(slug)s, %(name)s, %(description)s, %(expected_number_of_problems)s)
         ON CONFLICT (slug) DO UPDATE
         SET name = EXCLUDED.name,
-            description = EXCLUDED.description
+            description = EXCLUDED.description,
+            expected_number_of_problems = EXCLUDED.expected_number_of_problems
         RETURNING id;
         """
         return execute_insert(self.cursor, self.connection, sql, study_plan.to_dict())
@@ -134,13 +183,19 @@ def get_problem_by_slug(self, slug: str) -> Problem | None:
         :return: The Problem object with the given slug, or None if not found.
         """
         sql = """
-        SELECT * FROM leetcode.problems WHERE slug = %(slug)s;
+        SELECT id, title, content, difficulty, topics, companies, hints
+        FROM leetcode.problems
+        WHERE slug = %(slug)s;
         """
         self.cursor.execute(sql, {"slug": slug})
         result = self.cursor.fetchone()
+
         if result is None:
             return None
-        return Problem(*result)
+
+        problem = _create_problem_from_result(result)
+
+        return problem
 
     def get_study_plan_by_slug(self, slug: str) -> StudyPlan | None:
         """
@@ -149,25 +204,45 @@ def get_study_plan_by_slug(self, slug: str) -> StudyPlan | None:
         :return: The StudyPlan object with the given slug, or None if not found.
         """
         sql = """
-        SELECT sp.slug, sp.name, sp.description,
+        SELECT sp.slug, sp.name, sp.description, sp.expected_number_of_problems,
                COUNT(DISTINCT spp.problem_id) AS number_of_problems,
                COUNT(DISTINCT spp.category_name) AS number_of_categories
         FROM leetcode.study_plans sp
         LEFT JOIN leetcode.study_plan_problems spp ON sp.id = spp.study_plan_id
         WHERE sp.slug = %(slug)s
-        GROUP BY sp.slug, sp.name, sp.description;
+        GROUP BY sp.slug, sp.name, sp.description, sp.expected_number_of_problems;
         """
         self.cursor.execute(sql, {"slug": slug})
         result = self.cursor.fetchone()
+
         if result is None:
             return None
-        return StudyPlan(
-            slug=result[0],
-            name=result[1],
-            description=result[2],
-            number_of_problems=result[3],
-            number_of_categories=result[4],
-        )
+
+        return _create_study_plan_from_result(result)
+
+    def get_problems_by_study_plan_slug(self, slug: str) -> list[Problem]:
+        """
+        Get a list of problems for a study plan.
+        :param slug: The slug of the study plan.
+        :return: A list of Problem objects.
+        """
+        sql = """
+        SELECT p.id, p.title, p.content, p.difficulty, p.topics, p.companies, p.hints
+        FROM leetcode.problems p
+        JOIN leetcode.study_plan_problems spp ON p.id = spp.problem_id
+        JOIN leetcode.study_plans sp ON spp.study_plan_id = sp.id
+        WHERE sp.slug = %(slug)s;
+        """
+        self.cursor.execute(sql, {"slug": slug})
+        results = self.cursor.fetchall()
+
+        problems = []
+        for result in results:
+            problem = _create_problem_from_result(result)
+
+            problems.append(problem)
+
+        return problems
 
     def does_problem_exist(self, slug: str) -> bool:
         """
@@ -220,6 +295,46 @@ def does_company_exist(self, company: str) -> bool:
         except Exception:
             return False
 
+    def get_problem_count_by_study_plan(self, slug: str) -> int:
+        """
+        Get the number of problems in a study plan.
+        :param slug: The slug of the study plan.
+        :return: The number of problems in the study plan.
+        """
+        sql = """
+        SELECT COUNT(DISTINCT spp.problem_id)
+        FROM leetcode.study_plan_problems spp
+        JOIN leetcode.study_plans sp ON spp.study_plan_id = sp.id
+        WHERE sp.slug = %(slug)s;
+        """
+        self.cursor.execute(sql, {"slug": slug})
+        result = self.cursor.fetchone()
+
+        if result is None:
+            return 0
+
+        return result[0]
+
+    def get_category_count_by_study_plan(self, slug: str) -> int:
+        """
+        Get the number of categories in a study plan.
+        :param slug: The slug of the study plan.
+        :return: The number of categories in the study plan.
+        """
+        sql = """
+        SELECT COUNT(DISTINCT spp.category_name)
+        FROM leetcode.study_plan_problems spp
+        JOIN leetcode.study_plans sp ON spp.study_plan_id = sp.id
+        WHERE sp.slug = %(slug)s;
+        """
+        self.cursor.execute(sql, {"slug": slug})
+        result = self.cursor.fetchone()
+
+        if result is None:
+            return 0
+
+        return result[0]
+
     def get_problems_by_company(self, company: str) -> list[Problem]:
         """
         Get a list of problems by a specific company.
@@ -233,7 +348,11 @@ def get_problems_by_company(self, company: str) -> list[Problem]:
         """
         self.cursor.execute(sql, {"company": company})
         results = self.cursor.fetchall()
-        return [Problem(*result) for result in results]
+
+        if results is None:
+            return []
+
+        return [_create_problem_from_result(result) for result in results]
 
     def close(self):
         self.cursor.close()

leetcode/api/client.py

@@ -87,7 +87,7 @@ def get_study_plan_details(self, plan_slug: str) -> Dict[str, Any]:
         Fetch study plan details from LeetCode's GraphQL API using the provided plan slug.
 
         :param plan_slug: The slug of the study plan.
-        :return: A dictionary containing the study plan details.
+        :return: A dictionary containing the study plan details and the total number of problems.
         :raises Exception: If the API request fails or the response does not contain expected data.
         """
         api_url = "https://leetcode.com/graphql"
@@ -162,7 +162,17 @@ def get_study_plan_details(self, plan_slug: str) -> Dict[str, Any]:
         ):
             raise Exception("Study plan not found or invalid response format")
 
-        return response_data["data"]["studyPlanV2Detail"]
+        study_plan_details = response_data["data"]["studyPlanV2Detail"]
+
+        # Calculate the total number of problems in the study plan
+        total_problems = sum(
+            group.get("questionNum", 0) for group in study_plan_details["planSubGroups"]
+        )
+
+        # Add the total number of problems to the study plan details
+        study_plan_details["totalProblems"] = total_problems
+
+        return study_plan_details
 
     def get_recent_questions_for_company(
         self,