From d3547b7d44512032ee01fddd12d67e30bd884d16 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 10 Jun 2024 08:35:04 +0200 Subject: [PATCH] Generate --- .codegen/_openapi_sha | 2 +- experimental/mocks/mock_workspace_client.go | 18 ++-- .../catalog/mock_endpoints_interface.go | 95 ------------------- service/catalog/api.go | 84 ---------------- service/catalog/catalogs_usage_test.go | 4 - service/catalog/impl.go | 33 ------- service/catalog/interface.go | 13 --- service/iam/model.go | 6 +- service/pkg.go | 6 +- service/serving/api.go | 54 ++++++++++- service/serving/impl.go | 53 +++++++++++ service/serving/interface.go | 8 ++ service/vectorsearch/model.go | 2 + workspace_client.go | 9 +- 14 files changed, 137 insertions(+), 250 deletions(-) diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 2efdf11e6..10657f9db 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -fc26589775a1c48b14211cd11b6903d1b4ed071c \ No newline at end of file +f9166063cb5a6677c83877de2a2bb95c82958dd5 \ No newline at end of file diff --git a/experimental/mocks/mock_workspace_client.go b/experimental/mocks/mock_workspace_client.go index d5736f5d4..0f97337cc 100755 --- a/experimental/mocks/mock_workspace_client.go +++ b/experimental/mocks/mock_workspace_client.go @@ -61,7 +61,6 @@ func NewMockWorkspaceClient(t interface { DataSources: sql.NewMockDataSourcesInterface(t), Dbfs: files.NewMockDbfsInterface(t), DbsqlPermissions: sql.NewMockDbsqlPermissionsInterface(t), - Endpoints: catalog.NewMockEndpointsInterface(t), Experiments: ml.NewMockExperimentsInterface(t), ExternalLocations: catalog.NewMockExternalLocationsInterface(t), Files: files.NewMockFilesInterface(t), @@ -104,6 +103,7 @@ func NewMockWorkspaceClient(t interface { Secrets: workspace.NewMockSecretsInterface(t), ServicePrincipals: iam.NewMockServicePrincipalsInterface(t), ServingEndpoints: serving.NewMockServingEndpointsInterface(t), + ServingEndpointsDataPlane: serving.NewMockServingEndpointsDataPlaneInterface(t), Settings: settings.NewMockSettingsInterface(t), Shares: sharing.NewMockSharesInterface(t), StatementExecution: sql.NewMockStatementExecutionInterface(t), @@ -360,14 +360,6 @@ func (m *MockWorkspaceClient) GetMockDbsqlPermissionsAPI() *sql.MockDbsqlPermiss return api } -func (m *MockWorkspaceClient) GetMockEndpointsAPI() *catalog.MockEndpointsInterface { - api, ok := m.WorkspaceClient.Endpoints.(*catalog.MockEndpointsInterface) - if !ok { - panic(fmt.Sprintf("expected Endpoints to be *catalog.MockEndpointsInterface, actual was %T", m.WorkspaceClient.Endpoints)) - } - return api -} - func (m *MockWorkspaceClient) GetMockExperimentsAPI() *ml.MockExperimentsInterface { api, ok := m.WorkspaceClient.Experiments.(*ml.MockExperimentsInterface) if !ok { @@ -704,6 +696,14 @@ func (m *MockWorkspaceClient) GetMockServingEndpointsAPI() *serving.MockServingE return api } +func (m *MockWorkspaceClient) GetMockServingEndpointsDataPlaneAPI() *serving.MockServingEndpointsDataPlaneInterface { + api, ok := m.WorkspaceClient.ServingEndpointsDataPlane.(*serving.MockServingEndpointsDataPlaneInterface) + if !ok { + panic(fmt.Sprintf("expected ServingEndpointsDataPlane to be *serving.MockServingEndpointsDataPlaneInterface, actual was %T", m.WorkspaceClient.ServingEndpointsDataPlane)) + } + return api +} + func (m *MockWorkspaceClient) GetMockSettingsAPI() *settings.MockSettingsInterface { api, ok := m.WorkspaceClient.Settings.(*settings.MockSettingsInterface) if !ok { diff --git a/experimental/mocks/service/catalog/mock_endpoints_interface.go b/experimental/mocks/service/catalog/mock_endpoints_interface.go index c2f86f08b..47d826255 100644 --- a/experimental/mocks/service/catalog/mock_endpoints_interface.go +++ b/experimental/mocks/service/catalog/mock_endpoints_interface.go @@ -294,101 +294,6 @@ func (_c *MockEndpointsInterface_GetByName_Call) RunAndReturn(run func(context.C return _c } -// Impl provides a mock function with given fields: -func (_m *MockEndpointsInterface) Impl() catalog.EndpointsService { - ret := _m.Called() - - if len(ret) == 0 { - panic("no return value specified for Impl") - } - - var r0 catalog.EndpointsService - if rf, ok := ret.Get(0).(func() catalog.EndpointsService); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(catalog.EndpointsService) - } - } - - return r0 -} - -// MockEndpointsInterface_Impl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Impl' -type MockEndpointsInterface_Impl_Call struct { - *mock.Call -} - -// Impl is a helper method to define mock.On call -func (_e *MockEndpointsInterface_Expecter) Impl() *MockEndpointsInterface_Impl_Call { - return &MockEndpointsInterface_Impl_Call{Call: _e.mock.On("Impl")} -} - -func (_c *MockEndpointsInterface_Impl_Call) Run(run func()) *MockEndpointsInterface_Impl_Call { - _c.Call.Run(func(args mock.Arguments) { - run() - }) - return _c -} - -func (_c *MockEndpointsInterface_Impl_Call) Return(_a0 catalog.EndpointsService) *MockEndpointsInterface_Impl_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockEndpointsInterface_Impl_Call) RunAndReturn(run func() catalog.EndpointsService) *MockEndpointsInterface_Impl_Call { - _c.Call.Return(run) - return _c -} - -// WithImpl provides a mock function with given fields: impl -func (_m *MockEndpointsInterface) WithImpl(impl catalog.EndpointsService) catalog.EndpointsInterface { - ret := _m.Called(impl) - - if len(ret) == 0 { - panic("no return value specified for WithImpl") - } - - var r0 catalog.EndpointsInterface - if rf, ok := ret.Get(0).(func(catalog.EndpointsService) catalog.EndpointsInterface); ok { - r0 = rf(impl) - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(catalog.EndpointsInterface) - } - } - - return r0 -} - -// MockEndpointsInterface_WithImpl_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WithImpl' -type MockEndpointsInterface_WithImpl_Call struct { - *mock.Call -} - -// WithImpl is a helper method to define mock.On call -// - impl catalog.EndpointsService -func (_e *MockEndpointsInterface_Expecter) WithImpl(impl interface{}) *MockEndpointsInterface_WithImpl_Call { - return &MockEndpointsInterface_WithImpl_Call{Call: _e.mock.On("WithImpl", impl)} -} - -func (_c *MockEndpointsInterface_WithImpl_Call) Run(run func(impl catalog.EndpointsService)) *MockEndpointsInterface_WithImpl_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(catalog.EndpointsService)) - }) - return _c -} - -func (_c *MockEndpointsInterface_WithImpl_Call) Return(_a0 catalog.EndpointsInterface) *MockEndpointsInterface_WithImpl_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockEndpointsInterface_WithImpl_Call) RunAndReturn(run func(catalog.EndpointsService) catalog.EndpointsInterface) *MockEndpointsInterface_WithImpl_Call { - _c.Call.Return(run) - return _c -} - // NewMockEndpointsInterface creates a new instance of MockEndpointsInterface. 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 NewMockEndpointsInterface(t interface { diff --git a/service/catalog/api.go b/service/catalog/api.go index 344a4e07b..04cbc1362 100755 --- a/service/catalog/api.go +++ b/service/catalog/api.go @@ -1108,90 +1108,6 @@ func (a *ConnectionsAPI) Update(ctx context.Context, request UpdateConnection) ( return a.impl.Update(ctx, request) } -type EndpointsInterface 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 MockEndpointsInterface instead. - WithImpl(impl EndpointsService) EndpointsInterface - - // Impl returns low-level Endpoints API implementation - // Deprecated: use MockEndpointsInterface instead. - Impl() EndpointsService - - // Create an Endpoint. - Create(ctx context.Context, request CreateEndpointRequest) (*Endpoint, error) - - // Delete an Endpoint. - Delete(ctx context.Context, request DeleteEndpointRequest) error - - // Delete an Endpoint. - DeleteByName(ctx context.Context, name string) error - - // Get an Endpoint. - Get(ctx context.Context, request GetEndpointRequest) (*Endpoint, error) - - // Get an Endpoint. - GetByName(ctx context.Context, name string) (*Endpoint, error) -} - -func NewEndpoints(client *client.DatabricksClient) *EndpointsAPI { - return &EndpointsAPI{ - impl: &endpointsImpl{ - client: client, - }, - } -} - -// Endpoints are used to connect to PG clusters. -type EndpointsAPI struct { - // impl contains low-level REST API interface, that could be overridden - // through WithImpl(EndpointsService) - impl EndpointsService -} - -// 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 MockEndpointsInterface instead. -func (a *EndpointsAPI) WithImpl(impl EndpointsService) EndpointsInterface { - a.impl = impl - return a -} - -// Impl returns low-level Endpoints API implementation -// Deprecated: use MockEndpointsInterface instead. -func (a *EndpointsAPI) Impl() EndpointsService { - return a.impl -} - -// Create an Endpoint. -func (a *EndpointsAPI) Create(ctx context.Context, request CreateEndpointRequest) (*Endpoint, error) { - return a.impl.Create(ctx, request) -} - -// Delete an Endpoint. -func (a *EndpointsAPI) Delete(ctx context.Context, request DeleteEndpointRequest) error { - return a.impl.Delete(ctx, request) -} - -// Delete an Endpoint. -func (a *EndpointsAPI) DeleteByName(ctx context.Context, name string) error { - return a.impl.Delete(ctx, DeleteEndpointRequest{ - Name: name, - }) -} - -// Get an Endpoint. -func (a *EndpointsAPI) Get(ctx context.Context, request GetEndpointRequest) (*Endpoint, error) { - return a.impl.Get(ctx, request) -} - -// Get an Endpoint. -func (a *EndpointsAPI) GetByName(ctx context.Context, name string) (*Endpoint, error) { - return a.impl.Get(ctx, GetEndpointRequest{ - Name: name, - }) -} - type ExternalLocationsInterface interface { // WithImpl could be used to override low-level API implementations for unit // testing purposes with [github.com/golang/mock] or other mocking frameworks. diff --git a/service/catalog/catalogs_usage_test.go b/service/catalog/catalogs_usage_test.go index cfe15d6d6..e2d1a696c 100755 --- a/service/catalog/catalogs_usage_test.go +++ b/service/catalog/catalogs_usage_test.go @@ -273,12 +273,8 @@ func ExampleCatalogsAPI_Update_catalogWorkspaceBindings() { logger.Infof(ctx, "found %v", created) _, err = w.Catalogs.Update(ctx, catalog.UpdateCatalog{ -<<<<<<< HEAD - Name: created.Name, -======= Name: created.Name, IsolationMode: catalog.IsolationModeIsolationModeIsolated, ->>>>>>> main }) if err != nil { panic(err) diff --git a/service/catalog/impl.go b/service/catalog/impl.go index 93a6a4928..312d9522c 100755 --- a/service/catalog/impl.go +++ b/service/catalog/impl.go @@ -294,39 +294,6 @@ func (a *connectionsImpl) Update(ctx context.Context, request UpdateConnection) return &connectionInfo, err } -// unexported type that holds implementations of just Endpoints API methods -type endpointsImpl struct { - client *client.DatabricksClient -} - -func (a *endpointsImpl) Create(ctx context.Context, request CreateEndpointRequest) (*Endpoint, error) { - var endpoint Endpoint - path := "/api/2.0/hadron-endpoints" - headers := make(map[string]string) - headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" - err := a.client.Do(ctx, http.MethodPost, path, headers, request, &endpoint) - return &endpoint, err -} - -func (a *endpointsImpl) Delete(ctx context.Context, request DeleteEndpointRequest) error { - var deleteResponse DeleteResponse - path := fmt.Sprintf("/api/2.0/hadron-endpoints/%v", request.Name) - headers := make(map[string]string) - headers["Accept"] = "application/json" - err := a.client.Do(ctx, http.MethodDelete, path, headers, request, &deleteResponse) - return err -} - -func (a *endpointsImpl) Get(ctx context.Context, request GetEndpointRequest) (*Endpoint, error) { - var endpoint Endpoint - path := fmt.Sprintf("/api/2.0/hadron-endpoints/%v", request.Name) - headers := make(map[string]string) - headers["Accept"] = "application/json" - err := a.client.Do(ctx, http.MethodGet, path, headers, request, &endpoint) - return &endpoint, err -} - // unexported type that holds implementations of just ExternalLocations API methods type externalLocationsImpl struct { client *client.DatabricksClient diff --git a/service/catalog/interface.go b/service/catalog/interface.go index 3aab6d738..f5c406a46 100755 --- a/service/catalog/interface.go +++ b/service/catalog/interface.go @@ -234,19 +234,6 @@ type ConnectionsService interface { Update(ctx context.Context, request UpdateConnection) (*ConnectionInfo, error) } -// Endpoints are used to connect to PG clusters. -type EndpointsService interface { - - // Create an Endpoint. - Create(ctx context.Context, request CreateEndpointRequest) (*Endpoint, error) - - // Delete an Endpoint. - Delete(ctx context.Context, request DeleteEndpointRequest) error - - // Get an Endpoint. - Get(ctx context.Context, request GetEndpointRequest) (*Endpoint, error) -} - // An external location is an object that combines a cloud storage path with a // storage credential that authorizes access to the cloud storage path. Each // external location is subject to Unity Catalog access-control policies that diff --git a/service/iam/model.go b/service/iam/model.go index c2f2067d1..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. @@ -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/pkg.go b/service/pkg.go index e5688291b..ab8d6c4c3 100644 --- a/service/pkg.go +++ b/service/pkg.go @@ -64,8 +64,6 @@ // // - [provisioning.EncryptionKeysAPI]: These APIs manage encryption key configurations for this workspace (optional). // -// - [catalog.EndpointsAPI]: Endpoints are used to connect to PG clusters. -// // - [settings.EnhancedSecurityMonitoringAPI]: Controls whether enhanced security monitoring is enabled for the current workspace. // // - [settings.EsmEnablementAccountAPI]: The enhanced security monitoring setting at the account level controls whether to enable the feature on new workspaces. @@ -182,6 +180,8 @@ // // - [serving.ServingEndpointsAPI]: The Serving Endpoints API allows you to create, update, and delete model serving endpoints. // +// - [serving.ServingEndpointsDataPlaneAPI]: Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints for Serving endpoints service. +// // - [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. @@ -289,7 +289,6 @@ var ( _ *sql.DbsqlPermissionsAPI = nil _ *settings.DefaultNamespaceAPI = nil _ *provisioning.EncryptionKeysAPI = nil - _ *catalog.EndpointsAPI = nil _ *settings.EnhancedSecurityMonitoringAPI = nil _ *settings.EsmEnablementAccountAPI = nil _ *ml.ExperimentsAPI = nil @@ -348,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 64e30bc0d..5de0600e6 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,55 @@ 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, + controlPlane *ServingEndpointsAPI, +) *ServingEndpointsDataPlaneAPI { + return &ServingEndpointsDataPlaneAPI{ + impl: &servingEndpointsDataPlaneImpl{ + client: client, + controlPlane: controlPlane, + }, + } +} + +// Serving endpoints DataPlane provides a set of operations to interact with +// data plane endpoints for Serving endpoints service. +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 c267d2002..949ebd875 100755 --- a/service/serving/impl.go +++ b/service/serving/impl.go @@ -4,10 +4,15 @@ package serving import ( "context" + "errors" "fmt" "net/http" "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/httpclient" + goauth "golang.org/x/oauth2" + + "github.com/databricks/databricks-sdk-go/service/oauth2" ) // unexported type that holds implementations of just Apps API methods @@ -264,3 +269,51 @@ 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 { + oauth2.DataPlaneHelper + controlPlane *ServingEndpointsAPI + client *client.DatabricksClient +} + +func (a *servingEndpointsDataPlaneImpl) Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) { + getRequest := GetServingEndpointRequest{ + Name: request.Name, + } + token, err := a.client.Config.GetToken() + if err != nil { + return nil, err + } + infoGetter := func() (*oauth2.DataPlaneInfo, error) { + response, err := a.controlPlane.Get(ctx, getRequest) + if err != nil { + return nil, err + } + if response.DataPlaneInfo == nil { + return nil, errors.New("resource does not support direct Data Plane access") + } + return response.DataPlaneInfo.QueryInfo, nil + } + refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) { + return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token) + } + getParams := []string{ + request.Name, + } + endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("Query", getParams, refresh, infoGetter) + if err != nil { + return nil, err + } + headers := make(map[string]string) + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + opts := []httpclient.DoOption{} + opts = append(opts, httpclient.WithRequestHeaders(headers)) + var queryEndpointResponse QueryEndpointResponse + opts = append(opts, httpclient.WithRequestData(request)) + opts = append(opts, httpclient.WithResponseUnmarshal(&queryEndpointResponse)) + opts = append(opts, httpclient.WithToken(dataPlaneToken)) + err = a.client.ApiClient().Do(ctx, http.MethodPost, endpointUrl, opts...) + return &queryEndpointResponse, err +} diff --git a/service/serving/interface.go b/service/serving/interface.go index 61350c543..5de0c60ee 100755 --- a/service/serving/interface.go +++ b/service/serving/interface.go @@ -167,3 +167,11 @@ type ServingEndpointsService interface { // inherit permissions from their root object. UpdatePermissions(ctx context.Context, request ServingEndpointPermissionsRequest) (*ServingEndpointPermissions, error) } + +// Serving endpoints DataPlane provides a set of operations to interact with +// data plane endpoints for Serving endpoints service. +type ServingEndpointsDataPlaneService interface { + + // Query a serving endpoint. + Query(ctx context.Context, request QueryEndpointInput) (*QueryEndpointResponse, error) +} diff --git a/service/vectorsearch/model.go b/service/vectorsearch/model.go index 063c94254..b49f15634 100755 --- a/service/vectorsearch/model.go +++ b/service/vectorsearch/model.go @@ -559,6 +559,8 @@ type QueryVectorIndexRequest struct { NumResults int `json:"num_results,omitempty"` // Query text. Required for Delta Sync Index using model endpoint. QueryText string `json:"query_text,omitempty"` + // The query type to use. Choices are `ANN` and `HYBRID`. Defaults to `ANN`. + QueryType string `json:"query_type,omitempty"` // Query vector. Required for Direct Vector Access Index and Delta Sync // Index using self-managed vectors. QueryVector []float64 `json:"query_vector,omitempty"` diff --git a/workspace_client.go b/workspace_client.go index 282991492..10c476aa5 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -215,9 +215,6 @@ type WorkspaceClient struct { // permissions (superset of `CAN_RUN`) DbsqlPermissions sql.DbsqlPermissionsInterface - // Endpoints are used to connect to PG clusters. - Endpoints catalog.EndpointsInterface - // Experiments are the primary unit of organization in MLflow; all MLflow // runs belong to an experiment. Each experiment lets you visualize, search, // and compare runs, as well as download run artifacts or metadata for @@ -715,6 +712,10 @@ type WorkspaceClient struct { // applied to each served entity. ServingEndpoints serving.ServingEndpointsInterface + // Serving endpoints DataPlane provides a set of operations to interact with + // data plane endpoints for Serving endpoints service. + ServingEndpointsDataPlane serving.ServingEndpointsDataPlaneInterface + // Workspace Settings API allows users to manage settings at the workspace // level. Settings settings.SettingsInterface @@ -1029,7 +1030,6 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { DataSources: sql.NewDataSources(databricksClient), Dbfs: files.NewDbfs(databricksClient), DbsqlPermissions: sql.NewDbsqlPermissions(databricksClient), - Endpoints: catalog.NewEndpoints(databricksClient), Experiments: ml.NewExperiments(databricksClient), ExternalLocations: catalog.NewExternalLocations(databricksClient), Files: files.NewFiles(databricksClient), @@ -1072,6 +1072,7 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) { Secrets: workspace.NewSecrets(databricksClient), ServicePrincipals: iam.NewServicePrincipals(databricksClient), ServingEndpoints: servingEndpoints, + ServingEndpointsDataPlane: serving.NewServingEndpointsDataPlane(databricksClient, servingEndpoints), Settings: settings.NewSettings(databricksClient), Shares: sharing.NewShares(databricksClient), StatementExecution: sql.NewStatementExecution(databricksClient),