From 0fca4753c5902ca70dc25a8112684ccac67b0310 Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Tue, 31 Jul 2018 09:21:33 -0400
Subject: [PATCH] add script for ingesting PE numbers and necessary script
 config for database

---
 config/base.ini             |  1 +
 script/ingest_pe_numbers.py | 40 +++++++++++++++++++++++++++++++++++++
 script/setup                |  3 +++
 script/update               |  3 +++
 4 files changed, 47 insertions(+)
 create mode 100644 script/ingest_pe_numbers.py

diff --git a/config/base.ini b/config/base.ini
index 285eb11f..8b3a6881 100644
--- a/config/base.ini
+++ b/config/base.ini
@@ -7,6 +7,7 @@ FUNDZ_BASE_URL= http://localhost:8004
 COOKIE_SECRET = some-secret-please-replace
 SECRET = change_me_into_something_secret
 CAC_URL = https://localhost:8001
+PE_NUMBER_CSV_URL = http://c95e1ebb198426ee57b8-174bb05a294821bedbf46b6384fe9b1f.r31.cf5.rackcdn.com/penumbers.csv
 REDIS_URI = redis://localhost:6379
 SESSION_TTL_SECONDS = 600
 PGAPPNAME = atst
diff --git a/script/ingest_pe_numbers.py b/script/ingest_pe_numbers.py
new file mode 100644
index 00000000..1df5b132
--- /dev/null
+++ b/script/ingest_pe_numbers.py
@@ -0,0 +1,40 @@
+from urllib.request import urlopen
+import csv
+
+from sqlalchemy.dialects.postgresql import insert
+
+# Add root project dir to the python path
+import os
+import sys
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.append(parent_dir)
+
+from atst.app import make_deps, make_config
+from atst.models import PENumber
+
+
+def get_pe_numbers(url):
+    response = urlopen(url)
+    t = response.read().decode("utf-8")
+    return list(csv.reader(t.split("\r\n")))
+
+
+def insert_pe_numbers(db, list_of_pe_numbers):
+    stmt = insert(PENumber).values(list_of_pe_numbers)
+    do_update = stmt.on_conflict_do_update(
+        index_elements=["number"],
+        set_=dict(description=stmt.excluded.description)
+    )
+    db.execute(do_update)
+    db.commit()
+
+
+if __name__ == "__main__":
+    config = make_config()
+    deps = make_deps(config)
+    db = deps["db_session"]
+    url = config["default"]['PE_NUMBER_CSV_URL']
+    print("Fetching PE numbers from {}".format(url))
+    pe_numbers = get_pe_numbers(url)
+    print("Inserting {} PE numbers".format(len(pe_numbers)))
+    insert_pe_numbers(db, pe_numbers)
diff --git a/script/setup b/script/setup
index 5cd11461..bf83b19a 100755
--- a/script/setup
+++ b/script/setup
@@ -16,3 +16,6 @@ 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"
diff --git a/script/update b/script/update
index 1baef9a8..46d5d549 100755
--- a/script/update
+++ b/script/update
@@ -9,3 +9,6 @@ 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"