python/oauth2client/django_orm.py

   1 # Copyright (C) 2010 Google Inc.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #      http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 """OAuth 2.0 utilities for Django.
  16
  17 Utilities for using OAuth 2.0 in conjunction with
  18 the Django datastore.
  19 """
  20
  21 __author__ = 'jcgregorio@google.com (Joe Gregorio)'
  22
  23 import oauth2client
  24 import base64
  25 import pickle
  26
  27 from django.db import models
  28 from oauth2client.client import Storage as BaseStorage
  29
  30 class CredentialsField(models.Field):
  31
  32   __metaclass__ = models.SubfieldBase
  33
  34   def __init__(self, *args, **kwargs):
  35     if 'null' not in kwargs:
  36       kwargs['null'] = True
  37     super(CredentialsField, self).__init__(*args, **kwargs)
  38
  39   def get_internal_type(self):
  40     return "TextField"
  41
  42   def to_python(self, value):
  43     if value is None:
  44       return None
  45     if isinstance(value, oauth2client.client.Credentials):
  46       return value
  47     return pickle.loads(base64.b64decode(value))
  48
  49   def get_db_prep_value(self, value, connection, prepared=False):
  50     if value is None:
  51       return None
  52     return base64.b64encode(pickle.dumps(value))
  53
  54
  55 class FlowField(models.Field):
  56
  57   __metaclass__ = models.SubfieldBase
  58
  59   def __init__(self, *args, **kwargs):
  60     if 'null' not in kwargs:
  61       kwargs['null'] = True
  62     super(FlowField, self).__init__(*args, **kwargs)
  63
  64   def get_internal_type(self):
  65     return "TextField"
  66
  67   def to_python(self, value):
  68     if value is None:
  69       return None
  70     if isinstance(value, oauth2client.client.Flow):
  71       return value
  72     return pickle.loads(base64.b64decode(value))
  73
  74   def get_db_prep_value(self, value, connection, prepared=False):
  75     if value is None:
  76       return None
  77     return base64.b64encode(pickle.dumps(value))
  78
  79
  80 class Storage(BaseStorage):
  81   """Store and retrieve a single credential to and from
  82   the datastore.
  83
  84   This Storage helper presumes the Credentials
  85   have been stored as a CredenialsField
  86   on a db model class.
  87   """
  88
  89   def __init__(self, model_class, key_name, key_value, property_name):
  90     """Constructor for Storage.
  91
  92     Args:
  93       model: db.Model, model class
  94       key_name: string, key name for the entity that has the credentials
  95       key_value: string, key value for the entity that has the credentials
  96       property_name: string, name of the property that is an CredentialsProperty
  97     """
  98     self.model_class = model_class
  99     self.key_name = key_name
 100     self.key_value = key_value
 101     self.property_name = property_name
 102
 103   def locked_get(self):
 104     """Retrieve Credential from datastore.
 105
 106     Returns:
 107       oauth2client.Credentials
 108     """
 109     credential = None
 110
 111     query = {self.key_name: self.key_value}
 112     entities = self.model_class.objects.filter(**query)
 113     if len(entities) > 0:
 114       credential = getattr(entities[0], self.property_name)
 115       if credential and hasattr(credential, 'set_store'):
 116         credential.set_store(self)
 117     return credential
 118
 119   def locked_put(self, credentials):
 120     """Write a Credentials to the datastore.
 121
 122     Args:
 123       credentials: Credentials, the credentials to store.
 124     """
 125     args = {self.key_name: self.key_value}
 126     entity = self.model_class(**args)
 127     setattr(entity, self.property_name, credentials)
 128     entity.save()
 129
 130   def locked_delete(self):
 131     """Delete Credentials from the datastore."""
 132
 133     query = {self.key_name: self.key_value}
 134     entities = self.model_class.objects.filter(**query).delete()