Skip to content

Commit 510f600

Browse files
alexottalexott
authored
Use Jobs API 2.1 with name filter for databricks_job data source (#1744)
This patch switches to 2.1 version of Jobs API for List operation so we can more efficiently handle workspace with huge number of jobs (> 3000). It also uses new `name` parameter to search jobs by name in the `databricks_job` data source.
1 parent 87e590d commit 510f600

File tree

10 files changed

+153
-51
lines changed

10 files changed

+153
-51
lines changed

‎exporter/exporter_test.go‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -341,8 +341,8 @@ func TestImportingUsersGroupsSecretScopes(t *testing.T) {
341341
},
342342
{
343343
Method: "GET",
344-
Resource: "/api/2.0/jobs/list",
345-
Response: jobs.JobList{},
344+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
345+
Response: jobs.JobListResponse{},
346346
},
347347
{
348348
Method: "GET",
@@ -441,8 +441,8 @@ func TestImportingNoResourcesError(t *testing.T) {
441441
},
442442
{
443443
Method: "GET",
444-
Resource: "/api/2.0/jobs/list",
445-
Response: jobs.JobList{},
444+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
445+
Response: jobs.JobListResponse{},
446446
},
447447
{
448448
Method: "GET",
@@ -484,8 +484,8 @@ func TestImportingClusters(t *testing.T) {
484484
},
485485
{
486486
Method: "GET",
487-
Resource: "/api/2.0/jobs/list",
488-
Response: jobs.JobList{},
487+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
488+
Response: jobs.JobListResponse{},
489489
},
490490
{
491491
Method: "GET",
@@ -628,8 +628,8 @@ func TestImportingJobs_JobList(t *testing.T) {
628628
emptyRepos,
629629
{
630630
Method: "GET",
631-
Resource: "/api/2.0/jobs/list",
632-
Response: jobs.JobList{
631+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
632+
Response: jobs.JobListResponse{
633633
Jobs: []jobs.Job{
634634
{
635635
JobID: 14,
@@ -842,8 +842,8 @@ func TestImportingJobs_JobListMultiTask(t *testing.T) {
842842
emptyRepos,
843843
{
844844
Method: "GET",
845-
Resource: "/api/2.0/jobs/list",
846-
Response: jobs.JobList{
845+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
846+
Response: jobs.JobListResponse{
847847
Jobs: []jobs.Job{
848848
{
849849
JobID: 14,
@@ -1085,8 +1085,8 @@ func TestImportingSecrets(t *testing.T) {
10851085
},
10861086
{
10871087
Method: "GET",
1088-
Resource: "/api/2.0/jobs/list",
1089-
Response: jobs.JobList{},
1088+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
1089+
Response: jobs.JobListResponse{},
10901090
},
10911091
{
10921092
Method: "GET",

‎exporter/importables.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -445,10 +445,10 @@ var resourcesMap map[string]importable = map[string]importable{
445445
}
446446
}
447447
if task.DbtTask != nil {
448-
if task.SqlTask.WarehouseID != "" {
448+
if task.DbtTask.WarehouseId != "" {
449449
ic.Emit(&resource{
450450
Resource: "databricks_sql_endpoint",
451-
ID: task.SqlTask.WarehouseID,
451+
ID: task.DbtTask.WarehouseId,
452452
})
453453
}
454454
}

‎exporter/importables_test.go‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -227,12 +227,10 @@ func TestClusterList_NoNameMatch(t *testing.T) {
227227
func TestJobListNoNameMatch(t *testing.T) {
228228
ic := importContextForTest()
229229
ic.match = "bcd"
230-
ic.importJobs(jobs.JobList{
231-
Jobs: []jobs.Job{
232-
{
233-
Settings: &jobs.JobSettings{
234-
Name: "abc",
235-
},
230+
ic.importJobs([]jobs.Job{
231+
{
232+
Settings: &jobs.JobSettings{
233+
Name: "abc",
236234
},
237235
},
238236
})
@@ -251,13 +249,11 @@ func TestJobList_FailGetRuns(t *testing.T) {
251249
ic := importContextForTest()
252250
ic.Client = client
253251
ic.Context = ctx
254-
ic.importJobs(jobs.JobList{
255-
Jobs: []jobs.Job{
256-
{
257-
JobID: 1,
258-
Settings: &jobs.JobSettings{
259-
Name: "abc",
260-
},
252+
ic.importJobs([]jobs.Job{
253+
{
254+
JobID: 1,
255+
Settings: &jobs.JobSettings{
256+
Name: "abc",
261257
},
262258
},
263259
})

‎exporter/util.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -281,12 +281,12 @@ func eitherString(a any, b any) string {
281281
return ""
282282
}
283283

284-
func (ic *importContext) importJobs(l jobs.JobList) {
284+
func (ic *importContext) importJobs(l []jobs.Job) {
285285
nowSeconds := time.Now().Unix()
286286
a := jobs.NewJobsAPI(ic.Context, ic.Client)
287287
starterAfter := (nowSeconds - (ic.lastActiveDays * 24 * 60 * 60)) * 1000
288288
i := 0
289-
for _, job := range l.Jobs {
289+
for _, job := range l {
290290
if !ic.MatchesName(job.Settings.Name) {
291291
log.Printf("[INFO] Job name %s doesn't match selection %s", job.Settings.Name, ic.match)
292292
continue
@@ -317,7 +317,7 @@ func (ic *importContext) importJobs(l jobs.JobList) {
317317
ID: job.ID(),
318318
})
319319
i++
320-
log.Printf("[INFO] Imported %d of total %d jobs", i, len(l.Jobs))
320+
log.Printf("[INFO] Imported %d of total %d jobs", i, len(l))
321321
}
322322
}
323323

‎jobs/data_job.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package jobs
33
import (
44
"context"
55
"fmt"
6+
67
"github.com/databricks/terraform-provider-databricks/common"
78
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
89
)
@@ -16,11 +17,17 @@ func DataSourceJob() *schema.Resource {
1617
return common.DataResource(queryableJobData{}, func(ctx context.Context, e any, c *common.DatabricksClient) error {
1718
data := e.(*queryableJobData)
1819
jobsAPI := NewJobsAPI(ctx, c)
19-
list, err := jobsAPI.List()
20+
var list []Job
21+
var err error
22+
if data.Name != "" {
23+
list, err = jobsAPI.ListByName(data.Name, false)
24+
} else {
25+
list, err = jobsAPI.List()
26+
}
2027
if err != nil {
2128
return err
2229
}
23-
for _, job := range list.Jobs {
30+
for _, job := range list {
2431
currentJob := job // De-referencing the temp variable used by the loop
2532
currentJobId := currentJob.ID()
2633
currentJobName := currentJob.Settings.Name

‎jobs/data_job_test.go‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,22 @@
11
package jobs
22

33
import (
4-
"github.com/databricks/terraform-provider-databricks/qa"
4+
"fmt"
55
"testing"
6+
7+
"github.com/databricks/terraform-provider-databricks/qa"
68
)
79

8-
func commonFixtures() []qa.HTTPFixture {
10+
func commonFixtures(name string) []qa.HTTPFixture {
11+
resource := "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0"
12+
if name != "" {
13+
resource = fmt.Sprintf("/api/2.1/jobs/list?expand_tasks=false&limit=25&name=%s&offset=0", name)
14+
}
915
return []qa.HTTPFixture{
1016
{
1117
Method: "GET",
12-
Resource: "/api/2.0/jobs/list",
13-
Response: JobList{
18+
Resource: resource,
19+
Response: JobListResponse{
1420
Jobs: []Job{
1521
{
1622
JobID: 123,
@@ -32,7 +38,7 @@ func commonFixtures() []qa.HTTPFixture {
3238
}
3339
func TestDataSourceQueryableJobMatchesId(t *testing.T) {
3440
qa.ResourceFixture{
35-
Fixtures: commonFixtures(),
41+
Fixtures: commonFixtures(""),
3642
Resource: DataSourceJob(),
3743
Read: true,
3844
New: true,
@@ -47,7 +53,7 @@ func TestDataSourceQueryableJobMatchesId(t *testing.T) {
4753

4854
func TestDataSourceQueryableJobMatchesName(t *testing.T) {
4955
qa.ResourceFixture{
50-
Fixtures: commonFixtures(),
56+
Fixtures: commonFixtures("First"),
5157
Resource: DataSourceJob(),
5258
Read: true,
5359
NonWritable: true,
@@ -61,7 +67,15 @@ func TestDataSourceQueryableJobMatchesName(t *testing.T) {
6167

6268
func TestDataSourceQueryableJobNoMatchName(t *testing.T) {
6369
qa.ResourceFixture{
64-
Fixtures: commonFixtures(),
70+
Fixtures: []qa.HTTPFixture{
71+
{
72+
Method: "GET",
73+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&name=Third&offset=0",
74+
Response: JobListResponse{
75+
Jobs: []Job{},
76+
},
77+
},
78+
},
6579
Resource: DataSourceJob(),
6680
Read: true,
6781
NonWritable: true,
@@ -72,7 +86,7 @@ func TestDataSourceQueryableJobNoMatchName(t *testing.T) {
7286

7387
func TestDataSourceQueryableJobNoMatchId(t *testing.T) {
7488
qa.ResourceFixture{
75-
Fixtures: commonFixtures(),
89+
Fixtures: commonFixtures(""),
7690
Resource: DataSourceJob(),
7791
Read: true,
7892
NonWritable: true,

‎jobs/data_jobs.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ func DataSourceJobs() *schema.Resource {
2020
return err
2121
}
2222
response.Ids = map[string]string{}
23-
for _, v := range list.Jobs {
23+
for _, v := range list {
2424
name := v.Settings.Name
2525
_, duplicateName := response.Ids[name]
2626
if duplicateName {

‎jobs/data_jobs_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ func TestJobsData(t *testing.T) {
1111
Fixtures: []qa.HTTPFixture{
1212
{
1313
Method: "GET",
14-
Resource: "/api/2.0/jobs/list",
15-
Response: JobList{
14+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
15+
Response: JobListResponse{
1616
Jobs: []Job{
1717
{
1818
JobID: 123,

‎jobs/resource_job.go‎

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -225,9 +225,10 @@ func (js *JobSettings) sortWebhooksByID() {
225225
js.WebhookNotifications.Sort()
226226
}
227227

228-
// JobList returns a list of all jobs
229-
type JobList struct {
230-
Jobs []Job `json:"jobs"`
228+
// JobListResponse returns a list of all jobs
229+
type JobListResponse struct {
230+
Jobs []Job `json:"jobs"`
231+
HasMore bool `json:"has_more,omitempty"`
231232
}
232233

233234
// Job contains the information when using a GET request from the Databricks Jobs api
@@ -307,9 +308,38 @@ type JobsAPI struct {
307308
context context.Context
308309
}
309310

311+
// List all jobs matching the name. If name is empty, returns all jobs
312+
func (a JobsAPI) ListByName(name string, expandTasks bool) ([]Job, error) {
313+
jobs := []Job{}
314+
params := map[string]interface{}{
315+
"limit": 25,
316+
"expand_tasks": expandTasks,
317+
}
318+
if name != "" {
319+
params["name"] = name
320+
}
321+
offset := 0
322+
323+
ctx := context.WithValue(a.context, common.Api, common.API_2_1)
324+
for {
325+
var resp JobListResponse
326+
params["offset"] = offset
327+
err := a.client.Get(ctx, "/jobs/list", params, &resp)
328+
if err != nil {
329+
return nil, err
330+
}
331+
jobs = append(jobs, resp.Jobs...)
332+
if !resp.HasMore {
333+
break
334+
}
335+
offset += len(resp.Jobs)
336+
}
337+
return jobs, nil
338+
}
339+
310340
// List all jobs
311-
func (a JobsAPI) List() (l JobList, err error) {
312-
err = a.client.Get(a.context, "/jobs/list", nil, &l)
341+
func (a JobsAPI) List() (l []Job, err error) {
342+
l, err = a.ListByName("", false)
313343
return
314344
}
315345

‎jobs/resource_job_test.go‎

Lines changed: 58 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1323,8 +1323,8 @@ func TestJobsAPIList(t *testing.T) {
13231323
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
13241324
{
13251325
Method: "GET",
1326-
Resource: "/api/2.0/jobs/list",
1327-
Response: JobList{
1326+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
1327+
Response: JobListResponse{
13281328
Jobs: []Job{
13291329
{
13301330
JobID: 1,
@@ -1336,7 +1336,62 @@ func TestJobsAPIList(t *testing.T) {
13361336
a := NewJobsAPI(ctx, client)
13371337
l, err := a.List()
13381338
require.NoError(t, err)
1339-
assert.Len(t, l.Jobs, 1)
1339+
assert.Len(t, l, 1)
1340+
})
1341+
}
1342+
1343+
func TestJobsAPIListMultiplePages(t *testing.T) {
1344+
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
1345+
{
1346+
Method: "GET",
1347+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=0",
1348+
Response: JobListResponse{
1349+
Jobs: []Job{
1350+
{
1351+
JobID: 1,
1352+
},
1353+
},
1354+
HasMore: true,
1355+
},
1356+
},
1357+
{
1358+
Method: "GET",
1359+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&offset=1",
1360+
Response: JobListResponse{
1361+
Jobs: []Job{
1362+
{
1363+
JobID: 2,
1364+
},
1365+
},
1366+
HasMore: false,
1367+
},
1368+
},
1369+
}, func(ctx context.Context, client *common.DatabricksClient) {
1370+
a := NewJobsAPI(ctx, client)
1371+
l, err := a.List()
1372+
require.NoError(t, err)
1373+
assert.Len(t, l, 2)
1374+
})
1375+
}
1376+
1377+
func TestJobsAPIListByName(t *testing.T) {
1378+
qa.HTTPFixturesApply(t, []qa.HTTPFixture{
1379+
{
1380+
Method: "GET",
1381+
Resource: "/api/2.1/jobs/list?expand_tasks=false&limit=25&name=test&offset=0",
1382+
Response: JobListResponse{
1383+
Jobs: []Job{
1384+
{
1385+
JobID: 1,
1386+
},
1387+
},
1388+
},
1389+
},
1390+
}, func(ctx context.Context, client *common.DatabricksClient) {
1391+
a := NewJobsAPI(ctx, client)
1392+
l, err := a.ListByName("test", false)
1393+
require.NoError(t, err)
1394+
assert.Len(t, l, 1)
13401395
})
13411396
}
13421397

0 commit comments

Comments
 (0)