Package app: Wrap oauth2.TokenSource to ensure datastore is always updated.

diff --git a/app/user.go b/app/user.go
index 8465c2c..e47a5cc 100644 (file)
--- a/app/user.go
+++ b/app/user.go
@@ -3,11 +3,13 @@ package app
 import (
        "context"
        "fmt"
+       "net/http"
 
        "github.com/google/uuid"
        legacy_context "golang.org/x/net/context"
        "golang.org/x/oauth2"
        "google.golang.org/appengine/datastore"
+       "google.golang.org/appengine/log"
 )
 
 type User struct {
@@ -78,7 +80,47 @@ func (u *User) Token(ctx context.Context, svc string) (*oauth2.Token, error) {
 }
 
 func (u *User) SetToken(ctx context.Context, svc string, tok *oauth2.Token) error {
-       key := datastore.NewKey(ctx, "Token", "Fitbit", 0, u.key)
+       key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
        _, err := datastore.Put(ctx, key, tok)
        return err
 }
+
+func (u *User) OAuthClient(ctx context.Context, svc string, cfg *oauth2.Config) (*http.Client, error) {
+       key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
+
+       var tok oauth2.Token
+       if err := datastore.Get(ctx, key, &tok); err != nil {
+               return nil, err
+       }
+
+       src := cfg.TokenSource(ctx, &tok)
+       return oauth2.NewClient(ctx, &persistingTokenSource{
+               ctx: ctx,
+               t:   &tok,
+               src: src,
+               key: key,
+       }), nil
+}
+
+type persistingTokenSource struct {
+       ctx context.Context
+       t   *oauth2.Token
+       src oauth2.TokenSource
+       key *datastore.Key
+}
+
+func (s *persistingTokenSource) Token() (*oauth2.Token, error) {
+       tok, err := s.src.Token()
+       if err != nil {
+               return nil, err
+       }
+
+       if s.t.RefreshToken != tok.RefreshToken {
+               if _, err := datastore.Put(s.ctx, s.key, tok); err != nil {
+                       log.Errorf(s.ctx, "persisting OAuth token in datastore failed: %v", err)
+               }
+       }
+
+       s.t = tok
+       return tok, nil
+}

diff --git a/fitbit/fitbit.go b/fitbit/fitbit.go
index bdfc3dd..81d4ebc 100644 (file)
--- a/fitbit/fitbit.go
+++ b/fitbit/fitbit.go
@@ -127,36 +127,15 @@ func NewClient(ctx context.Context, fitbitUserID string, u *app.User) (*Client,
                fitbitUserID = "-"
        }
 
-       storedToken, err := u.Token(ctx, "Fitbit")
+       c, err := u.OAuthClient(ctx, "Fitbit", oauth2Config)
        if err != nil {
                return nil, err
        }
 
-       // The oauth2 package will refresh the token when it is valid for less
-       // than 10 seconds. To avoid a race with later calls (which would
-       // refresh the token but the new RefreshToken wouldn't make it back
-       // into datastore), we refresh earlier than that. The Fitbit tokens are
-       // quite long-lived (six hours?); the additional load this puts on the
-       // backends is negligible.
-       if storedToken.Expiry.Round(0).Add(-5 * time.Minute).Before(time.Now()) {
-               storedToken.Expiry = time.Now()
-       }
-
-       refreshedToken, err := oauth2Config.TokenSource(ctx, storedToken).Token()
-       if err != nil {
-               return nil, err
-       }
-
-       if refreshedToken.RefreshToken != storedToken.RefreshToken {
-               if err := u.SetToken(ctx, "Fitbit", refreshedToken); err != nil {
-                       return nil, err
-               }
-       }
-
        return &Client{
                fitbitUserID: fitbitUserID,
                appUser:      u,
-               client:       oauth2Config.Client(ctx, refreshedToken),
+               client:       c,
        }, nil
 }