diff --git a/atst/domain/environments.py b/atst/domain/environments.py index 41824dc5..829fe3fa 100644 --- a/atst/domain/environments.py +++ b/atst/domain/environments.py @@ -1,7 +1,7 @@ from sqlalchemy import text, func, or_ from sqlalchemy.orm.exc import NoResultFound -from sqlalchemy.orm import load_only from typing import List +from uuid import UUID from atst.database import db from atst.models import Environment, Application, Portfolio, TaskOrder, CLIN @@ -97,7 +97,7 @@ class Environments(object): @classmethod def base_provision_query(cls, now): return ( - db.session.query(Environment) + db.session.query(Environment.id) .join(Application) .join(Portfolio) .join(TaskOrder) @@ -110,37 +110,40 @@ class Environments(object): Environment.claimed_until <= func.now(), ) ) - # select only these columns - .options(load_only("id", "creator_id")) ) @classmethod - def get_environments_pending_creation(cls, now) -> List[Environment]: + def get_environments_pending_creation(cls, now) -> List[UUID]: """ Any environment with an active CLIN that doesn't yet have a `cloud_id`. """ - return cls.base_provision_query(now).filter(Environment.cloud_id == None).all() + results = ( + cls.base_provision_query(now).filter(Environment.cloud_id == None).all() + ) + return [id_ for id_, in results] @classmethod - def get_environments_pending_atat_user_creation(cls, now) -> List[Environment]: + def get_environments_pending_atat_user_creation(cls, now) -> List[UUID]: """ Any environment with an active CLIN that has a cloud_id but no `root_user_info`. """ - return ( + results = ( cls.base_provision_query(now) .filter(Environment.cloud_id != None) .filter(Environment.root_user_info == text("'null'")) ).all() + return [id_ for id_, in results] @classmethod - def get_environments_pending_baseline_creation(cls, now) -> List[Environment]: + def get_environments_pending_baseline_creation(cls, now) -> List[UUID]: """ Any environment with an active CLIN that has a `cloud_id` and `root_user_info` but no `baseline_info`. """ - return ( + results = ( cls.base_provision_query(now) .filter(Environment.cloud_id != None) .filter(Environment.root_user_info != text("'null'")) .filter(Environment.baseline_info == text("'null'")) ).all() + return [id_ for id_, in results] diff --git a/atst/jobs.py b/atst/jobs.py index 017b98d4..ec639306 100644 --- a/atst/jobs.py +++ b/atst/jobs.py @@ -109,7 +109,7 @@ def do_work(fn, task, csp, **kwargs): @celery.task(bind=True) -def create_environment(self, environment_id=None, atat_user_id=None): +def create_environment(self, environment_id=None): do_work(do_create_environment, self, app.csp.cloud, environment_id=environment_id) @@ -132,23 +132,23 @@ def create_environment_baseline(self, environment_id=None): @celery.task(bind=True) def dispatch_create_environment(self): - for environment in Environments.get_environments_pending_creation(pendulum.now()): - create_environment.delay( - environment_id=environment.id, atat_user_id=environment.creator_id - ) + for environment_id in Environments.get_environments_pending_creation( + pendulum.now() + ): + create_environment.delay(environment_id=environment_id) @celery.task(bind=True) def dispatch_create_atat_admin_user(self): - for environment in Environments.get_environments_pending_atat_user_creation( + for environment_id in Environments.get_environments_pending_atat_user_creation( pendulum.now() ): - create_atat_admin_user.delay(environment_id=environment.id) + create_atat_admin_user.delay(environment_id=environment_id) @celery.task(bind=True) def dispatch_create_environment_baseline(self): - for environment in Environments.get_environments_pending_baseline_creation( + for environment_id in Environments.get_environments_pending_baseline_creation( pendulum.now() ): - create_environment_baseline.delay(environment_id=environment.id) + create_environment_baseline.delay(environment_id=environment_id) diff --git a/tests/test_jobs.py b/tests/test_jobs.py index ab36a6c9..537716e4 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -138,9 +138,7 @@ def test_dispatch_create_environment(session, monkeypatch): dispatch_create_environment.run() - mock.delay.assert_called_once_with( - environment_id=environment.id, atat_user_id=environment.creator_id - ) + mock.delay.assert_called_once_with(environment_id=environment.id) def test_dispatch_create_atat_admin_user(session, monkeypatch):