Skip to content

Commit 315ab74

Browse files
dannyjdevdannyjdev
authored
[Storage Migration] V3 jobs (#10228)
1 parent 7316c09 commit 315ab74

12 files changed

Lines changed: 191 additions & 56 deletions

File tree

‎src/Catalog/Persistence/AzureStorage.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ private static ICloudBlobDirectory GetCloudBlobDirectoryUri(Uri storageBaseUri)
9494

9595
var blobEndpoint = new Uri(storageBaseUri.GetComponents(UriComponents.SchemeAndServer, UriFormat.Unescaped));
9696
// Create BlobServiceClient with anonymous credentials
97-
var blobServiceClient = new BlobServiceClientFactory(blobEndpoint, new DefaultAzureCredential());
97+
var blobServiceClient = new BlobServiceClientFactory(blobEndpoint);
9898

9999
string containerName = pathSegments[0];
100100
string pathInContainer = string.Join("/", pathSegments.Skip(1));

‎src/Ng/CommandHelpers.cs‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
using Azure;
1212
using Azure.Core;
1313
using Azure.Identity;
14+
using Azure.Storage;
1415
using Azure.Storage.Blobs;
1516
using Azure.Storage.Queues;
1617
using Microsoft.Extensions.Logging;
@@ -40,6 +41,8 @@ public static class CommandHelpers
4041
};
4142
private static readonly IDictionary<string, string> ArgumentNames = new Dictionary<string, string>
4243
{
44+
{ Arguments.UseManagedIdentity, Arguments.UseManagedIdentity },
45+
{ Arguments.ClientId, Arguments.ClientId},
4346
{ Arguments.StorageBaseAddress, Arguments.StorageBaseAddress },
4447
{ Arguments.StorageAccountName, Arguments.StorageAccountName },
4548
{ Arguments.StorageKeyValue, Arguments.StorageKeyValue },
@@ -52,7 +55,6 @@ public static class CommandHelpers
5255
{ Arguments.StorageOperationMaxExecutionTimeInSeconds, Arguments.StorageOperationMaxExecutionTimeInSeconds },
5356
{ Arguments.StorageServerTimeoutInSeconds, Arguments.StorageServerTimeoutInSeconds },
5457
{ Arguments.StorageInitializeContainer, Arguments.StorageInitializeContainer },
55-
{ Arguments.ClientId, Arguments.ClientId },
5658
};
5759

5860
public static IDictionary<string, string> GetArguments(string[] args, int start, out ICachingSecretInjector secretInjector)
@@ -164,6 +166,8 @@ public static CatalogStorageFactory CreateSuffixedStorageFactory(
164166

165167
IDictionary<string, string> names = new Dictionary<string, string>
166168
{
169+
{ Arguments.UseManagedIdentity, Arguments.UseManagedIdentity },
170+
{ Arguments.ClientId, Arguments.ClientId},
167171
{ Arguments.StorageBaseAddress, Arguments.StorageBaseAddress + suffix },
168172
{ Arguments.StorageAccountName, Arguments.StorageAccountName + suffix },
169173
{ Arguments.StorageUseManagedIdentity, Arguments.StorageUseManagedIdentity + suffix },
@@ -198,7 +202,6 @@ private static CatalogStorageFactory CreateStorageFactoryImpl(
198202
if (!string.IsNullOrEmpty(storageBaseAddressStr))
199203
{
200204
storageBaseAddressStr = storageBaseAddressStr.TrimEnd('/') + "/";
201-
202205
storageBaseAddress = new Uri(storageBaseAddressStr);
203206
}
204207

@@ -427,7 +430,16 @@ private static IBlobServiceClientFactory GetBlobServiceClient(
427430
IDictionary<string, string> argumentNameMap)
428431
{
429432
bool storageUseManagedIdentity = arguments.GetOrDefault(argumentNameMap[Arguments.StorageUseManagedIdentity], defaultValue: false);
430-
if (storageUseManagedIdentity)
433+
bool useManagedIdentity = storageUseManagedIdentity || arguments.GetOrDefault(argumentNameMap[Arguments.UseManagedIdentity], defaultValue: false);
434+
435+
var storageKeyValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageKeyValue]);
436+
var storageSasValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageSasValue]);
437+
438+
bool hasStorageKeyOrSas = !string.IsNullOrEmpty(storageKeyValue) || !string.IsNullOrEmpty(storageSasValue);
439+
440+
// This comparison is due to some jobs using both global and china storages in a single instance.
441+
// They require MSI auth for global storage and SAS/SAK auth for china storage.
442+
if (useManagedIdentity && !hasStorageKeyOrSas)
431443
{
432444
var managedIdentityClientId = arguments.GetOrThrow<string>(argumentNameMap[Arguments.ClientId]);
433445
var identity = new ManagedIdentityCredential(managedIdentityClientId);
@@ -444,7 +456,16 @@ private static QueueServiceClient GetQueueServiceClient(
444456
IDictionary<string, string> argumentNameMap)
445457
{
446458
bool storageUseManagedIdentity = arguments.GetOrDefault(argumentNameMap[Arguments.StorageUseManagedIdentity], defaultValue: false);
447-
if (storageUseManagedIdentity)
459+
bool useManagedIdentity = storageUseManagedIdentity || arguments.GetOrDefault(argumentNameMap[Arguments.UseManagedIdentity], defaultValue: false);
460+
461+
var storageKeyValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageKeyValue]);
462+
var storageSasValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageSasValue]);
463+
464+
bool hasStorageKeyOrSas = !string.IsNullOrEmpty(storageKeyValue) || !string.IsNullOrEmpty(storageSasValue);
465+
466+
// This comparison is due to some jobs using both global and china storages in a single instance.
467+
// They require MSI auth for global storage and SAS/SAK auth for china storage.
468+
if (useManagedIdentity && !hasStorageKeyOrSas)
448469
{
449470
var managedIdentityClientId = arguments.GetOrThrow<string>(argumentNameMap[Arguments.ClientId]);
450471
var identity = new ManagedIdentityCredential(managedIdentityClientId);

‎src/Ng/Jobs/Catalog2DnxJob.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -45,7 +45,7 @@ public override string GetUsage()
4545
+ $"-{Arguments.StoragePath} <path> "
4646
+ $"[-{Arguments.VaultName} <keyvault-name> "
4747
+ $"-{Arguments.UseManagedIdentity} true|false "
48-
+ $"-{Arguments.ClientId} <keyvault-client-id> Should not be set if {Arguments.UseManagedIdentity} is true"
48+
+ $"-{Arguments.ClientId} <client-id> If {Arguments.UseManagedIdentity} is true this is used for managed identity authentication, if false, is used for KeyVault certificate authentication"
4949
+ $"-{Arguments.CertificateThumbprint} <keyvault-certificate-thumbprint> Should not be set if {Arguments.UseManagedIdentity} is true"
5050
+ $"[-{Arguments.ValidateCertificate} true|false]]] "
5151
+ $"[-{Arguments.Verbose} true|false] "
@@ -122,4 +122,4 @@ protected override async Task RunInternalAsync(CancellationToken cancellationTok
122122
}
123123
}
124124
}
125-
}
125+
}

‎src/Ng/Jobs/Db2CatalogJob.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -63,7 +63,7 @@ public override string GetUsage()
6363
+ $"-{Arguments.StoragePath} <path> "
6464
+ $"[-{Arguments.VaultName} <keyvault-name> "
6565
+ $"-{Arguments.UseManagedIdentity} true|false "
66-
+ $"-{Arguments.ClientId} <keyvault-client-id> Should not be set if {Arguments.UseManagedIdentity} is true"
66+
+ $"-{Arguments.ClientId} <client-id> If {Arguments.UseManagedIdentity} is true this is used for managed identity authentication, if false, is used for KeyVault certificate authentication"
6767
+ $"-{Arguments.CertificateThumbprint} <keyvault-certificate-thumbprint> Should not be set if {Arguments.UseManagedIdentity} is true"
6868
+ $"[-{Arguments.ValidateCertificate} true|false]]] "
6969
+ $"-{Arguments.StorageTypeAuditing} file|azure "
@@ -413,4 +413,4 @@ private async Task<DateTime> Deletes2Catalog(
413413
return lastDeleted;
414414
}
415415
}
416-
}
416+
}

‎src/Ng/Jobs/FixCatalogCachingJob.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -52,7 +52,7 @@ public override string GetUsage()
5252
+ $"-{Arguments.StoragePath} <path> "
5353
+ $"[-{Arguments.VaultName} <keyvault-name> "
5454
+ $"-{Arguments.UseManagedIdentity} true|false "
55-
+ $"-{Arguments.ClientId} <keyvault-client-id> Should not be set if {Arguments.UseManagedIdentity} is true"
55+
+ $"-{Arguments.ClientId} <client-id> If {Arguments.UseManagedIdentity} is true this is used for managed identity authentication, if false, is used for KeyVault certificate authentication"
5656
+ $"-{Arguments.CertificateThumbprint} <keyvault-certificate-thumbprint> Should not be set if {Arguments.UseManagedIdentity} is true"
5757
+ $"[-{Arguments.ValidateCertificate} true|false ]]";
5858
}

‎src/Ng/Jobs/LightningJob.cs‎

Lines changed: 53 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -400,8 +400,10 @@ private void AddStorageCredentialArgument(StringBuilder argument, string sasToke
400400
{
401401
AppendArgument(argument, sasTokenArgument);
402402
}
403-
404-
AppendArgument(argument, storageKeyArgument);
403+
else if (!string.IsNullOrEmpty(_arguments.GetOrDefault<string>(storageKeyArgument)))
404+
{
405+
AppendArgument(argument, storageKeyArgument);
406+
}
405407
}
406408

407409
private async Task StrikeAsync()
@@ -504,32 +506,55 @@ private IContainer GetAutofacContainer()
504506

505507
services.Configure<Catalog2RegistrationConfiguration>(config =>
506508
{
509+
config.StorageUseManagedIdentity = _arguments.GetOrDefault<bool>(Arguments.UseManagedIdentity);
510+
config.StorageManagedIdentityClientId = _arguments.GetOrDefault<string>(Arguments.ClientId);
511+
507512
config.LegacyBaseUrl = _arguments.GetOrDefault<string>(Arguments.StorageBaseAddress);
508513
config.LegacyStorageContainer = _arguments.GetOrDefault<string>(Arguments.StorageContainer);
509-
config.StorageConnectionString = GetConnectionString(
510-
config.StorageConnectionString,
511-
Arguments.StorageAccountName,
512-
Arguments.StorageKeyValue,
513-
Arguments.StorageSasValue,
514-
Arguments.StorageSuffix);
515514

516515
config.GzippedBaseUrl = _arguments.GetOrDefault<string>(Arguments.CompressedStorageBaseAddress);
517516
config.GzippedStorageContainer = _arguments.GetOrDefault<string>(Arguments.CompressedStorageContainer);
518-
config.StorageConnectionString = GetConnectionString(
519-
config.StorageConnectionString,
520-
Arguments.CompressedStorageAccountName,
521-
Arguments.CompressedStorageKeyValue,
522-
Arguments.CompressedStorageSasValue,
523-
Arguments.StorageSuffix);
524517

525518
config.SemVer2BaseUrl = _arguments.GetOrDefault<string>(Arguments.SemVer2StorageBaseAddress);
526519
config.SemVer2StorageContainer = _arguments.GetOrDefault<string>(Arguments.SemVer2StorageContainer);
527-
config.StorageConnectionString = GetConnectionString(
528-
config.StorageConnectionString,
529-
Arguments.SemVer2StorageAccountName,
530-
Arguments.SemVer2StorageKeyValue,
531-
Arguments.SemVer2StorageSasValue,
532-
Arguments.StorageSuffix);
520+
521+
config.HasSasToken = new List<string>()
522+
{
523+
_arguments.GetOrDefault<string>(Arguments.StorageSasValue),
524+
_arguments.GetOrDefault<string>(Arguments.CompressedStorageSasValue),
525+
_arguments.GetOrDefault<string>(Arguments.SemVer2StorageSasValue)
526+
}
527+
.All(t => !string.IsNullOrEmpty(t));
528+
529+
if (config.StorageUseManagedIdentity && !config.HasSasToken)
530+
{
531+
var storageAccountName = _arguments.GetOrDefault<string>(Arguments.StorageAccountName);
532+
var storageSuffix = _arguments.GetOrDefault(Arguments.StorageSuffix, "core.windows.net");
533+
534+
config.StorageServiceUrl = $"https://{storageAccountName}.blob.{storageSuffix}";
535+
config.StorageConnectionString = $"BlobEndpoint={config.StorageServiceUrl}";
536+
}
537+
else
538+
{
539+
config.StorageConnectionString = GetConnectionString(
540+
config.StorageConnectionString,
541+
Arguments.StorageAccountName,
542+
Arguments.StorageKeyValue,
543+
Arguments.StorageSasValue,
544+
Arguments.StorageSuffix);
545+
config.StorageConnectionString = GetConnectionString(
546+
config.StorageConnectionString,
547+
Arguments.CompressedStorageAccountName,
548+
Arguments.CompressedStorageKeyValue,
549+
Arguments.CompressedStorageSasValue,
550+
Arguments.StorageSuffix);
551+
config.StorageConnectionString = GetConnectionString(
552+
config.StorageConnectionString,
553+
Arguments.SemVer2StorageAccountName,
554+
Arguments.SemVer2StorageKeyValue,
555+
Arguments.SemVer2StorageSasValue,
556+
Arguments.StorageSuffix);
557+
}
533558

534559
config.GalleryBaseUrl = _arguments.GetOrThrow<string>(Arguments.GalleryBaseAddress);
535560
var contentBaseAddress = _arguments.GetOrThrow<string>(Arguments.ContentBaseAddress);
@@ -568,7 +593,13 @@ private string GetConnectionString(
568593
}
569594
else
570595
{
571-
builder.AppendFormat("SharedAccessSignature={0};", _arguments.GetOrThrow<string>(accountSasArgument));
596+
var sasToken = _arguments.GetOrThrow<string>(accountSasArgument);
597+
// workaround for https://github.com/Azure/azure-sdk-for-net/issues/44373
598+
if (sasToken.StartsWith("?"))
599+
{
600+
sasToken = sasToken.Substring(1);
601+
}
602+
builder.AppendFormat("SharedAccessSignature={0};", sasToken);
572603
}
573604

574605
builder.AppendFormat("EndpointSuffix={0}", _arguments.GetOrDefault(endpointSuffixArgument, "core.windows.net"));

‎src/Ng/Jobs/NgJob.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -51,7 +51,7 @@ public static string GetUsageBase()
5151
return "Usage: ng [" + string.Join("|", NgJobFactory.JobMap.Keys) + "] "
5252
+ $"[-{Arguments.VaultName} <keyvault-name> "
5353
+ $"-{Arguments.UseManagedIdentity} true|false "
54-
+ $"-{Arguments.ClientId} <keyvault-client-id> Should not be set if {Arguments.UseManagedIdentity} is true"
54+
+ $"-{Arguments.ClientId} <client-id> If {Arguments.UseManagedIdentity} is true this is used for managed identity authentication, if false, is used for KeyVault certificate authentication"
5555
+ $"-{Arguments.CertificateThumbprint} <keyvault-certificate-thumbprint> Should not be set if {Arguments.UseManagedIdentity} is true"
5656
+ $"[-{Arguments.ValidateCertificate} true|false]]";
5757
}
@@ -76,4 +76,4 @@ public virtual async Task RunAsync(IDictionary<string, string> arguments, Cancel
7676
await RunInternalAsync(cancellationToken);
7777
}
7878
}
79-
}
79+
}

‎src/NuGet.Jobs.Catalog2Registration/Catalog2RegistrationConfiguration.cs‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) .NET Foundation. All rights reserved.
1+
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

44
using System;
@@ -11,6 +11,26 @@ public class Catalog2RegistrationConfiguration : ICommitCollectorConfiguration
1111
{
1212
private static readonly int DefaultMaxConcurrentHivesPerId = Enum.GetValues(typeof(HiveType)).Length;
1313

14+
/// <summary>
15+
/// Whether or not managed identity will be used as credential.
16+
/// </summary>
17+
public bool StorageUseManagedIdentity { get; set; }
18+
19+
/// <summary>
20+
/// Specific manage identity client id.
21+
/// </summary>
22+
public string StorageManagedIdentityClientId { get; set; }
23+
24+
/// <summary>
25+
/// Whether or not any storage contains a sas token.
26+
/// </summary>
27+
public bool HasSasToken { get; set; }
28+
29+
/// <summary>
30+
/// Azure storage service uri. e.g. https://<storage>.blob.core.windows.net
31+
/// </summary>
32+
public string StorageServiceUrl { get; set; }
33+
1434
/// <summary>
1535
/// The connection string used to connect to an Azure Blob Storage account. The connection string specifies
1636
/// the account name, the endpoint suffix (e.g. Azure vs. Azure China), and authentication credential (e.g. storage

‎src/NuGet.Jobs.Catalog2Registration/DependencyInjectionExtensions.cs‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Collections.Generic;
66
using System.Net.Http;
77
using Autofac;
8+
using Azure.Identity;
89
using Azure.Storage.Blobs;
910
using Microsoft.Extensions.Configuration;
1011
using Microsoft.Extensions.DependencyInjection;
@@ -31,10 +32,20 @@ public static ContainerBuilder AddCatalog2Registration(this ContainerBuilder con
3132
containerBuilder.AddV3();
3233

3334
RegisterCursorStorage(containerBuilder);
34-
3535
containerBuilder
36-
.RegisterStorageAccount<Catalog2RegistrationConfiguration>(c => c.StorageConnectionString, requestTimeout: DefaultBlobRequestOptions.ServerTimeout)
37-
.As<ICloudBlobClient>();
36+
.Register<ICloudBlobClient>(c =>
37+
{
38+
var options = c.Resolve<IOptionsSnapshot<Catalog2RegistrationConfiguration>>();
39+
40+
if (options.Value.StorageUseManagedIdentity && !options.Value.HasSasToken)
41+
{
42+
return CloudBlobClientWrapper.UsingMsi(options.Value.StorageConnectionString, clientId: options.Value.StorageManagedIdentityClientId, requestTimeout: DefaultBlobRequestOptions.ServerTimeout);
43+
}
44+
45+
return new CloudBlobClientWrapper(
46+
options.Value.StorageConnectionString,
47+
requestTimeout: DefaultBlobRequestOptions.ServerTimeout);
48+
});
3849

3950
containerBuilder.Register(c => new Catalog2RegistrationCommand(
4051
c.Resolve<ICollector>(),
@@ -57,7 +68,12 @@ private static void RegisterCursorStorage(ContainerBuilder containerBuilder)
5768
{
5869
var options = c.Resolve<IOptionsSnapshot<Catalog2RegistrationConfiguration>>();
5970

60-
// workaround for https://github.com/Azure/azure-sdk-for-net/issues/44373
71+
if (options.Value.StorageUseManagedIdentity && !options.Value.HasSasToken)
72+
{
73+
var credential = new ManagedIdentityCredential(options.Value.StorageManagedIdentityClientId);
74+
75+
return new BlobServiceClientFactory(new Uri(options.Value.StorageServiceUrl), credential);
76+
}
6177
var connectionString = options.Value.StorageConnectionString.Replace("SharedAccessSignature=?", "SharedAccessSignature=");
6278

6379
return new BlobServiceClientFactory(connectionString);

0 commit comments

Comments
 (0)