Consistent pagination experience across all APIs (#86)

# Goal: present consistent entity listing for the end-user

Only 30% of APIs have pagination, though there are two different
implementations of it. This PR proposes generating wrapper methods to
return consistent type for a collection of all results. It also adds
frequently used wrappers, like name-to-id mapping, which could later be
used to generate client-level upserts.

Decisions to discuss:

## Return type: slice vs iterator

* Current implementation proposes returning slice + error
* theoretically, with added complexity, we can move to iterator + error
implementation, though Go ecosystem didn't settle on the standard
iterator yet: golang/go#54245
* other ecosystems, like Python, Java, and JavaScript should lean
towards iterators

## Naming: ListQueriesAll (suffix) vs ListAllQueries (infix) vs all
methods returning slices (might be invasive)

* current template has "All" name suffix for methods returning all
results
* 18% of APIs return slice from the platform.
* 36% of APIs return all results, but in a single-field response.

### Out of scope of this PR
* splitting high-level generated interface from low-level direct mapping #87
* telemetry for higher-level generated code #88

.gitignore

@@ -17,6 +17,7 @@ vendor/
 coverage.txt
 
 .vscode/*.lograw-specs/**
+.vscode/*.log
 
 raw-specs/**
 out/**

databricks/openapi/code/load_test.go

@@ -67,8 +67,11 @@ func TestBasicDebug(t *testing.T) {
 	batch, err := NewFromFile("/tmp/processed-databricks-workspace-all.json", []string{})
 	assert.NoError(t, err)
 
-	m := batch.Packages["unitycatalog"].services["Catalogs"].methods["create"]
+	m := batch.Packages["dbsql"].services["Alerts"].methods["updateAlert"]
 	t.Log(m)
 
+	e := batch.Packages["dbsql"].types["EditAlert"]
+	t.Log(e)
+
 	assert.Len(t, batch.Packages, 1)
 }

databricks/openapi/code/method.go

@@ -34,6 +34,7 @@ type Pagination struct {
 	Limit     *Field
 	Results   *Field
 	Entity    *Entity
+	Token     *Binding
 	Increment int
 }
 
@@ -141,9 +142,17 @@ func (m *Method) Pagination() *Pagination {
 		// we assume that method already returns body-as-array
 		return nil
 	}
+	var token *Binding
+	if m.pagination.Token != nil {
+		token = &Binding{ // reuse the same datastructure as for waiters
+			PollField: m.Request.Field(m.pagination.Token.Request),
+			Bind:      m.Response.Field(m.pagination.Token.Response),
+		}
+	}
 	results := m.Response.Field(m.pagination.Results)
 	return &Pagination{
 		Results:   results,
+		Token:     token,
 		Entity:    results.Entity.ArrayValue,
 		Offset:    m.Request.Field(m.pagination.Offset),
 		Limit:     m.Request.Field(m.pagination.Limit),
@@ -162,6 +171,10 @@ func (m *Method) paginationItem() *Entity {
 	return m.Pagination().Entity
 }
 
+func (p *Pagination) MultiRequest() bool {
+	return p.Offset != nil || p.Token != nil
+}
+
 func (m *Method) NamedIdMap() *NamedIdMap {
 	entity := m.paginationItem()
 	if entity == nil {
@@ -188,18 +201,13 @@ func (m *Method) NamedIdMap() *NamedIdMap {
 	}
 }
 
-// GetByName returns method from the same service with x-databricks-crud:read
-func (m *Method) GetByName() *Method {
-	if m.NamedIdMap() == nil {
+// GetByName returns entity from the same service with x-databricks-crud:read
+func (m *Method) GetByName() *Entity {
+	n := m.NamedIdMap()
+	if n == nil {
 		return nil
 	}
-	for _, v := range m.Service.methods {
-		if strings.ToLower(v.operation.Crud) != "read" {
-			continue
-		}
-		return v
-	}
-	return nil
+	return n.Entity
 }
 
 func (m *Method) CanHaveResponseBody() bool {

databricks/openapi/code/package.go

@@ -51,6 +51,15 @@ func (pkg *Package) EmptyTypes() (types []*Named) {
 	return types
 }
 
+func (pkg *Package) HasPagination() bool {
+	for _, v := range pkg.services {
+		if v.HasPagination() {
+			return true
+		}
+	}
+	return false
+}
+
 func (pkg *Package) schemaToEntity(s *openapi.Schema, path []string, hasName bool) *Entity {
 	if s.IsRef() {
 		// if schema is src, load it to this package
@@ -180,7 +189,8 @@ func (pkg *Package) definedEntity(name string, s *openapi.Schema) *Entity {
 	if e.Name == "" {
 		e.Named = Named{name, s.Description}
 	}
-	return pkg.define(e)
+	pkg.define(e)
+	return pkg.types[e.Name]
 }
 
 func (pkg *Package) define(entity *Entity) *Entity {

databricks/openapi/code/service.go

@@ -31,6 +31,19 @@ func (svc *Service) Methods() (methods []*Method) {
 	return methods
 }
 
+func (svc *Service) HasPagination() bool {
+	for _, v := range svc.methods {
+		p := v.pagination
+		if p == nil {
+			continue
+		}
+		if p.Offset != "" || p.Token != nil {
+			return true
+		}
+	}
+	return false
+}
+
 func (svc *Service) paramToField(op *openapi.Operation, param openapi.Parameter) *Field {
 	named := Named{param.Name, param.Description}
 	return &Field{

databricks/openapi/model.go

@@ -106,11 +106,12 @@ func (o *Operation) Name() string {
 }
 
 type Pagination struct {
-	Offset    string `json:"offset,omitempty"`
-	Limit     string `json:"limit,omitempty"`
-	Results   string `json:"results,omitempty"`
-	Increment int    `json:"increment,omitempty"`
-	Inline    bool   `json:"inline,omitempty"`
+	Offset    string   `json:"offset,omitempty"`
+	Limit     string   `json:"limit,omitempty"`
+	Results   string   `json:"results,omitempty"`
+	Increment int      `json:"increment,omitempty"`
+	Inline    bool     `json:"inline,omitempty"`
+	Token     *Binding `json:"token,omitempty"`
 }
 
 type Wait struct {

internal/dbsql_test.go

@@ -38,3 +38,60 @@ func TestAccQueries(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, query.Query, loaded.Query)
 }
+
+func TestAccDashboards(t *testing.T) {
+	env := GetEnvOrSkipTest(t, "CLOUD_ENV")
+	t.Log(env)
+	ctx := context.Background()
+	wsc := workspaces.New()
+
+	all, err := wsc.Dashboards.ListDashboardsAll(ctx, dbsql.ListDashboardsRequest{})
+	require.NoError(t, err)
+	t.Log(len(all))
+}
+
+func TestAccQueriesList(t *testing.T) {
+	env := GetEnvOrSkipTest(t, "CLOUD_ENV")
+	t.Log(env)
+	ctx := context.Background()
+	wsc := workspaces.New()
+
+	srcs, err := wsc.DataSources.ListDataSources(ctx)
+	require.NoError(t, err)
+	if len(srcs) == 0 {
+		t.Skipf("no sql warehouses found")
+	}
+	for i := 0; i < 34; i++ {
+		query, err := wsc.Queries.CreateQuery(ctx, dbsql.QueryPostContent{
+			Name:         RandomName("go-sdk/test/"),
+			DataSourceId: srcs[0].Id,
+			Description:  "test query from Go SDK",
+			Query:        "SHOW TABLES",
+		})
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			err := wsc.Queries.DeleteQueryByQueryId(ctx, query.Id)
+			require.NoError(t, err)
+		})
+	}
+	var qs1, qs2, qs3 []dbsql.Query
+	{
+		result, err := wsc.Queries.ListQueries(ctx, dbsql.ListQueriesRequest{PageSize: 10})
+		require.NoError(t, err)
+		qs1 = result.Results
+	}
+	{
+		result, err := wsc.Queries.ListQueriesAll(ctx, dbsql.ListQueriesRequest{})
+		require.NoError(t, err)
+		qs2 = result
+	}
+	{
+		result, err := wsc.Queries.ListQueriesAll(ctx, dbsql.ListQueriesRequest{PageSize: 10})
+		require.NoError(t, err)
+		qs3 = result
+	}
+	for i := 0; i < len(qs1) && i < len(qs2); i++ {
+		assert.Equal(t, qs1[i], qs2[i], "Query at index %d not equal", i)
+		assert.Equal(t, qs1[i], qs3[i], "Query at index %d not equal", i)
+	}
+}

internal/jobs_test.go

@@ -196,3 +196,62 @@ func TestAccJobsApiFullIntegration(t *testing.T) {
 		assert.NotEqual(t, job.JobId, createdJob.JobId)
 	}
 }
+
+func TestAccJobsListAllNoDuplicates(t *testing.T) {
+	env := GetEnvOrSkipTest(t, "CLOUD_ENV")
+	t.Log(env)
+	ctx := context.Background()
+	w := workspaces.New()
+
+	// Fetch list of spark runtime versions
+	sparkVersions, err := w.Clusters.SparkVersions(ctx)
+	require.NoError(t, err)
+
+	// Select the latest LTS version
+	latestLTS, err := sparkVersions.Select(clusters.SparkVersionRequest{
+		Latest:          true,
+		LongTermSupport: true,
+	})
+	require.NoError(t, err)
+
+	// Fetch list of available node types
+	nodeTypes, err := w.Clusters.ListNodeTypes(ctx)
+	require.NoError(t, err)
+
+	// Select the smallest node type id
+	smallestWithDisk, err := nodeTypes.Smallest(clusters.NodeTypeRequest{
+		LocalDisk: true,
+	})
+	require.NoError(t, err)
+
+	for i := 0; i < 34; i++ {
+		createdJob, err := w.Jobs.Create(ctx, jobs.CreateJob{
+			Name: RandomName(t.Name()),
+			Tasks: []jobs.JobTaskSettings{{
+				Description: "test",
+				NewCluster: &jobs.NewCluster{
+					SparkVersion: latestLTS,
+					NodeTypeId:   smallestWithDisk,
+					NumWorkers:   1,
+				},
+				TaskKey:        "test",
+				TimeoutSeconds: 0,
+				NotebookTask: &jobs.NotebookTask{
+					NotebookPath: "/foo/bar",
+				},
+			}},
+		})
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			err := w.Jobs.DeleteByJobId(ctx, createdJob.JobId)
+			require.NoError(t, err)
+		})
+	}
+	all, err := w.Jobs.ListAll(ctx, jobs.ListRequest{})
+	require.NoError(t, err)
+	ids := map[int64]bool{}
+	for _, v := range all {
+		ids[v.JobId] = true
+	}
+	assert.Equal(t, len(all), len(ids), "Listing produced duplicate results")
+}

internal/repos_test.go

@@ -50,6 +50,7 @@ func TestAccRepos(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	_, err = wsc.Repos.List(ctx, repos.ListRequest{})
+	all, err := wsc.Repos.ListAll(ctx, repos.ListRequest{})
 	require.NoError(t, err)
+	assert.True(t, len(all) >= 1)
 }

service/.codegen/api.go.tmpl

@@ -8,6 +8,8 @@ import (
 	{{if .HasWaits}}"time"
 
 	"github.com/databricks/databricks-sdk-go/retries"{{end}}
+	{{if .HasPagination}}"github.com/databricks/databricks-sdk-go/databricks/useragent"
+	{{- end}}
 	"github.com/databricks/databricks-sdk-go/databricks/client"
 )
 {{range .Services}}
@@ -22,6 +24,10 @@ type {{.PascalName}}API struct {
 }
 
 {{range .Methods}}{{.Comment "// " 80}}
+{{- if .Pagination}}
+// 
+// Use {{.PascalName}}All() to get all {{.Pagination.Entity.PascalName}} instances{{if .Pagination.MultiRequest}}, which will iterate over every result page.{{end}}
+{{- end}}
 func (a *{{.Service.Name}}API) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{if .Response}}({{if .Response.ArrayValue}}[]{{.Response.ArrayValue.PascalName}}{{else}}*{{.Response.PascalName}}{{end}}, error){{else}}error{{end}} {
 	{{if .Response}}var {{.Response.CamelName}} {{if .Response.ArrayValue}}[]{{.Response.ArrayValue.PascalName}}{{else}}{{.Response.PascalName}}{{end}}
 	{{end -}}
@@ -63,9 +69,15 @@ func (a *{{.Service.Name}}API) {{.PascalName}}AndWait(ctx context.Context{{if .R
 	})
 }
 {{end}}{{if .Pagination}}
+// {{.PascalName}}All returns all {{.Pagination.Entity.PascalName}} instances{{if .Pagination.MultiRequest}} by calling {{.PascalName}} for every result page{{else}}. This method exists for consistency purposes.{{end}}
+// 
+// This method is generated by Databricks SDK Code Generator.
 func (a *{{.Service.Name}}API) {{.PascalName}}All(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) ([]{{.Pagination.Entity.PascalName}}, error) {
-	{{if .Pagination.Offset}}var results []{{.Pagination.Entity.PascalName}}
-	for {
+	{{if .Pagination.MultiRequest}}var results []{{.Pagination.Entity.PascalName}}
+	ctx = useragent.InContext(ctx, "sdk-feature", "pagination")
+	{{if eq .Pagination.Increment 1 -}}
+	request.{{.Pagination.Offset.PascalName}} = 1 // start iterating from the first page
+	{{end}}for {
 		response, err := a.{{.PascalName}}(ctx{{if .Request}}, request{{end}})
 		if err != nil {
 			return nil, err
@@ -76,7 +88,16 @@ func (a *{{.Service.Name}}API) {{.PascalName}}All(ctx context.Context{{if .Reque
 		for _, v := range response.{{.Pagination.Results.PascalName}} {
 			results = append(results, v)
 		}
-		request.{{.Pagination.Offset.PascalName}} += {{.Pagination.Increment}}
+		{{if .Pagination.Token -}}
+		request.{{.Pagination.Token.PollField.PascalName}} = response.{{.Pagination.Token.Bind.PascalName}}
+		if response.{{.Pagination.Token.Bind.PascalName}} == "" {
+			break
+		}
+		{{- else if eq .Pagination.Increment 1 -}}
+		request.{{.Pagination.Offset.PascalName}}++
+		{{- else -}}
+		request.{{.Pagination.Offset.PascalName}} += {{template "type" .Pagination.Offset.Entity}}(len(response.{{.Pagination.Results.PascalName}}))
+		{{- end}}
 	}
 	return results, nil{{else if .Pagination.Results}}response, err := a.{{.PascalName}}(ctx{{if .Request}}, request{{end}})
 	if err != nil {
@@ -97,7 +118,7 @@ func (a *{{.Service.Name}}API) {{.NamedIdMap.Entity.PascalName}}{{.NamedIdMap.Na
 	return mapping, nil
 }
 {{end}}{{if .GetByName}}
-func (a *{{.Service.Name}}API) Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}}(ctx context.Context, name string) (*{{.GetByName.Response.PascalName}}, error) {
+func (a *{{.Service.Name}}API) Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}}(ctx context.Context, name string) (*{{.GetByName.PascalName}}, error) {
 	result, err := a.{{.PascalName}}{{if not .NamedIdMap.Direct}}All{{end}}(ctx{{if .Request}}, {{.Request.PascalName}}{}{{end}})
 	if err != nil {
 		return nil, err

service/.codegen/interface.go.tmpl

@@ -16,6 +16,10 @@ import (
 type {{.PascalName}}Service interface {
 	{{range .Methods}}
     {{.Comment "    // " 80}}
+	{{- if .Pagination}}
+	// 
+	// Use {{.PascalName}}All() to get all {{.Pagination.Entity.PascalName}} instances{{if .Pagination.MultiRequest}}, which will iterate over every result page.{{end}}
+	{{- end}}
     {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{if .Response}}({{if .Response.ArrayValue}}[]{{.Response.ArrayValue.PascalName}}{{else}}*{{.Response.PascalName}}{{end}}, error){{else}}error{{end}}
 	{{if .Wait}}
 	// {{.PascalName}}AndWait calls {{.PascalName}}() and waits to reach {{range $i, $e := .Wait.Success}}{{if $i}} or {{end}}{{.Content}}{{end}} state
@@ -39,10 +43,10 @@ type {{.PascalName}}Service interface {
 	// This method is generated by Databricks SDK Code Generator.
 	{{.NamedIdMap.Entity.PascalName}}{{.NamedIdMap.Name.PascalName}}To{{.NamedIdMap.Id.PascalName}}Map(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) (map[string]{{template "type" .NamedIdMap.Id.Entity}}, error){{end}}
 	{{if .GetByName}}
-	// Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}} retrieves {{.GetByName.Response.PascalName}} by name.
+	// Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}} retrieves {{.GetByName.PascalName}} by name.
 	// 
 	// This method is generated by Databricks SDK Code Generator.
-	Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}}(ctx context.Context, name string) (*{{.GetByName.Response.PascalName}}, error){{end}}
+	Get{{.NamedIdMap.Entity.PascalName}}By{{.NamedIdMap.Name.PascalName}}(ctx context.Context, name string) (*{{.GetByName.PascalName}}, error){{end}}
 	{{- if and .Shortcut .Wait}}
 	// {{.Shortcut.PascalName}}AndWait calls {{.Shortcut.PascalName}} and waits until {{.Wait.Poll.Response.PascalName}} is in desired state.
 	//