1+ // ----------------------------------------------------------------------------------
2+ //
3+ // Copyright Microsoft Corporation
4+ // Licensed under the Apache License, Version 2.0 (the "License");
5+ // you may not use this file except in compliance with the License.
6+ // You may obtain a copy of the License at
7+ // http://www.apache.org/licenses/LICENSE-2.0
8+ // Unless required by applicable law or agreed to in writing, software
9+ // distributed under the License is distributed on an "AS IS" BASIS,
10+ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+ // See the License for the specific language governing permissions and
12+ // limitations under the License.
13+ // ----------------------------------------------------------------------------------
14+
15+ namespace Microsoft . Azure . Commands . ResourceManager . Common . ArgumentCompleters
16+ {
17+ using Commands . Common . Authentication ;
18+ using Commands . Common . Authentication . Abstractions ;
19+ using Management . Internal . Resources ;
20+ using Management . Internal . Resources . Models ;
21+ using Properties ;
22+ using Rest . Azure ;
23+ using System ;
24+ using System . Collections . Concurrent ;
25+ using System . Collections . Generic ;
26+ using System . Linq ;
27+ using System . Management . Automation ;
28+
29+ /// <summary>
30+ /// This attribute will allow the user to autocomplete the -ResourceGroup parameter of a cmdlet with valid resource groups
31+ /// </summary>
32+ public class ResourceGroupCompleterAttribute : ArgumentCompleterAttribute
33+ {
34+ private static IDictionary < int , IList < String > > _resourceGroupNamesDictionary = new ConcurrentDictionary < int , IList < string > > ( ) ;
35+ private static readonly object _lock = new object ( ) ;
36+
37+ protected static IList < String > ResourceGroupNames
38+ {
39+ get
40+ {
41+ lock ( _lock )
42+ {
43+ IAzureContext context = AzureRmProfileProvider . Instance . Profile . DefaultContext ;
44+ var contextHash = HashContext ( context ) ;
45+ if ( ! _resourceGroupNamesDictionary . ContainsKey ( contextHash ) )
46+ {
47+ var tempResourceGroupList = new List < string > ( ) ;
48+ try
49+ {
50+ var instance = AzureSession . Instance ;
51+ var client = instance . ClientFactory . CreateCustomArmClient < ResourceManagementClient > (
52+ context . Environment . GetEndpointAsUri ( AzureEnvironment . Endpoint . ResourceManager ) ,
53+ instance . AuthenticationFactory . GetServiceClientCredentials ( context , AzureEnvironment . Endpoint . ResourceManager ) ,
54+ instance . ClientFactory . GetCustomHandlers ( ) ) ;
55+ client . SubscriptionId = context . Subscription . Id ;
56+ // Retrieve only the first page of ResourceGroups to display to the user
57+ var resourceGroups = client . ResourceGroups . ListAsync ( ) ;
58+ if ( resourceGroups . Wait ( TimeSpan . FromSeconds ( 5 ) ) )
59+ {
60+ tempResourceGroupList = CreateResourceGroupList ( resourceGroups . Result ) ;
61+ if ( resourceGroups . Result != null )
62+ {
63+ _resourceGroupNamesDictionary [ contextHash ] = tempResourceGroupList ;
64+ }
65+ }
66+ #if DEBUG
67+ else
68+ {
69+ throw new Exception ( "client.ResourceGroups call timed out" ) ;
70+ }
71+ #endif
72+ }
73+
74+ catch ( Exception ex )
75+ {
76+ #if DEBUG
77+ throw ex ;
78+ #endif
79+ }
80+
81+ return tempResourceGroupList ;
82+ }
83+
84+ else
85+ {
86+ return _resourceGroupNamesDictionary [ contextHash ] ;
87+ }
88+ }
89+ }
90+ }
91+
92+ /// <summary>
93+ /// This class will provide a list of resource groups that are available to the user (with default resource group first if it exists). This will then be available to the user to tab through.
94+ /// Example: [Parameter(ParameterSetName = ListByNameInTenantParameterSet, ValueFromPipelineByPropertyName = true, Mandatory = false), ResourceGroupCompleter()]
95+ /// </summary>
96+ /// <param name="resourceTypes"></param>
97+ public ResourceGroupCompleterAttribute ( ) : base ( CreateScriptBlock ( ) )
98+ {
99+ }
100+
101+ public static string [ ] GetResourceGroups ( )
102+ {
103+ IAzureContext context = AzureRmProfileProvider . Instance . Profile . DefaultContext ;
104+ var resourceGroupNamesCopy = ResourceGroupNames ;
105+ if ( context . IsPropertySet ( Resources . DefaultResourceGroupKey ) )
106+ {
107+ return GetResourceGroups ( resourceGroupNamesCopy , context . ExtendedProperties [ Resources . DefaultResourceGroupKey ] ) ;
108+ }
109+ return GetResourceGroups ( resourceGroupNamesCopy , null ) ;
110+ }
111+
112+ public static string [ ] GetResourceGroups ( IList < string > resourceGroupNames , string defaultResourceGroup )
113+ {
114+ if ( defaultResourceGroup != null )
115+ {
116+ if ( resourceGroupNames . Contains ( defaultResourceGroup ) )
117+ {
118+ resourceGroupNames . Remove ( defaultResourceGroup ) ;
119+ }
120+ resourceGroupNames . Insert ( 0 , defaultResourceGroup ) ;
121+ }
122+ return resourceGroupNames . ToArray ( ) ;
123+ }
124+
125+ private static ScriptBlock CreateScriptBlock ( )
126+ {
127+ string script = "param($commandName, $parameterName, $wordToComplete, $commandAst, $fakeBoundParameter)\n " +
128+ "$locations = [Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters.ResourceGroupCompleterAttribute]::GetResourceGroups()\n " +
129+ "$locations | Where-Object { $_ -Like \" $wordToComplete*\" } | ForEach-Object { [System.Management.Automation.CompletionResult]::new($_, $_, 'ParameterValue', $_) }" ;
130+ ScriptBlock scriptBlock = ScriptBlock . Create ( script ) ;
131+ return scriptBlock ;
132+ }
133+
134+ private static int HashContext ( IAzureContext context )
135+ {
136+ return ( context . Account . Id + context . Environment . Name + context . Subscription . Id + context . Tenant . Id ) . GetHashCode ( ) ;
137+ }
138+
139+ public static List < string > CreateResourceGroupList ( IPage < ResourceGroup > result )
140+ {
141+ var tempResourceGroupList = new List < string > ( ) ;
142+ if ( result != null )
143+ {
144+ foreach ( ResourceGroup resourceGroup in result )
145+ {
146+ tempResourceGroupList . Add ( resourceGroup . Name ) ;
147+ }
148+ }
149+ #if DEBUG
150+ else
151+ {
152+ throw new Exception ( "Result from client.ResourceGroups is null" ) ;
153+ }
154+ #endif
155+ return tempResourceGroupList ;
156+ }
157+ }
158+ }
0 commit comments