def CreateTaskStorage(self, task, task_storage_format):
"""Creates a task storage.
The task storage is used to store attribute containers created by the task.
Args:
task (Task): task.
task_storage_format (str): storage format used to store task results.
Returns:
StorageWriter: storage writer.
Raises:
OSError: if the storage type or storage format is not supported.
IOError: if the storage type or storage format is not supported.
"""
if self._storage_type != definitions.STORAGE_TYPE_SESSION:
raise IOError('Unsupported storage type: {0:s}'.format(
self._storage_type))
if task_storage_format == definitions.STORAGE_FORMAT_SQLITE:
writer = self._CreateTaskStorageWriter(task)
elif task_storage_format == definitions.STORAGE_FORMAT_REDIS:
writer = redis_writer.RedisStorageWriter(
self._session, storage_type=definitions.STORAGE_TYPE_TASK, task=task)
else:
raise IOError('Unsupported storage format: {0:s}'.format(
task_storage_format))
task.storage_format = task_storage_format
return writer
def _CreateTaskStore(self, session, redis_client):
"""Creates a task store for testing.
Args:
session (Session): session the task store is part of.
redis_client (Redis): Redis client to query. If specified, no
new client will be created.
Returns:
Task: the the task for the store that was created.
"""
task = tasks.Task(session_identifier=session.identifier)
task_store = redis_writer.RedisStorageWriter(
session, storage_type=definitions.STORAGE_TYPE_TASK, task=task)
task_store.Open(redis_client=redis_client)
for event, event_data in containers_test_lib.CreateEventsFromValues(
self._TEST_EVENTS):
task_store.AddEventData(event_data)
event.SetEventDataIdentifier(event_data.GetIdentifier())
task_store.AddEvent(event)
task_store.Close()
return task
def testScanForProcessedTasks(self):
"""Tests the ScanForProcessedTasks method"""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
# There should be no processed task identifiers initially.
task_identifiers, _ = redis_store.RedisStore.ScanForProcessedTasks(
session.identifier, redis_client=redis_client)
self.assertEqual([], task_identifiers)
# Opening and closing a writer for a task should cause the task to be marked
# as complete.
storage_writer = writer.RedisStorageWriter(
storage_type=definitions.STORAGE_TYPE_TASK)
storage_writer.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
storage_writer.Close()
# The task store is now marked as processed.
task_identifiers, _ = redis_store.RedisStore.ScanForProcessedTasks(
session.identifier, redis_client=redis_client)
self.assertEqual([task.identifier], task_identifiers)
def testMarkTaskAsMerging(self):
"""Tests the MarkTaskAsMerging method"""
redis_client = self._CreateRedisClient()
session = sessions.Session()
task = tasks.Task(session_identifier=session.identifier)
# Trying to mark a task as merging without finalizing it raises an error.
with self.assertRaises(IOError):
redis_store.RedisStore.MarkTaskAsMerging(task.identifier,
session.identifier,
redis_client=redis_client)
# Opening and closing a writer for a task should cause the task to be marked
# as complete.
storage_writer = writer.RedisStorageWriter(
storage_type=definitions.STORAGE_TYPE_TASK)
storage_writer.Open(redis_client=redis_client,
session_identifier=task.session_identifier,
task_identifier=task.identifier)
storage_writer.Close()
redis_store.RedisStore.MarkTaskAsMerging(task.identifier,
session.identifier,
redis_client=redis_client)
def CreateStorageWriter(cls, storage_format):
"""Creates a storage writer.
Args:
storage_format (str): storage format.
Returns:
StorageWriter: a storage writer or None if the storage file cannot be
opened or the storage format is not supported.
"""
if storage_format == definitions.STORAGE_FORMAT_SQLITE:
return sqlite_writer.SQLiteStorageFileWriter()
if storage_format == definitions.STORAGE_FORMAT_REDIS:
return redis_writer.RedisStorageWriter()
return None
def CreateStorageWriter(cls, storage_format, session, path):
"""Creates a storage writer.
Args:
storage_format (str): storage format.
session (Session): session the storage changes are part of.
path (str): path to the storage file.
Returns:
StorageWriter: a storage writer or None if the storage file cannot be
opened or the storage format is not supported.
"""
if storage_format == definitions.STORAGE_FORMAT_SQLITE:
return sqlite_writer.SQLiteStorageFileWriter(session, path)
if storage_format == definitions.STORAGE_FORMAT_REDIS:
return redis_writer.RedisStorageWriter(session)
return None
