Add workload restore command

Adding _FindAllReferenceWorkload target to find workloads in project.
Discover the projects and then pipe all arguements to workload install
command

src/Cli/dotnet/Parser.cs

@@ -69,19 +69,18 @@ public static class Parser
         // Argument
         public static readonly Argument DotnetSubCommand = new Argument<string>() { Arity = ArgumentArity.ExactlyOne, IsHidden = true };
 
-        private static Command ConfigureCommandLine(Command rootCommand, bool includeWorkloadCommands = false)
+        private static Command ConfigureCommandLine(Command rootCommand)
         {
             // Add subcommands
             foreach (var subcommand in Subcommands)
             {
                 rootCommand.AddCommand(subcommand);
             }
 
-            // Workload command is behind a feature flag during development
-            rootCommand.AddCommand(WorkloadCommandParser.GetCommand(includeWorkloadCommands || Env.GetEnvironmentVariableAsBool("DEVENABLEWORKLOADCOMMAND", defaultValue: false)));
+            rootCommand.AddCommand(WorkloadCommandParser.GetCommand());
 
             //Add internal commands
-           rootCommand.AddCommand(InstallSuccessCommand);
+            rootCommand.AddCommand(InstallSuccessCommand);
 
             // Add options
             rootCommand.AddOption(DiagOption);
@@ -112,7 +111,7 @@ private static CommandLineBuilder DisablePosixBinding(this CommandLineBuilder bu
             .DisablePosixBinding()
             .Build();
 
-        public static System.CommandLine.Parsing.Parser GetWorkloadsInstance { get; } = new CommandLineBuilder(ConfigureCommandLine(new RootCommand(), true))
+        public static System.CommandLine.Parsing.Parser GetWorkloadsInstance { get; } = new CommandLineBuilder(ConfigureCommandLine(new RootCommand()))
             .UseExceptionHandler(ExceptionHandler)
             .UseHelp()
             .UseHelpBuilder(context => new DotnetHelpBuilder(context.Console))

src/Cli/dotnet/commands/dotnet-workload/WorkloadCommandParser.cs

@@ -8,7 +8,7 @@ namespace Microsoft.DotNet.Cli
 {
     internal static class WorkloadCommandParser
     {
-        public static Command GetCommand(bool includeAllCommands)
+        public static Command GetCommand()
         {
             var command = new Command("workload", LocalizableStrings.CommandDescription);
 
@@ -18,11 +18,7 @@ public static Command GetCommand(bool includeAllCommands)
             command.AddCommand(WorkloadSearchCommandParser.GetCommand());
             command.AddCommand(WorkloadUninstallCommandParser.GetCommand());
             command.AddCommand(WorkloadRepairCommandParser.GetCommand());
-            if (includeAllCommands)
-            {
-                command.AddCommand(WorkloadRestoreCommandParser.GetCommand());
-            }
-
+            command.AddCommand(WorkloadRestoreCommandParser.GetCommand());
             command.AddCommand(WorkloadElevateCommandParser.GetCommand());
 
             return command;

src/Cli/dotnet/commands/dotnet-workload/install/WorkloadInstallCommand.cs

@@ -57,7 +57,8 @@ public WorkloadInstallCommand(
             string dotnetDir = null,
             string userHome = null,
             string tempDirPath = null,
-            string version = null)
+            string version = null,
+            IReadOnlyCollection<string> workloadIds = null)
             : base(parseResult)
         {
             _reporter = reporter ?? Reporter.Output;
@@ -66,7 +67,7 @@ public WorkloadInstallCommand(
             _printDownloadLinkOnly = parseResult.ValueForOption<bool>(WorkloadInstallCommandParser.PrintDownloadLinkOnlyOption);
             _fromCacheOption = parseResult.ValueForOption<string>(WorkloadInstallCommandParser.FromCacheOption);
             _downloadToCacheOption = parseResult.ValueForOption<string>(WorkloadInstallCommandParser.DownloadToCacheOption);
-            _workloadIds = parseResult.ValueForArgument<IEnumerable<string>>(WorkloadInstallCommandParser.WorkloadIdArgument).ToList().AsReadOnly();
+            _workloadIds = workloadIds ?? parseResult.ValueForArgument<IEnumerable<string>>(WorkloadInstallCommandParser.WorkloadIdArgument).ToList().AsReadOnly();
             _verbosity = parseResult.ValueForOption<VerbosityOptions>(WorkloadInstallCommandParser.VerbosityOption);
             _dotnetPath = dotnetDir ?? Path.GetDirectoryName(Environment.ProcessPath);
             _sdkVersion = WorkloadOptionsExtensions.GetValidatedSdkVersion(parseResult.ValueForOption<string>(WorkloadInstallCommandParser.VersionOption), version, _dotnetPath);

src/Cli/dotnet/commands/dotnet-workload/install/WorkloadInstallCommandParser.cs

@@ -57,6 +57,13 @@ public static Command GetCommand()
             var command = new Command("install", LocalizableStrings.CommandDescription);
 
             command.AddArgument(WorkloadIdArgument);
+            AddWorkloadInstallCommandOptions(command);
+
+            return command;
+        }
+
+        internal static void AddWorkloadInstallCommandOptions(Command command)
+        {
             command.AddOption(VersionOption);
             command.AddOption(ConfigOption);
             command.AddOption(SourceOption);
@@ -68,8 +75,6 @@ public static Command GetCommand()
             command.AddOption(TempDirOption);
             command.AddWorkloadCommandNuGetRestoreActionConfigOptions();
             command.AddOption(VerbosityOption);
-
-            return command;
         }
     }
 }

src/Cli/dotnet/commands/dotnet-workload/restore/LocalizableStrings.resx

@@ -120,49 +120,13 @@
   <data name="CommandDescription" xml:space="preserve">
     <value>Restore workloads required for the project.</value>
   </data>
-  <data name="InvalidPackageWarning" xml:space="preserve">
-    <value>Warning: workload package '{0}' is invalid</value>
+  <data name="InstallingWorkloads" xml:space="preserve">
+    <value>Installing workloads: {0}</value>
   </data>
-  <data name="PackageIdColumn" xml:space="preserve">
-    <value>Package Id</value>
+  <data name="CouldNotFindAProject" xml:space="preserve">
+    <value>Couldn't find a project. Ensure a project exists in {0}, or pass the path to the project using {1}.</value>
   </data>
-  <data name="VersionColumn" xml:space="preserve">
-    <value>Version</value>
-  </data>
-  <data name="CommandsColumn" xml:space="preserve">
-    <value>Commands</value>
-  </data>
-  <data name="AddSourceOptionDescription" xml:space="preserve">
-    <value>Add an additional NuGet package source to use during installation.</value>
-  </data>
-  <data name="AddSourceOptionName" xml:space="preserve">
-    <value>SOURCE</value>
-  </data>
-  <data name="ConfigFileOptionDescription" xml:space="preserve">
-    <value>The NuGet configuration file to use.</value>
-  </data>
-  <data name="ConfigFileOptionName" xml:space="preserve">
-    <value>FILE</value>
-  </data>
-  <data name="VersionOptionDescription" xml:space="preserve">
-    <value>The version of the workload package to install.</value>
-  </data>
-  <data name="VersionOptionName" xml:space="preserve">
-    <value>VERSION</value>
-  </data>
-  <data name="PackageFailedToRestore" xml:space="preserve">
-    <value>Package "{0}" failed to restore, due to {1}</value>
-  </data>
-  <data name="RestoreSuccessful" xml:space="preserve">
-    <value>Workload '{0}' (version '{1}') was restored. Available commands: {2}</value>
-  </data>
-  <data name="RestorePartiallyFailed" xml:space="preserve">
-    <value>Restore partially failed.</value>
-  </data>
-  <data name="RestoreFailed" xml:space="preserve">
-    <value>Restore failed.</value>
-  </data>
-  <data name="NoWorkloadsWereRestored" xml:space="preserve">
-    <value>No workloads were restored.</value>
+  <data name="FailedToRunTarget" xml:space="preserve">
+    <value>Failed to restore workload for project {0}: Failed to run msbuild target _FindAllReferenceWorkload.</value>
   </data>
 </root>

src/Cli/dotnet/commands/dotnet-workload/restore/WorkloadRestoreCommand.cs

@@ -2,43 +2,128 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
 using System.CommandLine.Parsing;
+using System.IO;
+using System.Linq;
+using Microsoft.Build.Construction;
+using Microsoft.Build.Execution;
+using Microsoft.Build.Framework;
+using Microsoft.Build.Logging;
 using Microsoft.DotNet.Cli;
 using Microsoft.DotNet.Cli.Utils;
+using Microsoft.DotNet.Workloads.Workload.Install;
 using Microsoft.Extensions.EnvironmentAbstractions;
+using Microsoft.NET.Sdk.WorkloadManifestReader;
 
 namespace Microsoft.DotNet.Workloads.Workload.Restore
 {
     internal class WorkloadRestoreCommand : CommandBase
     {
-        private readonly string _configFilePath;
-        private readonly IReporter _errorReporter;
-        private readonly IFileSystem _fileSystem;
+        private readonly ParseResult _result;
         private readonly IReporter _reporter;
-        private readonly string[] _sources;
-        private readonly string _verbosity;
+        private readonly IEnumerable<string> _slnOrProjectArgument;
 
         public WorkloadRestoreCommand(
             ParseResult result,
             IFileSystem fileSystem = null,
             IReporter reporter = null)
             : base(result)
         {
-            _fileSystem = fileSystem ?? new FileSystemWrapper();
-
+            _result = result;
             _reporter = reporter ?? Reporter.Output;
-            _errorReporter = reporter ?? Reporter.Error;
-
-            _configFilePath = result.ValueForOption<string>(WorkloadRestoreCommandParser.ConfigOption);
-            _sources = result.ValueForOption<string[]>(WorkloadRestoreCommandParser.SourceOption);
-            _verbosity =
-                Enum.GetName(result.ValueForOption<VerbosityOptions>(WorkloadRestoreCommandParser.VerbosityOption));
+            _slnOrProjectArgument =
+                result.ValueForArgument<IEnumerable<string>>(RestoreCommandParser.SlnOrProjectArgument);
         }
 
         public override int Execute()
         {
-            _reporter.WriteLine("WIP workload restore stub");
+            var allProjects = DiscoverAllProjects(Directory.GetCurrentDirectory(), _slnOrProjectArgument).Distinct();
+
+            List<WorkloadId> allWorkloadId = RunTargetToGetWorkloadIds(allProjects);
+
+            _reporter.WriteLine(string.Format(LocalizableStrings.InstallingWorkloads, string.Join(" ", allWorkloadId)));
+
+            var workloadInstallCommand = new WorkloadInstallCommand(_result,
+                workloadIds: new ReadOnlyCollection<string>(allWorkloadId.Select(a => a.ToString()).ToList()));
+
+            workloadInstallCommand.Execute();
             return 0;
         }
+
+        private static List<WorkloadId> RunTargetToGetWorkloadIds(IEnumerable<string> allProjects)
+        {
+            var globalProperties = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase)
+            {
+                {Constants.MSBuildExtensionsPath, AppContext.BaseDirectory},
+                {"SkipResolvePackageAssets", "true"}
+            };
+
+            var allWorkloadId = new List<WorkloadId>();
+            foreach (string projectFile in allProjects)
+            {
+                var project = new ProjectInstance(projectFile, globalProperties, null);
+
+                bool buildResult = project.Build(new[] {"_FindAllReferenceWorkload"},
+                    loggers: Enumerable.Empty<ILogger>(),
+                    remoteLoggers: Enumerable.Empty<ForwardingLoggerRecord>(),
+                    targetOutputs: out var targetOutputs);
+
+                if (buildResult == false)
+                {
+                    throw new GracefulException(
+                        string.Format(
+                            LocalizableStrings.FailedToRunTarget,
+                            projectFile));
+                }
+
+                var targetResult = targetOutputs["_FindAllReferenceWorkload"];
+                allWorkloadId.AddRange(targetResult.Items.Select(item => new WorkloadId(item.ItemSpec)));
+            }
+
+            allWorkloadId = allWorkloadId.Distinct().ToList();
+            return allWorkloadId;
+        }
+
+        internal static List<string> DiscoverAllProjects(string currentDirectory,
+            IEnumerable<string> slnOrProjectArgument = null)
+        {
+            var slnFiles = new List<string>();
+            var projectFiles = new List<string>();
+            if (slnOrProjectArgument == null || !slnOrProjectArgument.Any())
+            {
+                slnFiles = Directory.GetFiles(currentDirectory, "*.sln").ToList();
+                projectFiles.AddRange(Directory.GetFiles(currentDirectory, "*.*proj"));
+            }
+            else
+            {
+                slnFiles = slnOrProjectArgument
+                    .Where(s => Path.GetExtension(s).Equals(".sln", StringComparison.OrdinalIgnoreCase))
+                    .Select(Path.GetFullPath).ToList();
+                projectFiles = slnOrProjectArgument
+                    .Where(s => Path.GetExtension(s).EndsWith("proj", StringComparison.OrdinalIgnoreCase))
+                    .Select(Path.GetFullPath).ToList();
+            }
+
+            foreach (string file in slnFiles)
+            {
+                var solutionFile = SolutionFile.Parse(file);
+                var projects = solutionFile.ProjectsInOrder;
+                foreach (var p in projects)
+                {
+                    projectFiles.Add(p.AbsolutePath);
+                }
+            }
+
+            if (projectFiles.Count == 0)
+            {
+                throw new GracefulException(
+                    LocalizableStrings.CouldNotFindAProject,
+                    currentDirectory, "--project");
+            }
+
+            return projectFiles;
+        }
     }
 }

src/Cli/dotnet/commands/dotnet-workload/restore/WorkloadRestoreCommandParser.cs

@@ -8,21 +8,12 @@ namespace Microsoft.DotNet.Cli
 {
     internal static class WorkloadRestoreCommandParser
     {
-        public static readonly Option ConfigOption = WorkloadInstallCommandParser.ConfigOption;
-
-        public static readonly Option SourceOption = WorkloadInstallCommandParser.SourceOption;
-
-        public static readonly Option VerbosityOption = WorkloadInstallCommandParser.VerbosityOption;
-
         public static Command GetCommand()
         {
             Command command = new Command("restore", LocalizableStrings.CommandDescription);
 
-            command.AddOption(ConfigOption);
-            command.AddOption(SourceOption);
-            command.AddWorkloadCommandNuGetRestoreActionConfigOptions();
-            command.AddOption(VerbosityOption);
-
+            command.AddArgument(RestoreCommandParser.SlnOrProjectArgument);
+            WorkloadInstallCommandParser.AddWorkloadInstallCommandOptions(command);
             return command;
         }
     }