def get_datastore(config=None, archive_access=False):
from assemblyline.datastore.helper import AssemblylineDatastore
from assemblyline.datastore.store import ESStore
if not config:
config = get_config()
if archive_access:
return AssemblylineDatastore(
ESStore(config.datastore.hosts, archive_access=True))
else:
return AssemblylineDatastore(
ESStore(config.datastore.hosts, archive_access=False))
def get_datastore(config=None, archive_access=False):
from assemblyline.datastore.helper import AssemblylineDatastore
if not config:
config = get_config(static=True)
if config.datastore.type == "elasticsearch":
from assemblyline.datastore.stores.es_store import ESStore
if archive_access:
return AssemblylineDatastore(ESStore(config.datastore.hosts, archive_access=True))
else:
return AssemblylineDatastore(ESStore(config.datastore.hosts, archive_access=False))
else:
from assemblyline.datastore.exceptions import DataStoreException
raise DataStoreException(f"Invalid datastore type: {config.datastore.type}")
def datastore_connection():
store = ESStore(config.datastore.hosts)
ret_val = store.ping()
if not ret_val:
pytest.skip("Could not connect to datastore")
return AssemblylineDatastore(store)
def archive_connection(config):
store = ESStore(config.datastore.hosts, archive_access=True)
ret_val = store.ping()
if not ret_val:
pytest.skip("Could not connect to datastore")
return AssemblylineDatastore(store)
def ingest_harness(clean_redis):
""""Setup a test environment.
By using a fake redis and datastore, we:
a) ensure that this test runs regardless of any code errors in replaced modules
b) ensure that the datastore and redis is EMPTY every time the test runs
isolating this test from any other test run at the same time
"""
datastore = AssemblylineDatastore(MockDatastore())
ingester = IngesterInput(datastore=datastore,
redis=clean_redis,
persistent_redis=clean_redis)
ingester.running = TrueCountTimes(1)
return datastore, ingester, ingester.ingester.ingest_queue
def ds():
return AssemblylineDatastore(MockDatastore())
def function_clean_datastore(datastore_connection: AssemblylineDatastore):
for name in datastore_connection.ds.get_models():
datastore_connection.get_collection(name).wipe()
return datastore_connection