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_session() -> Generator:
db_session = None
try:
db_session = create_session()
yield db_session
finally:
if db_session is not None:
db_session.close()
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()
async def initialize_scheduler():
global scheduler
scheduler = Scheduler()
db_session = None
try:
db_session = create_session()
await scheduler.start(db_session,)
finally:
db_session.close()
async def initialize_scheduler():
global scheduler
scheduler = Scheduler()
db_session = None
try:
db_session = create_session()
await scheduler.start(
db_session,
mlrun.api.schemas.AuthInfo(session=mlrun.config.config.httpdb.
projects.iguazio_access_key),
)
finally:
db_session.close()
def _initialize_db():
global db
if config.httpdb.db_type == "filedb":
logger.info("using FileRunDB")
db = FileDB(config.httpdb.dirpath)
db.initialize(None)
else:
logger.info("using SQLDB")
db = SQLDB(config.httpdb.dsn)
db_session = None
try:
db_session = create_session()
db.initialize(db_session)
finally:
db_session.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
def initialize_db(override_db=None):
global db
if override_db:
db = override_db
return
if config.httpdb.db_type == "filedb":
logger.info("Creating file db")
db = FileDB(config.httpdb.dirpath)
db.initialize(None)
else:
logger.info("Creating sql db")
db = SQLDB(config.httpdb.dsn)
db_session = None
try:
db_session = create_session()
db.initialize(db_session)
finally:
db_session.close()
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
def connect(self, secrets=None):
if not self.session:
self.session = create_session()
self.db = SQLAPIDB(self.dsn)
return self
