atst/script/remove_sample_data.py
leigh-mil 9412ccffea Reorder items to delete so that join models are deleted before the models they join
Add application invitations to list of items to be deleted
Clean up reference to workspace
2019-05-22 19:41:30 -04:00

118 lines
3.7 KiB
Python

# Add root application 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)
import ctypes
import sqlalchemy
from sqlalchemy import event as sqlalchemy_event
from atst.database import db
from atst.app import make_config, make_app
from atst.models.audit_event import AuditEvent
from atst.models.environment import Environment
from atst.models.environment_role import EnvironmentRole
from atst.models.application import Application
from atst.models.permission_set import PermissionSet
from atst.models.user import User
from atst.models.portfolio_role import PortfolioRole
from atst.models.portfolio import Portfolio
from atst.models.mixins import AuditableMixin
from atst.domain.environments import Environments
from atst.domain.exceptions import NotFoundError
from atst.domain.csp.reports import MockReportingProvider
from atst.domain.portfolios import Portfolios
from atst.domain.users import Users
dod_ids = [
"1234567890",
"2345678901",
"3456789012",
"4567890123",
"5678901234",
"6789012345",
"2342342342", # Andy
"3453453453", # Sally
"4564564564", # Betty
"6786786786",
]
def remove_sample_data(all_users=False):
query = db.session.query(User)
if not all_users:
query = query.filter(User.dod_id.in_(dod_ids))
users = query.all()
delete_listeners = [
k
for k in sqlalchemy_event.registry._key_to_collection
if k[1] == "after_delete"
]
for listener in delete_listeners:
[class_id, identifier, _] = listener
model = ctypes.cast(class_id, ctypes.py_object).value
sqlalchemy_event.remove(model, identifier, AuditableMixin.audit_delete)
for user in users:
all_portfolios = Portfolios.for_user(user)
portfolios = [p for p in all_portfolios if p.owner == user]
portfolio_audit = (
db.session.query(AuditEvent)
.filter(AuditEvent.portfolio_id.in_([w.id for w in portfolios]))
.all()
)
portfolio_roles = [role for portfolio in portfolios for role in portfolio.roles]
task_orders = [to for portfolio in portfolios for to in portfolio.task_orders]
portfolio_invites = [
invite for role in portfolio_roles for invite in role.invitations
]
applications = [p for portfolio in portfolios for p in portfolio.applications]
application_roles = [
a for application in applications for a in application.roles
]
application_invites = [
invite for role in application_roles for invite in role.invitations
]
environments = (
db.session.query(Environment)
.filter(Environment.application_id.in_([p.id for p in applications]))
.all()
)
roles = [role for env in environments for role in env.roles]
for set_of_things in [
portfolio_invites,
application_invites,
roles,
environments,
application_roles,
applications,
portfolio_roles,
task_orders,
portfolio_audit,
]:
for thing in set_of_things:
db.session.delete(thing)
db.session.commit()
query = "DELETE FROM portfolios WHERE portfolios.id = ANY(:ids);"
db.session.connection().execute(
sqlalchemy.text(query), ids=[w.id for w in portfolios]
)
db.session.commit()
if __name__ == "__main__":
config = make_config({"DISABLE_CRL_CHECK": True, "DEBUG": False})
app = make_app(config)
with app.app_context():
remove_sample_data()