class InMemoryTokenBackend(BaseStorageBackend): """ In memory backend (based on SQLite) for token. """ def __init__(self): super(InMemoryTokenBackend, self).__init__() if settings.DATABASES: self.db_manager = settings.DATABASES['default']['manager'] else: self.db_manager = SQLiteManager(name=IN_MEMORY) self.__create_models() def __create_models(self): self.db_manager.execute_script(SQL_CREATE_TOKEN_TABLE) def get(self, token): """ Get token from the storage. :param token: token as string. """ args = (token, ) row = self.db_manager.execute_sql_and_fetchone(SQL_TOKEN_GET, args) if row: token_object = convert_db_row_to_dict(row, TOKEN_MODEL_FIELDS) else: token_object = {} return token_object def get_token_by_username(self, token_name, username): """ Get token from the storage by token_name and username. :param token: token as string. """ args = (token_name, username) row = self.db_manager.execute_sql_and_fetchone( SQL_TOKEN_GET_BY_TOKEN_USERNAME, args ) if row: token_object = convert_db_row_to_dict(row, TOKEN_MODEL_FIELDS) else: token_object = {} return token_object def save(self, token_name, token, expired=None, user_id=None): """ Save token in the storage. :param user: instance of User class. :param token: token as string. """ args = (token_name, token, expired, user_id) self.db_manager.execute_sql(SQL_TOKEN_ADD, args)
class InMemoryTokenBackend(BaseStorageBackend): """ In memory backend (based on SQLite) for token. """ def __init__(self): super(InMemoryTokenBackend, self).__init__() if settings.DATABASES: self.db_manager = settings.DATABASES['default']['manager'] else: self.db_manager = SQLiteManager(name=IN_MEMORY) self.__create_models() def __create_models(self): self.db_manager.execute_script(SQL_CREATE_TOKEN_TABLE) def get(self, token): """ Get token from the storage. :param token: token as string. """ args = (token, ) row = self.db_manager.execute_sql_and_fetchone(SQL_TOKEN_GET, args) if row: token_object = convert_db_row_to_dict(row, TOKEN_MODEL_FIELDS) else: token_object = {} return token_object def get_token_by_username(self, token_name, username): """ Get token from the storage by token_name and username. :param token: token as string. """ args = (token_name, username) row = self.db_manager.execute_sql_and_fetchone( SQL_TOKEN_GET_BY_TOKEN_USERNAME, args) if row: token_object = convert_db_row_to_dict(row, TOKEN_MODEL_FIELDS) else: token_object = {} return token_object def save(self, token_name, token, expired=None, user_id=None): """ Save token in the storage. :param user: instance of User class. :param token: token as string. """ args = (token_name, token, expired, user_id) self.db_manager.execute_sql(SQL_TOKEN_ADD, args)
class SQLiteManagerTestCase(unittest.TestCase): def setUp(self): super(SQLiteManagerTestCase, self).setUp() self.manager = SQLiteManager(name=IN_MEMORY) def test_execute_sql(self): self.manager.execute_sql(SQL_CREATE_USER_TABLE) def test_execute_sql_and_fetchone(self): username = '******' self.manager.execute_sql(SQL_CREATE_USER_TABLE) self.manager.execute_sql(SQL_USER_ADD, (username, )) row = self.manager.execute_sql_and_fetchone(SQL_USER_GET, (username, )) self.assertEqual(row, (1, username)) def test_execute_sql_from_file(self): self.manager.execute_sql(SQL_CREATE_USER_TABLE) with tempfile.NamedTemporaryFile() as tmpfile: file_sql = bytes(""" INSERT INTO test_user_table (`username`) VALUES ('test_user'); """, encoding='utf-8') tmpfile.write(file_sql) self.manager.execute_sql_from_file(tmpfile.name) def test_execute_script(self): self.manager.execute_sql(SQL_CREATE_USER_TABLE) self.manager.execute_script(""" INSERT INTO test_user_table (`username`) VALUES ('test_user'); """)