diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 8c62ac620..2dd60127b 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -7eb5ad9a2ed3e3f1055968a2d1014ac92c06fe92 \ No newline at end of file +universe:/Users/hector.castejon/universe \ No newline at end of file diff --git a/.codegen/accounts.go.tmpl b/.codegen/accounts.go.tmpl index dbc4f3df5..d0487a06f 100644 --- a/.codegen/accounts.go.tmpl +++ b/.codegen/accounts.go.tmpl @@ -42,9 +42,20 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { if err != nil { return nil, err } + {{range .Services}}{{if and .IsAccounts (not .HasParent) (.HasDataPlaneMethods) (not .IsDataPlane)}} + {{.CamelName}} := {{.Package.Name}}.New{{.Name}}(apiClient) + {{- end}}{{end}} return &AccountClient{ Config: cfg, - {{range .Services}}{{if and .IsAccounts (not .HasParent)}} - {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient),{{end}}{{end}} + {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane) (not .HasDataPlaneMethods)}} + {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient), + {{- end -}} + {{if and .IsAccounts (not .HasParent) (not .IsDataPlane) (.HasDataPlaneMethods)}} + {{(.TrimPrefix "account").PascalName}}: {{.CamelName}}, + {{- end -}} + {{if and .IsAccounts (not .HasParent) .IsDataPlane}} + {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient, {{.ControlPlaneService.CamelName}}), + {{- end -}} + {{end}} }, nil } diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 11f089775..0be5b70cc 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -105,10 +105,17 @@ type {{.PascalName}}Interface interface { {{end -}} } -func New{{.PascalName}}(client *client.DatabricksClient) *{{.PascalName}}API { +func New{{.PascalName}}(client *client.DatabricksClient, +{{- if .IsDataPlane}} +controlPlane *{{.ControlPlaneService.PascalName}}API, +{{end -}} +) *{{.PascalName}}API { return &{{.PascalName}}API{ impl: &{{.CamelName}}Impl{ client: client, + {{- if .IsDataPlane}} + controlPlane: controlPlane, + {{end}} }, {{range .Subservices}} {{.CamelName}}: New{{.PascalName}}(client), @@ -148,6 +155,8 @@ func (a *{{.PascalName}}API) Impl() {{.PascalName}}Service { return a.impl } + + {{range .Waits}} // {{.PascalName}} repeatedly calls [{{.Method.Service.Name}}API.{{.Poll.PascalName}}] and waits to reach {{range $i, $e := .Success}}{{if $i}} or {{end}}{{.Content}}{{end}} state func (a *{{.Method.Service.PascalName}}API) {{.PascalName}}(ctx context.Context{{range .Binding}}, {{.PollField.CamelName}} {{template "type" .PollField.Entity}}{{end}}, @@ -427,4 +436,4 @@ func (a *{{.Service.Name}}API) {{.Shortcut.PascalName}}AndWait(ctx context.Conte {{- define "wait-response-type" -}} (*{{.Wait.PascalName}}[{{with .Response}}{{if .IsEmpty}}struct{}{{else}}{{.PascalName}}{{end}}{{end}}], error) -{{- end -}} +{{- end -}} \ No newline at end of file diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 811111567..e2359e7d8 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -4,7 +4,9 @@ package {{.Name}} import ( "context" + "errors" "fmt" + goauth "golang.org/x/oauth2" "time" "io" "net/http" @@ -16,10 +18,59 @@ import ( {{range .Services}} // unexported type that holds implementations of just {{.Name}} API methods type {{.CamelName}}Impl struct { + {{- if .IsDataPlane}} + oauth2.DataPlaneHelper + controlPlane *{{.ControlPlaneService.PascalName}}API + {{end -}} client *client.DatabricksClient } {{range .Methods}} +{{if .Service.IsDataPlane}} +func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} { + getRequest := {{.Service.DataPlaneInfoMethod.Request.PascalName}}{ + {{- range .Service.DataPlaneInfoMethod.Request.Fields}} + {{.PascalName}}: request.{{.PascalName}}, + {{end}} + } + token, err := a.client.Config.GetToken() + if err != nil { + return nil, err + } + infoGetter := func() (*oauth2.DataPlaneInfo, error) { + response, err := a.controlPlane.{{.Service.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) + if err != nil { + return nil, err + } + if response.{{(index .DataPlaneInfoFields 0).PascalName}} == nil { + return nil, errors.New("resource does not support direct Data Plane access") + } + return response{{range .DataPlaneInfoFields}}.{{.PascalName}}{{end}}, nil + } + refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) + } + getParams := []string{ + {{- range .Service.DataPlaneInfoMethod.Request.Fields}} + request.{{.PascalName}}, + {{end -}} + } + endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("{{.PascalName}}", getParams, refresh, infoGetter) + if err != nil { + return nil, err + } + {{ template "make-header" . }} + opts := []httpclient.DoOption{} + opts = append(opts, httpclient.WithRequestHeaders(headers)) + {{- template "response-var" . }} + {{if .Response}}opts = append(opts, httpclient.WithRequestData(request)){{end}} + {{if .Response}}opts = append(opts, httpclient.WithResponseUnmarshal(&{{.Response.CamelName}})){{end}} + opts = append(opts, httpclient.WithToken(dataPlaneToken)) + err = a.client.ApiClient().Do(ctx, http.Method{{.TitleVerb}}, endpointUrl, opts...) + return {{ template "response" . }} +} + +{{else}} func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} { {{- template "response-var" . }} path := {{ template "path" . }} @@ -27,6 +78,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re err := a.client.Do(ctx, http.Method{{.TitleVerb}}, path, headers, {{ template "request-param" . }}, {{if .Response}}&{{.Response.CamelName}}{{else}}nil{{end}}) return {{ template "response" . }} } +{{end}} {{end -}} {{end}} diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index 03189f9e5..a2eb01566 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -21,4 +21,6 @@ type {{.PascalName}}Service interface { {{end}} } + + {{end}} \ No newline at end of file diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index d4db23cfc..00fecb949 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.codegen/workspaces.go.tmpl @@ -20,19 +20,6 @@ type WorkspaceClient struct { {{end}}{{end}} } -// Returns a new OAuth scoped to the authorization details provided. -// It will return an error if the CredentialStrategy does not support OAuth tokens. -// -// **NOTE:** Experimental: This API may change or be removed in a future release -// without warning. -func (a *WorkspaceClient) GetOAuthToken(ctx context.Context, authorizationDetails string) (*credentials.OAuthToken, error) { - originalToken, err := a.Config.GetToken() - if err != nil { - return nil, err - } - return a.apiClient.GetOAuthToken(ctx, authorizationDetails, originalToken) -} - var ErrNotWorkspaceClient = errors.New("invalid Databricks Workspace configuration") // NewWorkspaceClient creates new Databricks SDK client for Workspaces or @@ -61,11 +48,21 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { if err != nil { return nil, err } + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (.HasDataPlaneMethods) (not .IsDataPlane)}} + {{.CamelName}} := {{.Package.Name}}.New{{.Name}}(databricksClient) + {{- end}}{{end}} return &WorkspaceClient{ Config: cfg, apiClient: apiClient, - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane) (not .HasDataPlaneMethods)}} {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient), - {{- end}}{{end}} + {{- end -}} + {{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane) (.HasDataPlaneMethods)}} + {{.Name}}: {{.CamelName}}, + {{- end -}} + {{if and (not .IsAccounts) (not .HasParent) .IsDataPlane}} + {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient, {{.ControlPlaneService.CamelName}}), + {{- end -}} + {{end}} }, nil } diff --git a/.gitattributes b/.gitattributes index 9c869d613..e77f5a5c4 100644 --- a/.gitattributes +++ b/.gitattributes @@ -79,6 +79,7 @@ experimental/mocks/service/provisioning/mock_storage_interface.go linguist-gener experimental/mocks/service/provisioning/mock_vpc_endpoints_interface.go linguist-generated=true experimental/mocks/service/provisioning/mock_workspaces_interface.go linguist-generated=true experimental/mocks/service/serving/mock_apps_interface.go linguist-generated=true +experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go linguist-generated=true experimental/mocks/service/serving/mock_serving_endpoints_interface.go linguist-generated=true experimental/mocks/service/settings/mock_account_ip_access_lists_interface.go linguist-generated=true experimental/mocks/service/settings/mock_account_settings_interface.go linguist-generated=true diff --git a/account_client.go b/account_client.go index 2596a826e..a226ec7df 100755 --- a/account_client.go +++ b/account_client.go @@ -289,6 +289,7 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { if err != nil { return nil, err } + return &AccountClient{ Config: cfg, diff --git a/client/client.go b/client/client.go index 671ccf09a..77cd7e146 100644 --- a/client/client.go +++ b/client/client.go @@ -7,6 +7,7 @@ import ( "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/httpclient" + "golang.org/x/oauth2" ) func New(cfg *config.Config) (*DatabricksClient, error) { @@ -46,6 +47,20 @@ func (c *DatabricksClient) ConfiguredAccountID() string { return c.Config.AccountID } +// Returns the inner Api Client. +func (c *DatabricksClient) ApiClient() *httpclient.ApiClient { + return c.client +} + +// Returns a new OAuth token using the provided token. The token must be a JWT token. +// The resulting token is scoped to the authorization details provided. +// +// **NOTE:** Experimental: This API may change or be removed in a future release +// without warning. +func (c *DatabricksClient) GetOAuthToken(ctx context.Context, authDetails string, token *oauth2.Token) (*oauth2.Token, error) { + return c.client.GetOAuthToken(ctx, authDetails, token) +} + // Do sends an HTTP request against path. func (c *DatabricksClient) Do(ctx context.Context, method, path string, headers map[string]string, request, response any, diff --git a/experimental/mocks/mock_workspace_client.go b/experimental/mocks/mock_workspace_client.go index 79f03eccf..b2c43aa79 100755 --- a/experimental/mocks/mock_workspace_client.go +++ b/experimental/mocks/mock_workspace_client.go @@ -102,6 +102,7 @@ func NewMockWorkspaceClient(t interface { Schemas: catalog.NewMockSchemasInterface(t), Secrets: workspace.NewMockSecretsInterface(t), ServicePrincipals: iam.NewMockServicePrincipalsInterface(t), + ServingEndpointDataPlane: serving.NewMockServingEndpointDataPlaneInterface(t), ServingEndpoints: serving.NewMockServingEndpointsInterface(t), Settings: settings.NewMockSettingsInterface(t), Shares: sharing.NewMockSharesInterface(t), @@ -687,6 +688,14 @@ func (m *MockWorkspaceClient) GetMockServicePrincipalsAPI() *iam.MockServicePrin return api } +func (m *MockWorkspaceClient) GetMockServingEndpointDataPlaneAPI() *serving.MockServingEndpointDataPlaneInterface { + api, ok := m.WorkspaceClient.ServingEndpointDataPlane.(*serving.MockServingEndpointDataPlaneInterface) + if !ok { + panic(fmt.Sprintf("expected ServingEndpointDataPlane to be *serving.MockServingEndpointDataPlaneInterface, actual was %T", m.WorkspaceClient.ServingEndpointDataPlane)) + } + return api +} + func (m *MockWorkspaceClient) GetMockServingEndpointsAPI() *serving.MockServingEndpointsInterface { api, ok := m.WorkspaceClient.ServingEndpoints.(*serving.MockServingEndpointsInterface) if !ok { diff --git a/experimental/mocks/service/catalog/mock_account_storage_credentials_interface.go b/experimental/mocks/service/catalog/mock_account_storage_credentials_interface.go index cc623fca9..ff4db744d 100644 --- a/experimental/mocks/service/catalog/mock_account_storage_credentials_interface.go +++ b/experimental/mocks/service/catalog/mock_account_storage_credentials_interface.go @@ -7,6 +7,8 @@ import ( catalog "github.com/databricks/databricks-sdk-go/service/catalog" + listing "github.com/databricks/databricks-sdk-go/listing" + mock "github.com/stretchr/testify/mock" ) @@ -296,65 +298,6 @@ func (_c *MockAccountStorageCredentialsInterface_GetByMetastoreIdAndStorageCrede return _c } -// GetByName provides a mock function with given fields: ctx, name -func (_m *MockAccountStorageCredentialsInterface) GetByName(ctx context.Context, name string) (*catalog.StorageCredentialInfo, error) { - ret := _m.Called(ctx, name) - - if len(ret) == 0 { - panic("no return value specified for GetByName") - } - - var r0 *catalog.StorageCredentialInfo - var r1 error - if rf, ok := ret.Get(0).(func(context.Context, string) (*catalog.StorageCredentialInfo, error)); ok { - return rf(ctx, name) - } - if rf, ok := ret.Get(0).(func(context.Context, string) *catalog.StorageCredentialInfo); ok { - r0 = rf(ctx, name) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(*catalog.StorageCredentialInfo) - } - } - - if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { - r1 = rf(ctx, name) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// MockAccountStorageCredentialsInterface_GetByName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetByName' -type MockAccountStorageCredentialsInterface_GetByName_Call struct { - *mock.Call -} - -// GetByName is a helper method to define mock.On call -// - ctx context.Context -// - name string -func (_e *MockAccountStorageCredentialsInterface_Expecter) GetByName(ctx interface{}, name interface{}) *MockAccountStorageCredentialsInterface_GetByName_Call { - return &MockAccountStorageCredentialsInterface_GetByName_Call{Call: _e.mock.On("GetByName", ctx, name)} -} - -func (_c *MockAccountStorageCredentialsInterface_GetByName_Call) Run(run func(ctx context.Context, name string)) *MockAccountStorageCredentialsInterface_GetByName_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(string)) - }) - return _c -} - -func (_c *MockAccountStorageCredentialsInterface_GetByName_Call) Return(_a0 *catalog.StorageCredentialInfo, _a1 error) *MockAccountStorageCredentialsInterface_GetByName_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockAccountStorageCredentialsInterface_GetByName_Call) RunAndReturn(run func(context.Context, string) (*catalog.StorageCredentialInfo, error)) *MockAccountStorageCredentialsInterface_GetByName_Call { - _c.Call.Return(run) - return _c -} - // Impl provides a mock function with given fields: func (_m *MockAccountStorageCredentialsInterface) Impl() catalog.AccountStorageCredentialsService { ret := _m.Called() @@ -403,33 +346,23 @@ func (_c *MockAccountStorageCredentialsInterface_Impl_Call) RunAndReturn(run fun } // List provides a mock function with given fields: ctx, request -func (_m *MockAccountStorageCredentialsInterface) List(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error) { +func (_m *MockAccountStorageCredentialsInterface) List(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest) listing.Iterator[catalog.StorageCredentialInfo] { ret := _m.Called(ctx, request) if len(ret) == 0 { panic("no return value specified for List") } - var r0 []catalog.StorageCredentialInfo - var r1 error - if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error)); ok { - return rf(ctx, request) - } - if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) []catalog.StorageCredentialInfo); ok { + var r0 listing.Iterator[catalog.StorageCredentialInfo] + if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) listing.Iterator[catalog.StorageCredentialInfo]); ok { r0 = rf(ctx, request) } else { if ret.Get(0) != nil { - r0 = ret.Get(0).([]catalog.StorageCredentialInfo) + r0 = ret.Get(0).(listing.Iterator[catalog.StorageCredentialInfo]) } } - if rf, ok := ret.Get(1).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) error); ok { - r1 = rf(ctx, request) - } else { - r1 = ret.Error(1) - } - - return r0, r1 + return r0 } // MockAccountStorageCredentialsInterface_List_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'List' @@ -451,39 +384,39 @@ func (_c *MockAccountStorageCredentialsInterface_List_Call) Run(run func(ctx con return _c } -func (_c *MockAccountStorageCredentialsInterface_List_Call) Return(_a0 []catalog.StorageCredentialInfo, _a1 error) *MockAccountStorageCredentialsInterface_List_Call { - _c.Call.Return(_a0, _a1) +func (_c *MockAccountStorageCredentialsInterface_List_Call) Return(_a0 listing.Iterator[catalog.StorageCredentialInfo]) *MockAccountStorageCredentialsInterface_List_Call { + _c.Call.Return(_a0) return _c } -func (_c *MockAccountStorageCredentialsInterface_List_Call) RunAndReturn(run func(context.Context, catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error)) *MockAccountStorageCredentialsInterface_List_Call { +func (_c *MockAccountStorageCredentialsInterface_List_Call) RunAndReturn(run func(context.Context, catalog.ListAccountStorageCredentialsRequest) listing.Iterator[catalog.StorageCredentialInfo]) *MockAccountStorageCredentialsInterface_List_Call { _c.Call.Return(run) return _c } -// ListByMetastoreId provides a mock function with given fields: ctx, metastoreId -func (_m *MockAccountStorageCredentialsInterface) ListByMetastoreId(ctx context.Context, metastoreId string) ([]catalog.StorageCredentialInfo, error) { - ret := _m.Called(ctx, metastoreId) +// ListAll provides a mock function with given fields: ctx, request +func (_m *MockAccountStorageCredentialsInterface) ListAll(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error) { + ret := _m.Called(ctx, request) if len(ret) == 0 { - panic("no return value specified for ListByMetastoreId") + panic("no return value specified for ListAll") } var r0 []catalog.StorageCredentialInfo var r1 error - if rf, ok := ret.Get(0).(func(context.Context, string) ([]catalog.StorageCredentialInfo, error)); ok { - return rf(ctx, metastoreId) + if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error)); ok { + return rf(ctx, request) } - if rf, ok := ret.Get(0).(func(context.Context, string) []catalog.StorageCredentialInfo); ok { - r0 = rf(ctx, metastoreId) + if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) []catalog.StorageCredentialInfo); ok { + r0 = rf(ctx, request) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]catalog.StorageCredentialInfo) } } - if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { - r1 = rf(ctx, metastoreId) + if rf, ok := ret.Get(1).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) error); ok { + r1 = rf(ctx, request) } else { r1 = ret.Error(1) } @@ -491,58 +424,58 @@ func (_m *MockAccountStorageCredentialsInterface) ListByMetastoreId(ctx context. return r0, r1 } -// MockAccountStorageCredentialsInterface_ListByMetastoreId_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListByMetastoreId' -type MockAccountStorageCredentialsInterface_ListByMetastoreId_Call struct { +// MockAccountStorageCredentialsInterface_ListAll_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListAll' +type MockAccountStorageCredentialsInterface_ListAll_Call struct { *mock.Call } -// ListByMetastoreId is a helper method to define mock.On call +// ListAll is a helper method to define mock.On call // - ctx context.Context -// - metastoreId string -func (_e *MockAccountStorageCredentialsInterface_Expecter) ListByMetastoreId(ctx interface{}, metastoreId interface{}) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { - return &MockAccountStorageCredentialsInterface_ListByMetastoreId_Call{Call: _e.mock.On("ListByMetastoreId", ctx, metastoreId)} +// - request catalog.ListAccountStorageCredentialsRequest +func (_e *MockAccountStorageCredentialsInterface_Expecter) ListAll(ctx interface{}, request interface{}) *MockAccountStorageCredentialsInterface_ListAll_Call { + return &MockAccountStorageCredentialsInterface_ListAll_Call{Call: _e.mock.On("ListAll", ctx, request)} } -func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) Run(run func(ctx context.Context, metastoreId string)) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { +func (_c *MockAccountStorageCredentialsInterface_ListAll_Call) Run(run func(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest)) *MockAccountStorageCredentialsInterface_ListAll_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(string)) + run(args[0].(context.Context), args[1].(catalog.ListAccountStorageCredentialsRequest)) }) return _c } -func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) Return(_a0 []catalog.StorageCredentialInfo, _a1 error) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { +func (_c *MockAccountStorageCredentialsInterface_ListAll_Call) Return(_a0 []catalog.StorageCredentialInfo, _a1 error) *MockAccountStorageCredentialsInterface_ListAll_Call { _c.Call.Return(_a0, _a1) return _c } -func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) RunAndReturn(run func(context.Context, string) ([]catalog.StorageCredentialInfo, error)) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { +func (_c *MockAccountStorageCredentialsInterface_ListAll_Call) RunAndReturn(run func(context.Context, catalog.ListAccountStorageCredentialsRequest) ([]catalog.StorageCredentialInfo, error)) *MockAccountStorageCredentialsInterface_ListAll_Call { _c.Call.Return(run) return _c } -// StorageCredentialInfoNameToIdMap provides a mock function with given fields: ctx, request -func (_m *MockAccountStorageCredentialsInterface) StorageCredentialInfoNameToIdMap(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest) (map[string]string, error) { - ret := _m.Called(ctx, request) +// ListByMetastoreId provides a mock function with given fields: ctx, metastoreId +func (_m *MockAccountStorageCredentialsInterface) ListByMetastoreId(ctx context.Context, metastoreId string) (*catalog.ListAccountStorageCredentialsResponse, error) { + ret := _m.Called(ctx, metastoreId) if len(ret) == 0 { - panic("no return value specified for StorageCredentialInfoNameToIdMap") + panic("no return value specified for ListByMetastoreId") } - var r0 map[string]string + var r0 *catalog.ListAccountStorageCredentialsResponse var r1 error - if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) (map[string]string, error)); ok { - return rf(ctx, request) + if rf, ok := ret.Get(0).(func(context.Context, string) (*catalog.ListAccountStorageCredentialsResponse, error)); ok { + return rf(ctx, metastoreId) } - if rf, ok := ret.Get(0).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) map[string]string); ok { - r0 = rf(ctx, request) + if rf, ok := ret.Get(0).(func(context.Context, string) *catalog.ListAccountStorageCredentialsResponse); ok { + r0 = rf(ctx, metastoreId) } else { if ret.Get(0) != nil { - r0 = ret.Get(0).(map[string]string) + r0 = ret.Get(0).(*catalog.ListAccountStorageCredentialsResponse) } } - if rf, ok := ret.Get(1).(func(context.Context, catalog.ListAccountStorageCredentialsRequest) error); ok { - r1 = rf(ctx, request) + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, metastoreId) } else { r1 = ret.Error(1) } @@ -550,31 +483,31 @@ func (_m *MockAccountStorageCredentialsInterface) StorageCredentialInfoNameToIdM return r0, r1 } -// MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StorageCredentialInfoNameToIdMap' -type MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call struct { +// MockAccountStorageCredentialsInterface_ListByMetastoreId_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListByMetastoreId' +type MockAccountStorageCredentialsInterface_ListByMetastoreId_Call struct { *mock.Call } -// StorageCredentialInfoNameToIdMap is a helper method to define mock.On call +// ListByMetastoreId is a helper method to define mock.On call // - ctx context.Context -// - request catalog.ListAccountStorageCredentialsRequest -func (_e *MockAccountStorageCredentialsInterface_Expecter) StorageCredentialInfoNameToIdMap(ctx interface{}, request interface{}) *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call { - return &MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call{Call: _e.mock.On("StorageCredentialInfoNameToIdMap", ctx, request)} +// - metastoreId string +func (_e *MockAccountStorageCredentialsInterface_Expecter) ListByMetastoreId(ctx interface{}, metastoreId interface{}) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { + return &MockAccountStorageCredentialsInterface_ListByMetastoreId_Call{Call: _e.mock.On("ListByMetastoreId", ctx, metastoreId)} } -func (_c *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call) Run(run func(ctx context.Context, request catalog.ListAccountStorageCredentialsRequest)) *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call { +func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) Run(run func(ctx context.Context, metastoreId string)) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(catalog.ListAccountStorageCredentialsRequest)) + run(args[0].(context.Context), args[1].(string)) }) return _c } -func (_c *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call) Return(_a0 map[string]string, _a1 error) *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call { +func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) Return(_a0 *catalog.ListAccountStorageCredentialsResponse, _a1 error) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { _c.Call.Return(_a0, _a1) return _c } -func (_c *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call) RunAndReturn(run func(context.Context, catalog.ListAccountStorageCredentialsRequest) (map[string]string, error)) *MockAccountStorageCredentialsInterface_StorageCredentialInfoNameToIdMap_Call { +func (_c *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call) RunAndReturn(run func(context.Context, string) (*catalog.ListAccountStorageCredentialsResponse, error)) *MockAccountStorageCredentialsInterface_ListByMetastoreId_Call { _c.Call.Return(run) return _c } diff --git a/experimental/mocks/service/marketplace/mock_consumer_listings_interface.go b/experimental/mocks/service/marketplace/mock_consumer_listings_interface.go index e1b389f9f..fb7b8598f 100644 --- a/experimental/mocks/service/marketplace/mock_consumer_listings_interface.go +++ b/experimental/mocks/service/marketplace/mock_consumer_listings_interface.go @@ -24,6 +24,65 @@ func (_m *MockConsumerListingsInterface) EXPECT() *MockConsumerListingsInterface return &MockConsumerListingsInterface_Expecter{mock: &_m.Mock} } +// GEt provides a mock function with given fields: ctx, request +func (_m *MockConsumerListingsInterface) GEt(ctx context.Context, request marketplace.BatchGetListingsRequest) (*marketplace.BatchGetListingsResponse, error) { + ret := _m.Called(ctx, request) + + if len(ret) == 0 { + panic("no return value specified for GEt") + } + + var r0 *marketplace.BatchGetListingsResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, marketplace.BatchGetListingsRequest) (*marketplace.BatchGetListingsResponse, error)); ok { + return rf(ctx, request) + } + if rf, ok := ret.Get(0).(func(context.Context, marketplace.BatchGetListingsRequest) *marketplace.BatchGetListingsResponse); ok { + r0 = rf(ctx, request) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*marketplace.BatchGetListingsResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, marketplace.BatchGetListingsRequest) error); ok { + r1 = rf(ctx, request) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockConsumerListingsInterface_GEt_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GEt' +type MockConsumerListingsInterface_GEt_Call struct { + *mock.Call +} + +// GEt is a helper method to define mock.On call +// - ctx context.Context +// - request marketplace.BatchGetListingsRequest +func (_e *MockConsumerListingsInterface_Expecter) GEt(ctx interface{}, request interface{}) *MockConsumerListingsInterface_GEt_Call { + return &MockConsumerListingsInterface_GEt_Call{Call: _e.mock.On("GEt", ctx, request)} +} + +func (_c *MockConsumerListingsInterface_GEt_Call) Run(run func(ctx context.Context, request marketplace.BatchGetListingsRequest)) *MockConsumerListingsInterface_GEt_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(marketplace.BatchGetListingsRequest)) + }) + return _c +} + +func (_c *MockConsumerListingsInterface_GEt_Call) Return(_a0 *marketplace.BatchGetListingsResponse, _a1 error) *MockConsumerListingsInterface_GEt_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockConsumerListingsInterface_GEt_Call) RunAndReturn(run func(context.Context, marketplace.BatchGetListingsRequest) (*marketplace.BatchGetListingsResponse, error)) *MockConsumerListingsInterface_GEt_Call { + _c.Call.Return(run) + return _c +} + // Get provides a mock function with given fields: ctx, request func (_m *MockConsumerListingsInterface) Get(ctx context.Context, request marketplace.GetListingRequest) (*marketplace.GetListingResponse, error) { ret := _m.Called(ctx, request) diff --git a/experimental/mocks/service/marketplace/mock_consumer_providers_interface.go b/experimental/mocks/service/marketplace/mock_consumer_providers_interface.go index 4aad119a2..1ad86b63f 100644 --- a/experimental/mocks/service/marketplace/mock_consumer_providers_interface.go +++ b/experimental/mocks/service/marketplace/mock_consumer_providers_interface.go @@ -24,6 +24,65 @@ func (_m *MockConsumerProvidersInterface) EXPECT() *MockConsumerProvidersInterfa return &MockConsumerProvidersInterface_Expecter{mock: &_m.Mock} } +// GEt provides a mock function with given fields: ctx, request +func (_m *MockConsumerProvidersInterface) GEt(ctx context.Context, request marketplace.BatchGetProvidersRequest) (*marketplace.BatchGetProvidersResponse, error) { + ret := _m.Called(ctx, request) + + if len(ret) == 0 { + panic("no return value specified for GEt") + } + + var r0 *marketplace.BatchGetProvidersResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, marketplace.BatchGetProvidersRequest) (*marketplace.BatchGetProvidersResponse, error)); ok { + return rf(ctx, request) + } + if rf, ok := ret.Get(0).(func(context.Context, marketplace.BatchGetProvidersRequest) *marketplace.BatchGetProvidersResponse); ok { + r0 = rf(ctx, request) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*marketplace.BatchGetProvidersResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, marketplace.BatchGetProvidersRequest) error); ok { + r1 = rf(ctx, request) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockConsumerProvidersInterface_GEt_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GEt' +type MockConsumerProvidersInterface_GEt_Call struct { + *mock.Call +} + +// GEt is a helper method to define mock.On call +// - ctx context.Context +// - request marketplace.BatchGetProvidersRequest +func (_e *MockConsumerProvidersInterface_Expecter) GEt(ctx interface{}, request interface{}) *MockConsumerProvidersInterface_GEt_Call { + return &MockConsumerProvidersInterface_GEt_Call{Call: _e.mock.On("GEt", ctx, request)} +} + +func (_c *MockConsumerProvidersInterface_GEt_Call) Run(run func(ctx context.Context, request marketplace.BatchGetProvidersRequest)) *MockConsumerProvidersInterface_GEt_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(marketplace.BatchGetProvidersRequest)) + }) + return _c +} + +func (_c *MockConsumerProvidersInterface_GEt_Call) Return(_a0 *marketplace.BatchGetProvidersResponse, _a1 error) *MockConsumerProvidersInterface_GEt_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockConsumerProvidersInterface_GEt_Call) RunAndReturn(run func(context.Context, marketplace.BatchGetProvidersRequest) (*marketplace.BatchGetProvidersResponse, error)) *MockConsumerProvidersInterface_GEt_Call { + _c.Call.Return(run) + return _c +} + // Get provides a mock function with given fields: ctx, request func (_m *MockConsumerProvidersInterface) Get(ctx context.Context, request marketplace.GetProviderRequest) (*marketplace.GetProviderResponse, error) { ret := _m.Called(ctx, request) diff --git a/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go b/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go new file mode 100644 index 000000000..4f581953e --- /dev/null +++ b/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go @@ -0,0 +1,191 @@ +// Code generated by mockery v2.43.0. DO NOT EDIT. + +package serving + +import ( + context "context" + + serving "github.com/databricks/databricks-sdk-go/service/serving" + mock "github.com/stretchr/testify/mock" +) + +// MockServingEndpointDataPlaneInterface is an autogenerated mock type for the ServingEndpointDataPlaneInterface type +type MockServingEndpointDataPlaneInterface struct { + mock.Mock +} + +type MockServingEndpointDataPlaneInterface_Expecter struct { + mock *mock.Mock +} + +func (_m *MockServingEndpointDataPlaneInterface) EXPECT() *MockServingEndpointDataPlaneInterface_Expecter { + return &MockServingEndpointDataPlaneInterface_Expecter{mock: &_m.Mock} +} + +// Impl provides a mock function with given fields: +func (_m *MockServingEndpointDataPlaneInterface) Impl() serving.ServingEndpointDataPlaneService { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Impl") + } + + var r0 serving.ServingEndpointDataPlaneService + if rf, ok := ret.Get(0).(func() serving.ServingEndpointDataPlaneService); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serving.ServingEndpointDataPlaneService) + } + } + + return r0 +} + +// MockServingEndpointDataPlaneInterface_Impl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Impl' +type MockServingEndpointDataPlaneInterface_Impl_Call struct { + *mock.Call +} + +// Impl is a helper method to define mock.On call +func (_e *MockServingEndpointDataPlaneInterface_Expecter) Impl() *MockServingEndpointDataPlaneInterface_Impl_Call { + return &MockServingEndpointDataPlaneInterface_Impl_Call{Call: _e.mock.On("Impl")} +} + +func (_c *MockServingEndpointDataPlaneInterface_Impl_Call) Run(run func()) *MockServingEndpointDataPlaneInterface_Impl_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_Impl_Call) Return(_a0 serving.ServingEndpointDataPlaneService) *MockServingEndpointDataPlaneInterface_Impl_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_Impl_Call) RunAndReturn(run func() serving.ServingEndpointDataPlaneService) *MockServingEndpointDataPlaneInterface_Impl_Call { + _c.Call.Return(run) + return _c +} + +// Query provides a mock function with given fields: ctx, request +func (_m *MockServingEndpointDataPlaneInterface) Query(ctx context.Context, request serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error) { + ret := _m.Called(ctx, request) + + if len(ret) == 0 { + panic("no return value specified for Query") + } + + var r0 *serving.QueryEndpointResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)); ok { + return rf(ctx, request) + } + if rf, ok := ret.Get(0).(func(context.Context, serving.QueryEndpointInput) *serving.QueryEndpointResponse); ok { + r0 = rf(ctx, request) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*serving.QueryEndpointResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, serving.QueryEndpointInput) error); ok { + r1 = rf(ctx, request) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockServingEndpointDataPlaneInterface_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query' +type MockServingEndpointDataPlaneInterface_Query_Call struct { + *mock.Call +} + +// Query is a helper method to define mock.On call +// - ctx context.Context +// - request serving.QueryEndpointInput +func (_e *MockServingEndpointDataPlaneInterface_Expecter) Query(ctx interface{}, request interface{}) *MockServingEndpointDataPlaneInterface_Query_Call { + return &MockServingEndpointDataPlaneInterface_Query_Call{Call: _e.mock.On("Query", ctx, request)} +} + +func (_c *MockServingEndpointDataPlaneInterface_Query_Call) Run(run func(ctx context.Context, request serving.QueryEndpointInput)) *MockServingEndpointDataPlaneInterface_Query_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(serving.QueryEndpointInput)) + }) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_Query_Call) Return(_a0 *serving.QueryEndpointResponse, _a1 error) *MockServingEndpointDataPlaneInterface_Query_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_Query_Call) RunAndReturn(run func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)) *MockServingEndpointDataPlaneInterface_Query_Call { + _c.Call.Return(run) + return _c +} + +// WithImpl provides a mock function with given fields: impl +func (_m *MockServingEndpointDataPlaneInterface) WithImpl(impl serving.ServingEndpointDataPlaneService) serving.ServingEndpointDataPlaneInterface { + ret := _m.Called(impl) + + if len(ret) == 0 { + panic("no return value specified for WithImpl") + } + + var r0 serving.ServingEndpointDataPlaneInterface + if rf, ok := ret.Get(0).(func(serving.ServingEndpointDataPlaneService) serving.ServingEndpointDataPlaneInterface); ok { + r0 = rf(impl) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serving.ServingEndpointDataPlaneInterface) + } + } + + return r0 +} + +// MockServingEndpointDataPlaneInterface_WithImpl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WithImpl' +type MockServingEndpointDataPlaneInterface_WithImpl_Call struct { + *mock.Call +} + +// WithImpl is a helper method to define mock.On call +// - impl serving.ServingEndpointDataPlaneService +func (_e *MockServingEndpointDataPlaneInterface_Expecter) WithImpl(impl interface{}) *MockServingEndpointDataPlaneInterface_WithImpl_Call { + return &MockServingEndpointDataPlaneInterface_WithImpl_Call{Call: _e.mock.On("WithImpl", impl)} +} + +func (_c *MockServingEndpointDataPlaneInterface_WithImpl_Call) Run(run func(impl serving.ServingEndpointDataPlaneService)) *MockServingEndpointDataPlaneInterface_WithImpl_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(serving.ServingEndpointDataPlaneService)) + }) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_WithImpl_Call) Return(_a0 serving.ServingEndpointDataPlaneInterface) *MockServingEndpointDataPlaneInterface_WithImpl_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_WithImpl_Call) RunAndReturn(run func(serving.ServingEndpointDataPlaneService) serving.ServingEndpointDataPlaneInterface) *MockServingEndpointDataPlaneInterface_WithImpl_Call { + _c.Call.Return(run) + return _c +} + +// NewMockServingEndpointDataPlaneInterface creates a new instance of MockServingEndpointDataPlaneInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockServingEndpointDataPlaneInterface(t interface { + mock.TestingT + Cleanup(func()) +}) *MockServingEndpointDataPlaneInterface { + mock := &MockServingEndpointDataPlaneInterface{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/httpclient/oauth_token.go b/httpclient/oauth_token.go index 691877a53..142afd48d 100644 --- a/httpclient/oauth_token.go +++ b/httpclient/oauth_token.go @@ -3,6 +3,7 @@ package httpclient import ( "context" "net/http" + "time" "github.com/databricks/databricks-sdk-go/credentials" "golang.org/x/oauth2" @@ -26,7 +27,7 @@ type GetOAuthTokenRequest struct { // // **NOTE:** Experimental: This API may change or be removed in a future release // without warning. -func (c *ApiClient) GetOAuthToken(ctx context.Context, authDetails string, token *oauth2.Token) (*credentials.OAuthToken, error) { +func (c *ApiClient) GetOAuthToken(ctx context.Context, authDetails string, token *oauth2.Token) (*oauth2.Token, error) { path := "/oidc/v1/token" data := GetOAuthTokenRequest{ GrantType: JWTGrantType, @@ -42,6 +43,9 @@ func (c *ApiClient) GetOAuthToken(ctx context.Context, authDetails string, token if err != nil { return nil, err } - - return &response, nil + return &oauth2.Token{ + AccessToken: response.AccessToken, + TokenType: response.TokenType, + Expiry: time.Now().Add(time.Duration(response.ExpiresIn) * time.Second), + }, nil } diff --git a/httpclient/request.go b/httpclient/request.go index 4da4e0039..5bb6623d8 100644 --- a/httpclient/request.go +++ b/httpclient/request.go @@ -33,6 +33,17 @@ func WithRequestHeaders(headers map[string]string) DoOption { }) } +// WithToken uses the specified golang.org/x/oauth2 token on a request +func WithToken(token *oauth2.Token) DoOption { + visitor := WithRequestVisitor(func(r *http.Request) error { + auth := fmt.Sprintf("%s %s", token.TokenType, token.AccessToken) + r.Header.Set("Authorization", auth) + return nil + }) + visitor.isAuthOption = true + return visitor +} + // WithTokenSource uses the specified golang.org/x/oauth2 token source on a request func WithTokenSource(ts oauth2.TokenSource) DoOption { return WithRequestVisitor(func(r *http.Request) error { diff --git a/openapi/code/load_test.go b/openapi/code/load_test.go index 115538942..9b254441e 100644 --- a/openapi/code/load_test.go +++ b/openapi/code/load_test.go @@ -129,3 +129,12 @@ func TestMethodsReport(t *testing.T) { assert.Equal(t, len(batch.packages), 1) } + +func TestDataPlane(t *testing.T) { + ctx := context.Background() + batch, err := NewFromFile(ctx, "../testdata/spec_dataplane.json") + require.NoError(t, err) + dataPlane := batch.packages["model"].services["Model"].methods["query"].DataPlane + require.Equal(t, "get", dataPlane.ConfigMethod) + require.Equal(t, "dataplaneInfo", dataPlane.Fields[0]) +} diff --git a/openapi/code/method.go b/openapi/code/method.go index 02ef11c71..57db1d14a 100644 --- a/openapi/code/method.go +++ b/openapi/code/method.go @@ -46,9 +46,31 @@ type Method struct { wait *openapi.Wait pagination *openapi.Pagination Operation *openapi.Operation + DataPlane *openapi.DataPlane shortcut bool } +func (m *Method) HasDataPlaneAPI() bool { + return m.DataPlane != nil +} + +// Returns the fields which contains the DataPlane info. Each field is nested in the previous one. +func (m *Method) DataPlaneInfoFields() []*Field { + if m.DataPlane == nil { + return nil + } + method := m.Service.DataPlaneInfoMethod() + fieldNames := m.DataPlane.Fields + currentLevelFields := method.Response.fields + fields := []*Field{} + for _, name := range fieldNames { + field := currentLevelFields[name] + fields = append(fields, field) + currentLevelFields = field.Entity.fields + } + return fields +} + // Shortcut holds definition of "shortcut" methods, that are generated for // methods with request entities only with required fields. type Shortcut struct { diff --git a/openapi/code/package.go b/openapi/code/package.go index 4ee549fa6..03d7320f9 100644 --- a/openapi/code/package.go +++ b/openapi/code/package.go @@ -420,6 +420,13 @@ func (pkg *Package) Load(ctx context.Context, spec *openapi.Specification, tag o svc.methods[method.Name] = method } } + + // Generate DataPlane service + if svc.HasDataPlaneMethods() { + dataPlaneService := svc.generateDataPlaneService() + pkg.services[dataPlaneService.Name] = dataPlaneService + } + return nil } diff --git a/openapi/code/service.go b/openapi/code/service.go index 49f8504ff..1bdb590a4 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -30,7 +30,86 @@ type Service struct { subservices map[string]*Service ByPathParamsMethods []*Shortcut ParentService *Service + ControlPlaneService *Service tag *openapi.Tag + IsDataPlane bool +} + +// Returns whether any method supports direct DataPlane access. +func (svc *Service) HasDataPlaneMethods() bool { + return len(svc.dataPlaneMethods()) > 0 +} + +// Returns the method in the Control Plane which contains the DataInfo object +func (svc *Service) DataPlaneInfoMethod() *Method { + methodName := "" + for _, m := range svc.methods { + if m.DataPlane != nil { + methodName = m.DataPlane.ConfigMethod + } + } + return svc.ControlPlaneService.methods[methodName] +} + +// Returns the corresponding service for DataPlane APIs. +func (svc *Service) generateDataPlaneService() *Service { + name := svc.Named.Singular().Name + "DataPlane" + description := fmt.Sprintf("%s provides a set of operations to interact with DataPlane endpoints for %s service.", name, svc.Name) + named := Named{ + Name: name, + Description: description, + } + s := &Service{ + Named: named, + Package: svc.Package, + ControlPlaneService: svc, + methods: svc.dataPlaneMethods(), + tag: &openapi.Tag{ + Node: svc.tag.Node, + Package: svc.tag.Package, + PathStyle: svc.tag.PathStyle, + Service: svc.tag.Service, + ParentService: svc.tag.ParentService, + IsAccounts: svc.tag.IsAccounts, + Name: named.Name, + }, + ByPathParamsMethods: svc.ByPathParamsMethods, + IsDataPlane: true, + } + for _, m := range s.methods { + m.Service = s + } + s.tag.Service = s.Name + return s +} + +// Copies methods which implement a DataPlane API. +func (svc *Service) dataPlaneMethods() map[string]*Method { + methods := map[string]*Method{} + for _, m := range svc.methods { + if m.DataPlane != nil { + methods[m.Name] = &Method{ + Named: m.Named, + Verb: m.Verb, + Path: m.Path, + Request: m.Request, + PathParts: m.PathParts, + Response: m.Response, + PathStyle: m.PathStyle, + NameFieldPath: m.NameFieldPath, + IdFieldPath: m.IdFieldPath, + RequestBodyField: m.RequestBodyField, + ResponseBodyField: m.ResponseBodyField, + FixedRequestHeaders: m.FixedRequestHeaders, + wait: m.wait, + Operation: m.Operation, + pagination: m.pagination, + shortcut: m.shortcut, + DataPlane: m.DataPlane, + } + } + } + return methods } // FullName holds package name and service name @@ -528,6 +607,7 @@ func (svc *Service) newMethod(verb, path string, params []openapi.Parameter, op Operation: op, pagination: op.Pagination, shortcut: op.Shortcut, + DataPlane: op.DataPlane, }, nil } diff --git a/openapi/extensions.go b/openapi/extensions.go index 047577f39..9e4d225c7 100644 --- a/openapi/extensions.go +++ b/openapi/extensions.go @@ -29,3 +29,9 @@ type Binding struct { Request string `json:"request,omitempty"` Response string `json:"response,omitempty"` } + +// DataPlane is the Databricks OpenAPI Extension for direct access to DataPlane APIs +type DataPlane struct { + ConfigMethod string `json:"configMethod"` + Fields []string `json:"field"` +} diff --git a/openapi/model.go b/openapi/model.go index 87c01b5d2..d63bbeac1 100644 --- a/openapi/model.go +++ b/openapi/model.go @@ -149,6 +149,7 @@ type Operation struct { Node Wait *Wait `json:"x-databricks-wait,omitempty"` Pagination *Pagination `json:"x-databricks-pagination,omitempty"` + DataPlane *DataPlane `json:"x-databricks-dataplane,omitempty"` Shortcut bool `json:"x-databricks-shortcut,omitempty"` Crud string `json:"x-databricks-crud,omitempty"` JsonOnly bool `json:"x-databricks-cli-json-only,omitempty"` diff --git a/openapi/testdata/spec_dataplane.json b/openapi/testdata/spec_dataplane.json new file mode 100644 index 000000000..32dd7224c --- /dev/null +++ b/openapi/testdata/spec_dataplane.json @@ -0,0 +1,229 @@ +{ + + "openapi": "3.0.0", + + "tags": [ + + { + + "name": "Model", + + "x-databricks-package": "model", + + "x-databricks-service": "Model" + + } + + ], + + "paths": { + + "/api/1.2/model/{name}": { + + "get": { + + "operationId": "Model.get", + + "parameters": [ + + { + + "name": "name", + + "in": "path", + + "required": true, + + "schema": { + + "type": "string" + + } + + } + + ], + + "responses": { + + "200": { + + "content": { + + "application/json": { + + "schema": { "$ref": "#/components/schemas/Model" } + + } + + }, + + "description": "Model was returned successfully." + + } + + }, + + "summary": "Get the model", + + "tags": [ + + "Model" + + ] + + } + + }, + + "/api/1.2/model-query/{name}": { + + "get": { + + "operationId": "Model.query", + + "parameters": [ + + { + + "name": "name", + + "in": "path", + + "required": true, + + "schema": { + + "type": "string" + + } + + } + + ], + + "responses": { + + "200": { + + "content": { + + "application/json": {} + + }, + + "description": "Model was queried successfully." + + } + + }, + + "summary": "Query the model", + + "tags": [ + + "Model" + + ], + + "x-databricks-dataplane" : { + + "configMethod": "get", + + "field": ["dataplaneInfo"] + + }, + + "x-databricks-wait": { + + "bind": "commandId", + + "failure": [ + + "Error" + + ], + + "field": [ + + "status" + + ], + + "message": [ + + "results", + + "cause" + + ], + + "poll": "commandStatus", + + "success": [ + + "Cancelled" + + ] + + } + + } + + } + + }, + + "components": { + + "schemas": { + + "Model": { + + "type": "object", + + "properties": { + + "name": { + + "type": "string" + + }, + + "dataplaneInfo": { + + "$ref": "#/components/schemas/DataPlaneInfo" + + } + + } + + }, + + "DataPlaneInfo": { + + "type": "object", + + "properties": { + + "endpointUrl": { + + "type": "string" + + }, + + "authorizationDetails": { + + "type": "string" + + } + + } + + } + + } + + } + + } \ No newline at end of file diff --git a/service/catalog/api.go b/service/catalog/api.go index fb20a6ef4..fcfa5423b 100755 --- a/service/catalog/api.go +++ b/service/catalog/api.go @@ -431,31 +431,23 @@ type AccountStorageCredentialsInterface interface { // // Gets a list of all storage credentials that have been assigned to given // metastore. - List(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) - - // StorageCredentialInfoNameToIdMap calls [AccountStorageCredentialsAPI.List] and creates a map of results with [StorageCredentialInfo].Name as key and [StorageCredentialInfo].Id as value. - // - // Returns an error if there's more than one [StorageCredentialInfo] with the same .Name. - // - // Note: All [StorageCredentialInfo] instances are loaded into memory before creating a map. // // This method is generated by Databricks SDK Code Generator. - StorageCredentialInfoNameToIdMap(ctx context.Context, request ListAccountStorageCredentialsRequest) (map[string]string, error) + List(ctx context.Context, request ListAccountStorageCredentialsRequest) listing.Iterator[StorageCredentialInfo] - // GetByName calls [AccountStorageCredentialsAPI.StorageCredentialInfoNameToIdMap] and returns a single [StorageCredentialInfo]. - // - // Returns an error if there's more than one [StorageCredentialInfo] with the same .Name. + // Get all storage credentials assigned to a metastore. // - // Note: All [StorageCredentialInfo] instances are loaded into memory before returning matching by name. + // Gets a list of all storage credentials that have been assigned to given + // metastore. // // This method is generated by Databricks SDK Code Generator. - GetByName(ctx context.Context, name string) (*StorageCredentialInfo, error) + ListAll(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) // Get all storage credentials assigned to a metastore. // // Gets a list of all storage credentials that have been assigned to given // metastore. - ListByMetastoreId(ctx context.Context, metastoreId string) ([]StorageCredentialInfo, error) + ListByMetastoreId(ctx context.Context, metastoreId string) (*ListAccountStorageCredentialsResponse, error) // Updates a storage credential. // @@ -552,68 +544,42 @@ func (a *AccountStorageCredentialsAPI) GetByMetastoreIdAndStorageCredentialName( // // Gets a list of all storage credentials that have been assigned to given // metastore. -func (a *AccountStorageCredentialsAPI) List(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) { - return a.impl.List(ctx, request) -} - -// StorageCredentialInfoNameToIdMap calls [AccountStorageCredentialsAPI.List] and creates a map of results with [StorageCredentialInfo].Name as key and [StorageCredentialInfo].Id as value. -// -// Returns an error if there's more than one [StorageCredentialInfo] with the same .Name. -// -// Note: All [StorageCredentialInfo] instances are loaded into memory before creating a map. // // This method is generated by Databricks SDK Code Generator. -func (a *AccountStorageCredentialsAPI) StorageCredentialInfoNameToIdMap(ctx context.Context, request ListAccountStorageCredentialsRequest) (map[string]string, error) { - ctx = useragent.InContext(ctx, "sdk-feature", "name-to-id") - mapping := map[string]string{} - result, err := a.List(ctx, request) - if err != nil { - return nil, err +func (a *AccountStorageCredentialsAPI) List(ctx context.Context, request ListAccountStorageCredentialsRequest) listing.Iterator[StorageCredentialInfo] { + + getNextPage := func(ctx context.Context, req ListAccountStorageCredentialsRequest) (*ListAccountStorageCredentialsResponse, error) { + ctx = useragent.InContext(ctx, "sdk-feature", "pagination") + return a.impl.List(ctx, req) } - for _, v := range result { - key := v.Name - _, duplicate := mapping[key] - if duplicate { - return nil, fmt.Errorf("duplicate .Name: %s", key) - } - mapping[key] = v.Id + getItems := func(resp *ListAccountStorageCredentialsResponse) []StorageCredentialInfo { + return resp.StorageCredentials } - return mapping, nil + + iterator := listing.NewIterator( + &request, + getNextPage, + getItems, + nil) + return iterator } -// GetByName calls [AccountStorageCredentialsAPI.StorageCredentialInfoNameToIdMap] and returns a single [StorageCredentialInfo]. -// -// Returns an error if there's more than one [StorageCredentialInfo] with the same .Name. +// Get all storage credentials assigned to a metastore. // -// Note: All [StorageCredentialInfo] instances are loaded into memory before returning matching by name. +// Gets a list of all storage credentials that have been assigned to given +// metastore. // // This method is generated by Databricks SDK Code Generator. -func (a *AccountStorageCredentialsAPI) GetByName(ctx context.Context, name string) (*StorageCredentialInfo, error) { - ctx = useragent.InContext(ctx, "sdk-feature", "get-by-name") - result, err := a.List(ctx, ListAccountStorageCredentialsRequest{}) - if err != nil { - return nil, err - } - tmp := map[string][]StorageCredentialInfo{} - for _, v := range result { - key := v.Name - tmp[key] = append(tmp[key], v) - } - alternatives, ok := tmp[name] - if !ok || len(alternatives) == 0 { - return nil, fmt.Errorf("StorageCredentialInfo named '%s' does not exist", name) - } - if len(alternatives) > 1 { - return nil, fmt.Errorf("there are %d instances of StorageCredentialInfo named '%s'", len(alternatives), name) - } - return &alternatives[0], nil +func (a *AccountStorageCredentialsAPI) ListAll(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) { + iterator := a.List(ctx, request) + return listing.ToSlice[StorageCredentialInfo](ctx, iterator) } // Get all storage credentials assigned to a metastore. // // Gets a list of all storage credentials that have been assigned to given // metastore. -func (a *AccountStorageCredentialsAPI) ListByMetastoreId(ctx context.Context, metastoreId string) ([]StorageCredentialInfo, error) { +func (a *AccountStorageCredentialsAPI) ListByMetastoreId(ctx context.Context, metastoreId string) (*ListAccountStorageCredentialsResponse, error) { return a.impl.List(ctx, ListAccountStorageCredentialsRequest{ MetastoreId: metastoreId, }) diff --git a/service/catalog/impl.go b/service/catalog/impl.go index b9e6905be..312d9522c 100755 --- a/service/catalog/impl.go +++ b/service/catalog/impl.go @@ -147,13 +147,13 @@ func (a *accountStorageCredentialsImpl) Get(ctx context.Context, request GetAcco return &accountsStorageCredentialInfo, err } -func (a *accountStorageCredentialsImpl) List(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) { - var storageCredentialInfoList []StorageCredentialInfo +func (a *accountStorageCredentialsImpl) List(ctx context.Context, request ListAccountStorageCredentialsRequest) (*ListAccountStorageCredentialsResponse, error) { + var listAccountStorageCredentialsResponse ListAccountStorageCredentialsResponse path := fmt.Sprintf("/api/2.0/accounts/%v/metastores/%v/storage-credentials", a.client.ConfiguredAccountID(), request.MetastoreId) headers := make(map[string]string) headers["Accept"] = "application/json" - err := a.client.Do(ctx, http.MethodGet, path, headers, request, &storageCredentialInfoList) - return storageCredentialInfoList, err + err := a.client.Do(ctx, http.MethodGet, path, headers, request, &listAccountStorageCredentialsResponse) + return &listAccountStorageCredentialsResponse, err } func (a *accountStorageCredentialsImpl) Update(ctx context.Context, request AccountsUpdateStorageCredential) (*AccountsStorageCredentialInfo, error) { diff --git a/service/catalog/interface.go b/service/catalog/interface.go index 170b8f656..f5c406a46 100755 --- a/service/catalog/interface.go +++ b/service/catalog/interface.go @@ -108,7 +108,9 @@ type AccountStorageCredentialsService interface { // // Gets a list of all storage credentials that have been assigned to given // metastore. - List(ctx context.Context, request ListAccountStorageCredentialsRequest) ([]StorageCredentialInfo, error) + // + // Use ListAll() to get all StorageCredentialInfo instances + List(ctx context.Context, request ListAccountStorageCredentialsRequest) (*ListAccountStorageCredentialsResponse, error) // Updates a storage credential. // diff --git a/service/catalog/model.go b/service/catalog/model.go index 6298e532b..2519046c1 100755 --- a/service/catalog/model.go +++ b/service/catalog/model.go @@ -2408,6 +2408,11 @@ type ListAccountStorageCredentialsRequest struct { MetastoreId string `json:"-" url:"-"` } +type ListAccountStorageCredentialsResponse struct { + // An array of metastore storage credentials. + StorageCredentials []StorageCredentialInfo `json:"storage_credentials,omitempty"` +} + // List catalogs type ListCatalogsRequest struct { // Whether to include catalogs in the response for which the principal can diff --git a/service/compute/model.go b/service/compute/model.go index f849f7998..354002147 100755 --- a/service/compute/model.go +++ b/service/compute/model.go @@ -455,7 +455,8 @@ type ClusterAttributes struct { // ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating // from legacy Passthrough on high concurrency clusters. * // `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy - // Passthrough on standard clusters. + // Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This + // mode provides a way that doesn’t have UC nor passthrough enabled. DataSecurityMode DataSecurityMode `json:"data_security_mode,omitempty"` DockerImage *DockerImage `json:"docker_image,omitempty"` @@ -613,7 +614,8 @@ type ClusterDetails struct { // ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating // from legacy Passthrough on high concurrency clusters. * // `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy - // Passthrough on standard clusters. + // Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This + // mode provides a way that doesn’t have UC nor passthrough enabled. DataSecurityMode DataSecurityMode `json:"data_security_mode,omitempty"` // Tags that are added by Databricks regardless of any `custom_tags`, // including: @@ -1156,7 +1158,8 @@ type ClusterSpec struct { // ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating // from legacy Passthrough on high concurrency clusters. * // `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy - // Passthrough on standard clusters. + // Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This + // mode provides a way that doesn’t have UC nor passthrough enabled. DataSecurityMode DataSecurityMode `json:"data_security_mode,omitempty"` DockerImage *DockerImage `json:"docker_image,omitempty"` @@ -1455,7 +1458,8 @@ type CreateCluster struct { // ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating // from legacy Passthrough on high concurrency clusters. * // `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy - // Passthrough on standard clusters. + // Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This + // mode provides a way that doesn’t have UC nor passthrough enabled. DataSecurityMode DataSecurityMode `json:"data_security_mode,omitempty"` DockerImage *DockerImage `json:"docker_image,omitempty"` @@ -1821,7 +1825,9 @@ func (f *DataPlaneEventDetailsEventType) Type() string { // * `LEGACY_TABLE_ACL`: This mode is for users migrating from legacy Table ACL // clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating from // legacy Passthrough on high concurrency clusters. * `LEGACY_SINGLE_USER`: This -// mode is for users migrating from legacy Passthrough on standard clusters. +// mode is for users migrating from legacy Passthrough on standard clusters. * +// `LEGACY_SINGLE_USER_STANDARD`: This mode provides a way that doesn’t have +// UC nor passthrough enabled. type DataSecurityMode string // This mode is for users migrating from legacy Passthrough on high concurrency @@ -1832,6 +1838,9 @@ const DataSecurityModeLegacyPassthrough DataSecurityMode = `LEGACY_PASSTHROUGH` // clusters. const DataSecurityModeLegacySingleUser DataSecurityMode = `LEGACY_SINGLE_USER` +// This mode provides a way that doesn’t have UC nor passthrough enabled. +const DataSecurityModeLegacySingleUserStandard DataSecurityMode = `LEGACY_SINGLE_USER_STANDARD` + // This mode is for users migrating from legacy Table ACL clusters. const DataSecurityModeLegacyTableAcl DataSecurityMode = `LEGACY_TABLE_ACL` @@ -1858,11 +1867,11 @@ func (f *DataSecurityMode) String() string { // Set raw string value and validate it against allowed values func (f *DataSecurityMode) Set(v string) error { switch v { - case `LEGACY_PASSTHROUGH`, `LEGACY_SINGLE_USER`, `LEGACY_TABLE_ACL`, `NONE`, `SINGLE_USER`, `USER_ISOLATION`: + case `LEGACY_PASSTHROUGH`, `LEGACY_SINGLE_USER`, `LEGACY_SINGLE_USER_STANDARD`, `LEGACY_TABLE_ACL`, `NONE`, `SINGLE_USER`, `USER_ISOLATION`: *f = DataSecurityMode(v) return nil default: - return fmt.Errorf(`value "%s" is not one of "LEGACY_PASSTHROUGH", "LEGACY_SINGLE_USER", "LEGACY_TABLE_ACL", "NONE", "SINGLE_USER", "USER_ISOLATION"`, v) + return fmt.Errorf(`value "%s" is not one of "LEGACY_PASSTHROUGH", "LEGACY_SINGLE_USER", "LEGACY_SINGLE_USER_STANDARD", "LEGACY_TABLE_ACL", "NONE", "SINGLE_USER", "USER_ISOLATION"`, v) } } @@ -2146,7 +2155,8 @@ type EditCluster struct { // ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating // from legacy Passthrough on high concurrency clusters. * // `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy - // Passthrough on standard clusters. + // Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This + // mode provides a way that doesn’t have UC nor passthrough enabled. DataSecurityMode DataSecurityMode `json:"data_security_mode,omitempty"` DockerImage *DockerImage `json:"docker_image,omitempty"` diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..c2f2067d1 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -329,7 +329,7 @@ type Group struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks group ID - Id string `json:"id,omitempty"` + Id string `json:"id,omitempty" url:"-"` Members []ComplexValue `json:"members,omitempty"` // Container for the group identifier. Workspace local versus account. @@ -1231,7 +1231,7 @@ type ServicePrincipal struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks service principal ID. - Id string `json:"id,omitempty"` + Id string `json:"id,omitempty" url:"-"` // Corresponds to AWS instance profile/arn role. Roles []ComplexValue `json:"roles,omitempty"` // The schema of the List response. @@ -1317,7 +1317,7 @@ type User struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks user ID. This is automatically set by Databricks. Any value // provided by the client will be ignored. - Id string `json:"id,omitempty" url:"-"` + Id string `json:"id,omitempty"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/jobs/model.go b/service/jobs/model.go index 5fc900323..3981a1f18 100755 --- a/service/jobs/model.go +++ b/service/jobs/model.go @@ -650,6 +650,8 @@ type ForEachTaskErrorMessageStats struct { Count int `json:"count,omitempty"` // Describes the error message occured during the iterations. ErrorMessage string `json:"error_message,omitempty"` + // Describes the termination reason for the error message. + TerminationCategory string `json:"termination_category,omitempty"` ForceSendFields []string `json:"-"` } diff --git a/service/marketplace/api.go b/service/marketplace/api.go index d6ea7e048..cc2749bd7 100755 --- a/service/marketplace/api.go +++ b/service/marketplace/api.go @@ -440,6 +440,12 @@ type ConsumerListingsInterface interface { // Deprecated: use MockConsumerListingsInterface instead. Impl() ConsumerListingsService + // Get one batch of listings. One may specify up to 50 IDs per request. + // + // Batch get a published listing in the Databricks Marketplace that the consumer + // has access to. + GEt(ctx context.Context, request BatchGetListingsRequest) (*BatchGetListingsResponse, error) + // Get listing. // // Get a published listing in the Databricks Marketplace that the consumer has @@ -535,6 +541,14 @@ func (a *ConsumerListingsAPI) Impl() ConsumerListingsService { return a.impl } +// Get one batch of listings. One may specify up to 50 IDs per request. +// +// Batch get a published listing in the Databricks Marketplace that the consumer +// has access to. +func (a *ConsumerListingsAPI) GEt(ctx context.Context, request BatchGetListingsRequest) (*BatchGetListingsResponse, error) { + return a.impl.GEt(ctx, request) +} + // Get listing. // // Get a published listing in the Databricks Marketplace that the consumer has @@ -836,6 +850,12 @@ type ConsumerProvidersInterface interface { // Deprecated: use MockConsumerProvidersInterface instead. Impl() ConsumerProvidersService + // Get one batch of providers. One may specify up to 50 IDs per request. + // + // Batch get a provider in the Databricks Marketplace with at least one visible + // listing. + GEt(ctx context.Context, request BatchGetProvidersRequest) (*BatchGetProvidersResponse, error) + // Get a provider. // // Get a provider in the Databricks Marketplace with at least one visible @@ -912,6 +932,14 @@ func (a *ConsumerProvidersAPI) Impl() ConsumerProvidersService { return a.impl } +// Get one batch of providers. One may specify up to 50 IDs per request. +// +// Batch get a provider in the Databricks Marketplace with at least one visible +// listing. +func (a *ConsumerProvidersAPI) GEt(ctx context.Context, request BatchGetProvidersRequest) (*BatchGetProvidersResponse, error) { + return a.impl.GEt(ctx, request) +} + // Get a provider. // // Get a provider in the Databricks Marketplace with at least one visible diff --git a/service/marketplace/impl.go b/service/marketplace/impl.go index a9c85828e..af972dc87 100755 --- a/service/marketplace/impl.go +++ b/service/marketplace/impl.go @@ -90,6 +90,15 @@ type consumerListingsImpl struct { client *client.DatabricksClient } +func (a *consumerListingsImpl) GEt(ctx context.Context, request BatchGetListingsRequest) (*BatchGetListingsResponse, error) { + var batchGetListingsResponse BatchGetListingsResponse + path := "/api/2.1/marketplace-consumer/listings:batchGet" + headers := make(map[string]string) + headers["Accept"] = "application/json" + err := a.client.Do(ctx, http.MethodGet, path, headers, request, &batchGetListingsResponse) + return &batchGetListingsResponse, err +} + func (a *consumerListingsImpl) Get(ctx context.Context, request GetListingRequest) (*GetListingResponse, error) { var getListingResponse GetListingResponse path := fmt.Sprintf("/api/2.1/marketplace-consumer/listings/%v", request.Id) @@ -155,6 +164,15 @@ type consumerProvidersImpl struct { client *client.DatabricksClient } +func (a *consumerProvidersImpl) GEt(ctx context.Context, request BatchGetProvidersRequest) (*BatchGetProvidersResponse, error) { + var batchGetProvidersResponse BatchGetProvidersResponse + path := "/api/2.1/marketplace-consumer/providers:batchGet" + headers := make(map[string]string) + headers["Accept"] = "application/json" + err := a.client.Do(ctx, http.MethodGet, path, headers, request, &batchGetProvidersResponse) + return &batchGetProvidersResponse, err +} + func (a *consumerProvidersImpl) Get(ctx context.Context, request GetProviderRequest) (*GetProviderResponse, error) { var getProviderResponse GetProviderResponse path := fmt.Sprintf("/api/2.1/marketplace-consumer/providers/%v", request.Id) diff --git a/service/marketplace/interface.go b/service/marketplace/interface.go index 51263c1e9..26f056d93 100755 --- a/service/marketplace/interface.go +++ b/service/marketplace/interface.go @@ -72,6 +72,12 @@ type ConsumerInstallationsService interface { // products that are available for consumption. type ConsumerListingsService interface { + // Get one batch of listings. One may specify up to 50 IDs per request. + // + // Batch get a published listing in the Databricks Marketplace that the + // consumer has access to. + GEt(ctx context.Context, request BatchGetListingsRequest) (*BatchGetListingsResponse, error) + // Get listing. // // Get a published listing in the Databricks Marketplace that the consumer @@ -122,6 +128,12 @@ type ConsumerPersonalizationRequestsService interface { // Providers are the entities that publish listings to the Marketplace. type ConsumerProvidersService interface { + // Get one batch of providers. One may specify up to 50 IDs per request. + // + // Batch get a provider in the Databricks Marketplace with at least one + // visible listing. + GEt(ctx context.Context, request BatchGetProvidersRequest) (*BatchGetProvidersResponse, error) + // Get a provider. // // Get a provider in the Databricks Marketplace with at least one visible diff --git a/service/marketplace/model.go b/service/marketplace/model.go index 5c1045d5d..efc78aa57 100755 --- a/service/marketplace/model.go +++ b/service/marketplace/model.go @@ -53,6 +53,24 @@ func (f *AssetType) Type() string { return "AssetType" } +// Get one batch of listings. One may specify up to 50 IDs per request. +type BatchGetListingsRequest struct { + Ids []string `json:"-" url:"ids,omitempty"` +} + +type BatchGetListingsResponse struct { + Listings []Listing `json:"listings,omitempty"` +} + +// Get one batch of providers. One may specify up to 50 IDs per request. +type BatchGetProvidersRequest struct { + Ids []string `json:"-" url:"ids,omitempty"` +} + +type BatchGetProvidersResponse struct { + Providers []ProviderInfo `json:"providers,omitempty"` +} + type Category string const CategoryAdvertisingAndMarketing Category = `ADVERTISING_AND_MARKETING` diff --git a/service/oauth2/data_plane.go b/service/oauth2/data_plane.go new file mode 100644 index 000000000..fec3ef155 --- /dev/null +++ b/service/oauth2/data_plane.go @@ -0,0 +1,50 @@ +package oauth2 + +import ( + "strings" + + "golang.org/x/oauth2" +) + +type DataPlaneHelper struct { + infos map[string]*DataPlaneInfo + tokens map[string]*oauth2.Token +} + +func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *oauth2.Token, error) { + if o.infos == nil { + o.infos = make(map[string]*DataPlaneInfo) + } + if o.tokens == nil { + o.tokens = make(map[string]*oauth2.Token) + } + key := o.generateKey(method, params) + info, infoOk := o.infos[key] + token, tokenOk := o.tokens[key] + if infoOk && tokenOk && token.Valid() { + return info.EndpointUrl, token, nil + } + if !infoOk { + newInfo, err := infoGetter() + if err != nil { + return "", nil, err + } + o.infos[key] = newInfo + info = newInfo + } + if !tokenOk || !token.Valid() { + newToken, err := refresh(info) + if err != nil { + return "", nil, err + } + o.tokens[key] = newToken + token = newToken + } + return info.EndpointUrl, token, nil +} + +func (o *DataPlaneHelper) generateKey(method string, params []string) string { + allElements := []string{method} + allElements = append(allElements, params...) + return strings.Join(allElements, "/") +} diff --git a/service/oauth2/data_plane_test.go b/service/oauth2/data_plane_test.go new file mode 100644 index 000000000..53647606c --- /dev/null +++ b/service/oauth2/data_plane_test.go @@ -0,0 +1,164 @@ +package oauth2 + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "golang.org/x/oauth2" +) + +type infoMock struct { + called bool + info *DataPlaneInfo + err error +} + +func (i *infoMock) DataPlaneInfoGetter() (*DataPlaneInfo, error) { + i.called = true + return i.info, i.err +} + +type tokenRefreshSpy struct { + called bool + expectedInfo *DataPlaneInfo + token *oauth2.Token + err error +} + +func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*oauth2.Token, error) { + t.expectedInfo = info + t.called = true + return t.token, t.err +} + +func TestTokenNotCached(t *testing.T) { + info := infoMock{ + info: &DataPlaneInfo{ + EndpointUrl: "url", + AuthorizationDetails: "authDetails", + }, + err: nil, + } + s := tokenRefreshSpy{ + token: &oauth2.Token{ + AccessToken: "token", + TokenType: "type", + Expiry: time.Now().Add(1 * time.Hour), + }, + err: nil, + } + c := DataPlaneHelper{} + url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + assert.NoError(t, err) + assert.Equal(t, "url", url) + assert.Equal(t, "token", token.AccessToken) + assert.Equal(t, "type", token.TokenType) + assert.True(t, token.Valid()) + assert.True(t, info.called) + assert.True(t, s.called) +} + +func TestTokenCached(t *testing.T) { + info := infoMock{ + info: &DataPlaneInfo{ + EndpointUrl: "url", + AuthorizationDetails: "authDetails", + }, + err: nil, + } + s := tokenRefreshSpy{ + token: &oauth2.Token{ + AccessToken: "token", + TokenType: "type", + Expiry: time.Now().Add(1 * time.Hour), + }, + err: nil, + } + c := DataPlaneHelper{} + c.infos = make(map[string]*DataPlaneInfo) + c.tokens = make(map[string]*oauth2.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) + assert.NoError(t, err) + assert.Equal(t, "url", url) + assert.Equal(t, "token", token.AccessToken) + assert.Equal(t, "type", token.TokenType) + assert.True(t, token.Valid()) + assert.False(t, info.called) + assert.False(t, s.called) +} + +func TestTokenExpired(t *testing.T) { + info := infoMock{ + info: &DataPlaneInfo{ + EndpointUrl: "url", + AuthorizationDetails: "authDetails", + }, + err: nil, + } + + expired := &oauth2.Token{ + AccessToken: "oldToken", + TokenType: "type", + Expiry: time.Now().Add(-1 * time.Hour), + } + s := tokenRefreshSpy{ + token: &oauth2.Token{ + AccessToken: "token", + TokenType: "type", + Expiry: time.Now().Add(1 * time.Hour), + }, + err: nil, + } + c := DataPlaneHelper{} + c.infos = make(map[string]*DataPlaneInfo) + c.tokens = make(map[string]*oauth2.Token) + c.infos["method/params"] = info.info + c.tokens["method/params"] = expired + url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + assert.NoError(t, err) + assert.Equal(t, "url", url) + assert.Equal(t, "token", token.AccessToken) + assert.Equal(t, "type", token.TokenType) + assert.True(t, token.Valid()) + assert.False(t, info.called) + assert.True(t, s.called) +} + +func TestTokenInfoError(t *testing.T) { + info := infoMock{ + info: nil, + err: assert.AnError, + } + s := tokenRefreshSpy{} + c := DataPlaneHelper{} + url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + assert.ErrorIs(t, err, assert.AnError) + assert.Empty(t, url) + assert.Nil(t, token) + assert.True(t, info.called) + assert.False(t, s.called) +} + +func TestTokenRefreshError(t *testing.T) { + info := infoMock{ + info: &DataPlaneInfo{ + EndpointUrl: "url", + AuthorizationDetails: "authDetails", + }, + err: nil, + } + s := tokenRefreshSpy{ + token: nil, + err: assert.AnError, + } + c := DataPlaneHelper{} + url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + assert.ErrorIs(t, err, assert.AnError) + assert.Empty(t, url) + assert.Nil(t, token) + assert.True(t, info.called) + assert.True(t, s.called) +} diff --git a/service/oauth2/model.go b/service/oauth2/model.go index defc930e2..b313af902 100755 --- a/service/oauth2/model.go +++ b/service/oauth2/model.go @@ -114,6 +114,23 @@ 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) +} + type DeleteCustomAppIntegrationOutput struct { } diff --git a/service/pkg.go b/service/pkg.go index a4e8181a9..e97b630c3 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -178,6 +178,8 @@ // // - [iam.AccountServicePrincipalsAPI]: Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD platforms. // +// - [serving.ServingEndpointDataPlaneAPI]: ServingEndpointDataPlane provides a set of operations to interact with DataPlane endpoints for ServingEndpoints service. +// // - [serving.ServingEndpointsAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. // // - [settings.SettingsAPI]: Workspace Settings API allows users to manage settings at the workspace level. @@ -344,6 +346,7 @@ var ( _ *oauth2.ServicePrincipalSecretsAPI = nil _ *iam.ServicePrincipalsAPI = nil _ *iam.AccountServicePrincipalsAPI = nil + _ *serving.ServingEndpointDataPlaneAPI = nil _ *serving.ServingEndpointsAPI = nil _ *settings.SettingsAPI = nil _ *settings.AccountSettingsAPI = nil diff --git a/service/serving/api.go b/service/serving/api.go index 2771bdba0..efb4e9a18 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -1,6 +1,6 @@ // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. -// These APIs allow you to manage Apps, Serving Endpoints, etc. +// These APIs allow you to manage Apps, Serving Endpoint Data Plane, Serving Endpoints, etc. package serving import ( @@ -552,6 +552,58 @@ func (a *AppsAPI) Update(ctx context.Context, request UpdateAppRequest) (*App, e return a.impl.Update(ctx, request) } +type ServingEndpointDataPlaneInterface interface { + // WithImpl could be used to override low-level API implementations for unit + // testing purposes with [github.com/golang/mock] or other mocking frameworks. + // Deprecated: use MockServingEndpointDataPlaneInterface instead. + WithImpl(impl ServingEndpointDataPlaneService) ServingEndpointDataPlaneInterface + + // Impl returns low-level ServingEndpointDataPlane API implementation + // Deprecated: use MockServingEndpointDataPlaneInterface instead. + Impl() ServingEndpointDataPlaneService + + // Query a serving endpoint. + Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) +} + +func NewServingEndpointDataPlane(client *client.DatabricksClient, + controlPlane *ServingEndpointsAPI, +) *ServingEndpointDataPlaneAPI { + return &ServingEndpointDataPlaneAPI{ + impl: &servingEndpointDataPlaneImpl{ + client: client, + controlPlane: controlPlane, + }, + } +} + +// ServingEndpointDataPlane provides a set of operations to interact with +// DataPlane endpoints for ServingEndpoints service. +type ServingEndpointDataPlaneAPI struct { + // impl contains low-level REST API interface, that could be overridden + // through WithImpl(ServingEndpointDataPlaneService) + impl ServingEndpointDataPlaneService +} + +// WithImpl could be used to override low-level API implementations for unit +// testing purposes with [github.com/golang/mock] or other mocking frameworks. +// Deprecated: use MockServingEndpointDataPlaneInterface instead. +func (a *ServingEndpointDataPlaneAPI) WithImpl(impl ServingEndpointDataPlaneService) ServingEndpointDataPlaneInterface { + a.impl = impl + return a +} + +// Impl returns low-level ServingEndpointDataPlane API implementation +// Deprecated: use MockServingEndpointDataPlaneInterface instead. +func (a *ServingEndpointDataPlaneAPI) Impl() ServingEndpointDataPlaneService { + return a.impl +} + +// Query a serving endpoint. +func (a *ServingEndpointDataPlaneAPI) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + return a.impl.Query(ctx, request) +} + type ServingEndpointsInterface interface { // WithImpl could be used to override low-level API implementations for unit // testing purposes with [github.com/golang/mock] or other mocking frameworks. diff --git a/service/serving/impl.go b/service/serving/impl.go index bbaee1b8f..b49ad6fd3 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -4,10 +4,15 @@ package serving import ( "context" + "errors" "fmt" "net/http" "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 Apps API methods @@ -109,6 +114,54 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, return &app, err } +// unexported type that holds implementations of just ServingEndpointDataPlane API methods +type servingEndpointDataPlaneImpl struct { + oauth2.DataPlaneHelper + controlPlane *ServingEndpointsAPI + client *client.DatabricksClient +} + +func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + getRequest := GetServingEndpointRequest{ + Name: request.Name, + } + token, err := a.client.Config.GetToken() + if err != nil { + return nil, err + } + infoGetter := func() (*oauth2.DataPlaneInfo, error) { + response, err := a.controlPlane.Get(ctx, getRequest) + if err != nil { + return nil, err + } + if response.DataPlaneInfo == nil { + return nil, errors.New("resource does not support direct Data Plane access") + } + return response.DataPlaneInfo.QueryInfo, nil + } + refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) + } + getParams := []string{ + request.Name, + } + endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("Query", getParams, refresh, infoGetter) + if err != nil { + return nil, err + } + headers := make(map[string]string) + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + opts := []httpclient.DoOption{} + opts = append(opts, httpclient.WithRequestHeaders(headers)) + var queryEndpointResponse QueryEndpointResponse + opts = append(opts, httpclient.WithRequestData(request)) + opts = append(opts, httpclient.WithResponseUnmarshal(&queryEndpointResponse)) + opts = append(opts, httpclient.WithToken(dataPlaneToken)) + err = a.client.ApiClient().Do(ctx, http.MethodPost, endpointUrl, opts...) + return &queryEndpointResponse, err +} + // unexported type that holds implementations of just ServingEndpoints API methods type servingEndpointsImpl struct { client *client.DatabricksClient diff --git a/service/serving/interface.go b/service/serving/interface.go index 6ee7f076f..dc29deb4f 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -67,6 +67,14 @@ type AppsService interface { Update(ctx context.Context, request UpdateAppRequest) (*App, error) } +// ServingEndpointDataPlane provides a set of operations to interact with +// DataPlane endpoints for ServingEndpoints service. +type ServingEndpointDataPlaneService interface { + + // Query a serving endpoint. + Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) +} + // The Serving Endpoints API allows you to create, update, and delete model // serving endpoints. // diff --git a/service/serving/model.go b/service/serving/model.go index 410e4e8ab..562dd08a2 100755 --- a/service/serving/model.go +++ b/service/serving/model.go @@ -7,6 +7,7 @@ import ( "io" "github.com/databricks/databricks-sdk-go/marshal" + "github.com/databricks/databricks-sdk-go/service/oauth2" ) type Ai21LabsConfig struct { @@ -966,6 +967,11 @@ type LogsRequest struct { ServedModelName string `json:"-" url:"-"` } +type ModelDataPlaneInfo struct { + // Information required to query DataPlane API 'query' endpoint. + QueryInfo *oauth2.DataPlaneInfo `json:"query_info,omitempty"` +} + type OpenAiConfig struct { // This field is only required for Azure AD OpenAI and is the Microsoft // Entra Client ID. @@ -1815,6 +1821,8 @@ type ServingEndpointDetailed struct { CreationTimestamp int64 `json:"creation_timestamp,omitempty"` // The email of the user who created the serving endpoint. Creator string `json:"creator,omitempty"` + // Information required to query DataPlane APIs. + DataPlaneInfo *ModelDataPlaneInfo `json:"data_plane_info,omitempty"` // Endpoint invocation url if route optimization is enabled for endpoint EndpointUrl string `json:"endpoint_url,omitempty"` // System-generated ID of the endpoint. This is used to refer to the diff --git a/service/sql/model.go b/service/sql/model.go index 431cdbdea..345658ff6 100755 --- a/service/sql/model.go +++ b/service/sql/model.go @@ -287,6 +287,7 @@ func (s ChannelInfo) MarshalJSON() ([]byte, error) { return marshal.Marshal(s) } +// Name of the channel type ChannelName string const ChannelNameChannelNameCurrent ChannelName = `CHANNEL_NAME_CURRENT` @@ -2352,8 +2353,7 @@ type QueryInfo struct { CanSubscribeToLiveQuery bool `json:"canSubscribeToLiveQuery,omitempty"` // Channel information for the SQL warehouse at the time of query execution ChannelUsed *ChannelInfo `json:"channel_used,omitempty"` - // Total execution time of the query from the client’s point of view, in - // milliseconds. + // Total execution time of the statement ( excluding result fetch time ). Duration int `json:"duration,omitempty"` // Alias for `warehouse_id`. EndpointId string `json:"endpoint_id,omitempty"` diff --git a/workspace_client.go b/workspace_client.go index 9d176242b..b0faa739b 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -3,12 +3,10 @@ package databricks import ( - "context" "errors" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/credentials" "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/service/catalog" @@ -698,6 +696,10 @@ type WorkspaceClient struct { // data by accident. ServicePrincipals iam.ServicePrincipalsInterface + // ServingEndpointDataPlane provides a set of operations to interact with + // DataPlane endpoints for ServingEndpoints service. + ServingEndpointDataPlane serving.ServingEndpointDataPlaneInterface + // The Serving Endpoints API allows you to create, update, and delete model // serving endpoints. // @@ -972,19 +974,6 @@ type WorkspaceClient struct { WorkspaceConf settings.WorkspaceConfInterface } -// Returns a new OAuth scoped to the authorization details provided. -// It will return an error if the CredentialStrategy does not support OAuth tokens. -// -// **NOTE:** Experimental: This API may change or be removed in a future release -// without warning. -func (a *WorkspaceClient) GetOAuthToken(ctx context.Context, authorizationDetails string) (*credentials.OAuthToken, error) { - originalToken, err := a.Config.GetToken() - if err != nil { - return nil, err - } - return a.apiClient.GetOAuthToken(ctx, authorizationDetails, originalToken) -} - var ErrNotWorkspaceClient = errors.New("invalid Databricks Workspace configuration") // NewWorkspaceClient creates new Databricks SDK client for Workspaces or @@ -1013,6 +1002,8 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { if err != nil { return nil, err } + + servingEndpoints := serving.NewServingEndpoints(databricksClient) return &WorkspaceClient{ Config: cfg, apiClient: apiClient, @@ -1080,7 +1071,8 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { Schemas: catalog.NewSchemas(databricksClient), Secrets: workspace.NewSecrets(databricksClient), ServicePrincipals: iam.NewServicePrincipals(databricksClient), - ServingEndpoints: serving.NewServingEndpoints(databricksClient), + ServingEndpointDataPlane: serving.NewServingEndpointDataPlane(databricksClient, servingEndpoints), + ServingEndpoints: servingEndpoints, Settings: settings.NewSettings(databricksClient), Shares: sharing.NewShares(databricksClient), StatementExecution: sql.NewStatementExecution(databricksClient),