app/user.go

   1 package app
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "net/http"
   7         "sync"
   8
   9         "github.com/google/uuid"
  10         legacy_context "golang.org/x/net/context"
  11         "golang.org/x/oauth2"
  12         "google.golang.org/appengine/datastore"
  13         "google.golang.org/appengine/log"
  14 )
  15
  16 type User struct {
  17         key *datastore.Key
  18
  19         ID    string
  20         Email string
  21 }
  22
  23 type dbUser struct {
  24         ID string
  25 }
  26
  27 func NewUser(ctx context.Context, email string) (*User, error) {
  28         var id string
  29         err := datastore.RunInTransaction(ctx, func(ctx legacy_context.Context) error {
  30                 key := datastore.NewKey(ctx, "User", email, 0, nil)
  31
  32                 var u dbUser
  33                 err := datastore.Get(ctx, key, &u)
  34                 if err != nil && err != datastore.ErrNoSuchEntity {
  35                         return err
  36                 }
  37                 if err == nil {
  38                         id = u.ID
  39                         return nil
  40                 }
  41
  42                 id = uuid.New().String()
  43                 _, err = datastore.Put(ctx, key, &dbUser{
  44                         ID: id,
  45                 })
  46                 return err
  47         }, nil)
  48         if err != nil {
  49                 return nil, err
  50         }
  51
  52         return &User{
  53                 key:   datastore.NewKey(ctx, "User", email, 0, nil),
  54                 ID:    id,
  55                 Email: email,
  56         }, nil
  57 }
  58
  59 func UserByID(ctx context.Context, id string) (*User, error) {
  60         q := datastore.NewQuery("User").Filter("ID=", id).KeysOnly()
  61         keys, err := q.GetAll(ctx, nil)
  62         if err != nil {
  63                 return nil, fmt.Errorf("datastore.Query.GetAll(): %v", err)
  64         }
  65         if len(keys) != 1 {
  66                 return nil, fmt.Errorf("len(keys) = %d, want 1", len(keys))
  67         }
  68
  69         return &User{
  70                 key:   keys[0],
  71                 ID:    id,
  72                 Email: keys[0].StringID(),
  73         }, nil
  74 }
  75
  76 func (u *User) Token(ctx context.Context, svc string) (*oauth2.Token, error) {
  77         key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
  78
  79         var tok oauth2.Token
  80         if err := datastore.Get(ctx, key, &tok); err != nil {
  81                 return nil, err
  82         }
  83
  84         return &tok, nil
  85 }
  86
  87 func (u *User) SetToken(ctx context.Context, svc string, tok *oauth2.Token) error {
  88         key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
  89         _, err := datastore.Put(ctx, key, tok)
  90         return err
  91 }
  92
  93 func (u *User) DeleteToken(ctx context.Context, svc string) error {
  94         key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
  95         return datastore.Delete(ctx, key)
  96 }
  97
  98 func (u *User) OAuthClient(ctx context.Context, svc string, cfg *oauth2.Config) (*http.Client, error) {
  99         key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
 100
 101         var tok oauth2.Token
 102         if err := datastore.Get(ctx, key, &tok); err != nil {
 103                 return nil, err
 104         }
 105
 106         src := cfg.TokenSource(ctx, &tok)
 107         return oauth2.NewClient(ctx, &persistingTokenSource{
 108                 ctx: ctx,
 109                 t:   &tok,
 110                 src: src,
 111                 key: key,
 112         }), nil
 113 }
 114
 115 func (u *User) String() string {
 116         return u.Email
 117 }
 118
 119 type persistingTokenSource struct {
 120         ctx context.Context
 121         t   *oauth2.Token
 122         src oauth2.TokenSource
 123         key *datastore.Key
 124
 125         sync.Mutex
 126 }
 127
 128 func (s *persistingTokenSource) Token() (*oauth2.Token, error) {
 129         s.Lock()
 130         defer s.Unlock()
 131
 132         tok, err := s.src.Token()
 133         if err != nil {
 134                 return nil, err
 135         }
 136
 137         if s.t.AccessToken != tok.AccessToken ||
 138                 s.t.TokenType != tok.TokenType ||
 139                 s.t.RefreshToken != tok.RefreshToken ||
 140                 !s.t.Expiry.Equal(tok.Expiry) {
 141                 if _, err := datastore.Put(s.ctx, s.key, tok); err != nil {
 142                         log.Errorf(s.ctx, "persisting OAuth token in datastore failed: %v", err)
 143                 }
 144         }
 145
 146         s.t = tok
 147         return tok, nil
 148 }