Add creation date field to stages and backfill them to the feature creation date (#3107)

* Add creation date field to stages and backfill them to the feature creation date

api/converters.py

@@ -115,12 +115,9 @@ def _prep_stage_info(
 
   # Get all stages associated with the feature, sorted by stage type.
   if prefetched_stages is not None:
-    prefetched_stages.sort(key=lambda s: s.stage_type)
     stages = prefetched_stages
   else:
-    stages = Stage.query(
-        Stage.feature_id == fe.key.integer_id()).order(Stage.stage_type)
-
+    stages = Stage.query(Stage.feature_id == fe.key.integer_id())
   stage_info: StagePrepResponse = {
       'proto': None,
       'dev_trial': None,
@@ -160,6 +157,7 @@ def _prep_stage_info(
       stage_info['rollout'] = s
     stage_info['all_stages'].append(stage_dict)
 
+  stage_info['all_stages'].sort(key=lambda s: (s['stage_type'], s['created']))
   return stage_info
 
 
@@ -174,6 +172,7 @@ def stage_to_json_dict(
 
   d: StageDict = {
     'id': stage.key.integer_id(),
+    'created': str(stage.created),
     'feature_id': stage.feature_id,
     'stage_type': stage.stage_type,
     'display_name': stage.display_name,

api/stages_api_test.py

@@ -16,6 +16,7 @@
 import testing_config  # Must be imported before the module under test.
 
 import flask
+from datetime import datetime
 from unittest import mock
 from google.cloud import ndb  # type: ignore
 import werkzeug.exceptions
@@ -30,6 +31,7 @@ test_app = flask.Flask(__name__)
 class StagesAPITest(testing_config.CustomTestCase):
 
   def setUp(self):
+    self.now = datetime.now()
     self.feature_owner = AppUser(email='feature_owner@example.com')
     self.feature_owner.put()
 
@@ -45,24 +47,27 @@ class StagesAPITest(testing_config.CustomTestCase):
         ux_emails=['ux_person@example.com'],
         intent_thread_url='https://example.com/intent',
         milestones=MilestoneSet(desktop_first=100),
-        experiment_goals='To be the very best.')
+        experiment_goals='To be the very best.',
+        created=self.now)
     self.stage_1.put()
     # Shipping stage.
-    self.stage_2 = Stage(id=11, feature_id=1, stage_type=160)
+    self.stage_2 = Stage(id=11, feature_id=1, stage_type=160, created=self.now)
     self.stage_2.put()
 
     self.stage_3 = Stage(id=30, feature_id=99, stage_type=150, browser='Chrome',
         ux_emails=['ux_person@example.com'],
         intent_thread_url='https://example.com/intent',
         milestones=MilestoneSet(desktop_first=100),
-        experiment_goals='To be the very best.')
+        experiment_goals='To be the very best.',
+        created=self.now)
     self.stage_3.put()
 
     self.stage_4 = Stage(id=40, feature_id=1, stage_type=150, browser='Chrome',
         ux_emails=['ux_person@example.com'],
         intent_thread_url='https://example.com/intent',
         milestones=MilestoneSet(desktop_first=100),
-        experiment_goals='To be the very best.')
+        experiment_goals='To be the very best.',
+        created=self.now)
     self.stage_4.put()
 
     self.stage_5 = Stage(id=50, feature_id=1, stage_type=150, browser='Chrome',
@@ -70,13 +75,15 @@ class StagesAPITest(testing_config.CustomTestCase):
         ux_emails=['ux_person@example.com'],
         intent_thread_url='https://example.com/intent',
         milestones=MilestoneSet(desktop_first=100),
-        experiment_goals='To be the very best.')
+        experiment_goals='To be the very best.',
+        created=self.now)
     self.stage_5.put()
 
     self.expected_stage_1 = {
         'android_first': None,
         'android_last': None,
         'announcement_url': None,
+        'created': str(self.now),
         'desktop_first': 100,
         'desktop_last': None,
         'display_name': None,
@@ -147,6 +154,7 @@ class StagesAPITest(testing_config.CustomTestCase):
     """Returns stage data with extension if requesting a valid stage ID."""
     extension = {
         'id': 50,
+        'created': str(self.now),
         'feature_id': 1,
         'stage_type': 150,
         'intent_stage': 3,
@@ -181,6 +189,7 @@ class StagesAPITest(testing_config.CustomTestCase):
 
     expect = {
         'id': 40,
+        'created': str(self.now),
         'feature_id': 1,
         'stage_type': 150,
         'intent_stage': 3,

internals/core_models.py

@@ -287,3 +287,4 @@ class Stage(ndb.Model):
   enterprise_policies = ndb.StringProperty(repeated=True)
 
   archived = ndb.BooleanProperty(default=False)
+  created = ndb.DateTimeProperty(auto_now_add=True)

internals/data_types.py

@@ -23,6 +23,7 @@ from typing import TypedDict
 # JSON representation of Stage entity data.
 class StageDict(TypedDict):
   id: int
+  created: str
   feature_id: int
   stage_type: int
   display_name: str

internals/maintenance_scripts.py

@@ -193,3 +193,26 @@ class BackfillRespondedOn(FlaskHandler):
 
     ndb.put_multi(batch)
     return f'{count} Gates entities updated.'
+
+class BackfillStageCreated(FlaskHandler):
+  def get_template_data(self, **kwargs):
+    """Backfill created dates for existing stages."""
+    self.require_cron_header()
+    count = 0
+    batch = []
+    BATCH_SIZE = 100
+    stages: ndb.Query = Stage.query()
+    for stage in stages:
+      feature_entry = FeatureEntry.get_by_id(stage.feature_id)
+      if feature_entry == None or stage.created != None:
+        continue
+      stage.created = feature_entry.created
+      batch.append(stage)
+      count += 1
+      if len(batch) > BATCH_SIZE:
+        ndb.put_multi(batch)
+        logging.info(f'Finished a batch of {BATCH_SIZE}')
+        batch = []
+
+    ndb.put_multi(batch)
+    return f'{count} Stages entities updated of {stages.count()} available stages.'

main.py

@@ -248,6 +248,8 @@ internals_routes: list[Route] = [
         maintenance_scripts.MigrateGeckoViews),
   Route('/scripts/backfill_responded_on',
         maintenance_scripts.BackfillRespondedOn),
+  Route('/scripts/backfill_stage_created',
+        maintenance_scripts.BackfillStageCreated),
 ]
 
 dev_routes: list[Route] = []