diff --git a/internals/maintenance_scripts.py b/internals/maintenance_scripts.py index 628776bde238..2c6d999a4634 100644 --- a/internals/maintenance_scripts.py +++ b/internals/maintenance_scripts.py @@ -743,3 +743,61 @@ def get_template_data(self, **kwargs) -> str: ndb.put_multi(batch) return f'{count} Features entities updated.' + + +class BackfillGateDates(FlaskHandler): + + def get_template_data(self, **kwargs) -> str: + """Backfill resolved_on and needs_work_started_on for all Gates.""" + self.require_cron_header() + + count = 0 + batch: list[Gate] = [] + BATCH_SIZE = 100 + votes_by_gate = collections.defaultdict(list) + for vote in Vote.query(): + votes_by_gate[vote.gate_id].append(vote) + for gate in Gate.query(): + gate_votes = votes_by_gate.get(gate.key.integer_id()) or [] + if self.calc_dates(gate, gate_votes): + batch.append(gate) + count += 1 + if len(batch) > BATCH_SIZE: + ndb.put_multi(batch) + batch = [] + + ndb.put_multi(batch) + return f'{count} Gate entities updated.' + + def calc_dates(self, gate: Gate, votes: list[Vote]) -> bool: + """Set resolved_on and needs_work_started_on if needed.""" + if not votes: + return False + new_resolved_on = self.calc_resolved_on(gate, votes) + new_needs_work_started_on = self.calc_needs_work_started_on(gate, votes) + if new_resolved_on is not None: + gate.resolved_on = new_resolved_on + if new_needs_work_started_on is not None: + gate.needs_work_started_on = new_needs_work_started_on + return bool(new_resolved_on or new_needs_work_started_on) + + def calc_resolved_on(self, gate: Gate, votes: list[Vote]) -> datetime | None: + """Return the date on which the gate was resolved, or None.""" + if gate.state not in Vote.FINAL_STATES: + return None + if gate.resolved_on: + return None + + return max(v.set_on for v in votes + if v.state in Vote.FINAL_STATES) + + def calc_needs_work_started_on( + self, gate: Gate, votes: list[Vote]) -> datetime | None: + """Return the latest date on which the gate entered NEEDS_WORK.""" + if gate.state != Vote.NEEDS_WORK: + return None + if gate.needs_work_started_on: + return None + + return max(v.set_on for v in votes + if v.state == Vote.NEEDS_WORK) diff --git a/internals/maintenance_scripts_test.py b/internals/maintenance_scripts_test.py index fba255b66581..7ba690c97847 100644 --- a/internals/maintenance_scripts_test.py +++ b/internals/maintenance_scripts_test.py @@ -757,3 +757,104 @@ def test_calc_all_shipping_years__some(self, mock_gasswm: mock.MagicMock): actual = self.handler.calc_all_shipping_years() expected = {22222: 2023, 33333: 2024, 44444: 2030} self.assertEqual(expected, actual) + + +class BackfillGateDatesTest(testing_config.CustomTestCase): + + def setUp(self): + self.gate = Gate( + feature_id=1, stage_id=2, + gate_type=core_enums.GATE_API_EXTEND_ORIGIN_TRIAL, + state=Gate.PREPARING) + self.handler = maintenance_scripts.BackfillGateDates() + + def test_calc_resolved_on__not_resolved(self): + """If a gate is not resolved, don't set a resolved_on date.""" + self.assertIsNone( + self.handler.calc_resolved_on(self.gate, [])) + + self.gate.state = Vote.REVIEW_REQUESTED + self.assertIsNone( + self.handler.calc_resolved_on(self.gate, [])) + + self.gate.state = Vote.NA_REQUESTED + self.assertIsNone( + self.handler.calc_resolved_on(self.gate, [])) + + self.gate.state = Vote.REVIEW_STARTED + self.assertIsNone( + self.handler.calc_resolved_on(self.gate, [])) + + self.gate.state = Vote.NEEDS_WORK + self.assertIsNone( + self.handler.calc_resolved_on(self.gate, [])) + + def test_calc_resolved_on__resolved(self): + """If a gate was resolved, resolved_on is the last approval.""" + self.gate.state = Vote.APPROVED + gate_id = 1234 + v1 = Vote(gate_id=gate_id, set_by='feature_owner@example.com', + state=Vote.REVIEW_REQUESTED, + set_on=datetime(2023, 1, 1, 12, 30, 0)) + v2 = Vote(gate_id=gate_id, set_by='reviewer_a@example.com', + state=Vote.REVIEW_STARTED, + set_on=datetime(2023, 1, 2, 12, 30, 0)) + v3 = Vote(gate_id=gate_id, set_by='reviewer_b@example.com', + state=Vote.APPROVED, + set_on=datetime(2023, 1, 3, 12, 30, 0)) + v4 = Vote(gate_id=gate_id, set_by='reviewer_c@example.com', + state=Vote.APPROVED, + set_on=datetime(2023, 1, 4, 12, 30, 0)) + v5 = Vote(gate_id=gate_id, set_by='reviewer_d@example.com', + state=Vote.REVIEW_STARTED, + set_on=datetime(2023, 1, 5, 12, 30, 0)) + + self.assertEqual( + self.handler.calc_resolved_on(self.gate, [v1, v2, v3, v4, v5]), + v4.set_on) + + def test_calc_needs_work_started_on__not_needed(self): + """If a gate is not NEEDS_WORK, don't set a needs_work_started_on date.""" + self.assertIsNone( + self.handler.calc_needs_work_started_on(self.gate, [])) + + self.gate.state = Vote.REVIEW_REQUESTED + self.assertIsNone( + self.handler.calc_needs_work_started_on(self.gate, [])) + + self.gate.state = Vote.NA_REQUESTED + self.assertIsNone( + self.handler.calc_needs_work_started_on(self.gate, [])) + + self.gate.state = Vote.REVIEW_STARTED + self.assertIsNone( + self.handler.calc_needs_work_started_on(self.gate, [])) + + self.gate.state = Vote.APPROVED + self.assertIsNone( + self.handler.calc_needs_work_started_on(self.gate, [])) + + def test_calc_needs_work_started_on__needed(self): + """If a gate is NEEDS_WORK, it started on the last NEEDS_WORK vote.""" + self.gate.state = Vote.NEEDS_WORK + gate_id = 1234 + v1 = Vote(gate_id=gate_id, set_by='feature_owner@example.com', + state=Vote.REVIEW_REQUESTED, + set_on=datetime(2023, 1, 1, 12, 30, 0)) + v2 = Vote(gate_id=gate_id, set_by='reviewer_a@example.com', + state=Vote.NEEDS_WORK, + set_on=datetime(2023, 1, 2, 12, 30, 0)) + v3 = Vote(gate_id=gate_id, set_by='reviewer_b@example.com', + state=Vote.APPROVED, + set_on=datetime(2023, 1, 3, 12, 30, 0)) + v4 = Vote(gate_id=gate_id, set_by='reviewer_c@example.com', + state=Vote.NEEDS_WORK, + set_on=datetime(2023, 1, 4, 12, 30, 0)) + v5 = Vote(gate_id=gate_id, set_by='reviewer_d@example.com', + state=Vote.REVIEW_STARTED, + set_on=datetime(2023, 1, 5, 12, 30, 0)) + + self.assertEqual( + self.handler.calc_needs_work_started_on( + self.gate, [v1, v2, v3, v4, v5]), + v4.set_on) diff --git a/main.py b/main.py index 232c68403f2b..0354fafedbdc 100644 --- a/main.py +++ b/main.py @@ -346,6 +346,8 @@ class Route: maintenance_scripts.DeleteEmptyExtensionStages), Route('/scripts/backfill_shipping_year', maintenance_scripts.BackfillShippingYear), + Route('/scripts/backfill_gate_dates', + maintenance_scripts.BackfillGateDates), ] dev_routes: list[Route] = []