From 49a1a219ae36d23ca59239eaf1139e16464070dd Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Thu, 23 Jan 2020 06:25:35 -0500
Subject: [PATCH] Script for setting up database user, schema, and seed data.

This script is for bootstrapping the initial database. It can be run via
a container, but requires that a Postgres superuser's credentials be
provided via our normal config. That way the superuser can provision a
less-privileged user for the application's database connection.
---
 atst/domain/users.py     |  7 +++-
 script/database_setup.py | 80 ++++++++++++++++++++++++++++++++++++++++
 script/reset_database.py |  4 +-
 3 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 script/database_setup.py

diff --git a/atst/domain/users.py b/atst/domain/users.py
index 5e09ce22..e5fdbad7 100644
--- a/atst/domain/users.py
+++ b/atst/domain/users.py
@@ -93,10 +93,13 @@ class Users(object):
         return user
 
     @classmethod
-    def give_ccpo_perms(cls, user):
+    def give_ccpo_perms(cls, user, commit=True):
         user.permission_sets = PermissionSets.get_all()
         db.session.add(user)
-        db.session.commit()
+
+        if commit:
+            db.session.commit()
+
         return user
 
     @classmethod
diff --git a/script/database_setup.py b/script/database_setup.py
new file mode 100644
index 00000000..bab95890
--- /dev/null
+++ b/script/database_setup.py
@@ -0,0 +1,80 @@
+# Add root application dir to the python path
+import os
+import sys
+from contextlib import contextmanager
+
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(parent_dir)
+
+import sqlalchemy
+from alembic import config as alembic_config
+import yaml
+
+from atst.app import make_config, make_app
+from atst.database import db
+from atst.domain.users import Users
+from atst.models import User
+from reset_database import reset_database
+
+
+def database_setup(username, password, dbname, ccpo_users):
+    """docstring for database_setup"""
+    print(
+        f"Creating Postgres user role for '{username}' and granting all privileges to database '{dbname}'."
+    )
+    try:
+        _create_database_user(username, password, dbname)
+    except sqlalchemy.exc.ProgrammingError as err:
+        raise err
+        print(f"Postgres user role '{username}' already exists.")
+
+    print("Applying schema and seeding roles and permissions.")
+    reset_database()
+    print("Creating initial set of CCPO users.")
+    _add_ccpo_users(ccpo_users)
+
+
+def _create_database_user(username, password, dbname):
+    conn = db.engine.connect()
+
+    meta = sqlalchemy.MetaData(bind=conn)
+    meta.reflect()
+
+    trans = conn.begin()
+    engine = trans.connection.engine
+
+    engine.execute(
+        f"CREATE ROLE {username} WITH LOGIN NOSUPERUSER INHERIT NOCREATEDB NOCREATEROLE NOREPLICATION PASSWORD '{password}';\n"
+        + f"GRANT ALL PRIVILEGES ON DATABASE {dbname} TO {username};\n"
+        + f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL PRIVILEGES ON TABLES TO {username}; \n"
+        + f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL PRIVILEGES ON SEQUENCES TO {username}; \n"
+        + f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL PRIVILEGES ON FUNCTIONS TO {username}; \n"
+    )
+
+    trans.commit()
+
+
+def _add_ccpo_users(ccpo_users):
+    for user_data in ccpo_users:
+        user = User(**user_data)
+        Users.give_ccpo_perms(user, commit=False)
+        db.session.add(user)
+
+    db.session.commit()
+
+
+def _load_yaml(file_):
+    with open(file_) as f:
+        return yaml.safe_load(f)
+
+
+if __name__ == "__main__":
+    config = make_config({"DISABLE_CRL_CHECK": True, "DEBUG": False})
+    app = make_app(config)
+    with app.app_context():
+        dbname = config.get("PGDATABASE", "atat")
+        username = sys.argv[1]
+        password = sys.argv[2]
+        ccpo_user_file = sys.argv[3]
+        ccpo_users = _load_yaml(ccpo_user_file)
+        database_setup(username, password, dbname, ccpo_users)
diff --git a/script/reset_database.py b/script/reset_database.py
index cfa63298..dda1c1ba 100644
--- a/script/reset_database.py
+++ b/script/reset_database.py
@@ -16,7 +16,9 @@ from atst.app import make_config, make_app
 def reset_database():
     conn = db.engine.connect()
 
-    meta = sqlalchemy.MetaData(bind=conn, reflect=True)
+    meta = sqlalchemy.MetaData(bind=conn)
+    meta.reflect()
+
     trans = conn.begin()
 
     # drop all tables