From 014215155819937d616cab181537f27e68380641 Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Sat, 8 Feb 2020 12:58:18 -0500
Subject: [PATCH] Database user needs to own tables and sequences.

This change allows the newly made database user to apply migrations.

It also includes a very Azure-specific change. Say we have an Azure
Postgres database user "root", which is the user making the database
connections for this script, and it is creating an "atat" user/role.
That root user will be a member of the azure_pg_admin group. In order
for root to change the ownership of the tables in the database to
atat, it needs to have membership in the atat role. To achieve this we
grant azure_pg_admin the atat role.
---
 script/database_setup.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/script/database_setup.py b/script/database_setup.py
index 7784be05..e4964516 100644
--- a/script/database_setup.py
+++ b/script/database_setup.py
@@ -16,16 +16,14 @@ from reset_database import reset_database
 
 
 def database_setup(username, password, dbname, ccpo_users):
+    print("Applying schema and seeding roles and permissions.")
+    reset_database()
+
     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:
-        print(f"Postgres user role '{username}' already exists.")
+    _create_database_user(username, password, dbname)
 
-    print("Applying schema and seeding roles and permissions.")
-    reset_database()
     print("Creating initial set of CCPO users.")
     _add_ccpo_users(ccpo_users)
 
@@ -47,6 +45,22 @@ def _create_database_user(username, password, dbname):
         f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL PRIVILEGES ON FUNCTIONS TO {username}; \n"
     )
 
+    try:
+        # TODO: make this more configurable
+        engine.execute(f"GRANT {username} TO azure_pg_admin;")
+    except sqlalchemy.exc.ProgrammingError as err:
+        print(f"Cannot grant new role {username} to azure_pg_admin")
+
+    for table in meta.tables:
+        engine.execute(f"ALTER TABLE {table} OWNER TO {username};\n")
+
+    sequence_results = engine.execute(
+        "SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';"
+    ).fetchall()
+    sequences = [p[0] for p in sequence_results]
+    for sequence in sequences:
+        engine.execute(f"ALTER SEQUENCE {sequence} OWNER TO {username};\n")
+
     trans.commit()