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