diff --git a/oauth2client/multistore_file.py b/oauth2client/multistore_file.py index e1b39f7..4d2c091 100644 --- a/oauth2client/multistore_file.py +++ b/oauth2client/multistore_file.py @@ -125,6 +125,43 @@ def get_credential_storage_custom_key( An object derived from client.Storage for getting/setting the credential. """ + multistore = _get_multistore(filename, warn_on_readonly=warn_on_readonly) + key = util.dict_to_tuple_key(key_dict) + return multistore._get_storage(key) + + +@util.positional(1) +def get_all_credential_keys(filename, warn_on_readonly=True): + """Gets all the registered credential keys in the given Multistore. + + Args: + filename: The JSON file storing a set of credentials + warn_on_readonly: if True, log a warning if the store is readonly + + Returns: + A list of the credential keys present in the file. They are returned as + dictionaries that can be passed into get_credential_storage_custom_key to + get the actual credentials. + """ + multistore = _get_multistore(filename, warn_on_readonly=warn_on_readonly) + multistore._lock() + try: + return multistore._get_all_credential_keys() + finally: + multistore._unlock() + + +@util.positional(1) +def _get_multistore(filename, warn_on_readonly=True): + """A helper method to initialize the multistore with proper locking. + + Args: + filename: The JSON file storing a set of credentials + warn_on_readonly: if True, log a warning if the store is readonly + + Returns: + A multistore object + """ filename = os.path.expanduser(filename) _multistores_lock.acquire() try: @@ -132,8 +169,7 @@ def get_credential_storage_custom_key( filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly)) finally: _multistores_lock.release() - key = util.dict_to_tuple_key(key_dict) - return multistore._get_storage(key) + return multistore class _MultiStore(object): @@ -356,6 +392,14 @@ class _MultiStore(object): raw_creds.append({'key': raw_key, 'credential': raw_cred}) self._locked_json_write(raw_data) + def _get_all_credential_keys(self): + """Gets all the registered credential keys in the multistore. + + Returns: + A list of dictionaries corresponding to all the keys currently registered + """ + return [dict(key) for key in self._data.keys()] + def _get_credential(self, key): """Get a credential from the multistore. diff --git a/tests/test_oauth2client_file.py b/tests/test_oauth2client_file.py index c954d5e..786d708 100644 --- a/tests/test_oauth2client_file.py +++ b/tests/test_oauth2client_file.py @@ -60,14 +60,13 @@ class OAuth2ClientFileTests(unittest.TestCase): except OSError: pass - def create_test_credentials(self): + def create_test_credentials(self, client_id='some_client_id'): 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, @@ -285,5 +284,40 @@ class OAuth2ClientFileTests(unittest.TestCase): self.assertEqual(credentials.access_token, stored_credentials.access_token) + + def test_multistore_file_get_all_keys(self): + # start with no keys + keys = multistore_file.get_all_credential_keys(FILENAME) + self.assertEquals([], keys) + + # store credentials + credentials = self.create_test_credentials(client_id='client1') + custom_key = {'myapp': 'testing', 'clientid': 'client1'} + store1 = multistore_file.get_credential_storage_custom_key( + FILENAME, custom_key) + store1.put(credentials) + + keys = multistore_file.get_all_credential_keys(FILENAME) + self.assertEquals([custom_key], keys) + + # store more credentials + credentials = self.create_test_credentials(client_id='client2') + string_key = 'string_key' + store2 = multistore_file.get_credential_storage_custom_string_key( + FILENAME, string_key) + store2.put(credentials) + + keys = multistore_file.get_all_credential_keys(FILENAME) + self.assertEquals(2, len(keys)) + self.assertTrue(custom_key in keys) + self.assertTrue({'key': string_key} in keys) + + # back to no keys + store1.delete() + store2.delete() + keys = multistore_file.get_all_credential_keys(FILENAME) + self.assertEquals([], keys) + + if __name__ == '__main__': unittest.main()