Skip to content

Commit

Permalink
Fix pagination
Browse files Browse the repository at this point in the history
  • Loading branch information
fjakobs committed Jul 13, 2023
1 parent a340c23 commit bbccb60
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 66 deletions.
2 changes: 1 addition & 1 deletion packages/databricks-sdk-js/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
33 changes: 21 additions & 12 deletions packages/databricks-sdk-js/src/.codegen/api.ts.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand All @@ -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}};
Expand Down
28 changes: 6 additions & 22 deletions packages/databricks-sdk-js/src/Pagination.integ.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<number>();
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 = [];
Expand Down Expand Up @@ -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;
}
}
Expand Down
38 changes: 11 additions & 27 deletions packages/databricks-sdk-js/src/Pagination.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand All @@ -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;
Expand All @@ -91,15 +81,9 @@ describe(__filename, () => {
};

const items = [];
const seen = new Set<number>();
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;
}
Expand Down Expand Up @@ -146,7 +130,7 @@ describe(__filename, () => {
items.push(job);
}

assert.equal(items.length, 50);
assert.equal(items.length, 25);
});

// sql dashboards list
Expand Down
14 changes: 14 additions & 0 deletions packages/databricks-sdk-js/src/apis/jobs/api.ts

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

24 changes: 22 additions & 2 deletions packages/databricks-sdk-js/src/apis/sql/api.ts

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions packages/databricks-sdk-js/src/apis/workspace/model.ts

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit bbccb60

Please sign in to comment.