wehub-resource-sync
1228 lines
36 KiB
Python
1228 lines
36 KiB
Python
import os
|
|
import re
|
|
from contextlib import contextmanager
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from alembic import command
|
|
|
|
from mlflow.store.db.utils import _get_alembic_config
|
|
from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
|
|
|
|
_LEGACY_REGISTERED_MODEL_TAGS = sa.table(
|
|
"registered_model_tags",
|
|
sa.column("key"),
|
|
sa.column("value"),
|
|
sa.column("name"),
|
|
)
|
|
|
|
_LEGACY_MODEL_VERSION_TAGS = sa.table(
|
|
"model_version_tags",
|
|
sa.column("key"),
|
|
sa.column("value"),
|
|
sa.column("name"),
|
|
sa.column("version"),
|
|
)
|
|
|
|
_LEGACY_REGISTERED_MODEL_ALIASES = sa.table(
|
|
"registered_model_aliases",
|
|
sa.column("alias"),
|
|
sa.column("version"),
|
|
sa.column("name"),
|
|
)
|
|
|
|
_LEGACY_EVALUATION_DATASETS = sa.table(
|
|
"evaluation_datasets",
|
|
sa.column("dataset_id"),
|
|
sa.column("name"),
|
|
sa.column("schema"),
|
|
sa.column("profile"),
|
|
sa.column("digest"),
|
|
sa.column("created_time"),
|
|
sa.column("last_update_time"),
|
|
sa.column("created_by"),
|
|
sa.column("last_updated_by"),
|
|
)
|
|
|
|
_LEGACY_SECRETS = sa.table(
|
|
"secrets",
|
|
sa.column("secret_id"),
|
|
sa.column("secret_name"),
|
|
sa.column("encrypted_value"),
|
|
sa.column("wrapped_dek"),
|
|
sa.column("kek_version"),
|
|
sa.column("masked_value"),
|
|
sa.column("provider"),
|
|
sa.column("auth_config"),
|
|
sa.column("description"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
)
|
|
|
|
_LEGACY_ENDPOINTS = sa.table(
|
|
"endpoints",
|
|
sa.column("endpoint_id"),
|
|
sa.column("name"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
)
|
|
|
|
_LEGACY_MODEL_DEFINITIONS = sa.table(
|
|
"model_definitions",
|
|
sa.column("model_definition_id"),
|
|
sa.column("name"),
|
|
sa.column("secret_id"),
|
|
sa.column("provider"),
|
|
sa.column("model_name"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
)
|
|
|
|
_WORKSPACE_TABLES = (
|
|
"experiments",
|
|
"registered_models",
|
|
"model_versions",
|
|
"registered_model_tags",
|
|
"model_version_tags",
|
|
"registered_model_aliases",
|
|
"evaluation_datasets",
|
|
"secrets",
|
|
"endpoints",
|
|
"model_definitions",
|
|
"webhooks",
|
|
"jobs",
|
|
)
|
|
|
|
_REGISTERED_MODEL_TAGS = sa.table(
|
|
"registered_model_tags",
|
|
sa.column("workspace"),
|
|
sa.column("key"),
|
|
sa.column("value"),
|
|
sa.column("name"),
|
|
)
|
|
|
|
_MODEL_VERSION_TAGS = sa.table(
|
|
"model_version_tags",
|
|
sa.column("workspace"),
|
|
sa.column("key"),
|
|
sa.column("value"),
|
|
sa.column("name"),
|
|
sa.column("version"),
|
|
)
|
|
|
|
_REGISTERED_MODEL_ALIASES = sa.table(
|
|
"registered_model_aliases",
|
|
sa.column("workspace"),
|
|
sa.column("name"),
|
|
sa.column("alias"),
|
|
sa.column("version"),
|
|
)
|
|
|
|
_EVALUATION_DATASETS = sa.table(
|
|
"evaluation_datasets",
|
|
sa.column("dataset_id"),
|
|
sa.column("name"),
|
|
sa.column("schema"),
|
|
sa.column("profile"),
|
|
sa.column("digest"),
|
|
sa.column("created_time"),
|
|
sa.column("last_update_time"),
|
|
sa.column("created_by"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("workspace"),
|
|
)
|
|
|
|
REVISION = "1b5f0d9ad7c1"
|
|
PREVIOUS_REVISION = "c8d9e0f1a2b3"
|
|
|
|
DB_URI = os.environ.get("MLFLOW_TRACKING_URI")
|
|
USE_EXTERNAL_DB = DB_URI is not None and not DB_URI.startswith("sqlite")
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def _upgrade_external_db_to_head_after_suite():
|
|
"""
|
|
When running under Docker (i.e., with a shared external DB), make sure the DB ends up on the
|
|
latest revision once this module finishes. Individual tests intentionally downgrade to the
|
|
pre-workspace schema, so without this hook, the subsequent suites in the database workflow would
|
|
run against an outdated schema after new migrations land.
|
|
"""
|
|
yield
|
|
if USE_EXTERNAL_DB:
|
|
config = _get_alembic_config(DB_URI)
|
|
command.upgrade(config, "head")
|
|
|
|
|
|
@contextmanager
|
|
def _identity_insert(conn, table_name: str):
|
|
if conn.dialect.name != "mssql":
|
|
yield
|
|
return
|
|
|
|
conn.execute(sa.text(f"SET IDENTITY_INSERT {table_name} ON"))
|
|
try:
|
|
yield
|
|
finally:
|
|
conn.execute(sa.text(f"SET IDENTITY_INSERT {table_name} OFF"))
|
|
|
|
|
|
def _insert_table_row(conn, table, **values):
|
|
conn.execute(sa.insert(table).values(**values))
|
|
|
|
|
|
def _assert_workspace_column(inspector, table_name: str, expected_default: str):
|
|
columns = inspector.get_columns(table_name)
|
|
workspace = next((col for col in columns if col["name"] == "workspace"), None)
|
|
assert workspace is not None, f"{table_name} lacks workspace column"
|
|
assert not workspace.get("nullable", False)
|
|
default_value = _get_workspace_default(workspace)
|
|
assert default_value == expected_default
|
|
|
|
|
|
def _assert_workspace_columns(inspector, expected_default: str = "default"):
|
|
for table in _WORKSPACE_TABLES:
|
|
_assert_workspace_column(inspector, table, expected_default)
|
|
|
|
|
|
def _has_index(inspector, table: str, index_name: str, columns: list[str]):
|
|
indexes = inspector.get_indexes(table)
|
|
return any(
|
|
index["name"] == index_name and index.get("column_names") == columns for index in indexes
|
|
)
|
|
|
|
|
|
def _prepare_database(tmp_path):
|
|
if USE_EXTERNAL_DB:
|
|
engine = sa.create_engine(DB_URI)
|
|
with engine.begin() as conn:
|
|
metadata = sa.MetaData()
|
|
metadata.reflect(bind=conn)
|
|
metadata.drop_all(bind=conn)
|
|
InitialBase.metadata.create_all(conn)
|
|
config = _get_alembic_config(DB_URI)
|
|
else:
|
|
db_path = tmp_path / "workspace_migration.sqlite"
|
|
url = f"sqlite:///{db_path}"
|
|
engine = sa.create_engine(url)
|
|
InitialBase.metadata.create_all(engine)
|
|
config = _get_alembic_config(url)
|
|
command.upgrade(config, PREVIOUS_REVISION)
|
|
return engine, config
|
|
|
|
|
|
def _seed_pre_workspace_entities(conn):
|
|
# This intentionally uses raw SQL matching the legacy schema (no workspace columns)
|
|
# so the migration under test is fully responsible for adding/backfilling the new
|
|
# fields. The helper insert functions below operate on the post-migration schema
|
|
# and therefore cannot be reused here.
|
|
with _identity_insert(conn, "experiments"):
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO experiments (
|
|
experiment_id,
|
|
name,
|
|
artifact_location,
|
|
lifecycle_stage,
|
|
creation_time,
|
|
last_update_time
|
|
)
|
|
VALUES (
|
|
:experiment_id,
|
|
:name,
|
|
:artifact_location,
|
|
:lifecycle_stage,
|
|
:creation_time,
|
|
:last_update_time
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"experiment_id": 1,
|
|
"name": "exp-default",
|
|
"artifact_location": "path",
|
|
"lifecycle_stage": "active",
|
|
"creation_time": 0,
|
|
"last_update_time": 0,
|
|
},
|
|
)
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO runs (
|
|
run_uuid,
|
|
name,
|
|
source_type,
|
|
source_name,
|
|
entry_point_name,
|
|
user_id,
|
|
status,
|
|
start_time,
|
|
end_time,
|
|
source_version,
|
|
lifecycle_stage,
|
|
artifact_uri,
|
|
experiment_id
|
|
)
|
|
VALUES (
|
|
:run_uuid,
|
|
:name,
|
|
:source_type,
|
|
:source_name,
|
|
:entry_point_name,
|
|
:user_id,
|
|
:status,
|
|
:start_time,
|
|
:end_time,
|
|
:source_version,
|
|
:lifecycle_stage,
|
|
:artifact_uri,
|
|
:experiment_id
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"run_uuid": "run-default",
|
|
"name": "upgrade-validation-run",
|
|
"source_type": "LOCAL",
|
|
"source_name": "script.py",
|
|
"entry_point_name": "main",
|
|
"user_id": "user",
|
|
"status": "FINISHED",
|
|
"start_time": 0,
|
|
"end_time": 1,
|
|
"source_version": "abc123",
|
|
"lifecycle_stage": "active",
|
|
"artifact_uri": "path/artifacts",
|
|
"experiment_id": 1,
|
|
},
|
|
)
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO registered_models (name, creation_time, last_updated_time, description)
|
|
VALUES (:name, :creation_time, :last_updated_time, :description)
|
|
"""
|
|
),
|
|
{"name": "rm-default", "creation_time": 0, "last_updated_time": 0, "description": "desc"},
|
|
)
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO model_versions (
|
|
name,
|
|
version,
|
|
creation_time,
|
|
last_updated_time,
|
|
user_id,
|
|
current_stage,
|
|
description,
|
|
source,
|
|
run_id,
|
|
status,
|
|
status_message,
|
|
run_link,
|
|
storage_location
|
|
)
|
|
VALUES (
|
|
:name,
|
|
:version,
|
|
:creation_time,
|
|
:last_updated_time,
|
|
:user_id,
|
|
:current_stage,
|
|
:description,
|
|
:source,
|
|
:run_id,
|
|
:status,
|
|
:status_message,
|
|
:run_link,
|
|
:storage_location
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"name": "rm-default",
|
|
"version": 1,
|
|
"creation_time": 0,
|
|
"last_updated_time": 0,
|
|
"user_id": "user",
|
|
"current_stage": "None",
|
|
"description": "desc",
|
|
"source": "source",
|
|
"run_id": "run-id",
|
|
"status": "READY",
|
|
"status_message": "message",
|
|
"run_link": "link",
|
|
"storage_location": "location",
|
|
},
|
|
)
|
|
_insert_table_row(
|
|
conn,
|
|
_LEGACY_REGISTERED_MODEL_TAGS,
|
|
key="tag",
|
|
value="value",
|
|
name="rm-default",
|
|
)
|
|
_insert_table_row(
|
|
conn,
|
|
_LEGACY_MODEL_VERSION_TAGS,
|
|
key="tag",
|
|
value="value",
|
|
name="rm-default",
|
|
version=1,
|
|
)
|
|
_insert_table_row(
|
|
conn,
|
|
_LEGACY_REGISTERED_MODEL_ALIASES,
|
|
alias="alias",
|
|
version=1,
|
|
name="rm-default",
|
|
)
|
|
_insert_table_row(
|
|
conn,
|
|
_LEGACY_EVALUATION_DATASETS,
|
|
dataset_id="ds-default",
|
|
name="Dataset",
|
|
schema="schema",
|
|
profile="profile",
|
|
digest="digest",
|
|
created_time=0,
|
|
last_update_time=0,
|
|
created_by="user",
|
|
last_updated_by="user",
|
|
)
|
|
|
|
|
|
def _get_workspace_default(column_info):
|
|
default = column_info.get("default") or column_info.get("server_default")
|
|
if default is None:
|
|
return None
|
|
value = str(default).strip()
|
|
if value.startswith("(") and value.endswith(")"):
|
|
value = value[1:-1]
|
|
value = value.strip()
|
|
value = value.strip("'\"")
|
|
if "::" in value:
|
|
value = value.split("::", 1)[0]
|
|
return value.strip("'\"")
|
|
|
|
|
|
def _add_workspace(conn, name: str, description: str):
|
|
conn.execute(
|
|
sa.text("INSERT INTO workspaces (name, description) VALUES (:name, :description)"),
|
|
{"name": name, "description": description},
|
|
)
|
|
|
|
|
|
def _insert_experiment(
|
|
conn,
|
|
*,
|
|
experiment_id: int,
|
|
name: str,
|
|
workspace: str,
|
|
artifact_location: str = "path",
|
|
lifecycle_stage: str = "active",
|
|
):
|
|
with _identity_insert(conn, "experiments"):
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO experiments (
|
|
experiment_id,
|
|
name,
|
|
artifact_location,
|
|
lifecycle_stage,
|
|
creation_time,
|
|
last_update_time,
|
|
workspace
|
|
)
|
|
VALUES (
|
|
:experiment_id,
|
|
:name,
|
|
:artifact_location,
|
|
:lifecycle_stage,
|
|
:creation_time,
|
|
:last_update_time,
|
|
:workspace
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"experiment_id": experiment_id,
|
|
"name": name,
|
|
"artifact_location": artifact_location,
|
|
"lifecycle_stage": lifecycle_stage,
|
|
"creation_time": 0,
|
|
"last_update_time": 0,
|
|
"workspace": workspace,
|
|
},
|
|
)
|
|
|
|
|
|
def _insert_run(
|
|
conn,
|
|
*,
|
|
run_uuid: str,
|
|
experiment_id: int,
|
|
name: str = "run",
|
|
artifact_uri: str = "path/artifacts",
|
|
):
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO runs (
|
|
run_uuid,
|
|
name,
|
|
source_type,
|
|
source_name,
|
|
entry_point_name,
|
|
user_id,
|
|
status,
|
|
start_time,
|
|
end_time,
|
|
source_version,
|
|
lifecycle_stage,
|
|
artifact_uri,
|
|
experiment_id
|
|
)
|
|
VALUES (
|
|
:run_uuid,
|
|
:name,
|
|
:source_type,
|
|
:source_name,
|
|
:entry_point_name,
|
|
:user_id,
|
|
:status,
|
|
:start_time,
|
|
:end_time,
|
|
:source_version,
|
|
:lifecycle_stage,
|
|
:artifact_uri,
|
|
:experiment_id
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"run_uuid": run_uuid,
|
|
"name": name,
|
|
"source_type": "LOCAL",
|
|
"source_name": "script.py",
|
|
"entry_point_name": "main",
|
|
"user_id": "user",
|
|
"status": "FINISHED",
|
|
"start_time": 0,
|
|
"end_time": 1,
|
|
"source_version": "abc123",
|
|
"lifecycle_stage": "active",
|
|
"artifact_uri": artifact_uri,
|
|
"experiment_id": experiment_id,
|
|
},
|
|
)
|
|
|
|
|
|
def _insert_registered_model(
|
|
conn,
|
|
*,
|
|
name: str,
|
|
workspace: str,
|
|
description: str = "desc",
|
|
creation_time: int = 0,
|
|
):
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO registered_models (
|
|
name,
|
|
creation_time,
|
|
last_updated_time,
|
|
description,
|
|
workspace
|
|
)
|
|
VALUES (
|
|
:name,
|
|
:creation_time,
|
|
:last_updated_time,
|
|
:description,
|
|
:workspace
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"name": name,
|
|
"creation_time": creation_time,
|
|
"last_updated_time": creation_time,
|
|
"description": description,
|
|
"workspace": workspace,
|
|
},
|
|
)
|
|
|
|
|
|
def _insert_model_version(
|
|
conn,
|
|
*,
|
|
name: str,
|
|
version: int,
|
|
workspace: str,
|
|
run_id: str = "run-id",
|
|
storage_location: str = "location",
|
|
):
|
|
conn.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO model_versions (
|
|
name,
|
|
version,
|
|
creation_time,
|
|
last_updated_time,
|
|
user_id,
|
|
current_stage,
|
|
description,
|
|
source,
|
|
run_id,
|
|
status,
|
|
status_message,
|
|
run_link,
|
|
storage_location,
|
|
workspace
|
|
)
|
|
VALUES (
|
|
:name,
|
|
:version,
|
|
:creation_time,
|
|
:last_updated_time,
|
|
:user_id,
|
|
:current_stage,
|
|
:description,
|
|
:source,
|
|
:run_id,
|
|
:status,
|
|
:status_message,
|
|
:run_link,
|
|
:storage_location,
|
|
:workspace
|
|
)
|
|
"""
|
|
),
|
|
{
|
|
"name": name,
|
|
"version": version,
|
|
"creation_time": 0,
|
|
"last_updated_time": 0,
|
|
"user_id": "user",
|
|
"current_stage": "None",
|
|
"description": "desc",
|
|
"source": "source",
|
|
"run_id": run_id,
|
|
"status": "READY",
|
|
"status_message": "message",
|
|
"run_link": "link",
|
|
"storage_location": storage_location,
|
|
"workspace": workspace,
|
|
},
|
|
)
|
|
|
|
|
|
def _insert_registered_model_tag(
|
|
conn,
|
|
*,
|
|
workspace: str,
|
|
name: str,
|
|
key: str,
|
|
value: str = "value",
|
|
):
|
|
_insert_table_row(
|
|
conn,
|
|
_REGISTERED_MODEL_TAGS,
|
|
workspace=workspace,
|
|
key=key,
|
|
value=value,
|
|
name=name,
|
|
)
|
|
|
|
|
|
def _insert_model_version_tag(
|
|
conn,
|
|
*,
|
|
workspace: str,
|
|
name: str,
|
|
version: int,
|
|
key: str,
|
|
value: str = "value",
|
|
):
|
|
_insert_table_row(
|
|
conn,
|
|
_MODEL_VERSION_TAGS,
|
|
workspace=workspace,
|
|
key=key,
|
|
value=value,
|
|
name=name,
|
|
version=version,
|
|
)
|
|
|
|
|
|
def _insert_registered_model_alias(
|
|
conn,
|
|
*,
|
|
workspace: str,
|
|
name: str,
|
|
alias: str,
|
|
version: int = 1,
|
|
):
|
|
_insert_table_row(
|
|
conn,
|
|
_REGISTERED_MODEL_ALIASES,
|
|
workspace=workspace,
|
|
name=name,
|
|
alias=alias,
|
|
version=version,
|
|
)
|
|
|
|
|
|
def _insert_evaluation_dataset(
|
|
conn,
|
|
*,
|
|
dataset_id: str,
|
|
workspace: str,
|
|
name: str = "Dataset",
|
|
digest: str = "digest",
|
|
):
|
|
conn.execute(
|
|
sa.insert(_EVALUATION_DATASETS).values(
|
|
dataset_id=dataset_id,
|
|
name=name,
|
|
schema="schema",
|
|
profile="profile",
|
|
digest=digest,
|
|
created_time=0,
|
|
last_update_time=0,
|
|
created_by="user",
|
|
last_updated_by="user",
|
|
workspace=workspace,
|
|
)
|
|
)
|
|
|
|
|
|
_SECRETS = sa.table(
|
|
"secrets",
|
|
sa.column("secret_id"),
|
|
sa.column("secret_name"),
|
|
sa.column("encrypted_value"),
|
|
sa.column("wrapped_dek"),
|
|
sa.column("kek_version"),
|
|
sa.column("masked_value"),
|
|
sa.column("provider"),
|
|
sa.column("auth_config"),
|
|
sa.column("description"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
sa.column("workspace"),
|
|
)
|
|
|
|
|
|
def _insert_secret(
|
|
conn,
|
|
*,
|
|
secret_id: str,
|
|
secret_name: str,
|
|
workspace: str,
|
|
):
|
|
conn.execute(
|
|
sa.insert(_SECRETS).values(
|
|
secret_id=secret_id,
|
|
secret_name=secret_name,
|
|
encrypted_value=b"encrypted",
|
|
wrapped_dek=b"dek",
|
|
kek_version=1,
|
|
masked_value="***",
|
|
provider="openai",
|
|
auth_config=None,
|
|
description=None,
|
|
created_by="user",
|
|
created_at=0,
|
|
last_updated_by="user",
|
|
last_updated_at=0,
|
|
workspace=workspace,
|
|
)
|
|
)
|
|
|
|
|
|
_ENDPOINTS = sa.table(
|
|
"endpoints",
|
|
sa.column("endpoint_id"),
|
|
sa.column("name"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
sa.column("workspace"),
|
|
)
|
|
|
|
|
|
def _insert_endpoint(
|
|
conn,
|
|
*,
|
|
endpoint_id: str,
|
|
name: str,
|
|
workspace: str,
|
|
):
|
|
conn.execute(
|
|
sa.insert(_ENDPOINTS).values(
|
|
endpoint_id=endpoint_id,
|
|
name=name,
|
|
created_by="user",
|
|
created_at=0,
|
|
last_updated_by="user",
|
|
last_updated_at=0,
|
|
workspace=workspace,
|
|
)
|
|
)
|
|
|
|
|
|
_MODEL_DEFINITIONS = sa.table(
|
|
"model_definitions",
|
|
sa.column("model_definition_id"),
|
|
sa.column("name"),
|
|
sa.column("secret_id"),
|
|
sa.column("provider"),
|
|
sa.column("model_name"),
|
|
sa.column("created_by"),
|
|
sa.column("created_at"),
|
|
sa.column("last_updated_by"),
|
|
sa.column("last_updated_at"),
|
|
sa.column("workspace"),
|
|
)
|
|
|
|
|
|
def _insert_model_definition(
|
|
conn,
|
|
*,
|
|
model_definition_id: str,
|
|
name: str,
|
|
workspace: str,
|
|
):
|
|
conn.execute(
|
|
sa.insert(_MODEL_DEFINITIONS).values(
|
|
model_definition_id=model_definition_id,
|
|
name=name,
|
|
secret_id=None,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
created_by="user",
|
|
created_at=0,
|
|
last_updated_by="user",
|
|
last_updated_at=0,
|
|
workspace=workspace,
|
|
)
|
|
)
|
|
|
|
|
|
def _fetch_conflicts(conn, table_name: str, columns: tuple[str, ...]):
|
|
metadata = sa.MetaData()
|
|
table = sa.Table(table_name, metadata, autoload_with=conn)
|
|
group_columns = [table.c[column] for column in columns]
|
|
stmt = sa.select(*group_columns).group_by(*group_columns).having(sa.func.count() > 1)
|
|
return conn.execute(stmt).fetchall()
|
|
|
|
|
|
def test_workspace_migration_upgrade_adds_columns_and_backfills(tmp_path):
|
|
engine, config = _prepare_database(tmp_path)
|
|
try:
|
|
with engine.begin() as conn:
|
|
_seed_pre_workspace_entities(conn)
|
|
|
|
command.upgrade(config, REVISION)
|
|
inspector = sa.inspect(engine)
|
|
|
|
_assert_workspace_columns(inspector, "default")
|
|
|
|
with engine.connect() as conn:
|
|
assert conn.execute(
|
|
sa.text(
|
|
"SELECT experiment_id, name, workspace FROM experiments ORDER BY experiment_id"
|
|
)
|
|
).fetchall() == [(1, "exp-default", "default")]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT run_uuid, experiment_id FROM runs ORDER BY run_uuid")
|
|
).fetchall() == [("run-default", 1)]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT name, workspace FROM registered_models")
|
|
).fetchall() == [("rm-default", "default")]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT name, version, workspace FROM model_versions")
|
|
).fetchall() == [("rm-default", 1, "default")]
|
|
|
|
assert conn.execute(
|
|
sa.select(
|
|
_REGISTERED_MODEL_TAGS.c.workspace,
|
|
_REGISTERED_MODEL_TAGS.c.name,
|
|
_REGISTERED_MODEL_TAGS.c.key,
|
|
)
|
|
).fetchall() == [("default", "rm-default", "tag")]
|
|
|
|
assert conn.execute(
|
|
sa.select(
|
|
_MODEL_VERSION_TAGS.c.workspace,
|
|
_MODEL_VERSION_TAGS.c.name,
|
|
_MODEL_VERSION_TAGS.c.version,
|
|
_MODEL_VERSION_TAGS.c.key,
|
|
)
|
|
).fetchall() == [("default", "rm-default", 1, "tag")]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT workspace, name, alias FROM registered_model_aliases")
|
|
).fetchall() == [("default", "rm-default", "alias")]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT dataset_id, workspace FROM evaluation_datasets")
|
|
).fetchall() == [("ds-default", "default")]
|
|
|
|
assert conn.execute(
|
|
sa.text("SELECT name, description FROM workspaces ORDER BY name")
|
|
).fetchall() == [("default", "Default workspace for legacy resources")]
|
|
|
|
pk_registered_models = inspector.get_pk_constraint("registered_models")
|
|
assert pk_registered_models["constrained_columns"] == ["workspace", "name"]
|
|
|
|
pk_model_versions = inspector.get_pk_constraint("model_versions")
|
|
assert pk_model_versions["constrained_columns"] == [
|
|
"workspace",
|
|
"name",
|
|
"version",
|
|
]
|
|
|
|
pk_registered_model_tags = inspector.get_pk_constraint("registered_model_tags")
|
|
assert pk_registered_model_tags["constrained_columns"] == [
|
|
"workspace",
|
|
"key",
|
|
"name",
|
|
]
|
|
|
|
pk_model_version_tags = inspector.get_pk_constraint("model_version_tags")
|
|
assert pk_model_version_tags["constrained_columns"] == [
|
|
"workspace",
|
|
"key",
|
|
"name",
|
|
"version",
|
|
]
|
|
|
|
pk_model_aliases = inspector.get_pk_constraint("registered_model_aliases")
|
|
assert pk_model_aliases["constrained_columns"] == [
|
|
"workspace",
|
|
"name",
|
|
"alias",
|
|
]
|
|
|
|
try:
|
|
unique_experiments = inspector.get_unique_constraints("experiments")
|
|
except NotImplementedError:
|
|
if inspector.bind.dialect.name == "mssql":
|
|
unique_experiments = None
|
|
else:
|
|
raise
|
|
if unique_experiments is not None:
|
|
assert any(
|
|
{"workspace", "name"} == set(constraint.get("column_names", []))
|
|
for constraint in unique_experiments
|
|
)
|
|
|
|
fk_model_versions = inspector.get_foreign_keys("model_versions")
|
|
assert any(
|
|
fk.get("constrained_columns") == ["workspace", "name"]
|
|
and fk.get("referred_table") == "registered_models"
|
|
for fk in fk_model_versions
|
|
)
|
|
|
|
assert _has_index(inspector, "experiments", "idx_experiments_workspace", ["workspace"])
|
|
assert _has_index(
|
|
inspector,
|
|
"experiments",
|
|
"idx_experiments_workspace_creation_time",
|
|
["workspace", "creation_time"],
|
|
)
|
|
assert _has_index(
|
|
inspector, "registered_models", "idx_registered_models_workspace", ["workspace"]
|
|
)
|
|
assert _has_index(
|
|
inspector, "evaluation_datasets", "idx_evaluation_datasets_workspace", ["workspace"]
|
|
)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
def test_workspace_migration_downgrade_reverts_schema(tmp_path):
|
|
engine, config = _prepare_database(tmp_path)
|
|
try:
|
|
command.upgrade(config, REVISION)
|
|
with engine.begin() as conn:
|
|
_add_workspace(conn, "team-a", "Team A")
|
|
_insert_experiment(conn, experiment_id=1, name="exp-default", workspace="default")
|
|
_insert_run(
|
|
conn,
|
|
run_uuid="run-default",
|
|
experiment_id=1,
|
|
name="downgrade-validation-run",
|
|
)
|
|
_insert_experiment(conn, experiment_id=2, name="exp-team-a", workspace="team-a")
|
|
|
|
command.downgrade(config, PREVIOUS_REVISION)
|
|
inspector = sa.inspect(engine)
|
|
|
|
tables = inspector.get_table_names()
|
|
assert "workspaces" not in tables
|
|
|
|
for table in (
|
|
"experiments",
|
|
"registered_models",
|
|
"model_versions",
|
|
"registered_model_tags",
|
|
"model_version_tags",
|
|
"registered_model_aliases",
|
|
"evaluation_datasets",
|
|
"webhooks",
|
|
"jobs",
|
|
):
|
|
column_names = {col["name"] for col in inspector.get_columns(table)}
|
|
assert "workspace" not in column_names
|
|
|
|
with engine.connect() as conn:
|
|
assert conn.execute(
|
|
sa.text("SELECT experiment_id, name FROM experiments ORDER BY experiment_id")
|
|
).fetchall() == [(1, "exp-default"), (2, "exp-team-a")]
|
|
assert conn.execute(
|
|
sa.text("SELECT run_uuid, experiment_id FROM runs ORDER BY run_uuid")
|
|
).fetchall() == [("run-default", 1)]
|
|
|
|
pk_registered_models = inspector.get_pk_constraint("registered_models")
|
|
assert pk_registered_models["constrained_columns"] == ["name"]
|
|
|
|
pk_model_versions = inspector.get_pk_constraint("model_versions")
|
|
assert pk_model_versions["constrained_columns"] == ["name", "version"]
|
|
|
|
pk_registered_model_tags = inspector.get_pk_constraint("registered_model_tags")
|
|
assert pk_registered_model_tags["constrained_columns"] == ["key", "name"]
|
|
|
|
pk_model_version_tags = inspector.get_pk_constraint("model_version_tags")
|
|
assert pk_model_version_tags["constrained_columns"] == ["key", "name", "version"]
|
|
|
|
pk_registered_model_aliases = inspector.get_pk_constraint("registered_model_aliases")
|
|
assert pk_registered_model_aliases["constrained_columns"] == ["name", "alias"]
|
|
|
|
try:
|
|
unique_experiments = inspector.get_unique_constraints("experiments")
|
|
except NotImplementedError:
|
|
if inspector.bind.dialect.name == "mssql":
|
|
unique_experiments = None
|
|
else:
|
|
raise
|
|
if unique_experiments is not None:
|
|
assert any(
|
|
set(constraint.get("column_names", [])) == {"name"}
|
|
for constraint in unique_experiments
|
|
)
|
|
|
|
fk_model_versions = inspector.get_foreign_keys("model_versions")
|
|
assert any(
|
|
fk.get("constrained_columns") == ["name"]
|
|
and fk.get("referred_table") == "registered_models"
|
|
for fk in fk_model_versions
|
|
)
|
|
|
|
fk_registered_model_tags = inspector.get_foreign_keys("registered_model_tags")
|
|
assert any(
|
|
fk.get("constrained_columns") == ["name"]
|
|
and fk.get("referred_table") == "registered_models"
|
|
for fk in fk_registered_model_tags
|
|
)
|
|
|
|
fk_model_version_tags = inspector.get_foreign_keys("model_version_tags")
|
|
assert any(
|
|
fk.get("constrained_columns") == ["name", "version"]
|
|
and fk.get("referred_table") == "model_versions"
|
|
for fk in fk_model_version_tags
|
|
)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
def _setup_experiment_conflict(conn):
|
|
_insert_experiment(conn, experiment_id=1, name="duplicate-exp", workspace="default")
|
|
_insert_run(conn, run_uuid="run-exp-default", experiment_id=1)
|
|
_insert_experiment(conn, experiment_id=2, name="duplicate-exp", workspace="team-a")
|
|
|
|
|
|
def _setup_registered_model_conflict(conn):
|
|
_insert_registered_model(conn, name="duplicate-model", workspace="default")
|
|
_insert_registered_model(conn, name="duplicate-model", workspace="team-a")
|
|
|
|
|
|
def _setup_model_version_conflict(conn):
|
|
_insert_registered_model(conn, name="mv-model", workspace="default")
|
|
_insert_registered_model(conn, name="mv-model", workspace="team-a")
|
|
_insert_model_version(conn, name="mv-model", version=1, workspace="default")
|
|
_insert_model_version(conn, name="mv-model", version=1, workspace="team-a")
|
|
|
|
|
|
def _setup_registered_model_tag_conflict(conn):
|
|
_insert_registered_model(conn, name="tag-model", workspace="default")
|
|
_insert_registered_model(conn, name="tag-model", workspace="team-a")
|
|
_insert_registered_model_tag(conn, workspace="default", name="tag-model", key="tag-key")
|
|
_insert_registered_model_tag(conn, workspace="team-a", name="tag-model", key="tag-key")
|
|
|
|
|
|
def _setup_model_version_tag_conflict(conn):
|
|
_insert_registered_model(conn, name="mvt-model", workspace="default")
|
|
_insert_registered_model(conn, name="mvt-model", workspace="team-a")
|
|
_insert_model_version(conn, name="mvt-model", version=1, workspace="default")
|
|
_insert_model_version(conn, name="mvt-model", version=1, workspace="team-a")
|
|
_insert_model_version_tag(
|
|
conn, workspace="default", name="mvt-model", version=1, key="mv-tag-key"
|
|
)
|
|
_insert_model_version_tag(
|
|
conn, workspace="team-a", name="mvt-model", version=1, key="mv-tag-key"
|
|
)
|
|
|
|
|
|
def _setup_registered_model_alias_conflict(conn):
|
|
_insert_registered_model(conn, name="alias-model", workspace="default")
|
|
_insert_registered_model(conn, name="alias-model", workspace="team-a")
|
|
_insert_registered_model_alias(conn, workspace="default", name="alias-model", alias="latest")
|
|
_insert_registered_model_alias(conn, workspace="team-a", name="alias-model", alias="latest")
|
|
|
|
|
|
def _setup_evaluation_dataset_conflict(conn):
|
|
_insert_evaluation_dataset(
|
|
conn, dataset_id="ds-default", name="duplicate-ds", workspace="default"
|
|
)
|
|
_insert_evaluation_dataset(
|
|
conn, dataset_id="ds-team-a", name="duplicate-ds", workspace="team-a"
|
|
)
|
|
|
|
|
|
def _setup_secret_conflict(conn):
|
|
_insert_secret(conn, secret_id="s-default", secret_name="duplicate-secret", workspace="default")
|
|
_insert_secret(conn, secret_id="s-team-a", secret_name="duplicate-secret", workspace="team-a")
|
|
|
|
|
|
def _setup_endpoint_conflict(conn):
|
|
_insert_endpoint(conn, endpoint_id="e-default", name="duplicate-endpoint", workspace="default")
|
|
_insert_endpoint(conn, endpoint_id="e-team-a", name="duplicate-endpoint", workspace="team-a")
|
|
|
|
|
|
def _setup_model_definition_conflict(conn):
|
|
_insert_model_definition(
|
|
conn, model_definition_id="md-default", name="duplicate-def", workspace="default"
|
|
)
|
|
_insert_model_definition(
|
|
conn, model_definition_id="md-team-a", name="duplicate-def", workspace="team-a"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("setup_conflict", "expected_fragment", "case_slug"),
|
|
[
|
|
(_setup_experiment_conflict, "duplicate experiments with the same name", "experiments"),
|
|
(
|
|
_setup_registered_model_conflict,
|
|
"duplicate registered models with the same name",
|
|
"models",
|
|
),
|
|
(
|
|
_setup_evaluation_dataset_conflict,
|
|
"duplicate evaluation datasets with the same name",
|
|
"evaluation_datasets",
|
|
),
|
|
(
|
|
_setup_secret_conflict,
|
|
"duplicate secrets with the same name",
|
|
"secrets",
|
|
),
|
|
(
|
|
_setup_endpoint_conflict,
|
|
"duplicate endpoints with the same name",
|
|
"endpoints",
|
|
),
|
|
(
|
|
_setup_model_definition_conflict,
|
|
"duplicate model definitions with the same name",
|
|
"model_definitions",
|
|
),
|
|
],
|
|
)
|
|
def test_workspace_migration_downgrade_detects_conflicts(
|
|
tmp_path, setup_conflict, expected_fragment, case_slug
|
|
):
|
|
case_dir = tmp_path / f"conflict_{case_slug}"
|
|
case_dir.mkdir()
|
|
engine, config = _prepare_database(case_dir)
|
|
try:
|
|
command.upgrade(config, REVISION)
|
|
with engine.begin() as conn:
|
|
_add_workspace(conn, "team-a", "Team A")
|
|
setup_conflict(conn)
|
|
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match=re.escape(expected_fragment),
|
|
):
|
|
command.downgrade(config, PREVIOUS_REVISION)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("setup_conflict", "table_name", "columns", "case_slug"),
|
|
[
|
|
(
|
|
_setup_model_version_conflict,
|
|
"model_versions",
|
|
("name", "version"),
|
|
"model_versions",
|
|
),
|
|
(
|
|
_setup_registered_model_tag_conflict,
|
|
"registered_model_tags",
|
|
("name", "key"),
|
|
"registered_model_tags",
|
|
),
|
|
(
|
|
_setup_model_version_tag_conflict,
|
|
"model_version_tags",
|
|
("name", "version", "key"),
|
|
"model_version_tags",
|
|
),
|
|
(
|
|
_setup_registered_model_alias_conflict,
|
|
"registered_model_aliases",
|
|
("name", "alias"),
|
|
"registered_model_aliases",
|
|
),
|
|
],
|
|
)
|
|
def test_workspace_migration_conflict_detection_queries(
|
|
tmp_path, setup_conflict, table_name, columns, case_slug
|
|
):
|
|
case_dir = tmp_path / f"conflict_query_{case_slug}"
|
|
case_dir.mkdir()
|
|
engine, config = _prepare_database(case_dir)
|
|
try:
|
|
command.upgrade(config, REVISION)
|
|
with engine.begin() as conn:
|
|
_add_workspace(conn, "team-a", "Team A")
|
|
setup_conflict(conn)
|
|
conflicts = _fetch_conflicts(conn, table_name, columns)
|
|
assert conflicts, f"Expected conflicts for {table_name}, found none"
|
|
finally:
|
|
engine.dispose()
|