opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-08-30 17:11:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): migrate legacy model types during upgrade (#41326)

Browse Source
This commit is contained in:
-LAN-
2026-08-27 02:00:33 +00:00
committed by GitHub
parent 0d05eae517
commit 82520c144f
7 changed files with 550 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/core/helper/model_provider_cache.py
+1 -1
View File
@@ -14,7 +14,7 @@ class ProviderCredentialsCacheType(StrEnum):
class ProviderCredentialsCache:
def __init__(self, tenant_id: str, identity_id: str, cache_type: ProviderCredentialsCacheType):
self.cache_key = f"{cache_type}_credentials:tenant_id:{tenant_id}:id:{identity_id}"
self.cache_key = f"{cache_type}_credentials:v2:tenant_id:{tenant_id}:id:{identity_id}"
def get(self) -> dict[str, Any] | None:
"""
api/core/provider_manager.py
+1 -1
View File
@@ -73,7 +73,7 @@ _credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
_PROVIDER_CONFIGURATION_CACHE_TTL_SECONDS = 300
_PROVIDER_CONFIGURATION_CACHE_VERSION_TTL_SECONDS = 360
_PROVIDER_CONFIGURATION_CACHE_VERSION_KEY = "provider_configurations:tenant:{tenant_id}:source:{source}:version"
_PROVIDER_CONFIGURATION_CACHE_SOURCE_KEY = "provider_configurations:tenant:{tenant_id}:source:{source}:v:{version}"
_PROVIDER_CONFIGURATION_CACHE_SOURCE_KEY = "provider_configurations:v2:tenant:{tenant_id}:source:{source}:v:{version}"
class ProviderConfigurationCacheSource(StrEnum):
api/migrations/versions/2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py
+267
View File
@@ -0,0 +1,267 @@
"""migrate legacy model types
Revision ID: 5578e028b2f2
Revises: 9b7c6d5e4f3a
Create Date: 2026-08-27 12:00:00.000000
"""
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision = "5578e028b2f2"
down_revision = "9b7c6d5e4f3a"
branch_labels = None
depends_on = None
_AFFECTED_TABLES = (
"provider_models",
"provider_model_credentials",
"tenant_default_models",
"provider_model_settings",
"load_balancing_model_configs",
)
_LEGACY_MODEL_TYPES = ("text-generation", "embeddings", "reranking")
_MAPPED_MODEL_TYPES = (
"text-generation",
"llm",
"embeddings",
"text-embedding",
"reranking",
"rerank",
)
_CREDENTIAL_MERGES_TABLE = "tmp_5578e028b2f2_credential_merges"
def _canonical_model_type(alias: str) -> str:
return f"""CASE {alias}.model_type
WHEN 'text-generation' THEN 'llm'
WHEN 'embeddings' THEN 'text-embedding'
WHEN 'reranking' THEN 'rerank'
ELSE {alias}.model_type
END"""
def _mapped_model_types_sql() -> str:
return ", ".join(f"'{model_type}'" for model_type in _MAPPED_MODEL_TYPES)
def _legacy_model_types_sql() -> str:
return ", ".join(f"'{model_type}'" for model_type in _LEGACY_MODEL_TYPES)
def _same_business_key(left_alias: str, right_alias: str, key_columns: tuple[str, ...]) -> str:
key_condition = "\n AND ".join(f"{left_alias}.{column} = {right_alias}.{column}" for column in key_columns)
return f"""{key_condition}
AND {_canonical_model_type(left_alias)} = {_canonical_model_type(right_alias)}"""
def _delete_duplicates(
table_name: str,
key_columns: tuple[str, ...],
*,
extra_condition: str | None = None,
require_legacy_row: bool = True,
) -> None:
dialect_name = op.get_context().dialect.name
scoped_condition = f"\n AND {extra_condition}" if extra_condition else ""
common_condition = f"""{_same_business_key("loser", "winner", key_columns)}
AND loser.model_type IN ({_mapped_model_types_sql()})
AND winner.model_type IN ({_mapped_model_types_sql()})
AND (
loser.updated_at < winner.updated_at
OR (loser.updated_at = winner.updated_at AND loser.id < winner.id)
){scoped_condition}"""
legacy_table = f", {table_name} AS legacy" if require_legacy_row else ""
legacy_join = ""
if require_legacy_row:
legacy_join = f"""
AND {_same_business_key("loser", "legacy", key_columns)}
AND legacy.model_type IN ({_legacy_model_types_sql()})"""
if dialect_name == "postgresql":
op.execute(
sa.text(
f"""DELETE FROM {table_name} AS loser
USING {table_name} AS winner{legacy_table}
WHERE {common_condition}{legacy_join}"""
)
)
return
if dialect_name in {"mysql", "mariadb"}:
legacy_table = ""
if require_legacy_row:
legacy_table = f"""
INNER JOIN {table_name} AS legacy
ON {_same_business_key("loser", "legacy", key_columns)}
AND legacy.model_type IN ({_legacy_model_types_sql()})"""
op.execute(
sa.text(
f"""DELETE loser
FROM {table_name} AS loser
INNER JOIN {table_name} AS winner
ON {common_condition}{legacy_table}"""
)
)
return
raise RuntimeError(f"unsupported database dialect: {dialect_name}")
def _create_credential_merges() -> None:
dialect_name = op.get_context().dialect.name
if dialect_name == "postgresql":
table_options = "ON COMMIT DROP"
elif dialect_name in {"mysql", "mariadb"}:
# MySQL does not roll back temporary-table DDL. Remove a table left by
# an in-process retry without risking a permanent table of the same name.
op.execute(sa.text(f"DROP TEMPORARY TABLE IF EXISTS {_CREDENTIAL_MERGES_TABLE}"))
table_options = ""
else:
raise RuntimeError(f"unsupported database dialect: {dialect_name}")
op.execute(
sa.text(
f"""CREATE TEMPORARY TABLE {_CREDENTIAL_MERGES_TABLE} {table_options} AS
SELECT id AS loser_id, winner_id
FROM (
SELECT
id,
FIRST_VALUE(id) OVER (
PARTITION BY
tenant_id,
provider_name,
model_name,
credential_name,
{_canonical_model_type("provider_model_credentials")}
ORDER BY updated_at DESC, id DESC
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
) AS winner_id,
SUM(CASE WHEN model_type IN ({_legacy_model_types_sql()}) THEN 1 ELSE 0 END) OVER (
PARTITION BY
tenant_id,
provider_name,
model_name,
credential_name,
{_canonical_model_type("provider_model_credentials")}
) AS legacy_count
FROM provider_model_credentials
WHERE model_type IN ({_mapped_model_types_sql()})
) AS ranked_credentials
WHERE id <> winner_id AND legacy_count > 0"""
)
)
def _rewrite_credential_references() -> None:
dialect_name = op.get_context().dialect.name
if dialect_name == "postgresql":
op.execute(
sa.text(
f"""UPDATE provider_models AS model
SET credential_id = merges.winner_id
FROM {_CREDENTIAL_MERGES_TABLE} AS merges
WHERE model.credential_id = merges.loser_id"""
)
)
op.execute(
sa.text(
f"""UPDATE load_balancing_model_configs AS config
SET
credential_id = merges.winner_id,
name = winner.credential_name,
encrypted_config = winner.encrypted_config
FROM {_CREDENTIAL_MERGES_TABLE} AS merges
INNER JOIN provider_model_credentials AS winner ON winner.id = merges.winner_id
WHERE config.credential_id = merges.loser_id"""
)
)
return
if dialect_name in {"mysql", "mariadb"}:
op.execute(
sa.text(
f"""UPDATE provider_models AS model
INNER JOIN {_CREDENTIAL_MERGES_TABLE} AS merges ON model.credential_id = merges.loser_id
SET model.credential_id = merges.winner_id"""
)
)
op.execute(
sa.text(
f"""UPDATE load_balancing_model_configs AS config
INNER JOIN {_CREDENTIAL_MERGES_TABLE} AS merges ON config.credential_id = merges.loser_id
INNER JOIN provider_model_credentials AS winner ON winner.id = merges.winner_id
SET
config.credential_id = merges.winner_id,
config.name = winner.credential_name,
config.encrypted_config = winner.encrypted_config"""
)
)
return
raise RuntimeError(f"unsupported database dialect: {dialect_name}")
def _delete_merged_credentials() -> None:
dialect_name = op.get_context().dialect.name
if dialect_name == "postgresql":
op.execute(
sa.text(
f"""DELETE FROM provider_model_credentials AS credential
USING {_CREDENTIAL_MERGES_TABLE} AS merges
WHERE credential.id = merges.loser_id"""
)
)
op.execute(sa.text(f"DROP TABLE {_CREDENTIAL_MERGES_TABLE}"))
return
if dialect_name in {"mysql", "mariadb"}:
op.execute(
sa.text(
f"""DELETE credential
FROM provider_model_credentials AS credential
INNER JOIN {_CREDENTIAL_MERGES_TABLE} AS merges ON credential.id = merges.loser_id"""
)
)
op.execute(sa.text(f"DROP TEMPORARY TABLE {_CREDENTIAL_MERGES_TABLE}"))
return
raise RuntimeError(f"unsupported database dialect: {dialect_name}")
def _canonicalize_model_types() -> None:
for table_name in _AFFECTED_TABLES:
op.execute(
sa.text(
f"""UPDATE {table_name}
SET model_type = {_canonical_model_type(table_name)}
WHERE model_type IN ('text-generation', 'embeddings', 'reranking')"""
)
)
def upgrade() -> None:
# Preserve the established manual migration policy: the newest row wins
# when legacy and canonical values collapse onto the same business key.
_delete_duplicates("provider_models", ("tenant_id", "provider_name", "model_name"))
_delete_duplicates("tenant_default_models", ("tenant_id",))
_delete_duplicates("provider_model_settings", ("tenant_id", "provider_name", "model_name"))
_delete_duplicates(
"load_balancing_model_configs",
("tenant_id", "provider_name", "model_name"),
extra_condition="loser.name = '__inherit__' AND winner.name = '__inherit__'",
require_legacy_row=False,
)
_create_credential_merges()
_rewrite_credential_references()
_delete_merged_credentials()
_canonicalize_model_types()
def downgrade() -> None:
# Canonical rows created after the enum rename cannot be distinguished from
# rows changed here, so reversing this data migration would corrupt valid data.
pass
api/tests/test_containers_integration_tests/commands/test_legacy_model_type_migration.py
+192
View File
@@ -1,13 +1,17 @@
from __future__ import annotations
import importlib
import importlib.util
import io
import json
from collections.abc import Generator
from datetime import datetime, timedelta
from pathlib import Path
import pytest
import sqlalchemy as sa
from alembic.migration import MigrationContext
from alembic.operations import Operations
from tests.helpers.legacy_model_type_migration import (
assert_tenant_rows_use_only_canonical_model_types,
@@ -16,6 +20,28 @@ from tests.helpers.legacy_model_type_migration import (
seed_legacy_model_type_dirty_data,
)
_ALEMBIC_MIGRATION_PATH = (
Path(__file__).resolve().parents[3]
/ "migrations/versions/2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py"
)
def _run_legacy_model_type_alembic_upgrade(engine: sa.Engine) -> None:
spec = importlib.util.spec_from_file_location("migrate_legacy_model_types", _ALEMBIC_MIGRATION_PATH)
if spec is None or spec.loader is None:
raise RuntimeError("failed to load legacy model type migration")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with engine.begin() as connection:
operations = Operations(MigrationContext.configure(connection))
original_op = module.__dict__["op"]
module.__dict__["op"] = operations
try:
module.__dict__["upgrade"]()
finally:
module.__dict__["op"] = original_op
def _parse_json_lines(output: io.StringIO) -> list[dict[str, object]]:
return [json.loads(line) for line in output.getvalue().splitlines() if line.strip()]
@@ -223,6 +249,172 @@ def test_legacy_model_type_migration_end_to_end_across_supported_backends(
assert second_apply_state == first_apply_state
def test_legacy_model_type_alembic_upgrade_across_supported_backends(
container_engine: tuple[str, sa.Engine],
) -> None:
_, engine = container_engine
helper_module = importlib.import_module("tests.helpers.legacy_model_type_migration")
helper_module.drop_minimal_legacy_model_type_schema(engine)
fixture = seed_legacy_model_type_dirty_data(engine)
canonical_provider_model_id = "00000000-0000-0000-0000-00000000ca01"
canonical_default_model_id = "00000000-0000-0000-0000-00000000ca02"
canonical_credential_ids = {
"00000000-0000-0000-0000-00000000ca03",
"00000000-0000-0000-0000-00000000ca04",
}
older_inherit_id = "00000000-0000-0000-0000-00000000ca05"
newer_inherit_id = "00000000-0000-0000-0000-00000000ca06"
now = datetime(2025, 1, 1, 12, 0, 0)
with engine.begin() as connection:
connection.execute(
sa.text(
"""
INSERT INTO provider_models
(
id, tenant_id, provider_name, model_name, model_type,
credential_id, is_valid, created_at, updated_at
)
VALUES
(
:id, :tenant_id, 'openai', 'gpt-4o-mini', 'llm',
:credential_id, :is_valid, :created_at, :updated_at
)
"""
),
{
"id": canonical_provider_model_id,
"tenant_id": fixture.primary.tenant_id,
"credential_id": fixture.primary.winner_credential_id,
"is_valid": True,
"created_at": now - timedelta(days=2),
"updated_at": now - timedelta(hours=7),
},
)
connection.execute(
sa.text(
"""
INSERT INTO tenant_default_models
(id, tenant_id, provider_name, model_name, model_type, created_at, updated_at)
VALUES
(:id, :tenant_id, 'openai', 'gpt-4o-mini', 'llm', :created_at, :updated_at)
"""
),
{
"id": canonical_default_model_id,
"tenant_id": fixture.primary.tenant_id,
"created_at": now - timedelta(days=2),
"updated_at": now,
},
)
connection.execute(
sa.text(
"""
INSERT INTO provider_model_credentials
(
id, tenant_id, provider_name, model_name, model_type,
credential_name, encrypted_config, created_at, updated_at
)
VALUES
(:older_id, :tenant_id, 'openai', 'gpt-4o-mini', 'llm',
'canonical-only', '{"api_key":"older"}', :created_at, :older_updated_at),
(:newer_id, :tenant_id, 'openai', 'gpt-4o-mini', 'llm',
'canonical-only', '{"api_key":"newer"}', :created_at, :newer_updated_at)
"""
),
{
"older_id": min(canonical_credential_ids),
"newer_id": max(canonical_credential_ids),
"tenant_id": fixture.primary.tenant_id,
"created_at": now - timedelta(days=2),
"older_updated_at": now - timedelta(hours=2),
"newer_updated_at": now - timedelta(hours=1),
},
)
_insert_load_balancing_model_config(
engine,
row_id=older_inherit_id,
tenant_id=fixture.primary.tenant_id,
provider_name="openai",
model_name="gpt-4o-mini",
model_type="llm",
name="__inherit__",
encrypted_config='{"api_key":"older-inherit"}',
credential_id=fixture.primary.winner_credential_id,
enabled=True,
created_at=now - timedelta(days=2),
updated_at=now - timedelta(hours=2),
)
_insert_load_balancing_model_config(
engine,
row_id=newer_inherit_id,
tenant_id=fixture.primary.tenant_id,
provider_name="openai",
model_name="gpt-4o-mini",
model_type="text-generation",
name="__inherit__",
encrypted_config='{"api_key":"newer-inherit"}',
credential_id=fixture.primary.distinct_credential_id,
enabled=True,
created_at=now - timedelta(days=2),
updated_at=now - timedelta(hours=1),
)
_run_legacy_model_type_alembic_upgrade(engine)
for tenant_id in (fixture.primary.tenant_id, fixture.secondary.tenant_id):
assert_tenant_rows_use_only_canonical_model_types(engine, tenant_id)
table_names = (
"provider_models",
"tenant_default_models",
"provider_model_settings",
"load_balancing_model_configs",
"provider_model_credentials",
)
first_apply_state = {table_name: fetch_table_rows(engine, table_name) for table_name in table_names}
primary_provider_models = [
row
for row in first_apply_state["provider_models"]
if row["tenant_id"] == fixture.primary.tenant_id and row["model_name"] == "gpt-4o-mini"
]
assert [row["id"] for row in primary_provider_models] == [fixture.primary.provider_model_id]
assert primary_provider_models[0]["credential_id"] == fixture.primary.winner_credential_id
primary_defaults = [
row
for row in first_apply_state["tenant_default_models"]
if row["tenant_id"] == fixture.primary.tenant_id and row["model_type"] == "llm"
]
assert [row["id"] for row in primary_defaults] == [canonical_default_model_id]
primary_credential_ids = {
row["id"]
for row in first_apply_state["provider_model_credentials"]
if row["tenant_id"] == fixture.primary.tenant_id
}
assert canonical_credential_ids <= primary_credential_ids
assert count_rows(engine, "provider_model_credentials", tenant_id=fixture.primary.tenant_id) == 4
primary_load_balancing_config = next(
row
for row in first_apply_state["load_balancing_model_configs"]
if row["id"] == fixture.primary.load_balancing_config_id
)
assert primary_load_balancing_config["credential_id"] == fixture.primary.winner_credential_id
assert primary_load_balancing_config["encrypted_config"] == fixture.primary.winner_encrypted_config
primary_inherit_ids = {
row["id"]
for row in first_apply_state["load_balancing_model_configs"]
if row["tenant_id"] == fixture.primary.tenant_id and row["name"] == "__inherit__"
}
assert primary_inherit_ids == {newer_inherit_id}
_run_legacy_model_type_alembic_upgrade(engine)
second_apply_state = {table_name: fetch_table_rows(engine, table_name) for table_name in table_names}
assert second_apply_state == first_apply_state
def test_load_balancing_inherit_deduplication_is_applied_consistently_across_supported_backends(
migration_module,
container_engine: tuple[str, sa.Engine],
api/tests/unit_tests/core/helper/test_model_provider_cache.py
+1
View File
@@ -16,6 +16,7 @@ def test_model_provider_credentials_cache_get_returns_decoded_dict(mocker: Mocke
redis_client_mock.get.return_value = json.dumps(payload).encode("utf-8")
assert cache.cache_key == "provider_credentials:v2:tenant_id:tenant:id:identity"
assert cache.get() == payload
api/tests/unit_tests/core/test_provider_manager.py
+2 -2
View File
@@ -1143,8 +1143,8 @@ def test_provider_configuration_cache_skips_write_when_version_changes_during_lo
assert version_bumped is True
assert fake_redis.store[version_key] == "1"
assert "provider_configurations:tenant:tenant-id:source:provider_model_credentials:v:0" not in fake_redis.store
assert "provider_configurations:tenant:tenant-id:source:provider_model_credentials:v:1" not in fake_redis.store
assert "provider_configurations:v2:tenant:tenant-id:source:provider_model_credentials:v:0" not in fake_redis.store
assert "provider_configurations:v2:tenant:tenant-id:source:provider_model_credentials:v:1" not in fake_redis.store
assert result["openai"][0].credential_name == "primary"
api/tests/unit_tests/migrations/test_migrate_legacy_model_types.py
+86
View File
@@ -0,0 +1,86 @@
from __future__ import annotations
import importlib.util
from io import StringIO
from pathlib import Path
from types import ModuleType
import pytest
from alembic.migration import MigrationContext
from alembic.operations import Operations
_MIGRATION_PATH = (
Path(__file__).resolve().parents[3]
/ "migrations/versions/2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py"
)
def _load_migration_module() -> ModuleType:
spec = importlib.util.spec_from_file_location("migrate_legacy_model_types", _MIGRATION_PATH)
if spec is None or spec.loader is None:
raise RuntimeError("failed to load migration module")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
@pytest.mark.parametrize(
("dialect_name", "duplicate_delete", "reference_update", "temporary_drop"),
[
(
"postgresql",
"DELETE FROM provider_models AS loser USING provider_models AS winner",
"UPDATE provider_models AS model SET credential_id = merges.winner_id FROM",
"DROP TABLE tmp_5578e028b2f2_credential_merges",
),
(
"mysql",
"DELETE loser FROM provider_models AS loser INNER JOIN provider_models AS winner",
"UPDATE provider_models AS model INNER JOIN tmp_5578e028b2f2_credential_merges",
"DROP TEMPORARY TABLE tmp_5578e028b2f2_credential_merges",
),
],
)
def test_upgrade_emits_collision_safe_sql_for_supported_databases(
dialect_name: str,
duplicate_delete: str,
reference_update: str,
temporary_drop: str,
) -> None:
module = _load_migration_module()
output = StringIO()
migration_context = MigrationContext.configure(
dialect_name=dialect_name,
opts={"as_sql": True, "literal_binds": True, "output_buffer": output},
)
operations = Operations(migration_context)
original_op = module.__dict__["op"]
module.__dict__["op"] = operations
try:
module.__dict__["upgrade"]()
finally:
module.__dict__["op"] = original_op
generated_sql = " ".join(output.getvalue().split())
assert duplicate_delete in generated_sql
assert reference_update in generated_sql
assert temporary_drop in generated_sql
assert "ORDER BY updated_at DESC, id DESC" in generated_sql
assert "legacy.model_type IN ('text-generation', 'embeddings', 'reranking')" in generated_sql
assert "WHERE id <> winner_id AND legacy_count > 0" in generated_sql
for table_name in (
"provider_models",
"provider_model_credentials",
"tenant_default_models",
"provider_model_settings",
"load_balancing_model_configs",
):
assert f"UPDATE {table_name} SET model_type = CASE {table_name}.model_type" in generated_sql
for old_value, new_value in (
("text-generation", "llm"),
("embeddings", "text-embedding"),
("reranking", "rerank"),
):
assert f"WHEN '{old_value}' THEN '{new_value}'" in generated_sql