Add company-specific problem fetching

nicholasadamou · nicholasadamou · commit a58397a61186 · 2024-08-23T15:14:15.000-04:00
Introduced functionality to fetch LeetCode problems specific to a particular company. Updated `main.py` to handle a new `--company` argument and enhanced the `Database` and `LeetCode` classes with methods to check for and fetch company-related problems.
diff --git a/database/database.py b/database/database.py
@@ -201,6 +201,38 @@ def does_study_plan_exist(self, slug: str) -> bool:
         except Exception:
             return False
 
+    def does_company_exist(self, company: str) -> bool:
+        """
+        Check if a company exists in the database.
+        :param company: The name of the company.
+        :return: True if the company exists, False otherwise.
+        """
+        sql = """
+        SELECT EXISTS(SELECT 1 FROM leetcode.companies WHERE name = %(company)s);
+        """
+        self.cursor.execute(sql, {"company": company})
+
+        try:
+            result = self.cursor.fetchone()
+            return result[0]
+        except Exception:
+            return False
+
+    def get_problems_by_company(self, company: str) -> list[Problem]:
+        """
+        Get a list of problems by a specific company.
+        :param company: The name of the company.
+        :return: A list of Problem objects.
+        """
+        sql = """
+        SELECT p.id, p.question_id, p.title, p.slug, p.content, p.difficulty, p.topics, p.companies, p.hints, p.link
+        FROM leetcode.problems p
+        WHERE %(company)s = ANY(p.companies);
+        """
+        self.cursor.execute(sql, {"company": company})
+        results = self.cursor.fetchall()
+        return [Problem(*result) for result in results]
+
     def close(self):
         self.cursor.close()
         self.connection.close()
diff --git a/leetcode/api/client.py b/leetcode/api/client.py
@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, List
 
 import requests
 
@@ -163,3 +163,71 @@ 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"]
+
+    def get_recent_questions_for_company(
+        self,
+        company_slug: str,
+        timeframe: str = "six-months",
+        difficulties=None,
+        top_n: int = 50,
+    ) -> Dict[str, Any]:
+        """
+        Fetch the top N most recently asked questions for a specific company within a given timeframe.
+
+        :param company_slug: The slug of the company (e.g., 'microsoft').
+        :param timeframe: The timeframe for questions (e.g., 'last-30-days', 'three-months', 'six-months').
+        :param difficulties: The list of difficulties to filter the questions (default is ['EASY', 'MEDIUM']).
+        :param top_n: The number of top questions to retrieve (default is 50).
+        :return: A dictionary containing the list of questions.
+        :raises Exception: If the API request fails or the response does not contain expected data.
+
+        """
+        if difficulties is None:
+            difficulties = ["EASY", "MEDIUM"]
+
+        api_url = "https://leetcode.com/graphql"
+        favorite_slug = f"{company_slug}-{timeframe}"
+        query = """
+        query favoriteQuestionListForCompany($favoriteSlug: String!, $filter: FavoriteQuestionFilterInput) {
+          favoriteQuestionList(favoriteSlug: $favoriteSlug, filter: $filter) {
+            questions {
+              difficulty
+              id
+              paidOnly
+              questionFrontendId
+              status
+              title
+              titleSlug
+              translatedTitle
+              isInMyFavorites
+              frequency
+              topicTags {
+                name
+                nameTranslated
+                slug
+              }
+            }
+          }
+        }
+        """
+        variables = {
+            "favoriteSlug": favorite_slug,
+            "filter": {"difficultyList": difficulties, "positionRoleTagSlug": ""},
+        }
+        headers = self._get_headers()
+
+        response = requests.post(
+            api_url, json={"query": query, "variables": variables}, headers=headers
+        )
+        response.raise_for_status()  # Raise an exception for HTTP errors
+
+        response_data = response.json()
+        if (
+            "data" not in response_data
+            or "favoriteQuestionList" not in response_data["data"]
+            or "questions" not in response_data["data"]["favoriteQuestionList"]
+        ):
+            raise Exception("Questions not found or invalid response format")
+
+        # Limit to top N questions
+        return response_data["data"]["favoriteQuestionList"]["questions"][:top_n]
diff --git a/leetcode/leetcode.py b/leetcode/leetcode.py
@@ -2,7 +2,7 @@
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Dict
+from typing import Dict, List
 
 from requests import HTTPError, RequestException
 
@@ -46,9 +46,11 @@ def __init__(self, client: Client, database: Database = None):
 
         self.problems: Dict[str, Problem] = {}
         self.study_plans: Dict[str, StudyPlan] = {}
+        self.companies: Dict[str, List[Problem]] = {}
 
         self.problems_lock = threading.Lock()
         self.study_plans_lock = threading.Lock()
+        self.companies_lock = threading.Lock()
         self.database_lock = threading.Lock()
 
     def _fetch_and_store_problem(self, slug: str) -> Problem:
@@ -129,6 +131,45 @@ def get_problem(self, slug: str) -> Problem:
         """
         return self.problems[slug] if slug in self.problems else None
 
+    def fetch_and_store_company_problems(self, company: str) -> List[Problem]:
+        """
+        Fetch problems from LeetCode by the company tag and store them in the companies' dictionary.
+
+        :param company: The company tag.
+        :return: A dictionary of fetched Problem objects with the company tag as the key.
+        """
+        with self.companies_lock:
+            if company in self.companies:
+                print(f"Company {company} problems already fetched")
+                return self.companies[company]
+
+        with self.database_lock:
+            if self.database.does_company_exist(company):
+                company_problems = self.database.get_problems_by_company(company)
+                print(f"Company {company} problems already fetched")
+                with self.companies_lock:
+                    self.companies[company] = company_problems
+                return company_problems
+
+        company_problems = []
+
+        questions = _fetch_with_retries(
+            lambda: self.client.get_recent_questions_for_company(company)
+        )
+
+        if not questions:
+            raise Exception("No problems found for the company")
+
+        for question in questions:
+            slug = question["titleSlug"]
+            problem = self.get_problem(slug) or self._fetch_and_store_problem(slug)
+            company_problems[slug].append(problem)
+
+        with self.companies_lock:
+            self.companies[company] = company_problems
+
+        return company_problems
+
     def fetch_and_store_study_plan(self, plan_slug: str) -> StudyPlan:
         """
         Fetch a study plan from LeetCode by its slug and store it in the study plan dictionary.
diff --git a/main.py b/main.py
@@ -13,7 +13,7 @@
 load_dotenv()
 
 
-def main(csrf_token, leetcode_session, plans):
+def main(csrf_token, leetcode_session, plans, company):
     if not leetcode_session and not csrf_token:
         print(
             "Using Non-Premium LeetCode account. Some premium only data will not be returned."
@@ -29,14 +29,26 @@ def main(csrf_token, leetcode_session, plans):
     client = Client(configuration)
     leetcode = LeetCode(client, database=Database())
 
-    for plan in plans:
-        print(f"Fetching study plan problems: {plan}")
+    if plans:
+        for plan in plans:
+            print(f"Fetching study plan problems: {plan}")
+            try:
+                study_plan = leetcode.fetch_and_store_study_plan(plan)
+                print(
+                    "====================================================================================================="
+                )
+                print(study_plan)
+            except Exception as e:
+                print(e)
+
+    if company:
+        print(f"Fetching company related problems: {company}")
         try:
-            study_plan = leetcode.fetch_and_store_study_plan(plan)
+            company_problems = leetcode.fetch_and_store_company_problems(company)
             print(
                 "====================================================================================================="
             )
-            print(study_plan)
+            print(company_problems)
         except Exception as e:
             print(e)
 
@@ -70,10 +82,16 @@ def main(csrf_token, leetcode_session, plans):
         "--plans",
         type=str,
         nargs="+",
-        default=["top-interview-150", "leetcode-75"],
+        default=None,
         help="Slugs of the study plans to fetch",
     )
+    parser.add_argument(
+        "--company",
+        type=str,
+        default=None,
+        help="Company name to filter problems by",
+    )
 
     args = parser.parse_args()
 
-    main(args.csrf_token, args.leetcode_session, args.plans)
+    main(args.csrf_token, args.leetcode_session, args.plans, args.company)
