def db(request) -> Generator:
if request.param == "sqldb":
dsn = "sqlite:///:memory:?check_same_thread=false"
config.httpdb.dsn = dsn
_init_engine()
# memory sqldb remove it self when all session closed, this session will keep it up during all test
db_session = create_session()
try:
init_data()
db = SQLDB(dsn)
db.initialize(db_session)
initialize_db(db)
initialize_project_member()
yield db
finally:
close_session(db_session)
elif request.param == "filedb":
db = FileDB(config.httpdb.dirpath)
db_session = create_session(request.param)
try:
db.initialize(db_session)
yield db
finally:
shutil.rmtree(config.httpdb.dirpath,
ignore_errors=True,
onerror=None)
close_session(db_session)
else:
raise Exception("Unknown db type")
def db() -> Generator:
"""
This fixture initialize the db singleton (so it will be accessible using mlrun.api.singletons.get_db()
and generates a db session that can be used by the test
"""
db_file = NamedTemporaryFile(suffix="-mlrun.db")
logger.info(f"Created temp db file: {db_file.name}")
config.httpdb.db_type = "sqldb"
dsn = f"sqlite:///{db_file.name}?check_same_thread=false"
config.httpdb.dsn = dsn
# TODO: make it simpler - doesn't make sense to call 3 different functions to initialize the db
# we need to force re-init the engine cause otherwise it is cached between tests
_init_engine(config.httpdb.dsn)
# forcing from scratch because we created an empty file for the db
init_data(from_scratch=True)
initialize_db()
initialize_project_member()
# we're also running client code in tests so set dbpath as well
# note that setting this attribute triggers connection to the run db therefore must happen after the initialization
config.dbpath = dsn
yield create_session()
logger.info(f"Removing temp db file: {db_file.name}")
db_file.close()
async def move_api_to_online():
logger.info("Moving api to online")
initialize_project_member()
await initialize_scheduler()
# periodic cleanup is not needed if we're not inside kubernetes cluster
if get_k8s_helper(silent=True).is_running_inside_kubernetes_cluster():
_start_periodic_cleanup()
_start_periodic_runs_monitoring()
async def _initialize_singletons():
initialize_logs_dir()
initialize_db()
initialize_project_member()
await initialize_scheduler()