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(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 main():
init_data()
uvicorn.run(
'mlrun.api.main:app',
host='0.0.0.0',
port=config.httpdb.port,
debug=config.httpdb.debug,
)
def db() -> Generator:
db_file = NamedTemporaryFile(suffix="-mlrun.db")
logger.info(f"Created temp db file: {db_file.name}")
_init_engine(f"sqlite:///{db_file.name}?check_same_thread=false")
init_data()
yield create_session()
logger.info(f"Removing temp db file: {db_file.name}")
db_file.close()
def main():
init_data()
uvicorn.run(
"mlrun.api.main:app",
host="0.0.0.0",
port=config.httpdb.port,
debug=config.httpdb.debug,
access_log=False,
)
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
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
