Update GHES generate script to accept --use-github-storage option

src/OctoshiftCLI.IntegrationTests/GhesToGithub.cs

@@ -25,9 +25,11 @@ public sealed class GhesToGithub : IDisposable
     private readonly GithubClient _sourceGithubClient;
     private readonly GithubApi _sourceGithubApi;
     private readonly BlobServiceClient _blobServiceClient;
-    private readonly Dictionary<string, string> _tokens;
     private readonly DateTime _startTime;
     private readonly ArchiveUploader _archiveUploader;
+    private readonly string _sourceGithubToken;
+    private readonly string _targetGithubToken;
+    private readonly string _azureStorageConnectionString;
 
     public GhesToGithub(ITestOutputHelper output)
     {
@@ -36,45 +38,60 @@ public GhesToGithub(ITestOutputHelper output)
 
         var logger = new OctoLogger(_ => { }, x => _output.WriteLine(x), _ => { }, _ => { });
 
-        var sourceGithubToken = Environment.GetEnvironmentVariable("GHES_PAT");
-        var targetGithubToken = Environment.GetEnvironmentVariable("GHEC_PAT");
-        var azureStorageConnectionString = Environment.GetEnvironmentVariable($"AZURE_STORAGE_CONNECTION_STRING_GHES_{TestHelper.GetOsName().ToUpper()}");
-        _tokens = new Dictionary<string, string>
-        {
-            ["GH_SOURCE_PAT"] = sourceGithubToken,
-            ["GH_PAT"] = targetGithubToken,
-            ["AZURE_STORAGE_CONNECTION_STRING"] = azureStorageConnectionString
-        };
+        _sourceGithubToken = Environment.GetEnvironmentVariable("GHES_PAT");
+        _targetGithubToken = Environment.GetEnvironmentVariable("GHEC_PAT");
+        _azureStorageConnectionString = Environment.GetEnvironmentVariable($"AZURE_STORAGE_CONNECTION_STRING_GHES_{TestHelper.GetOsName().ToUpper()}");
 
         _versionClient = new HttpClient();
         _archiveUploader = new ArchiveUploader(_targetGithubClient, logger);
 
         _sourceGithubHttpClient = new HttpClient();
-        _sourceGithubClient = new GithubClient(logger, _sourceGithubHttpClient, new VersionChecker(_versionClient, logger), new RetryPolicy(logger), new DateTimeProvider(), sourceGithubToken);
+        _sourceGithubClient = new GithubClient(logger, _sourceGithubHttpClient, new VersionChecker(_versionClient, logger), new RetryPolicy(logger), new DateTimeProvider(), _sourceGithubToken);
         _sourceGithubApi = new GithubApi(_sourceGithubClient, GHES_API_URL, new RetryPolicy(logger), _archiveUploader);
 
         _targetGithubHttpClient = new HttpClient();
-        _targetGithubClient = new GithubClient(logger, _targetGithubHttpClient, new VersionChecker(_versionClient, logger), new RetryPolicy(logger), new DateTimeProvider(), targetGithubToken);
+        _targetGithubClient = new GithubClient(logger, _targetGithubHttpClient, new VersionChecker(_versionClient, logger), new RetryPolicy(logger), new DateTimeProvider(), _targetGithubToken);
         _targetGithubApi = new GithubApi(_targetGithubClient, "https://api.github.com", new RetryPolicy(logger), _archiveUploader);
 
-        _blobServiceClient = new BlobServiceClient(azureStorageConnectionString);
+        _blobServiceClient = new BlobServiceClient(_azureStorageConnectionString);
 
         _sourceHelper = new TestHelper(_output, _sourceGithubApi, _sourceGithubClient) { GithubApiBaseUrl = GHES_API_URL };
         _targetHelper = new TestHelper(_output, _targetGithubApi, _targetGithubClient, _blobServiceClient);
     }
 
-    public async Task Basic()
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task Basic(bool useGithubStorage)
     {
         var githubSourceOrg = $"e2e-testing-{TestHelper.GetOsName()}";
         var githubTargetOrg = $"octoshift-e2e-ghes-{TestHelper.GetOsName()}";
         const string repo1 = "repo-1";
         const string repo2 = "repo-2";
 
+        var tokens = new Dictionary<string, string> { };
+
+        tokens = useGithubStorage
+            ? new Dictionary<string, string>
+            {
+                ["GH_SOURCE_PAT"] = _sourceGithubToken,
+                ["GH_PAT"] = _targetGithubToken,
+            }
+            : new Dictionary<string, string>
+            {
+                ["GH_SOURCE_PAT"] = _sourceGithubToken,
+                ["GH_PAT"] = _targetGithubToken,
+                ["AZURE_STORAGE_CONNECTION_STRING"] = _azureStorageConnectionString
+            };
+
         var retryPolicy = new RetryPolicy(null);
 
         await retryPolicy.Retry(async () =>
         {
-            await _targetHelper.ResetBlobContainers();
+            if (!useGithubStorage)
+            {
+                await _targetHelper.ResetBlobContainers();
+            }
 
             await _sourceHelper.ResetGithubTestEnvironment(githubSourceOrg);
             await _targetHelper.ResetGithubTestEnvironment(githubTargetOrg);
@@ -83,8 +100,14 @@ await retryPolicy.Retry(async () =>
             await _sourceHelper.CreateGithubRepo(githubSourceOrg, repo2);
         });
 
-        await _targetHelper.RunGeiCliMigration(
-            $"generate-script --github-source-org {githubSourceOrg} --github-target-org {githubTargetOrg} --ghes-api-url {GHES_API_URL} --download-migration-logs", _tokens);
+        // Build the command with conditional option
+        var command = $"generate-script --github-source-org {githubSourceOrg} --github-target-org {githubTargetOrg} --ghes-api-url {GHES_API_URL} --download-migration-logs";
+        if (useGithubStorage)
+        {
+            command += " --use-github-storage true";
+        }
+
+        await _targetHelper.RunGeiCliMigration(command, tokens);
 
         _targetHelper.AssertNoErrorInLogs(_startTime);
 
@@ -96,6 +119,7 @@ await _targetHelper.RunGeiCliMigration(
         _targetHelper.AssertMigrationLogFileExists(githubTargetOrg, repo1);
         _targetHelper.AssertMigrationLogFileExists(githubTargetOrg, repo2);
     }
+
     public void Dispose()
     {
         _sourceGithubHttpClient?.Dispose();

src/OctoshiftCLI.Tests/gei/Commands/GenerateScript/GenerateScriptCommandHandlerTests.cs

@@ -1102,6 +1102,63 @@ exit 1
             _script.Should().NotContain(expected);
         }
 
+        [Fact]
+        public async Task Sequential_Github_Single_Repo_With_UseGithubStorage()
+        {
+            // Arrange
+            _mockGithubApi
+                .Setup(m => m.GetRepos(SOURCE_ORG))
+                .ReturnsAsync(new[] { (REPO, "private") });
+
+            var expected = $"Exec {{ gh gei migrate-repo --github-source-org \"{SOURCE_ORG}\" --source-repo \"{REPO}\" --github-target-org \"{TARGET_ORG}\" --target-repo \"{REPO}\" --target-repo-visibility private --use-github-storage true }}";
+
+            // Act
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = SOURCE_ORG,
+                GithubTargetOrg = TARGET_ORG,
+                Output = new FileInfo("unit-test-output"),
+                Sequential = true,
+                UseGithubStorage = true
+            };
+            await _handler.Handle(args);
+
+            _script = TrimNonExecutableLines(_script);
+
+            // Assert
+            _script.Should().Be(expected);
+        }
+
+        [Fact]
+        public async Task Parallel_Github_Single_Repo_With_UseGithubStorage()
+        {
+            // Arrange
+            _mockGithubApi
+                .Setup(m => m.GetRepos(SOURCE_ORG))
+                .ReturnsAsync(new[] { (REPO, "private") });
+
+            var expected = new StringBuilder();
+            expected.AppendLine($"$MigrationID = ExecAndGetMigrationID {{ gh gei migrate-repo --github-source-org \"{SOURCE_ORG}\" --source-repo \"{REPO}\" --github-target-org \"{TARGET_ORG}\" --target-repo \"{REPO}\" --queue-only --target-repo-visibility private --use-github-storage true }}");
+            expected.AppendLine($"$RepoMigrations[\"{REPO}\"] = $MigrationID");
+            expected.Append($"if ($RepoMigrations[\"{REPO}\"]) {{ gh gei wait-for-migration --migration-id $RepoMigrations[\"{REPO}\"] }}");
+
+            // Act
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = SOURCE_ORG,
+                GithubTargetOrg = TARGET_ORG,
+                Output = new FileInfo("unit-test-output"),
+                UseGithubStorage = true
+            };
+            await _handler.Handle(args);
+
+            _script = TrimNonExecutableLines(_script, 19, 7);
+
+            // Assert
+            _script.Should().Be(expected.ToString());
+        }
+
+
         [Fact]
         public async Task Validates_Env_Vars_Blob_Storage_Not_Validated_When_GHES_3_8()
         {

src/OctoshiftCLI.Tests/gei/Commands/GenerateScript/GenerateScriptCommandTests.cs

@@ -39,7 +39,7 @@ public void Should_Have_Options()
             var command = new GenerateScriptCommand();
             command.Should().NotBeNull();
             command.Name.Should().Be("generate-script");
-            command.Options.Count.Should().Be(15);
+            command.Options.Count.Should().Be(16);
 
             TestHelpers.VerifyCommandOption(command.Options, "github-source-org", true);
             TestHelpers.VerifyCommandOption(command.Options, "github-target-org", true);
@@ -56,6 +56,7 @@ public void Should_Have_Options()
             TestHelpers.VerifyCommandOption(command.Options, "aws-region", false);
             TestHelpers.VerifyCommandOption(command.Options, "keep-archive", false);
             TestHelpers.VerifyCommandOption(command.Options, "target-api-url", false);
+            TestHelpers.VerifyCommandOption(command.Options, "use-github-storage", false, true);
         }
 
         [Fact]

src/gei/Commands/GenerateScript/GenerateScriptCommand.cs

@@ -36,6 +36,7 @@ public GenerateScriptCommand() : base(
             AddOption(GithubSourcePat);
             AddOption(Verbose);
             AddOption(KeepArchive);
+            AddOption(UseGithubStorage);
         }
         public Option<string> GithubSourceOrg { get; } = new("--github-source-org")
         {
@@ -99,6 +100,11 @@ public GenerateScriptCommand() : base(
         {
             Description = "The URL of the target API, if not migrating to github.com. Defaults to https://api.github.com"
         };
+        public Option<bool> UseGithubStorage { get; } = new("--use-github-storage")
+        {
+            IsHidden = true,
+            Description = "Enables multipart uploads to a GitHub owned storage for use during migration",
+        };
 
         public override GenerateScriptCommandHandler BuildHandler(GenerateScriptCommandArgs args, IServiceProvider sp)
         {

src/gei/Commands/GenerateScript/GenerateScriptCommandArgs.cs

@@ -23,12 +23,21 @@ public class GenerateScriptCommandArgs : CommandArgs
         public string GithubSourcePat { get; set; }
         public bool KeepArchive { get; set; }
         public string TargetApiUrl { get; set; }
+        public bool UseGithubStorage { get; set; }
 
         public override void Validate(OctoLogger log)
         {
-            if (AwsBucketName.HasValue() && GhesApiUrl.IsNullOrWhiteSpace())
+            if (AwsBucketName.HasValue())
             {
-                throw new OctoshiftCliException("--ghes-api-url must be specified when --aws-bucket-name is specified.");
+                if (GhesApiUrl.IsNullOrWhiteSpace())
+                {
+                    throw new OctoshiftCliException("--ghes-api-url must be specified when --aws-bucket-name is specified.");
+                }
+
+                if (UseGithubStorage)
+                {
+                    throw new OctoshiftCliException("The --use-github-storage flag was provided with an AWS S3 Bucket name. Archive cannot be uploaded to both locations.");
+                }
             }
 
             if (NoSslVerify && GhesApiUrl.IsNullOrWhiteSpace())
@@ -45,6 +54,11 @@ public override void Validate(OctoLogger log)
                 {
                     throw new OctoshiftCliException("--ghes-api-url is invalid. Please check URL before trying again.");
                 }
+
+                if (UseGithubStorage)
+                {
+                    throw new OctoshiftCliException("--ghes-api-url must be specified when --use-github-storage is specified.");
+                }
             }
         }
     }