make version of ts_mon compatible with appengine

appengine/chromium_build_logs/third_party/oauth2client/django_orm.py

-# Copyright (C) 2010 Google Inc.
+# Copyright 2014 Google Inc. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ skipping 12 matching lines @@
 import base64
 import pickle
 
 from django.db import models
 from oauth2client.client import Storage as BaseStorage
 
 class CredentialsField(models.Field):
 
   __metaclass__ = models.SubfieldBase
 
+  def __init__(self, *args, **kwargs):
+    if 'null' not in kwargs:
+      kwargs['null'] = True
+    super(CredentialsField, self).__init__(*args, **kwargs)
+
   def get_internal_type(self):
     return "TextField"
 
   def to_python(self, value):
-    if not value:
+    if value is None:
       return None
     if isinstance(value, oauth2client.client.Credentials):
       return value
     return pickle.loads(base64.b64decode(value))
 
   def get_db_prep_value(self, value, connection, prepared=False):
+    if value is None:
+      return None
     return base64.b64encode(pickle.dumps(value))
 
 
 class FlowField(models.Field):
 
   __metaclass__ = models.SubfieldBase
 
+  def __init__(self, *args, **kwargs):
+    if 'null' not in kwargs:
+      kwargs['null'] = True
+    super(FlowField, self).__init__(*args, **kwargs)
+
   def get_internal_type(self):
     return "TextField"
 
   def to_python(self, value):
     if value is None:
       return None
     if isinstance(value, oauth2client.client.Flow):
       return value
     return pickle.loads(base64.b64decode(value))
 
   def get_db_prep_value(self, value, connection, prepared=False):
+    if value is None:
+      return None
     return base64.b64encode(pickle.dumps(value))
 
 
 class Storage(BaseStorage):
   """Store and retrieve a single credential to and from
   the datastore.
 
   This Storage helper presumes the Credentials
   have been stored as a CredenialsField
   on a db model class.
@@ skipping 22 matching lines @@
     credential = None
 
     query = {self.key_name: self.key_value}
     entities = self.model_class.objects.filter(**query)
     if len(entities) > 0:
       credential = getattr(entities[0], self.property_name)
       if credential and hasattr(credential, 'set_store'):
         credential.set_store(self)
     return credential
 
-  def locked_put(self, credentials):
+  def locked_put(self, credentials, overwrite=False):
     """Write a Credentials to the datastore.
 
     Args:
       credentials: Credentials, the credentials to store.
+      overwrite: Boolean, indicates whether you would like these credentials to
+                          overwrite any existing stored credentials.
     """
     args = {self.key_name: self.key_value}
-    entity = self.model_class(**args)
+
+    if overwrite:
+      entity, unused_is_new = self.model_class.objects.get_or_create(**args)
+    else:
+      entity = self.model_class(**args)
+
     setattr(entity, self.property_name, credentials)
     entity.save()
+
+  def locked_delete(self):
+    """Delete Credentials from the datastore."""
+
+    query = {self.key_name: self.key_value}
+    entities = self.model_class.objects.filter(**query).delete()