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 data_migration_db(request) -> Generator:
# Data migrations performed before the API goes up, therefore there's no project member yet
# that's the only difference between this fixture and the db fixture. because of the parameterization it was hard to
# share code between them, we anyway going to remove filedb soon, then there won't be params, and we could re-use
# code
# TODO: fix duplication
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)
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()
def db():
global session_maker
dsn = "sqlite:///:memory:?check_same_thread=false"
db_session = None
try:
config.httpdb.dsn = dsn
_init_engine(dsn)
init_data()
initialize_db()
db_session = create_session()
db = SQLDB(dsn)
db.initialize(db_session)
finally:
if db_session is not None:
db_session.close()
mlrun.api.utils.singletons.db.initialize_db(db)
mlrun.api.utils.singletons.project_member.initialize_project_member()
return db
async def startup_event():
logger.info(
"configuration dump",
dumped_config=config.dump_yaml(),
version=mlrun.utils.version.Version().get(),
)
loop = asyncio.get_running_loop()
# Using python 3.8 default instead of 3.7 one - max(1, os.cpu_count()) * 5 cause it's causing to high memory
# consumption - https://bugs.python.org/issue35279
# TODO: remove when moving to python 3.8
max_workers = config.httpdb.max_workers or min(32, os.cpu_count() + 4)
loop.set_default_executor(
concurrent.futures.ThreadPoolExecutor(max_workers=max_workers))
initialize_logs_dir()
initialize_db()
if config.httpdb.state == mlrun.api.schemas.APIStates.online:
await move_api_to_online()
def db(request):
path = mkdtemp()
print(f"db fixture: path={path!r}")
if request.param == "sql":
db_file = f"{path}/mlrun.db"
dsn = f"sqlite:///{db_file}?check_same_thread=false"
config.httpdb.dsn = dsn
_init_engine(dsn)
init_data()
initialize_db()
db_session = create_session()
db = SQLDB(dsn, session=db_session)
elif request.param == "file":
db = FileRunDB(path)
else:
assert False, f"unknown db type - {request.param}"
db.connect()
if request.param == "sql":
mlrun.api.utils.singletons.db.initialize_db(db.db)
mlrun.api.utils.singletons.project_member.initialize_project_member()
return db
async def _initialize_singletons():
initialize_logs_dir()
initialize_db()
initialize_project_member()
await initialize_scheduler()
