diff --git a/util/Seeder/Scenes/MigrationCohortExportScene.cs b/util/Seeder/Scenes/MigrationCohortExportScene.cs new file mode 100644 index 000000000000..ecbc7f632069 --- /dev/null +++ b/util/Seeder/Scenes/MigrationCohortExportScene.cs @@ -0,0 +1,232 @@ +using System.ComponentModel.DataAnnotations; +using System.Globalization; +using AutoMapper; +using Bit.Core.Billing.Enums; +using Bit.Core.Billing.Organizations.PlanMigration.Repositories; +using Bit.Core.Entities; +using Bit.Core.Enums; +using Bit.Core.Services; +using Bit.Core.Utilities; +using Bit.Infrastructure.EntityFramework.Repositories; +using LinqToDB.EntityFrameworkCore; +using CoreCohort = Bit.Core.Billing.Organizations.PlanMigration.Entities.OrganizationPlanMigrationCohort; +using CoreCohortAssignment = Bit.Core.Billing.Organizations.PlanMigration.Entities.OrganizationPlanMigrationCohortAssignment; +using CoreOrganization = Bit.Core.AdminConsole.Entities.Organization; +using EfCohort = Bit.Infrastructure.EntityFramework.Billing.Models.OrganizationPlanMigrationCohort; +using EfCohortAssignment = Bit.Infrastructure.EntityFramework.Billing.Models.OrganizationPlanMigrationCohortAssignment; +using EfOrganization = Bit.Infrastructure.EntityFramework.AdminConsole.Models.Organization; +using EfPlayItem = Bit.Infrastructure.EntityFramework.Models.PlayItem; +using MigrationPath = Bit.Core.Billing.Organizations.PlanMigration.Enums.MigrationPathId; + +namespace Bit.Seeder.Scenes; + +public readonly struct MigrationCohortExportSceneResult +{ + public Guid CohortId { get; init; } + public bool CohortCreated { get; init; } + public int OrganizationsCreated { get; init; } + public int AssignmentsCreated { get; init; } +} + +/// +/// Seeds a migration cohort plus N inert organizations, each assigned to the cohort, so the +/// Admin Portal cohort-management table has a populated cohort ready to "Export CSV". +/// +/// +/// The HTTP/API equivalent of the PM-36965 SQL seed script. The organizations are throwaway +/// placeholders (Disabled, no billing, no users) and exist only because +/// OrganizationPlanMigrationCohortAssignment has a FK to Organization with a UNIQUE +/// constraint on OrganizationId -- every assignment needs its own real organization row. +/// +/// The cohort, organizations, and assignments are bulk-inserted via LinqToDB BulkCopy (the same +/// path uses) so counts in the tens of thousands stay fast, all +/// inside a single transaction so a mid-sequence failure leaves the database untouched rather than +/// partly seeded. Each organization's CreationDate is spread across distinct seconds. The +/// export, however, orders on the assignment's (CreationDate, Id), and assignments +/// created in one bulk batch share a near-identical CreationDate, so the cursor advances on the +/// Id tiebreaker. That is exactly the bulk-loaded case the export's stored procedure is +/// documented to handle, so it remains correct. +/// +/// Cleanup: BulkCopy bypasses the repository CreateAsync hook that normally records +/// play-id tracking rows, so this scene bulk-inserts the rows itself when the +/// request carries an x-play-id. That lets DELETE /seed/{playId} tear the seeded +/// organizations down (their assignments cascade away via the FK). The cohort row is NOT play-id +/// tracked -- DestroySceneCommand only deletes users/organizations -- so an emptied cohort is +/// left behind after a play-id delete; remove it by name if desired +/// (DELETE FROM dbo.OrganizationPlanMigrationCohort WHERE Name = '<CohortName>'). +/// +/// LOCAL / NON-PROD ONLY -- the SeederApi refuses to run in production. +/// +public class MigrationCohortExportScene( + DatabaseContext db, + IMapper mapper, + IPlayIdService playIdService, + IOrganizationPlanMigrationCohortRepository cohortRepository) + : IScene +{ + /// + /// Inert plan applied to seeded placeholder organizations. Arbitrary -- these orgs never + /// behave like live organizations (Enabled = false, no billing, no users). + /// + private const PlanType InertPlanType = PlanType.EnterpriseAnnually; + + /// Anchor for the spread CreationDate, matching the SQL seed script's baseline. + private static readonly DateTime _creationAnchor = new(2026, 1, 1, 0, 0, 0, DateTimeKind.Utc); + + private const int MinOrgCount = 1; + private const int MaxOrgCount = 100_000; + + public class Request + { + [Required] + public required string CohortName { get; set; } + + // The scene enforces this range explicitly in SeedAsync -- the SeederApi executes scenes by + // deserializing arguments and calling SeedAsync directly, so DataAnnotations here are not run + // by the framework. The attribute documents the contract and applies if validation is wired up. + [Required] + [Range(MinOrgCount, MaxOrgCount)] + public required int OrgCount { get; set; } + + /// + /// Migration path for the cohort. Null produces a Churn-only cohort. Defaults to + /// to mirror the SQL seed script. + /// + public MigrationPath? MigrationPathId { get; set; } = MigrationPath.Enterprise2020AnnualToCurrent; + + /// + /// Distinctive name prefix for the seeded organizations so they can be identified and + /// cleaned up. Must stay short enough that the suffixed name fits Organization.Name. + /// + [MaxLength(40)] + public string NamePrefix { get; set; } = "pm36965-seed-"; + } + + public async Task> SeedAsync(Request request) + { + if (request.OrgCount is < MinOrgCount or > MaxOrgCount) + { + throw new ArgumentOutOfRangeException( + nameof(request), request.OrgCount, + $"{nameof(request.OrgCount)} must be between {MinOrgCount} and {MaxOrgCount}."); + } + + var existingCohort = await cohortRepository.GetByNameAsync(request.CohortName); + var cohortCreated = existingCohort is null; + + var cohort = existingCohort ?? BuildCohort(request); + var organizations = BuildOrganizations(request); + var assignments = BuildAssignments(organizations, cohort.Id); + + // Wrap every write in one transaction so the scene is all-or-nothing: a mid-sequence failure + // (e.g. a timeout on a large assignment insert) rolls back the cohort and organizations rather + // than leaving orphaned rows or organizations without their play-id tracking. LinqToDB's + // BulkCopy honors the EF transaction opened here (see UserRepository for the same pattern). + await using var transaction = await db.Database.BeginTransactionAsync(); + + // BulkCopy maps Core entities to their EF counterparts (navigation properties are ignored). + if (cohortCreated) + { + db.BulkCopy([mapper.Map(cohort)]); + } + + db.BulkCopy(organizations.Select(mapper.Map)); + db.BulkCopy(assignments.Select(mapper.Map)); + + // Record play-id tracking rows for the organizations so DELETE /seed/{playId} can clean them + // up. BulkCopy bypasses the repository CreateAsync hook that normally does this, so the scene + // inserts the PlayItem rows directly -- and in bulk, to keep the fast path fast. + RecordOrganizationsForPlayId(organizations); + + await transaction.CommitAsync(); + + return new SceneResult( + result: new MigrationCohortExportSceneResult + { + CohortId = cohort.Id, + CohortCreated = cohortCreated, + OrganizationsCreated = organizations.Count, + AssignmentsCreated = assignments.Count, + }, + mangleMap: []); + } + + private static CoreCohort BuildCohort(Request request) + { + var cohort = new CoreCohort + { + Name = request.CohortName, + MigrationPathId = request.MigrationPathId, + IsActive = true, + }; + cohort.SetNewId(); + return cohort; + } + + private static List BuildOrganizations(Request request) + { + var organizations = new List(request.OrgCount); + + for (var n = 1; n <= request.OrgCount; n++) + { + // Spread CreationDate across distinct seconds so the export keyset cursor advances over + // real dates, exactly as the SQL seed script does. + var creationDate = _creationAnchor.AddSeconds(n); + var suffix = n.ToString("D6", CultureInfo.InvariantCulture); + + organizations.Add(new CoreOrganization + { + Id = CoreHelpers.GenerateComb(), + Name = $"{request.NamePrefix}{suffix}", + BillingEmail = $"{request.NamePrefix}{n}@example.com", + Plan = "Enterprise (Annually)", + PlanType = InertPlanType, + Status = OrganizationStatusType.Created, + Enabled = false, // never behaves like a live organization + CreationDate = creationDate, + RevisionDate = creationDate, + }); + } + + return organizations; + } + + private static List BuildAssignments( + IEnumerable organizations, Guid cohortId) => + organizations.Select(org => + { + // CreationDate defaults to DateTime.UtcNow at construction and is not settable here, so + // assignments in one bulk batch share a near-identical value. RevisionDate is settable, so + // align it for tidiness. + var assignment = new CoreCohortAssignment + { + OrganizationId = org.Id, + CohortId = cohortId, + }; + assignment.RevisionDate = assignment.CreationDate; + assignment.SetNewId(); + return assignment; + }).ToList(); + + /// + /// Bulk-inserts a row per organization when the request is part of a play + /// session, mirroring what the repository tracking decorators do on CreateAsync. No-op when + /// no x-play-id is set. Deleting an organization cascade-removes its + /// row, so DELETE /seed/{playId} stays self-consistent. + /// + private void RecordOrganizationsForPlayId(IEnumerable organizations) + { + if (!playIdService.InPlay(out var playId)) + { + return; + } + + var playItems = organizations.Select(org => + { + var playItem = PlayItem.Create(org, playId); + playItem.SetNewId(); // repository normally assigns the COMB id; BulkCopy bypasses it + return playItem; + }); + db.BulkCopy(playItems.Select(mapper.Map)); + } +}