From bbccb6074942d22d306e9635e16405e05da03b72 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Thu, 13 Jul 2023 15:47:24 +0200 Subject: [PATCH] Fix pagination --- packages/databricks-sdk-js/package.json | 2 +- .../src/.codegen/api.ts.tmpl | 33 ++++++++++------ .../databricks-sdk-js/src/Pagination.integ.ts | 28 +++----------- .../databricks-sdk-js/src/Pagination.test.ts | 38 ++++++------------- .../databricks-sdk-js/src/apis/jobs/api.ts | 14 +++++++ .../databricks-sdk-js/src/apis/sql/api.ts | 24 +++++++++++- .../src/apis/workspace/model.ts | 4 +- 7 files changed, 77 insertions(+), 66 deletions(-) diff --git a/packages/databricks-sdk-js/package.json b/packages/databricks-sdk-js/package.json index 2d390f607..df2a8a9f1 100644 --- a/packages/databricks-sdk-js/package.json +++ b/packages/databricks-sdk-js/package.json @@ -21,7 +21,7 @@ "watch": "tsc --build --watch", "clean": "rm -rf dist node_modules", "openapi:fetch": "./scripts/fetch_openapi.sh d4525bbc428d236f7508d3190f973fbf17fa5614", - "openapi:generate": "./scripts/generate_openapi.sh v0.12.0 && yarn run build", + "openapi:generate": "./scripts/generate_openapi.sh 3b4492b6d659ca3d03035e76e1bbc2c964d44f19 && yarn run build", "generate-notice": "../../scripts/generate_notice.sh", "fix": "eslint src --ext ts --fix && prettier . --write", "test:lint": "eslint src --ext ts && prettier . -c", diff --git a/packages/databricks-sdk-js/src/.codegen/api.ts.tmpl b/packages/databricks-sdk-js/src/.codegen/api.ts.tmpl index 860404abd..c56c2cebc 100644 --- a/packages/databricks-sdk-js/src/.codegen/api.ts.tmpl +++ b/packages/databricks-sdk-js/src/.codegen/api.ts.tmpl @@ -123,12 +123,15 @@ export class {{$s.PascalName}}Service { @context context?: Context ): AsyncIterable<{{$s.Package.Name}}.{{.Pagination.Entity.PascalName}}> { {{if .Pagination.MultiRequest}} - {{if .Pagination.NeedsOffsetDedupe -}} + {{if .NeedsOffsetDedupe -}} // deduplicate items that may have been added during iteration - const seen: Record<{{template "type" .Pagination.Entity.IdentifierField.Entity}}, boolean> = {}; + const seen: Record<{{template "type" .IdentifierField.Entity}}, boolean> = {}; {{end}}{{if eq .Pagination.Increment 1 -}} request.{{.Pagination.Offset.Name}} = 1; // start iterating from the first page {{end}} + {{- if and .Pagination.Token .Pagination.Limit -}} + let totalCount = 0; + {{ end -}} while(true) { const response = await this._{{.CamelName}}({{if .Request}}request,{{end}} context); if ( @@ -143,8 +146,8 @@ export class {{$s.PascalName}}Service { } for (const v of response.{{.Pagination.Results.Name}}) { - {{- if .Pagination.NeedsOffsetDedupe -}} - const id = v.{{.Pagination.Entity.IdentifierField.Name}}; + {{- if .NeedsOffsetDedupe -}} + const id = v.{{.IdentifierField.Name}}; if (id) { if (seen[id]) { // item was added during iteration @@ -162,15 +165,21 @@ export class {{$s.PascalName}}Service { } request = response.next_page; {{- else if .Pagination.Token -}} - request.{{.Pagination.Token.PollField.Name}} = response.{{.Pagination.Token.Bind.Name}} - if (!response.{{.Pagination.Token.Bind.Name}}) { - break; - } + request.{{.Pagination.Token.PollField.Name}} = response.{{.Pagination.Token.Bind.Name}} + if (!response.{{.Pagination.Token.Bind.Name}}) { + break; + } {{- else if eq .Pagination.Increment 1 -}} - request.{{.Pagination.Offset.Name}}+= 1; - {{- else -}} - request.{{.Pagination.Offset.Name}} = request.{{.Pagination.Offset.Name}} || 0; - request.{{.Pagination.Offset.Name}} += response.{{.Pagination.Results.Name}}.length; + // paginate by increments of 1 + request.{{.Pagination.Offset.Name}}+= 1; + {{- end}} + {{- if and .Pagination.Token .Pagination.Limit -}} + const count = response.{{.Pagination.Results.Name}}.length + totalCount += count + const limit = request.{{.Pagination.Limit.Name}} + if (limit && totalCount >= limit) { + break; + } {{- end}} }{{else if .Pagination.Results}} const response = (await this._{{.CamelName}}({{if .Request}}request,{{end}} context)).{{.Pagination.Results.Name}}; diff --git a/packages/databricks-sdk-js/src/Pagination.integ.ts b/packages/databricks-sdk-js/src/Pagination.integ.ts index a252c82c6..3ccbe62f1 100644 --- a/packages/databricks-sdk-js/src/Pagination.integ.ts +++ b/packages/databricks-sdk-js/src/Pagination.integ.ts @@ -28,26 +28,6 @@ describe.skip(__filename, function () { assert.ok(items.length > 0); }); - // jobs list - it("should paginate by token and dedupe results", async () => { - const items = []; - const seen = new Set(); - for await (const job of wsClient.jobs.list({})) { - items.push(job); - assert.ok(job.job_id); - if (seen.has(job.job_id!)) { - assert.fail(`job_id ${job.job_id} already seen`); - } else { - seen.add(job.job_id!); - } - if (items.length > 50) { - break; - } - } - - assert.ok(items.length > 0); - }); - // jobs list it("should paginate by offset", async () => { const items = []; @@ -77,12 +57,16 @@ describe.skip(__filename, function () { }); it("should paginate cluster events", async () => { + assert.ok( + process.env.TEST_DEFAULT_CLUSTER_ID, + "TEST_DEFAULT_CLUSTER_ID must be set" + ); const items = []; for await (const item of wsClient.clusters.events({ - cluster_id: process.env.DATABRICKS_CLUSTER_ID!, + cluster_id: process.env.TEST_DEFAULT_CLUSTER_ID, })) { items.push(item); - if (items.length > 50) { + if (items.length > 60) { break; } } diff --git a/packages/databricks-sdk-js/src/Pagination.test.ts b/packages/databricks-sdk-js/src/Pagination.test.ts index 8c6a364ef..a6bff8964 100644 --- a/packages/databricks-sdk-js/src/Pagination.test.ts +++ b/packages/databricks-sdk-js/src/Pagination.test.ts @@ -52,7 +52,7 @@ describe(__filename, () => { }); // jobs list - it("should paginate by token and dedupe results", async () => { + it("should paginate by token", async () => { const responses = [ { jobs: [...Array(20).keys()].map((i) => { @@ -61,28 +61,18 @@ describe(__filename, () => { creator_user_name: `test-user-${i}`, }; }), - has_more: true, - }, - { - jobs: [...Array(20).keys()] - .map((i) => { - return { - job_id: 43368721962707 + i + 20, - creator_user_name: `test-user-${i + 20}`, - }; - }) - .concat([ - { - // duplicate entry - job_id: 43368721962707, - creator_user_name: `test-user-0`, - }, - ]), - has_more: true, + next_page_token: "foo", }, { - has_more: false, + jobs: [...Array(20).keys()].map((i) => { + return { + job_id: 43368721962707 + i + 20, + creator_user_name: `test-user-${i + 20}`, + }; + }), + next_page_token: "bar", }, + {}, ]; const jobsApi = wsClient.jobs; @@ -91,15 +81,9 @@ describe(__filename, () => { }; const items = []; - const seen = new Set(); for await (const job of wsClient.jobs.list({})) { items.push(job); assert.ok(job.job_id); - if (seen.has(job.job_id!)) { - assert.fail(`job_id ${job.job_id} already seen`); - } else { - seen.add(job.job_id!); - } if (items.length > 50) { break; } @@ -146,7 +130,7 @@ describe(__filename, () => { items.push(job); } - assert.equal(items.length, 50); + assert.equal(items.length, 25); }); // sql dashboards list diff --git a/packages/databricks-sdk-js/src/apis/jobs/api.ts b/packages/databricks-sdk-js/src/apis/jobs/api.ts index b304283e6..01f3cb4f9 100755 --- a/packages/databricks-sdk-js/src/apis/jobs/api.ts +++ b/packages/databricks-sdk-js/src/apis/jobs/api.ts @@ -438,6 +438,7 @@ export class JobsService { request: jobs.ListJobsRequest, @context context?: Context ): AsyncIterable { + let totalCount = 0; while (true) { const response = await this._list(request, context); if ( @@ -459,6 +460,12 @@ export class JobsService { if (!response.next_page_token) { break; } + const count = response.jobs.length; + totalCount += count; + const limit = request.limit; + if (limit && totalCount >= limit) { + break; + } } } @@ -486,6 +493,7 @@ export class JobsService { request: jobs.ListRunsRequest, @context context?: Context ): AsyncIterable { + let totalCount = 0; while (true) { const response = await this._listRuns(request, context); if ( @@ -507,6 +515,12 @@ export class JobsService { if (!response.next_page_token) { break; } + const count = response.runs.length; + totalCount += count; + const limit = request.limit; + if (limit && totalCount >= limit) { + break; + } } } diff --git a/packages/databricks-sdk-js/src/apis/sql/api.ts b/packages/databricks-sdk-js/src/apis/sql/api.ts index d754464de..dec5b871b 100755 --- a/packages/databricks-sdk-js/src/apis/sql/api.ts +++ b/packages/databricks-sdk-js/src/apis/sql/api.ts @@ -298,8 +298,9 @@ export class DashboardsService { request: sql.ListDashboardsRequest, @context context?: Context ): AsyncIterable { + // deduplicate items that may have been added during iteration + const seen: Record = {}; request.page = 1; // start iterating from the first page - while (true) { const response = await this._list(request, context); if ( @@ -314,9 +315,18 @@ export class DashboardsService { } for (const v of response.results) { + const id = v.id; + if (id) { + if (seen[id]) { + // item was added during iteration + continue; + } + seen[id] = true; + } yield v; } + // paginate by increments of 1 request.page += 1; } } @@ -653,8 +663,9 @@ export class QueriesService { request: sql.ListQueriesRequest, @context context?: Context ): AsyncIterable { + // deduplicate items that may have been added during iteration + const seen: Record = {}; request.page = 1; // start iterating from the first page - while (true) { const response = await this._list(request, context); if ( @@ -669,9 +680,18 @@ export class QueriesService { } for (const v of response.results) { + const id = v.id; + if (id) { + if (seen[id]) { + // item was added during iteration + continue; + } + seen[id] = true; + } yield v; } + // paginate by increments of 1 request.page += 1; } } diff --git a/packages/databricks-sdk-js/src/apis/workspace/model.ts b/packages/databricks-sdk-js/src/apis/workspace/model.ts index 9817e3a2e..c458c73fe 100755 --- a/packages/databricks-sdk-js/src/apis/workspace/model.ts +++ b/packages/databricks-sdk-js/src/apis/workspace/model.ts @@ -188,12 +188,12 @@ export interface DeleteSecret { } export type ExportFormat = + | "AUTO" | "DBC" | "HTML" | "JUPYTER" | "R_MARKDOWN" - | "SOURCE" - | "AUTO"; + | "SOURCE"; /** * Export a workspace object