Modify oauth2client.multistore_file to store and retrieve credentials using an arbitrary key.

Reviewed in https://codereview.appspot.com/7067054/.

apiclient/http.py

@@ -1194,7 +1194,7 @@ class BatchHttpRequest(object):
 
     Args:
       http: httplib2.Http, an http object to be used in place of the one the
-        HttpRequest request object was constructed with.  If one isn't supplied
+        HttpRequest request object was constructed with. If one isn't supplied
         then use a http object from the requests in this batch.
 
     Returns:

oauth2client/client.py

@@ -234,7 +234,7 @@ class Flow(object):
 class Storage(object):
   """Base class for all Storage objects.
 
-  Store and retrieve a single credential.  This class supports locking
+  Store and retrieve a single credential. This class supports locking
   such that multiple processes and threads can operate on a single
   store.
   """
@@ -540,7 +540,7 @@ class OAuth2Credentials(Credentials):
     Args:
       store: Storage, an implementation of Stroage object.
         This is needed to store the latest access_token if it
-        has expired and been refreshed.  This implementation uses
+        has expired and been refreshed. This implementation uses
         locking to check for updates before updating the
         access_token.
     """
@@ -661,7 +661,7 @@ class AccessTokenCredentials(OAuth2Credentials):
 
   Credentials can be applied to an httplib2.Http object using the
   authorize() method, which then signs each request from that object
-  with the OAuth 2.0 access token.  This set of credentials is for the
+  with the OAuth 2.0 access token. This set of credentials is for the
   use case where you have acquired an OAuth 2.0 access_token from
   another place such as a JavaScript client or another web
   application, and wish to use it from Python. Because only the
@@ -723,7 +723,7 @@ class AssertionCredentials(OAuth2Credentials):
 
   This credential does not require a flow to instantiate because it
   represents a two legged flow, and therefore has all of the required
-  information to generate and refresh its own access tokens.  It must
+  information to generate and refresh its own access tokens. It must
   be subclassed to generate the appropriate assertion string.
 
   AssertionCredentials objects may be safely pickled and unpickled.

oauth2client/locked_file.py

@@ -64,7 +64,7 @@ class _Opener(object):
     return self._locked
 
   def file_handle(self):
-    """The file handle to the file.  Valid only after opened."""
+    """The file handle to the file. Valid only after opened."""
     return self._fh
 
   def filename(self):
@@ -198,7 +198,7 @@ try:
             raise e
           if e.errno != errno.EACCES:
             raise e
-          # We could not acquire the lock.  Try again.
+          # We could not acquire the lock. Try again.
           if (time.time() - start_time) >= timeout:
             logger.warn('Could not lock %s in %s seconds' % (
                 self._filename, timeout))
@@ -280,7 +280,7 @@ try:
           if e[0] != _Win32Opener.FILE_IN_USE_ERROR:
             raise
 
-          # We could not acquire the lock.  Try again.
+          # We could not acquire the lock. Try again.
           if (time.time() - start_time) >= timeout:
             logger.warn('Could not lock %s in %s seconds' % (
                 self._filename, timeout))

oauth2client/multistore_file.py

@@ -3,7 +3,7 @@
 """Multi-credential file store with lock support.
 
 This module implements a JSON credential store where multiple
-credentials can be stored in one file.  That file supports locking
+credentials can be stored in one file. That file supports locking
 both in a single process and across processes.
 
 The credential themselves are keyed off of:
@@ -72,6 +72,55 @@ def get_credential_storage(filename, client_id, user_agent, scope,
     scope: string or iterable of strings, Scope(s) being requested
     warn_on_readonly: if True, log a warning if the store is readonly
 
+  Returns:
+    An object derived from client.Storage for getting/setting the
+    credential.
+  """
+  # Recreate the legacy key with these specific parameters
+  key = {'clientId': client_id, 'userAgent': user_agent,
+         'scope': util.scopes_to_string(scope)}
+  return get_credential_storage_custom_key(
+      filename, key, warn_on_readonly=warn_on_readonly)
+
+
+@util.positional(2)
+def get_credential_storage_custom_string_key(
+    filename, key_string, warn_on_readonly=True):
+  """Get a Storage instance for a credential using a single string as a key.
+
+  Allows you to provide a string as a custom key that will be used for
+  credential storage and retrieval.
+
+  Args:
+    filename: The JSON file storing a set of credentials
+    key_string: A string to use as the key for storing this credential.
+    warn_on_readonly: if True, log a warning if the store is readonly
+
+  Returns:
+    An object derived from client.Storage for getting/setting the
+    credential.
+  """
+  # Create a key dictionary that can be used
+  key_dict = {'key': key_string}
+  return get_credential_storage_custom_key(
+      filename, key_dict, warn_on_readonly=warn_on_readonly)
+
+
+@util.positional(2)
+def get_credential_storage_custom_key(
+    filename, key_dict, warn_on_readonly=True):
+  """Get a Storage instance for a credential using a dictionary as a key.
+
+  Allows you to provide a dictionary as a custom key that will be used for
+  credential storage and retrieval.
+
+  Args:
+    filename: The JSON file storing a set of credentials
+    key_dict: A dictionary to use as the key for storing this credential. There
+      is no ordering of the keys in the dictionary. Logically equivalent
+      dictionaries will produce equivalent storage keys.
+    warn_on_readonly: if True, log a warning if the store is readonly
+
   Returns:
     An object derived from client.Storage for getting/setting the
     credential.
@@ -83,8 +132,8 @@ def get_credential_storage(filename, client_id, user_agent, scope,
         filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly))
   finally:
     _multistores_lock.release()
-  scope = util.scopes_to_string(scope)
+  key = util.dict_to_tuple_key(key_dict)
-  return multistore._get_storage(client_id, user_agent, scope)
+  return multistore._get_storage(key)
 
 
 class _MultiStore(object):
@@ -103,11 +152,11 @@ class _MultiStore(object):
 
     self._create_file_if_needed()
 
-    # Cache of deserialized store.  This is only valid after the
+    # Cache of deserialized store. This is only valid after the
-    # _MultiStore is locked or _refresh_data_cache is called.  This is
+    # _MultiStore is locked or _refresh_data_cache is called. This is
     # of the form of:
     #
-    # (client_id, user_agent, scope) -> OAuth2Credential
+    # ((key, value), (key, value)...) -> OAuth2Credential
     #
     # If this is None, then the store hasn't been read yet.
     self._data = None
@@ -115,11 +164,9 @@ class _MultiStore(object):
   class _Storage(BaseStorage):
     """A Storage object that knows how to read/write a single credential."""
 
-    def __init__(self, multistore, client_id, user_agent, scope):
+    def __init__(self, multistore, key):
       self._multistore = multistore
-      self._client_id = client_id
+      self._key = key
-      self._user_agent = user_agent
-      self._scope = scope
 
     def acquire_lock(self):
       """Acquires any lock necessary to access this Storage.
@@ -144,8 +191,7 @@ class _MultiStore(object):
       Returns:
         oauth2client.client.Credentials
       """
-      credential = self._multistore._get_credential(
+      credential = self._multistore._get_credential(self._key)
-          self._client_id, self._user_agent, self._scope)
       if credential:
         credential.set_store(self)
       return credential
@@ -158,7 +204,7 @@ class _MultiStore(object):
       Args:
         credentials: Credentials, the credentials to store.
       """
-      self._multistore._update_credential(credentials, self._scope)
+      self._multistore._update_credential(self._key, credentials)
 
     def locked_delete(self):
       """Delete a credential.
@@ -168,8 +214,7 @@ class _MultiStore(object):
       Args:
         credentials: Credentials, the credentials to store.
       """
-      self._multistore._delete_credential(self._client_id, self._user_agent,
+      self._multistore._delete_credential(self._key)
-          self._scope)
 
   def _create_file_if_needed(self):
     """Create an empty file if necessary.
@@ -201,9 +246,9 @@ class _MultiStore(object):
       self._write()
     elif not self._read_only or self._data is None:
       # Only refresh the data if we are read/write or we haven't
-      # cached the data yet.  If we are readonly, we assume is isn't
+      # cached the data yet. If we are readonly, we assume is isn't
       # changing out from under us and that we only have to read it
-      # once.  This prevents us from whacking any new access keys that
+      # once. This prevents us from whacking any new access keys that
       # we have cached in memory but were unable to write out.
       self._refresh_data_cache()
 
@@ -292,10 +337,7 @@ class _MultiStore(object):
         OAuth2Credential object.
     """
     raw_key = cred_entry['key']
-    client_id = raw_key['clientId']
+    key = util.dict_to_tuple_key(raw_key)
-    user_agent = raw_key['userAgent']
-    scope = raw_key['scope']
-    key = (client_id, user_agent, scope)
     credential = None
     credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credential']))
     return (key, credential)
@@ -309,73 +351,59 @@ class _MultiStore(object):
     raw_creds = []
     raw_data['data'] = raw_creds
     for (cred_key, cred) in self._data.items():
-      raw_key = {
+      raw_key = dict(cred_key)
-          'clientId': cred_key[0],
-          'userAgent': cred_key[1],
-          'scope': cred_key[2]
-          }
       raw_cred = simplejson.loads(cred.to_json())
       raw_creds.append({'key': raw_key, 'credential': raw_cred})
     self._locked_json_write(raw_data)
 
-  def _get_credential(self, client_id, user_agent, scope):
+  def _get_credential(self, key):
     """Get a credential from the multistore.
 
     The multistore must be locked.
 
     Args:
-      client_id: The client_id for the credential
+      key: The key used to retrieve the credential
-      user_agent: The user agent for the credential
-      scope: A string for the scope(s) being requested
 
     Returns:
       The credential specified or None if not present
     """
-    key = (client_id, user_agent, scope)
-
     return self._data.get(key, None)
 
-  def _update_credential(self, cred, scope):
+  def _update_credential(self, key, cred):
     """Update a credential and write the multistore.
 
     This must be called when the multistore is locked.
 
     Args:
+      key: The key used to retrieve the credential
       cred: The OAuth2Credential to update/set
-      scope: The scope(s) that this credential covers
     """
-    key = (cred.client_id, cred.user_agent, scope)
     self._data[key] = cred
     self._write()
 
-  def _delete_credential(self, client_id, user_agent, scope):
+  def _delete_credential(self, key):
     """Delete a credential and write the multistore.
 
     This must be called when the multistore is locked.
 
     Args:
-      client_id: The client_id for the credential
+      key: The key used to retrieve the credential
-      user_agent: The user agent for the credential
-      scope: The scope(s) that this credential covers
     """
-    key = (client_id, user_agent, scope)
     try:
       del self._data[key]
     except KeyError:
       pass
     self._write()
 
-  def _get_storage(self, client_id, user_agent, scope):
+  def _get_storage(self, key):
     """Get a Storage object to get/set a credential.
 
     This Storage is a 'view' into the multistore.
 
     Args:
-      client_id: The client_id for the credential
+      key: The key used to retrieve the credential
-      user_agent: The user agent for the credential
-      scope: A string for the scope(s) being requested
 
     Returns:
       A Storage object that can be used to get/set this cred
     """
-    return self._Storage(self, client_id, user_agent, scope)
+    return self._Storage(self, key)

oauth2client/util.py

@@ -42,7 +42,7 @@ def positional(max_positional_args):
   """A decorator to declare that only the first N arguments my be positional.
 
   This decorator makes it easy to support Python 3 style key-word only
-  parameters.  For example, in Python 3 it is possible to write:
+  parameters. For example, in Python 3 it is possible to write:
 
     def fn(pos1, *, kwonly1=None, kwonly1=None):
       ...
@@ -92,7 +92,7 @@ def positional(max_positional_args):
   respectively, if a declaration is violated.
 
   Args:
-    max_positional_arguments: Maximum number of positional arguments.  All
+    max_positional_arguments: Maximum number of positional arguments. All
       parameters after the this index must be keyword only.
 
   Returns:
@@ -145,3 +145,18 @@ def scopes_to_string(scopes):
     return scopes
   else:
     return ' '.join(scopes)
+
+
+def dict_to_tuple_key(dictionary):
+  """Converts a dictionary to a tuple that can be used as an immutable key.
+
+  The resulting key is always sorted so that logically equivalent dictionaries
+  always produce an identical tuple for a key.
+
+  Args:
+    dictionary: the dictionary to use as the key.
+
+  Returns:
+    A tuple representing the dictionary in it's naturally sorted ordering.
+  """
+  return tuple(sorted(dictionary.items()))

tests/test_oauth2client_file.py

@@ -35,6 +35,7 @@ from apiclient.http import HttpMockSequence
 from oauth2client import file
 from oauth2client import locked_file
 from oauth2client import multistore_file
+from oauth2client import util
 from oauth2client.anyjson import simplejson
 from oauth2client.client import AccessTokenCredentials
 from oauth2client.client import AssertionCredentials
@@ -58,6 +59,21 @@ class OAuth2ClientFileTests(unittest.TestCase):
     except OSError:
       pass
 
+  def create_test_credentials(self):
+    access_token = 'foo'
+    client_secret = 'cOuDdkfjxxnv+'
+    refresh_token = '1/0/a.df219fjls0'
+    token_expiry = datetime.datetime.utcnow()
+    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
+    user_agent = 'refresh_checker/1.0'
+    client_id = 'some_client_id'
+
+    credentials = OAuth2Credentials(
+        access_token, client_id, client_secret,
+        refresh_token, token_expiry, token_uri,
+        user_agent)
+    return credentials
+
   def test_non_existent_file_storage(self):
     s = file.Storage(FILENAME)
     credentials = s.get()
@@ -78,18 +94,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
 
   def test_pickle_and_json_interop(self):
     # Write a file with a pickled OAuth2Credentials.
-    access_token = 'foo'
+    credentials = self.create_test_credentials()
-    client_id = 'some_client_id'
-    client_secret = 'cOuDdkfjxxnv+'
-    refresh_token = '1/0/a.df219fjls0'
-    token_expiry = datetime.datetime.utcnow()
-    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
-    user_agent = 'refresh_checker/1.0'
-
-    credentials = OAuth2Credentials(
-        access_token, client_id, client_secret,
-        refresh_token, token_expiry, token_uri,
-        user_agent)
 
     f = open(FILENAME, 'w')
     pickle.dump(credentials, f)
@@ -112,18 +117,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
     self.assertEquals(data['_module'], OAuth2Credentials.__module__)
 
   def test_token_refresh(self):
-    access_token = 'foo'
+    credentials = self.create_test_credentials()
-    client_id = 'some_client_id'
-    client_secret = 'cOuDdkfjxxnv+'
-    refresh_token = '1/0/a.df219fjls0'
-    token_expiry = datetime.datetime.utcnow()
-    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
-    user_agent = 'refresh_checker/1.0'
-
-    credentials = OAuth2Credentials(
-        access_token, client_id, client_secret,
-        refresh_token, token_expiry, token_uri,
-        user_agent)
 
     s = file.Storage(FILENAME)
     s.put(credentials)
@@ -136,18 +130,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
     self.assertEquals(credentials.access_token, 'bar')
 
   def test_credentials_delete(self):
-    access_token = 'foo'
+    credentials = self.create_test_credentials()
-    client_id = 'some_client_id'
-    client_secret = 'cOuDdkfjxxnv+'
-    refresh_token = '1/0/a.df219fjls0'
-    token_expiry = datetime.datetime.utcnow()
-    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
-    user_agent = 'refresh_checker/1.0'
-
-    credentials = OAuth2Credentials(
-        access_token, client_id, client_secret,
-        refresh_token, token_expiry, token_uri,
-        user_agent)
 
     s = file.Storage(FILENAME)
     s.put(credentials)
@@ -175,18 +158,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
       self.assertEquals('0600', oct(stat.S_IMODE(os.stat(FILENAME).st_mode)))
 
   def test_read_only_file_fail_lock(self):
-    access_token = 'foo'
+    credentials = self.create_test_credentials()
-    client_secret = 'cOuDdkfjxxnv+'
-    refresh_token = '1/0/a.df219fjls0'
-    token_expiry = datetime.datetime.utcnow()
-    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
-    user_agent = 'refresh_checker/1.0'
-    client_id = 'some_client_id'
-
-    credentials = OAuth2Credentials(
-        access_token, client_id, client_secret,
-        refresh_token, token_expiry, token_uri,
-        user_agent)
 
     open(FILENAME, 'a+b').close()
     os.chmod(FILENAME, 0400)
@@ -230,18 +202,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
     self.assertEquals(None, credentials)
 
   def test_multistore_file(self):
-    access_token = 'foo'
+    credentials = self.create_test_credentials()
-    client_secret = 'cOuDdkfjxxnv+'
-    refresh_token = '1/0/a.df219fjls0'
-    token_expiry = datetime.datetime.utcnow()
-    token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
-    user_agent = 'refresh_checker/1.0'
-    client_id = 'some_client_id'
-
-    credentials = OAuth2Credentials(
-        access_token, client_id, client_secret,
-        refresh_token, token_expiry, token_uri,
-        user_agent)
 
     store = multistore_file.get_credential_storage(
         FILENAME,
@@ -263,5 +224,65 @@ class OAuth2ClientFileTests(unittest.TestCase):
     if os.name == 'posix':
       self.assertEquals('0600', oct(stat.S_IMODE(os.stat(FILENAME).st_mode)))
 
+  def test_multistore_file_custom_key(self):
+    credentials = self.create_test_credentials()
+
+    custom_key = {'myapp': 'testing', 'clientid': 'some client'}
+    store = multistore_file.get_credential_storage_custom_key(
+        FILENAME, custom_key)
+
+    store.put(credentials)
+    stored_credentials = store.get()
+
+    self.assertNotEquals(None, stored_credentials)
+    self.assertEqual(credentials.access_token, stored_credentials.access_token)
+
+    store.delete()
+    stored_credentials = store.get()
+
+    self.assertEquals(None, stored_credentials)
+
+  def test_multistore_file_custom_string_key(self):
+    credentials = self.create_test_credentials()
+
+    # store with string key
+    store = multistore_file.get_credential_storage_custom_string_key(
+        FILENAME, 'mykey')
+
+    store.put(credentials)
+    stored_credentials = store.get()
+
+    self.assertNotEquals(None, stored_credentials)
+    self.assertEqual(credentials.access_token, stored_credentials.access_token)
+
+    # try retrieving with a dictionary
+    store_dict = multistore_file.get_credential_storage_custom_string_key(
+        FILENAME, {'key': 'mykey'})
+    stored_credentials = store.get()
+    self.assertNotEquals(None, stored_credentials)
+    self.assertEqual(credentials.access_token, stored_credentials.access_token)
+
+    store.delete()
+    stored_credentials = store.get()
+
+    self.assertEquals(None, stored_credentials)
+
+  def test_multistore_file_backwards_compatibility(self):
+    credentials = self.create_test_credentials()
+    scopes = ['scope1', 'scope2']
+
+    # store the credentials using the legacy key method
+    store = multistore_file.get_credential_storage(
+        FILENAME, 'client_id', 'user_agent', scopes)
+    store.put(credentials)
+
+    # retrieve the credentials using a custom key that matches the legacy key
+    key = {'clientId': 'client_id', 'userAgent': 'user_agent',
+           'scope': util.scopes_to_string(scopes)}
+    store = multistore_file.get_credential_storage_custom_key(FILENAME, key)
+    stored_credentials = store.get()
+
+    self.assertEqual(credentials.access_token, stored_credentials.access_token)
+
 if __name__ == '__main__':
   unittest.main()

tests/test_oauth2client_util.py

@@ -25,3 +25,20 @@ class ScopeToStringTests(unittest.TestCase):
     ]
     for expected, case in cases:
       self.assertEqual(expected, util.scopes_to_string(case))
+
+
+class KeyConversionTests(unittest.TestCase):
+
+  def test_key_conversions(self):
+    d = {'somekey': 'some value', 'another': 'something else', 'onemore': 'foo'}
+    tuple_key = util.dict_to_tuple_key(d)
+
+    # the resulting key should be naturally sorted
+    self.assertEqual(
+        (('another', 'something else'),
+         ('onemore', 'foo'),
+         ('somekey', 'some value')),
+        tuple_key)
+
+    # check we get the original dictionary back
+    self.assertEqual(d, dict(tuple_key))