diff --git a/features/legacy_migrator_coverage.feature b/features/legacy_migrator_coverage.feature index 0edec1c76..9eb9ca743 100644 --- a/features/legacy_migrator_coverage.feature +++ b/features/legacy_migrator_coverage.feature @@ -203,4 +203,10 @@ Feature: Legacy Data Migrator Coverage Scenario: Handle file system errors during backup Given I have a legacy project with read-only JSON files When I run the legacy data migration - Then the migration should handle backup errors gracefully \ No newline at end of file + Then the migration should handle backup errors gracefully + Scenario: get_all_for_project is called only once for multiple plans + Given I have a legacy project with multiple plans to migrate efficiently + When I run the legacy data migration with query tracking + Then get_all_for_project should be called exactly once + And the migration should return true + And all plans should be migrated correctly diff --git a/features/steps/legacy_migrator_steps.py b/features/steps/legacy_migrator_steps.py index d33a79e20..55a7bb213 100644 --- a/features/steps/legacy_migrator_steps.py +++ b/features/steps/legacy_migrator_steps.py @@ -1012,3 +1012,76 @@ def step_check_backup_error_handling(context: Context) -> None: """Check that migration handled backup errors gracefully.""" # The migration may fail or succeed, but shouldn't crash assert context.migration_result in [True, False] + + +@given("I have a legacy project with multiple plans to migrate efficiently") +def step_create_project_multiple_plans_efficiency(context: Context) -> None: + """Create a legacy project with multiple plans for efficiency testing.""" + project_dir = context.temp_dir / "efficiency_test" + project_dir.mkdir(exist_ok=True) + + cleveragents_dir = project_dir / ".cleveragents" + cleveragents_dir.mkdir(exist_ok=True) + + (cleveragents_dir / "project.name").write_text("efficiency_test") + + # Create multiple plans to ensure the loop runs more than once + plans_data = { + "plan_alpha": {"prompt": "Alpha plan", "status": "pending", "current": False}, + "plan_beta": {"prompt": "Beta plan", "status": "built", "current": False}, + "plan_gamma": {"prompt": "Gamma plan", "status": "applied", "current": True}, + "plan_delta": {"prompt": "Delta plan", "status": "pending", "current": False}, + } + (cleveragents_dir / "plans.json").write_text(json.dumps(plans_data)) + (cleveragents_dir / "current").write_text("plan_gamma") + + context.project_path = project_dir + + +@when("I run the legacy data migration with query tracking") +def step_run_migration_with_tracking(context: Context) -> None: + """Run migration while tracking how many times get_all_for_project is called.""" + # Find the plans repository class to patch at the class level + with context.unit_of_work.transaction() as ctx: + repo_class = type(ctx.plans) + + call_tracker = {"count": 0} + original_repo_method = repo_class.get_all_for_project + + def tracking_method(self, project_id): + call_tracker["count"] += 1 + return original_repo_method(self, project_id) + + repo_class.get_all_for_project = tracking_method + try: + migrator = LegacyDataMigrator(context.unit_of_work) + context.migration_result = migrator.migrate_project_data(context.project_path) + finally: + repo_class.get_all_for_project = original_repo_method + + context.get_all_for_project_call_count = call_tracker["count"] + + +@then("get_all_for_project should be called exactly once") +def step_check_get_all_called_once(context: Context) -> None: + """Verify get_all_for_project was called exactly once (not once per plan).""" + call_count = context.get_all_for_project_call_count + assert call_count == 1, ( + f"Expected get_all_for_project to be called exactly once, " + f"but it was called {call_count} times. " + f"This indicates an N+1 query pattern is still present." + ) + + +@then("all plans should be migrated correctly") +def step_check_all_plans_migrated(context: Context) -> None: + """Verify all plans from the efficiency test were migrated.""" + with context.unit_of_work.transaction() as ctx: + project = ctx.projects.get_by_name("efficiency_test") + assert project is not None, "Project should have been created" + plans = ctx.plans.get_all_for_project(project.id) + plan_names = {p.name for p in plans} + expected_names = {"plan_alpha", "plan_beta", "plan_gamma", "plan_delta"} + assert plan_names == expected_names, ( + f"Expected plans {expected_names}, got {plan_names}" + ) diff --git a/src/cleveragents/infrastructure/database/legacy_migrator.py b/src/cleveragents/infrastructure/database/legacy_migrator.py index 2e34d2ec3..7e5309860 100644 --- a/src/cleveragents/infrastructure/database/legacy_migrator.py +++ b/src/cleveragents/infrastructure/database/legacy_migrator.py @@ -101,9 +101,11 @@ class LegacyDataMigrator: with open(plans_json) as f: plans_data = json.load(f) + # Fetch all existing plans once before the loop to avoid + # N+1 database queries (one query per plan iteration). + all_plans = ctx.plans.get_all_for_project(project.id) for plan_name, plan_info in plans_data.items(): # Check if plan already exists - all_plans = ctx.plans.get_all_for_project(project.id) existing_plan = next( (p for p in all_plans if p.name == plan_name), None )