def test_list_runs_name_filter(db: DBInterface, db_session: Session):
run_name_1 = "run_name_1"
run_name_2 = "run_name_2"
run_1 = {"metadata": {"name": run_name_1}, "status": {"bla": "blabla"}}
run_2 = {"metadata": {"name": run_name_2}, "status": {"bla": "blabla"}}
# run with no name
run_3 = {"metadata": {}, "status": {"bla": "blabla"}}
run_uid_1 = "run_uid_1"
run_uid_2 = "run_uid_2"
run_uid_3 = "run_uid_3"
db.store_run(
db_session, run_1, run_uid_1,
)
db.store_run(
db_session, run_2, run_uid_2,
)
db.store_run(
db_session, run_3, run_uid_3,
)
runs = db.list_runs(db_session)
assert len(runs) == 3
runs = db.list_runs(db_session, name=run_name_1)
assert len(runs) == 1
assert runs[0]["metadata"]["name"] == run_name_1
runs = db.list_runs(db_session, name=run_name_2)
assert len(runs) == 1
assert runs[0]["metadata"]["name"] == run_name_2
runs = db.list_runs(db_session, name="run_name")
assert len(runs) == 2
def _update_ui_url(
self,
db: DBInterface,
db_session: Session,
project: str,
uid: str,
crd_object,
run: Dict = None,
leader_session: Optional[str] = None,
):
app_state = (crd_object.get("status", {}).get("applicationState",
{}).get("state"))
state = SparkApplicationStates.spark_application_state_to_run_state(
app_state)
ui_url = None
if state == RunStates.running:
ui_url = (crd_object.get("status",
{}).get("driverInfo",
{}).get("webUIIngressAddress"))
db_ui_url = run.get("status", {}).get("ui_url")
if db_ui_url == ui_url:
return
run.setdefault("status", {})["ui_url"] = ui_url
db.store_run(db_session,
run,
uid,
project,
leader_session=leader_session)
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_store_artifact_tagging(db: DBInterface, db_session: Session):
artifact_1_key = "artifact_key_1"
artifact_1_body = _generate_artifact(artifact_1_key)
artifact_1_kind = ChartArtifact.kind
artifact_1_with_kind_body = _generate_artifact(artifact_1_key,
kind=artifact_1_kind)
artifact_1_uid = "artifact_uid"
artifact_1_with_kind_uid = "artifact_uid_2"
db.store_artifact(
db_session,
artifact_1_key,
artifact_1_body,
artifact_1_uid,
)
db.store_artifact(
db_session,
artifact_1_key,
artifact_1_with_kind_body,
artifact_1_with_kind_uid,
)
artifact = db.read_artifact(db_session, artifact_1_key, tag="latest")
assert artifact["kind"] == artifact_1_kind
artifact = db.read_artifact(db_session, artifact_1_key, tag=artifact_1_uid)
assert artifact.get("kind") is None
artifacts = db.list_artifacts(db_session, artifact_1_key, tag="latest")
assert len(artifacts) == 1
artifacts = db.list_artifacts(db_session,
artifact_1_key,
tag=artifact_1_uid)
assert len(artifacts) == 1
def test_list_artifact_name_filter(db: DBInterface, db_session: Session):
artifact_name_1 = "artifact_name_1"
artifact_name_2 = "artifact_name_2"
artifact_1 = _generate_artifact(artifact_name_1)
artifact_2 = _generate_artifact(artifact_name_2)
uid = "artifact_uid"
db.store_artifact(
db_session,
artifact_name_1,
artifact_1,
uid,
)
db.store_artifact(
db_session,
artifact_name_2,
artifact_2,
uid,
)
artifacts = db.list_artifacts(db_session)
assert len(artifacts) == 2
artifacts = db.list_artifacts(db_session, name=artifact_name_1)
assert len(artifacts) == 1
assert artifacts[0]["metadata"]["name"] == artifact_name_1
artifacts = db.list_artifacts(db_session, name=artifact_name_2)
assert len(artifacts) == 1
assert artifacts[0]["metadata"]["name"] == artifact_name_2
artifacts = db.list_artifacts(db_session, name="artifact_name")
assert len(artifacts) == 2
def _is_runtime_resource_run_in_transient_state(
self, db: DBInterface, db_session: Session,
runtime_resource: Dict) -> bool:
"""
A runtime can have different underlying resources (like pods or CRDs) - to generalize we call it runtime
resource. This function will verify whether the Run object related to this runtime resource is in transient
state. This is useful in order to determine whether an object can be removed. for example, a kubejob's pod
might be in completed state, but we would like to verify that the run is completed as well to verify the logs
were collected before we're removing the pod.
"""
project, uid = self._resolve_runtime_resource_run(runtime_resource)
# if no uid, assume in stable state
if not uid:
return False
run = db.read_run(db_session, uid, project)
if run.get('status',
{}).get('state') not in FunctionStates.stable_phases():
return True
# give some grace period
now = datetime.now(timezone.utc)
last_update_str = run.get('status', {}).get('last_update', now)
last_update = datetime.fromisoformat(last_update_str)
if last_update + timedelta(seconds=float(
config.runtime_resources_deletion_grace_period)) > now:
return True
return False
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_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_store_artifact_restoring_multiple_tags(db: DBInterface, db_session: Session):
artifact_key = "artifact_key_1"
artifact_1_uid = "artifact_uid_1"
artifact_2_uid = "artifact_uid_2"
artifact_1_body = _generate_artifact(artifact_key, uid=artifact_1_uid)
artifact_2_body = _generate_artifact(artifact_key, uid=artifact_2_uid)
artifact_1_tag = "artifact_tag_1"
artifact_2_tag = "artifact_tag_2"
db.store_artifact(
db_session, artifact_key, artifact_1_body, artifact_1_uid, tag=artifact_1_tag,
)
db.store_artifact(
db_session, artifact_key, artifact_2_body, artifact_2_uid, tag=artifact_2_tag,
)
artifacts = db.list_artifacts(db_session, artifact_key, tag="*")
assert len(artifacts) == 2
expected_uids = [artifact_1_uid, artifact_2_uid]
uids = [artifact["metadata"]["uid"] for artifact in artifacts]
assert deepdiff.DeepDiff(expected_uids, uids, ignore_order=True,) == {}
expected_tags = [artifact_1_tag, artifact_2_tag]
tags = [artifact["tag"] for artifact in artifacts]
assert deepdiff.DeepDiff(expected_tags, tags, ignore_order=True,) == {}
artifact = db.read_artifact(db_session, artifact_key, tag=artifact_1_tag)
assert artifact["metadata"]["uid"] == artifact_1_uid
assert artifact["tag"] == artifact_1_tag
artifact = db.read_artifact(db_session, artifact_key, tag=artifact_2_tag)
assert artifact["metadata"]["uid"] == artifact_2_uid
assert artifact["tag"] == artifact_2_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_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_delete_project(db: DBInterface, db_session: Session):
project_to_keep = "project_to_keep"
project_to_remove = "project_to_remove"
_create_resources_of_all_kinds(db, db_session, project_to_keep)
_create_resources_of_all_kinds(db, db_session, project_to_remove)
project_to_keep_table_name_records_count_map_before_project_removal = _assert_resources_in_project(
db, db_session, project_to_keep)
_assert_resources_in_project(db, db_session, project_to_remove)
db.delete_project(db_session, project_to_remove)
project_to_keep_table_name_records_count_map_after_project_removal = _assert_resources_in_project(
db, db_session, project_to_keep)
_assert_resources_in_project(db,
db_session,
project_to_remove,
assert_no_resources=True)
assert (deepdiff.DeepDiff(
project_to_keep_table_name_records_count_map_before_project_removal,
project_to_keep_table_name_records_count_map_after_project_removal,
ignore_order=True,
) == {})
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_list_artifact_iter_parameter(db: DBInterface, db_session: Session):
artifact_name_1 = "artifact_name_1"
artifact_name_2 = "artifact_name_2"
artifact_1 = _generate_artifact(artifact_name_1)
artifact_2 = _generate_artifact(artifact_name_2)
uid = "artifact_uid"
# Use iters with multiple digits, to make sure filtering them via regex works
test_iters = [0, 5, 9, 42, 219, 2102]
for iter in test_iters:
artifact_1["iter"] = artifact_2["iter"] = iter
db.store_artifact(db_session, artifact_name_1, artifact_1, uid, iter)
db.store_artifact(db_session, artifact_name_2, artifact_2, uid, iter)
# No filter on iter. All are expected
artifacts = db.list_artifacts(db_session)
assert len(artifacts) == len(test_iters) * 2
# Look for the various iteration numbers. Note that 0 is a special case due to the DB structure
for iter in test_iters:
artifacts = db.list_artifacts(db_session, iter=iter)
assert len(artifacts) == 2
for artifact in artifacts:
assert artifact["iter"] == iter
# Negative test
artifacts = db.list_artifacts(db_session, iter=666)
assert len(artifacts) == 0
# Iter filter and a name filter, make sure query composition works
artifacts = db.list_artifacts(db_session, name=artifact_name_1, iter=2102)
assert len(artifacts) == 1
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_list_artifacts_exact_name_match(db: DBInterface, db_session: Session):
artifact_1_key = "pre_artifact_key_suffix"
artifact_2_key = "pre-artifact-key-suffix"
artifact_1_uid = "artifact_uid_1"
artifact_2_uid = "artifact_uid_2"
artifact_1_body = _generate_artifact(artifact_1_key, uid=artifact_1_uid)
artifact_2_body = _generate_artifact(artifact_2_key, uid=artifact_2_uid)
# Store each twice - once with no iter, and once with an iter
db.store_artifact(
db_session, artifact_1_key, artifact_1_body, artifact_1_uid,
)
artifact_1_body["iter"] = 42
db.store_artifact(
db_session, artifact_1_key, artifact_1_body, artifact_1_uid, iter=42,
)
db.store_artifact(
db_session, artifact_2_key, artifact_2_body, artifact_2_uid,
)
artifact_2_body["iter"] = 42
db.store_artifact(
db_session, artifact_2_key, artifact_2_body, artifact_2_uid, iter=42,
)
def _list_and_assert_count(key, count, iter=None):
results = db.list_artifacts(db_session, name=key, iter=iter)
assert len(results) == count
return results
# Ensure fuzzy query works, and we have everything we need
_list_and_assert_count("~key", count=4)
# Do an exact match with underscores in the name - must escape the _ do it doesn't do a like query
list_results = _list_and_assert_count(artifact_1_key, count=2)
for artifact in list_results:
assert artifact["metadata"]["name"] == artifact_1_key
_list_and_assert_count("%key%", count=0)
# Verify we don't get artifacts whose name is "%-suffix" due to the like query used in the DB
_list_and_assert_count("suffix", count=0)
# This should also be filtered, since the prefix is "pre" which is 3 chars. There's a known caveat if
# prefix is 1 or 2 chars long.
_list_and_assert_count("artifact-key-suffix", count=0)
_list_and_assert_count(artifact_1_key, iter=42, count=1)
_list_and_assert_count("~key", iter=42, count=2)
_list_and_assert_count("~key", iter=666, count=0)
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_data_migration_fill_project_state(
db: DBInterface, db_session: sqlalchemy.orm.Session,
):
for i in range(10):
project_name = f"project-name-{i}"
_generate_and_insert_pre_060_record(db_session, project_name)
projects = db.list_projects(db_session)
for project in projects.projects:
# getting default value from the schema
assert project.spec.desired_state == mlrun.api.schemas.ProjectState.online
assert project.status.state is None
mlrun.api.initial_data._fill_project_state(db, db_session)
projects = db.list_projects(db_session)
for project in projects.projects:
assert project.spec.desired_state == mlrun.api.schemas.ProjectState.online
assert project.status.state == project.spec.desired_state
# verify not storing for no reason
db.store_project = unittest.mock.Mock()
mlrun.api.initial_data._fill_project_state(db, db_session)
assert db.store_project.call_count == 0
def test_create_feature_set(db: DBInterface, db_session: Session):
name = "dummy"
feature_set = _create_feature_set(name)
project = "proj_test"
feature_set = schemas.FeatureSet(**feature_set)
db.create_feature_set(db_session, project, feature_set, versioned=True)
db.get_feature_set(db_session, project, name)
feature_set_res = db.list_feature_sets(db_session, project)
assert len(feature_set_res.feature_sets) == 1
features_res = db.list_features(db_session, project, "time")
assert len(features_res.features) == 1
def _ensure_runtime_resource_run_logs_collected(self, db: DBInterface,
db_session: Session,
runtime_resource: Dict):
project, uid = self._resolve_runtime_resource_run(runtime_resource)
# if cannot resolve related run, assume collected
if not uid:
return
# import here to avoid circular imports
import mlrun.api.crud as crud
log_file_exists = crud.Logs.log_file_exists(project, uid)
store_log = False
if not log_file_exists:
store_log = True
else:
log_mtime = crud.Logs.get_log_mtime(project, uid)
log_mtime_datetime = datetime.fromtimestamp(
log_mtime, timezone.utc)
now = datetime.now(timezone.utc)
run = db.read_run(db_session, uid, project)
last_update_str = run.get('status', {}).get('last_update', now)
last_update = datetime.fromisoformat(last_update_str)
# this function is used to verify that logs collected from runtime resources before deleting them
# here we're using the knowledge that the function is called only after a it was verified that the runtime
# resource run is not in transient state, so we're assuming the run's last update is the last one, so if the
# log file was modified after it, we're considering it as all logs collected
if log_mtime_datetime < last_update:
store_log = True
if store_log:
logger.debug(f'Storing runtime resource log before deletion')
logs_from_k8s, _ = crud.Logs.get_log(db_session,
project,
uid,
source=LogSources.K8S)
crud.Logs.store_log(logs_from_k8s, project, uid, append=False)
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_list_artifact_category_filter(db: DBInterface, db_session: Session):
artifact_name_1 = "artifact_name_1"
artifact_kind_1 = ChartArtifact.kind
artifact_name_2 = "artifact_name_2"
artifact_kind_2 = PlotArtifact.kind
artifact_name_3 = "artifact_name_3"
artifact_kind_3 = ModelArtifact.kind
artifact_name_4 = "artifact_name_4"
artifact_kind_4 = DatasetArtifact.kind
artifact_1 = _generate_artifact(artifact_name_1, kind=artifact_kind_1)
artifact_2 = _generate_artifact(artifact_name_2, kind=artifact_kind_2)
artifact_3 = _generate_artifact(artifact_name_3, kind=artifact_kind_3)
artifact_4 = _generate_artifact(artifact_name_4, kind=artifact_kind_4)
uid = "artifact_uid"
db.store_artifact(
db_session,
artifact_name_1,
artifact_1,
uid,
)
db.store_artifact(
db_session,
artifact_name_2,
artifact_2,
uid,
)
db.store_artifact(
db_session,
artifact_name_3,
artifact_3,
uid,
)
db.store_artifact(
db_session,
artifact_name_4,
artifact_4,
uid,
)
artifacts = db.list_artifacts(db_session)
assert len(artifacts) == 4
artifacts = db.list_artifacts(db_session,
category=schemas.ArtifactCategories.model)
assert len(artifacts) == 1
assert artifacts[0]["metadata"]["name"] == artifact_name_3
artifacts = db.list_artifacts(db_session,
category=schemas.ArtifactCategories.dataset)
assert len(artifacts) == 1
assert artifacts[0]["metadata"]["name"] == artifact_name_4
artifacts = db.list_artifacts(db_session,
category=schemas.ArtifactCategories.other)
assert len(artifacts) == 2
assert artifacts[0]["metadata"]["name"] == artifact_name_1
assert artifacts[1]["metadata"]["name"] == artifact_name_2
