Add pageable support for RegisteredModels/List and ModelVersion/List (#…

…212)

Signed-off-by: Jason Wang <[email protected]>

csharp/Microsoft.Azure.Databricks.Client.Sample/SampleProgram.UnityCatalog.cs

@@ -75,7 +75,7 @@ private static async Task TestUnityCatalogApi(DatabricksClient client)
         Console.WriteLine("Catalog deleted");
 
         Console.WriteLine("Listing Regsitered Models");
-        var registeredModels = await client.UnityCatalog.RegisteredModels.ListRegisteredModels();
+        var (registeredModels, _) = await client.UnityCatalog.RegisteredModels.List();
         foreach (var model in registeredModels)
         {
             Console.WriteLine($"\t{model.MetastoreId}, {model.Name}");
@@ -84,7 +84,7 @@ private static async Task TestUnityCatalogApi(DatabricksClient client)
 
         var fullName = "main.default.revenue_forecasting";
         Console.WriteLine("Listing Regsitered Models");
-        var modelVersionsList = await client.UnityCatalog.ModelVersion.ListModelVersions(fullName);
+        var (modelVersionsList, _) = await client.UnityCatalog.ModelVersion.List(fullName);
         foreach (var model in modelVersionsList)
         {
             Console.WriteLine($"\t{model.MetastoreId}, {model.CatalogName}");

csharp/Microsoft.Azure.Databricks.Client.Test/UnityCatalog/ModelVersionApiClientTests.cs

@@ -11,10 +11,10 @@ namespace Microsoft.Azure.Databricks.Client.Test.UnityCatalog;
 public class ModelVersionApiClientTests : UnityCatalogApiClientTest
 {
     [TestMethod]
-    public async Task ListModelVersionsTest()
+    public async Task ListTest()
     {
         var full_name = "main.default.revenue_forecasting";
-        var requestUri = $"{BaseApiUri}models/{full_name}/versions";
+        var requestUri = $"{BaseApiUri}models/{full_name}/versions?";
 
         var expectedResponse = @"
         {
@@ -68,16 +68,17 @@ public async Task ListModelVersionsTest()
         var mockClient = handler.CreateClient();
         mockClient.BaseAddress = ApiClientTest.BaseApiUri;
 
-        using var client = new ModelVersionApiClient(mockClient);
-        var response = await client.ListModelVersions(full_name);
+        using var client = new ModelVersionApiClient(mockClient);
+
+        var (actual, token) = await client.List(full_name);
 
-        var responseJson = JsonSerializer.Serialize(response, Options);
-        CollectionAssert.AreEqual(expected?.ToList(), response?.ToList());
+        CollectionAssert.AreEqual(expected!.ToArray(), actual.ToArray());
+        Assert.AreEqual("some-page-token", token);
     }
 
 
     [TestMethod]
-    public async Task GetModelVersionTest()
+    public async Task GetTest()
     {
         var expectedResponse = @"
         {
@@ -113,14 +114,14 @@ public async Task GetModelVersionTest()
         mockClient.BaseAddress = ApiClientTest.BaseApiUri;
 
         using var client = new ModelVersionApiClient(mockClient);
-        var response = await client.GetModelVersion(full_name, version);
+        var response = await client.Get(full_name, version);
 
         var responseJson = JsonSerializer.Serialize(response, Options);
         AssertJsonDeepEquals(expectedResponse, responseJson);
     }
 
     [TestMethod]
-    public async Task GetModelVersionByAliasTest()
+    public async Task GetByAliasTest()
     {
         var full_name = "main.default.revenue_forecasting_model";
         var alias = "champion";
@@ -165,7 +166,7 @@ public async Task GetModelVersionByAliasTest()
 
         using var client = new ModelVersionApiClient(mockClient);
 
-        var response = await client.GetModelVersionByAlias(full_name, alias);
+        var response = await client.GetByAlias(full_name, alias);
         var responseJson = JsonSerializer.Serialize(response, Options);
         AssertJsonDeepEquals(expectedResponse, responseJson);
     }

csharp/Microsoft.Azure.Databricks.Client.Test/UnityCatalog/RegisteredModelsApiClientTests.cs

@@ -13,9 +13,9 @@ namespace Microsoft.Azure.Databricks.Client.Test.UnityCatalog;
 public class RegisteredModelsApiClientTests : UnityCatalogApiClientTest
 {
     [TestMethod]
-    public async Task ListRegisteredModelsTest()
+    public async Task ListTest()
     {
-        var requestUri = $"{BaseApiUri}models";
+        var requestUri = $"{BaseApiUri}models?";
         var expectedResponse = @"
         {
             ""registered_models"": [
@@ -57,6 +57,7 @@ public async Task ListRegisteredModelsTest()
             ""next_page_token"": ""some-page-token""
             }";
 
+        var expected = JsonNode.Parse(expectedResponse)?["registered_models"].Deserialize<IEnumerable<RegisteredModel>>(Options);
 
         var handler = CreateMockHandler();
         handler
@@ -67,16 +68,14 @@ public async Task ListRegisteredModelsTest()
         mockClient.BaseAddress = ApiClientTest.BaseApiUri;
 
         using var client = new RegisteredModelsApiClient(mockClient);
-        var response = await client.ListRegisteredModels();
-
-        var responseJson = JsonSerializer.Serialize(response, Options);
-        var expected = JsonNode.Parse(expectedResponse)?["registered_models"].Deserialize<IEnumerable<RegisteredModel>>(Options);
-        CollectionAssert.AreEqual(expected?.ToList(), response?.ToList());
+        var (actual, token) = await client.List();
+        CollectionAssert.AreEqual(expected!.ToArray(), actual.ToArray());
+        Assert.AreEqual("some-page-token", token);
     }
 
 
     [TestMethod]
-    public void GetRegisteredModelTest()
+    public async Task GetTest()
     {
         var expectedResponse = @"
         {
@@ -110,14 +109,14 @@ public void GetRegisteredModelTest()
         mockClient.BaseAddress = ApiClientTest.BaseApiUri;
 
         using var client = new RegisteredModelsApiClient(mockClient);
-        var response = client.GetRegisteredModel(full_name).Result;
+        var response = await client.Get(full_name);
 
         var responseJson = JsonSerializer.Serialize(response, Options);
         AssertJsonDeepEquals(expectedResponse, responseJson);
     }
 
     [TestMethod]
-    public async Task SetRegisteredModelAliasTest()
+    public async Task SetTest()
     {
         var full_name = "main.default.revenue_forecasting_model";
         var alias = "champion";
@@ -126,15 +125,15 @@ public async Task SetRegisteredModelAliasTest()
 
 
         var expectedRequest = @"
-    {
-        ""version_num"": 2
-    }";
+        {
+            ""version_num"": 2
+        }";
 
         var expectedResponse = @"
-    {
-        ""alias_name"": ""champion"",
-        ""version_num"": 2
-    }";
+        {
+            ""alias_name"": ""champion"",
+            ""version_num"": 2
+        }";
 
         var handler = CreateMockHandler();
         handler
@@ -146,7 +145,7 @@ public async Task SetRegisteredModelAliasTest()
         mockClient.BaseAddress = ApiClientTest.BaseApiUri;
 
         using var client = new RegisteredModelsApiClient(mockClient);
-        var response = await client.SetRegisteredModelAlias(full_name, alias, version_num);
+        var response = await client.SetAlias(full_name, alias, version_num);
         var responseJson = JsonSerializer.Serialize(response, Options);
         AssertJsonDeepEquals(expectedResponse, responseJson);
 

csharp/Microsoft.Azure.Databricks.Client/Pageable.cs

@@ -14,7 +14,7 @@ public AsyncPageable(Func<string, Task<(List<T>, bool, string)>> getNextPage)
             this._getNextPage = getNextPage;
         }
 
-        public override async IAsyncEnumerable<Page<T>> AsPages(string continuationToken = null, int? pageSizeHint = null)
+        public override async IAsyncEnumerable<Page<T>> AsPages(string continuationToken = default, int? pageSizeHint = default)
         {
             var nextPageToken = continuationToken;
             var hasNextPage = true;

csharp/Microsoft.Azure.Databricks.Client/ReposApiClient.cs

@@ -40,7 +40,7 @@ public async Task<Repo> Get(long repoId, CancellationToken cancellationToken = d
         return await HttpGet<Repo>(this.HttpClient, requestUri, cancellationToken).ConfigureAwait(false);
     }
 
-    public async Task<(IEnumerable<Repo>, string)> List(string pathPrefix = null, string pageToken = null, CancellationToken cancellationToken = default)
+    public async Task<(IEnumerable<Repo>, string)> List(string pathPrefix = default, string pageToken = default, CancellationToken cancellationToken = default)
     {
         var requestUri = this._apiBaseUrl + "?";
         if (pathPrefix != null)

csharp/Microsoft.Azure.Databricks.Client/UnityCatalog/Interfaces/IModelVersionApi.cs

@@ -12,13 +12,15 @@ public interface IModelVersionApi : IDisposable
     /// List model versions. You can list model versions under a particular schema, or list all model versions in the current metastore.
     /// The returned models are filtered based on the privileges of the calling user. For example, the metastore admin is able to list all the model versions. A regular user needs to be the owner or have the EXECUTE privilege on the parent registered model to recieve the model versions in the response. For the latter case, the caller must also be the owner or have the USE_CATALOG privilege on the parent catalog and the USE_SCHEMA privilege on the parent schema.
     /// </summary>
-    Task<IEnumerable<ModelVersion>> ListModelVersions(string full_name, int max_results = 0, CancellationToken cancellationToken = default);
+    Task<(IEnumerable<ModelVersion>, string)> List(string full_name, int max_results = 0, string pageToken = default, CancellationToken cancellationToken = default);
+
+    global::Azure.AsyncPageable<ModelVersion> ListPageable(string full_name, int max_results = 0, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Get a model version.
     /// The caller must be a metastore admin or an owner of (or have the EXECUTE privilege on) the parent registered model. For the latter case, the caller must also be the owner or have the USE_CATALOG privilege on the parent catalog and the USE_SCHEMA privilege on the parent schema.
     /// </summary>
-    Task<ModelVersion> GetModelVersion(
+    Task<ModelVersion> Get(
         string full_name,
         int version,
         CancellationToken cancellationToken = default
@@ -28,5 +30,5 @@ Task<ModelVersion> GetModelVersion(
     /// Get a model version by alias
     /// The caller must be a metastore admin or an owner of (or have the EXECUTE privilege on) the registered model. For the latter case, the caller must also be the owner or have the USE_CATALOG privilege on the parent catalog and the USE_SCHEMA privilege on the parent schema.
     /// </summary>
-    Task<ModelVersion> GetModelVersionByAlias(string full_name, string alias, CancellationToken cancellationToken = default);
+    Task<ModelVersion> GetByAlias(string full_name, string alias, CancellationToken cancellationToken = default);
 }

csharp/Microsoft.Azure.Databricks.Client/UnityCatalog/Interfaces/IRegisteredModelsApi.cs

@@ -11,9 +11,17 @@ public interface IRegisteredModelsApi : IDisposable
     /// <summary>
     /// List registered models. You can list registered models under a particular schema, or list all registered models in the current metastore.
     /// </summary>
-    Task<IEnumerable<RegisteredModel>> ListRegisteredModels(
-        string catalog_name = null,
-        string schema_name = null,
+    Task<(IEnumerable<RegisteredModel>, string)> List(
+        string catalog_name = default,
+        string schema_name = default,
+        int max_results = 0,
+        string pageToken = default,
+        CancellationToken cancellationToken = default
+        );
+
+    global::Azure.AsyncPageable<RegisteredModel> ListPageable(
+        string catalog_name = default,
+        string schema_name = default,
         int max_results = 0,
         CancellationToken cancellationToken = default
         );
@@ -22,16 +30,16 @@ Task<IEnumerable<RegisteredModel>> ListRegisteredModels(
     ///Get a registered model
     /// The caller must be a metastore admin or an owner of(or have the EXECUTE privilege on) the registered model.For the latter case, the caller must also be the owner or have the USE_CATALOG privilege on the parent catalog and the USE_SCHEMA privilege on the parent schema.
     /// </summary>
-    Task<RegisteredModel> GetRegisteredModel(string full_name, CancellationToken cancellationToken = default);
+    Task<RegisteredModel> Get(string full_name, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Set an alias on the specified registered model.
     /// The caller must be a metastore admin or an owner of the registered model.For the latter case, the caller must also be the owner or have the USE_CATALOG privilege on the parent catalog and the USE_SCHEMA privilege on the parent schema.
     /// </summary>
-    Task<RegisteredModelAlias> SetRegisteredModelAlias(
+    Task<RegisteredModelAlias> SetAlias(
         string full_name,
         string alias,
         int version_num,
         CancellationToken cancellationToken = default
-    );
+        );
 }

csharp/Microsoft.Azure.Databricks.Client/UnityCatalog/ModelVersionApiClient.cs

@@ -14,40 +14,51 @@ public class ModelVersionApiClient : ApiClient, IModelVersionApi
     {
         public ModelVersionApiClient(HttpClient httpClient) : base(httpClient) { }
 
-        public async Task<IEnumerable<ModelVersion>> ListModelVersions(
+        public async Task<(IEnumerable<ModelVersion>, string)> List(
               string full_name,
               int max_results = 0,
+              string pageToken = default,
               CancellationToken cancellationToken = default)
         {
-            var requestUriSb = new StringBuilder($"{BaseUnityCatalogUri}/models/{full_name}/versions");
+            var requestUriSb = new StringBuilder($"{BaseUnityCatalogUri}/models/{full_name}/versions?");
             if (max_results > 0)
             {
-                requestUriSb.Append($"?max_results={max_results}");
+                requestUriSb.Append($"&max_results={max_results}");
+            }
+
+            if (pageToken != null)
+            {
+                requestUriSb.Append($"&page_token={pageToken}");
             }
 
             var requestUri = requestUriSb.ToString();
-            var modelVersionsJson = await HttpGet<JsonObject>(HttpClient, requestUri, cancellationToken).ConfigureAwait(false);
-            modelVersionsJson.TryGetPropertyValue("model_versions", out var modelVersions);
-            return modelVersions?.Deserialize<IEnumerable<ModelVersion>>(Options) ?? Enumerable.Empty<ModelVersion>();
-        }
+            var response = await HttpGet<JsonObject>(HttpClient, requestUri, cancellationToken).ConfigureAwait(false);
+            response.TryGetPropertyValue("model_versions", out var modelVersions);
+            response.TryGetPropertyValue("next_page_token", out var nextPageTokenNode);
+
+            var versions = modelVersions?.Deserialize<IEnumerable<ModelVersion>>(Options) ?? Enumerable.Empty<ModelVersion>();
+            var nextPageToken = nextPageTokenNode?.GetValue<string>() ?? string.Empty;
 
+            return (versions, nextPageToken);
+        }
 
+        public global::Azure.AsyncPageable<ModelVersion> ListPageable(string full_name, int max_results = 0, CancellationToken cancellationToken = default)
+        {
+            return new AsyncPageable<ModelVersion>(async (pageToken) =>
+            {
+                var (versions, nextPageToken) = await List(full_name, max_results, pageToken, cancellationToken).ConfigureAwait(false);
+                return (versions.ToList(), !string.IsNullOrEmpty(nextPageToken), nextPageToken);
+            });
+        }
 
-        public async Task<ModelVersion> GetModelVersion(
-           string full_name,
-           int version,
-           CancellationToken cancellationToken = default
-        )
+        public async Task<ModelVersion> Get(string full_name, int version, CancellationToken cancellationToken = default)
         {
             var requestUriSb = new StringBuilder($"{BaseUnityCatalogUri}/models/{full_name}/versions/{version}");
             var requestUri = requestUriSb.ToString();
             return await HttpGet<ModelVersion>(HttpClient, requestUri, cancellationToken).ConfigureAwait(false);
         }
 
-        public async Task<ModelVersion> GetModelVersionByAlias(
-            string full_name,
-            string alias,
-            CancellationToken cancellationToken = default)
+        public async Task<ModelVersion> GetByAlias(string full_name, string alias, CancellationToken cancellationToken = default)
         {
             var requestUri = $"{BaseUnityCatalogUri}/models/{full_name}/aliases/{alias}";
             return await HttpGet<ModelVersion>(HttpClient, requestUri, cancellationToken).ConfigureAwait(false);