Moving to ORM so can switch between mysql and sqlite easily

.gitignore

@@ -3,3 +3,4 @@ venv
 
 *.orig
 *.log
+*.pyc

config/config.example.ini

@@ -7,4 +7,4 @@ DB_NAME=quarter_data
 
 [VARS]
 # Example: QUARTERS=["SP19", "WI19"]
-QUARTERS=["SP20"]
+QUARTERS=["SP20", "WI18", "SP17"]

database/db_interface.py

@@ -0,0 +1,4 @@
+from dictalchemy import DictableModel
+from sqlalchemy.ext.declarative import declarative_base
+
+Base = declarative_base(cls=DictableModel)

scraper_impl/course_scraper.py

@@ -12,9 +12,12 @@
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.select import Select
 from selenium.webdriver.support.wait import WebDriverWait
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session, sessionmaker
 
+from utils.models import Department
 from utils.scraper_util import Browser
-from settings import COURSES_HTML_PATH, QUARTERS_TO_SCRAPE
+from settings import COURSES_HTML_PATH, SQLITE_STR
 from settings import DATABASE_PATH, DATABASE_DIR
 from settings import SCHEDULE_OF_CLASSES_URL
 from settings import TIMEOUT, DEPT_SEARCH_TIMEOUT
@@ -160,21 +163,18 @@ def save_page(self, department, quarter, page_contents, num_page):
 
 
 class CourseScraper:
-    def __init__(self):
-        # Connecting to the database for the list of departments
-        os.makedirs(DATABASE_DIR, exist_ok=True)
-        self.database = sqlite3.connect(DATABASE_PATH)
-        self.cursor = self.database.cursor()
+    def __init__(self, quarter):
+        self.engine = create_engine(SQLITE_STR)
+        self.quarter = quarter
+
+        self.session: Session
+        self.session = sessionmaker(bind=self.engine)()
 
         self.department_queue = queue.Queue()
 
-        for quarter in QUARTERS_TO_SCRAPE:
-            self.cursor.execute("SELECT DEPT_CODE FROM DEPARTMENT WHERE QUARTER=?", (quarter,))
-            # fetching the data returns a tuple with one element,
-            # so using list comprehension to convert the data
-            self.departments = [i[0] for i in self.cursor.fetchall()]
-            for department in self.departments:
-                self.department_queue.put({"department": department, "quarter": quarter})
+        departments = [i.dept_code for i in self.session.query(Department).filter(Department.quarter == self.quarter)]
+        for department in departments:
+            self.department_queue.put({"department": department, "quarter": quarter})
 
         # Recreate top level folder
         if os.path.exists(COURSES_HTML_PATH):

scraper_impl/department_scraper.py

@@ -1,72 +1,71 @@
-import os
-import sqlite3
+import logging
+from typing import List
 
+from selenium.common.exceptions import NoSuchElementException
 from selenium.webdriver.support.select import Select
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
 
-from settings import DATABASE_PATH, DATABASE_DIR, QUARTERS_TO_SCRAPE, DEPARTMENT_URL
+from settings import DEPARTMENT_URL, SQLITE_STR
+from utils.models import Department
 from utils.scraper_util import get_browser
 
 
 class DepartmentScraper:
     INFO_MAX_INDEX = 4
 
-    def __init__(self):
+    def __init__(self, quarter):
         # Start up the browser
         self.browser = get_browser()
+        self.quarter = quarter
+        self.engine = create_engine(SQLITE_STR)
+        self.session = sessionmaker(bind=self.engine)()
 
         # Add an implicit wait so that the department options load
         self.browser.implicitly_wait(15)
 
-        # Establish database connection
-        os.makedirs(DATABASE_DIR, exist_ok=True)
-        self.database = sqlite3.connect(DATABASE_PATH)
-        self.database.row_factory = sqlite3.Row
-        self.cursor = self.database.cursor()
-
-    def create_table(self):
-        self.cursor.execute('DROP TABLE IF EXISTS DEPARTMENT')
-        self.cursor.execute('CREATE TABLE IF NOT EXISTS DEPARTMENT (QUARTER TEXT, DEPT_CODE TEXT)')
-
     def scrape(self):
         print("Beginning department scraping.")
-        self.create_table()
-
-        for quarter in QUARTERS_TO_SCRAPE:
-            print("Scraping departments for %s" % quarter)
-            self.search(quarter)
-            departments = self.get_departments()
-            self.insert_departments(quarter, departments)
-            print("Finished scraping departments for %s" % quarter)
+        print("Scraping departments for %s" % self.quarter)
 
+        self.search()
+        self.create_tables()
+        departments = self.get_departments()
+        self.session.add_all(departments)
         self.close()
+
+        print("Finished scraping departments for %s" % self.quarter)
         print("Finished department scraping.")
 
-    def search(self, quarter):
+    def create_tables(self):
+        Department.__table__.create(self.engine, checkfirst=True)
+        self.session.query(Department).filter(Department.quarter == self.quarter).delete()
+
+    def search(self):
         self.browser.get(DEPARTMENT_URL)
         select = Select(self.browser.find_element_by_id('selectedTerm'))
-        select.select_by_value(quarter)
+        try:
+            select.select_by_value(self.quarter)
+        except NoSuchElementException:
+            logging.error(f"Could not find data for {self.quarter}")
 
-    def get_departments(self):
+    def get_departments(self) -> List[Department]:
+        ret: List[Department]
         ret = []
         departments = self.browser.find_element_by_id('selectedSubjects') \
             .find_elements_by_tag_name('option')
         for department in departments:
             department = department.text
-            # Get first four elements
+            # Get first four characters
             department = department[:DepartmentScraper.INFO_MAX_INDEX]
             # Making sure department is in the correct format
             ret.append(self.normalize_department(department))
 
         return ret
 
-    def insert_departments(self, quarter, departments):
-        for department in departments:
-            self.cursor.execute('INSERT INTO DEPARTMENT(QUARTER, DEPT_CODE) VALUES(?, ?)', (quarter, department))
-
-    def normalize_department(self, department):
-        return department.strip()
+    def normalize_department(self, department) -> Department:
+        return Department(quarter=self.quarter, dept_code=department.strip())
 
     def close(self):
-        self.database.commit()
-        self.database.close()
+        self.session.commit()
         self.browser.close()

sd_cleaner/course_cleaner.py

@@ -1,11 +1,12 @@
-import copy
 import itertools
-import sqlite3
 from itertools import groupby
 from typing import List, Dict
 
-from sd_parser.course_parser import ClassRow
-from settings import DATABASE_PATH, QUARTERS_TO_SCRAPE
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, Session
+
+from settings import SQLITE_STR
+from utils.models import ClassRow, Department
 from utils.timeutils import TimeIntervalCollection
 
 """
@@ -14,49 +15,43 @@
 
 
 class CourseCleaner:
-    def __init__(self):
-        self.database = sqlite3.connect(DATABASE_PATH)
-        # will set the return value to a dict
-        self.database.row_factory = sqlite3.Row
-        self.cursor = self.database.cursor()
+    def __init__(self, quarter):
+        self.engine = create_engine(SQLITE_STR)
+        self.quarter = quarter
+
+        self.session: Session
+        self.session = sessionmaker(bind=self.engine)()
 
-    def clean(self, data: Dict[str, List[ClassRow]]):
+    def clean(self, data: List[ClassRow]):
         print('Begin cleaning database.')
-        self.setup_tables()
+        self.create_tables()
         self._clean(data)
-        self.close()
+        self.session.commit()
         print('Finished cleaning database.')
 
-    def setup_tables(self):
-        for quarter in QUARTERS_TO_SCRAPE:
-            self.cursor.execute("DROP TABLE IF EXISTS {}".format(quarter))
-            self.cursor.execute("CREATE TABLE {}"
-                                "(DEPARTMENT TEXT, COURSE_NUM TEXT, SECTION_ID TEXT, COURSE_ID TEXT,"
-                                "SECTION_TYPE TEXT, DAYS TEXT, TIME TEXT, LOCATION TEXT, ROOM TEXT, "
-                                "INSTRUCTOR TEXT, DESCRIPTION TEXT, UNITS TEXT, FOREIGN KEY (DEPARTMENT)"
-                                "REFERENCES DEPARTMENT(DEPT_CODE))".format(quarter))
+    def create_tables(self):
+        Department.__table__.create(self.engine, checkfirst=True)
+        ClassRow.__table__.create(self.engine, checkfirst=True)
+        self.session.query(ClassRow).filter(ClassRow.quarter == self.quarter).delete()
 
-    def _clean(self, data: Dict[str, List[ClassRow]]):
+    def _clean(self, data: List[ClassRow]):
         # getting list of departments
-        self.cursor.execute("SELECT * FROM DEPARTMENT")
-        departments = [i["DEPT_CODE"] for i in self.cursor.fetchall()]
+        departments = [i.dept_code for i in self.session.query(Department).filter(Department.quarter == self.quarter)]
 
-        # handle each department and quarter separately
-        for quarter in QUARTERS_TO_SCRAPE:
-            quarter_data = data[quarter]
-            print("Cleaning quarter {}".format(quarter))
-            for department in departments:
-                classes = self.process_department(department, quarter_data)
-                self.save_classes(classes, quarter)
+        print("Cleaning quarter {}".format(self.quarter))
+        for department in departments:
+            classes = self.process_department(department, data)
+            self.save_classes(classes)
 
     """
     Will store in format with partitions for the courseNums in the same format : i.e CSE3$0 means section 0 of CSE3.
     """
 
     def process_department(self, department, data: List[ClassRow]) -> List[ClassRow]:
+        visible_classes: List[ClassRow]
         visible_classes = [row for row in data if row.department == department]
         # doing this so fast_ptr knows where to stop
-        visible_classes.append(ClassRow(course_num=None))
+        visible_classes.append(ClassRow(quarter=self.quarter, course_num=None))
 
         # blank class list for ones to insert
         classes_to_insert = []
@@ -88,26 +83,8 @@ def process_department(self, department, data: List[ClassRow]) -> List[ClassRow]
             fast_ptr += 1
         return classes_to_insert
 
-    def save_classes(self, classes_to_insert: List[ClassRow], quarter):
-        sql_str = """\
-                      INSERT INTO {}(DEPARTMENT, COURSE_NUM, SECTION_ID, \
-                      COURSE_ID, SECTION_TYPE, DAYS, TIME, LOCATION, ROOM, INSTRUCTOR, DESCRIPTION, UNITS)  \
-                      VALUES 
-                      (:department, 
-                      :course_num, 
-                      :section_id, 
-                      :course_id,
-                      :section_type,
-                      :days,
-                      :times, 
-                      :location,
-                      :room,
-                      :instructor,
-                      :description,
-                      :units) \
-                    """.format(quarter)
-        for c in classes_to_insert:
-            self.cursor.execute(sql_str, vars(c))
+    def save_classes(self, classes_to_insert: List[ClassRow]):
+        self.session.add_all(classes_to_insert)
 
     def sanitize_classes(self, classes: List[ClassRow]):
         return [c for c in classes if not c.is_cancelled()]
@@ -150,7 +127,7 @@ def process_current_course(self, classes: List[ClassRow], section_count: Dict[st
 
                     # class with type is a class
                     for class_with_type in type_group:
-                        new_class = copy.copy(class_with_type)
+                        new_class = ClassRow(**class_with_type.asdict())
                         # always guaranteed to have at least one element in the list
                         new_class.course_id = replica[0].course_id
                         # handle the passing of variable information through rows here
@@ -219,15 +196,10 @@ def split_into_subsections(self, section: ClassRow):
                                                         len(days) - 1]))
         for entry in day_time_pairs:
             # each subsection has mostly the same info as the section
-            subsection = copy.copy(section)
+            subsection = ClassRow(**section.asdict())
             subsection.days = entry[0]
             subsection.times = entry[1]
             ret.append(subsection)
         return ret
 
-    def close(self):
-        self.database.commit()
-        self.database.execute("VACUUM")
-        self.database.close()
-
 # Cleaner().clean()