Skip to content

Commit 1757750

Browse files
authored
Added unmarshal attribute for expires_in for device flow auth token (#4319)
Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>
1 parent d0ed6c4 commit 1757750

3 files changed

Lines changed: 8 additions & 3 deletions

File tree

flyteidl/clients/go/admin/deviceflow/payload.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,6 @@ type DeviceAccessTokenRequest struct {
4040

4141
type DeviceAccessTokenResponse struct {
4242
oauth2.Token
43-
Error string `json:"error"`
43+
Error string `json:"error"`
44+
ExpiresIn int64 `json:"expires_in"` // relative seconds from now
4445
}

flyteidl/clients/go/admin/deviceflow/token_orchestrator.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,12 +125,14 @@ func (t TokenOrchestrator) PollTokenEndpoint(ctx context.Context, tokReq DeviceA
125125
// Unmarshalled response if it contains an error then check if we need to increase the polling interval
126126
if len(tokResp.Error) > 0 {
127127
if tokResp.Error == errSlowDown || tokResp.Error == errAuthPending {
128-
pollInterval = pollInterval * 2
129-
128+
logger.Debugf(ctx, "going to poll again due to error %v", tokResp.Error)
130129
} else {
131130
return nil, fmt.Errorf("oauth error : %v", tokResp.Error)
132131
}
133132
} else {
133+
if secs := tokResp.ExpiresIn; secs > 0 {
134+
tokResp.Token.Expiry = time.Now().Add(time.Duration(secs) * time.Second)
135+
}
134136
// Got the auth token in the response and save it in the cache
135137
err = t.TokenCache.SaveToken(&tokResp.Token)
136138
// Saving into the cache is only considered to be a warning in this case.

flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ func TestFetchFromAuthFlow(t *testing.T) {
8585
Token: oauth2.Token{
8686
AccessToken: "access_token",
8787
},
88+
ExpiresIn: 300,
8889
}
8990
darBytes, err := json.Marshal(dar)
9091
assert.Nil(t, err)
@@ -119,5 +120,6 @@ func TestFetchFromAuthFlow(t *testing.T) {
119120
assert.Nil(t, err)
120121
assert.NotNil(t, authToken)
121122
assert.Equal(t, "access_token", authToken.AccessToken)
123+
assert.True(t, authToken.Expiry.After(time.Now().Add(time.Second*200)))
122124
})
123125
}

0 commit comments

Comments
 (0)