matusthemostlygreat - added accessors for providers and quotas

Azure · Jan 27, 2021 · 4abba19 · 4abba19
1 parent 91e3dc4
commit 4abba19
Showing 1 changed file with 42 additions and 7 deletions.
diff --git a/sdk/quantum/Azure.Quantum.Jobs/src/QuantumJobClient.cs b/sdk/quantum/Azure.Quantum.Jobs/src/QuantumJobClient.cs
@@ -38,14 +38,13 @@ public QuantumJobClient(
             }
 
             var authPolicy = new BearerTokenAuthenticationPolicy(credential, "https://quantum.microsoft.com");
+            var diagnostics = new ClientDiagnostics(options);
+            var pipeline = HttpPipelineBuilder.Build(options, authPolicy);
+            var endpoint = new Uri($"https://{location}.quantum.azure.com");
 
-            _jobs = new JobsRestClient(
-                new ClientDiagnostics(options),
-                HttpPipelineBuilder.Build(options, authPolicy),
-                subscriptionId,
-                resourceGroupName,
-                workspaceName,
-                new Uri($"https://{location}.quantum.azure.com"));
+            _jobs = new JobsRestClient(diagnostics, pipeline, subscriptionId, resourceGroupName, workspaceName, endpoint);
+            _providers = new ProvidersRestClient(diagnostics, pipeline, subscriptionId, resourceGroupName, workspaceName, endpoint);
+            _quotas = new QuotasRestClient(diagnostics, pipeline, subscriptionId, resourceGroupName, workspaceName, endpoint);
         }
 
         /// <summary> Get job by id. </summary>
@@ -116,6 +115,34 @@ public virtual async Task<Response> CancelAsync(string jobId, CancellationToken
             return await _jobs.CancelAsync(jobId, cancellationToken).ConfigureAwait(false);
         }
 
+        /// <summary> Get provider status. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        public virtual Pageable<ProviderStatus> GetProviderStatus(CancellationToken cancellationToken = default)
+        {
+            return PageResponseEnumerator.CreateEnumerable(cont => ToPage(string.IsNullOrEmpty(cont) ? _providers.GetStatus() : _providers.GetStatusNextPage(cont)));
+        }
+
+        /// <summary> Get provider status. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        public virtual AsyncPageable<ProviderStatus> GetProviderStatusAsync(CancellationToken cancellationToken = default)
+        {
+            return PageResponseEnumerator.CreateAsyncEnumerable(async cont => ToPage(string.IsNullOrEmpty(cont) ? await _providers.GetStatusAsync().ConfigureAwait(false) : await _providers.GetStatusNextPageAsync(cont).ConfigureAwait(false)));
+        }
+
+        /// <summary> Get quota status. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        public virtual Pageable<QuantumJobQuota> GetQuantumJobQuotas(CancellationToken cancellationToken = default)
+        {
+            return PageResponseEnumerator.CreateEnumerable(cont => ToPage(string.IsNullOrEmpty(cont) ? _quotas.List() : _quotas.ListNextPage(cont)));
+        }
+
+        /// <summary> Get quota status. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        public virtual AsyncPageable<QuantumJobQuota> GetQuantumJobQuotasAsync(CancellationToken cancellationToken = default)
+        {
+            return PageResponseEnumerator.CreateAsyncEnumerable(async cont => ToPage(string.IsNullOrEmpty(cont) ? await _quotas.ListAsync().ConfigureAwait(false) : await _quotas.ListNextPageAsync(cont).ConfigureAwait(false)));
+        }
+
         /// <summary> Initializes a new instance of QuantumJobClient for mocking. </summary>
         protected QuantumJobClient()
         {
@@ -124,6 +151,14 @@ protected QuantumJobClient()
         private static Page<JobDetails> ToPage(Response<JobDetailsList> list) =>
             Page.FromValues(list.Value.Values, list.Value.NextLink, list.GetRawResponse());
 
+        private static Page<ProviderStatus> ToPage(Response<ProviderStatusList> list) =>
+            Page.FromValues(list.Value.Values, list.Value.NextLink, list.GetRawResponse());
+
+        private static Page<QuantumJobQuota> ToPage(Response<QuantumJobQuotaList> list) =>
+            Page.FromValues(list.Value.Values, list.Value.NextLink, list.GetRawResponse());
+
         private JobsRestClient _jobs;
+        private ProvidersRestClient _providers;
+        private QuotasRestClient _quotas;
     }
 }