Update jobs to latest SQL AAD (#485)

src/ArchivePackages/ArchivePackages.Job.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel.Design;
+using System.Data.SqlClient;
 using System.Diagnostics.Tracing;
 using System.Linq;
 using System.Threading.Tasks;
@@ -12,8 +13,6 @@
 using Microsoft.WindowsAzure.Storage.Blob;
 using Newtonsoft.Json.Linq;
 using NuGet.Jobs;
-using NuGet.Services.KeyVault;
-using NuGet.Services.Sql;
 
 namespace ArchivePackages
 {
@@ -53,22 +52,20 @@ public class Job : JobBase
         /// Blob containing the cursor data. Cursor data comprises of cursorDateTime
         /// </summary>
         public string CursorBlobName { get; set; }
-
-        private ISqlConnectionFactory _packageDbConnectionFactory;
 
         protected CloudBlobContainer SourceContainer { get; private set; }
 
         protected CloudBlobContainer PrimaryDestinationContainer { get; private set; }
 
         protected CloudBlobContainer SecondaryDestinationContainer { get; private set; }
 
+        private SqlConnectionStringBuilder PackageDatabase { get; set; }
+
         public Job() : base(JobEventSource.Log) { }
 
         public override void Init(IServiceContainer serviceContainer, IDictionary<string, string> jobArgsDictionary)
         {
-            var secretInjector = (ISecretInjector)serviceContainer.GetService(typeof(ISecretInjector));
-            var packageDbConnectionString = JobConfigurationManager.GetArgument(jobArgsDictionary, JobArgumentNames.PackageDatabase);
-            _packageDbConnectionFactory = new AzureSqlConnectionFactory(packageDbConnectionString, secretInjector);
+            PackageDatabase = RegisterDatabase(serviceContainer, jobArgsDictionary, JobArgumentNames.PackageDatabase);
 
             Source = CloudStorageAccount.Parse(
                         JobConfigurationManager.GetArgument(jobArgsDictionary, JobArgumentNames.Source));
@@ -92,13 +89,18 @@ public override void Init(IServiceContainer serviceContainer, IDictionary<string
 
         public override async Task Run()
         {
-            JobEventSourceLog.PreparingToArchive(Source.Credentials.AccountName, SourceContainer.Name, PrimaryDestination.Credentials.AccountName, PrimaryDestinationContainer.Name, _packageDbConnectionFactory.DataSource, _packageDbConnectionFactory.InitialCatalog);
+            JobEventSourceLog.PreparingToArchive(
+                Source.Credentials.AccountName, SourceContainer.Name,
+                PrimaryDestination.Credentials.AccountName, PrimaryDestinationContainer.Name,
+                PackageDatabase.DataSource, PackageDatabase.InitialCatalog);
+
             await Archive(PrimaryDestinationContainer);
 
             // todo: consider reusing package query for primary and secondary archives
             if (SecondaryDestinationContainer != null)
             {
                 JobEventSourceLog.PreparingToArchive2(SecondaryDestination.Credentials.AccountName, SecondaryDestinationContainer.Name);
+
                 await Archive(SecondaryDestinationContainer);
             }
         }
@@ -124,9 +126,10 @@ private async Task Archive(CloudBlobContainer destinationContainer)
 
             JobEventSourceLog.CursorData(cursorDateTime.ToString(DateTimeFormatSpecifier));
 
-            JobEventSourceLog.GatheringPackagesToArchiveFromDb(_packageDbConnectionFactory.DataSource, _packageDbConnectionFactory.InitialCatalog);
+            JobEventSourceLog.GatheringPackagesToArchiveFromDb(PackageDatabase.DataSource, PackageDatabase.InitialCatalog);
+
             List<PackageRef> packages;
-            using (var connection = await _packageDbConnectionFactory.CreateAsync())
+            using (var connection = await OpenSqlConnectionAsync(JobArgumentNames.PackageDatabase))
             {
                 packages = (await connection.QueryAsync<PackageRef>(@"
 			    SELECT pr.Id, p.NormalizedVersion AS Version, p.Hash, p.LastEdited, p.Published
@@ -135,7 +138,8 @@ FROM Packages p
 			    WHERE Published > @cursorDateTime OR LastEdited > @cursorDateTime", new { cursorDateTime = cursorDateTime }))
                     .ToList();
             }
-            JobEventSourceLog.GatheredPackagesToArchiveFromDb(packages.Count, _packageDbConnectionFactory.DataSource, _packageDbConnectionFactory.InitialCatalog);
+
+            JobEventSourceLog.GatheredPackagesToArchiveFromDb(packages.Count, PackageDatabase.DataSource, PackageDatabase.InitialCatalog);
 
             var archiveSet = packages
                 .AsParallel()

src/ArchivePackages/ArchivePackages.csproj

@@ -75,9 +75,6 @@
     <PackageReference Include="Newtonsoft.Json">
       <Version>9.0.1</Version>
     </PackageReference>
-    <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
-    </PackageReference>
     <PackageReference Include="System.Net.Http">
       <Version>4.3.3</Version>
     </PackageReference>

src/Gallery.CredentialExpiration/Gallery.CredentialExpiration.csproj

@@ -92,9 +92,6 @@
     <PackageReference Include="Newtonsoft.Json">
       <Version>9.0.1</Version>
     </PackageReference>
-    <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30263</Version>
-    </PackageReference>
     <PackageReference Include="NuGet.Services.Storage">
       <Version>2.1.3</Version>
     </PackageReference>

src/Gallery.CredentialExpiration/GalleryCredentialExpiration.cs

@@ -6,20 +6,22 @@
 using System.Data.SqlClient;
 using System.Linq;
 using System.Threading.Tasks;
-using NuGet.Services.Sql;
 using Gallery.CredentialExpiration.Models;
 
 namespace Gallery.CredentialExpiration
 {
     public class GalleryCredentialExpiration : ICredentialExpirationExporter
     {
         private readonly CredentialExpirationJobMetadata _jobMetadata;
-        private readonly ISqlConnectionFactory _galleryDatabase;
 
-        public GalleryCredentialExpiration(CredentialExpirationJobMetadata jobMetadata, ISqlConnectionFactory galleryDatabase)
+        private Func<Task<SqlConnection>> OpenGallerySqlConnectionAsync { get; }
+
+        public GalleryCredentialExpiration(
+            CredentialExpirationJobMetadata jobMetadata,
+            Func<Task<SqlConnection>> openGallerySqlConnectionAsync)
         {
             _jobMetadata = jobMetadata;
-            _galleryDatabase = galleryDatabase;
+            OpenGallerySqlConnectionAsync = openGallerySqlConnectionAsync;
         }
 
         /// <summary>
@@ -51,7 +53,7 @@ public async Task<List<ExpiredCredentialData>> GetCredentialsAsync(TimeSpan time
             var minNotificationDate = ConvertToString(GetMinNotificationDate());
 
             // Connect to database
-            using (var galleryConnection = await _galleryDatabase.CreateAsync())
+            using (var galleryConnection = await OpenGallerySqlConnectionAsync())
             {
                 // Fetch credentials that expire in _warnDaysBeforeExpiration days 
                 // + the user's e-mail address

src/Gallery.CredentialExpiration/Job.cs

@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.ComponentModel.Design;
 using System.Data.SqlClient;
@@ -15,10 +14,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.WindowsAzure.Storage;
 using Newtonsoft.Json;
-using Newtonsoft.Json.Linq;
 using NuGet.Jobs;
-using NuGet.Services.KeyVault;
-using NuGet.Services.Sql;
 using NuGet.Services.Storage;
 
 namespace Gallery.CredentialExpiration
@@ -34,8 +30,6 @@ public class Job : JobBase
         private string _galleryBrand;
         private string _galleryAccountUrl;
 
-        private ISqlConnectionFactory _galleryDatabase;
-
         private string _mailFrom;
         private SmtpClient _smtpClient;
 
@@ -47,9 +41,7 @@ public override void Init(IServiceContainer serviceContainer, IDictionary<string
         {
             _whatIf = JobConfigurationManager.TryGetBoolArgument(jobArgsDictionary, JobArgumentNames.WhatIf);
 
-            var secretInjector = (ISecretInjector)serviceContainer.GetService(typeof(ISecretInjector));
-            var databaseConnectionString = JobConfigurationManager.GetArgument(jobArgsDictionary, JobArgumentNames.GalleryDatabase);
-            _galleryDatabase = new AzureSqlConnectionFactory(databaseConnectionString, secretInjector);
+            RegisterDatabase(serviceContainer, jobArgsDictionary, JobArgumentNames.GalleryDatabase);
 
             _galleryBrand = JobConfigurationManager.GetArgument(jobArgsDictionary, MyJobArgumentNames.GalleryBrand);
             _galleryAccountUrl = JobConfigurationManager.GetArgument(jobArgsDictionary, MyJobArgumentNames.GalleryAccountUrl);
@@ -71,12 +63,17 @@ public override void Init(IServiceContainer serviceContainer, IDictionary<string
             _storage = storageFactory.Create();
         }
 
+        public Task<SqlConnection> OpenGallerySqlConnectionAsync()
+        {
+            return OpenSqlConnectionAsync(JobArgumentNames.GalleryDatabase);
+        }
+
         public override async Task Run()
         {
             var jobRunTime = DateTimeOffset.UtcNow;
             // Default values
             var jobCursor = new JobRunTimeCursor( jobCursorTime: jobRunTime, maxProcessedCredentialsTime: jobRunTime );
-            var galleryCredentialExpiration = new GalleryCredentialExpiration(new CredentialExpirationJobMetadata(jobRunTime, _warnDaysBeforeExpiration, jobCursor), _galleryDatabase);
+            var galleryCredentialExpiration = new GalleryCredentialExpiration(new CredentialExpirationJobMetadata(jobRunTime, _warnDaysBeforeExpiration, jobCursor), OpenGallerySqlConnectionAsync);
 
             try
             {
@@ -89,7 +86,7 @@ public override async Task Run()
                     // Load from cursor
                     // Throw if the schema is not correct to ensure that not-intended emails are sent.
                     jobCursor = JsonConvert.DeserializeObject<JobRunTimeCursor>(content, new JsonSerializerSettings() { MissingMemberHandling = MissingMemberHandling.Error });
-                    galleryCredentialExpiration = new GalleryCredentialExpiration(new CredentialExpirationJobMetadata(jobRunTime, _warnDaysBeforeExpiration, jobCursor), _galleryDatabase);
+                    galleryCredentialExpiration = new GalleryCredentialExpiration(new CredentialExpirationJobMetadata(jobRunTime, _warnDaysBeforeExpiration, jobCursor), OpenGallerySqlConnectionAsync);
                 }
 
                 // Connect to database

src/Gallery.Maintenance/DeleteExpiredVerificationKeysTask.cs

@@ -9,6 +9,7 @@
 using System.Threading.Tasks;
 using Gallery.Maintenance.Models;
 using Microsoft.Extensions.Logging;
+using NuGet.Jobs;
 
 namespace Gallery.Maintenance
 {
@@ -37,7 +38,7 @@ public override async Task RunAsync(Job job)
         {
             IEnumerable<PackageVerificationKey> expiredKeys;
 
-            using (var connection = await job.GalleryDatabase.CreateAsync())
+            using (var connection = await job.OpenSqlConnectionAsync(JobArgumentNames.GalleryDatabase))
             {
                 expiredKeys = await connection.QueryWithRetryAsync<PackageVerificationKey>(
                     SelectQuery,
@@ -59,7 +60,7 @@ public override async Task RunAsync(Job job)
 
             if (expectedRowCount > 0)
             {
-                using (var connection = await job.GalleryDatabase.CreateAsync())
+                using (var connection = await job.OpenSqlConnectionAsync(JobArgumentNames.GalleryDatabase))
                 {
                     using (var transaction = connection.BeginTransaction())
                     {

src/Gallery.Maintenance/Gallery.Maintenance.csproj

@@ -67,9 +67,6 @@
     <PackageReference Include="Newtonsoft.Json">
       <Version>9.0.1</Version>
     </PackageReference>
-    <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30263</Version>
-    </PackageReference>
     <PackageReference Include="System.Net.Http">
       <Version>4.3.3</Version>
     </PackageReference>

src/Gallery.Maintenance/Job.cs

@@ -4,13 +4,11 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel.Design;
-using System.Data;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using NuGet.Jobs;
 using NuGet.Services.KeyVault;
-using NuGet.Services.Sql;
 
 namespace Gallery.Maintenance
 {
@@ -19,15 +17,9 @@ namespace Gallery.Maintenance
     /// </summary>
     public class Job : JobBase
     {
-
-        public ISqlConnectionFactory GalleryDatabase { get; private set; }
-
         public override void Init(IServiceContainer serviceContainer, IDictionary<string, string> jobArgsDictionary)
         {
-            var secretInjector = (ISecretInjector)serviceContainer.GetService(typeof(ISecretInjector));
-            var databaseConnectionString = JobConfigurationManager.GetArgument(jobArgsDictionary, JobArgumentNames.GalleryDatabase);
-
-            GalleryDatabase = new AzureSqlConnectionFactory(databaseConnectionString, secretInjector);
+            RegisterDatabase(serviceContainer, jobArgsDictionary, JobArgumentNames.GalleryDatabase);
         }
 
         public override async Task Run()