def test_list_artifact_tags(db: SQLDB, db_session: Session):
db.store_artifact(db_session, "k1", {}, "1", tag="t1", project="p1")
db.store_artifact(db_session, "k1", {}, "2", tag="t2", project="p1")
db.store_artifact(db_session, "k1", {}, "2", tag="t2", project="p2")
tags = db.list_artifact_tags(db_session, "p1")
assert {"t1", "t2"} == set(tags), "bad tags"
def test_list_artifact_tags(db: SQLDB, db_session: Session):
db.store_artifact(db_session, 'k1', {}, '1', tag='t1', project='p1')
db.store_artifact(db_session, 'k1', {}, '2', tag='t2', project='p1')
db.store_artifact(db_session, 'k1', {}, '2', tag='t2', project='p2')
tags = db.list_artifact_tags(db_session, 'p1')
assert {'t1', 't2'} == set(tags), 'bad tags'
def test_list_artifact_tags(db: SQLDB, db_session: Session):
db.store_artifact(db_session, "k1", {}, "1", tag="t1", project="p1")
db.store_artifact(db_session, "k1", {}, "2", tag="t2", project="p1")
db.store_artifact(db_session, "k1", {}, "2", tag="t2", project="p2")
tags = db.list_artifact_tags(db_session, "p1")
assert [("p1", "k1", "t1"), ("p1", "k1", "t2")] == tags
def test_list_projects(db: SQLDB, db_session: Session):
for i in range(10):
run = new_run("s1", {"l1": "v1", "l2": "v2"}, x=1)
db.store_run(db_session, run, "u7", project=f"prj{i % 3}", iter=i)
assert {"prj0", "prj1",
"prj2"} == {p.name
for p in db.list_projects(db_session)}
def test_schedules(db: SQLDB, db_session: Session):
count = 7
for i in range(count):
data = {'i': i}
db.store_schedule(db_session, data)
scheds = list(db.list_schedules(db_session))
assert count == len(scheds), 'wrong number of schedules'
assert set(range(count)) == set(s['i'] for s in scheds), 'bad scheds'
def _tag_objs(db: SQLDB, db_session: Session, count, project, tags):
by_tag = defaultdict(list)
for i in range(count):
cls = _tagged[i % len(_tagged)]
obj = cls()
by_tag[tags[i % len(tags)]].append(obj)
db_session.add(obj)
db_session.commit()
for tag, objs in by_tag.items():
db.tag_objects(db_session, objs, project, tag)
def test_list_artifact_date(db: SQLDB, db_session: Session):
t1 = datetime(2020, 2, 16)
t2 = t1 - timedelta(days=7)
t3 = t2 - timedelta(days=7)
prj = "p7"
db.store_artifact(db_session, "k1", {"updated": t1}, "u1", project=prj)
db.store_artifact(db_session, "k2", {"updated": t2}, "u2", project=prj)
db.store_artifact(db_session, "k3", {"updated": t3}, "u3", project=prj)
arts = db.list_artifacts(db_session, project=prj, since=t3, tag="*")
assert 3 == len(arts), "since t3"
arts = db.list_artifacts(db_session, project=prj, since=t2, tag="*")
assert 2 == len(arts), "since t2"
arts = db.list_artifacts(db_session,
project=prj,
since=t1 + timedelta(days=1),
tag="*")
assert not arts, "since t1+"
arts = db.list_artifacts(db_session, project=prj, until=t2, tag="*")
assert 2 == len(arts), "until t2"
arts = db.list_artifacts(db_session,
project=prj,
since=t2,
until=t2,
tag="*")
assert 1 == len(arts), "since/until t2"
def test_list_artifact_date(db: SQLDB, db_session: Session):
t1 = datetime(2020, 2, 16)
t2 = t1 - timedelta(days=7)
t3 = t2 - timedelta(days=7)
prj = 'p7'
db.store_artifact(db_session, 'k1', {'updated': t1}, 'u1', project=prj)
db.store_artifact(db_session, 'k2', {'updated': t2}, 'u2', project=prj)
db.store_artifact(db_session, 'k3', {'updated': t3}, 'u3', project=prj)
arts = db.list_artifacts(db_session, project=prj, since=t3, tag='*')
assert 3 == len(arts), 'since t3'
arts = db.list_artifacts(db_session, project=prj, since=t2, tag='*')
assert 2 == len(arts), 'since t2'
arts = db.list_artifacts(
db_session, project=prj, since=t1 + timedelta(days=1), tag='*')
assert not arts, 'since t1+'
arts = db.list_artifacts(db_session, project=prj, until=t2, tag='*')
assert 2 == len(arts), 'until t2'
arts = db.list_artifacts(db_session, project=prj, since=t2, until=t2, tag='*')
assert 1 == len(arts), 'since/until t2'
def test_projects_crud(db: SQLDB, db_session: Session):
project = mlrun.api.schemas.Project(
metadata=mlrun.api.schemas.ProjectMetadata(name="p1"),
spec=mlrun.api.schemas.ProjectSpec(description="banana",
other_field="value"),
status=mlrun.api.schemas.ObjectStatus(state="active"),
)
db.create_project(db_session, project)
project_output = db.get_project(db_session, name=project.metadata.name)
assert (deepdiff.DeepDiff(
project.dict(),
project_output.dict(exclude={"id"}),
ignore_order=True,
) == {})
project_patch = {"spec": {"description": "lemon"}}
db.patch_project(db_session, project.metadata.name, project_patch)
project_output = db.get_project(db_session, name=project.metadata.name)
assert project_output.spec.description == project_patch["spec"][
"description"]
project_2 = mlrun.api.schemas.Project(
metadata=mlrun.api.schemas.ProjectMetadata(name="p2"), )
db.create_project(db_session, project_2)
projects_output = db.list_projects(
db_session, format_=mlrun.api.schemas.Format.name_only)
assert [project.metadata.name,
project_2.metadata.name] == projects_output.projects
def test_list_projects(db: SQLDB, db_session: Session):
for idx in range(10):
run = new_run("s1", {"l1": "v1", "l2": "v2"}, x=1)
db.store_run(db_session, run, "u7", project=f"prj{idx % 3}", iter=idx)
projects_output = db.list_projects(db_session)
assert {"prj0", "prj1", "prj2"} == {
project.metadata.name
for project in projects_output.projects
}
def _tag_objs(db: SQLDB, db_session: Session, count, project, tags):
tagged = [tagged_model for tagged_model in _tagged if tagged_model != Run]
by_tag = defaultdict(list)
for idx in range(count):
cls = tagged[idx % len(tagged)]
obj = cls()
by_tag[tags[idx % len(tags)]].append(obj)
db_session.add(obj)
db_session.commit()
for tag, objs in by_tag.items():
db.tag_objects(db_session, objs, project, tag)
def db():
global session_maker
dsn = "sqlite:///:memory:?check_same_thread=false"
try:
session_maker = init_sqldb(dsn)
db_session = session_maker()
db = SQLDB(dsn)
db.initialize(db_session)
finally:
db_session.close()
return db
def test_projects(db: SQLDB, db_session: Session):
prj1 = {
"name": "p1",
"description": "banana",
# 'users': ['u1', 'u2'],
"spec": {
"company": "ACME"
},
"state": "active",
"created": datetime.now(),
}
pid1 = db.add_project(db_session, prj1)
p1 = db.get_project(db_session, project_id=pid1)
assert p1, f"project {pid1} not found"
out = {
"name": p1.name,
"description": p1.description,
# 'users': sorted(u.name for u in p1.users),
"spec": p1.spec,
"state": p1.state,
"created": p1.created,
}
assert prj1 == out, "bad project"
data = {"description": "lemon"}
db.update_project(db_session, p1.name, data)
p1 = db.get_project(db_session, project_id=pid1)
assert data["description"] == p1.description, "bad update"
prj2 = {"name": "p2"}
db.add_project(db_session, prj2)
prjs = {p.name for p in db.list_projects(db_session)}
assert {prj1["name"], prj2["name"]} == prjs, "list"
def test_projects(db: SQLDB, db_session: Session):
prj1 = {
'name': 'p1',
'description': 'banana',
# 'users': ['u1', 'u2'],
'spec': {'company': 'ACME'},
'state': 'active',
'created': datetime.now(),
}
pid1 = db.add_project(db_session, prj1)
p1 = db.get_project(db_session, project_id=pid1)
assert p1, f'project {pid1} not found'
out = {
'name': p1.name,
'description': p1.description,
# 'users': sorted(u.name for u in p1.users),
'spec': p1.spec,
'state': p1.state,
'created': p1.created,
}
assert prj1 == out, 'bad project'
data = {'description': 'lemon'}
db.update_project(db_session, p1.name, data)
p1 = db.get_project(db_session, project_id=pid1)
assert data['description'] == p1.description, 'bad update'
prj2 = {'name': 'p2'}
db.add_project(db_session, prj2)
prjs = {p.name for p in db.list_projects(db_session)}
assert {prj1['name'], prj2['name']} == prjs, 'list'
def test_list_tags(db: SQLDB, db_session: Session):
p1, tags1 = "prj1", ["a", "b", "c"]
_tag_objs(db, db_session, 17, p1, tags1)
p2, tags2 = "prj2", ["b", "c", "d", "e"]
_tag_objs(db, db_session, 11, p2, tags2)
tags = db.list_tags(db_session, p1)
assert set(tags) == set(tags1), "tags"
def test_list_tags(db: SQLDB, db_session: Session):
p1, tags1 = 'prj1', ['a', 'b', 'c']
_tag_objs(db, db_session, 17, p1, tags1)
p2, tags2 = 'prj2', ['b', 'c', 'd', 'e']
_tag_objs(db, db_session, 11, p2, tags2)
tags = db.list_tags(db_session, p1)
assert set(tags) == set(tags1), 'tags'
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(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 test_artifacts_latest(db: SQLDB, db_session: Session):
k1, u1, art1 = 'k1', 'u1', {'a': 1}
prj = 'p38'
db.store_artifact(db_session, k1, art1, u1, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag='latest')
assert art1['a'] == arts[0]['a'], 'bad artifact'
u2, art2 = 'u2', {'a': 17}
db.store_artifact(db_session, k1, art2, u2, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag='latest')
assert 1 == len(arts), 'count'
assert art2['a'] == arts[0]['a'], 'bad artifact'
k2, u3, art3 = 'k2', 'u3', {'a': 99}
db.store_artifact(db_session, k2, art3, u3, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag='latest')
assert 2 == len(arts), 'number'
assert {17, 99} == set(art['a'] for art in arts), 'latest'
def test_artifacts_latest(db: SQLDB, db_session: Session):
k1, u1, art1 = "k1", "u1", {"a": 1}
prj = "p38"
db.store_artifact(db_session, k1, art1, u1, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag="latest")
assert art1["a"] == arts[0]["a"], "bad artifact"
u2, art2 = "u2", {"a": 17}
db.store_artifact(db_session, k1, art2, u2, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag="latest")
assert 1 == len(arts), "count"
assert art2["a"] == arts[0]["a"], "bad artifact"
k2, u3, art3 = "k2", "u3", {"a": 99}
db.store_artifact(db_session, k2, art3, u3, project=prj)
arts = db.list_artifacts(db_session, project=prj, tag="latest")
assert 2 == len(arts), "number"
assert {17, 99} == set(art["a"] for art in arts), "latest"
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 test_tags(db: SQLDB, db_session: Session, cls):
p1, n1 = 'prj1', 'name1'
obj1, obj2, obj3 = cls(), cls(), cls()
db_session.add(obj1)
db_session.add(obj2)
db_session.add(obj3)
db_session.commit()
db.tag_objects(db_session, [obj1, obj2], p1, n1)
objs = db.find_tagged(db_session, p1, n1)
assert {obj1, obj2} == set(objs), 'find tags'
db.del_tag(db_session, p1, n1)
objs = db.find_tagged(db_session, p1, n1)
assert [] == objs, 'find tags after del'
def test_tags(db: SQLDB, db_session: Session, cls):
p1, n1 = "prj1", "name1"
object_identifier = "name"
if cls == Artifact:
object_identifier = "key"
obj1, obj2, obj3 = cls(), cls(), cls()
for index, obj in enumerate([obj1, obj2, obj3]):
setattr(obj, object_identifier, f"obj-identifier-{index}")
db_session.add(obj1)
db_session.add(obj2)
db_session.add(obj3)
db_session.commit()
db.tag_objects(db_session, [obj1, obj2], p1, n1)
objs = db.find_tagged(db_session, p1, n1)
assert {obj1, obj2} == set(objs)
db.del_tag(db_session, p1, n1)
objs = db.find_tagged(db_session, p1, n1)
assert [] == objs
def test_cache_projects(db: SQLDB, db_session: Session):
assert 0 == len(db._projects), "empty cache"
name = "prj348"
db.add_project(db_session, {"name": name})
assert {name} == db._projects, "project"
mock = Mock()
with patch(db, add_project=mock):
db._ensure_project(db_session, name)
mock.assert_not_called()
mock = Mock()
with patch(db, add_project=mock):
db._ensure_project(db_session, name + "-new")
mock.assert_called_once()
project_2_name = "project-2"
db.add_project(db_session, {"name": project_2_name})
db._projects = set()
mock = Mock()
with patch(db, add_project=mock):
db._ensure_project(db_session, name)
mock.assert_not_called()
def test_run_iter0(db: SQLDB, db_session: Session):
uid, prj = "uid39", "lemon"
run = new_run("s1", {"l1": "v1", "l2": "v2"}, x=1)
for i in range(7):
db.store_run(db_session, run, uid, prj, i)
db._get_run(db_session, uid, prj, 0) # See issue 140
def test_list_projects(db: SQLDB, db_session: Session):
for i in range(10):
run = new_run('s1', {'l1': 'v1', 'l2': 'v2'}, x=1)
db.store_run(db_session, run, 'u7', project=f'prj{i % 3}', iter=i)
assert {'prj0', 'prj1', 'prj2'} == {p.name for p in db.list_projects(db_session)}
def test_cache_projects(db: SQLDB, db_session: Session):
assert 0 == len(db._projects), 'empty cache'
name = 'prj348'
db.add_project(db_session, {'name': name})
assert {name} == db._projects, 'project'
mock = Mock()
with patch(db, add_project=mock):
db._create_project_if_not_exists(db_session, name)
mock.assert_not_called()
mock = Mock()
with patch(db, add_project=mock):
db._create_project_if_not_exists(db_session, name + '-new')
mock.assert_called_once()
project_2_name = "project-2"
db.add_project(db_session, {'name': project_2_name})
db._projects = set()
mock = Mock()
with patch(db, add_project=mock):
db._create_project_if_not_exists(db_session, name)
mock.assert_not_called()
def test_run_iter0(db: SQLDB, db_session: Session):
uid, prj = 'uid39', 'lemon'
run = new_run('s1', {'l1': 'v1', 'l2': 'v2'}, x=1)
for i in range(7):
db.store_run(db_session, run, uid, prj, i)
db._get_run(db_session, uid, prj, 0) # See issue 140
