Use new TFM logic for package ingestion in backfill command (#8431)

New TFM logic for package ingestion and backfill command

Author: drewgillies

src/GalleryTools/Commands/BackfillCommand.cs

@@ -5,21 +5,25 @@
 using CsvHelper;
 using CsvHelper.Configuration;
 using CsvHelper.TypeConversion;
+using Knapcode.MiniZip;
 using Microsoft.Extensions.CommandLineUtils;
 using NuGet.Packaging;
 using NuGet.Services.Entities;
 using NuGetGallery;
-using NuGetGallery.Configuration;
 using System;
+using System.Collections.Generic;
 using System.Data.Entity;
+using System.Data.SqlClient;
 using System.IO;
 using System.Linq;
+using System.Linq.Expressions;
 using System.Net.Http;
 using System.Text;
 using System.Threading.Tasks;
 using System.Xml;
 using System.Xml.Linq;
 using GalleryTools.Utils;
+using NuGet.Services.Sql;
 
 namespace GalleryTools.Commands
 {
@@ -35,6 +39,15 @@ public abstract class BackfillCommand<TMetadata>
 
         protected virtual int UpdateBatchSize => 100;
 
+        protected virtual int LimitTo => 0;
+
+        protected virtual MetadataSourceType SourceType => MetadataSourceType.NuspecOnly;
+
+        protected virtual Expression<Func<Package, object>> QueryIncludes => null;
+
+
+        protected IPackageService _packageService;
+
         public static void Configure<TCommand>(CommandLineApplication config) where TCommand : BackfillCommand<TMetadata>, new()
         {
             config.Description = "Backfill metadata for packages in the gallery";
@@ -53,10 +66,12 @@ public abstract class BackfillCommand<TMetadata>
                 builder.RegisterAssemblyModules(typeof(DefaultDependenciesModule).Assembly);
                 var container = builder.Build();
 
-                var connectionString = container.Resolve<IAppConfiguration>().SqlConnectionString;
+                var sqlConnectionFactory = container.Resolve<ISqlConnectionFactory>();
+                var sqlConnection = await sqlConnectionFactory.CreateAsync();
                 var serviceDiscoveryUriValue = new Uri(serviceDiscoveryUri.Value());
 
                 var command = new TCommand();
+                command._packageService = container.Resolve<IPackageService>();
 
                 var metadataFileName = fileName.HasValue() ? fileName.Value() : command.MetadataFileName;
 
@@ -75,54 +90,90 @@ public abstract class BackfillCommand<TMetadata>
                         }
                     }
 
-                    await command.Collect(connectionString, serviceDiscoveryUriValue, lastCreateTime, metadataFileName);
+                    await command.Collect(sqlConnection, serviceDiscoveryUriValue, lastCreateTime, metadataFileName);
                 }
 
                 if (updateDB.HasValue())
                 {
-                    await command.Update(connectionString, metadataFileName);
+                    await command.Update(sqlConnection, metadataFileName);
                 }
 
                 return 0;
             });
         }
 
-        public async Task Collect(string connectionString, Uri serviceDiscoveryUri, DateTime? lastCreateTime, string fileName)
+        public async Task Collect(SqlConnection connection, Uri serviceDiscoveryUri, DateTime? lastCreateTime, string fileName)
         {
-            using (var context = new EntitiesContext(connectionString, readOnly: true))
+            using (var context = new EntitiesContext(connection, readOnly: true))
             using (var cursor = new FileCursor(CursorFileName))
             using (var logger = new Logger(ErrorsFileName))
             {
+                context.SetCommandTimeout(300); // large query
+
                 var startTime = await cursor.Read();
 
                 logger.Log($"Starting metadata collection - Cursor time: {startTime:u}");
 
                 var repository = new EntityRepository<Package>(context);
 
                 var packages = repository.GetAll()
-                    .Include(p => p.PackageRegistration)
+                    .Include(p => p.PackageRegistration);
+                if (QueryIncludes != null)
+                {
+                    packages = packages.Include(QueryIncludes);
+                }
+                
+                packages = packages
                     .Where(p => p.Created < lastCreateTime && p.Created > startTime)
-                    .Where(p => p.PackageStatusKey == PackageStatus.Available || p.PackageStatusKey == PackageStatus.Validating)
-                    .OrderBy(p => p.Created);
+                    .OrderBy(p => p.PackageRegistration.Id);
+                if (LimitTo > 0)
+                {
+                    packages = packages.Take(LimitTo);
+                }
 
                 var flatContainerUri = await GetFlatContainerUri(serviceDiscoveryUri);
 
                 using (var csv = CreateCsvWriter(fileName))
                 using (var http = new HttpClient())
                 {
+                    // We want these downloads ignored by stats pipelines - this user agent is automatically skipped.
+                    // See https://github.com/NuGet/NuGet.Jobs/blob/262da48ed05d0366613bbf1c54f47879aad96dcd/src/Stats.ImportAzureCdnStatistics/StatisticsParser.cs#L41
+                    http.DefaultRequestHeaders.TryAddWithoutValidation("User-Agent", 
+                        "Mozilla/5.0 (compatible; MSIE 9.0; Windows NT 6.1; Trident/5.0; AppInsights)   Backfill Job: NuGet.Gallery GalleryTools");
+
                     var counter = 0;
                     var lastCreatedDate = default(DateTime?);
 
                     foreach (var package in packages)
                     {
                         var id = package.PackageRegistration.Id;
                         var version = package.NormalizedVersion;
-
-                        var nuspecUri = $"{flatContainerUri}/{id.ToLowerInvariant()}/{version.ToLowerInvariant()}/{id.ToLowerInvariant()}.nuspec";
+                        var idLowered = id.ToLowerInvariant();
+                        var versionLowered = version.ToLowerInvariant();
 
                         try
                         {
-                            var metadata = await FetchMetadata(http, nuspecUri);
+                            var metadata = default(TMetadata);
+
+                            var nuspecUri =
+                                $"{flatContainerUri}/{idLowered}/{versionLowered}/{idLowered}.nuspec";
+                            using (var nuspecStream = await http.GetStreamAsync(nuspecUri))
+                            {
+                                var document = LoadDocument(nuspecStream);
+
+                                var nuspecReader = new NuspecReader(document);
+
+                                if (SourceType == MetadataSourceType.NuspecOnly)
+                                {
+                                    metadata = ReadMetadata(nuspecReader);
+                                }
+                                else if (SourceType == MetadataSourceType.Nupkg)
+                                {
+                                    var nupkgUri =
+                                        $"{flatContainerUri}/{idLowered}/{versionLowered}/{idLowered}.{versionLowered}.nupkg";
+                                    metadata = await FetchMetadataAsync(http, nupkgUri, nuspecReader, id, version, logger);
+                                }
+                            }
 
                             if (ShouldWriteMetadata(metadata))
                             {
@@ -132,7 +183,7 @@ public async Task Collect(string connectionString, Uri serviceDiscoveryUri, Date
 
                                 await csv.NextRecordAsync();
 
-                                logger.LogPackage(id, version, "Metadata saved.");
+                                logger.LogPackage(id, version, $"Metadata saved");
                             }
                         }
                         catch (Exception e)
@@ -163,14 +214,14 @@ public async Task Collect(string connectionString, Uri serviceDiscoveryUri, Date
             }
         }
 
-        public async Task Update(string connectionString, string fileName)
+        public async Task Update(SqlConnection connection, string fileName)
         {
             if (!File.Exists(fileName))
             {
                 throw new ArgumentException($"File '{fileName}' doesn't exist");
             }
 
-            using (var context = new EntitiesContext(connectionString, readOnly: false))
+            using (var context = new EntitiesContext(connection, readOnly: false))
             using (var cursor = new FileCursor(CursorFileName))
             using (var logger = new Logger(ErrorsFileName))
             {
@@ -199,9 +250,7 @@ public async Task Update(string connectionString, string fileName)
 
                             if (package != null)
                             {
-
-                                UpdatePackage(package, metadata.Metadata);
-
+                                UpdatePackage(package, metadata.Metadata, context);
                                 logger.LogPackage(metadata.Id, metadata.Version, "Metadata updated.");
 
                                 counter++;
@@ -234,13 +283,15 @@ public async Task Update(string connectionString, string fileName)
             }
         }
 
-        protected abstract TMetadata ReadMetadata(NuspecReader reader);
+        protected virtual TMetadata ReadMetadata(NuspecReader reader) => default;
+
+        protected virtual TMetadata ReadMetadata(IList<string> files, NuspecReader nuspecReader) => default;
 
         protected abstract bool ShouldWriteMetadata(TMetadata metadata);
 
         protected abstract void ConfigureClassMap(PackageMetadataClassMap map);
 
-        protected abstract void UpdatePackage(Package package, TMetadata metadata);
+        protected abstract void UpdatePackage(Package package, TMetadata metadata, EntitiesContext context);
 
         private static async Task<string> GetFlatContainerUri(Uri serviceDiscoveryUri)
         {
@@ -251,15 +302,26 @@ private static async Task<string> GetFlatContainerUri(Uri serviceDiscoveryUri)
             return result.First().AbsoluteUri.TrimEnd('/');
         }
 
-        private async Task<TMetadata> FetchMetadata(HttpClient httpClient, string nuspecUri)
+        private async Task<TMetadata> FetchMetadataAsync(
+            HttpClient httpClient, string nupkgUri, NuspecReader nuspecReader, string id, string version, Logger logger)
         {
-            using (var nuspecStream = await httpClient.GetStreamAsync(nuspecUri))
-            {
-                var document = LoadDocument(nuspecStream);
-
-                var reader = new NuspecReader(document);
+            var httpZipProvider = new HttpZipProvider(httpClient);
 
-                return ReadMetadata(reader);
+            try
+            {
+                var zipDirectoryReader = await httpZipProvider.GetReaderAsync(new Uri(nupkgUri));
+                var zipDirectory = await zipDirectoryReader.ReadAsync();
+                var files = zipDirectory
+                    .Entries
+                    .Select(x => x.GetName())
+                    .ToList();
+
+                return ReadMetadata(files, nuspecReader);
+            }
+            catch (Exception e)
+            {
+                await logger.LogPackageError(id, version, e);
+                return default;
             }
         }
 
@@ -460,5 +522,21 @@ public void Dispose()
                 Writer.Dispose();
             }
         }
+
+        /// <summary>
+        /// This enum allows our logic to respond to a package's need for only a nupsec to determine metadata, or whether
+        /// it needs access to the .nupkg for analysis of the package
+        /// </summary>
+        public enum MetadataSourceType
+        {
+            /// <summary>
+            /// Just the nuspec will suffice for metadata extraction
+            /// </summary>
+            NuspecOnly,
+            /// <summary>
+            /// We need to dig deeper into the bupkg for the metadata
+            /// </summary>
+            Nupkg
+        }
     }
 }

src/GalleryTools/Commands/BackfillDevelopmentDependencyMetadataCommand.cs

@@ -4,6 +4,7 @@
 using Microsoft.Extensions.CommandLineUtils;
 using NuGet.Packaging;
 using NuGet.Services.Entities;
+using NuGetGallery;
 
 namespace GalleryTools.Commands
 {
@@ -32,7 +33,7 @@ protected override void ConfigureClassMap(PackageMetadataClassMap map)
             map.Map(x => x.Metadata).Index(3);
         }
 
-        protected override void UpdatePackage(Package package, bool metadata)
+        protected override void UpdatePackage(Package package, bool metadata, EntitiesContext context)
         {
             package.DevelopmentDependency = metadata;
         }

src/GalleryTools/Commands/BackfillRepositoryMetadataCommand.cs

@@ -5,6 +5,7 @@
 using NuGet.Packaging;
 using NuGet.Packaging.Core;
 using NuGet.Services.Entities;
+using NuGetGallery;
 
 namespace GalleryTools.Commands
 {
@@ -38,7 +39,7 @@ protected override void ConfigureClassMap(PackageMetadataClassMap map)
             map.Map(x => x.Metadata.Commit).Index(6);
         }
 
-        protected override void UpdatePackage(Package package, RepositoryMetadata metadata)
+        protected override void UpdatePackage(Package package, RepositoryMetadata metadata, EntitiesContext context)
         {
             package.RepositoryUrl = metadata.Url;
 

src/GalleryTools/Commands/BackfillTfmMetadataCommand.cs

@@ -0,0 +1,75 @@
+// 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.Linq;
+using System.Linq.Expressions;
+using Microsoft.Extensions.CommandLineUtils;
+using NuGet.Packaging;
+using NuGet.Services.Entities;
+using NuGetGallery;
+
+namespace GalleryTools.Commands
+{
+    public sealed class BackfillTfmMetadataCommand : BackfillCommand<List<string>>
+    {
+        protected override string MetadataFileName => "tfmMetadata.txt";
+        
+        protected override MetadataSourceType SourceType => MetadataSourceType.Nupkg;
+        
+        protected override Expression<Func<Package, object>> QueryIncludes => p => p.SupportedFrameworks;
+
+        public static void Configure(CommandLineApplication config)
+        {
+            Configure<BackfillTfmMetadataCommand>(config);
+        }
+
+        protected override List<string> ReadMetadata(IList<string> files, NuspecReader nuspecReader)
+        {
+            var supportedTFMs = new List<string>();
+            var supportedFrameworks = _packageService.GetSupportedFrameworks(nuspecReader, files);
+            foreach (var tfm in supportedFrameworks)
+            {
+                supportedTFMs.Add(tfm.GetShortFolderName());
+            }
+
+            return supportedTFMs;
+        }
+
+        protected override bool ShouldWriteMetadata(List<string> metadata) => true;
+
+        protected override void ConfigureClassMap(PackageMetadataClassMap map)
+        {
+            map.Map(x => x.Metadata).Index(3);
+        }
+
+        protected override void UpdatePackage(Package package, List<string> metadata, EntitiesContext context)
+        {
+            var existingTFMs = package.SupportedFrameworks == null
+                ? Enumerable.Empty<string>()
+                : package.SupportedFrameworks.Select(f => f.FrameworkName.GetShortFolderName()).OrderBy(f => f);
+
+            var newTFMs = metadata == null || metadata.Count == 0
+                ? Enumerable.Empty<string>() 
+                : metadata.OrderBy(f => f);
+
+            if (Enumerable.SequenceEqual(existingTFMs, newTFMs))
+            {
+                return; // nothing to change
+            }
+
+            // clean out the old (which will be left unattached in table otherwise) before adding new
+            if (package.SupportedFrameworks != null)
+            {
+                foreach (var supportedFramework in package.SupportedFrameworks.ToList())
+                {
+                    package.SupportedFrameworks.Remove(supportedFramework);
+                    context.PackageFrameworks.Remove(supportedFramework);
+                }
+            }
+
+            package.SupportedFrameworks = newTFMs.Select(f => new PackageFramework {Package = package, TargetFramework = f}).ToList();
+        }
+    }
+}