def test_list_functions_no_tags(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}}
function_2 = {'bla2': 'blabla', 'status': {'bla': 'blabla'}}
function_name_1 = 'function_name_1'
# It is impossible to create a function without tag - only to create with a tag, and then tag another function with
# the same tag
tag = 'some_tag'
function_1_hash_key = db.store_function(db_session,
function_1,
function_name_1,
tag=tag,
versioned=True)
function_2_hash_key = db.store_function(db_session,
function_2,
function_name_1,
tag=tag,
versioned=True)
assert function_1_hash_key != function_2_hash_key
functions = db.list_functions(db_session, function_name_1)
assert len(functions) == 2
# Verify function 1 without tag and has not status
for function in functions:
if function['metadata']['hash'] == function_1_hash_key:
assert function['metadata']['tag'] == ''
assert function['status'] is None
def test_list_functions_no_tags(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla", "status": {"bla": "blabla"}}
function_2 = {"bla2": "blabla", "status": {"bla": "blabla"}}
function_name_1 = "function_name_1"
# It is impossible to create a function without tag - only to create with a tag, and then tag another function with
# the same tag
tag = "some_tag"
function_1_hash_key = db.store_function(db_session,
function_1,
function_name_1,
tag=tag,
versioned=True)
function_2_hash_key = db.store_function(db_session,
function_2,
function_name_1,
tag=tag,
versioned=True)
assert function_1_hash_key != function_2_hash_key
functions = db.list_functions(db_session, function_name_1)
assert len(functions) == 2
# Verify function 1 without tag and has not status
for function in functions:
if function["metadata"]["hash"] == function_1_hash_key:
assert function["metadata"]["tag"] == ""
assert function["status"] is None
def test_get_function_not_found(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla", "status": {"bla": "blabla"}}
function_name_1 = "function_name_1"
db.store_function(db_session, function_1, function_name_1, versioned=True)
with pytest.raises(mlrun.errors.MLRunNotFoundError):
db.get_function(db_session, function_name_1, tag="inexistent_tag")
with pytest.raises(mlrun.errors.MLRunNotFoundError):
db.get_function(db_session,
function_name_1,
hash_key="inexistent_hash_key")
def test_list_functions_multiple_tags(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}}
function_name_1 = 'function_name_1'
tags = ['some_tag', 'some_tag2', 'some_tag3']
for tag in tags:
db.store_function(db_session,
function_1,
function_name_1,
tag=tag,
versioned=True)
functions = db.list_functions(db_session, function_name_1)
assert len(functions) == len(tags)
for function in functions:
function_tag = function['metadata']['tag']
tags.remove(function_tag)
assert len(tags) == 0
def test_list_functions_by_tag(db: DBInterface, db_session: Session):
tag = "function_name_1"
names = ["some_name", "some_name2", "some_name3"]
for name in names:
function_body = {"metadata": {"name": name}}
db.store_function(db_session,
function_body,
name,
tag=tag,
versioned=True)
functions = db.list_functions(db_session, tag=tag)
assert len(functions) == len(names)
for function in functions:
function_name = function["metadata"]["name"]
names.remove(function_name)
assert len(names) == 0
def test_list_functions_multiple_tags(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla", "status": {"bla": "blabla"}}
function_name_1 = "function_name_1"
tags = ["some_tag", "some_tag2", "some_tag3"]
for tag in tags:
db.store_function(db_session,
function_1,
function_name_1,
tag=tag,
versioned=True)
functions = db.list_functions(db_session, function_name_1)
assert len(functions) == len(tags)
for function in functions:
function_tag = function["metadata"]["tag"]
tags.remove(function_tag)
assert len(tags) == 0
def test_list_functions_filtering_unversioned_untagged(db: DBInterface,
db_session: Session):
function_1 = {"bla": "blabla"}
function_2 = {"bla": "blablablabla"}
function_name_1 = "function_name_1"
tag = "some_tag"
db.store_function(db_session,
function_1,
function_name_1,
versioned=False,
tag=tag)
tagged_function_hash_key = db.store_function(db_session,
function_2,
function_name_1,
versioned=True,
tag=tag)
functions = db.list_functions(db_session, function_name_1)
# First we stored to the tag without versioning (unversioned instance) then we stored to the tag with version
# so the unversioned instance remained untagged, verifying we're not getting it
assert len(functions) == 1
assert functions[0]["metadata"]["hash"] == tagged_function_hash_key
def test_store_function_not_versioned(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla"}
function_name_1 = "function_name_1"
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=False)
function_result_1 = db.get_function(db_session, function_name_1)
assert function_result_1 is not None
assert function_result_1["metadata"]["tag"] == "latest"
# not versioned so not queryable by hash key
with pytest.raises(mlrun.errors.MLRunNotFoundError):
db.get_function(db_session,
function_name_1,
hash_key=function_hash_key)
function_2 = {"bla": "blabla", "bla2": "blabla2"}
db.store_function(db_session, function_2, function_name_1, versioned=False)
functions = db.list_functions(db_session, function_name_1)
# Verifying only the latest version was saved
assert len(functions) == 1
def test_store_function_not_versioned(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla'}
function_name_1 = 'function_name_1'
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=False)
function_result_1 = db.get_function(db_session, function_name_1)
assert function_result_1 is not None
assert function_result_1['metadata']['tag'] == 'latest'
# not versioned so not queryable by hash key
function_result_2 = db.get_function(db_session,
function_name_1,
hash_key=function_hash_key)
assert function_result_2 is None
function_2 = {'bla': 'blabla', 'bla2': 'blabla2'}
db.store_function(db_session, function_2, function_name_1, versioned=False)
functions = db.list_functions(db_session, function_name_1)
# Verifying only the latest version was saved
assert len(functions) == 1
def test_store_function_versioned(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla'}
function_name_1 = 'function_name_1'
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_without_hash_key = db.get_function(
db_session, function_name_1)
assert function_queried_without_hash_key is not None
assert function_queried_without_hash_key['metadata']['tag'] == 'latest'
# Verifying versioned function is queryable by hash_key
function_queried_with_hash_key = db.get_function(
db_session, function_name_1, hash_key=function_hash_key)
function_queried_with_hash_key_hash = function_queried_with_hash_key[
'metadata']['hash']
assert function_queried_with_hash_key is not None
assert function_queried_with_hash_key['metadata']['tag'] == ''
assert function_queried_with_hash_key_hash == function_hash_key
function_2 = {'bla': 'blabla', 'bla2': 'blabla2'}
function_name_1 = 'function_name_1'
db.store_function(db_session, function_2, function_name_1, versioned=True)
functions = db.list_functions(db_session, function_name_1)
# Verifying both versions of the functions were saved
assert len(functions) == 2
tagged_count = 0
for function in functions:
if function['metadata']['tag'] == 'latest':
tagged_count += 1
# but only one was tagged
assert tagged_count == 1
def test_store_function_versioned(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla"}
function_name_1 = "function_name_1"
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_without_hash_key = db.get_function(
db_session, function_name_1)
assert function_queried_without_hash_key is not None
assert function_queried_without_hash_key["metadata"]["tag"] == "latest"
# Verifying versioned function is queryable by hash_key
function_queried_with_hash_key = db.get_function(
db_session, function_name_1, hash_key=function_hash_key)
function_queried_with_hash_key_hash = function_queried_with_hash_key[
"metadata"]["hash"]
assert function_queried_with_hash_key is not None
assert function_queried_with_hash_key["metadata"]["tag"] == ""
assert function_queried_with_hash_key_hash == function_hash_key
function_2 = {"bla": "blabla", "bla2": "blabla2"}
function_name_1 = "function_name_1"
db.store_function(db_session, function_2, function_name_1, versioned=True)
functions = db.list_functions(db_session, function_name_1)
# Verifying both versions of the functions were saved
assert len(functions) == 2
tagged_count = 0
for function in functions:
if function["metadata"]["tag"] == "latest":
tagged_count += 1
# but only one was tagged
assert tagged_count == 1
def test_get_function_by_hash_key(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}}
function_name_1 = 'function_name_1'
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_without_hash_key = db.get_function(
db_session, function_name_1)
assert function_queried_without_hash_key is not None
# Verifying function is queryable by hash_key
function_queried_with_hash_key = db.get_function(
db_session, function_name_1, hash_key=function_hash_key)
assert function_queried_with_hash_key is not None
# function queried by hash shouldn't have tag
assert function_queried_without_hash_key['metadata']['tag'] == 'latest'
assert function_queried_with_hash_key['metadata']['tag'] == ''
def test_get_function_by_hash_key(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla", "status": {"bla": "blabla"}}
function_name_1 = "function_name_1"
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_without_hash_key = db.get_function(
db_session, function_name_1)
assert function_queried_without_hash_key is not None
# Verifying function is queryable by hash_key
function_queried_with_hash_key = db.get_function(
db_session, function_name_1, hash_key=function_hash_key)
assert function_queried_with_hash_key is not None
# function queried by hash shouldn't have tag
assert function_queried_without_hash_key["metadata"]["tag"] == "latest"
assert function_queried_with_hash_key["metadata"]["tag"] == ""
def test_store_function_default_to_latest(db: DBInterface,
db_session: Session):
function_1 = {'bla': 'blabla'}
function_name_1 = 'function_name_1'
function_hash_key = db.store_function(db_session, function_1,
function_name_1)
assert function_hash_key is not None
function_queried_without_tag = db.get_function(db_session, function_name_1)
function_queried_without_tag_hash = function_queried_without_tag[
'metadata']['hash']
assert function_hash_key == function_queried_without_tag_hash
assert function_queried_without_tag['metadata']['tag'] == 'latest'
function_queried_with_tag = db.get_function(db_session,
function_name_1,
tag='latest')
function_queried_without_tag_hash = function_queried_with_tag['metadata'][
'hash']
assert function_queried_with_tag is not None
assert function_queried_with_tag['metadata']['tag'] == 'latest'
assert function_queried_without_tag_hash == function_queried_without_tag_hash
def test_get_function_by_tag(db: DBInterface, db_session: Session):
function_1 = {"bla": "blabla", "status": {"bla": "blabla"}}
function_name_1 = "function_name_1"
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_by_tag = db.get_function(db_session,
function_name_1,
tag="latest")
function_queried_by_hash_key = db.get_function(db_session,
function_name_1,
hash_key=function_hash_key)
function_not_queried_by_tag_hash = function_queried_by_hash_key[
"metadata"]["hash"]
assert function_hash_key == function_not_queried_by_tag_hash
# function not queried by tag shouldn't have status
assert function_queried_by_tag["status"] is not None
assert function_queried_by_hash_key["status"] is None
def test_store_function_default_to_latest(db: DBInterface,
db_session: Session):
function_1 = {"bla": "blabla"}
function_name_1 = "function_name_1"
function_hash_key = db.store_function(db_session, function_1,
function_name_1)
assert function_hash_key is not None
function_queried_without_tag = db.get_function(db_session, function_name_1)
function_queried_without_tag_hash = function_queried_without_tag[
"metadata"]["hash"]
assert function_hash_key == function_queried_without_tag_hash
assert function_queried_without_tag["metadata"]["tag"] == "latest"
function_queried_with_tag = db.get_function(db_session,
function_name_1,
tag="latest")
function_queried_without_tag_hash = function_queried_with_tag["metadata"][
"hash"]
assert function_queried_with_tag is not None
assert function_queried_with_tag["metadata"]["tag"] == "latest"
assert function_queried_without_tag_hash == function_queried_without_tag_hash
def test_get_function_by_tag(db: DBInterface, db_session: Session):
function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}}
function_name_1 = 'function_name_1'
function_hash_key = db.store_function(db_session,
function_1,
function_name_1,
versioned=True)
function_queried_by_tag = db.get_function(db_session,
function_name_1,
tag='latest')
function_queried_by_hash_key = db.get_function(db_session,
function_name_1,
hash_key=function_hash_key)
function_not_queried_by_tag_hash = function_queried_by_hash_key[
'metadata']['hash']
assert function_hash_key == function_not_queried_by_tag_hash
# function not queried by tag shouldn't have status
assert function_queried_by_tag['status'] is not None
assert function_queried_by_hash_key['status'] is None
def test_delete_function(db: DBInterface, db_session: Session):
labels = {
"name": "value",
"name2": "value2",
}
function = {
"bla": "blabla",
"metadata": {
"labels": labels
},
"status": {
"bla": "blabla"
},
}
function_name = "function_name_1"
project = "bla"
tags = ["some_tag", "some_tag2", "some_tag3"]
function_hash_key = None
for tag in tags:
function_hash_key = db.store_function(db_session,
function,
function_name,
project,
tag=tag,
versioned=True)
# if not exploding then function exists
for tag in tags:
db.get_function(db_session, function_name, project, tag=tag)
db.get_function(db_session,
function_name,
project,
hash_key=function_hash_key)
assert len(tags) == len(
db.list_functions(db_session, function_name, project))
number_of_tags = (db_session.query(Function.Tag).filter_by(
project=project, obj_name=function_name).count())
number_of_labels = db_session.query(Function.Label).count()
assert len(tags) == number_of_tags
assert len(labels) == number_of_labels
db.delete_function(db_session, project, function_name)
for tag in tags:
with pytest.raises(mlrun.errors.MLRunNotFoundError):
db.get_function(db_session, function_name, project, tag=tag)
with pytest.raises(mlrun.errors.MLRunNotFoundError):
db.get_function(db_session,
function_name,
project,
hash_key=function_hash_key)
assert 0 == len(db.list_functions(db_session, function_name, project))
# verifying tags and labels (different table) records were removed
number_of_tags = (db_session.query(Function.Tag).filter_by(
project=project, obj_name=function_name).count())
number_of_labels = db_session.query(Function.Label).count()
assert number_of_tags == 0
assert number_of_labels == 0
def _create_resources_of_all_kinds(db: DBInterface, db_session: Session,
project: str):
# Create several functions with several tags
labels = {
"name": "value",
"name2": "value2",
}
function = {
"bla": "blabla",
"metadata": {
"labels": labels
},
"status": {
"bla": "blabla"
},
}
function_names = ["function_name_1", "function_name_2", "function_name_3"]
function_tags = ["some_tag", "some_tag2", "some_tag3"]
for function_name in function_names:
for function_tag in function_tags:
db.store_function(
db_session,
function,
function_name,
project,
tag=function_tag,
versioned=True,
)
# Create several artifacts with several tags
artifact = {
"bla": "blabla",
"labels": labels,
"status": {
"bla": "blabla"
},
}
artifact_keys = ["artifact_key_1", "artifact_key_2", "artifact_key_3"]
artifact_uids = ["some_uid", "some_uid2", "some_uid3"]
artifact_tags = ["some_tag", "some_tag2", "some_tag3"]
for artifact_key in artifact_keys:
for artifact_uid in artifact_uids:
for artifact_tag in artifact_tags:
for artifact_iter in range(3):
db.store_artifact(
db_session,
artifact_key,
artifact,
artifact_uid,
artifact_iter,
artifact_tag,
project,
)
# Create several runs
run = {
"bla": "blabla",
"metadata": {
"labels": labels
},
"status": {
"bla": "blabla"
},
}
run_uids = ["some_uid", "some_uid2", "some_uid3"]
for run_uid in run_uids:
for run_iter in range(3):
db.store_run(db_session, run, run_uid, project, run_iter)
# Create several logs
log = b"some random log"
log_uids = ["some_uid", "some_uid2", "some_uid3"]
for log_uid in log_uids:
db.store_log(db_session, log_uid, project, log)
# Create several schedule
schedule = {
"bla": "blabla",
"status": {
"bla": "blabla"
},
}
schedule_cron_trigger = schemas.ScheduleCronTrigger(year=1999)
schedule_names = ["schedule_name_1", "schedule_name_2", "schedule_name_3"]
for schedule_name in schedule_names:
db.create_schedule(
db_session,
project,
schedule_name,
schemas.ScheduleKinds.job,
schedule,
schedule_cron_trigger,
labels,
)
feature_set = schemas.FeatureSet(
metadata=schemas.ObjectMetadata(name="dummy",
tag="latest",
labels={"owner": "nobody"}),
spec=schemas.FeatureSetSpec(
entities=[
schemas.Entity(name="ent1",
value_type="str",
labels={"label": "1"})
],
features=[
schemas.Feature(name="feat1",
value_type="str",
labels={"label": "1"})
],
),
status={},
)
db.create_feature_set(db_session, project, feature_set)
