move pe number ingestion for setup into a migration

alembic/versions/2c2a2af465d3_add_pe_numbers.py

@@ -0,0 +1,44 @@
+"""add PE numbers
+
+Revision ID: 2c2a2af465d3
+Revises: 0845b2f0f401
+Create Date: 2018-08-27 16:26:51.707146
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import insert
+from urllib.request import urlopen
+import csv
+
+from atst.app import make_config
+from atst.models.pe_number import PENumber
+
+
+# revision identifiers, used by Alembic.
+revision = '2c2a2af465d3'
+down_revision = '0845b2f0f401'
+branch_labels = None
+depends_on = None
+
+
+def get_pe_numbers(url):
+    response = urlopen(url)
+    t = response.read().decode("utf-8")
+    return list(csv.reader(t.split("\r\n")))
+
+
+def upgrade():
+    config = make_config()
+    pe_numbers = get_pe_numbers(config["PE_NUMBER_CSV_URL"])
+    db = op.get_bind()
+    stmt = insert(PENumber).values(pe_numbers)
+    do_update = stmt.on_conflict_do_update(
+        index_elements=["number"], set_=dict(description=stmt.excluded.description)
+    )
+    db.execute(do_update)
+
+
+def downgrade():
+    db = op.get_bind()
+    db.execute("TRUNCATE TABLE pe_number")

script/setup

@@ -14,8 +14,5 @@ RESET_DB="true"
 # Run the shared setup script
 source ./script/include/run_setup
 
-# Fetch and import the PE numbers
-run_command "python script/ingest_pe_numbers.py"
-
 # Compile assets and generate hash-named static files
 yarn build

script/update

@@ -9,6 +9,3 @@ MIGRATE_DB="true"
 
 # Run the shared update script
 source ./script/include/run_update
-
-# Fetch and import/update the PE numbers
-run_command "python script/ingest_pe_numbers.py"