Only fetch environment_id in provisioning queries

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]

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()):
+    for environment_id in Environments.get_environments_pending_creation(
-        create_environment.delay(
+        pendulum.now()
-            environment_id=environment.id, atat_user_id=environment.creator_id
+    ):
-        )
+        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)

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(
+    mock.delay.assert_called_once_with(environment_id=environment.id)
-        environment_id=environment.id, atat_user_id=environment.creator_id
-    )
 
 
 def test_dispatch_create_atat_admin_user(session, monkeypatch):