From 8406a839916ce8c2163489c2a3c8b76014f526bd Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Fri, 3 Jan 2025 11:09:30 +0100 Subject: [PATCH 1/4] update --- .codegen/_openapi_sha | 2 +- .gitattributes | 2 + .../mocks/service/jobs/mock_jobs_interface.go | 59 ------------- service/apps/model.go | 12 +++ service/catalog/model.go | 6 +- service/jobs/api.go | 28 +++---- service/jobs/interface.go | 14 +++- service/jobs/model.go | 83 ++++++++++++++++--- service/oauth2/model.go | 28 +++++-- service/pipelines/model.go | 39 +++++++++ service/serving/api.go | 3 +- service/{oauth2 => serving}/ext_data_plane.go | 21 ++--- .../ext_data_plane_test.go | 51 ++++++------ service/serving/impl.go | 2 +- 14 files changed, 213 insertions(+), 137 deletions(-) rename service/{oauth2 => serving}/ext_data_plane.go (69%) rename service/{oauth2 => serving}/ext_data_plane_test.go (76%) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 8622b29ca..36a12857c 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -a6a317df8327c9b1e5cb59a03a42ffa2aabeef6d \ No newline at end of file +adfb709c0e78bfd6e62e53756cf4f5ec3f258e9a \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index da818989a..91ba6180a 100644 --- a/.gitattributes +++ b/.gitattributes @@ -77,9 +77,11 @@ experimental/mocks/service/marketplace/mock_provider_provider_analytics_dashboar experimental/mocks/service/marketplace/mock_provider_providers_interface.go linguist-generated=true experimental/mocks/service/ml/mock_experiments_interface.go linguist-generated=true experimental/mocks/service/ml/mock_model_registry_interface.go linguist-generated=true +experimental/mocks/service/oauth2/mock_account_federation_policy_interface.go linguist-generated=true experimental/mocks/service/oauth2/mock_custom_app_integration_interface.go linguist-generated=true experimental/mocks/service/oauth2/mock_o_auth_published_apps_interface.go linguist-generated=true experimental/mocks/service/oauth2/mock_published_app_integration_interface.go linguist-generated=true +experimental/mocks/service/oauth2/mock_service_principal_federation_policy_interface.go linguist-generated=true experimental/mocks/service/oauth2/mock_service_principal_secrets_interface.go linguist-generated=true experimental/mocks/service/pipelines/mock_pipelines_interface.go linguist-generated=true experimental/mocks/service/provisioning/mock_credentials_interface.go linguist-generated=true diff --git a/experimental/mocks/service/jobs/mock_jobs_interface.go b/experimental/mocks/service/jobs/mock_jobs_interface.go index c00b20329..1a353e910 100644 --- a/experimental/mocks/service/jobs/mock_jobs_interface.go +++ b/experimental/mocks/service/jobs/mock_jobs_interface.go @@ -753,65 +753,6 @@ func (_c *MockJobsInterface_Get_Call) RunAndReturn(run func(context.Context, job return _c } -// GetByJobId provides a mock function with given fields: ctx, jobId -func (_m *MockJobsInterface) GetByJobId(ctx context.Context, jobId int64) (*jobs.Job, error) { - ret := _m.Called(ctx, jobId) - - if len(ret) == 0 { - panic("no return value specified for GetByJobId") - } - - var r0 *jobs.Job - var r1 error - if rf, ok := ret.Get(0).(func(context.Context, int64) (*jobs.Job, error)); ok { - return rf(ctx, jobId) - } - if rf, ok := ret.Get(0).(func(context.Context, int64) *jobs.Job); ok { - r0 = rf(ctx, jobId) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(*jobs.Job) - } - } - - if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok { - r1 = rf(ctx, jobId) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// MockJobsInterface_GetByJobId_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetByJobId' -type MockJobsInterface_GetByJobId_Call struct { - *mock.Call -} - -// GetByJobId is a helper method to define mock.On call -// - ctx context.Context -// - jobId int64 -func (_e *MockJobsInterface_Expecter) GetByJobId(ctx interface{}, jobId interface{}) *MockJobsInterface_GetByJobId_Call { - return &MockJobsInterface_GetByJobId_Call{Call: _e.mock.On("GetByJobId", ctx, jobId)} -} - -func (_c *MockJobsInterface_GetByJobId_Call) Run(run func(ctx context.Context, jobId int64)) *MockJobsInterface_GetByJobId_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(int64)) - }) - return _c -} - -func (_c *MockJobsInterface_GetByJobId_Call) Return(_a0 *jobs.Job, _a1 error) *MockJobsInterface_GetByJobId_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockJobsInterface_GetByJobId_Call) RunAndReturn(run func(context.Context, int64) (*jobs.Job, error)) *MockJobsInterface_GetByJobId_Call { - _c.Call.Return(run) - return _c -} - // GetBySettingsName provides a mock function with given fields: ctx, name func (_m *MockJobsInterface) GetBySettingsName(ctx context.Context, name string) (*jobs.BaseJob, error) { ret := _m.Called(ctx, name) diff --git a/service/apps/model.go b/service/apps/model.go index 7eaa5f90b..0bced9bce 100755 --- a/service/apps/model.go +++ b/service/apps/model.go @@ -605,6 +605,18 @@ type CreateAppDeploymentRequest struct { // Create an app type CreateAppRequest struct { App *App `json:"app,omitempty"` + // If true, the app will not be started after creation. + NoCompute bool `json:"-" url:"no_compute,omitempty"` + + ForceSendFields []string `json:"-"` +} + +func (s *CreateAppRequest) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, s) +} + +func (s CreateAppRequest) MarshalJSON() ([]byte, error) { + return marshal.Marshal(s) } // Delete an app diff --git a/service/catalog/model.go b/service/catalog/model.go index fce115ed1..4493e9b3e 100755 --- a/service/catalog/model.go +++ b/service/catalog/model.go @@ -4657,6 +4657,8 @@ type ProvisioningInfoState string const ProvisioningInfoStateActive ProvisioningInfoState = `ACTIVE` +const ProvisioningInfoStateDegraded ProvisioningInfoState = `DEGRADED` + const ProvisioningInfoStateDeleting ProvisioningInfoState = `DELETING` const ProvisioningInfoStateFailed ProvisioningInfoState = `FAILED` @@ -4673,11 +4675,11 @@ func (f *ProvisioningInfoState) String() string { // Set raw string value and validate it against allowed values func (f *ProvisioningInfoState) Set(v string) error { switch v { - case `ACTIVE`, `DELETING`, `FAILED`, `PROVISIONING`, `UPDATING`: + case `ACTIVE`, `DEGRADED`, `DELETING`, `FAILED`, `PROVISIONING`, `UPDATING`: *f = ProvisioningInfoState(v) return nil default: - return fmt.Errorf(`value "%s" is not one of "ACTIVE", "DELETING", "FAILED", "PROVISIONING", "UPDATING"`, v) + return fmt.Errorf(`value "%s" is not one of "ACTIVE", "DEGRADED", "DELETING", "FAILED", "PROVISIONING", "UPDATING"`, v) } } diff --git a/service/jobs/api.go b/service/jobs/api.go index b34dbfec2..c8f3ed241 100755 --- a/service/jobs/api.go +++ b/service/jobs/api.go @@ -81,12 +81,13 @@ type JobsInterface interface { // Get a single job. // // Retrieves the details for a single job. - Get(ctx context.Context, request GetJobRequest) (*Job, error) - - // Get a single job. // - // Retrieves the details for a single job. - GetByJobId(ctx context.Context, jobId int64) (*Job, error) + // In Jobs API 2.2, requests for a single job support pagination of `tasks` and + // `job_clusters` when either exceeds 100 elements. Use the `next_page_token` + // field to check for more results and pass its value as the `page_token` in + // subsequent requests. Arrays with fewer than 100 elements in a page will be + // empty on later pages. + Get(ctx context.Context, request GetJobRequest) (*Job, error) // Get job permission levels. // @@ -112,7 +113,13 @@ type JobsInterface interface { // Get a single job run. // - // Retrieve the metadata of a run. + // Retrieves the metadata of a run. + // + // In Jobs API 2.2, requests for a single job run support pagination of `tasks` + // and `job_clusters` when either exceeds 100 elements. Use the + // `next_page_token` field to check for more results and pass its value as the + // `page_token` in subsequent requests. Arrays with fewer than 100 elements in a + // page will be empty on later pages. GetRun(ctx context.Context, request GetRunRequest) (*Run, error) // Get the output for a single run. @@ -427,15 +434,6 @@ func (a *JobsAPI) DeleteRunByRunId(ctx context.Context, runId int64) error { }) } -// Get a single job. -// -// Retrieves the details for a single job. -func (a *JobsAPI) GetByJobId(ctx context.Context, jobId int64) (*Job, error) { - return a.jobsImpl.Get(ctx, GetJobRequest{ - JobId: jobId, - }) -} - // Get job permission levels. // // Gets the permission levels that a user can have on an object. diff --git a/service/jobs/interface.go b/service/jobs/interface.go index b1ce89b9e..be9b5518d 100755 --- a/service/jobs/interface.go +++ b/service/jobs/interface.go @@ -61,6 +61,12 @@ type JobsService interface { // Get a single job. // // Retrieves the details for a single job. + // + // In Jobs API 2.2, requests for a single job support pagination of `tasks` + // and `job_clusters` when either exceeds 100 elements. Use the + // `next_page_token` field to check for more results and pass its value as + // the `page_token` in subsequent requests. Arrays with fewer than 100 + // elements in a page will be empty on later pages. Get(ctx context.Context, request GetJobRequest) (*Job, error) // Get job permission levels. @@ -76,7 +82,13 @@ type JobsService interface { // Get a single job run. // - // Retrieve the metadata of a run. + // Retrieves the metadata of a run. + // + // In Jobs API 2.2, requests for a single job run support pagination of + // `tasks` and `job_clusters` when either exceeds 100 elements. Use the + // `next_page_token` field to check for more results and pass its value as + // the `page_token` in subsequent requests. Arrays with fewer than 100 + // elements in a page will be empty on later pages. GetRun(ctx context.Context, request GetRunRequest) (*Run, error) // Get the output for a single run. diff --git a/service/jobs/model.go b/service/jobs/model.go index 641d98afa..e23efe94c 100755 --- a/service/jobs/model.go +++ b/service/jobs/model.go @@ -23,6 +23,11 @@ type BaseJob struct { // based on accessible budget policies of the run_as identity on job // creation or modification. EffectiveBudgetPolicyId string `json:"effective_budget_policy_id,omitempty"` + // Indicates if the job has more sub-resources (`tasks`, `job_clusters`) + // that are not shown. They can be accessed via :method:jobs/get endpoint. + // It is only relevant for API 2.2 :method:jobs/list requests with + // `expand_tasks=true`. + HasMore bool `json:"has_more,omitempty"` // The canonical identifier for this job. JobId int64 `json:"job_id,omitempty"` // Settings for this job and all of its runs. These settings can be updated @@ -89,9 +94,16 @@ type BaseRun struct { // Note: dbt and SQL File tasks support only version-controlled sources. If // dbt or SQL File tasks are used, `git_source` must be defined on the job. GitSource *GitSource `json:"git_source,omitempty"` + // Indicates if the run has more sub-resources (`tasks`, `job_clusters`) + // that are not shown. They can be accessed via :method:jobs/getrun + // endpoint. It is only relevant for API 2.2 :method:jobs/listruns requests + // with `expand_tasks=true`. + HasMore bool `json:"has_more,omitempty"` // A list of job cluster specifications that can be shared and reused by // tasks of this job. Libraries cannot be declared in a shared job cluster. - // You must declare dependent libraries in task settings. + // You must declare dependent libraries in task settings. If more than 100 + // job clusters are available, you can paginate through them using + // :method:jobs/getrun. JobClusters []JobCluster `json:"job_clusters,omitempty"` // The canonical identifier of the job that contains this run. JobId int64 `json:"job_id,omitempty"` @@ -153,7 +165,10 @@ type BaseRun struct { // The current status of the run Status *RunStatus `json:"status,omitempty"` // The list of tasks performed by the run. Each task has its own `run_id` - // which you can use to call `JobsGetOutput` to retrieve the run resutls. + // which you can use to call `JobsGetOutput` to retrieve the run resutls. If + // more than 100 tasks are available, you can paginate through them using + // :method:jobs/getrun. Use the `next_page_token` field at the object root + // to determine if more results are available. Tasks []RunTask `json:"tasks,omitempty"` // The type of trigger that fired this run. // @@ -543,7 +558,9 @@ type CreateJob struct { Health *JobsHealthRules `json:"health,omitempty"` // A list of job cluster specifications that can be shared and reused by // tasks of this job. Libraries cannot be declared in a shared job cluster. - // You must declare dependent libraries in task settings. + // You must declare dependent libraries in task settings. If more than 100 + // job clusters are available, you can paginate through them using + // :method:jobs/get. JobClusters []JobCluster `json:"job_clusters,omitempty"` // An optional maximum allowed number of concurrent runs of the job. Set // this value if you want to be able to execute multiple runs of the same @@ -583,7 +600,10 @@ type CreateJob struct { // limitations as cluster tags. A maximum of 25 tags can be added to the // job. Tags map[string]string `json:"tags,omitempty"` - // A list of task specifications to be executed by this job. + // A list of task specifications to be executed by this job. If more than + // 100 tasks are available, you can paginate through them using + // :method:jobs/get. Use the `next_page_token` field at the object root to + // determine if more results are available. Tasks []Task `json:"tasks,omitempty"` // An optional timeout applied to each run of this job. A value of `0` means // no timeout. @@ -961,6 +981,19 @@ type GetJobRequest struct { // The canonical identifier of the job to retrieve information about. This // field is required. JobId int64 `json:"-" url:"job_id"` + // Use `next_page_token` returned from the previous GetJob to request the + // next page of the job's sub-resources. + PageToken string `json:"-" url:"page_token,omitempty"` + + ForceSendFields []string `json:"-"` +} + +func (s *GetJobRequest) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, s) +} + +func (s GetJobRequest) MarshalJSON() ([]byte, error) { + return marshal.Marshal(s) } // Get job policy compliance @@ -1005,8 +1038,8 @@ type GetRunRequest struct { IncludeHistory bool `json:"-" url:"include_history,omitempty"` // Whether to include resolved parameter values in the response. IncludeResolvedValues bool `json:"-" url:"include_resolved_values,omitempty"` - // To list the next page of job tasks, set this field to the value of the - // `next_page_token` returned in the GetJob response. + // Use `next_page_token` returned from the previous GetRun to request the + // next page of the run's sub-resources. PageToken string `json:"-" url:"page_token,omitempty"` // The canonical identifier of the run for which to retrieve the metadata. // This field is required. @@ -1139,8 +1172,15 @@ type Job struct { // based on accessible budget policies of the run_as identity on job // creation or modification. EffectiveBudgetPolicyId string `json:"effective_budget_policy_id,omitempty"` + // Indicates if the job has more sub-resources (`tasks`, `job_clusters`) + // that are not shown. They can be accessed via :method:jobs/get endpoint. + // It is only relevant for API 2.2 :method:jobs/list requests with + // `expand_tasks=true`. + HasMore bool `json:"has_more,omitempty"` // The canonical identifier for this job. JobId int64 `json:"job_id,omitempty"` + // A token that can be used to list the next page of sub-resources. + NextPageToken string `json:"next_page_token,omitempty"` // The email of an active workspace user or the application ID of a service // principal that the job runs as. This value can be changed by setting the // `run_as` field when creating or updating a job. @@ -1583,7 +1623,9 @@ type JobSettings struct { Health *JobsHealthRules `json:"health,omitempty"` // A list of job cluster specifications that can be shared and reused by // tasks of this job. Libraries cannot be declared in a shared job cluster. - // You must declare dependent libraries in task settings. + // You must declare dependent libraries in task settings. If more than 100 + // job clusters are available, you can paginate through them using + // :method:jobs/get. JobClusters []JobCluster `json:"job_clusters,omitempty"` // An optional maximum allowed number of concurrent runs of the job. Set // this value if you want to be able to execute multiple runs of the same @@ -1623,7 +1665,10 @@ type JobSettings struct { // limitations as cluster tags. A maximum of 25 tags can be added to the // job. Tags map[string]string `json:"tags,omitempty"` - // A list of task specifications to be executed by this job. + // A list of task specifications to be executed by this job. If more than + // 100 tasks are available, you can paginate through them using + // :method:jobs/get. Use the `next_page_token` field at the object root to + // determine if more results are available. Tasks []Task `json:"tasks,omitempty"` // An optional timeout applied to each run of this job. A value of `0` means // no timeout. @@ -1862,7 +1907,9 @@ func (s ListJobComplianceRequest) MarshalJSON() ([]byte, error) { // List jobs type ListJobsRequest struct { - // Whether to include task and cluster details in the response. + // Whether to include task and cluster details in the response. Note that in + // API 2.2, only the first 100 elements will be shown. Use :method:jobs/get + // to paginate through all tasks and clusters. ExpandTasks bool `json:"-" url:"expand_tasks,omitempty"` // The number of jobs to return. This value must be greater than 0 and less // or equal to 100. The default value is 20. @@ -1924,7 +1971,9 @@ type ListRunsRequest struct { // results; otherwise, lists both active and completed runs. This field // cannot be `true` when active_only is `true`. CompletedOnly bool `json:"-" url:"completed_only,omitempty"` - // Whether to include task and cluster details in the response. + // Whether to include task and cluster details in the response. Note that in + // API 2.2, only the first 100 elements will be shown. Use + // :method:jobs/getrun to paginate through all tasks and clusters. ExpandTasks bool `json:"-" url:"expand_tasks,omitempty"` // The job for which to list runs. If omitted, the Jobs service lists runs // from all jobs. @@ -2557,12 +2606,19 @@ type Run struct { // Note: dbt and SQL File tasks support only version-controlled sources. If // dbt or SQL File tasks are used, `git_source` must be defined on the job. GitSource *GitSource `json:"git_source,omitempty"` + // Indicates if the run has more sub-resources (`tasks`, `job_clusters`) + // that are not shown. They can be accessed via :method:jobs/getrun + // endpoint. It is only relevant for API 2.2 :method:jobs/listruns requests + // with `expand_tasks=true`. + HasMore bool `json:"has_more,omitempty"` // Only populated by for-each iterations. The parent for-each task is // located in tasks array. Iterations []RunTask `json:"iterations,omitempty"` // A list of job cluster specifications that can be shared and reused by // tasks of this job. Libraries cannot be declared in a shared job cluster. - // You must declare dependent libraries in task settings. + // You must declare dependent libraries in task settings. If more than 100 + // job clusters are available, you can paginate through them using + // :method:jobs/getrun. JobClusters []JobCluster `json:"job_clusters,omitempty"` // The canonical identifier of the job that contains this run. JobId int64 `json:"job_id,omitempty"` @@ -2626,7 +2682,10 @@ type Run struct { // The current status of the run Status *RunStatus `json:"status,omitempty"` // The list of tasks performed by the run. Each task has its own `run_id` - // which you can use to call `JobsGetOutput` to retrieve the run resutls. + // which you can use to call `JobsGetOutput` to retrieve the run resutls. If + // more than 100 tasks are available, you can paginate through them using + // :method:jobs/getrun. Use the `next_page_token` field at the object root + // to determine if more results are available. Tasks []RunTask `json:"tasks,omitempty"` // The type of trigger that fired this run. // diff --git a/service/oauth2/model.go b/service/oauth2/model.go index 5ce0f66d5..8c4bc6c6d 100755 --- a/service/oauth2/model.go +++ b/service/oauth2/model.go @@ -9,8 +9,9 @@ import ( // Create account federation policy type CreateAccountFederationPolicyRequest struct { Policy *FederationPolicy `json:"policy,omitempty"` - // The identifier for the federation policy. If unspecified, the id will be - // assigned by Databricks. + // The identifier for the federation policy. The identifier must contain + // only lowercase alphanumeric characters, numbers, hyphens, and slashes. If + // unspecified, the id will be assigned by Databricks. PolicyId string `json:"-" url:"policy_id,omitempty"` ForceSendFields []string `json:"-"` @@ -105,8 +106,9 @@ func (s CreatePublishedAppIntegrationOutput) MarshalJSON() ([]byte, error) { // Create service principal federation policy type CreateServicePrincipalFederationPolicyRequest struct { Policy *FederationPolicy `json:"policy,omitempty"` - // The identifier for the federation policy. If unspecified, the id will be - // assigned by Databricks. + // The identifier for the federation policy. The identifier must contain + // only lowercase alphanumeric characters, numbers, hyphens, and slashes. If + // unspecified, the id will be assigned by Databricks. PolicyId string `json:"-" url:"policy_id,omitempty"` // The service principal id for the federation policy. ServicePrincipalId int64 `json:"-" url:"-"` @@ -172,6 +174,7 @@ func (s DataPlaneInfo) MarshalJSON() ([]byte, error) { // Delete account federation policy type DeleteAccountFederationPolicyRequest struct { + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` } @@ -196,6 +199,7 @@ type DeleteResponse struct { // Delete service principal federation policy type DeleteServicePrincipalFederationPolicyRequest struct { + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` // The service principal id for the federation policy. ServicePrincipalId int64 `json:"-" url:"-"` @@ -214,9 +218,13 @@ type FederationPolicy struct { CreateTime string `json:"create_time,omitempty"` // Description of the federation policy. Description string `json:"description,omitempty"` - // Name of the federation policy. The name must contain only lowercase - // alphanumeric characters, numbers, and hyphens. It must be unique within - // the account. + // Resource name for the federation policy. Example values include + // `accounts//federationPolicies/my-federation-policy` for + // Account Federation Policies, and + // `accounts//servicePrincipals//federationPolicies/my-federation-policy` + // for Service Principal Federation Policies. Typically an output parameter, + // which does not need to be specified in create or update requests. If + // specified in a request, must match the value in the request URL. Name string `json:"name,omitempty"` // Specifies the policy to use for validating OIDC claims in your federated // tokens. @@ -239,6 +247,7 @@ func (s FederationPolicy) MarshalJSON() ([]byte, error) { // Get account federation policy type GetAccountFederationPolicyRequest struct { + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` } @@ -366,6 +375,7 @@ func (s GetPublishedAppsOutput) MarshalJSON() ([]byte, error) { // Get service principal federation policy type GetServicePrincipalFederationPolicyRequest struct { + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` // The service principal id for the federation policy. ServicePrincipalId int64 `json:"-" url:"-"` @@ -630,7 +640,7 @@ func (s TokenAccessPolicy) MarshalJSON() ([]byte, error) { // Update account federation policy type UpdateAccountFederationPolicyRequest struct { Policy *FederationPolicy `json:"policy,omitempty"` - + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` // Field mask is required to be passed into the PATCH request. Field mask // specifies which fields of the setting payload will be updated. The field @@ -663,7 +673,7 @@ type UpdatePublishedAppIntegrationOutput struct { // Update service principal federation policy type UpdateServicePrincipalFederationPolicyRequest struct { Policy *FederationPolicy `json:"policy,omitempty"` - + // The identifier for the federation policy. PolicyId string `json:"-" url:"-"` // The service principal id for the federation policy. ServicePrincipalId int64 `json:"-" url:"-"` diff --git a/service/pipelines/model.go b/service/pipelines/model.go index 3a67b316e..04f110841 100755 --- a/service/pipelines/model.go +++ b/service/pipelines/model.go @@ -56,6 +56,13 @@ type CreatePipeline struct { Photon bool `json:"photon,omitempty"` // Restart window of this pipeline. RestartWindow *RestartWindow `json:"restart_window,omitempty"` + // Write-only setting, available only in Create/Update calls. Specifies the + // user or service principal that the pipeline runs as. If not specified, + // the pipeline runs as the user who created the pipeline. + // + // Only `user_name` or `service_principal_name` can be specified. If both + // are specified, an error is thrown. + RunAs *RunAs `json:"run_as,omitempty"` // The default schema (database) where tables are read from or published to. // The presence of this field implies that the pipeline is in direct // publishing mode. @@ -219,6 +226,13 @@ type EditPipeline struct { PipelineId string `json:"pipeline_id,omitempty" url:"-"` // Restart window of this pipeline. RestartWindow *RestartWindow `json:"restart_window,omitempty"` + // Write-only setting, available only in Create/Update calls. Specifies the + // user or service principal that the pipeline runs as. If not specified, + // the pipeline runs as the user who created the pipeline. + // + // Only `user_name` or `service_principal_name` can be specified. If both + // are specified, an error is thrown. + RunAs *RunAs `json:"run_as,omitempty"` // The default schema (database) where tables are read from or published to. // The presence of this field implies that the pipeline is in direct // publishing mode. @@ -1376,6 +1390,31 @@ func (f *RestartWindowDaysOfWeek) Type() string { return "RestartWindowDaysOfWeek" } +// Write-only setting, available only in Create/Update calls. Specifies the user +// or service principal that the pipeline runs as. If not specified, the +// pipeline runs as the user who created the pipeline. +// +// Only `user_name` or `service_principal_name` can be specified. If both are +// specified, an error is thrown. +type RunAs struct { + // Application ID of an active service principal. Setting this field + // requires the `servicePrincipal/user` role. + ServicePrincipalName string `json:"service_principal_name,omitempty"` + // The email of an active workspace user. Users can only set this field to + // their own email. + UserName string `json:"user_name,omitempty"` + + ForceSendFields []string `json:"-"` +} + +func (s *RunAs) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, s) +} + +func (s RunAs) MarshalJSON() ([]byte, error) { + return marshal.Marshal(s) +} + type SchemaSpec struct { // Required. Destination catalog to store tables. DestinationCatalog string `json:"destination_catalog,omitempty"` diff --git a/service/serving/api.go b/service/serving/api.go index cddbdd737..5a2606d6b 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -11,7 +11,6 @@ import ( "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/retries" - "github.com/databricks/databricks-sdk-go/service/oauth2" "github.com/databricks/databricks-sdk-go/useragent" ) @@ -468,7 +467,7 @@ func NewServingEndpointsDataPlane(client *client.DatabricksClient, return &ServingEndpointsDataPlaneAPI{ servingEndpointsDataPlaneImpl: servingEndpointsDataPlaneImpl{ client: client, - dataPlaneService: oauth2.NewDataPlaneService(), + dataPlaneService: NewDataPlaneService(), controlPlane: controlPlane, }, } diff --git a/service/oauth2/ext_data_plane.go b/service/serving/ext_data_plane.go similarity index 69% rename from service/oauth2/ext_data_plane.go rename to service/serving/ext_data_plane.go index aa9eeddcb..c4ec977b4 100644 --- a/service/oauth2/ext_data_plane.go +++ b/service/serving/ext_data_plane.go @@ -1,27 +1,28 @@ -package oauth2 +package serving import ( "strings" "sync" - "golang.org/x/oauth2" + "github.com/databricks/databricks-sdk-go/service/oauth2" + goauth "golang.org/x/oauth2" ) // DataPlaneService is an interface for services that access DataPlane. type DataPlaneService interface { - GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *oauth2.Token, error) + GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error) } func NewDataPlaneService() DataPlaneService { return &dataPlaneServiceImpl{ - infos: make(map[string]*DataPlaneInfo), - tokens: make(map[string]*oauth2.Token), + infos: make(map[string]*oauth2.DataPlaneInfo), + tokens: make(map[string]*goauth.Token), } } type dataPlaneServiceImpl struct { - infos map[string]*DataPlaneInfo - tokens map[string]*oauth2.Token + infos map[string]*oauth2.DataPlaneInfo + tokens map[string]*goauth.Token // This class can be shared across multiple threads. // This mutex is used to synchronize access to the infos and tokens maps. mu sync.Mutex @@ -29,7 +30,7 @@ type dataPlaneServiceImpl struct { // GetDataPlaneDetails returns the endpoint URL and token. It returns a cached token if it is valid, // otherwise it refreshes the token and returns the new token. -func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *oauth2.Token, error) { +func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error) { key := o.generateKey(method, params) info, err := o.getInfo(key, infoGetter) if err != nil { @@ -42,7 +43,7 @@ func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []strin return info.EndpointUrl, token, nil } -func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*DataPlaneInfo, error)) (*DataPlaneInfo, error) { +func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*oauth2.DataPlaneInfo, error)) (*oauth2.DataPlaneInfo, error) { info, infoOk := o.infos[key] if !infoOk { o.mu.Lock() @@ -60,7 +61,7 @@ func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*DataPlane return info, nil } -func (o *dataPlaneServiceImpl) refreshToken(key string, info *DataPlaneInfo, refresh func(*DataPlaneInfo) (*oauth2.Token, error)) (*oauth2.Token, error) { +func (o *dataPlaneServiceImpl) refreshToken(key string, info *oauth2.DataPlaneInfo, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error)) (*goauth.Token, error) { token, tokenOk := o.tokens[key] if !tokenOk || !token.Valid() { o.mu.Lock() diff --git a/service/oauth2/ext_data_plane_test.go b/service/serving/ext_data_plane_test.go similarity index 76% rename from service/oauth2/ext_data_plane_test.go rename to service/serving/ext_data_plane_test.go index ca622a456..4fa3ca3dd 100644 --- a/service/oauth2/ext_data_plane_test.go +++ b/service/serving/ext_data_plane_test.go @@ -1,32 +1,33 @@ -package oauth2 +package serving import ( "testing" "time" + "github.com/databricks/databricks-sdk-go/service/oauth2" "github.com/stretchr/testify/assert" - "golang.org/x/oauth2" + goauth "golang.org/x/oauth2" ) type infoMock struct { called bool - info *DataPlaneInfo + info *oauth2.DataPlaneInfo err error } -func (i *infoMock) DataPlaneInfoGetter() (*DataPlaneInfo, error) { +func (i *infoMock) DataPlaneInfoGetter() (*oauth2.DataPlaneInfo, error) { i.called = true return i.info, i.err } type tokenRefreshSpy struct { called bool - expectedInfo *DataPlaneInfo - token *oauth2.Token + expectedInfo *oauth2.DataPlaneInfo + token *goauth.Token err error } -func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*oauth2.Token, error) { +func (t *tokenRefreshSpy) TokenRefresh(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { t.expectedInfo = info t.called = true return t.token, t.err @@ -34,14 +35,14 @@ func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*oauth2.Token, erro func TestTokenNotCached(t *testing.T) { info := infoMock{ - info: &DataPlaneInfo{ + info: &oauth2.DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, err: nil, } s := tokenRefreshSpy{ - token: &oauth2.Token{ + token: &goauth.Token{ AccessToken: "token", TokenType: "type", Expiry: time.Now().Add(1 * time.Hour), @@ -49,8 +50,8 @@ func TestTokenNotCached(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{ - infos: make(map[string]*DataPlaneInfo), - tokens: make(map[string]*oauth2.Token), + infos: make(map[string]*oauth2.DataPlaneInfo), + tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) assert.NoError(t, err) @@ -64,14 +65,14 @@ func TestTokenNotCached(t *testing.T) { func TestTokenCached(t *testing.T) { info := infoMock{ - info: &DataPlaneInfo{ + info: &oauth2.DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, err: nil, } s := tokenRefreshSpy{ - token: &oauth2.Token{ + token: &goauth.Token{ AccessToken: "token", TokenType: "type", Expiry: time.Now().Add(1 * time.Hour), @@ -79,8 +80,8 @@ func TestTokenCached(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{} - c.infos = make(map[string]*DataPlaneInfo) - c.tokens = make(map[string]*oauth2.Token) + c.infos = make(map[string]*oauth2.DataPlaneInfo) + c.tokens = make(map[string]*goauth.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = s.token url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) @@ -95,20 +96,20 @@ func TestTokenCached(t *testing.T) { func TestTokenExpired(t *testing.T) { info := infoMock{ - info: &DataPlaneInfo{ + info: &oauth2.DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, err: nil, } - expired := &oauth2.Token{ + expired := &goauth.Token{ AccessToken: "oldToken", TokenType: "type", Expiry: time.Now().Add(-1 * time.Hour), } s := tokenRefreshSpy{ - token: &oauth2.Token{ + token: &goauth.Token{ AccessToken: "token", TokenType: "type", Expiry: time.Now().Add(1 * time.Hour), @@ -116,8 +117,8 @@ func TestTokenExpired(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{} - c.infos = make(map[string]*DataPlaneInfo) - c.tokens = make(map[string]*oauth2.Token) + c.infos = make(map[string]*oauth2.DataPlaneInfo) + c.tokens = make(map[string]*goauth.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = expired url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) @@ -137,8 +138,8 @@ func TestTokenInfoError(t *testing.T) { } s := tokenRefreshSpy{} c := dataPlaneServiceImpl{ - infos: make(map[string]*DataPlaneInfo), - tokens: make(map[string]*oauth2.Token), + infos: make(map[string]*oauth2.DataPlaneInfo), + tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) assert.ErrorIs(t, err, assert.AnError) @@ -150,7 +151,7 @@ func TestTokenInfoError(t *testing.T) { func TestTokenRefreshError(t *testing.T) { info := infoMock{ - info: &DataPlaneInfo{ + info: &oauth2.DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -161,8 +162,8 @@ func TestTokenRefreshError(t *testing.T) { err: assert.AnError, } c := dataPlaneServiceImpl{ - infos: make(map[string]*DataPlaneInfo), - tokens: make(map[string]*oauth2.Token), + infos: make(map[string]*oauth2.DataPlaneInfo), + tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) assert.ErrorIs(t, err, assert.AnError) diff --git a/service/serving/impl.go b/service/serving/impl.go index 0b211b214..1261c1b7b 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -183,7 +183,7 @@ func (a *servingEndpointsImpl) UpdatePermissions(ctx context.Context, request Se // unexported type that holds implementations of just ServingEndpointsDataPlane API methods type servingEndpointsDataPlaneImpl struct { - dataPlaneService oauth2.DataPlaneService + dataPlaneService DataPlaneService controlPlane *ServingEndpointsAPI client *client.DatabricksClient } From d4383aeaa3dcb6f71d309060499f79bf3528aedc Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Fri, 3 Jan 2025 16:45:35 +0100 Subject: [PATCH 2/4] update --- .codegen/_openapi_sha | 2 +- .../mocks/service/jobs/mock_jobs_interface.go | 59 +++++++++++++++++++ service/jobs/api.go | 26 ++++++++ service/pkg.go | 6 +- 4 files changed, 89 insertions(+), 4 deletions(-) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 36a12857c..784f606d8 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -adfb709c0e78bfd6e62e53756cf4f5ec3f258e9a \ No newline at end of file +821ab6511c9a977bae10cdcd2cf5aa0ad6ab8f8f \ No newline at end of file diff --git a/experimental/mocks/service/jobs/mock_jobs_interface.go b/experimental/mocks/service/jobs/mock_jobs_interface.go index 1a353e910..c00b20329 100644 --- a/experimental/mocks/service/jobs/mock_jobs_interface.go +++ b/experimental/mocks/service/jobs/mock_jobs_interface.go @@ -753,6 +753,65 @@ func (_c *MockJobsInterface_Get_Call) RunAndReturn(run func(context.Context, job return _c } +// GetByJobId provides a mock function with given fields: ctx, jobId +func (_m *MockJobsInterface) GetByJobId(ctx context.Context, jobId int64) (*jobs.Job, error) { + ret := _m.Called(ctx, jobId) + + if len(ret) == 0 { + panic("no return value specified for GetByJobId") + } + + var r0 *jobs.Job + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, int64) (*jobs.Job, error)); ok { + return rf(ctx, jobId) + } + if rf, ok := ret.Get(0).(func(context.Context, int64) *jobs.Job); ok { + r0 = rf(ctx, jobId) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*jobs.Job) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok { + r1 = rf(ctx, jobId) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockJobsInterface_GetByJobId_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetByJobId' +type MockJobsInterface_GetByJobId_Call struct { + *mock.Call +} + +// GetByJobId is a helper method to define mock.On call +// - ctx context.Context +// - jobId int64 +func (_e *MockJobsInterface_Expecter) GetByJobId(ctx interface{}, jobId interface{}) *MockJobsInterface_GetByJobId_Call { + return &MockJobsInterface_GetByJobId_Call{Call: _e.mock.On("GetByJobId", ctx, jobId)} +} + +func (_c *MockJobsInterface_GetByJobId_Call) Run(run func(ctx context.Context, jobId int64)) *MockJobsInterface_GetByJobId_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(int64)) + }) + return _c +} + +func (_c *MockJobsInterface_GetByJobId_Call) Return(_a0 *jobs.Job, _a1 error) *MockJobsInterface_GetByJobId_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockJobsInterface_GetByJobId_Call) RunAndReturn(run func(context.Context, int64) (*jobs.Job, error)) *MockJobsInterface_GetByJobId_Call { + _c.Call.Return(run) + return _c +} + // GetBySettingsName provides a mock function with given fields: ctx, name func (_m *MockJobsInterface) GetBySettingsName(ctx context.Context, name string) (*jobs.BaseJob, error) { ret := _m.Called(ctx, name) diff --git a/service/jobs/api.go b/service/jobs/api.go index c8f3ed241..fb4f0ef56 100755 --- a/service/jobs/api.go +++ b/service/jobs/api.go @@ -89,6 +89,17 @@ type JobsInterface interface { // empty on later pages. Get(ctx context.Context, request GetJobRequest) (*Job, error) + // Get a single job. + // + // Retrieves the details for a single job. + // + // In Jobs API 2.2, requests for a single job support pagination of `tasks` and + // `job_clusters` when either exceeds 100 elements. Use the `next_page_token` + // field to check for more results and pass its value as the `page_token` in + // subsequent requests. Arrays with fewer than 100 elements in a page will be + // empty on later pages. + GetByJobId(ctx context.Context, jobId int64) (*Job, error) + // Get job permission levels. // // Gets the permission levels that a user can have on an object. @@ -434,6 +445,21 @@ func (a *JobsAPI) DeleteRunByRunId(ctx context.Context, runId int64) error { }) } +// Get a single job. +// +// Retrieves the details for a single job. +// +// In Jobs API 2.2, requests for a single job support pagination of `tasks` and +// `job_clusters` when either exceeds 100 elements. Use the `next_page_token` +// field to check for more results and pass its value as the `page_token` in +// subsequent requests. Arrays with fewer than 100 elements in a page will be +// empty on later pages. +func (a *JobsAPI) GetByJobId(ctx context.Context, jobId int64) (*Job, error) { + return a.jobsImpl.Get(ctx, GetJobRequest{ + JobId: jobId, + }) +} + // Get job permission levels. // // Gets the permission levels that a user can have on an object. diff --git a/service/pkg.go b/service/pkg.go index bd49b7d6b..b891a9547 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -48,10 +48,10 @@ // // - [marketplace.ConsumerProvidersAPI]: Providers are the entities that publish listings to the Marketplace. // -// - [catalog.CredentialsAPI]: A credential represents an authentication and authorization mechanism for accessing services on your cloud tenant. -// // - [provisioning.CredentialsAPI]: These APIs manage credential configurations for this workspace. // +// - [catalog.CredentialsAPI]: A credential represents an authentication and authorization mechanism for accessing services on your cloud tenant. +// // - [settings.CredentialsManagerAPI]: Credentials manager interacts with with Identity Providers to to perform token exchanges using stored credentials and refresh tokens. // // - [settings.CspEnablementAccountAPI]: The compliance security profile settings at the account level control whether to enable it for new workspaces. @@ -323,8 +323,8 @@ var ( _ *marketplace.ConsumerListingsAPI = nil _ *marketplace.ConsumerPersonalizationRequestsAPI = nil _ *marketplace.ConsumerProvidersAPI = nil - _ *catalog.CredentialsAPI = nil _ *provisioning.CredentialsAPI = nil + _ *catalog.CredentialsAPI = nil _ *settings.CredentialsManagerAPI = nil _ *settings.CspEnablementAccountAPI = nil _ *iam.CurrentUserAPI = nil From 5c3d7a45bd6899dbf94e4d4a4ea51f7998789bfa Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Mon, 6 Jan 2025 15:44:49 +0100 Subject: [PATCH 3/4] update --- .codegen/_openapi_sha | 2 +- service/oauth2/model.go | 17 ---------------- service/serving/ext_data_plane.go | 13 ++++++------- service/serving/ext_data_plane_test.go | 27 +++++++++++++------------- service/serving/impl.go | 6 ++---- service/serving/model.go | 20 +++++++++++++++++-- 6 files changed, 40 insertions(+), 45 deletions(-) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 784f606d8..3cf20d57b 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -821ab6511c9a977bae10cdcd2cf5aa0ad6ab8f8f \ No newline at end of file +- \ No newline at end of file diff --git a/service/oauth2/model.go b/service/oauth2/model.go index 8c4bc6c6d..08de08f93 100755 --- a/service/oauth2/model.go +++ b/service/oauth2/model.go @@ -155,23 +155,6 @@ func (s CreateServicePrincipalSecretResponse) MarshalJSON() ([]byte, error) { return marshal.Marshal(s) } -type DataPlaneInfo struct { - // Authorization details as a string. - AuthorizationDetails string `json:"authorization_details,omitempty"` - // The URL of the endpoint for this operation in the dataplane. - EndpointUrl string `json:"endpoint_url,omitempty"` - - ForceSendFields []string `json:"-"` -} - -func (s *DataPlaneInfo) UnmarshalJSON(b []byte) error { - return marshal.Unmarshal(b, s) -} - -func (s DataPlaneInfo) MarshalJSON() ([]byte, error) { - return marshal.Marshal(s) -} - // Delete account federation policy type DeleteAccountFederationPolicyRequest struct { // The identifier for the federation policy. diff --git a/service/serving/ext_data_plane.go b/service/serving/ext_data_plane.go index c4ec977b4..37c80aed1 100644 --- a/service/serving/ext_data_plane.go +++ b/service/serving/ext_data_plane.go @@ -4,24 +4,23 @@ import ( "strings" "sync" - "github.com/databricks/databricks-sdk-go/service/oauth2" goauth "golang.org/x/oauth2" ) // DataPlaneService is an interface for services that access DataPlane. type DataPlaneService interface { - GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error) + GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *goauth.Token, error) } func NewDataPlaneService() DataPlaneService { return &dataPlaneServiceImpl{ - infos: make(map[string]*oauth2.DataPlaneInfo), + infos: make(map[string]*DataPlaneInfo), tokens: make(map[string]*goauth.Token), } } type dataPlaneServiceImpl struct { - infos map[string]*oauth2.DataPlaneInfo + infos map[string]*DataPlaneInfo tokens map[string]*goauth.Token // This class can be shared across multiple threads. // This mutex is used to synchronize access to the infos and tokens maps. @@ -30,7 +29,7 @@ type dataPlaneServiceImpl struct { // GetDataPlaneDetails returns the endpoint URL and token. It returns a cached token if it is valid, // otherwise it refreshes the token and returns the new token. -func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error) { +func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *goauth.Token, error) { key := o.generateKey(method, params) info, err := o.getInfo(key, infoGetter) if err != nil { @@ -43,7 +42,7 @@ func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []strin return info.EndpointUrl, token, nil } -func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*oauth2.DataPlaneInfo, error)) (*oauth2.DataPlaneInfo, error) { +func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*DataPlaneInfo, error)) (*DataPlaneInfo, error) { info, infoOk := o.infos[key] if !infoOk { o.mu.Lock() @@ -61,7 +60,7 @@ func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*oauth2.Da return info, nil } -func (o *dataPlaneServiceImpl) refreshToken(key string, info *oauth2.DataPlaneInfo, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error)) (*goauth.Token, error) { +func (o *dataPlaneServiceImpl) refreshToken(key string, info *DataPlaneInfo, refresh func(*DataPlaneInfo) (*goauth.Token, error)) (*goauth.Token, error) { token, tokenOk := o.tokens[key] if !tokenOk || !token.Valid() { o.mu.Lock() diff --git a/service/serving/ext_data_plane_test.go b/service/serving/ext_data_plane_test.go index 4fa3ca3dd..27a0972ac 100644 --- a/service/serving/ext_data_plane_test.go +++ b/service/serving/ext_data_plane_test.go @@ -4,30 +4,29 @@ import ( "testing" "time" - "github.com/databricks/databricks-sdk-go/service/oauth2" "github.com/stretchr/testify/assert" goauth "golang.org/x/oauth2" ) type infoMock struct { called bool - info *oauth2.DataPlaneInfo + info *DataPlaneInfo err error } -func (i *infoMock) DataPlaneInfoGetter() (*oauth2.DataPlaneInfo, error) { +func (i *infoMock) DataPlaneInfoGetter() (*DataPlaneInfo, error) { i.called = true return i.info, i.err } type tokenRefreshSpy struct { called bool - expectedInfo *oauth2.DataPlaneInfo + expectedInfo *DataPlaneInfo token *goauth.Token err error } -func (t *tokenRefreshSpy) TokenRefresh(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { +func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*goauth.Token, error) { t.expectedInfo = info t.called = true return t.token, t.err @@ -35,7 +34,7 @@ func (t *tokenRefreshSpy) TokenRefresh(info *oauth2.DataPlaneInfo) (*goauth.Toke func TestTokenNotCached(t *testing.T) { info := infoMock{ - info: &oauth2.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -50,7 +49,7 @@ func TestTokenNotCached(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{ - infos: make(map[string]*oauth2.DataPlaneInfo), + infos: make(map[string]*DataPlaneInfo), tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) @@ -65,7 +64,7 @@ func TestTokenNotCached(t *testing.T) { func TestTokenCached(t *testing.T) { info := infoMock{ - info: &oauth2.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -80,7 +79,7 @@ func TestTokenCached(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{} - c.infos = make(map[string]*oauth2.DataPlaneInfo) + c.infos = make(map[string]*DataPlaneInfo) c.tokens = make(map[string]*goauth.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = s.token @@ -96,7 +95,7 @@ func TestTokenCached(t *testing.T) { func TestTokenExpired(t *testing.T) { info := infoMock{ - info: &oauth2.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -117,7 +116,7 @@ func TestTokenExpired(t *testing.T) { err: nil, } c := dataPlaneServiceImpl{} - c.infos = make(map[string]*oauth2.DataPlaneInfo) + c.infos = make(map[string]*DataPlaneInfo) c.tokens = make(map[string]*goauth.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = expired @@ -138,7 +137,7 @@ func TestTokenInfoError(t *testing.T) { } s := tokenRefreshSpy{} c := dataPlaneServiceImpl{ - infos: make(map[string]*oauth2.DataPlaneInfo), + infos: make(map[string]*DataPlaneInfo), tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) @@ -151,7 +150,7 @@ func TestTokenInfoError(t *testing.T) { func TestTokenRefreshError(t *testing.T) { info := infoMock{ - info: &oauth2.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -162,7 +161,7 @@ func TestTokenRefreshError(t *testing.T) { err: assert.AnError, } c := dataPlaneServiceImpl{ - infos: make(map[string]*oauth2.DataPlaneInfo), + infos: make(map[string]*DataPlaneInfo), tokens: make(map[string]*goauth.Token), } url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) diff --git a/service/serving/impl.go b/service/serving/impl.go index 1261c1b7b..77476ef7f 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -11,8 +11,6 @@ import ( "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/httpclient" goauth "golang.org/x/oauth2" - - "github.com/databricks/databricks-sdk-go/service/oauth2" ) // unexported type that holds implementations of just ServingEndpoints API methods @@ -196,7 +194,7 @@ func (a *servingEndpointsDataPlaneImpl) Query(ctx context.Context, request Query if err != nil { return nil, err } - infoGetter := func() (*oauth2.DataPlaneInfo, error) { + infoGetter := func() (*DataPlaneInfo, error) { response, err := a.controlPlane.Get(ctx, getRequest) if err != nil { return nil, err @@ -206,7 +204,7 @@ func (a *servingEndpointsDataPlaneImpl) Query(ctx context.Context, request Query } return response.DataPlaneInfo.QueryInfo, nil } - refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + refresh := func(info *DataPlaneInfo) (*goauth.Token, error) { return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) } getParams := []string{ diff --git a/service/serving/model.go b/service/serving/model.go index 6ac7eb49d..ade652782 100755 --- a/service/serving/model.go +++ b/service/serving/model.go @@ -7,7 +7,6 @@ import ( "io" "github.com/databricks/databricks-sdk-go/marshal" - "github.com/databricks/databricks-sdk-go/service/oauth2" ) type Ai21LabsConfig struct { @@ -498,6 +497,23 @@ func (s CreateServingEndpoint) MarshalJSON() ([]byte, error) { return marshal.Marshal(s) } +type DataPlaneInfo struct { + // Authorization details as a string. + AuthorizationDetails string `json:"authorization_details,omitempty"` + // The URL of the endpoint for this operation in the dataplane. + EndpointUrl string `json:"endpoint_url,omitempty"` + + ForceSendFields []string `json:"-"` +} + +func (s *DataPlaneInfo) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, s) +} + +func (s DataPlaneInfo) MarshalJSON() ([]byte, error) { + return marshal.Marshal(s) +} + type DatabricksModelServingConfig struct { // The Databricks secret key reference for a Databricks API token that // corresponds to a user or service principal with Can Query access to the @@ -981,7 +997,7 @@ type LogsRequest struct { type ModelDataPlaneInfo struct { // Information required to query DataPlane API 'query' endpoint. - QueryInfo *oauth2.DataPlaneInfo `json:"query_info,omitempty"` + QueryInfo *DataPlaneInfo `json:"query_info,omitempty"` } type OpenAiConfig struct { From 4e48a6722c4c3d2645f336aad017e1af47510f6f Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Mon, 6 Jan 2025 19:03:50 +0100 Subject: [PATCH 4/4] update --- .codegen/_openapi_sha | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 3cf20d57b..e43e9607f 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -- \ No newline at end of file +1668b0db17e23605f8c9d29fb3b674c01590732d \ No newline at end of file