Script to associate origin trials with existing stages (#3286)

* Create script to associate trials with stages

* Update test

* move to cron job

* comments and cleanup

* changes suggested by @jrobbins

cron.yaml

@@ -28,4 +28,9 @@ cron:
   schedule: 1st monday of month 9:00
 - description: Update all feature links that are staled.
   url: /cron/update_all_feature_links
-  schedule: every tuesday 05:00
+  schedule: every tuesday 05:00
+- description: |
+    Check for origin trials and associate them with their respective
+    ChromeStatus feature entry.
+  url: /cron/associate_origin_trials
+  schedule: every day 6:00

framework/origin_trials_client_test.py

@@ -37,7 +37,14 @@ class OriginTrialsClientTest(testing_config.CustomTestCase):
         'chromestatusUrl': 'https://example.com/chromestatus',
         'startMilestone': '123',
         'endMilestone': '456',
+        'originalEndMilestone': '450',
         'endTime': '2025-01-01T00:00:00Z',
+        'feedbackUrl': 'https://example.com/feedback',
+        'documentationUrl': 'https://example.com/docs',
+        'intentToExperimentUrl': 'https://example.com/intent',
+        'type': 'ORIGIN_TRIAL',
+        'allowThirdPartyOrigins': True,
+        'trialExtensions': [{}],
       },
       {
         'id': '3611886901151137793',
@@ -88,6 +95,13 @@ class OriginTrialsClientTest(testing_config.CustomTestCase):
         'chromestatus_url': 'https://example.com/chromestatus',
         'start_milestone': '123',
         'end_milestone': '456',
+        'original_end_milestone': '450',
+        'feedback_url': 'https://example.com/feedback',
+        'documentation_url': 'https://example.com/docs',
+        'intent_to_experiment_url': 'https://example.com/intent',
+        'trial_extensions': [{}],
+        'type': 'ORIGIN_TRIAL',
+        'allow_third_party_origins': True,
         'end_time': '2025-01-01T00:00:00Z',
       },
     ]

internals/data_types.py

@@ -275,24 +275,38 @@ class VerboseFeatureDict(TypedDict):
 @dataclass
 class OriginTrialInfo():
   def __init__(self, api_trial):
-    self.id = api_trial.get('id', '')
-    self.display_name = api_trial.get('displayName', '')
-    self.description = api_trial.get('description', '')
-    self.origin_trial_feature_name = api_trial.get('originTrialFeatureName', '')
-    self.enabled = api_trial.get('enabled', None)
-    self.status = api_trial.get('status', '')
-    self.chromestatus_url = api_trial.get('chromestatusUrl', '')
-    self.start_milestone = api_trial.get('startMilestone', '')
-    self.end_milestone = api_trial.get('endMilestone', '')
-    self.end_time = api_trial.get('endTime', '')
+    self.id = api_trial.get('id', None)
+    self.display_name = api_trial.get('displayName', None)
+    self.description = api_trial.get('description', None)
+    self.origin_trial_feature_name = api_trial.get('originTrialFeatureName', None)
+    self.enabled = api_trial.get('enabled', False)
+    self.status = api_trial.get('status', None)
+    self.chromestatus_url = api_trial.get('chromestatusUrl', None)
+    self.start_milestone = api_trial.get('startMilestone', None)
+    self.end_milestone = api_trial.get('endMilestone', None)
+    self.original_end_milestone = api_trial.get('originalEndMilestone', None)
+    self.end_time = api_trial.get('endTime', None)
+    self.documentation_url = api_trial.get('documentationUrl', None)
+    self.feedback_url = api_trial.get('feedbackUrl', None)
+    self.intent_to_experiment_url = api_trial.get('intentToExperimentUrl', None)
+    self.trial_extensions = api_trial.get('trialExtensions', None)
+    self.type = api_trial.get('type', None)
+    self.allow_third_party_origins = api_trial.get('allowThirdPartyOrigins', False)
 
-  id: str
-  display_name: str
-  description: str
-  origin_trial_feature_name: str
-  enabled: bool|None
-  status: str
-  chromestatus_url: str
-  start_milestone: str
-  end_milestone: str
-  end_time: str
+  id: str|None
+  display_name: str|None
+  description: str|None
+  origin_trial_feature_name: str|None
+  enabled: bool
+  status: str|None
+  chromestatus_url: str|None
+  start_milestone: str|None
+  end_milestone: str|None
+  original_end_milestone: str|None
+  end_time: str|None
+  documentation_url: str|None
+  feedback_url: str|None
+  intent_to_experiment_url: str|None
+  trial_extensions: list|None
+  type: str|None
+  allow_third_party_origins: bool

internals/maintenance_scripts.py

@@ -14,11 +14,13 @@
 
 import datetime
 import logging
+from typing import Any
 from google.cloud import ndb  # type: ignore
 
 from framework.basehandlers import FlaskHandler
+from framework import origin_trials_client
 from internals import approval_defs
-from internals.core_models import FeatureEntry, Stage
+from internals.core_models import FeatureEntry, MilestoneSet, Stage
 from internals.review_models import Gate, Vote, Activity
 from internals.core_enums import *
 from internals.feature_links import batch_index_feature_entries
@@ -224,4 +226,140 @@ class BackfillFeatureLinks(FlaskHandler):
     all_feature_entries = FeatureEntry.query().fetch()
     count = batch_index_feature_entries(all_feature_entries, True)
     return f'{len(all_feature_entries)} FeatureEntry entities backfilled of {count} feature links.'
-  
+
+
+class AssociateOTs(FlaskHandler):
+
+  def write_fields_for_trial_stage(self, trial_stage: Stage, trial_data: dict[str, Any]):
+    """Check if any OT stage fields are unfilled and populate them with
+    the matching trial data.
+    """
+    if trial_stage.origin_trial_id is None:
+      trial_stage.origin_trial_id = trial_data['id']
+
+    if trial_stage.ot_chromium_trial_name is None:
+      trial_stage.ot_chromium_trial_name = trial_data['origin_trial_feature_name']
+
+    if trial_stage.milestones is None:
+      trial_stage.milestones = MilestoneSet()
+    if (trial_stage.milestones.desktop_first is None and
+        trial_data['start_milestone'] is not None):
+      trial_stage.milestones.desktop_first = int(trial_data['start_milestone'])
+    if trial_stage.milestones.desktop_last is None:
+      # An original end milestone is kept if the trial has had extensions.
+      # TODO(DanielRyanSmith): Extension milestones in the trial data
+      # should be associated with new extension stages for data accuracy.
+      if trial_data['original_end_milestone'] is not None:
+        trial_stage.milestones.desktop_last = (
+            int(trial_data['original_end_milestone']))
+      elif trial_data['end_milestone'] is not None:
+        trial_stage.milestones.desktop_last = (
+            int(trial_data['end_milestone']))
+
+    if trial_stage.display_name is None:
+      trial_stage.display_name = trial_data['display_name']
+
+    if trial_stage.intent_thread_url is None:
+      trial_stage.intent_thread_url = trial_data['intent_to_experiment_url']
+
+    if trial_stage.origin_trial_feedback_url is None:
+      trial_stage.origin_trial_feedback_url = trial_data['feedback_url']
+
+    if trial_stage.ot_documentation_url is None:
+      trial_stage.ot_documentation_url = trial_data['documentation_url']
+
+    if trial_stage.ot_has_third_party_support:
+      trial_stage.ot_has_third_party_support = trial_data['allow_third_party_origins']
+
+    if not trial_stage.ot_is_deprecation_trial:
+      trial_stage.ot_is_deprecation_trial = trial_data['type'] == 'DEPRECATION'
+
+  def parse_feature_id(self, chromestatus_url: str|None) -> int|None:
+      if chromestatus_url is None:
+        return None
+      # The ChromeStatus feature ID is pulled out of the ChromeStatus URL.
+      chromestatus_id_start = chromestatus_url.rfind('/')
+      if chromestatus_id_start == -1:
+        logging.info(f'Bad ChromeStatus URL: {chromestatus_url}')
+        return None
+      # Add 1 to index, which is the start index of the ID.
+      chromestatus_id_start += 1
+      chromestatus_id_str = chromestatus_url[chromestatus_id_start:]
+      try:
+        chromestatus_id = int(chromestatus_id_str)
+      except ValueError:
+        logging.info(
+            f'Unable to parse ID from ChromeStatus URL: {chromestatus_url}')
+        return None
+      return chromestatus_id
+
+  def find_trial_stage(self, feature_id: int) -> Stage|None:
+      fe: FeatureEntry|None = FeatureEntry.get_by_id(feature_id)
+      if fe is None:
+        logging.info(f'No feature found for ChromeStatus ID: {feature_id}')
+        return None
+
+      trial_stage_type = STAGE_TYPES_ORIGIN_TRIAL[fe.feature_type]
+      trial_stages = Stage.query(
+          Stage.stage_type == trial_stage_type,
+          Stage.feature_id == feature_id).fetch()
+      # If there are no OT stages for the feature, we can't associate the
+      # trial with any stages.
+      if len(trial_stages) == 0:
+        logging.info(f'No OT stages found for feature ID: {feature_id}')
+        return None
+      # If there is currently more than one origin trial stage for the
+      # feature, we don't know which one represents the given trial.
+      if len(trial_stages) > 1:
+        logging.info('Multiple origin trial stages found for feature '
+                     f'{feature_id}. Cannot discern which stage to associate '
+                     'trial with.')
+        return None
+      return trial_stages[0]
+
+  def get_template_data(self, **kwargs):
+    """Link existing origin trials with their ChromeStatus entry"""
+    self.require_cron_header()
+
+    trials_list = origin_trials_client.get_trials_list()
+    entities_to_write: list[Stage] = []
+    trials_with_no_feature: list[str] = []
+    for trial_data in trials_list:
+      stage = Stage.query(
+          Stage.origin_trial_id == trial_data['id']).get()
+      # If this trial is already associated with a ChromeStatus stage,
+      # just see if any unfilled fields need to be populated.
+      if stage:
+        self.write_fields_for_trial_stage(stage, trial_data)
+        entities_to_write.append(stage)
+        continue
+
+      feature_id = self.parse_feature_id(trial_data['chromestatus_url'])
+      if feature_id is None:
+        trials_with_no_feature.append(trial_data)
+        continue
+
+      ot_stage = self.find_trial_stage(feature_id)
+      if ot_stage is None:
+        trials_with_no_feature.append(trial_data)
+        continue
+
+      self.write_fields_for_trial_stage(ot_stage, trial_data)
+      entities_to_write.append(ot_stage)
+
+    # List any origin trials that did not get associated with a feature entry.
+    if len(trials_with_no_feature) > 0:
+      logging.info('Trials not associated with a ChromeStatus feature:')
+    else:
+      logging.info('All trials associated with a ChromeStatus feature!')
+    for trial_data in trials_with_no_feature:
+      logging.info(f'{trial_data["id"]} {trial_data["display_name"]}')
+
+    # Update all the stages at the end. Note that there is a chance
+    # the stage entities have not changed any values if all fields already
+    # had a value.
+    logging.info(f'{len(entities_to_write)} stages to update.')
+    if len(entities_to_write) > 0:
+      ndb.put_multi(entities_to_write)
+
+    return f'{len(entities_to_write)} Stages updated with trial data.'

main.py

@@ -241,6 +241,7 @@ internals_routes: list[Route] = [
       inactive_users.RemoveInactiveUsersHandler),
   Route('/cron/reindex_all', search_fulltext.ReindexAllFeatures),
   Route('/cron/update_all_feature_links', feature_links.UpdateAllFeatureLinksHandlers),
+  Route('/cron/associate_origin_trials', maintenance_scripts.AssociateOTs),
 
   Route('/admin/find_stop_words', search_fulltext.FindStopWords),