atst/script/remove_sample_data.py

# 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)

import ctypes
import sqlalchemy
from sqlalchemy import or_, 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.project import Project
from atst.models.request import Request
from atst.models.request_revision import RequestRevision
from atst.models.request_status_event import RequestStatus, RequestStatusEvent
from atst.models.role import Role
from atst.models.user import User
from atst.models.workspace_role import WorkspaceRole
from atst.models.workspace import Workspace
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.requests import Requests
from atst.domain.users import Users
from atst.domain.workspaces import Workspaces
from tests.factories import RequestFactory, LegacyTaskOrderFactory


dod_ids = [
    "1234567890",
    "2345678901",
    "3456789012",
    "4567890123",
    "5678901234",
    "6789012345",
    "2342342342", # Andy
    "3453453453", # Sally
    "4564564564", # Betty
    "6786786786",
]


def create_demo_workspace(name, data):
    try:
        workspace_owner = Users.get_by_dod_id("678678678") # Other
        auditor = Users.get_by_dod_id("3453453453") # Sally
    except NotFoundError:
        print("Could not find demo users; will not create demo workspace {}".format(name))
        return

    request = RequestFactory.build(creator=workspace_owner)
    request.legacy_task_order = LegacyTaskOrderFactory.build()
    request = Requests.update(
        request.id, {"financial_verification": RequestFactory.mock_financial_data()}
    )
    approved_request = Requests.set_status(request, RequestStatus.APPROVED)

    workspace = Requests.approve_and_create_workspace(request)
    Workspaces.update(workspace, { "name": name })

    for mock_project in data["projects"]:
        project = Project(workspace=workspace, name=mock_project.name, description='')
        env_names = [env.name for env in mock_project.environments]
        envs = Environments.create_many(project, env_names)
        db.session.add(project)
        db.session.commit()


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:
        requests = (
            db.session.query(Request)
            .filter(Request.id.in_([r.id for r in user.owned_requests]))
            .all()
        )
        request_audit = (
            db.session.query(AuditEvent)
            .filter(AuditEvent.request_id.in_([r.id for r in requests]))
            .all()
        )
        events = [ev for r in requests for ev in r.status_events]
        revisions = [rev for r in requests for rev in r.revisions]
        workspaces = [r.workspace for r in requests if r.workspace]
        ws_audit = (
            db.session.query(AuditEvent)
            .filter(AuditEvent.workspace_id.in_([w.id for w in workspaces]))
            .all()
        )
        workspace_roles = [role for workspace in workspaces for role in workspace.roles]
        invites = [invite for role in workspace_roles for invite in role.invitations]
        projects = [p for workspace in workspaces for p in workspace.projects]
        environments = (
            db.session.query(Environment)
            .filter(Environment.project_id.in_([p.id for p in projects]))
            .all()
        )
        roles = [role for env in environments for role in env.roles]

        for set_of_things in [
            roles,
            environments,
            projects,
            invites,
            workspace_roles,
            ws_audit,
            events,
            revisions,
            request_audit,
        ]:
            for thing in set_of_things:
                db.session.delete(thing)

        db.session.commit()

        query = "DELETE FROM workspaces WHERE workspaces.id = ANY(:ids);"
        db.session.connection().execute(
            sqlalchemy.text(query), ids=[w.id for w in workspaces]
        )

        query = "DELETE FROM requests WHERE requests.id = ANY(:ids);"
        db.session.connection().execute(
            sqlalchemy.text(query), ids=[r.id for r in requests]
        )

        db.session.commit()


if __name__ == "__main__":
    config = make_config()
    app = make_app(config)
    with app.app_context():
        remove_sample_data()
        create_demo_workspace('Aardvark', MockReportingProvider.REPORT_FIXTURE_MAP["Aardvark"])
        create_demo_workspace('Beluga', MockReportingProvider.REPORT_FIXTURE_MAP["Beluga"])