Add DataPlane API Support

.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,21 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) {
 	if err != nil {
 		return nil, err
 	}
+	{{range .Services}}{{if and .IsAccounts (not .HasParent) .HasDataPlaneAPI}}
+	{{.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)}} 
+				{{- if .HasDataPlaneAPI}}
+					{{(.TrimPrefix "account").PascalName}}: {{.CamelName}},
+				{{- else if .IsDataPlane}}
+					{{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient,  {{.ControlPlaneService.CamelName}}),
+				{{- else }}
+					{{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient),
+				{{- end -}}
+			{{- end -}}
+		{{end}}
 	}, nil
 }

.codegen/api.go.tmpl

@@ -1,6 +1,6 @@
 // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
 
-// {{if .Description}}{{.Description}}{{else}}These APIs allow you to manage {{range .Services}}{{if (not .IsDataPlane)}}{{.TitleName}}, {{end}}{{end}}etc.{{end}}
+// {{if .Description}}{{.Description}}{{else}}These APIs allow you to manage {{range .Services}}{{.TitleName}}, {{end}}etc.{{end}}
 package {{.Name}}
 
 import (
@@ -22,7 +22,6 @@ import (
 		"Workspace"
 }}
 {{range .Services}}
-{{if not .IsDataPlane}}
 
 type {{.PascalName}}Interface interface {
 	{{if (in $extendedInterface .PascalName) }}
@@ -106,10 +105,18 @@ type {{.PascalName}}Interface interface {
 {{end -}}
 }
 
-func New{{.PascalName}}(client *client.DatabricksClient) *{{.PascalName}}API {
+func New{{.PascalName}}(client *client.DatabricksClient,
+{{- if .IsDataPlane}}
+controlPlane *{{.ControlPlaneService.PascalName}}API,
+{{end -}}
+) *{{.PascalName}}API {
 	return &{{.PascalName}}API{
 		impl: &{{.CamelName}}Impl{
 			client: client,
+			{{- if .IsDataPlane}}
+			DataPlaneService: oauth2.NewDataPlaneService(),
+			controlPlane: controlPlane,
+			{{end}}
 		},
 		{{range .Subservices}}
 		{{.CamelName}}: New{{.PascalName}}(client),
@@ -149,6 +156,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}},
@@ -412,7 +421,6 @@ func (a *{{.Service.Name}}API) {{.Shortcut.PascalName}}AndWait(ctx context.Conte
 }
 {{end}}
 {{end -}}
-{{end}}
 {{- end}}
 {{- define "field-path" -}}
 	{{- range .}}.{{.PascalName}}{{end}}
@@ -428,4 +436,4 @@ func (a *{{.Service.Name}}API) {{.Shortcut.PascalName}}AndWait(ctx context.Conte
 
 {{- define "wait-response-type" -}}
 (*{{.Wait.PascalName}}[{{with .Response}}{{if .IsEmpty}}struct{}{{else}}{{.PascalName}}{{end}}{{end}}], error)
-{{- end -}}
+{{- end -}}

.codegen/impl.go.tmpl

@@ -4,7 +4,9 @@ package {{.Name}}
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	goauth "golang.org/x/oauth2"
 	"time"
 	"io"
 	"net/http"
@@ -13,22 +15,71 @@ import (
 	{{range .ImportedPackages}}
 	"github.com/databricks/databricks-sdk-go/service/{{.}}"{{end}}
 )
-{{range .Services}} {{if not .IsDataPlane}}
+{{range .Services}}
 // unexported type that holds implementations of just {{.Name}} API methods
 type {{.CamelName}}Impl struct {
+	{{- if .IsDataPlane}}
+	oauth2.DataPlaneService
+	controlPlane *{{.ControlPlaneService.PascalName}}API
+	{{end -}}
 	client *client.DatabricksClient
 }
 
 {{range .Methods}}
+{{if .Service.IsDataPlane}}
+func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} {
+	getRequest := {{.Service.DataPlaneInfoMethod.Request.PascalName}}{
+	{{- range .Service.DataPlaneInfoMethod.Request.Fields}}
+		{{.PascalName}}: request.{{.PascalName}},
+	{{end}}
+	}
+	token, err := a.client.Config.GetToken()
+	if err != nil {
+		return nil, err
+	}
+	infoGetter := func() (*oauth2.DataPlaneInfo, error) {
+		response, err := a.controlPlane.{{.Service.DataPlaneInfoMethod.PascalName}}(ctx, getRequest)
+		if err != nil {
+			return nil, err
+		}
+		if response.{{(index .DataPlaneInfoFields 0).PascalName}} == nil {
+			return nil, errors.New("resource does not support direct Data Plane access")
+		}
+		return response{{range .DataPlaneInfoFields}}.{{.PascalName}}{{end}}, nil 
+	}
+	refresh := func(info *oauth2.DataPlaneInfo) (*goauth.Token, error) {
+		return a.client.GetOAuthToken(ctx, info.AuthorizationDetails, token)
+	}
+	getParams := []string{ 
+	{{- range .Service.DataPlaneInfoMethod.Request.Fields}}
+		request.{{.PascalName}},
+	{{end -}}
+	}
+	endpointUrl, dataPlaneToken, err := a.GetDataPlaneDetails("{{.PascalName}}", getParams, refresh, infoGetter) 
+	if err != nil {
+		return nil, err
+	}
+	{{ template "make-header" . }}
+	opts := []httpclient.DoOption{}
+	opts = append(opts, httpclient.WithRequestHeaders(headers))
+	{{- template "response-var" . }}
+	{{if .Request}}opts = append(opts, httpclient.WithRequestData(request)){{end}}
+	{{if .Response}}opts = append(opts, httpclient.WithResponseUnmarshal(&{{.Response.CamelName}})){{end}}
+	opts = append(opts, httpclient.WithToken(dataPlaneToken))
+	err = a.client.ApiClient().Do(ctx, http.Method{{.TitleVerb}}, endpointUrl, opts...)
+	return {{ template "response" . }}
+}
+
+{{else}}
 func (a *{{.Service.CamelName}}Impl) {{.PascalName}}(ctx context.Context{{if .Request}}, request {{.Request.PascalName}}{{end}}) {{ template "response-type" . }} {
 	{{- template "response-var" . }}
 	path := {{ template "path" . }}
 	{{ 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 -}}
 {{end}}
 
 {{- define "path" -}}

.codegen/interface.go.tmpl

@@ -8,7 +8,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/retries"
 )
 
-{{range .Services}} {{if not .IsDataPlane}}
+{{range .Services}}
 {{.Comment "// " 80}}
 type {{.PascalName}}Service interface {
 	{{range .Methods}}
@@ -21,5 +21,6 @@ type {{.PascalName}}Service interface {
 	{{end}}
 }
 
+
+
 {{end}}
-{{end}}

.codegen/mock_account_client.go.tmpl

@@ -22,12 +22,12 @@ 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}}
 		},
 	}
-	{{range .ServicesSortedByParent}} {{if and .IsAccounts (not .IsDataPlane)}}
+	{{range .ServicesSortedByParent}} {{if and .IsAccounts }}
 	{{if  .HasSubservices}}
 	mock{{.PascalName}}API := cli.GetMock{{.PascalName}}API()
 	{{end}}
@@ -40,7 +40,7 @@ func NewMockAccountClient(t interface {
 	return cli
 }
 
-{{range .Services}}{{if and (.IsAccounts) (.HasParent) (not .IsDataPlane)}}
+{{range .Services}}{{if and .IsAccounts .HasParent }}
 func(m *MockAccountClient) GetMock{{.PascalName}}API() *{{ template "mock-interface-name" . }} {
 	api, ok := m.GetMock{{.ParentService.PascalName}}API().{{.PascalName}}().(*{{ template "mock-interface-name" . }})
 	if !ok {
@@ -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 {

.codegen/mock_workspace_client.go.tmpl

@@ -22,12 +22,12 @@ 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}}
 		},
 	}
-	{{range .ServicesSortedByParent}} {{if and (not .IsAccounts) (not .IsDataPlane)}}
+	{{range .ServicesSortedByParent}} {{if and (not .IsAccounts) }}
 	{{if  .HasSubservices}}
 	mock{{.CamelName}}API := cli.GetMock{{.PascalName}}API()
 	{{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 {

.codegen/service-pkg.go.tmpl

@@ -1,6 +1,6 @@
 // Databricks SDK for Go APIs
 // 
-//{{range .Services}}{{if and .Description (not .IsDataPlane)}} - [{{.Package.Name}}.{{.PascalName}}API]: {{.Summary}}
+//{{range .Services}}{{if and .Description }} - [{{.Package.Name}}.{{.PascalName}}API]: {{.Summary}}
 //
 //{{end}}{{end}}
 package service
@@ -15,6 +15,6 @@ import (
 // https://pkg.go.dev/github.com/databricks/databricks-sdk-go/service
 // See: https://pkg.go.dev/golang.org/x/tools/internal/imports#ImportPathToAssumedName
 var (
-{{range .Services}}{{if (not .IsDataPlane)}} _ *{{.Package.Name}}.{{.PascalName}}API = nil{{end}}
+{{range .Services}} _ *{{.Package.Name}}.{{.PascalName}}API = nil
 {{end}}
 )

.codegen/workspaces.go.tmpl

@@ -14,25 +14,12 @@ 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}}
 }
 
-// Returns a new OAuth scoped to the authorization details provided.
-// It will return an error if the CredentialStrategy does not support OAuth tokens.
-//
-// **NOTE:** Experimental: This API may change or be removed in a future release
-// without warning.
-func (a *WorkspaceClient) GetOAuthToken(ctx context.Context, authorizationDetails string) (*credentials.OAuthToken, error) {
-	originalToken, err := a.Config.GetToken()
-	if err != nil {
-		return nil, err
-	}
-	return a.apiClient.GetOAuthToken(ctx, authorizationDetails, originalToken)
-}
-
 var ErrNotWorkspaceClient = errors.New("invalid Databricks Workspace configuration")
 
 // NewWorkspaceClient creates new Databricks SDK client for Workspaces or
@@ -61,11 +48,22 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) {
 	if err != nil {
 		return nil, err
 	}
+	{{range .Services}}{{if and (not .IsAccounts) (not .HasParent) .HasDataPlaneAPI}}
+	{{.CamelName}} := {{.Package.Name}}.New{{.Name}}(databricksClient)
+	{{- end}}{{end}}
 	return &WorkspaceClient{
 		Config: cfg,
 		apiClient: apiClient,
-		{{range .Services}}{{if and (not .IsAccounts) (not .HasParent) (not .IsDataPlane)}}
-		{{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient),
-		{{- end}}{{end}}
+		{{range .Services}}
+			{{- if and (not .IsAccounts) (not .HasParent)}}
+				{{- if .HasDataPlaneAPI}}
+					{{.Name}}: {{.CamelName}},
+				{{- else if .IsDataPlane}}
+					{{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient,  {{.ControlPlaneService.CamelName}}),
+				{{- else }}
+					{{.Name}}: {{.Package.Name}}.New{{.Name}}(databricksClient),
+				{{- end -}}
+			{{- end -}}
+		{{end}}
 	}, nil
 }

.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_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

client/client.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/httpclient"
+	"golang.org/x/oauth2"
 )
 
 func New(cfg *config.Config) (*DatabricksClient, error) {
@@ -46,6 +47,20 @@ func (c *DatabricksClient) ConfiguredAccountID() string {
 	return c.Config.AccountID
 }
 
+// ApiClient returns the inner Api Client.
+func (c *DatabricksClient) ApiClient() *httpclient.ApiClient {
+	return c.client
+}
+
+// GetOAuthToken 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,