From 46f2a00ee33ffa29bcc031f13c69efbe9c6aaa0f Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Wed, 29 Aug 2018 11:05:58 -0400
Subject: [PATCH] move pe number ingestion for setup into a migration

---
 .../versions/2c2a2af465d3_add_pe_numbers.py   | 44 +++++++++++++++++++
 script/setup                                  |  3 --
 script/update                                 |  3 --
 3 files changed, 44 insertions(+), 6 deletions(-)
 create mode 100644 alembic/versions/2c2a2af465d3_add_pe_numbers.py

diff --git a/alembic/versions/2c2a2af465d3_add_pe_numbers.py b/alembic/versions/2c2a2af465d3_add_pe_numbers.py
new file mode 100644
index 00000000..eeb50fec
--- /dev/null
+++ b/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")
diff --git a/script/setup b/script/setup
index d43e34c6..2e2dbcc6 100755
--- a/script/setup
+++ b/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
diff --git a/script/update b/script/update
index 46d5d549..1baef9a8 100755
--- a/script/update
+++ b/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"