From 0ed887f973d84052b2d8f4f7452dc8517d5fad43 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 27 May 2024 09:51:27 +0200 Subject: [PATCH 01/24] WIP --- .codegen/interface.go.tmpl | 13 ++ openapi/code/load_test.go | 9 ++ openapi/code/method.go | 5 + openapi/code/resource_client.go | 14 ++ openapi/code/service.go | 17 ++ openapi/extensions.go | 6 + openapi/model.go | 1 + openapi/testdata/spec_dataplane.json | 229 +++++++++++++++++++++++++++ 8 files changed, 294 insertions(+) create mode 100644 openapi/code/resource_client.go create mode 100644 openapi/testdata/spec_dataplane.json diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index 03189f9e5..e1f97cc48 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -21,4 +21,17 @@ type {{.PascalName}}Service interface { {{end}} } +{{if .HasDataPlaneMethods}} +type {{.Singular.PascalName}}Service interface { + {{range .DataPlaneMethods}} + {{.Comment " // " 80}} + {{- if .Pagination}} + // + // Use {{.PascalName}}All() to get all {{.Pagination.Entity.PascalName}} instances{{if .Pagination.MultiRequest}}, which will iterate over every result page.{{end}} + {{- end}} + {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} + {{end}} +} +{{end}} + {{end}} \ No newline at end of file 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..f00ca1fe0 100644 --- a/openapi/code/method.go +++ b/openapi/code/method.go @@ -46,9 +46,14 @@ 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 +} + // 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/resource_client.go b/openapi/code/resource_client.go new file mode 100644 index 000000000..6782641c4 --- /dev/null +++ b/openapi/code/resource_client.go @@ -0,0 +1,14 @@ +package code + +import "github.com/databricks/databricks-sdk-go/openapi" + +type ResourceClient struct { + Named + + PathStyle openapi.PathStyle + IsAccounts bool + Package *Package + methods map[string]*Method + ByPathParamsMethods []*Shortcut + tag *openapi.Tag +} diff --git a/openapi/code/service.go b/openapi/code/service.go index 49f8504ff..a40fda74f 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -33,6 +33,22 @@ type Service struct { tag *openapi.Tag } +// Returns whether any method supports direct DataPlane access. +func (svc *Service) HasDataPlaneMethods() bool { + return len(svc.DataPlaneMethods()) > 0 +} + +// Returns a sorted slice of methods which support direct DataPlane access. +func (svc *Service) DataPlaneMethods() (methods []*Method) { + for _, v := range svc.methods { + if v.DataPlane != nil { + methods = append(methods, v) + } + } + pascalNameSort(methods) + return methods +} + // FullName holds package name and service name func (svc *Service) FullName() string { return fmt.Sprintf("%s.%s", svc.Package.FullName(), svc.PascalName()) @@ -528,6 +544,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 From 773ea11e8dc7054d5612f4059073f4034fa6dfcc Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 27 May 2024 13:10:24 +0200 Subject: [PATCH 02/24] Generate --- .codegen/_openapi_sha | 2 +- ...k_account_storage_credentials_interface.go | 171 ++++++------------ service/catalog/api.go | 88 +++------ service/catalog/impl.go | 8 +- service/catalog/interface.go | 4 +- service/catalog/model.go | 5 + service/compute/model.go | 26 ++- service/iam/model.go | 4 +- service/sql/model.go | 3 +- 9 files changed, 113 insertions(+), 198 deletions(-) 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/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/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..b4aa4acd5 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. @@ -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/sql/model.go b/service/sql/model.go index 431cdbdea..e08608ba1 100755 --- a/service/sql/model.go +++ b/service/sql/model.go @@ -2352,8 +2352,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"` From 2407299eb791a77b50e2ff56465cd569f7ce5234 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 27 May 2024 14:27:14 +0200 Subject: [PATCH 03/24] Test --- .codegen/api.go.tmpl | 4 +- .codegen/interface.go.tmpl | 13 +- .gitattributes | 2 + ...k_serving_endpoint_data_plane_interface.go | 254 ++++++++++++++++++ .../mock_serving_endpoint_interface.go | 96 +++++++ openapi/code/package.go | 3 + openapi/code/service.go | 45 +++- service/iam/model.go | 4 +- service/oauth2/model.go | 17 ++ service/serving/model.go | 8 + 10 files changed, 425 insertions(+), 21 deletions(-) create mode 100644 experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go create mode 100644 experimental/mocks/service/serving/mock_serving_endpoint_interface.go diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 11f089775..57dcceecd 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -148,6 +148,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 +429,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/interface.go.tmpl b/.codegen/interface.go.tmpl index e1f97cc48..a2eb01566 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -21,17 +21,6 @@ type {{.PascalName}}Service interface { {{end}} } -{{if .HasDataPlaneMethods}} -type {{.Singular.PascalName}}Service interface { - {{range .DataPlaneMethods}} - {{.Comment " // " 80}} - {{- if .Pagination}} - // - // Use {{.PascalName}}All() to get all {{.Pagination.Entity.PascalName}} instances{{if .Pagination.MultiRequest}}, which will iterate over every result page.{{end}} - {{- end}} - {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} - {{end}} -} -{{end}} + {{end}} \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index 9c869d613..e95df8c99 100644 --- a/.gitattributes +++ b/.gitattributes @@ -79,6 +79,8 @@ 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_endpoint_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/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..25d305850 --- /dev/null +++ b/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go @@ -0,0 +1,254 @@ +// 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" + + time "time" +) + +// 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 +} + +// WaitGetServingEndpointNotUpdating provides a mock function with given fields: ctx, name, timeout, callback +func (_m *MockServingEndpointDataPlaneInterface) WaitGetServingEndpointNotUpdating(ctx context.Context, name string, timeout time.Duration, callback func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { + ret := _m.Called(ctx, name, timeout, callback) + + if len(ret) == 0 { + panic("no return value specified for WaitGetServingEndpointNotUpdating") + } + + var r0 *serving.ServingEndpointDetailed + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error)); ok { + return rf(ctx, name, timeout, callback) + } + if rf, ok := ret.Get(0).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) *serving.ServingEndpointDetailed); ok { + r0 = rf(ctx, name, timeout, callback) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*serving.ServingEndpointDetailed) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) error); ok { + r1 = rf(ctx, name, timeout, callback) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WaitGetServingEndpointNotUpdating' +type MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call struct { + *mock.Call +} + +// WaitGetServingEndpointNotUpdating is a helper method to define mock.On call +// - ctx context.Context +// - name string +// - timeout time.Duration +// - callback func(*serving.ServingEndpointDetailed) +func (_e *MockServingEndpointDataPlaneInterface_Expecter) WaitGetServingEndpointNotUpdating(ctx interface{}, name interface{}, timeout interface{}, callback interface{}) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { + return &MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call{Call: _e.mock.On("WaitGetServingEndpointNotUpdating", ctx, name, timeout, callback)} +} + +func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) Run(run func(ctx context.Context, name string, timeout time.Duration, callback func(*serving.ServingEndpointDetailed))) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(time.Duration), args[3].(func(*serving.ServingEndpointDetailed))) + }) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) Return(_a0 *serving.ServingEndpointDetailed, _a1 error) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) RunAndReturn(run func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error)) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_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/experimental/mocks/service/serving/mock_serving_endpoint_interface.go b/experimental/mocks/service/serving/mock_serving_endpoint_interface.go new file mode 100644 index 000000000..38b187ccb --- /dev/null +++ b/experimental/mocks/service/serving/mock_serving_endpoint_interface.go @@ -0,0 +1,96 @@ +// 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" +) + +// MockServingEndpointInterface is an autogenerated mock type for the ServingEndpointInterface type +type MockServingEndpointInterface struct { + mock.Mock +} + +type MockServingEndpointInterface_Expecter struct { + mock *mock.Mock +} + +func (_m *MockServingEndpointInterface) EXPECT() *MockServingEndpointInterface_Expecter { + return &MockServingEndpointInterface_Expecter{mock: &_m.Mock} +} + +// Query provides a mock function with given fields: ctx, request +func (_m *MockServingEndpointInterface) 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 +} + +// MockServingEndpointInterface_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query' +type MockServingEndpointInterface_Query_Call struct { + *mock.Call +} + +// Query is a helper method to define mock.On call +// - ctx context.Context +// - request serving.QueryEndpointInput +func (_e *MockServingEndpointInterface_Expecter) Query(ctx interface{}, request interface{}) *MockServingEndpointInterface_Query_Call { + return &MockServingEndpointInterface_Query_Call{Call: _e.mock.On("Query", ctx, request)} +} + +func (_c *MockServingEndpointInterface_Query_Call) Run(run func(ctx context.Context, request serving.QueryEndpointInput)) *MockServingEndpointInterface_Query_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(serving.QueryEndpointInput)) + }) + return _c +} + +func (_c *MockServingEndpointInterface_Query_Call) Return(_a0 *serving.QueryEndpointResponse, _a1 error) *MockServingEndpointInterface_Query_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockServingEndpointInterface_Query_Call) RunAndReturn(run func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)) *MockServingEndpointInterface_Query_Call { + _c.Call.Return(run) + return _c +} + +// NewMockServingEndpointInterface creates a new instance of MockServingEndpointInterface. 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 NewMockServingEndpointInterface(t interface { + mock.TestingT + Cleanup(func()) +}) *MockServingEndpointInterface { + mock := &MockServingEndpointInterface{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/openapi/code/package.go b/openapi/code/package.go index 4ee549fa6..04eb9c6a0 100644 --- a/openapi/code/package.go +++ b/openapi/code/package.go @@ -32,6 +32,9 @@ func (pkg *Package) FullName() string { func (pkg *Package) Services() (types []*Service) { for _, v := range pkg.services { types = append(types, v) + if v.HasDataPlaneMethods() { + types = append(types, v.DataPlaneService()) + } } pascalNameSort(types) return types diff --git a/openapi/code/service.go b/openapi/code/service.go index a40fda74f..d5810b627 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -35,17 +35,50 @@ type Service struct { // Returns whether any method supports direct DataPlane access. func (svc *Service) HasDataPlaneMethods() bool { - return len(svc.DataPlaneMethods()) > 0 + return len(svc.dataPlaneMethods()) > 0 +} + +// Returns the corresponding service for DataPlane APIs. +func (svc *Service) DataPlaneService() *Service { + s := &Service{ + Named: Named{svc.Name + "DataPlane", svc.Description}, + Package: svc.Package, + methods: svc.dataPlaneMethods(), + tag: svc.tag, + ByPathParamsMethods: svc.ByPathParamsMethods, + } + for _, m := range s.methods { + m.Service = s + } + return s } // Returns a sorted slice of methods which support direct DataPlane access. -func (svc *Service) DataPlaneMethods() (methods []*Method) { - for _, v := range svc.methods { - if v.DataPlane != nil { - methods = append(methods, v) +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, + } } } - pascalNameSort(methods) return methods } diff --git a/service/iam/model.go b/service/iam/model.go index b4aa4acd5..55e4ccf6b 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" url:"-"` + Id string `json:"id,omitempty"` Members []ComplexValue `json:"members,omitempty"` // Container for the group identifier. Workspace local versus account. @@ -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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. 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/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 From dc6b88230a003a321f03f4872548c55167de9a8d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 08:20:02 +0200 Subject: [PATCH 04/24] Test2 --- .codegen/accounts.go.tmpl | 4 +- .codegen/api.go.tmpl | 4 + .codegen/interface.go.tmpl | 4 + .codegen/mock_account_client.go.tmpl | 4 +- .codegen/mock_workspace_client.go.tmpl | 6 +- .codegen/workspaces.go.tmpl | 4 +- .gitattributes | 1 + ..._serving_endpoints_data_plane_interface.go | 191 ++++++++++++++++++ openapi/code/service.go | 2 + service/iam/model.go | 4 +- service/pkg.go | 3 + service/serving/api.go | 50 ++++- service/serving/impl.go | 15 ++ service/serving/interface.go | 7 + service/sql/model.go | 1 + 15 files changed, 288 insertions(+), 12 deletions(-) create mode 100644 experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go diff --git a/.codegen/accounts.go.tmpl b/.codegen/accounts.go.tmpl index dbc4f3df5..6138f84bb 100644 --- a/.codegen/accounts.go.tmpl +++ b/.codegen/accounts.go.tmpl @@ -12,7 +12,7 @@ import ( type AccountClient struct { Config *config.Config - {{range .Services}}{{if and .IsAccounts (not .HasParent)}} + {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} {{.Comment " // " 80}} {{(.TrimPrefix "account").PascalName}} {{.Package.Name}}.{{.Name}}Interface {{end}}{{end}} @@ -44,7 +44,7 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { } return &AccountClient{ Config: cfg, - {{range .Services}}{{if and .IsAccounts (not .HasParent)}} + {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient),{{end}}{{end}} }, nil } diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 57dcceecd..dbee3ec31 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -116,7 +116,11 @@ func New{{.PascalName}}(client *client.DatabricksClient) *{{.PascalName}}API { } } +{{- if not .IsDataPlane}} {{.Comment "// " 80}} +{{else}} +// TODO comment +{{end -}} type {{.PascalName}}API struct { // impl contains low-level REST API interface, that could be overridden // through WithImpl({{.PascalName}}Service) diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index a2eb01566..082d10d20 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -9,7 +9,11 @@ import ( ) {{range .Services}} +{{- if not .IsDataPlane}} {{.Comment "// " 80}} +{{else}} +// TODO comment +{{end -}} type {{.PascalName}}Service interface { {{range .Methods}} {{.Comment " // " 80}} diff --git a/.codegen/mock_account_client.go.tmpl b/.codegen/mock_account_client.go.tmpl index 0f668d176..035cfc542 100644 --- a/.codegen/mock_account_client.go.tmpl +++ b/.codegen/mock_account_client.go.tmpl @@ -22,7 +22,7 @@ func NewMockAccountClient(t interface { cli := &MockAccountClient { AccountClient: &databricks.AccountClient{ Config: nil, - {{range .Services}}{{if and .IsAccounts (not .HasParent)}} + {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} {{(.TrimPrefix "account").PascalName}}: {{ template "mock-interface-constructor" . }}(t), {{- end}}{{end}} }, @@ -50,7 +50,7 @@ func(m *MockAccountClient) GetMock{{.PascalName}}API() *{{ template "mock-interf } {{end}}{{end}} -{{range .Services}}{{if and .IsAccounts (not .HasParent)}} +{{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} func(m *MockAccountClient) GetMock{{.Name}}API() *{{ template "mock-interface-name" . }} { api, ok := m.AccountClient.{{(.TrimPrefix "account").PascalName}}.(*{{ template "mock-interface-name" . }}) if !ok { diff --git a/.codegen/mock_workspace_client.go.tmpl b/.codegen/mock_workspace_client.go.tmpl index ef6d49e3b..ef4aad8f2 100644 --- a/.codegen/mock_workspace_client.go.tmpl +++ b/.codegen/mock_workspace_client.go.tmpl @@ -22,7 +22,7 @@ func NewMockWorkspaceClient(t interface { cli := &MockWorkspaceClient{ WorkspaceClient: &databricks.WorkspaceClient{ Config: nil, - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} {{.Name}}: {{ template "mock-interface-constructor" . }}(t), {{- end}}{{end}} }, @@ -40,7 +40,7 @@ func NewMockWorkspaceClient(t interface { return cli } -{{range .Services}}{{if and (not .IsAccounts) (.HasParent)}} +{{range .Services}}{{if and (not .IsAccounts) (.HasParent) (not .IsDataPlane)}} func(m *MockWorkspaceClient) GetMock{{.PascalName}}API() *{{ template "mock-interface-name" . }} { api, ok := m.GetMock{{.ParentService.Name}}API().{{.PascalName}}().(*{{ template "mock-interface-name" . }}) if !ok { @@ -51,7 +51,7 @@ func(m *MockWorkspaceClient) GetMock{{.PascalName}}API() *{{ template "mock-inte {{end}}{{end}} -{{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} +{{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} func(m *MockWorkspaceClient) GetMock{{.Name}}API() *{{ template "mock-interface-name" . }} { api, ok := m.WorkspaceClient.{{.Name}}.(*{{ template "mock-interface-name" . }}) if !ok { diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index d4db23cfc..4aa223dfa 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.codegen/workspaces.go.tmpl @@ -14,7 +14,7 @@ type WorkspaceClient struct { Config *config.Config apiClient *httpclient.ApiClient - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} {{.Comment " // " 80}} {{.Name}} {{.Package.Name}}.{{.Name}}Interface {{end}}{{end}} @@ -64,7 +64,7 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { return &WorkspaceClient{ Config: cfg, apiClient: apiClient, - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient), {{- end}}{{end}} }, nil diff --git a/.gitattributes b/.gitattributes index e95df8c99..2d7de9dad 100644 --- a/.gitattributes +++ b/.gitattributes @@ -81,6 +81,7 @@ experimental/mocks/service/provisioning/mock_workspaces_interface.go linguist-ge 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_endpoint_interface.go linguist-generated=true +experimental/mocks/service/serving/mock_serving_endpoints_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/experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go b/experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go new file mode 100644 index 000000000..e5441d026 --- /dev/null +++ b/experimental/mocks/service/serving/mock_serving_endpoints_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" +) + +// MockServingEndpointsDataPlaneInterface is an autogenerated mock type for the ServingEndpointsDataPlaneInterface type +type MockServingEndpointsDataPlaneInterface struct { + mock.Mock +} + +type MockServingEndpointsDataPlaneInterface_Expecter struct { + mock *mock.Mock +} + +func (_m *MockServingEndpointsDataPlaneInterface) EXPECT() *MockServingEndpointsDataPlaneInterface_Expecter { + return &MockServingEndpointsDataPlaneInterface_Expecter{mock: &_m.Mock} +} + +// Impl provides a mock function with given fields: +func (_m *MockServingEndpointsDataPlaneInterface) Impl() serving.ServingEndpointsDataPlaneService { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for Impl") + } + + var r0 serving.ServingEndpointsDataPlaneService + if rf, ok := ret.Get(0).(func() serving.ServingEndpointsDataPlaneService); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serving.ServingEndpointsDataPlaneService) + } + } + + return r0 +} + +// MockServingEndpointsDataPlaneInterface_Impl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Impl' +type MockServingEndpointsDataPlaneInterface_Impl_Call struct { + *mock.Call +} + +// Impl is a helper method to define mock.On call +func (_e *MockServingEndpointsDataPlaneInterface_Expecter) Impl() *MockServingEndpointsDataPlaneInterface_Impl_Call { + return &MockServingEndpointsDataPlaneInterface_Impl_Call{Call: _e.mock.On("Impl")} +} + +func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) Run(run func()) *MockServingEndpointsDataPlaneInterface_Impl_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) Return(_a0 serving.ServingEndpointsDataPlaneService) *MockServingEndpointsDataPlaneInterface_Impl_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) RunAndReturn(run func() serving.ServingEndpointsDataPlaneService) *MockServingEndpointsDataPlaneInterface_Impl_Call { + _c.Call.Return(run) + return _c +} + +// Query provides a mock function with given fields: ctx, request +func (_m *MockServingEndpointsDataPlaneInterface) 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 +} + +// MockServingEndpointsDataPlaneInterface_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query' +type MockServingEndpointsDataPlaneInterface_Query_Call struct { + *mock.Call +} + +// Query is a helper method to define mock.On call +// - ctx context.Context +// - request serving.QueryEndpointInput +func (_e *MockServingEndpointsDataPlaneInterface_Expecter) Query(ctx interface{}, request interface{}) *MockServingEndpointsDataPlaneInterface_Query_Call { + return &MockServingEndpointsDataPlaneInterface_Query_Call{Call: _e.mock.On("Query", ctx, request)} +} + +func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) Run(run func(ctx context.Context, request serving.QueryEndpointInput)) *MockServingEndpointsDataPlaneInterface_Query_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(serving.QueryEndpointInput)) + }) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) Return(_a0 *serving.QueryEndpointResponse, _a1 error) *MockServingEndpointsDataPlaneInterface_Query_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) RunAndReturn(run func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)) *MockServingEndpointsDataPlaneInterface_Query_Call { + _c.Call.Return(run) + return _c +} + +// WithImpl provides a mock function with given fields: impl +func (_m *MockServingEndpointsDataPlaneInterface) WithImpl(impl serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface { + ret := _m.Called(impl) + + if len(ret) == 0 { + panic("no return value specified for WithImpl") + } + + var r0 serving.ServingEndpointsDataPlaneInterface + if rf, ok := ret.Get(0).(func(serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface); ok { + r0 = rf(impl) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serving.ServingEndpointsDataPlaneInterface) + } + } + + return r0 +} + +// MockServingEndpointsDataPlaneInterface_WithImpl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WithImpl' +type MockServingEndpointsDataPlaneInterface_WithImpl_Call struct { + *mock.Call +} + +// WithImpl is a helper method to define mock.On call +// - impl serving.ServingEndpointsDataPlaneService +func (_e *MockServingEndpointsDataPlaneInterface_Expecter) WithImpl(impl interface{}) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { + return &MockServingEndpointsDataPlaneInterface_WithImpl_Call{Call: _e.mock.On("WithImpl", impl)} +} + +func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) Run(run func(impl serving.ServingEndpointsDataPlaneService)) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(serving.ServingEndpointsDataPlaneService)) + }) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) Return(_a0 serving.ServingEndpointsDataPlaneInterface) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) RunAndReturn(run func(serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { + _c.Call.Return(run) + return _c +} + +// NewMockServingEndpointsDataPlaneInterface creates a new instance of MockServingEndpointsDataPlaneInterface. 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 NewMockServingEndpointsDataPlaneInterface(t interface { + mock.TestingT + Cleanup(func()) +}) *MockServingEndpointsDataPlaneInterface { + mock := &MockServingEndpointsDataPlaneInterface{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/openapi/code/service.go b/openapi/code/service.go index d5810b627..f916acb7e 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -31,6 +31,7 @@ type Service struct { ByPathParamsMethods []*Shortcut ParentService *Service tag *openapi.Tag + IsDataPlane bool } // Returns whether any method supports direct DataPlane access. @@ -46,6 +47,7 @@ func (svc *Service) DataPlaneService() *Service { methods: svc.dataPlaneMethods(), tag: svc.tag, ByPathParamsMethods: svc.ByPathParamsMethods, + IsDataPlane: true, } for _, m := range s.methods { m.Service = s diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..52d3de756 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -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/pkg.go b/service/pkg.go index a4e8181a9..5637f43e9 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -180,6 +180,8 @@ // // - [serving.ServingEndpointsAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. // +// - [serving.ServingEndpointsDataPlaneAPI]: 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. // // - [settings.AccountSettingsAPI]: Accounts Settings API allows users to manage settings at the account level. @@ -345,6 +347,7 @@ var ( _ *iam.ServicePrincipalsAPI = nil _ *iam.AccountServicePrincipalsAPI = nil _ *serving.ServingEndpointsAPI = nil + _ *serving.ServingEndpointsDataPlaneAPI = nil _ *settings.SettingsAPI = nil _ *settings.AccountSettingsAPI = nil _ *sharing.SharesAPI = nil diff --git a/service/serving/api.go b/service/serving/api.go index 2771bdba0..8d8d0d456 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 Endpoints, Serving Endpoints Data Plane, etc. package serving import ( @@ -1103,3 +1103,51 @@ func (a *ServingEndpointsAPI) UpdateConfigAndWait(ctx context.Context, endpointC func (a *ServingEndpointsAPI) UpdatePermissions(ctx context.Context, request ServingEndpointPermissionsRequest) (*ServingEndpointPermissions, error) { return a.impl.UpdatePermissions(ctx, request) } + +type ServingEndpointsDataPlaneInterface 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 MockServingEndpointsDataPlaneInterface instead. + WithImpl(impl ServingEndpointsDataPlaneService) ServingEndpointsDataPlaneInterface + + // Impl returns low-level ServingEndpointsDataPlane API implementation + // Deprecated: use MockServingEndpointsDataPlaneInterface instead. + Impl() ServingEndpointsDataPlaneService + + // Query a serving endpoint. + Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) +} + +func NewServingEndpointsDataPlane(client *client.DatabricksClient) *ServingEndpointsDataPlaneAPI { + return &ServingEndpointsDataPlaneAPI{ + impl: &servingEndpointsDataPlaneImpl{ + client: client, + }, + } +} + +// TODO comment +type ServingEndpointsDataPlaneAPI struct { + // impl contains low-level REST API interface, that could be overridden + // through WithImpl(ServingEndpointsDataPlaneService) + impl ServingEndpointsDataPlaneService +} + +// 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 MockServingEndpointsDataPlaneInterface instead. +func (a *ServingEndpointsDataPlaneAPI) WithImpl(impl ServingEndpointsDataPlaneService) ServingEndpointsDataPlaneInterface { + a.impl = impl + return a +} + +// Impl returns low-level ServingEndpointsDataPlane API implementation +// Deprecated: use MockServingEndpointsDataPlaneInterface instead. +func (a *ServingEndpointsDataPlaneAPI) Impl() ServingEndpointsDataPlaneService { + return a.impl +} + +// Query a serving endpoint. +func (a *ServingEndpointsDataPlaneAPI) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + return a.impl.Query(ctx, request) +} diff --git a/service/serving/impl.go b/service/serving/impl.go index bbaee1b8f..774ca4791 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -264,3 +264,18 @@ func (a *servingEndpointsImpl) UpdatePermissions(ctx context.Context, request Se err := a.client.Do(ctx, http.MethodPatch, path, headers, request, &servingEndpointPermissions) return &servingEndpointPermissions, err } + +// unexported type that holds implementations of just ServingEndpointsDataPlane API methods +type servingEndpointsDataPlaneImpl struct { + client *client.DatabricksClient +} + +func (a *servingEndpointsDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + var queryEndpointResponse QueryEndpointResponse + path := fmt.Sprintf("/serving-endpoints/%v/invocations", request.Name) + headers := make(map[string]string) + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + err := a.client.Do(ctx, http.MethodPost, path, headers, request, &queryEndpointResponse) + return &queryEndpointResponse, err +} diff --git a/service/serving/interface.go b/service/serving/interface.go index 6ee7f076f..fc6744f4c 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -167,3 +167,10 @@ type ServingEndpointsService interface { // inherit permissions from their root object. UpdatePermissions(ctx context.Context, request ServingEndpointPermissionsRequest) (*ServingEndpointPermissions, error) } + +// TODO comment +type ServingEndpointsDataPlaneService interface { + + // Query a serving endpoint. + Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) +} diff --git a/service/sql/model.go b/service/sql/model.go index e08608ba1..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` From 9d1213586a925afc75bd7e5bb7d843a6136feced Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 08:40:14 +0200 Subject: [PATCH 05/24] singular --- .gitattributes | 3 - ...k_serving_endpoint_data_plane_interface.go | 63 ------ .../mock_serving_endpoint_interface.go | 96 --------- ..._serving_endpoints_data_plane_interface.go | 191 ------------------ openapi/code/service.go | 2 +- service/iam/model.go | 4 +- service/pkg.go | 6 +- service/serving/api.go | 98 ++++----- service/serving/impl.go | 30 +-- service/serving/interface.go | 14 +- 10 files changed, 77 insertions(+), 430 deletions(-) delete mode 100644 experimental/mocks/service/serving/mock_serving_endpoint_interface.go delete mode 100644 experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go diff --git a/.gitattributes b/.gitattributes index 2d7de9dad..9c869d613 100644 --- a/.gitattributes +++ b/.gitattributes @@ -79,9 +79,6 @@ 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_endpoint_interface.go linguist-generated=true -experimental/mocks/service/serving/mock_serving_endpoints_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/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go b/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go index 25d305850..4f581953e 100644 --- a/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go +++ b/experimental/mocks/service/serving/mock_serving_endpoint_data_plane_interface.go @@ -7,8 +7,6 @@ import ( serving "github.com/databricks/databricks-sdk-go/service/serving" mock "github.com/stretchr/testify/mock" - - time "time" ) // MockServingEndpointDataPlaneInterface is an autogenerated mock type for the ServingEndpointDataPlaneInterface type @@ -130,67 +128,6 @@ func (_c *MockServingEndpointDataPlaneInterface_Query_Call) RunAndReturn(run fun return _c } -// WaitGetServingEndpointNotUpdating provides a mock function with given fields: ctx, name, timeout, callback -func (_m *MockServingEndpointDataPlaneInterface) WaitGetServingEndpointNotUpdating(ctx context.Context, name string, timeout time.Duration, callback func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error) { - ret := _m.Called(ctx, name, timeout, callback) - - if len(ret) == 0 { - panic("no return value specified for WaitGetServingEndpointNotUpdating") - } - - var r0 *serving.ServingEndpointDetailed - var r1 error - if rf, ok := ret.Get(0).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error)); ok { - return rf(ctx, name, timeout, callback) - } - if rf, ok := ret.Get(0).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) *serving.ServingEndpointDetailed); ok { - r0 = rf(ctx, name, timeout, callback) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(*serving.ServingEndpointDetailed) - } - } - - if rf, ok := ret.Get(1).(func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) error); ok { - r1 = rf(ctx, name, timeout, callback) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WaitGetServingEndpointNotUpdating' -type MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call struct { - *mock.Call -} - -// WaitGetServingEndpointNotUpdating is a helper method to define mock.On call -// - ctx context.Context -// - name string -// - timeout time.Duration -// - callback func(*serving.ServingEndpointDetailed) -func (_e *MockServingEndpointDataPlaneInterface_Expecter) WaitGetServingEndpointNotUpdating(ctx interface{}, name interface{}, timeout interface{}, callback interface{}) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { - return &MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call{Call: _e.mock.On("WaitGetServingEndpointNotUpdating", ctx, name, timeout, callback)} -} - -func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) Run(run func(ctx context.Context, name string, timeout time.Duration, callback func(*serving.ServingEndpointDetailed))) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(string), args[2].(time.Duration), args[3].(func(*serving.ServingEndpointDetailed))) - }) - return _c -} - -func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) Return(_a0 *serving.ServingEndpointDetailed, _a1 error) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_Call) RunAndReturn(run func(context.Context, string, time.Duration, func(*serving.ServingEndpointDetailed)) (*serving.ServingEndpointDetailed, error)) *MockServingEndpointDataPlaneInterface_WaitGetServingEndpointNotUpdating_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) diff --git a/experimental/mocks/service/serving/mock_serving_endpoint_interface.go b/experimental/mocks/service/serving/mock_serving_endpoint_interface.go deleted file mode 100644 index 38b187ccb..000000000 --- a/experimental/mocks/service/serving/mock_serving_endpoint_interface.go +++ /dev/null @@ -1,96 +0,0 @@ -// 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" -) - -// MockServingEndpointInterface is an autogenerated mock type for the ServingEndpointInterface type -type MockServingEndpointInterface struct { - mock.Mock -} - -type MockServingEndpointInterface_Expecter struct { - mock *mock.Mock -} - -func (_m *MockServingEndpointInterface) EXPECT() *MockServingEndpointInterface_Expecter { - return &MockServingEndpointInterface_Expecter{mock: &_m.Mock} -} - -// Query provides a mock function with given fields: ctx, request -func (_m *MockServingEndpointInterface) 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 -} - -// MockServingEndpointInterface_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query' -type MockServingEndpointInterface_Query_Call struct { - *mock.Call -} - -// Query is a helper method to define mock.On call -// - ctx context.Context -// - request serving.QueryEndpointInput -func (_e *MockServingEndpointInterface_Expecter) Query(ctx interface{}, request interface{}) *MockServingEndpointInterface_Query_Call { - return &MockServingEndpointInterface_Query_Call{Call: _e.mock.On("Query", ctx, request)} -} - -func (_c *MockServingEndpointInterface_Query_Call) Run(run func(ctx context.Context, request serving.QueryEndpointInput)) *MockServingEndpointInterface_Query_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(serving.QueryEndpointInput)) - }) - return _c -} - -func (_c *MockServingEndpointInterface_Query_Call) Return(_a0 *serving.QueryEndpointResponse, _a1 error) *MockServingEndpointInterface_Query_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockServingEndpointInterface_Query_Call) RunAndReturn(run func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)) *MockServingEndpointInterface_Query_Call { - _c.Call.Return(run) - return _c -} - -// NewMockServingEndpointInterface creates a new instance of MockServingEndpointInterface. 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 NewMockServingEndpointInterface(t interface { - mock.TestingT - Cleanup(func()) -}) *MockServingEndpointInterface { - mock := &MockServingEndpointInterface{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -} diff --git a/experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go b/experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go deleted file mode 100644 index e5441d026..000000000 --- a/experimental/mocks/service/serving/mock_serving_endpoints_data_plane_interface.go +++ /dev/null @@ -1,191 +0,0 @@ -// 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" -) - -// MockServingEndpointsDataPlaneInterface is an autogenerated mock type for the ServingEndpointsDataPlaneInterface type -type MockServingEndpointsDataPlaneInterface struct { - mock.Mock -} - -type MockServingEndpointsDataPlaneInterface_Expecter struct { - mock *mock.Mock -} - -func (_m *MockServingEndpointsDataPlaneInterface) EXPECT() *MockServingEndpointsDataPlaneInterface_Expecter { - return &MockServingEndpointsDataPlaneInterface_Expecter{mock: &_m.Mock} -} - -// Impl provides a mock function with given fields: -func (_m *MockServingEndpointsDataPlaneInterface) Impl() serving.ServingEndpointsDataPlaneService { - ret := _m.Called() - - if len(ret) == 0 { - panic("no return value specified for Impl") - } - - var r0 serving.ServingEndpointsDataPlaneService - if rf, ok := ret.Get(0).(func() serving.ServingEndpointsDataPlaneService); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(serving.ServingEndpointsDataPlaneService) - } - } - - return r0 -} - -// MockServingEndpointsDataPlaneInterface_Impl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Impl' -type MockServingEndpointsDataPlaneInterface_Impl_Call struct { - *mock.Call -} - -// Impl is a helper method to define mock.On call -func (_e *MockServingEndpointsDataPlaneInterface_Expecter) Impl() *MockServingEndpointsDataPlaneInterface_Impl_Call { - return &MockServingEndpointsDataPlaneInterface_Impl_Call{Call: _e.mock.On("Impl")} -} - -func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) Run(run func()) *MockServingEndpointsDataPlaneInterface_Impl_Call { - _c.Call.Run(func(args mock.Arguments) { - run() - }) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) Return(_a0 serving.ServingEndpointsDataPlaneService) *MockServingEndpointsDataPlaneInterface_Impl_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_Impl_Call) RunAndReturn(run func() serving.ServingEndpointsDataPlaneService) *MockServingEndpointsDataPlaneInterface_Impl_Call { - _c.Call.Return(run) - return _c -} - -// Query provides a mock function with given fields: ctx, request -func (_m *MockServingEndpointsDataPlaneInterface) 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 -} - -// MockServingEndpointsDataPlaneInterface_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query' -type MockServingEndpointsDataPlaneInterface_Query_Call struct { - *mock.Call -} - -// Query is a helper method to define mock.On call -// - ctx context.Context -// - request serving.QueryEndpointInput -func (_e *MockServingEndpointsDataPlaneInterface_Expecter) Query(ctx interface{}, request interface{}) *MockServingEndpointsDataPlaneInterface_Query_Call { - return &MockServingEndpointsDataPlaneInterface_Query_Call{Call: _e.mock.On("Query", ctx, request)} -} - -func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) Run(run func(ctx context.Context, request serving.QueryEndpointInput)) *MockServingEndpointsDataPlaneInterface_Query_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(context.Context), args[1].(serving.QueryEndpointInput)) - }) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) Return(_a0 *serving.QueryEndpointResponse, _a1 error) *MockServingEndpointsDataPlaneInterface_Query_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_Query_Call) RunAndReturn(run func(context.Context, serving.QueryEndpointInput) (*serving.QueryEndpointResponse, error)) *MockServingEndpointsDataPlaneInterface_Query_Call { - _c.Call.Return(run) - return _c -} - -// WithImpl provides a mock function with given fields: impl -func (_m *MockServingEndpointsDataPlaneInterface) WithImpl(impl serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface { - ret := _m.Called(impl) - - if len(ret) == 0 { - panic("no return value specified for WithImpl") - } - - var r0 serving.ServingEndpointsDataPlaneInterface - if rf, ok := ret.Get(0).(func(serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface); ok { - r0 = rf(impl) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(serving.ServingEndpointsDataPlaneInterface) - } - } - - return r0 -} - -// MockServingEndpointsDataPlaneInterface_WithImpl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WithImpl' -type MockServingEndpointsDataPlaneInterface_WithImpl_Call struct { - *mock.Call -} - -// WithImpl is a helper method to define mock.On call -// - impl serving.ServingEndpointsDataPlaneService -func (_e *MockServingEndpointsDataPlaneInterface_Expecter) WithImpl(impl interface{}) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { - return &MockServingEndpointsDataPlaneInterface_WithImpl_Call{Call: _e.mock.On("WithImpl", impl)} -} - -func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) Run(run func(impl serving.ServingEndpointsDataPlaneService)) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(serving.ServingEndpointsDataPlaneService)) - }) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) Return(_a0 serving.ServingEndpointsDataPlaneInterface) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockServingEndpointsDataPlaneInterface_WithImpl_Call) RunAndReturn(run func(serving.ServingEndpointsDataPlaneService) serving.ServingEndpointsDataPlaneInterface) *MockServingEndpointsDataPlaneInterface_WithImpl_Call { - _c.Call.Return(run) - return _c -} - -// NewMockServingEndpointsDataPlaneInterface creates a new instance of MockServingEndpointsDataPlaneInterface. 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 NewMockServingEndpointsDataPlaneInterface(t interface { - mock.TestingT - Cleanup(func()) -}) *MockServingEndpointsDataPlaneInterface { - mock := &MockServingEndpointsDataPlaneInterface{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -} diff --git a/openapi/code/service.go b/openapi/code/service.go index f916acb7e..f88bd3399 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -42,7 +42,7 @@ func (svc *Service) HasDataPlaneMethods() bool { // Returns the corresponding service for DataPlane APIs. func (svc *Service) DataPlaneService() *Service { s := &Service{ - Named: Named{svc.Name + "DataPlane", svc.Description}, + Named: Named{svc.Named.Singular().Name + "DataPlane", svc.Description}, Package: svc.Package, methods: svc.dataPlaneMethods(), tag: svc.tag, diff --git a/service/iam/model.go b/service/iam/model.go index 52d3de756..b4aa4acd5 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" url:"-"` + Id string `json:"id,omitempty"` // Corresponds to AWS instance profile/arn role. Roles []ComplexValue `json:"roles,omitempty"` // The schema of the List response. diff --git a/service/pkg.go b/service/pkg.go index 5637f43e9..aaea39732 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -178,9 +178,9 @@ // // - [iam.AccountServicePrincipalsAPI]: Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD platforms. // -// - [serving.ServingEndpointsAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. +// - [serving.ServingEndpointDataPlaneAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. // -// - [serving.ServingEndpointsDataPlaneAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. +// - [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. // @@ -346,8 +346,8 @@ var ( _ *oauth2.ServicePrincipalSecretsAPI = nil _ *iam.ServicePrincipalsAPI = nil _ *iam.AccountServicePrincipalsAPI = nil + _ *serving.ServingEndpointDataPlaneAPI = nil _ *serving.ServingEndpointsAPI = nil - _ *serving.ServingEndpointsDataPlaneAPI = nil _ *settings.SettingsAPI = nil _ *settings.AccountSettingsAPI = nil _ *sharing.SharesAPI = nil diff --git a/service/serving/api.go b/service/serving/api.go index 8d8d0d456..e262683c5 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, Serving Endpoints Data Plane, etc. +// These APIs allow you to manage Apps, Serving Endpoint Data Plane, Serving Endpoints, etc. package serving import ( @@ -552,6 +552,54 @@ 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) *ServingEndpointDataPlaneAPI { + return &ServingEndpointDataPlaneAPI{ + impl: &servingEndpointDataPlaneImpl{ + client: client, + }, + } +} + +// TODO comment +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. @@ -1103,51 +1151,3 @@ func (a *ServingEndpointsAPI) UpdateConfigAndWait(ctx context.Context, endpointC func (a *ServingEndpointsAPI) UpdatePermissions(ctx context.Context, request ServingEndpointPermissionsRequest) (*ServingEndpointPermissions, error) { return a.impl.UpdatePermissions(ctx, request) } - -type ServingEndpointsDataPlaneInterface 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 MockServingEndpointsDataPlaneInterface instead. - WithImpl(impl ServingEndpointsDataPlaneService) ServingEndpointsDataPlaneInterface - - // Impl returns low-level ServingEndpointsDataPlane API implementation - // Deprecated: use MockServingEndpointsDataPlaneInterface instead. - Impl() ServingEndpointsDataPlaneService - - // Query a serving endpoint. - Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) -} - -func NewServingEndpointsDataPlane(client *client.DatabricksClient) *ServingEndpointsDataPlaneAPI { - return &ServingEndpointsDataPlaneAPI{ - impl: &servingEndpointsDataPlaneImpl{ - client: client, - }, - } -} - -// TODO comment -type ServingEndpointsDataPlaneAPI struct { - // impl contains low-level REST API interface, that could be overridden - // through WithImpl(ServingEndpointsDataPlaneService) - impl ServingEndpointsDataPlaneService -} - -// 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 MockServingEndpointsDataPlaneInterface instead. -func (a *ServingEndpointsDataPlaneAPI) WithImpl(impl ServingEndpointsDataPlaneService) ServingEndpointsDataPlaneInterface { - a.impl = impl - return a -} - -// Impl returns low-level ServingEndpointsDataPlane API implementation -// Deprecated: use MockServingEndpointsDataPlaneInterface instead. -func (a *ServingEndpointsDataPlaneAPI) Impl() ServingEndpointsDataPlaneService { - return a.impl -} - -// Query a serving endpoint. -func (a *ServingEndpointsDataPlaneAPI) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { - return a.impl.Query(ctx, request) -} diff --git a/service/serving/impl.go b/service/serving/impl.go index 774ca4791..2671f773e 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -109,6 +109,21 @@ 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 { + client *client.DatabricksClient +} + +func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + var queryEndpointResponse QueryEndpointResponse + path := fmt.Sprintf("/serving-endpoints/%v/invocations", request.Name) + headers := make(map[string]string) + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + err := a.client.Do(ctx, http.MethodPost, path, headers, request, &queryEndpointResponse) + return &queryEndpointResponse, err +} + // unexported type that holds implementations of just ServingEndpoints API methods type servingEndpointsImpl struct { client *client.DatabricksClient @@ -264,18 +279,3 @@ func (a *servingEndpointsImpl) UpdatePermissions(ctx context.Context, request Se err := a.client.Do(ctx, http.MethodPatch, path, headers, request, &servingEndpointPermissions) return &servingEndpointPermissions, err } - -// unexported type that holds implementations of just ServingEndpointsDataPlane API methods -type servingEndpointsDataPlaneImpl struct { - client *client.DatabricksClient -} - -func (a *servingEndpointsDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { - var queryEndpointResponse QueryEndpointResponse - path := fmt.Sprintf("/serving-endpoints/%v/invocations", request.Name) - headers := make(map[string]string) - headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" - err := a.client.Do(ctx, http.MethodPost, path, headers, request, &queryEndpointResponse) - return &queryEndpointResponse, err -} diff --git a/service/serving/interface.go b/service/serving/interface.go index fc6744f4c..c891f9407 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -67,6 +67,13 @@ type AppsService interface { Update(ctx context.Context, request UpdateAppRequest) (*App, error) } +// TODO comment +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. // @@ -167,10 +174,3 @@ type ServingEndpointsService interface { // inherit permissions from their root object. UpdatePermissions(ctx context.Context, request ServingEndpointPermissionsRequest) (*ServingEndpointPermissions, error) } - -// TODO comment -type ServingEndpointsDataPlaneService interface { - - // Query a serving endpoint. - Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) -} From 75345a37e51d5b21f681850cfe3f8806e2e24e1a Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 08:57:15 +0200 Subject: [PATCH 06/24] Client getter --- .codegen/api.go.tmpl | 12 ++++ .gitattributes | 1 + .../mock_serving_endpoints_interface.go | 57 +++++++++++++++++++ service/iam/model.go | 4 +- service/serving/api.go | 8 +++ 5 files changed, 80 insertions(+), 2 deletions(-) diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index dbee3ec31..6d8594b18 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -42,6 +42,11 @@ type {{.PascalName}}Interface interface { {{.PascalName}}() {{.PascalName}}Interface {{end}} + {{if and (not .IsDataPlane) .HasDataPlaneMethods}} + // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API + Get{{.Singular.PascalName}}DataPlaneClient() ({{.Singular.PascalName}}DataPlaneService, error) + {{end}} + {{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 {{.PascalName}}(ctx context.Context{{range .Binding}}, {{.PollField.CamelName}} {{template "type" .PollField.Entity}}{{end}}, @@ -138,6 +143,13 @@ func (a *{{.ParentService.PascalName}}API) {{.PascalName}}() {{.PascalName}}Inte } {{end}} +{{if and (not .IsDataPlane) .HasDataPlaneMethods}} + // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API +func (a *{{.PascalName}}API) Get{{.Singular.PascalName}}DataPlaneClient() ({{.Singular.PascalName}}DataPlaneService, error) { + return nil, nil +} +{{end}} + // 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 Mock{{.PascalName}}Interface instead. 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/experimental/mocks/service/serving/mock_serving_endpoints_interface.go b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go index 496304aec..957c1098a 100644 --- a/experimental/mocks/service/serving/mock_serving_endpoints_interface.go +++ b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go @@ -940,6 +940,63 @@ func (_c *MockServingEndpointsInterface_GetPermissionsByServingEndpointId_Call) return _c } +// GetServingEndpointDataPlaneClient provides a mock function with given fields: +func (_m *MockServingEndpointsInterface) GetServingEndpointDataPlaneClient() (serving.ServingEndpointDataPlaneService, error) { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for GetServingEndpointDataPlaneClient") + } + + var r0 serving.ServingEndpointDataPlaneService + var r1 error + if rf, ok := ret.Get(0).(func() (serving.ServingEndpointDataPlaneService, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() serving.ServingEndpointDataPlaneService); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serving.ServingEndpointDataPlaneService) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetServingEndpointDataPlaneClient' +type MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call struct { + *mock.Call +} + +// GetServingEndpointDataPlaneClient is a helper method to define mock.On call +func (_e *MockServingEndpointsInterface_Expecter) GetServingEndpointDataPlaneClient() *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { + return &MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call{Call: _e.mock.On("GetServingEndpointDataPlaneClient")} +} + +func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Run(run func()) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Return(_a0 serving.ServingEndpointDataPlaneService, _a1 error) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) RunAndReturn(run func() (serving.ServingEndpointDataPlaneService, error)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { + _c.Call.Return(run) + return _c +} + // Impl provides a mock function with given fields: func (_m *MockServingEndpointsInterface) Impl() serving.ServingEndpointsService { ret := _m.Called() diff --git a/service/iam/model.go b/service/iam/model.go index b4aa4acd5..55e4ccf6b 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" url:"-"` + Id string `json:"id,omitempty"` Members []ComplexValue `json:"members,omitempty"` // Container for the group identifier. Workspace local versus account. @@ -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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/serving/api.go b/service/serving/api.go index e262683c5..e9f9833fe 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -610,6 +610,9 @@ type ServingEndpointsInterface interface { // Deprecated: use MockServingEndpointsInterface instead. Impl() ServingEndpointsService + // Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API + GetServingEndpointDataPlaneClient() (ServingEndpointDataPlaneService, error) + // WaitGetServingEndpointNotUpdating repeatedly calls [ServingEndpointsAPI.Get] and waits to reach NOT_UPDATING state WaitGetServingEndpointNotUpdating(ctx context.Context, name string, timeout time.Duration, callback func(*ServingEndpointDetailed)) (*ServingEndpointDetailed, error) @@ -790,6 +793,11 @@ type ServingEndpointsAPI struct { impl ServingEndpointsService } +// Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API +func (a *ServingEndpointsAPI) GetServingEndpointDataPlaneClient() (ServingEndpointDataPlaneService, error) { + return nil, nil +} + // 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 MockServingEndpointsInterface instead. From be639aece630422526fbf68216ca192f517a0759 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 09:03:58 +0200 Subject: [PATCH 07/24] DP method --- openapi/code/service.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/openapi/code/service.go b/openapi/code/service.go index f88bd3399..643598d09 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -39,6 +39,16 @@ func (svc *Service) HasDataPlaneMethods() bool { return len(svc.dataPlaneMethods()) > 0 } +func (svc *Service) DataPlaneInfoMethod() *Method { + methodName := "" + for _, m := range svc.methods { + if m.DataPlane != nil { + methodName = m.DataPlane.ConfigMethod + } + } + return svc.methods[methodName] +} + // Returns the corresponding service for DataPlane APIs. func (svc *Service) DataPlaneService() *Service { s := &Service{ From ecd2bc0ee5dbc0116518f9617c79280f8742bc00 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 09:09:32 +0200 Subject: [PATCH 08/24] Request --- .codegen/api.go.tmpl | 4 +-- .../mock_serving_endpoints_interface.go | 29 ++++++++++--------- service/iam/model.go | 4 +-- service/serving/api.go | 4 +-- 4 files changed, 21 insertions(+), 20 deletions(-) diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 6d8594b18..817fe7329 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -44,7 +44,7 @@ type {{.PascalName}}Interface interface { {{if and (not .IsDataPlane) .HasDataPlaneMethods}} // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API - Get{{.Singular.PascalName}}DataPlaneClient() ({{.Singular.PascalName}}DataPlaneService, error) + Get{{.Singular.PascalName}}DataPlaneClient(request {{.DataPlaneInfoMethod.Request.PascalName}}) ({{.Singular.PascalName}}DataPlaneService, error) {{end}} {{range .Waits}} @@ -145,7 +145,7 @@ func (a *{{.ParentService.PascalName}}API) {{.PascalName}}() {{.PascalName}}Inte {{if and (not .IsDataPlane) .HasDataPlaneMethods}} // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API -func (a *{{.PascalName}}API) Get{{.Singular.PascalName}}DataPlaneClient() ({{.Singular.PascalName}}DataPlaneService, error) { +func (a *{{.PascalName}}API) Get{{.Singular.PascalName}}DataPlaneClient(request {{.DataPlaneInfoMethod.Request.PascalName}}) ({{.Singular.PascalName}}DataPlaneService, error) { return nil, nil } {{end}} diff --git a/experimental/mocks/service/serving/mock_serving_endpoints_interface.go b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go index 957c1098a..318a63006 100644 --- a/experimental/mocks/service/serving/mock_serving_endpoints_interface.go +++ b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go @@ -940,9 +940,9 @@ func (_c *MockServingEndpointsInterface_GetPermissionsByServingEndpointId_Call) return _c } -// GetServingEndpointDataPlaneClient provides a mock function with given fields: -func (_m *MockServingEndpointsInterface) GetServingEndpointDataPlaneClient() (serving.ServingEndpointDataPlaneService, error) { - ret := _m.Called() +// GetServingEndpointDataPlaneClient provides a mock function with given fields: request +func (_m *MockServingEndpointsInterface) GetServingEndpointDataPlaneClient(request serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error) { + ret := _m.Called(request) if len(ret) == 0 { panic("no return value specified for GetServingEndpointDataPlaneClient") @@ -950,19 +950,19 @@ func (_m *MockServingEndpointsInterface) GetServingEndpointDataPlaneClient() (se var r0 serving.ServingEndpointDataPlaneService var r1 error - if rf, ok := ret.Get(0).(func() (serving.ServingEndpointDataPlaneService, error)); ok { - return rf() + if rf, ok := ret.Get(0).(func(serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error)); ok { + return rf(request) } - if rf, ok := ret.Get(0).(func() serving.ServingEndpointDataPlaneService); ok { - r0 = rf() + if rf, ok := ret.Get(0).(func(serving.GetServingEndpointRequest) serving.ServingEndpointDataPlaneService); ok { + r0 = rf(request) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(serving.ServingEndpointDataPlaneService) } } - if rf, ok := ret.Get(1).(func() error); ok { - r1 = rf() + if rf, ok := ret.Get(1).(func(serving.GetServingEndpointRequest) error); ok { + r1 = rf(request) } else { r1 = ret.Error(1) } @@ -976,13 +976,14 @@ type MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call struct } // GetServingEndpointDataPlaneClient is a helper method to define mock.On call -func (_e *MockServingEndpointsInterface_Expecter) GetServingEndpointDataPlaneClient() *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { - return &MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call{Call: _e.mock.On("GetServingEndpointDataPlaneClient")} +// - request serving.GetServingEndpointRequest +func (_e *MockServingEndpointsInterface_Expecter) GetServingEndpointDataPlaneClient(request interface{}) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { + return &MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call{Call: _e.mock.On("GetServingEndpointDataPlaneClient", request)} } -func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Run(run func()) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { +func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Run(run func(request serving.GetServingEndpointRequest)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { _c.Call.Run(func(args mock.Arguments) { - run() + run(args[0].(serving.GetServingEndpointRequest)) }) return _c } @@ -992,7 +993,7 @@ func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) return _c } -func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) RunAndReturn(run func() (serving.ServingEndpointDataPlaneService, error)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { +func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) RunAndReturn(run func(serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { _c.Call.Return(run) return _c } diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..b4aa4acd5 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. @@ -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/serving/api.go b/service/serving/api.go index e9f9833fe..a3e7a66d5 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -611,7 +611,7 @@ type ServingEndpointsInterface interface { Impl() ServingEndpointsService // Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API - GetServingEndpointDataPlaneClient() (ServingEndpointDataPlaneService, error) + GetServingEndpointDataPlaneClient(request GetServingEndpointRequest) (ServingEndpointDataPlaneService, error) // WaitGetServingEndpointNotUpdating repeatedly calls [ServingEndpointsAPI.Get] and waits to reach NOT_UPDATING state WaitGetServingEndpointNotUpdating(ctx context.Context, name string, @@ -794,7 +794,7 @@ type ServingEndpointsAPI struct { } // Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API -func (a *ServingEndpointsAPI) GetServingEndpointDataPlaneClient() (ServingEndpointDataPlaneService, error) { +func (a *ServingEndpointsAPI) GetServingEndpointDataPlaneClient(request GetServingEndpointRequest) (ServingEndpointDataPlaneService, error) { return nil, nil } From c361d9d64d5bee6e1e95ed509a0f30aae9aaa496 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 11:26:33 +0200 Subject: [PATCH 09/24] CP client --- .codegen/api.go.tmpl | 9 ++++++++- .codegen/impl.go.tmpl | 3 +++ openapi/code/service.go | 1 + service/iam/model.go | 4 ++-- service/serving/api.go | 7 +++++-- service/serving/impl.go | 3 ++- service/sql/model.go | 1 - 7 files changed, 21 insertions(+), 7 deletions(-) diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 817fe7329..cf545a96e 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -110,10 +110,17 @@ type {{.PascalName}}Interface interface { {{end -}} } -func New{{.PascalName}}(client *client.DatabricksClient) *{{.PascalName}}API { +func New{{.PascalName}}(client *client.DatabricksClient, +{{- if .IsDataPlane}} +controlPlane *{{.ParentService.PascalName}}Interface, +{{end -}} +) *{{.PascalName}}API { return &{{.PascalName}}API{ impl: &{{.CamelName}}Impl{ client: client, + {{- if .IsDataPlane}} + controlPlane: controlPlane, + {{end}} }, {{range .Subservices}} {{.CamelName}}: New{{.PascalName}}(client), diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 811111567..7464a61df 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -17,6 +17,9 @@ import ( // unexported type that holds implementations of just {{.Name}} API methods type {{.CamelName}}Impl struct { client *client.DatabricksClient + {{- if .IsDataPlane}} + controlPlane *{{.ParentService.PascalName}}Interface + {{end -}} } {{range .Methods}} diff --git a/openapi/code/service.go b/openapi/code/service.go index 643598d09..8ca49a404 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -54,6 +54,7 @@ func (svc *Service) DataPlaneService() *Service { s := &Service{ Named: Named{svc.Named.Singular().Name + "DataPlane", svc.Description}, Package: svc.Package, + ParentService: svc, methods: svc.dataPlaneMethods(), tag: svc.tag, ByPathParamsMethods: svc.ByPathParamsMethods, diff --git a/service/iam/model.go b/service/iam/model.go index b4aa4acd5..c0e595616 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/serving/api.go b/service/serving/api.go index a3e7a66d5..6615fda39 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -566,10 +566,13 @@ type ServingEndpointDataPlaneInterface interface { Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) } -func NewServingEndpointDataPlane(client *client.DatabricksClient) *ServingEndpointDataPlaneAPI { +func NewServingEndpointDataPlane(client *client.DatabricksClient, + controlPlane *ServingEndpointsInterface, +) *ServingEndpointDataPlaneAPI { return &ServingEndpointDataPlaneAPI{ impl: &servingEndpointDataPlaneImpl{ - client: client, + client: client, + controlPlane: controlPlane, }, } } diff --git a/service/serving/impl.go b/service/serving/impl.go index 2671f773e..5c80828fe 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -111,7 +111,8 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, // unexported type that holds implementations of just ServingEndpointDataPlane API methods type servingEndpointDataPlaneImpl struct { - client *client.DatabricksClient + client *client.DatabricksClient + controlPlane *ServingEndpointsInterface } func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { diff --git a/service/sql/model.go b/service/sql/model.go index 345658ff6..e08608ba1 100755 --- a/service/sql/model.go +++ b/service/sql/model.go @@ -287,7 +287,6 @@ func (s ChannelInfo) MarshalJSON() ([]byte, error) { return marshal.Marshal(s) } -// Name of the channel type ChannelName string const ChannelNameChannelNameCurrent ChannelName = `CHANNEL_NAME_CURRENT` From ff60fc19454a1e992c83b5c17ba18e0f14a1d516 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 13:19:08 +0200 Subject: [PATCH 10/24] delete stuff --- openapi/code/resource_client.go | 14 -------------- 1 file changed, 14 deletions(-) delete mode 100644 openapi/code/resource_client.go diff --git a/openapi/code/resource_client.go b/openapi/code/resource_client.go deleted file mode 100644 index 6782641c4..000000000 --- a/openapi/code/resource_client.go +++ /dev/null @@ -1,14 +0,0 @@ -package code - -import "github.com/databricks/databricks-sdk-go/openapi" - -type ResourceClient struct { - Named - - PathStyle openapi.PathStyle - IsAccounts bool - Package *Package - methods map[string]*Method - ByPathParamsMethods []*Shortcut - tag *openapi.Tag -} From a12d718527a7c209bdff87e2671c67900271d037 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 13:54:17 +0200 Subject: [PATCH 11/24] cleanup --- .codegen/api.go.tmpl | 2 +- .codegen/impl.go.tmpl | 12 +++++++++- .codegen/workspaces.go.tmpl | 13 ----------- client/client.go | 10 ++++++++ data_plane/data_plane.go | 41 +++++++++++++++++++++++++++++++++ httpclient/oauth_token.go | 10 +++++--- openapi/code/resource_client.go | 14 +++++++++++ service/iam/model.go | 4 ++-- service/serving/api.go | 2 +- service/serving/impl.go | 2 +- workspace_client.go | 15 ------------ 11 files changed, 88 insertions(+), 37 deletions(-) create mode 100644 data_plane/data_plane.go create mode 100644 openapi/code/resource_client.go diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index cf545a96e..f00a1fbbd 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -112,7 +112,7 @@ type {{.PascalName}}Interface interface { func New{{.PascalName}}(client *client.DatabricksClient, {{- if .IsDataPlane}} -controlPlane *{{.ParentService.PascalName}}Interface, +controlPlane *{{.ParentService.PascalName}}API, {{end -}} ) *{{.PascalName}}API { return &{{.PascalName}}API{ diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 7464a61df..8dce64966 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -18,11 +18,12 @@ import ( type {{.CamelName}}Impl struct { client *client.DatabricksClient {{- if .IsDataPlane}} - controlPlane *{{.ParentService.PascalName}}Interface + controlPlane *{{.ParentService.PascalName}}API {{end -}} } {{range .Methods}} +{{if .Service.IsDataPlane}} func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} { {{- template "response-var" . }} path := {{ template "path" . }} @@ -30,6 +31,15 @@ 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" . }} } +{{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" . }} + {{ template "make-header" . }} + 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/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index 4aa223dfa..7880532e7 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 diff --git a/client/client.go b/client/client.go index 671ccf09a..140f16920 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,15 @@ func (c *DatabricksClient) ConfiguredAccountID() string { return c.Config.AccountID } +// 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/data_plane/data_plane.go b/data_plane/data_plane.go new file mode 100644 index 000000000..f7a58bda8 --- /dev/null +++ b/data_plane/data_plane.go @@ -0,0 +1,41 @@ +package dataplane + +import ( + "strings" + + dp "github.com/databricks/databricks-sdk-go/service/oauth2" + "golang.org/x/oauth2" +) + +type DataPlaneTokenCache struct { + infos map[string]*dp.DataPlaneInfo + tokens map[string]*oauth2.Token +} + +func (o *DataPlaneTokenCache) GetDataPlane(method string, params []string, refresh func(*dp.DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*dp.DataPlaneInfo, error)) (string, *oauth2.Token, error) { + 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 + } + newToken, err := refresh(info) + if err != nil { + return "", nil, err + } + o.tokens[method] = newToken + return info.EndpointUrl, newToken, nil +} + +func (o *DataPlaneTokenCache) generateKey(method string, params []string) string { + allElements := append(params, method) + return strings.Join(allElements, "/") +} 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/openapi/code/resource_client.go b/openapi/code/resource_client.go new file mode 100644 index 000000000..6782641c4 --- /dev/null +++ b/openapi/code/resource_client.go @@ -0,0 +1,14 @@ +package code + +import "github.com/databricks/databricks-sdk-go/openapi" + +type ResourceClient struct { + Named + + PathStyle openapi.PathStyle + IsAccounts bool + Package *Package + methods map[string]*Method + ByPathParamsMethods []*Shortcut + tag *openapi.Tag +} diff --git a/service/iam/model.go b/service/iam/model.go index c0e595616..55e4ccf6b 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" url:"-"` + Id string `json:"id,omitempty"` 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" url:"-"` + Id string `json:"id,omitempty"` // Corresponds to AWS instance profile/arn role. Roles []ComplexValue `json:"roles,omitempty"` // The schema of the List response. diff --git a/service/serving/api.go b/service/serving/api.go index 6615fda39..9e8342186 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -567,7 +567,7 @@ type ServingEndpointDataPlaneInterface interface { } func NewServingEndpointDataPlane(client *client.DatabricksClient, - controlPlane *ServingEndpointsInterface, + controlPlane *ServingEndpointsAPI, ) *ServingEndpointDataPlaneAPI { return &ServingEndpointDataPlaneAPI{ impl: &servingEndpointDataPlaneImpl{ diff --git a/service/serving/impl.go b/service/serving/impl.go index 5c80828fe..b41e2453d 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -112,7 +112,7 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, // unexported type that holds implementations of just ServingEndpointDataPlane API methods type servingEndpointDataPlaneImpl struct { client *client.DatabricksClient - controlPlane *ServingEndpointsInterface + controlPlane *ServingEndpointsAPI } func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { diff --git a/workspace_client.go b/workspace_client.go index 9d176242b..38b535e0f 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" @@ -972,19 +970,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 From 338b481cf82efc19f0e45bd551c905c54519be69 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 14:25:49 +0200 Subject: [PATCH 12/24] Method WIP --- .codegen/impl.go.tmpl | 32 +++++++++++++++++----- {data_plane => dataplane}/data_plane.go | 0 service/iam/model.go | 4 +-- service/serving/impl.go | 35 +++++++++++++++++++------ 4 files changed, 55 insertions(+), 16 deletions(-) rename {data_plane => dataplane}/data_plane.go (100%) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 8dce64966..a9f01042d 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -12,24 +12,44 @@ import ( "github.com/databricks/databricks-sdk-go/httpclient" {{range .ImportedPackages}} "github.com/databricks/databricks-sdk-go/service/{{.}}"{{end}} + goauth "golang.org/x/oauth2" ) {{range .Services}} // unexported type that holds implementations of just {{.Name}} API methods type {{.CamelName}}Impl struct { - client *client.DatabricksClient {{- if .IsDataPlane}} + dataplane.DataPlaneTokenCache controlPlane *{{.ParentService.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" . }} { - {{- template "response-var" . }} - path := {{ template "path" . }} - {{ template "make-header" . }} - err := a.client.Do(ctx, http.Method{{.TitleVerb}}, path, headers, {{ template "request-param" . }}, {{if .Response}}&{{.Response.CamelName}}{{else}}nil{{end}}) - return {{ template "response" . }} + getRequest := {{.Service.ParentService.DataPlaneInfoMethod.Request.PascalName}}{ + {{- range .Service.ParentService.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.ParentService.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) + if err != nil { + return nil, err + } + return response.DataPlaneInfo.QueryInfo, nil //TODO generate + } + refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) + } + endpointUrl, token, err := a.GetDataPlane("Query", []string{request.Name}, refresh, infoGetter) //TODO generate + fmt.Println(endpointUrl) + + return nil, nil } {{else}} func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} { diff --git a/data_plane/data_plane.go b/dataplane/data_plane.go similarity index 100% rename from data_plane/data_plane.go rename to dataplane/data_plane.go diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..b4aa4acd5 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. @@ -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/serving/impl.go b/service/serving/impl.go index b41e2453d..e3ce8ac51 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -8,6 +8,10 @@ import ( "net/http" "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/dataplane" + + "github.com/databricks/databricks-sdk-go/service/oauth2" + goauth "golang.org/x/oauth2" ) // unexported type that holds implementations of just Apps API methods @@ -111,18 +115,33 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, // unexported type that holds implementations of just ServingEndpointDataPlane API methods type servingEndpointDataPlaneImpl struct { - client *client.DatabricksClient + dataplane.DataPlaneTokenCache controlPlane *ServingEndpointsAPI + client *client.DatabricksClient } func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { - var queryEndpointResponse QueryEndpointResponse - path := fmt.Sprintf("/serving-endpoints/%v/invocations", request.Name) - headers := make(map[string]string) - headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" - err := a.client.Do(ctx, http.MethodPost, path, headers, request, &queryEndpointResponse) - return &queryEndpointResponse, err + 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 + } + return response.DataPlaneInfo.QueryInfo, nil //TODO generate + } + refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) + } + endpointUrl, token, err := a.GetDataPlane("Query", []string{request.Name}, refresh, infoGetter) //TODO generate + fmt.Println(endpointUrl) + + return nil, nil } // unexported type that holds implementations of just ServingEndpoints API methods From d71d58fb663269a216fd6366e98a7d7e59af4f1d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 14:49:33 +0200 Subject: [PATCH 13/24] more gen --- .codegen/impl.go.tmpl | 20 +++++++++++++++----- openapi/code/method.go | 16 ++++++++++++++++ service/iam/model.go | 4 ++-- service/serving/impl.go | 17 +++++++++++++---- 4 files changed, 46 insertions(+), 11 deletions(-) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index a9f01042d..a3382292d 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -41,16 +41,26 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re if err != nil { return nil, err } - return response.DataPlaneInfo.QueryInfo, nil //TODO generate + return response{{range .DataPlaneInfoFields}}.{{.PascalName}}{{end}}, nil } refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) } - endpointUrl, token, err := a.GetDataPlane("Query", []string{request.Name}, refresh, infoGetter) //TODO generate - fmt.Println(endpointUrl) - - return nil, nil + getParams := []string{ + {{- range .Service.ParentService.DataPlaneInfoMethod.Request.Fields}} + request.{{.PascalName}}, + {{end -}} + } + endpointUrl, token, err := a.GetDataPlane("{{.PascalName}}", getParams, refresh, infoGetter) + if err != nil { + return nil, err + } + {{- template "response-var" . }} + {{ template "make-header" . }} //TODO: visitors + err = a.client.Do(ctx, http.Method{{.TitleVerb}}, endpointUrl, headers, {{ template "request-param" . }}, {{if .Response}}&{{.Response.CamelName}}{{else}}nil{{end}}) + 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" . }} diff --git a/openapi/code/method.go b/openapi/code/method.go index f00ca1fe0..75f41f8cc 100644 --- a/openapi/code/method.go +++ b/openapi/code/method.go @@ -54,6 +54,22 @@ func (m *Method) HasDataPlaneAPI() bool { return m.DataPlane != nil } +func (m *Method) DataPlaneInfoFields() []*Field { + if m.DataPlane == nil { + return nil + } + method := m.Service.ParentService.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/service/iam/model.go b/service/iam/model.go index b4aa4acd5..52d3de756 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" url:"-"` + Id string `json:"id,omitempty"` 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. diff --git a/service/serving/impl.go b/service/serving/impl.go index e3ce8ac51..28e95bb80 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -138,10 +138,19 @@ func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryE refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) } - endpointUrl, token, err := a.GetDataPlane("Query", []string{request.Name}, refresh, infoGetter) //TODO generate - fmt.Println(endpointUrl) - - return nil, nil + getParams := []string{ + request.Name, + } + endpointUrl, token, err := a.GetDataPlane("Query", getParams, refresh, infoGetter) //TODO generate + if err != nil { + return nil, err + } + var queryEndpointResponse QueryEndpointResponse + headers := make(map[string]string) + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" //TODO: fix + err = a.client.Do(ctx, http.MethodPost, endpointUrl, headers, request, &queryEndpointResponse) + return &queryEndpointResponse, err } // unexported type that holds implementations of just ServingEndpoints API methods From ef348a1c405a57558da894437a29b60bfa7ef6ad Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 15:17:30 +0200 Subject: [PATCH 14/24] finish method --- .codegen/impl.go.tmpl | 11 ++++++++--- client/client.go | 5 +++++ httpclient/request.go | 11 +++++++++++ service/iam/model.go | 4 ++-- service/serving/impl.go | 16 +++++++++++----- 5 files changed, 37 insertions(+), 10 deletions(-) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index a3382292d..01047a0ad 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -51,13 +51,18 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re request.{{.PascalName}}, {{end -}} } - endpointUrl, token, err := a.GetDataPlane("{{.PascalName}}", getParams, refresh, infoGetter) + endpointUrl, dataPlaneToken, err := a.GetDataPlane("{{.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" . }} - {{ template "make-header" . }} //TODO: visitors - err = a.client.Do(ctx, http.Method{{.TitleVerb}}, endpointUrl, headers, {{ template "request-param" . }}, {{if .Response}}&{{.Response.CamelName}}{{else}}nil{{end}}) + {{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" . }} } diff --git a/client/client.go b/client/client.go index 140f16920..77cd7e146 100644 --- a/client/client.go +++ b/client/client.go @@ -47,6 +47,11 @@ 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. // 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/service/iam/model.go b/service/iam/model.go index 52d3de756..55e4ccf6b 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -1231,7 +1231,7 @@ type ServicePrincipal struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks service principal ID. - Id string `json:"id,omitempty" url:"-"` + Id string `json:"id,omitempty"` // 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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/serving/impl.go b/service/serving/impl.go index 28e95bb80..26cb315bd 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -9,6 +9,7 @@ import ( "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/dataplane" + "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/service/oauth2" goauth "golang.org/x/oauth2" @@ -133,7 +134,7 @@ func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryE if err != nil { return nil, err } - return response.DataPlaneInfo.QueryInfo, nil //TODO generate + return response.DataPlaneInfo.QueryInfo, nil } refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) @@ -141,15 +142,20 @@ func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryE getParams := []string{ request.Name, } - endpointUrl, token, err := a.GetDataPlane("Query", getParams, refresh, infoGetter) //TODO generate + endpointUrl, dataPlaneToken, err := a.GetDataPlane("Query", getParams, refresh, infoGetter) if err != nil { return nil, err } - var queryEndpointResponse QueryEndpointResponse headers := make(map[string]string) headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" //TODO: fix - err = a.client.Do(ctx, http.MethodPost, endpointUrl, headers, request, &queryEndpointResponse) + 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 } From 5e30cb1f0c547958debf5633823241d95ede1214 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 28 May 2024 15:38:06 +0200 Subject: [PATCH 15/24] Clean interface --- .codegen/accounts.go.tmpl | 15 +- .codegen/api.go.tmpl | 12 - .codegen/workspaces.go.tmpl | 16 +- account_client.go | 51 ++++ .../mock_serving_endpoints_interface.go | 58 ---- service/iam/model.go | 4 +- service/serving/api.go | 8 - workspace_client.go | 260 ++++++++++++------ 8 files changed, 260 insertions(+), 164 deletions(-) diff --git a/.codegen/accounts.go.tmpl b/.codegen/accounts.go.tmpl index 6138f84bb..7f7c0551d 100644 --- a/.codegen/accounts.go.tmpl +++ b/.codegen/accounts.go.tmpl @@ -12,7 +12,7 @@ import ( type AccountClient struct { Config *config.Config - {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} + {{range .Services}}{{if and .IsAccounts (not .HasParent)}} {{.Comment " // " 80}} {{(.TrimPrefix "account").PascalName}} {{.Package.Name}}.{{.Name}}Interface {{end}}{{end}} @@ -42,9 +42,22 @@ 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) (not .IsDataPlane)}} {{(.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, {{.ParentService.CamelName}}), + {{- end}} + {{end}} }, nil } diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index f00a1fbbd..24cc23236 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -42,11 +42,6 @@ type {{.PascalName}}Interface interface { {{.PascalName}}() {{.PascalName}}Interface {{end}} - {{if and (not .IsDataPlane) .HasDataPlaneMethods}} - // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API - Get{{.Singular.PascalName}}DataPlaneClient(request {{.DataPlaneInfoMethod.Request.PascalName}}) ({{.Singular.PascalName}}DataPlaneService, error) - {{end}} - {{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 {{.PascalName}}(ctx context.Context{{range .Binding}}, {{.PollField.CamelName}} {{template "type" .PollField.Entity}}{{end}}, @@ -150,13 +145,6 @@ func (a *{{.ParentService.PascalName}}API) {{.PascalName}}() {{.PascalName}}Inte } {{end}} -{{if and (not .IsDataPlane) .HasDataPlaneMethods}} - // Returns a client to interact with a specific {{.Singular.PascalName}} resource via the DataPlane API -func (a *{{.PascalName}}API) Get{{.Singular.PascalName}}DataPlaneClient(request {{.DataPlaneInfoMethod.Request.PascalName}}) ({{.Singular.PascalName}}DataPlaneService, error) { - return nil, nil -} -{{end}} - // 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 Mock{{.PascalName}}Interface instead. diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index 7880532e7..f404f037a 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.codegen/workspaces.go.tmpl @@ -14,7 +14,7 @@ type WorkspaceClient struct { Config *config.Config apiClient *httpclient.ApiClient - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} {{.Comment " // " 80}} {{.Name}} {{.Package.Name}}.{{.Name}}Interface {{end}}{{end}} @@ -48,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) (not .IsDataPlane)}} + {{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, {{.ParentService.CamelName}}), + {{- end}} + {{end}} }, nil } diff --git a/account_client.go b/account_client.go index 2596a826e..afbd6c50a 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, @@ -317,5 +318,55 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { VpcEndpoints: provisioning.NewVpcEndpoints(apiClient), WorkspaceAssignment: iam.NewWorkspaceAssignment(apiClient), Workspaces: provisioning.NewWorkspaces(apiClient), + + AccessControl: iam.NewAccountAccessControl(apiClient), + + BillableUsage: billing.NewBillableUsage(apiClient), + + Budgets: billing.NewBudgets(apiClient), + + Credentials: provisioning.NewCredentials(apiClient), + + CustomAppIntegration: oauth2.NewCustomAppIntegration(apiClient), + + EncryptionKeys: provisioning.NewEncryptionKeys(apiClient), + + Groups: iam.NewAccountGroups(apiClient), + + IpAccessLists: settings.NewAccountIpAccessLists(apiClient), + + LogDelivery: billing.NewLogDelivery(apiClient), + + MetastoreAssignments: catalog.NewAccountMetastoreAssignments(apiClient), + + Metastores: catalog.NewAccountMetastores(apiClient), + + NetworkConnectivity: settings.NewNetworkConnectivity(apiClient), + + Networks: provisioning.NewNetworks(apiClient), + + OAuthPublishedApps: oauth2.NewOAuthPublishedApps(apiClient), + + PrivateAccess: provisioning.NewPrivateAccess(apiClient), + + PublishedAppIntegration: oauth2.NewPublishedAppIntegration(apiClient), + + ServicePrincipalSecrets: oauth2.NewServicePrincipalSecrets(apiClient), + + ServicePrincipals: iam.NewAccountServicePrincipals(apiClient), + + Settings: settings.NewAccountSettings(apiClient), + + Storage: provisioning.NewStorage(apiClient), + + StorageCredentials: catalog.NewAccountStorageCredentials(apiClient), + + Users: iam.NewAccountUsers(apiClient), + + VpcEndpoints: provisioning.NewVpcEndpoints(apiClient), + + WorkspaceAssignment: iam.NewWorkspaceAssignment(apiClient), + + Workspaces: provisioning.NewWorkspaces(apiClient), }, nil } diff --git a/experimental/mocks/service/serving/mock_serving_endpoints_interface.go b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go index 318a63006..496304aec 100644 --- a/experimental/mocks/service/serving/mock_serving_endpoints_interface.go +++ b/experimental/mocks/service/serving/mock_serving_endpoints_interface.go @@ -940,64 +940,6 @@ func (_c *MockServingEndpointsInterface_GetPermissionsByServingEndpointId_Call) return _c } -// GetServingEndpointDataPlaneClient provides a mock function with given fields: request -func (_m *MockServingEndpointsInterface) GetServingEndpointDataPlaneClient(request serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error) { - ret := _m.Called(request) - - if len(ret) == 0 { - panic("no return value specified for GetServingEndpointDataPlaneClient") - } - - var r0 serving.ServingEndpointDataPlaneService - var r1 error - if rf, ok := ret.Get(0).(func(serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error)); ok { - return rf(request) - } - if rf, ok := ret.Get(0).(func(serving.GetServingEndpointRequest) serving.ServingEndpointDataPlaneService); ok { - r0 = rf(request) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(serving.ServingEndpointDataPlaneService) - } - } - - if rf, ok := ret.Get(1).(func(serving.GetServingEndpointRequest) error); ok { - r1 = rf(request) - } else { - r1 = ret.Error(1) - } - - return r0, r1 -} - -// MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetServingEndpointDataPlaneClient' -type MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call struct { - *mock.Call -} - -// GetServingEndpointDataPlaneClient is a helper method to define mock.On call -// - request serving.GetServingEndpointRequest -func (_e *MockServingEndpointsInterface_Expecter) GetServingEndpointDataPlaneClient(request interface{}) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { - return &MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call{Call: _e.mock.On("GetServingEndpointDataPlaneClient", request)} -} - -func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Run(run func(request serving.GetServingEndpointRequest)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(serving.GetServingEndpointRequest)) - }) - return _c -} - -func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) Return(_a0 serving.ServingEndpointDataPlaneService, _a1 error) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { - _c.Call.Return(_a0, _a1) - return _c -} - -func (_c *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call) RunAndReturn(run func(serving.GetServingEndpointRequest) (serving.ServingEndpointDataPlaneService, error)) *MockServingEndpointsInterface_GetServingEndpointDataPlaneClient_Call { - _c.Call.Return(run) - return _c -} - // Impl provides a mock function with given fields: func (_m *MockServingEndpointsInterface) Impl() serving.ServingEndpointsService { ret := _m.Called() diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..52d3de756 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -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/serving/api.go b/service/serving/api.go index 9e8342186..6f74b1c11 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -613,9 +613,6 @@ type ServingEndpointsInterface interface { // Deprecated: use MockServingEndpointsInterface instead. Impl() ServingEndpointsService - // Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API - GetServingEndpointDataPlaneClient(request GetServingEndpointRequest) (ServingEndpointDataPlaneService, error) - // WaitGetServingEndpointNotUpdating repeatedly calls [ServingEndpointsAPI.Get] and waits to reach NOT_UPDATING state WaitGetServingEndpointNotUpdating(ctx context.Context, name string, timeout time.Duration, callback func(*ServingEndpointDetailed)) (*ServingEndpointDetailed, error) @@ -796,11 +793,6 @@ type ServingEndpointsAPI struct { impl ServingEndpointsService } -// Returns a client to interact with a specific ServingEndpoint resource via the DataPlane API -func (a *ServingEndpointsAPI) GetServingEndpointDataPlaneClient(request GetServingEndpointRequest) (ServingEndpointDataPlaneService, error) { - return nil, nil -} - // 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 MockServingEndpointsInterface instead. diff --git a/workspace_client.go b/workspace_client.go index 38b535e0f..8eb8989bf 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -696,6 +696,22 @@ type WorkspaceClient struct { // data by accident. ServicePrincipals iam.ServicePrincipalsInterface + // The Serving Endpoints API allows you to create, update, and delete model + // serving endpoints. + // + // You can use a serving endpoint to serve models from the Databricks Model + // Registry or from Unity Catalog. Endpoints expose the underlying models as + // scalable REST API endpoints using serverless compute. This means the + // endpoints and associated compute resources are fully managed by + // Databricks and will not appear in your cloud account. A serving endpoint + // can consist of one or more MLflow models from the Databricks Model + // Registry, called served entities. A serving endpoint can have at most ten + // served entities. You can configure traffic settings to define how + // requests should be routed to your served entities behind an endpoint. + // Additionally, you can configure the scale of resources that should be + // applied to each served entity. + ServingEndpointDataPlane serving.ServingEndpointDataPlaneInterface + // The Serving Endpoints API allows you to create, update, and delete model // serving endpoints. // @@ -998,90 +1014,174 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { if err != nil { return nil, err } + + servingEndpoints := serving.NewServingEndpoints(databricksClient) return &WorkspaceClient{ Config: cfg, apiClient: apiClient, - AccountAccessControlProxy: iam.NewAccountAccessControlProxy(databricksClient), - Alerts: sql.NewAlerts(databricksClient), - Apps: serving.NewApps(databricksClient), - ArtifactAllowlists: catalog.NewArtifactAllowlists(databricksClient), - Catalogs: catalog.NewCatalogs(databricksClient), - CleanRooms: sharing.NewCleanRooms(databricksClient), - ClusterPolicies: compute.NewClusterPolicies(databricksClient), - Clusters: compute.NewClusters(databricksClient), - CommandExecution: compute.NewCommandExecution(databricksClient), - Connections: catalog.NewConnections(databricksClient), - ConsumerFulfillments: marketplace.NewConsumerFulfillments(databricksClient), - ConsumerInstallations: marketplace.NewConsumerInstallations(databricksClient), - ConsumerListings: marketplace.NewConsumerListings(databricksClient), - ConsumerPersonalizationRequests: marketplace.NewConsumerPersonalizationRequests(databricksClient), - ConsumerProviders: marketplace.NewConsumerProviders(databricksClient), - CredentialsManager: settings.NewCredentialsManager(databricksClient), - CurrentUser: iam.NewCurrentUser(databricksClient), - DashboardWidgets: sql.NewDashboardWidgets(databricksClient), - Dashboards: sql.NewDashboards(databricksClient), - DataSources: sql.NewDataSources(databricksClient), - Dbfs: files.NewDbfs(databricksClient), - DbsqlPermissions: sql.NewDbsqlPermissions(databricksClient), - Experiments: ml.NewExperiments(databricksClient), - ExternalLocations: catalog.NewExternalLocations(databricksClient), - Files: files.NewFiles(databricksClient), - Functions: catalog.NewFunctions(databricksClient), - GitCredentials: workspace.NewGitCredentials(databricksClient), - GlobalInitScripts: compute.NewGlobalInitScripts(databricksClient), - Grants: catalog.NewGrants(databricksClient), - Groups: iam.NewGroups(databricksClient), - InstancePools: compute.NewInstancePools(databricksClient), - InstanceProfiles: compute.NewInstanceProfiles(databricksClient), - IpAccessLists: settings.NewIpAccessLists(databricksClient), - Jobs: jobs.NewJobs(databricksClient), - Lakeview: dashboards.NewLakeview(databricksClient), - Libraries: compute.NewLibraries(databricksClient), - Metastores: catalog.NewMetastores(databricksClient), - ModelRegistry: ml.NewModelRegistry(databricksClient), - ModelVersions: catalog.NewModelVersions(databricksClient), - OnlineTables: catalog.NewOnlineTables(databricksClient), - PermissionMigration: iam.NewPermissionMigration(databricksClient), - Permissions: iam.NewPermissions(databricksClient), - Pipelines: pipelines.NewPipelines(databricksClient), - PolicyFamilies: compute.NewPolicyFamilies(databricksClient), - ProviderExchangeFilters: marketplace.NewProviderExchangeFilters(databricksClient), - ProviderExchanges: marketplace.NewProviderExchanges(databricksClient), - ProviderFiles: marketplace.NewProviderFiles(databricksClient), - ProviderListings: marketplace.NewProviderListings(databricksClient), - ProviderPersonalizationRequests: marketplace.NewProviderPersonalizationRequests(databricksClient), + AccountAccessControlProxy: iam.NewAccountAccessControlProxy(databricksClient), + + Alerts: sql.NewAlerts(databricksClient), + + Apps: serving.NewApps(databricksClient), + + ArtifactAllowlists: catalog.NewArtifactAllowlists(databricksClient), + + Catalogs: catalog.NewCatalogs(databricksClient), + + CleanRooms: sharing.NewCleanRooms(databricksClient), + + ClusterPolicies: compute.NewClusterPolicies(databricksClient), + + Clusters: compute.NewClusters(databricksClient), + + CommandExecution: compute.NewCommandExecution(databricksClient), + + Connections: catalog.NewConnections(databricksClient), + + ConsumerFulfillments: marketplace.NewConsumerFulfillments(databricksClient), + + ConsumerInstallations: marketplace.NewConsumerInstallations(databricksClient), + + ConsumerListings: marketplace.NewConsumerListings(databricksClient), + + ConsumerPersonalizationRequests: marketplace.NewConsumerPersonalizationRequests(databricksClient), + + ConsumerProviders: marketplace.NewConsumerProviders(databricksClient), + + CredentialsManager: settings.NewCredentialsManager(databricksClient), + + CurrentUser: iam.NewCurrentUser(databricksClient), + + DashboardWidgets: sql.NewDashboardWidgets(databricksClient), + + Dashboards: sql.NewDashboards(databricksClient), + + DataSources: sql.NewDataSources(databricksClient), + + Dbfs: files.NewDbfs(databricksClient), + + DbsqlPermissions: sql.NewDbsqlPermissions(databricksClient), + + Experiments: ml.NewExperiments(databricksClient), + + ExternalLocations: catalog.NewExternalLocations(databricksClient), + + Files: files.NewFiles(databricksClient), + + Functions: catalog.NewFunctions(databricksClient), + + GitCredentials: workspace.NewGitCredentials(databricksClient), + + GlobalInitScripts: compute.NewGlobalInitScripts(databricksClient), + + Grants: catalog.NewGrants(databricksClient), + + Groups: iam.NewGroups(databricksClient), + + InstancePools: compute.NewInstancePools(databricksClient), + + InstanceProfiles: compute.NewInstanceProfiles(databricksClient), + + IpAccessLists: settings.NewIpAccessLists(databricksClient), + + Jobs: jobs.NewJobs(databricksClient), + + Lakeview: dashboards.NewLakeview(databricksClient), + + Libraries: compute.NewLibraries(databricksClient), + + Metastores: catalog.NewMetastores(databricksClient), + + ModelRegistry: ml.NewModelRegistry(databricksClient), + + ModelVersions: catalog.NewModelVersions(databricksClient), + + OnlineTables: catalog.NewOnlineTables(databricksClient), + + PermissionMigration: iam.NewPermissionMigration(databricksClient), + + Permissions: iam.NewPermissions(databricksClient), + + Pipelines: pipelines.NewPipelines(databricksClient), + + PolicyFamilies: compute.NewPolicyFamilies(databricksClient), + + ProviderExchangeFilters: marketplace.NewProviderExchangeFilters(databricksClient), + + ProviderExchanges: marketplace.NewProviderExchanges(databricksClient), + + ProviderFiles: marketplace.NewProviderFiles(databricksClient), + + ProviderListings: marketplace.NewProviderListings(databricksClient), + + ProviderPersonalizationRequests: marketplace.NewProviderPersonalizationRequests(databricksClient), + ProviderProviderAnalyticsDashboards: marketplace.NewProviderProviderAnalyticsDashboards(databricksClient), - ProviderProviders: marketplace.NewProviderProviders(databricksClient), - Providers: sharing.NewProviders(databricksClient), - QualityMonitors: catalog.NewQualityMonitors(databricksClient), - Queries: sql.NewQueries(databricksClient), - QueryHistory: sql.NewQueryHistory(databricksClient), - QueryVisualizations: sql.NewQueryVisualizations(databricksClient), - RecipientActivation: sharing.NewRecipientActivation(databricksClient), - Recipients: sharing.NewRecipients(databricksClient), - RegisteredModels: catalog.NewRegisteredModels(databricksClient), - Repos: workspace.NewRepos(databricksClient), - Schemas: catalog.NewSchemas(databricksClient), - Secrets: workspace.NewSecrets(databricksClient), - ServicePrincipals: iam.NewServicePrincipals(databricksClient), - ServingEndpoints: serving.NewServingEndpoints(databricksClient), - Settings: settings.NewSettings(databricksClient), - Shares: sharing.NewShares(databricksClient), - StatementExecution: sql.NewStatementExecution(databricksClient), - StorageCredentials: catalog.NewStorageCredentials(databricksClient), - SystemSchemas: catalog.NewSystemSchemas(databricksClient), - TableConstraints: catalog.NewTableConstraints(databricksClient), - Tables: catalog.NewTables(databricksClient), - TokenManagement: settings.NewTokenManagement(databricksClient), - Tokens: settings.NewTokens(databricksClient), - Users: iam.NewUsers(databricksClient), - VectorSearchEndpoints: vectorsearch.NewVectorSearchEndpoints(databricksClient), - VectorSearchIndexes: vectorsearch.NewVectorSearchIndexes(databricksClient), - Volumes: catalog.NewVolumes(databricksClient), - Warehouses: sql.NewWarehouses(databricksClient), - Workspace: workspace.NewWorkspace(databricksClient), - WorkspaceBindings: catalog.NewWorkspaceBindings(databricksClient), - WorkspaceConf: settings.NewWorkspaceConf(databricksClient), + + ProviderProviders: marketplace.NewProviderProviders(databricksClient), + + Providers: sharing.NewProviders(databricksClient), + + QualityMonitors: catalog.NewQualityMonitors(databricksClient), + + Queries: sql.NewQueries(databricksClient), + + QueryHistory: sql.NewQueryHistory(databricksClient), + + QueryVisualizations: sql.NewQueryVisualizations(databricksClient), + + RecipientActivation: sharing.NewRecipientActivation(databricksClient), + + Recipients: sharing.NewRecipients(databricksClient), + + RegisteredModels: catalog.NewRegisteredModels(databricksClient), + + Repos: workspace.NewRepos(databricksClient), + + Schemas: catalog.NewSchemas(databricksClient), + + Secrets: workspace.NewSecrets(databricksClient), + + ServicePrincipals: iam.NewServicePrincipals(databricksClient), + + ServingEndpointDataPlane: serving.NewServingEndpointDataPlane(databricksClient, servingEndpoints), + + ServingEndpoints: servingEndpoints, + + Settings: settings.NewSettings(databricksClient), + + Shares: sharing.NewShares(databricksClient), + + StatementExecution: sql.NewStatementExecution(databricksClient), + + StorageCredentials: catalog.NewStorageCredentials(databricksClient), + + SystemSchemas: catalog.NewSystemSchemas(databricksClient), + + TableConstraints: catalog.NewTableConstraints(databricksClient), + + Tables: catalog.NewTables(databricksClient), + + TokenManagement: settings.NewTokenManagement(databricksClient), + + Tokens: settings.NewTokens(databricksClient), + + Users: iam.NewUsers(databricksClient), + + VectorSearchEndpoints: vectorsearch.NewVectorSearchEndpoints(databricksClient), + + VectorSearchIndexes: vectorsearch.NewVectorSearchIndexes(databricksClient), + + Volumes: catalog.NewVolumes(databricksClient), + + Warehouses: sql.NewWarehouses(databricksClient), + + Workspace: workspace.NewWorkspace(databricksClient), + + WorkspaceBindings: catalog.NewWorkspaceBindings(databricksClient), + + WorkspaceConf: settings.NewWorkspaceConf(databricksClient), }, nil } From bb425d53def75580eb076fa9f3fc7c9ceded1790 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 09:36:15 +0200 Subject: [PATCH 16/24] Tests --- dataplane/data_plane.go | 22 +++-- dataplane/data_plane_test.go | 165 +++++++++++++++++++++++++++++++++++ 2 files changed, 181 insertions(+), 6 deletions(-) create mode 100644 dataplane/data_plane_test.go diff --git a/dataplane/data_plane.go b/dataplane/data_plane.go index f7a58bda8..8c40a3d4d 100644 --- a/dataplane/data_plane.go +++ b/dataplane/data_plane.go @@ -13,6 +13,12 @@ type DataPlaneTokenCache struct { } func (o *DataPlaneTokenCache) GetDataPlane(method string, params []string, refresh func(*dp.DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*dp.DataPlaneInfo, error)) (string, *oauth2.Token, error) { + if o.infos == nil { + o.infos = make(map[string]*dp.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] @@ -27,15 +33,19 @@ func (o *DataPlaneTokenCache) GetDataPlane(method string, params []string, refre o.infos[key] = newInfo info = newInfo } - newToken, err := refresh(info) - if err != nil { - return "", nil, err + if !tokenOk || !token.Valid() { + newToken, err := refresh(info) + if err != nil { + return "", nil, err + } + o.tokens[key] = newToken + token = newToken } - o.tokens[method] = newToken - return info.EndpointUrl, newToken, nil + return info.EndpointUrl, token, nil } func (o *DataPlaneTokenCache) generateKey(method string, params []string) string { - allElements := append(params, method) + allElements := []string{method} + allElements = append(allElements, params...) return strings.Join(allElements, "/") } diff --git a/dataplane/data_plane_test.go b/dataplane/data_plane_test.go new file mode 100644 index 000000000..a002fe4ca --- /dev/null +++ b/dataplane/data_plane_test.go @@ -0,0 +1,165 @@ +package dataplane + +import ( + "testing" + "time" + + dp "github.com/databricks/databricks-sdk-go/service/oauth2" + "github.com/stretchr/testify/assert" + "golang.org/x/oauth2" +) + +type infoMock struct { + called bool + info *dp.DataPlaneInfo + err error +} + +func (i *infoMock) DataPlaneInfoGetter() (*dp.DataPlaneInfo, error) { + i.called = true + return i.info, i.err +} + +type tokenRefreshSpy struct { + called bool + expectedInfo *dp.DataPlaneInfo + token *oauth2.Token + err error +} + +func (t *tokenRefreshSpy) TokenRefresh(info *dp.DataPlaneInfo) (*oauth2.Token, error) { + t.expectedInfo = info + t.called = true + return t.token, t.err +} + +func TestTokenNotCached(t *testing.T) { + info := infoMock{ + info: &dp.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 := DataPlaneTokenCache{} + url, token, err := c.GetDataPlane("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: &dp.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 := DataPlaneTokenCache{} + c.infos = make(map[string]*dp.DataPlaneInfo) + c.tokens = make(map[string]*oauth2.Token) + c.infos["method/params"] = info.info + c.tokens["method/params"] = s.token + url, token, err := c.GetDataPlane("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: &dp.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 := DataPlaneTokenCache{} + c.infos = make(map[string]*dp.DataPlaneInfo) + c.tokens = make(map[string]*oauth2.Token) + c.infos["method/params"] = info.info + c.tokens["method/params"] = expired + url, token, err := c.GetDataPlane("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 := DataPlaneTokenCache{} + url, token, err := c.GetDataPlane("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: &dp.DataPlaneInfo{ + EndpointUrl: "url", + AuthorizationDetails: "authDetails", + }, + err: nil, + } + s := tokenRefreshSpy{ + token: nil, + err: assert.AnError, + } + c := DataPlaneTokenCache{} + url, token, err := c.GetDataPlane("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) +} From ae979701fc1617abb8fb9163a7401349e209c0a2 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 09:47:31 +0200 Subject: [PATCH 17/24] Error hadnling --- .codegen/impl.go.tmpl | 4 ++ .../mock_consumer_listings_interface.go | 59 +++++++++++++++++++ .../mock_consumer_providers_interface.go | 59 +++++++++++++++++++ service/iam/model.go | 4 +- service/jobs/model.go | 2 + service/marketplace/api.go | 28 +++++++++ service/marketplace/impl.go | 18 ++++++ service/marketplace/interface.go | 12 ++++ service/marketplace/model.go | 18 ++++++ service/serving/impl.go | 4 ++ 10 files changed, 206 insertions(+), 2 deletions(-) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 01047a0ad..5541597c8 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -4,6 +4,7 @@ package {{.Name}} import ( "context" + "errors" "fmt" "time" "io" @@ -41,6 +42,9 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re 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) { 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/service/iam/model.go b/service/iam/model.go index 52d3de756..b4aa4acd5 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" url:"-"` + Id string `json:"id,omitempty"` // Corresponds to AWS instance profile/arn role. Roles []ComplexValue `json:"roles,omitempty"` // The schema of the List response. 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/serving/impl.go b/service/serving/impl.go index 26cb315bd..26009b09d 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -4,6 +4,7 @@ package serving import ( "context" + "errors" "fmt" "net/http" @@ -134,6 +135,9 @@ func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryE 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) { From 579f6b5fcb7a33d5482fbbd4e47949ea84442bc0 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 10:03:04 +0200 Subject: [PATCH 18/24] Fixes --- .codegen/accounts.go.tmpl | 8 +- .codegen/api.go.tmpl | 2 +- .codegen/interface.go.tmpl | 2 +- .codegen/workspaces.go.tmpl | 6 +- account_client.go | 50 ------- openapi/code/resource_client.go | 14 -- service/iam/model.go | 6 +- service/serving/api.go | 2 +- service/serving/interface.go | 2 +- workspace_client.go | 243 +++++++++++--------------------- 10 files changed, 94 insertions(+), 241 deletions(-) delete mode 100644 openapi/code/resource_client.go diff --git a/.codegen/accounts.go.tmpl b/.codegen/accounts.go.tmpl index 7f7c0551d..d948b7acd 100644 --- a/.codegen/accounts.go.tmpl +++ b/.codegen/accounts.go.tmpl @@ -47,17 +47,15 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { {{- end}}{{end}} return &AccountClient{ Config: cfg, - {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} - {{(.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}} + {{- end -}} {{if and .IsAccounts (not .HasParent) (not .IsDataPlane) (.HasDataPlaneMethods)}} {{(.TrimPrefix "account").PascalName}}: {{.CamelName}}, - {{- end}} + {{- end -}} {{if and .IsAccounts (not .HasParent) .IsDataPlane}} {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient, {{.ParentService.CamelName}}), - {{- end}} + {{- end -}} {{end}} }, nil } diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 24cc23236..c035b4834 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -126,7 +126,7 @@ controlPlane *{{.ParentService.PascalName}}API, {{- if not .IsDataPlane}} {{.Comment "// " 80}} {{else}} -// TODO comment +// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ParentService.PascalName}} service. {{end -}} type {{.PascalName}}API struct { // impl contains low-level REST API interface, that could be overridden diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index 082d10d20..2c55dc227 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -12,7 +12,7 @@ import ( {{- if not .IsDataPlane}} {{.Comment "// " 80}} {{else}} -// TODO comment +// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ParentService.PascalName}} service. {{end -}} type {{.PascalName}}Service interface { {{range .Methods}} diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index f404f037a..f5ef4dc05 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.codegen/workspaces.go.tmpl @@ -56,13 +56,13 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { apiClient: apiClient, {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane) (not .HasDataPlaneMethods)}} {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient), - {{- end}} + {{- end -}} {{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane) (.HasDataPlaneMethods)}} {{.Name}}: {{.CamelName}}, - {{- end}} + {{- end -}} {{if and (not .IsAccounts) (not .HasParent) .IsDataPlane}} {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient, {{.ParentService.CamelName}}), - {{- end}} + {{- end -}} {{end}} }, nil } diff --git a/account_client.go b/account_client.go index afbd6c50a..a226ec7df 100755 --- a/account_client.go +++ b/account_client.go @@ -318,55 +318,5 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { VpcEndpoints: provisioning.NewVpcEndpoints(apiClient), WorkspaceAssignment: iam.NewWorkspaceAssignment(apiClient), Workspaces: provisioning.NewWorkspaces(apiClient), - - AccessControl: iam.NewAccountAccessControl(apiClient), - - BillableUsage: billing.NewBillableUsage(apiClient), - - Budgets: billing.NewBudgets(apiClient), - - Credentials: provisioning.NewCredentials(apiClient), - - CustomAppIntegration: oauth2.NewCustomAppIntegration(apiClient), - - EncryptionKeys: provisioning.NewEncryptionKeys(apiClient), - - Groups: iam.NewAccountGroups(apiClient), - - IpAccessLists: settings.NewAccountIpAccessLists(apiClient), - - LogDelivery: billing.NewLogDelivery(apiClient), - - MetastoreAssignments: catalog.NewAccountMetastoreAssignments(apiClient), - - Metastores: catalog.NewAccountMetastores(apiClient), - - NetworkConnectivity: settings.NewNetworkConnectivity(apiClient), - - Networks: provisioning.NewNetworks(apiClient), - - OAuthPublishedApps: oauth2.NewOAuthPublishedApps(apiClient), - - PrivateAccess: provisioning.NewPrivateAccess(apiClient), - - PublishedAppIntegration: oauth2.NewPublishedAppIntegration(apiClient), - - ServicePrincipalSecrets: oauth2.NewServicePrincipalSecrets(apiClient), - - ServicePrincipals: iam.NewAccountServicePrincipals(apiClient), - - Settings: settings.NewAccountSettings(apiClient), - - Storage: provisioning.NewStorage(apiClient), - - StorageCredentials: catalog.NewAccountStorageCredentials(apiClient), - - Users: iam.NewAccountUsers(apiClient), - - VpcEndpoints: provisioning.NewVpcEndpoints(apiClient), - - WorkspaceAssignment: iam.NewWorkspaceAssignment(apiClient), - - Workspaces: provisioning.NewWorkspaces(apiClient), }, nil } diff --git a/openapi/code/resource_client.go b/openapi/code/resource_client.go deleted file mode 100644 index 6782641c4..000000000 --- a/openapi/code/resource_client.go +++ /dev/null @@ -1,14 +0,0 @@ -package code - -import "github.com/databricks/databricks-sdk-go/openapi" - -type ResourceClient struct { - Named - - PathStyle openapi.PathStyle - IsAccounts bool - Package *Package - methods map[string]*Method - ByPathParamsMethods []*Shortcut - tag *openapi.Tag -} diff --git a/service/iam/model.go b/service/iam/model.go index b4aa4acd5..76321c4b8 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" url:"-"` + Id string `json:"id,omitempty"` 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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/serving/api.go b/service/serving/api.go index 6f74b1c11..f57d1dbf6 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -577,7 +577,7 @@ func NewServingEndpointDataPlane(client *client.DatabricksClient, } } -// TODO comment +// ServingEndpointDataPlaneService 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) diff --git a/service/serving/interface.go b/service/serving/interface.go index c891f9407..0d122b3c0 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -67,7 +67,7 @@ type AppsService interface { Update(ctx context.Context, request UpdateAppRequest) (*App, error) } -// TODO comment +// ServingEndpointDataPlaneService provides a set of operations to interact with DataPlane endpoints for ServingEndpoints service. type ServingEndpointDataPlaneService interface { // Query a serving endpoint. diff --git a/workspace_client.go b/workspace_client.go index 8eb8989bf..8ad3c424a 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -1020,168 +1020,87 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { Config: cfg, apiClient: apiClient, - AccountAccessControlProxy: iam.NewAccountAccessControlProxy(databricksClient), - - Alerts: sql.NewAlerts(databricksClient), - - Apps: serving.NewApps(databricksClient), - - ArtifactAllowlists: catalog.NewArtifactAllowlists(databricksClient), - - Catalogs: catalog.NewCatalogs(databricksClient), - - CleanRooms: sharing.NewCleanRooms(databricksClient), - - ClusterPolicies: compute.NewClusterPolicies(databricksClient), - - Clusters: compute.NewClusters(databricksClient), - - CommandExecution: compute.NewCommandExecution(databricksClient), - - Connections: catalog.NewConnections(databricksClient), - - ConsumerFulfillments: marketplace.NewConsumerFulfillments(databricksClient), - - ConsumerInstallations: marketplace.NewConsumerInstallations(databricksClient), - - ConsumerListings: marketplace.NewConsumerListings(databricksClient), - - ConsumerPersonalizationRequests: marketplace.NewConsumerPersonalizationRequests(databricksClient), - - ConsumerProviders: marketplace.NewConsumerProviders(databricksClient), - - CredentialsManager: settings.NewCredentialsManager(databricksClient), - - CurrentUser: iam.NewCurrentUser(databricksClient), - - DashboardWidgets: sql.NewDashboardWidgets(databricksClient), - - Dashboards: sql.NewDashboards(databricksClient), - - DataSources: sql.NewDataSources(databricksClient), - - Dbfs: files.NewDbfs(databricksClient), - - DbsqlPermissions: sql.NewDbsqlPermissions(databricksClient), - - Experiments: ml.NewExperiments(databricksClient), - - ExternalLocations: catalog.NewExternalLocations(databricksClient), - - Files: files.NewFiles(databricksClient), - - Functions: catalog.NewFunctions(databricksClient), - - GitCredentials: workspace.NewGitCredentials(databricksClient), - - GlobalInitScripts: compute.NewGlobalInitScripts(databricksClient), - - Grants: catalog.NewGrants(databricksClient), - - Groups: iam.NewGroups(databricksClient), - - InstancePools: compute.NewInstancePools(databricksClient), - - InstanceProfiles: compute.NewInstanceProfiles(databricksClient), - - IpAccessLists: settings.NewIpAccessLists(databricksClient), - - Jobs: jobs.NewJobs(databricksClient), - - Lakeview: dashboards.NewLakeview(databricksClient), - - Libraries: compute.NewLibraries(databricksClient), - - Metastores: catalog.NewMetastores(databricksClient), - - ModelRegistry: ml.NewModelRegistry(databricksClient), - - ModelVersions: catalog.NewModelVersions(databricksClient), - - OnlineTables: catalog.NewOnlineTables(databricksClient), - - PermissionMigration: iam.NewPermissionMigration(databricksClient), - - Permissions: iam.NewPermissions(databricksClient), - - Pipelines: pipelines.NewPipelines(databricksClient), - - PolicyFamilies: compute.NewPolicyFamilies(databricksClient), - - ProviderExchangeFilters: marketplace.NewProviderExchangeFilters(databricksClient), - - ProviderExchanges: marketplace.NewProviderExchanges(databricksClient), - - ProviderFiles: marketplace.NewProviderFiles(databricksClient), - - ProviderListings: marketplace.NewProviderListings(databricksClient), - - ProviderPersonalizationRequests: marketplace.NewProviderPersonalizationRequests(databricksClient), - + AccountAccessControlProxy: iam.NewAccountAccessControlProxy(databricksClient), + Alerts: sql.NewAlerts(databricksClient), + Apps: serving.NewApps(databricksClient), + ArtifactAllowlists: catalog.NewArtifactAllowlists(databricksClient), + Catalogs: catalog.NewCatalogs(databricksClient), + CleanRooms: sharing.NewCleanRooms(databricksClient), + ClusterPolicies: compute.NewClusterPolicies(databricksClient), + Clusters: compute.NewClusters(databricksClient), + CommandExecution: compute.NewCommandExecution(databricksClient), + Connections: catalog.NewConnections(databricksClient), + ConsumerFulfillments: marketplace.NewConsumerFulfillments(databricksClient), + ConsumerInstallations: marketplace.NewConsumerInstallations(databricksClient), + ConsumerListings: marketplace.NewConsumerListings(databricksClient), + ConsumerPersonalizationRequests: marketplace.NewConsumerPersonalizationRequests(databricksClient), + ConsumerProviders: marketplace.NewConsumerProviders(databricksClient), + CredentialsManager: settings.NewCredentialsManager(databricksClient), + CurrentUser: iam.NewCurrentUser(databricksClient), + DashboardWidgets: sql.NewDashboardWidgets(databricksClient), + Dashboards: sql.NewDashboards(databricksClient), + DataSources: sql.NewDataSources(databricksClient), + Dbfs: files.NewDbfs(databricksClient), + DbsqlPermissions: sql.NewDbsqlPermissions(databricksClient), + Experiments: ml.NewExperiments(databricksClient), + ExternalLocations: catalog.NewExternalLocations(databricksClient), + Files: files.NewFiles(databricksClient), + Functions: catalog.NewFunctions(databricksClient), + GitCredentials: workspace.NewGitCredentials(databricksClient), + GlobalInitScripts: compute.NewGlobalInitScripts(databricksClient), + Grants: catalog.NewGrants(databricksClient), + Groups: iam.NewGroups(databricksClient), + InstancePools: compute.NewInstancePools(databricksClient), + InstanceProfiles: compute.NewInstanceProfiles(databricksClient), + IpAccessLists: settings.NewIpAccessLists(databricksClient), + Jobs: jobs.NewJobs(databricksClient), + Lakeview: dashboards.NewLakeview(databricksClient), + Libraries: compute.NewLibraries(databricksClient), + Metastores: catalog.NewMetastores(databricksClient), + ModelRegistry: ml.NewModelRegistry(databricksClient), + ModelVersions: catalog.NewModelVersions(databricksClient), + OnlineTables: catalog.NewOnlineTables(databricksClient), + PermissionMigration: iam.NewPermissionMigration(databricksClient), + Permissions: iam.NewPermissions(databricksClient), + Pipelines: pipelines.NewPipelines(databricksClient), + PolicyFamilies: compute.NewPolicyFamilies(databricksClient), + ProviderExchangeFilters: marketplace.NewProviderExchangeFilters(databricksClient), + ProviderExchanges: marketplace.NewProviderExchanges(databricksClient), + ProviderFiles: marketplace.NewProviderFiles(databricksClient), + ProviderListings: marketplace.NewProviderListings(databricksClient), + ProviderPersonalizationRequests: marketplace.NewProviderPersonalizationRequests(databricksClient), ProviderProviderAnalyticsDashboards: marketplace.NewProviderProviderAnalyticsDashboards(databricksClient), - - ProviderProviders: marketplace.NewProviderProviders(databricksClient), - - Providers: sharing.NewProviders(databricksClient), - - QualityMonitors: catalog.NewQualityMonitors(databricksClient), - - Queries: sql.NewQueries(databricksClient), - - QueryHistory: sql.NewQueryHistory(databricksClient), - - QueryVisualizations: sql.NewQueryVisualizations(databricksClient), - - RecipientActivation: sharing.NewRecipientActivation(databricksClient), - - Recipients: sharing.NewRecipients(databricksClient), - - RegisteredModels: catalog.NewRegisteredModels(databricksClient), - - Repos: workspace.NewRepos(databricksClient), - - Schemas: catalog.NewSchemas(databricksClient), - - Secrets: workspace.NewSecrets(databricksClient), - - ServicePrincipals: iam.NewServicePrincipals(databricksClient), - - ServingEndpointDataPlane: serving.NewServingEndpointDataPlane(databricksClient, servingEndpoints), - - ServingEndpoints: servingEndpoints, - - Settings: settings.NewSettings(databricksClient), - - Shares: sharing.NewShares(databricksClient), - - StatementExecution: sql.NewStatementExecution(databricksClient), - - StorageCredentials: catalog.NewStorageCredentials(databricksClient), - - SystemSchemas: catalog.NewSystemSchemas(databricksClient), - - TableConstraints: catalog.NewTableConstraints(databricksClient), - - Tables: catalog.NewTables(databricksClient), - - TokenManagement: settings.NewTokenManagement(databricksClient), - - Tokens: settings.NewTokens(databricksClient), - - Users: iam.NewUsers(databricksClient), - - VectorSearchEndpoints: vectorsearch.NewVectorSearchEndpoints(databricksClient), - - VectorSearchIndexes: vectorsearch.NewVectorSearchIndexes(databricksClient), - - Volumes: catalog.NewVolumes(databricksClient), - - Warehouses: sql.NewWarehouses(databricksClient), - - Workspace: workspace.NewWorkspace(databricksClient), - - WorkspaceBindings: catalog.NewWorkspaceBindings(databricksClient), - - WorkspaceConf: settings.NewWorkspaceConf(databricksClient), + ProviderProviders: marketplace.NewProviderProviders(databricksClient), + Providers: sharing.NewProviders(databricksClient), + QualityMonitors: catalog.NewQualityMonitors(databricksClient), + Queries: sql.NewQueries(databricksClient), + QueryHistory: sql.NewQueryHistory(databricksClient), + QueryVisualizations: sql.NewQueryVisualizations(databricksClient), + RecipientActivation: sharing.NewRecipientActivation(databricksClient), + Recipients: sharing.NewRecipients(databricksClient), + RegisteredModels: catalog.NewRegisteredModels(databricksClient), + Repos: workspace.NewRepos(databricksClient), + Schemas: catalog.NewSchemas(databricksClient), + Secrets: workspace.NewSecrets(databricksClient), + ServicePrincipals: iam.NewServicePrincipals(databricksClient), + ServingEndpointDataPlane: serving.NewServingEndpointDataPlane(databricksClient, servingEndpoints), + ServingEndpoints: servingEndpoints, + Settings: settings.NewSettings(databricksClient), + Shares: sharing.NewShares(databricksClient), + StatementExecution: sql.NewStatementExecution(databricksClient), + StorageCredentials: catalog.NewStorageCredentials(databricksClient), + SystemSchemas: catalog.NewSystemSchemas(databricksClient), + TableConstraints: catalog.NewTableConstraints(databricksClient), + Tables: catalog.NewTables(databricksClient), + TokenManagement: settings.NewTokenManagement(databricksClient), + Tokens: settings.NewTokens(databricksClient), + Users: iam.NewUsers(databricksClient), + VectorSearchEndpoints: vectorsearch.NewVectorSearchEndpoints(databricksClient), + VectorSearchIndexes: vectorsearch.NewVectorSearchIndexes(databricksClient), + Volumes: catalog.NewVolumes(databricksClient), + Warehouses: sql.NewWarehouses(databricksClient), + Workspace: workspace.NewWorkspace(databricksClient), + WorkspaceBindings: catalog.NewWorkspaceBindings(databricksClient), + WorkspaceConf: settings.NewWorkspaceConf(databricksClient), }, nil } From e8b0d3cd7eb439e59db83686b3e3da9bd4d98293 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 10:06:37 +0200 Subject: [PATCH 19/24] fixes 2 --- .codegen/mock_account_client.go.tmpl | 4 ++-- .codegen/mock_workspace_client.go.tmpl | 6 +++--- experimental/mocks/mock_workspace_client.go | 9 +++++++++ service/iam/model.go | 6 +++--- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/.codegen/mock_account_client.go.tmpl b/.codegen/mock_account_client.go.tmpl index 035cfc542..0f668d176 100644 --- a/.codegen/mock_account_client.go.tmpl +++ b/.codegen/mock_account_client.go.tmpl @@ -22,7 +22,7 @@ func NewMockAccountClient(t interface { cli := &MockAccountClient { AccountClient: &databricks.AccountClient{ Config: nil, - {{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} + {{range .Services}}{{if and .IsAccounts (not .HasParent)}} {{(.TrimPrefix "account").PascalName}}: {{ template "mock-interface-constructor" . }}(t), {{- end}}{{end}} }, @@ -50,7 +50,7 @@ func(m *MockAccountClient) GetMock{{.PascalName}}API() *{{ template "mock-interf } {{end}}{{end}} -{{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane)}} +{{range .Services}}{{if and .IsAccounts (not .HasParent)}} func(m *MockAccountClient) GetMock{{.Name}}API() *{{ template "mock-interface-name" . }} { api, ok := m.AccountClient.{{(.TrimPrefix "account").PascalName}}.(*{{ template "mock-interface-name" . }}) if !ok { diff --git a/.codegen/mock_workspace_client.go.tmpl b/.codegen/mock_workspace_client.go.tmpl index ef4aad8f2..ef6d49e3b 100644 --- a/.codegen/mock_workspace_client.go.tmpl +++ b/.codegen/mock_workspace_client.go.tmpl @@ -22,7 +22,7 @@ func NewMockWorkspaceClient(t interface { cli := &MockWorkspaceClient{ WorkspaceClient: &databricks.WorkspaceClient{ Config: nil, - {{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} + {{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} {{.Name}}: {{ template "mock-interface-constructor" . }}(t), {{- end}}{{end}} }, @@ -40,7 +40,7 @@ func NewMockWorkspaceClient(t interface { return cli } -{{range .Services}}{{if and (not .IsAccounts) (.HasParent) (not .IsDataPlane)}} +{{range .Services}}{{if and (not .IsAccounts) (.HasParent)}} func(m *MockWorkspaceClient) GetMock{{.PascalName}}API() *{{ template "mock-interface-name" . }} { api, ok := m.GetMock{{.ParentService.Name}}API().{{.PascalName}}().(*{{ template "mock-interface-name" . }}) if !ok { @@ -51,7 +51,7 @@ func(m *MockWorkspaceClient) GetMock{{.PascalName}}API() *{{ template "mock-inte {{end}}{{end}} -{{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}} +{{range .Services}}{{if and (not .IsAccounts) (not .HasParent)}} func(m *MockWorkspaceClient) GetMock{{.Name}}API() *{{ template "mock-interface-name" . }} { api, ok := m.WorkspaceClient.{{.Name}}.(*{{ template "mock-interface-name" . }}) if !ok { 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/service/iam/model.go b/service/iam/model.go index 76321c4b8..b4aa4acd5 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" url:"-"` + Id string `json:"id,omitempty"` // 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. From bb5d6181187d0cc5de1c12fe44e5433e65f55b92 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 10:32:02 +0200 Subject: [PATCH 20/24] Cleanup --- .codegen/accounts.go.tmpl | 2 +- .codegen/api.go.tmpl | 4 ++-- .codegen/impl.go.tmpl | 10 +++++----- .codegen/interface.go.tmpl | 2 +- .codegen/workspaces.go.tmpl | 2 +- openapi/code/method.go | 2 +- openapi/code/package.go | 10 +++++++--- openapi/code/service.go | 7 ++++--- service/iam/model.go | 4 ++-- service/marketplace/model.go | 4 ++-- 10 files changed, 26 insertions(+), 21 deletions(-) diff --git a/.codegen/accounts.go.tmpl b/.codegen/accounts.go.tmpl index d948b7acd..d0487a06f 100644 --- a/.codegen/accounts.go.tmpl +++ b/.codegen/accounts.go.tmpl @@ -54,7 +54,7 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) { {{(.TrimPrefix "account").PascalName}}: {{.CamelName}}, {{- end -}} {{if and .IsAccounts (not .HasParent) .IsDataPlane}} - {{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient, {{.ParentService.CamelName}}), + {{(.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 c035b4834..7880c1d0b 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -107,7 +107,7 @@ type {{.PascalName}}Interface interface { func New{{.PascalName}}(client *client.DatabricksClient, {{- if .IsDataPlane}} -controlPlane *{{.ParentService.PascalName}}API, +controlPlane *{{.ControlPlaneService.PascalName}}API, {{end -}} ) *{{.PascalName}}API { return &{{.PascalName}}API{ @@ -126,7 +126,7 @@ controlPlane *{{.ParentService.PascalName}}API, {{- if not .IsDataPlane}} {{.Comment "// " 80}} {{else}} -// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ParentService.PascalName}} service. +// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ControlPlaneService.PascalName}} service. {{end -}} type {{.PascalName}}API struct { // impl contains low-level REST API interface, that could be overridden diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 5541597c8..73c37e87e 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -20,7 +20,7 @@ import ( type {{.CamelName}}Impl struct { {{- if .IsDataPlane}} dataplane.DataPlaneTokenCache - controlPlane *{{.ParentService.PascalName}}API + controlPlane *{{.ControlPlaneService.PascalName}}API {{end -}} client *client.DatabricksClient } @@ -28,8 +28,8 @@ type {{.CamelName}}Impl struct { {{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.ParentService.DataPlaneInfoMethod.Request.PascalName}}{ - {{- range .Service.ParentService.DataPlaneInfoMethod.Request.Fields}} + getRequest := {{.Service.ControlPlaneService.DataPlaneInfoMethod.Request.PascalName}}{ + {{- range .Service.ControlPlaneService.DataPlaneInfoMethod.Request.Fields}} {{.PascalName}}: request.{{.PascalName}}, {{end}} } @@ -38,7 +38,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re return nil, err } infoGetter := func() (*oauth2.DataPlaneInfo, error) { - response, err := a.controlPlane.{{.Service.ParentService.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) + response, err := a.controlPlane.{{.Service.ControlPlaneService.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) if err != nil { return nil, err } @@ -51,7 +51,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) } getParams := []string{ - {{- range .Service.ParentService.DataPlaneInfoMethod.Request.Fields}} + {{- range .Service.ControlPlaneService.DataPlaneInfoMethod.Request.Fields}} request.{{.PascalName}}, {{end -}} } diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index 2c55dc227..0eceba013 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -12,7 +12,7 @@ import ( {{- if not .IsDataPlane}} {{.Comment "// " 80}} {{else}} -// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ParentService.PascalName}} service. +// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ControlPlaneService.PascalName}} service. {{end -}} type {{.PascalName}}Service interface { {{range .Methods}} diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index f5ef4dc05..00fecb949 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.codegen/workspaces.go.tmpl @@ -61,7 +61,7 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { {{.Name}}: {{.CamelName}}, {{- end -}} {{if and (not .IsAccounts) (not .HasParent) .IsDataPlane}} - {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient, {{.ParentService.CamelName}}), + {{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient, {{.ControlPlaneService.CamelName}}), {{- end -}} {{end}} }, nil diff --git a/openapi/code/method.go b/openapi/code/method.go index 75f41f8cc..71158f300 100644 --- a/openapi/code/method.go +++ b/openapi/code/method.go @@ -58,7 +58,7 @@ func (m *Method) DataPlaneInfoFields() []*Field { if m.DataPlane == nil { return nil } - method := m.Service.ParentService.DataPlaneInfoMethod() + method := m.Service.ControlPlaneService.DataPlaneInfoMethod() fieldNames := m.DataPlane.Fields currentLevelFields := method.Response.fields fields := []*Field{} diff --git a/openapi/code/package.go b/openapi/code/package.go index 04eb9c6a0..03d7320f9 100644 --- a/openapi/code/package.go +++ b/openapi/code/package.go @@ -32,9 +32,6 @@ func (pkg *Package) FullName() string { func (pkg *Package) Services() (types []*Service) { for _, v := range pkg.services { types = append(types, v) - if v.HasDataPlaneMethods() { - types = append(types, v.DataPlaneService()) - } } pascalNameSort(types) return types @@ -423,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 8ca49a404..6e2c7d30e 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -30,6 +30,7 @@ type Service struct { subservices map[string]*Service ByPathParamsMethods []*Shortcut ParentService *Service + ControlPlaneService *Service tag *openapi.Tag IsDataPlane bool } @@ -50,11 +51,11 @@ func (svc *Service) DataPlaneInfoMethod() *Method { } // Returns the corresponding service for DataPlane APIs. -func (svc *Service) DataPlaneService() *Service { +func (svc *Service) generateDataPlaneService() *Service { s := &Service{ Named: Named{svc.Named.Singular().Name + "DataPlane", svc.Description}, Package: svc.Package, - ParentService: svc, + ControlPlaneService: svc, methods: svc.dataPlaneMethods(), tag: svc.tag, ByPathParamsMethods: svc.ByPathParamsMethods, @@ -66,7 +67,7 @@ func (svc *Service) DataPlaneService() *Service { return s } -// Returns a sorted slice of methods which support direct DataPlane access. +// Copies methods which implement a DataPlane API. func (svc *Service) dataPlaneMethods() map[string]*Method { methods := map[string]*Method{} for _, m := range svc.methods { diff --git a/service/iam/model.go b/service/iam/model.go index b4aa4acd5..52d3de756 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" url:"-"` + Id string `json:"id,omitempty"` 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. diff --git a/service/marketplace/model.go b/service/marketplace/model.go index efc78aa57..92e7d1939 100755 --- a/service/marketplace/model.go +++ b/service/marketplace/model.go @@ -637,9 +637,9 @@ func (s FileInfo) MarshalJSON() ([]byte, error) { } type FileParent struct { - FileParentType FileParentType `json:"file_parent_type,omitempty"` + FileParentType FileParentType `json:"file_parent_type,omitempty" url:"file_parent_type,omitempty"` // TODO make the following fields required - ParentId string `json:"parent_id,omitempty"` + ParentId string `json:"parent_id,omitempty" url:"parent_id,omitempty"` ForceSendFields []string `json:"-"` } From c161714816bfed79e0ee4f524d783625876930d7 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 10:42:06 +0200 Subject: [PATCH 21/24] c2 --- .codegen/impl.go.tmpl | 8 ++++---- openapi/code/method.go | 3 ++- openapi/code/service.go | 3 ++- service/iam/model.go | 4 ++-- service/marketplace/model.go | 4 ++-- service/sql/model.go | 1 + 6 files changed, 13 insertions(+), 10 deletions(-) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index 73c37e87e..f95b3eb8c 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -28,8 +28,8 @@ type {{.CamelName}}Impl struct { {{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.ControlPlaneService.DataPlaneInfoMethod.Request.PascalName}}{ - {{- range .Service.ControlPlaneService.DataPlaneInfoMethod.Request.Fields}} + getRequest := {{.Service.DataPlaneInfoMethod.Request.PascalName}}{ + {{- range .Service.DataPlaneInfoMethod.Request.Fields}} {{.PascalName}}: request.{{.PascalName}}, {{end}} } @@ -38,7 +38,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re return nil, err } infoGetter := func() (*oauth2.DataPlaneInfo, error) { - response, err := a.controlPlane.{{.Service.ControlPlaneService.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) + response, err := a.controlPlane.{{.Service.DataPlaneInfoMethod.PascalName}}(ctx, getRequest) if err != nil { return nil, err } @@ -51,7 +51,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) } getParams := []string{ - {{- range .Service.ControlPlaneService.DataPlaneInfoMethod.Request.Fields}} + {{- range .Service.DataPlaneInfoMethod.Request.Fields}} request.{{.PascalName}}, {{end -}} } diff --git a/openapi/code/method.go b/openapi/code/method.go index 71158f300..57db1d14a 100644 --- a/openapi/code/method.go +++ b/openapi/code/method.go @@ -54,11 +54,12 @@ 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.ControlPlaneService.DataPlaneInfoMethod() + method := m.Service.DataPlaneInfoMethod() fieldNames := m.DataPlane.Fields currentLevelFields := method.Response.fields fields := []*Field{} diff --git a/openapi/code/service.go b/openapi/code/service.go index 6e2c7d30e..ca666f797 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -40,6 +40,7 @@ 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 { @@ -47,7 +48,7 @@ func (svc *Service) DataPlaneInfoMethod() *Method { methodName = m.DataPlane.ConfigMethod } } - return svc.methods[methodName] + return svc.ControlPlaneService.methods[methodName] } // Returns the corresponding service for DataPlane APIs. diff --git a/service/iam/model.go b/service/iam/model.go index 52d3de756..55e4ccf6b 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -1231,7 +1231,7 @@ type ServicePrincipal struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks service principal ID. - Id string `json:"id,omitempty" url:"-"` + Id string `json:"id,omitempty"` // 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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/service/marketplace/model.go b/service/marketplace/model.go index 92e7d1939..efc78aa57 100755 --- a/service/marketplace/model.go +++ b/service/marketplace/model.go @@ -637,9 +637,9 @@ func (s FileInfo) MarshalJSON() ([]byte, error) { } type FileParent struct { - FileParentType FileParentType `json:"file_parent_type,omitempty" url:"file_parent_type,omitempty"` + FileParentType FileParentType `json:"file_parent_type,omitempty"` // TODO make the following fields required - ParentId string `json:"parent_id,omitempty" url:"parent_id,omitempty"` + ParentId string `json:"parent_id,omitempty"` ForceSendFields []string `json:"-"` } diff --git a/service/sql/model.go b/service/sql/model.go index e08608ba1..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` From 6eb49519aa5034fa80924e6a19dcf1005fadabaf Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 10:53:59 +0200 Subject: [PATCH 22/24] Rename --- .codegen/impl.go.tmpl | 4 ++-- dataplane/data_plane.go | 6 +++--- dataplane/data_plane_test.go | 20 ++++++++++---------- service/iam/model.go | 4 ++-- service/serving/impl.go | 4 ++-- service/sql/model.go | 1 - 6 files changed, 19 insertions(+), 20 deletions(-) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index f95b3eb8c..cf18bde2a 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -19,7 +19,7 @@ import ( // unexported type that holds implementations of just {{.Name}} API methods type {{.CamelName}}Impl struct { {{- if .IsDataPlane}} - dataplane.DataPlaneTokenCache + dataplane.DataPlaneHelper controlPlane *{{.ControlPlaneService.PascalName}}API {{end -}} client *client.DatabricksClient @@ -55,7 +55,7 @@ func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Re request.{{.PascalName}}, {{end -}} } - endpointUrl, dataPlaneToken, err := a.GetDataPlane("{{.PascalName}}", getParams, refresh, infoGetter) + endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("{{.PascalName}}", getParams, refresh, infoGetter) if err != nil { return nil, err } diff --git a/dataplane/data_plane.go b/dataplane/data_plane.go index 8c40a3d4d..d717ee306 100644 --- a/dataplane/data_plane.go +++ b/dataplane/data_plane.go @@ -7,12 +7,12 @@ import ( "golang.org/x/oauth2" ) -type DataPlaneTokenCache struct { +type DataPlaneHelper struct { infos map[string]*dp.DataPlaneInfo tokens map[string]*oauth2.Token } -func (o *DataPlaneTokenCache) GetDataPlane(method string, params []string, refresh func(*dp.DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*dp.DataPlaneInfo, error)) (string, *oauth2.Token, error) { +func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, refresh func(*dp.DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*dp.DataPlaneInfo, error)) (string, *oauth2.Token, error) { if o.infos == nil { o.infos = make(map[string]*dp.DataPlaneInfo) } @@ -44,7 +44,7 @@ func (o *DataPlaneTokenCache) GetDataPlane(method string, params []string, refre return info.EndpointUrl, token, nil } -func (o *DataPlaneTokenCache) generateKey(method string, params []string) string { +func (o *DataPlaneHelper) generateKey(method string, params []string) string { allElements := []string{method} allElements = append(allElements, params...) return strings.Join(allElements, "/") diff --git a/dataplane/data_plane_test.go b/dataplane/data_plane_test.go index a002fe4ca..c9b2c1145 100644 --- a/dataplane/data_plane_test.go +++ b/dataplane/data_plane_test.go @@ -49,8 +49,8 @@ func TestTokenNotCached(t *testing.T) { }, err: nil, } - c := DataPlaneTokenCache{} - url, token, err := c.GetDataPlane("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + 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) @@ -76,12 +76,12 @@ func TestTokenCached(t *testing.T) { }, err: nil, } - c := DataPlaneTokenCache{} + c := DataPlaneHelper{} c.infos = make(map[string]*dp.DataPlaneInfo) c.tokens = make(map[string]*oauth2.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = s.token - url, token, err := c.GetDataPlane("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + 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) @@ -113,12 +113,12 @@ func TestTokenExpired(t *testing.T) { }, err: nil, } - c := DataPlaneTokenCache{} + c := DataPlaneHelper{} c.infos = make(map[string]*dp.DataPlaneInfo) c.tokens = make(map[string]*oauth2.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = expired - url, token, err := c.GetDataPlane("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + 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) @@ -134,8 +134,8 @@ func TestTokenInfoError(t *testing.T) { err: assert.AnError, } s := tokenRefreshSpy{} - c := DataPlaneTokenCache{} - url, token, err := c.GetDataPlane("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + 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) @@ -155,8 +155,8 @@ func TestTokenRefreshError(t *testing.T) { token: nil, err: assert.AnError, } - c := DataPlaneTokenCache{} - url, token, err := c.GetDataPlane("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter) + 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) diff --git a/service/iam/model.go b/service/iam/model.go index 55e4ccf6b..52d3de756 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -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/serving/impl.go b/service/serving/impl.go index 26009b09d..d69f01d79 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -117,7 +117,7 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, // unexported type that holds implementations of just ServingEndpointDataPlane API methods type servingEndpointDataPlaneImpl struct { - dataplane.DataPlaneTokenCache + dataplane.DataPlaneHelper controlPlane *ServingEndpointsAPI client *client.DatabricksClient } @@ -146,7 +146,7 @@ func (a *servingEndpointDataPlaneImpl) Query(ctx context.Context, request QueryE getParams := []string{ request.Name, } - endpointUrl, dataPlaneToken, err := a.GetDataPlane("Query", getParams, refresh, infoGetter) + endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("Query", getParams, refresh, infoGetter) if err != nil { return nil, err } diff --git a/service/sql/model.go b/service/sql/model.go index 345658ff6..e08608ba1 100755 --- a/service/sql/model.go +++ b/service/sql/model.go @@ -287,7 +287,6 @@ func (s ChannelInfo) MarshalJSON() ([]byte, error) { return marshal.Marshal(s) } -// Name of the channel type ChannelName string const ChannelNameChannelNameCurrent ChannelName = `CHANNEL_NAME_CURRENT` From 2a4517d3adae014561edfa31e35dac9dd4d2ca4e Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 11:04:34 +0200 Subject: [PATCH 23/24] Move --- .codegen/impl.go.tmpl | 4 ++-- service/iam/model.go | 4 ++-- {dataplane => service/oauth2}/data_plane.go | 9 ++++---- .../oauth2}/data_plane_test.go | 23 +++++++++---------- service/serving/impl.go | 5 ++-- 5 files changed, 21 insertions(+), 24 deletions(-) rename {dataplane => service/oauth2}/data_plane.go (76%) rename {dataplane => service/oauth2}/data_plane_test.go (88%) diff --git a/.codegen/impl.go.tmpl b/.codegen/impl.go.tmpl index cf18bde2a..e2359e7d8 100644 --- a/.codegen/impl.go.tmpl +++ b/.codegen/impl.go.tmpl @@ -6,6 +6,7 @@ import ( "context" "errors" "fmt" + goauth "golang.org/x/oauth2" "time" "io" "net/http" @@ -13,13 +14,12 @@ import ( "github.com/databricks/databricks-sdk-go/httpclient" {{range .ImportedPackages}} "github.com/databricks/databricks-sdk-go/service/{{.}}"{{end}} - goauth "golang.org/x/oauth2" ) {{range .Services}} // unexported type that holds implementations of just {{.Name}} API methods type {{.CamelName}}Impl struct { {{- if .IsDataPlane}} - dataplane.DataPlaneHelper + oauth2.DataPlaneHelper controlPlane *{{.ControlPlaneService.PascalName}}API {{end -}} client *client.DatabricksClient diff --git a/service/iam/model.go b/service/iam/model.go index 52d3de756..55e4ccf6b 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -1231,7 +1231,7 @@ type ServicePrincipal struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks service principal ID. - Id string `json:"id,omitempty" url:"-"` + Id string `json:"id,omitempty"` // 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"` + Id string `json:"id,omitempty" url:"-"` Name *Name `json:"name,omitempty"` // Corresponds to AWS instance profile/arn role. diff --git a/dataplane/data_plane.go b/service/oauth2/data_plane.go similarity index 76% rename from dataplane/data_plane.go rename to service/oauth2/data_plane.go index d717ee306..fec3ef155 100644 --- a/dataplane/data_plane.go +++ b/service/oauth2/data_plane.go @@ -1,20 +1,19 @@ -package dataplane +package oauth2 import ( "strings" - dp "github.com/databricks/databricks-sdk-go/service/oauth2" "golang.org/x/oauth2" ) type DataPlaneHelper struct { - infos map[string]*dp.DataPlaneInfo + infos map[string]*DataPlaneInfo tokens map[string]*oauth2.Token } -func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, refresh func(*dp.DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*dp.DataPlaneInfo, error)) (string, *oauth2.Token, error) { +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]*dp.DataPlaneInfo) + o.infos = make(map[string]*DataPlaneInfo) } if o.tokens == nil { o.tokens = make(map[string]*oauth2.Token) diff --git a/dataplane/data_plane_test.go b/service/oauth2/data_plane_test.go similarity index 88% rename from dataplane/data_plane_test.go rename to service/oauth2/data_plane_test.go index c9b2c1145..53647606c 100644 --- a/dataplane/data_plane_test.go +++ b/service/oauth2/data_plane_test.go @@ -1,33 +1,32 @@ -package dataplane +package oauth2 import ( "testing" "time" - dp "github.com/databricks/databricks-sdk-go/service/oauth2" "github.com/stretchr/testify/assert" "golang.org/x/oauth2" ) type infoMock struct { called bool - info *dp.DataPlaneInfo + info *DataPlaneInfo err error } -func (i *infoMock) DataPlaneInfoGetter() (*dp.DataPlaneInfo, error) { +func (i *infoMock) DataPlaneInfoGetter() (*DataPlaneInfo, error) { i.called = true return i.info, i.err } type tokenRefreshSpy struct { called bool - expectedInfo *dp.DataPlaneInfo + expectedInfo *DataPlaneInfo token *oauth2.Token err error } -func (t *tokenRefreshSpy) TokenRefresh(info *dp.DataPlaneInfo) (*oauth2.Token, error) { +func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*oauth2.Token, error) { t.expectedInfo = info t.called = true return t.token, t.err @@ -35,7 +34,7 @@ func (t *tokenRefreshSpy) TokenRefresh(info *dp.DataPlaneInfo) (*oauth2.Token, e func TestTokenNotCached(t *testing.T) { info := infoMock{ - info: &dp.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -62,7 +61,7 @@ func TestTokenNotCached(t *testing.T) { func TestTokenCached(t *testing.T) { info := infoMock{ - info: &dp.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -77,7 +76,7 @@ func TestTokenCached(t *testing.T) { err: nil, } c := DataPlaneHelper{} - c.infos = make(map[string]*dp.DataPlaneInfo) + 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 @@ -93,7 +92,7 @@ func TestTokenCached(t *testing.T) { func TestTokenExpired(t *testing.T) { info := infoMock{ - info: &dp.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, @@ -114,7 +113,7 @@ func TestTokenExpired(t *testing.T) { err: nil, } c := DataPlaneHelper{} - c.infos = make(map[string]*dp.DataPlaneInfo) + c.infos = make(map[string]*DataPlaneInfo) c.tokens = make(map[string]*oauth2.Token) c.infos["method/params"] = info.info c.tokens["method/params"] = expired @@ -145,7 +144,7 @@ func TestTokenInfoError(t *testing.T) { func TestTokenRefreshError(t *testing.T) { info := infoMock{ - info: &dp.DataPlaneInfo{ + info: &DataPlaneInfo{ EndpointUrl: "url", AuthorizationDetails: "authDetails", }, diff --git a/service/serving/impl.go b/service/serving/impl.go index d69f01d79..b49ad6fd3 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -9,11 +9,10 @@ import ( "net/http" "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/dataplane" "github.com/databricks/databricks-sdk-go/httpclient" + goauth "golang.org/x/oauth2" "github.com/databricks/databricks-sdk-go/service/oauth2" - goauth "golang.org/x/oauth2" ) // unexported type that holds implementations of just Apps API methods @@ -117,7 +116,7 @@ func (a *appsImpl) Update(ctx context.Context, request UpdateAppRequest) (*App, // unexported type that holds implementations of just ServingEndpointDataPlane API methods type servingEndpointDataPlaneImpl struct { - dataplane.DataPlaneHelper + oauth2.DataPlaneHelper controlPlane *ServingEndpointsAPI client *client.DatabricksClient } From e86f8991968868d8d8814f0736a4f369947bcd4d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 29 May 2024 13:12:48 +0200 Subject: [PATCH 24/24] Extract description generation and tag --- .codegen/api.go.tmpl | 4 ---- .codegen/interface.go.tmpl | 4 ---- openapi/code/service.go | 19 +++++++++++++++++-- service/iam/model.go | 6 +++--- service/pkg.go | 2 +- service/serving/api.go | 3 ++- service/serving/interface.go | 3 ++- service/sql/model.go | 1 + workspace_client.go | 16 ++-------------- 9 files changed, 28 insertions(+), 30 deletions(-) diff --git a/.codegen/api.go.tmpl b/.codegen/api.go.tmpl index 7880c1d0b..0be5b70cc 100644 --- a/.codegen/api.go.tmpl +++ b/.codegen/api.go.tmpl @@ -123,11 +123,7 @@ controlPlane *{{.ControlPlaneService.PascalName}}API, } } -{{- if not .IsDataPlane}} {{.Comment "// " 80}} -{{else}} -// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ControlPlaneService.PascalName}} service. -{{end -}} type {{.PascalName}}API struct { // impl contains low-level REST API interface, that could be overridden // through WithImpl({{.PascalName}}Service) diff --git a/.codegen/interface.go.tmpl b/.codegen/interface.go.tmpl index 0eceba013..a2eb01566 100644 --- a/.codegen/interface.go.tmpl +++ b/.codegen/interface.go.tmpl @@ -9,11 +9,7 @@ import ( ) {{range .Services}} -{{- if not .IsDataPlane}} {{.Comment "// " 80}} -{{else}} -// {{.PascalName}}Service provides a set of operations to interact with DataPlane endpoints for {{.ControlPlaneService.PascalName}} service. -{{end -}} type {{.PascalName}}Service interface { {{range .Methods}} {{.Comment " // " 80}} diff --git a/openapi/code/service.go b/openapi/code/service.go index ca666f797..1bdb590a4 100644 --- a/openapi/code/service.go +++ b/openapi/code/service.go @@ -53,18 +53,33 @@ func (svc *Service) DataPlaneInfoMethod() *Method { // 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{svc.Named.Singular().Name + "DataPlane", svc.Description}, + Named: named, Package: svc.Package, ControlPlaneService: svc, methods: svc.dataPlaneMethods(), - tag: svc.tag, + 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 } 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/pkg.go b/service/pkg.go index aaea39732..e97b630c3 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -178,7 +178,7 @@ // // - [iam.AccountServicePrincipalsAPI]: Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD platforms. // -// - [serving.ServingEndpointDataPlaneAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. +// - [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. // diff --git a/service/serving/api.go b/service/serving/api.go index f57d1dbf6..efb4e9a18 100755 --- a/service/serving/api.go +++ b/service/serving/api.go @@ -577,7 +577,8 @@ func NewServingEndpointDataPlane(client *client.DatabricksClient, } } -// ServingEndpointDataPlaneService provides a set of operations to interact with DataPlane endpoints for ServingEndpoints service. +// 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) diff --git a/service/serving/interface.go b/service/serving/interface.go index 0d122b3c0..dc29deb4f 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -67,7 +67,8 @@ type AppsService interface { Update(ctx context.Context, request UpdateAppRequest) (*App, error) } -// ServingEndpointDataPlaneService provides a set of operations to interact with DataPlane endpoints for ServingEndpoints service. +// ServingEndpointDataPlane provides a set of operations to interact with +// DataPlane endpoints for ServingEndpoints service. type ServingEndpointDataPlaneService interface { // Query a serving endpoint. diff --git a/service/sql/model.go b/service/sql/model.go index e08608ba1..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` diff --git a/workspace_client.go b/workspace_client.go index 8ad3c424a..b0faa739b 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -696,20 +696,8 @@ type WorkspaceClient struct { // data by accident. ServicePrincipals iam.ServicePrincipalsInterface - // The Serving Endpoints API allows you to create, update, and delete model - // serving endpoints. - // - // You can use a serving endpoint to serve models from the Databricks Model - // Registry or from Unity Catalog. Endpoints expose the underlying models as - // scalable REST API endpoints using serverless compute. This means the - // endpoints and associated compute resources are fully managed by - // Databricks and will not appear in your cloud account. A serving endpoint - // can consist of one or more MLflow models from the Databricks Model - // Registry, called served entities. A serving endpoint can have at most ten - // served entities. You can configure traffic settings to define how - // requests should be routed to your served entities behind an endpoint. - // Additionally, you can configure the scale of resources that should be - // applied to each served entity. + // 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