Generate dataplane code

.codegen/_openapi_sha

@@ -1 +1 @@
-7eb5ad9a2ed3e3f1055968a2d1014ac92c06fe92
+universe:/Users/hector.castejon/universe

.codegen/accounts.go.tmpl

@@ -42,9 +42,20 @@ func NewAccountClient(c ...*Config) (*AccountClient, error) {
 	if err != nil {
 		return nil, err
 	}
+	{{range .Services}}{{if and .IsAccounts (not .HasParent) (.HasDataPlaneMethods) (not .IsDataPlane)}}
+	{{.CamelName}} := {{.Package.Name}}.New{{.Name}}(apiClient)
+	{{- end}}{{end}}
 	return &AccountClient{
 		Config: cfg,
-		{{range .Services}}{{if and .IsAccounts (not .HasParent)}}
-		{{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient),{{end}}{{end}}
+		{{range .Services}}{{if and .IsAccounts (not .HasParent) (not .IsDataPlane) (not .HasDataPlaneMethods)}}
+		{{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient),
+		{{- end -}}
+		{{if and .IsAccounts (not .HasParent) (not .IsDataPlane) (.HasDataPlaneMethods)}}
+		{{(.TrimPrefix "account").PascalName}}: {{.CamelName}},
+		{{- end -}}
+		{{if and .IsAccounts (not .HasParent) .IsDataPlane}}
+		{{(.TrimPrefix "account").PascalName}}: {{.Package.Name}}.New{{.Name}}(apiClient,  {{.ControlPlaneService.CamelName}}),
+		{{- end -}}
+		{{end}}
 	}, nil
 }

.codegen/api.go.tmpl

@@ -105,10 +105,17 @@ type {{.PascalName}}Interface interface {
 {{end -}}
 }
 
-func New{{.PascalName}}(client *client.DatabricksClient) *{{.PascalName}}API {
+func New{{.PascalName}}(client *client.DatabricksClient,
+{{- if .IsDataPlane}}
+controlPlane *{{.ControlPlaneService.PascalName}}API,
+{{end -}}
+) *{{.PascalName}}API {
 	return &{{.PascalName}}API{
 		impl: &{{.CamelName}}Impl{
 			client: client,
+			{{- if .IsDataPlane}}
+			controlPlane: controlPlane,
+			{{end}}
 		},
 		{{range .Subservices}}
 		{{.CamelName}}: New{{.PascalName}}(client),
@@ -148,6 +155,8 @@ func (a *{{.PascalName}}API) Impl() {{.PascalName}}Service {
 	return a.impl
 }
 
+
+
 {{range .Waits}}
 // {{.PascalName}} repeatedly calls [{{.Method.Service.Name}}API.{{.Poll.PascalName}}] and waits to reach {{range $i, $e := .Success}}{{if $i}} or {{end}}{{.Content}}{{end}} state
 func (a *{{.Method.Service.PascalName}}API) {{.PascalName}}(ctx context.Context{{range .Binding}}, {{.PollField.CamelName}} {{template "type" .PollField.Entity}}{{end}},
@@ -427,4 +436,4 @@ func (a *{{.Service.Name}}API) {{.Shortcut.PascalName}}AndWait(ctx context.Conte
 
 {{- define "wait-response-type" -}}
 (*{{.Wait.PascalName}}[{{with .Response}}{{if .IsEmpty}}struct{}{{else}}{{.PascalName}}{{end}}{{end}}], error)
-{{- end -}}
+{{- end -}}

.codegen/impl.go.tmpl

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

.codegen/interface.go.tmpl

@@ -21,4 +21,6 @@ type {{.PascalName}}Service interface {
 	{{end}}
 }
 
+
+
 {{end}}

.codegen/workspaces.go.tmpl

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

.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

client/client.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/httpclient"
+	"golang.org/x/oauth2"
 )
 
 func New(cfg *config.Config) (*DatabricksClient, error) {
@@ -46,6 +47,20 @@ func (c *DatabricksClient) ConfiguredAccountID() string {
 	return c.Config.AccountID
 }
 
+// Returns the inner Api Client.
+func (c *DatabricksClient) ApiClient() *httpclient.ApiClient {
+	return c.client
+}
+
+// Returns a new OAuth token using the provided token. The token must be a JWT token.
+// The resulting token is scoped to the authorization details provided.
+//
+// **NOTE:** Experimental: This API may change or be removed in a future release
+// without warning.
+func (c *DatabricksClient) GetOAuthToken(ctx context.Context, authDetails string, token *oauth2.Token) (*oauth2.Token, error) {
+	return c.client.GetOAuthToken(ctx, authDetails, token)
+}
+
 // Do sends an HTTP request against path.
 func (c *DatabricksClient) Do(ctx context.Context, method, path string,
 	headers map[string]string, request, response any,