Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
@parthban-db
parthban-db committed Jan 6, 2025
1 parent 4399da5 commit 5c3d7a4
Show file tree
Hide file tree
Showing 6 changed files with 40 additions and 45 deletions.
2 changes: 1 addition & 1 deletion .codegen/_openapi_sha
View file
Original file line number Diff line number Diff line change
@@ -1 +1 @@
821ab6511c9a977bae10cdcd2cf5aa0ad6ab8f8f
-
17 changes: 0 additions & 17 deletions service/oauth2/model.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

13 changes: 6 additions & 7 deletions service/serving/ext_data_plane.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,23 @@ import (
"strings"
"sync"

"github.com/databricks/databricks-sdk-go/service/oauth2"
goauth "golang.org/x/oauth2"
)

// DataPlaneService is an interface for services that access DataPlane.
type DataPlaneService interface {
GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error)
GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *goauth.Token, error)
}

func NewDataPlaneService() DataPlaneService {
return &dataPlaneServiceImpl{
infos: make(map[string]*oauth2.DataPlaneInfo),
infos: make(map[string]*DataPlaneInfo),
tokens: make(map[string]*goauth.Token),
}
}

type dataPlaneServiceImpl struct {
infos map[string]*oauth2.DataPlaneInfo
infos map[string]*DataPlaneInfo
tokens map[string]*goauth.Token
// This class can be shared across multiple threads.
// This mutex is used to synchronize access to the infos and tokens maps.
Expand All @@ -30,7 +29,7 @@ type dataPlaneServiceImpl struct {

// GetDataPlaneDetails returns the endpoint URL and token. It returns a cached token if it is valid,
// otherwise it refreshes the token and returns the new token.
func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*oauth2.DataPlaneInfo, error)) (string, *goauth.Token, error) {
func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*goauth.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *goauth.Token, error) {
key := o.generateKey(method, params)
info, err := o.getInfo(key, infoGetter)
if err != nil {
Expand All @@ -43,7 +42,7 @@ func (o *dataPlaneServiceImpl) GetDataPlaneDetails(method string, params []strin
return info.EndpointUrl, token, nil
}

func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*oauth2.DataPlaneInfo, error)) (*oauth2.DataPlaneInfo, error) {
func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*DataPlaneInfo, error)) (*DataPlaneInfo, error) {
info, infoOk := o.infos[key]
if !infoOk {
o.mu.Lock()
Expand All @@ -61,7 +60,7 @@ func (o *dataPlaneServiceImpl) getInfo(key string, infoGetter func() (*oauth2.Da
return info, nil
}

func (o *dataPlaneServiceImpl) refreshToken(key string, info *oauth2.DataPlaneInfo, refresh func(*oauth2.DataPlaneInfo) (*goauth.Token, error)) (*goauth.Token, error) {
func (o *dataPlaneServiceImpl) refreshToken(key string, info *DataPlaneInfo, refresh func(*DataPlaneInfo) (*goauth.Token, error)) (*goauth.Token, error) {
token, tokenOk := o.tokens[key]
if !tokenOk || !token.Valid() {
o.mu.Lock()
Expand Down
27 changes: 13 additions & 14 deletions service/serving/ext_data_plane_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,38 +4,37 @@ import (
"testing"
"time"

"github.com/databricks/databricks-sdk-go/service/oauth2"
"github.com/stretchr/testify/assert"
goauth "golang.org/x/oauth2"
)

type infoMock struct {
called bool
info *oauth2.DataPlaneInfo
info *DataPlaneInfo
err error
}

func (i *infoMock) DataPlaneInfoGetter() (*oauth2.DataPlaneInfo, error) {
func (i *infoMock) DataPlaneInfoGetter() (*DataPlaneInfo, error) {
i.called = true
return i.info, i.err
}

type tokenRefreshSpy struct {
called bool
expectedInfo *oauth2.DataPlaneInfo
expectedInfo *DataPlaneInfo
token *goauth.Token
err error
}

func (t *tokenRefreshSpy) TokenRefresh(info *oauth2.DataPlaneInfo) (*goauth.Token, error) {
func (t *tokenRefreshSpy) TokenRefresh(info *DataPlaneInfo) (*goauth.Token, error) {
t.expectedInfo = info
t.called = true
return t.token, t.err
}

func TestTokenNotCached(t *testing.T) {
info := infoMock{
info: &oauth2.DataPlaneInfo{
info: &DataPlaneInfo{
EndpointUrl: "url",
AuthorizationDetails: "authDetails",
},
Expand All @@ -50,7 +49,7 @@ func TestTokenNotCached(t *testing.T) {
err: nil,
}
c := dataPlaneServiceImpl{
infos: make(map[string]*oauth2.DataPlaneInfo),
infos: make(map[string]*DataPlaneInfo),
tokens: make(map[string]*goauth.Token),
}
url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter)
Expand All @@ -65,7 +64,7 @@ func TestTokenNotCached(t *testing.T) {

func TestTokenCached(t *testing.T) {
info := infoMock{
info: &oauth2.DataPlaneInfo{
info: &DataPlaneInfo{
EndpointUrl: "url",
AuthorizationDetails: "authDetails",
},
Expand All @@ -80,7 +79,7 @@ func TestTokenCached(t *testing.T) {
err: nil,
}
c := dataPlaneServiceImpl{}
c.infos = make(map[string]*oauth2.DataPlaneInfo)
c.infos = make(map[string]*DataPlaneInfo)
c.tokens = make(map[string]*goauth.Token)
c.infos["method/params"] = info.info
c.tokens["method/params"] = s.token
Expand All @@ -96,7 +95,7 @@ func TestTokenCached(t *testing.T) {

func TestTokenExpired(t *testing.T) {
info := infoMock{
info: &oauth2.DataPlaneInfo{
info: &DataPlaneInfo{
EndpointUrl: "url",
AuthorizationDetails: "authDetails",
},
Expand All @@ -117,7 +116,7 @@ func TestTokenExpired(t *testing.T) {
err: nil,
}
c := dataPlaneServiceImpl{}
c.infos = make(map[string]*oauth2.DataPlaneInfo)
c.infos = make(map[string]*DataPlaneInfo)
c.tokens = make(map[string]*goauth.Token)
c.infos["method/params"] = info.info
c.tokens["method/params"] = expired
Expand All @@ -138,7 +137,7 @@ func TestTokenInfoError(t *testing.T) {
}
s := tokenRefreshSpy{}
c := dataPlaneServiceImpl{
infos: make(map[string]*oauth2.DataPlaneInfo),
infos: make(map[string]*DataPlaneInfo),
tokens: make(map[string]*goauth.Token),
}
url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter)
Expand All @@ -151,7 +150,7 @@ func TestTokenInfoError(t *testing.T) {

func TestTokenRefreshError(t *testing.T) {
info := infoMock{
info: &oauth2.DataPlaneInfo{
info: &DataPlaneInfo{
EndpointUrl: "url",
AuthorizationDetails: "authDetails",
},
Expand All @@ -162,7 +161,7 @@ func TestTokenRefreshError(t *testing.T) {
err: assert.AnError,
}
c := dataPlaneServiceImpl{
infos: make(map[string]*oauth2.DataPlaneInfo),
infos: make(map[string]*DataPlaneInfo),
tokens: make(map[string]*goauth.Token),
}
url, token, err := c.GetDataPlaneDetails("method", []string{"params"}, s.TokenRefresh, info.DataPlaneInfoGetter)
Expand Down
6 changes: 2 additions & 4 deletions service/serving/impl.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

20 changes: 18 additions & 2 deletions service/serving/model.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 5c3d7a4

Please sign in to comment.