start db code into its own file, update migrations with backfill

3 files changed, 68 insertions, 73 deletions

diff --git a/poller/__init__.py b/poller/__init__.py
index 29852ba..f6e5259 100644
--- a/poller/__init__.py
+++ b/poller/__init__.py
@@ -12,6 +12,7 @@ import requests
 from bs4 import BeautifulSoup
 
 from .dedup import dedup
+from .db import create_tables, get_all_from_db, get_latest_from_db
 
 POOL_TIME = 5 * 60 # Seconds
 DASHBOARD_URL = 'https://rit.edu/ready/spring-dashboard'
@@ -22,31 +23,6 @@ db_lock = threading.Lock()
 if not os.path.exists('./data'):
     os.mkdir('./data')
 
-def interrupt():
-    global data_thread
-    data_thread.cancel()
-
-def create_tables():
-    with db_lock:
-        db_conn = sqlite3.connect('data/data.sqlite3')
-        c = db_conn.cursor()
-        sql = f'CREATE TABLE IF NOT EXISTS `alertlevel` (time DATETIME PRIMARY KEY NOT NULL, color CHAR(50) NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `total` (time DATETIME PRIMARY KEY NOT NULL, total_students INT NOT NULL, total_staff INT NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `new` (time DATETIME PRIMARY KEY NOT NULL, new_students INT NOT NULL, new_staff INT NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `quarantine` (time DATETIME PRIMARY KEY NOT NULL, quarantine_on_campus INT NOT NULL, quarantine_off_campus INT NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `isolation` (time DATETIME PRIMARY KEY NOT NULL, isolation_on_campus INT NOT NULL, isolation_off_campus INT NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `beds` (time DATETIME PRIMARY KEY NOT NULL, beds_available INT NOT NULL);'
-        c.execute(sql)
-        sql = f'CREATE TABLE IF NOT EXISTS `tests` (time DATETIME PRIMARY KEY NOT NULL, tests_administered INT NOT NULL);'
-        c.execute(sql)
-        db_conn.commit()
-        db_conn.close()
-
 def update_db():
     with db_lock:
         db_conn = sqlite3.connect('data/data.sqlite3')
@@ -69,47 +45,6 @@ def update_db():
         db_conn.close()
         dedup()
 
-def get_latest_from_db():
-    return get_all_from_db()[-1]
-
-def get_all_from_db():
-    with db_lock:
-        db_conn = sqlite3.connect('data/data.sqlite3')
-        c = db_conn.cursor()
-        sql = 'SELECT alertlevel.time, alertlevel.color, total.total_students, total.total_staff, new.new_students, new.new_staff, ' + \
-            'quarantine.quarantine_on_campus, quarantine.quarantine_off_campus, isolation.isolation_on_campus, isolation.isolation_off_campus, ' + \
-            'beds.beds_available, tests.tests_administered ' + \
-            'FROM `alertlevel` ' + \
-            'INNER JOIN `total` ' + \
-            'ON alertlevel.time = total.time ' + \
-            'INNER JOIN `new` ' + \
-            'ON alertlevel.time = new.time ' + \
-            'INNER JOIN `quarantine` ' + \
-            'ON alertlevel.time = quarantine.time ' + \
-            'INNER JOIN `isolation` ' + \
-            'ON alertlevel.time = isolation.time ' + \
-            'INNER JOIN `beds` ' + \
-            'ON alertlevel.time = beds.time ' + \
-            'INNER JOIN `tests` ' + \
-            'ON alertlevel.time = tests.time'
-        c.execute(sql)
-
-        data = [{
-            'alert_level': d[1],        
-            'total_students': d[2],
-            'total_staff': d[3],
-            'new_students': d[4],
-            'new_staff': d[5],
-            'quarantine_on_campus': d[6],
-            'quarantine_off_campus': d[7],
-            'isolation_on_campus': d[8],
-            'isolation_off_campus': d[9],
-            'beds_available': d[10],
-            'tests_administered': d[11],
-            'last_updated': d[0]
-        } for d in c.fetchall()]
-        return data
-
 def data_is_same(current_data):
     global LATEST_DATA
     if LATEST_DATA is None or current_data is None:
@@ -181,11 +116,6 @@ def get_data():
         update_db()
     return current_data
 
-
-get_data()
-# When you kill Flask (SIGTERM), clear the trigger for the next thread
-atexit.register(interrupt)
-
 APP = Flask(__name__)
 
 # Load file based configuration overrides if present
diff --git a/poller/db.py b/poller/db.py
new file mode 100644
index 0000000..d3836a7
--- /dev/null
+++ b/poller/db.py
@@ -0,0 +1,67 @@
+import threading
+import sqlite3
+
+db_lock = threading.Lock()
+
+def create_tables():
+    with db_lock:
+        db_conn = sqlite3.connect('data/data.sqlite3')
+        c = db_conn.cursor()
+        sql = f'CREATE TABLE IF NOT EXISTS `alertlevel` (time DATETIME PRIMARY KEY NOT NULL, color CHAR(50) NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `total` (time DATETIME PRIMARY KEY NOT NULL, total_students INT NOT NULL, total_staff INT NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `new` (time DATETIME PRIMARY KEY NOT NULL, new_students INT NOT NULL, new_staff INT NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `quarantine` (time DATETIME PRIMARY KEY NOT NULL, quarantine_on_campus INT NOT NULL, quarantine_off_campus INT NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `isolation` (time DATETIME PRIMARY KEY NOT NULL, isolation_on_campus INT NOT NULL, isolation_off_campus INT NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `beds` (time DATETIME PRIMARY KEY NOT NULL, beds_available INT NOT NULL);'
+        c.execute(sql)
+        sql = f'CREATE TABLE IF NOT EXISTS `tests` (time DATETIME PRIMARY KEY NOT NULL, tests_administered INT NOT NULL);'
+        c.execute(sql)
+        db_conn.commit()
+        db_conn.close()
+
+
+def get_latest_from_db():
+    return get_all_from_db()[-1]
+
+def get_all_from_db():
+    with db_lock:
+        db_conn = sqlite3.connect('data/data.sqlite3')
+        c = db_conn.cursor()
+        sql = 'SELECT alertlevel.time, alertlevel.color, total.total_students, total.total_staff, new.new_students, new.new_staff, ' + \
+            'quarantine.quarantine_on_campus, quarantine.quarantine_off_campus, isolation.isolation_on_campus, isolation.isolation_off_campus, ' + \
+            'beds.beds_available, tests.tests_administered ' + \
+            'FROM `alertlevel` ' + \
+            'INNER JOIN `total` ' + \
+            'ON alertlevel.time = total.time ' + \
+            'INNER JOIN `new` ' + \
+            'ON alertlevel.time = new.time ' + \
+            'INNER JOIN `quarantine` ' + \
+            'ON alertlevel.time = quarantine.time ' + \
+            'INNER JOIN `isolation` ' + \
+            'ON alertlevel.time = isolation.time ' + \
+            'INNER JOIN `beds` ' + \
+            'ON alertlevel.time = beds.time ' + \
+            'INNER JOIN `tests` ' + \
+            'ON alertlevel.time = tests.time'
+        c.execute(sql)
+
+        data = [{
+            'alert_level': d[1],        
+            'total_students': d[2],
+            'total_staff': d[3],
+            'new_students': d[4],
+            'new_staff': d[5],
+            'quarantine_on_campus': d[6],
+            'quarantine_off_campus': d[7],
+            'isolation_on_campus': d[8],
+            'isolation_off_campus': d[9],
+            'beds_available': d[10],
+            'tests_administered': d[11],
+            'last_updated': d[0]
+        } for d in c.fetchall()]
+        return data
diff --git a/poller/dedup.py b/poller/dedup.py
index 9db5886..ba68be8 100644
--- a/poller/dedup.py
+++ b/poller/dedup.py
@@ -69,5 +69,3 @@ def dedup():
         else:
             drop_by_date(data[i]['last_updated'])
             print('dropped ' + data[i]['last_updated'])
-
-dedup()
\ No newline at end of file