Jobs.Common helpers for NuGet.Services.Sql (#505)

NuGet · Aug 7, 2018 · 6c451f8 · 6c451f8
1 parent 32182ee
commit 6c451f8
Show file tree

Hide file tree

Showing 19 changed files with 429 additions and 36 deletions.
diff --git a/NuGet.Jobs.sln b/NuGet.Jobs.sln
@@ -151,6 +151,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Validation.Symbols.Tests",
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Validation.Symbols.Core.Tests", "tests\Validation.Symbols.Core.Tests\Validation.Symbols.Core.Tests.csproj", "{9ED642DF-4623-4EB2-8B72-52C6489BF9D6}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "NuGet.Jobs.Common.Tests", "tests\NuGet.Jobs.Common.Tests\NuGet.Jobs.Common.Tests.csproj", "{CE96428B-8138-4914-9999-2B391797FFF8}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -399,6 +401,10 @@ Global
 		{9ED642DF-4623-4EB2-8B72-52C6489BF9D6}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{9ED642DF-4623-4EB2-8B72-52C6489BF9D6}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{9ED642DF-4623-4EB2-8B72-52C6489BF9D6}.Release|Any CPU.Build.0 = Release|Any CPU
+		{CE96428B-8138-4914-9999-2B391797FFF8}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{CE96428B-8138-4914-9999-2B391797FFF8}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{CE96428B-8138-4914-9999-2B391797FFF8}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{CE96428B-8138-4914-9999-2B391797FFF8}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -464,6 +470,7 @@ Global
 		{2DD07A73-8C88-4429-BB24-C2813586EF92} = {678D7B14-F8BC-4193-99AF-2EE8AA390A02}
 		{640D29AB-4D1B-4FC7-AE67-AD12EE5AC503} = {6A776396-02B1-475D-A104-26940ADB04AB}
 		{9ED642DF-4623-4EB2-8B72-52C6489BF9D6} = {6A776396-02B1-475D-A104-26940ADB04AB}
+		{CE96428B-8138-4914-9999-2B391797FFF8} = {6A776396-02B1-475D-A104-26940ADB04AB}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {284A7AC3-FB43-4F1F-9C9C-2AF0E1F46C2B}

diff --git a/src/ArchivePackages/ArchivePackages.csproj b/src/ArchivePackages/ArchivePackages.csproj
@@ -76,7 +76,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="System.Net.Http">
       <Version>4.3.3</Version>

diff --git a/src/Gallery.CredentialExpiration/Gallery.CredentialExpiration.csproj b/src/Gallery.CredentialExpiration/Gallery.CredentialExpiration.csproj
@@ -93,7 +93,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30263</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Storage">
       <Version>2.1.3</Version>

diff --git a/src/Gallery.Maintenance/Gallery.Maintenance.csproj b/src/Gallery.Maintenance/Gallery.Maintenance.csproj
@@ -68,7 +68,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30263</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="System.Net.Http">
       <Version>4.3.3</Version>

diff --git a/src/NuGet.Jobs.Common/Extensions/SqlConnectionStringBuilderExtensions.cs b/src/NuGet.Jobs.Common/Extensions/SqlConnectionStringBuilderExtensions.cs
diff --git a/src/NuGet.Jobs.Common/JobBase.cs b/src/NuGet.Jobs.Common/JobBase.cs
@@ -1,11 +1,18 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
 using System.ComponentModel.Design;
+using System.Data.SqlClient;
 using System.Diagnostics.Tracing;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using NuGet.Jobs.Configuration;
+using NuGet.Services.KeyVault;
+using NuGet.Services.Sql;
 
 namespace NuGet.Jobs
 {
@@ -22,6 +29,7 @@ protected JobBase(EventSource jobEventSource)
         {
             JobName = GetType().ToString();
             _jobEventSource = jobEventSource;
+            SqlConnectionFactories = new Dictionary<string, ISqlConnectionFactory>();
         }
 
         public string JobName { get; private set; }
@@ -30,14 +38,180 @@ protected JobBase(EventSource jobEventSource)
 
         protected ILogger Logger { get; private set; }
 
+        private Dictionary<string, ISqlConnectionFactory> SqlConnectionFactories { get; }
+
         public void SetLogger(ILoggerFactory loggerFactory, ILogger logger)
         {
             LoggerFactory = loggerFactory;
             Logger = logger;
         }
 
+        /// <summary>
+        /// Initialize the job, provided the service container and configuration.
+        /// </summary>
         public abstract void Init(IServiceContainer serviceContainer, IDictionary<string, string> jobArgsDictionary);
 
+        /// <summary>
+        /// Run the job.
+        /// </summary>
         public abstract Task Run();
+
+
+        /// <summary>
+        /// Test connection early to fail fast, and log connection diagnostics.
+        /// </summary>
+        private async Task TestConnection(string name, ISqlConnectionFactory connectionFactory)
+        {
+            try
+            {
+                using (var connection = await connectionFactory.OpenAsync())
+                using (var cmd = new SqlCommand("SELECT CONCAT(CURRENT_USER, '/', SYSTEM_USER)", connection))
+                {
+                    var result = cmd.ExecuteScalar();
+                    var user = result.ToString();
+                    Logger.LogInformation("Verified CreateSqlConnectionAsync({name}) connects to database {DataSource}/{InitialCatalog} as {User}",
+                        name, connectionFactory.DataSource, connectionFactory.InitialCatalog, user);
+                }
+            }
+            catch (Exception e)
+            {
+                Logger.LogError(0, e, "Failed to connect to database {DataSource}/{InitialCatalog}",
+                    connectionFactory.DataSource, connectionFactory.InitialCatalog);
+
+                throw;
+            }
+        }
+
+        public SqlConnectionStringBuilder GetDatabaseRegistration<T>()
+            where T : IDbConfiguration
+        {
+            if (SqlConnectionFactories.TryGetValue(GetDatabaseKey<T>(), out var connectionFactory))
+            {
+                return ((AzureSqlConnectionFactory)connectionFactory).SqlConnectionStringBuilder;
+            }
+
+            return null;
+        }
+
+        /// <summary>
+        /// Initializes an <see cref="ISqlConnectionFactory"/>, for use by validation jobs.
+        /// </summary>
+        /// <returns>ConnectionStringBuilder, used for diagnostics.</returns>
+        public SqlConnectionStringBuilder RegisterDatabase<T>(
+            IServiceProvider serviceProvider,
+            bool testConnection = true)
+            where T : IDbConfiguration
+        {
+            if (serviceProvider == null)
+            {
+                throw new ArgumentNullException(nameof(serviceProvider));
+            }
+
+            var secretInjector = serviceProvider.GetRequiredService<ISecretInjector>();
+            var connectionString = serviceProvider.GetRequiredService<IOptionsSnapshot<T>>().Value.ConnectionString;
+            var connectionFactory = new AzureSqlConnectionFactory(connectionString, secretInjector);
+
+            return RegisterDatabase(GetDatabaseKey<T>(), connectionString, testConnection, secretInjector);
+        }
+
+        /// <summary>
+        /// Initializes an <see cref="ISqlConnectionFactory"/>, for use by non-validation jobs.
+        /// </summary>
+        /// <returns>ConnectionStringBuilder, used for diagnostics.</returns>
+        public SqlConnectionStringBuilder RegisterDatabase(
+            IServiceContainer serviceContainer,
+            IDictionary<string, string> jobArgsDictionary,
+            string connectionStringArgName,
+            bool testConnection = true)
+        {
+            if (serviceContainer == null)
+            {
+                throw new ArgumentNullException(nameof(serviceContainer));
+            }
+
+            if (jobArgsDictionary == null)
+            {
+                throw new ArgumentNullException(nameof(jobArgsDictionary));
+            }
+
+            if (string.IsNullOrEmpty(connectionStringArgName))
+            {
+                throw new ArgumentException("Argument cannot be null or empty.", nameof(connectionStringArgName));
+            }
+
+            var secretInjector = (ISecretInjector)serviceContainer.GetService(typeof(ISecretInjector));
+            var connectionString = JobConfigurationManager.GetArgument(jobArgsDictionary, connectionStringArgName);
+
+            return RegisterDatabase(connectionStringArgName, connectionString, testConnection, secretInjector);
+        }
+
+        /// <summary>
+        /// Register a job database at initialization time. Each call should overwrite any existing
+        /// registration because <see cref="JobRunner"/> calls <see cref="Init"/> on every iteration.
+        /// </summary>
+        /// <returns>ConnectionStringBuilder, used for diagnostics.</returns>
+        private SqlConnectionStringBuilder RegisterDatabase(
+            string name,
+            string connectionString,
+            bool testConnection,
+            ISecretInjector secretInjector)
+        {
+            var connectionFactory = new AzureSqlConnectionFactory(connectionString, secretInjector, Logger);
+            SqlConnectionFactories[name] = connectionFactory;
+
+            if (testConnection)
+            {
+                Task.Run(() => TestConnection(name, connectionFactory)).Wait();
+            }
+
+            return connectionFactory.SqlConnectionStringBuilder;
+        }
+
+        private static string GetDatabaseKey<T>()
+        {
+            return typeof(T).Name;
+        }
+
+        /// <summary>
+        /// Create a SqlConnection, for use by validation jobs.
+        /// </summary>
+        public Task<SqlConnection> CreateSqlConnectionAsync<T>()
+            where T : IDbConfiguration
+        {
+            var name = GetDatabaseKey<T>();
+            if (!SqlConnectionFactories.ContainsKey(name))
+            {
+                throw new InvalidOperationException($"Database {name} has not been registered.");
+            }
+
+            return SqlConnectionFactories[name].CreateAsync();
+        }
+
+        /// <summary>
+        /// Synchronous creation of a SqlConnection, for use by validation jobs.
+        /// </summary>
+        public SqlConnection CreateSqlConnection<T>()
+            where T : IDbConfiguration
+        {
+            return Task.Run(() => CreateSqlConnectionAsync<T>()).Result;
+        }
+
+        /// <summary>
+        /// Creates and opens a SqlConnection, for use by non-validation jobs.
+        /// </summary>
+        public Task<SqlConnection> OpenSqlConnectionAsync(string connectionStringArgName)
+        {
+            if (string.IsNullOrEmpty(connectionStringArgName))
+            {
+                throw new ArgumentException("Argument cannot be null or empty.", nameof(connectionStringArgName));
+            }
+
+            if (!SqlConnectionFactories.ContainsKey(connectionStringArgName))
+            {
+                throw new InvalidOperationException($"Database {connectionStringArgName} has not been registered.");
+            }
+
+            return SqlConnectionFactories[connectionStringArgName].OpenAsync();
+        }
     }
 }
diff --git a/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj b/src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj
@@ -50,7 +50,6 @@
     <Compile Include="Configuration\ValidationDbConfiguration.cs" />
     <Compile Include="Configuration\ValidationStorageConfiguration.cs" />
     <Compile Include="Extensions\LoggerExtensions.cs" />
-    <Compile Include="Extensions\SqlConnectionStringBuilderExtensions.cs" />
     <Compile Include="Extensions\XElementExtensions.cs" />
     <Compile Include="SecretReader\ISecretReaderFactory.cs" />
     <Compile Include="SecretReader\SecretReaderFactory.cs" />
@@ -83,6 +82,9 @@
     <PackageReference Include="NuGet.Services.Logging">
       <Version>2.27.0</Version>
     </PackageReference>
+    <PackageReference Include="NuGet.Services.Sql">
+      <Version>2.27.0</Version>
+    </PackageReference>
     <PackageReference Include="System.Net.Http">
       <Version>4.3.3</Version>
     </PackageReference>

diff --git a/src/NuGet.Services.Validation.Orchestrator/NuGet.Services.Validation.Orchestrator.csproj b/src/NuGet.Services.Validation.Orchestrator/NuGet.Services.Validation.Orchestrator.csproj
@@ -127,7 +127,7 @@
       <Version>1.2.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.26.0-master-34394</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Validation.Issues">
       <Version>2.27.0-master-35351</Version>

diff --git a/src/NuGet.SupportRequests.Notifications/NuGet.SupportRequests.Notifications.csproj b/src/NuGet.SupportRequests.Notifications/NuGet.SupportRequests.Notifications.csproj
@@ -107,7 +107,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
   </ItemGroup>
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />

diff --git a/src/Search.GenerateAuxiliaryData/Search.GenerateAuxiliaryData.csproj b/src/Search.GenerateAuxiliaryData/Search.GenerateAuxiliaryData.csproj
@@ -93,7 +93,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="WindowsAzure.Storage">
       <Version>7.1.2</Version>

diff --git a/src/Stats.AggregateCdnDownloadsInGallery/Stats.AggregateCdnDownloadsInGallery.csproj b/src/Stats.AggregateCdnDownloadsInGallery/Stats.AggregateCdnDownloadsInGallery.csproj
@@ -76,7 +76,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
   </ItemGroup>
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />

diff --git a/src/Stats.CreateAzureCdnWarehouseReports/Stats.CreateAzureCdnWarehouseReports.csproj b/src/Stats.CreateAzureCdnWarehouseReports/Stats.CreateAzureCdnWarehouseReports.csproj
@@ -100,7 +100,7 @@
       <Version>2.25.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Versioning">
       <Version>4.3.0-preview1-2524</Version>

diff --git a/src/Stats.ImportAzureCdnStatistics/Stats.ImportAzureCdnStatistics.csproj b/src/Stats.ImportAzureCdnStatistics/Stats.ImportAzureCdnStatistics.csproj
@@ -126,7 +126,7 @@
       <Version>2.25.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Versioning">
       <Version>4.3.0-preview1-2524</Version>

diff --git a/src/Stats.RefreshClientDimension/Stats.RefreshClientDimension.csproj b/src/Stats.RefreshClientDimension/Stats.RefreshClientDimension.csproj
@@ -84,7 +84,7 @@
       <Version>9.0.1</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
     <PackageReference Include="UAParser">
       <Version>1.2.0</Version>

diff --git a/src/Stats.RollUpDownloadFacts/Stats.RollUpDownloadFacts.csproj b/src/Stats.RollUpDownloadFacts/Stats.RollUpDownloadFacts.csproj
@@ -85,7 +85,7 @@
       <Version>2.25.0</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
   </ItemGroup>
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />

diff --git a/src/UpdateLicenseReports/UpdateLicenseReports.csproj b/src/UpdateLicenseReports/UpdateLicenseReports.csproj
@@ -75,7 +75,7 @@
       <Version>2.0.10</Version>
     </PackageReference>
     <PackageReference Include="NuGet.Services.Sql">
-      <Version>2.25.0-master-30453</Version>
+      <Version>2.27.0</Version>
     </PackageReference>
   </ItemGroup>
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />