fix(database): harden DM fresh deployment startup

This commit is contained in:
GuoQing Zhang
2026-08-28 14:58:33 +08:00
parent d4928cf31b
commit 614b9e0967
4 changed files with 242 additions and 98 deletions
+13 -1
View File
@@ -269,6 +269,7 @@ async def _init_default_root_department(session):
from bisheng.database.models.department import Department
from bisheng.database.models.tenant import Tenant
root_dept_id = "BS@root"
guest_dept_id = "BS@guest"
guest_name = "临时访客"
guest_sort_order = 2147483647
@@ -285,12 +286,13 @@ async def _init_default_root_department(session):
if root_dept is None:
root_dept = Department(
dept_id="BS@root",
dept_id=root_dept_id,
name="默认组织",
parent_id=None,
tenant_id=DEFAULT_TENANT_ID,
path="",
source="local",
external_id=root_dept_id,
status="active",
)
session.add(root_dept)
@@ -300,6 +302,12 @@ async def _init_default_root_department(session):
tenant.root_dept_id = root_dept.id
session.add(root_dept)
await session.commit()
elif root_dept.external_id != root_dept_id:
# DM8 treats repeated NULL values as duplicates for this composite
# unique key. Persist the internal identity before inserting guest.
root_dept.external_id = root_dept_id
session.add(root_dept)
await session.commit()
guest = (await session.exec(select(Department).where(Department.dept_id == guest_dept_id))).first()
if guest is None:
@@ -310,6 +318,7 @@ async def _init_default_root_department(session):
tenant_id=DEFAULT_TENANT_ID,
path=f"{root_dept.path}",
source="local",
external_id=guest_dept_id,
status="active",
sort_order=guest_sort_order,
)
@@ -327,6 +336,9 @@ async def _init_default_root_department(session):
if guest.sort_order != guest_sort_order:
guest.sort_order = guest_sort_order
changed = True
if guest.external_id != guest_dept_id:
guest.external_id = guest_dept_id
changed = True
expected_path_prefix = root_dept.path or ""
if not (guest.path or "").startswith(expected_path_prefix):
guest.path = f"{expected_path_prefix}{guest.id}/"
@@ -34,39 +34,37 @@ Merge revision:
step operators have to apply.
"""
from typing import List, Sequence, Union
from collections.abc import Sequence
from typing import Union
import sqlalchemy as sa
from alembic import op
from bisheng.core.database.dialect_helpers import index_exists
revision: str = 'f020_llm_tenant'
revision: str = "f020_llm_tenant"
down_revision: Union[str, Sequence[str], None] = (
'f015_reconcile_log_fields',
'f017_llm_token_log',
"f015_reconcile_log_fields",
"f017_llm_token_log",
)
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
_OLD_UNIQUE_INDEX_NAME = 'name' # SQLAlchemy default when unique=True is on the column
_NEW_UNIQUE_INDEX_NAME = 'uk_llm_server_tenant_name'
_TABLE = 'llm_server'
_OLD_UNIQUE_INDEX_NAME = "name" # SQLAlchemy default when unique=True is on the column
_NEW_UNIQUE_INDEX_NAME = "uk_llm_server_tenant_name"
_TABLE = "llm_server"
def _find_duplicates(conn) -> List[sa.engine.Row]:
def _find_duplicates(conn) -> list[sa.engine.Row]:
"""Return rows violating the (tenant_id, name) uniqueness invariant.
Separated out so tests can patch a connection whose ``execute``
returns a canned fetchall.
"""
result = conn.execute(sa.text(
'SELECT tenant_id, name, COUNT(*) AS cnt '
f'FROM {_TABLE} '
'GROUP BY tenant_id, name '
'HAVING cnt > 1'
))
result = conn.execute(
sa.text(f"SELECT tenant_id, name, COUNT(*) AS cnt FROM {_TABLE} GROUP BY tenant_id, name HAVING COUNT(*) > 1")
)
return result.fetchall()
@@ -75,15 +73,11 @@ def upgrade() -> None:
conflicts = _find_duplicates(conn)
if conflicts:
lines = [
f' tenant_id={row.tenant_id!r}, name={row.name!r}, count={row.cnt}'
for row in conflicts
]
lines = [f" tenant_id={row.tenant_id!r}, name={row.name!r}, count={row.cnt}" for row in conflicts]
raise RuntimeError(
'llm_server has duplicate (tenant_id, name) pairs — cannot '
'create composite UNIQUE index. Please deduplicate manually '
'(e.g. suffix ``-dup-{id}``) and rerun `alembic upgrade head`.\n'
+ '\n'.join(lines)
"llm_server has duplicate (tenant_id, name) pairs — cannot "
"create composite UNIQUE index. Please deduplicate manually "
"(e.g. suffix ``-dup-{id}``) and rerun `alembic upgrade head`.\n" + "\n".join(lines)
)
# Drop the old global UNIQUE(name). The index name SQLAlchemy gives a
@@ -97,7 +91,7 @@ def upgrade() -> None:
op.create_index(
_NEW_UNIQUE_INDEX_NAME,
_TABLE,
['tenant_id', 'name'],
["tenant_id", "name"],
unique=True,
)
@@ -115,6 +109,6 @@ def downgrade() -> None:
op.create_index(
_OLD_UNIQUE_INDEX_NAME,
_TABLE,
['name'],
["name"],
unique=True,
)
@@ -17,60 +17,65 @@ from unittest.mock import MagicMock, patch
import pytest
MIGRATION_MOD = 'bisheng.core.database.alembic.versions.v2_5_1_f020_llm_tenant'
MIGRATION_MOD = "bisheng.core.database.alembic.versions.v2_5_1_f020_llm_tenant"
def _build_conn(duplicate_rows, index_exists):
"""Build a connection whose ``execute`` dispatches based on the SQL.
def _build_conn(duplicate_rows):
"""Build a connection whose ``execute`` returns duplicate pre-check rows.
Parameters
----------
duplicate_rows : list
Rows returned by the GROUP BY pre-check (empty = no conflicts).
index_exists : dict[str, bool]
Maps ``INDEX_NAME`` bind param to the existence answer.
"""
executed_sql = []
def _execute(stmt, params=None):
def _execute(stmt, _params=None):
sql = str(stmt)
executed_sql.append(sql)
result = MagicMock()
if 'GROUP BY' in sql:
if "GROUP BY" in sql:
result.fetchall = MagicMock(return_value=duplicate_rows)
elif 'STATISTICS' in sql:
idx_name = (params or {}).get('i')
result.scalar = MagicMock(
return_value=1 if index_exists.get(idx_name) else 0
)
else:
result.scalar = MagicMock(return_value=0)
return result
conn = MagicMock()
conn.execute = _execute
conn.executed_sql = executed_sql
return conn
def _index_exists(indexes):
return lambda _conn, _table, index_name: indexes.get(index_name, False)
def test_upgrade_rejects_duplicate_tenant_name_pairs():
"""Pre-check finds duplicates → RuntimeError mentions them and aborts."""
import importlib
mig = importlib.import_module(MIGRATION_MOD)
row = MagicMock()
row.tenant_id = 5
row.name = 'Azure-GPT-4'
row.name = "Azure-GPT-4"
row.cnt = 2
conn = _build_conn(duplicate_rows=[row], index_exists={})
conn = _build_conn(duplicate_rows=[row])
with patch.object(mig.op, 'get_bind', return_value=conn), \
patch.object(mig.op, 'drop_index') as drop, \
patch.object(mig.op, 'create_index') as create:
with (
patch.object(mig.op, "get_bind", return_value=conn),
patch.object(mig.op, "drop_index") as drop,
patch.object(mig.op, "create_index") as create,
):
with pytest.raises(RuntimeError) as excinfo:
mig.upgrade()
msg = str(excinfo.value)
assert 'duplicate' in msg.lower() or 'duplicate' in msg
assert 'tenant_id=5' in msg
assert 'Azure-GPT-4' in msg
assert "duplicate" in msg.lower() or "duplicate" in msg
assert "tenant_id=5" in msg
assert "Azure-GPT-4" in msg
assert "HAVING COUNT(*) > 1" in conn.executed_sql[0]
assert "HAVING cnt > 1" not in conn.executed_sql[0]
# No DDL should have been issued on the failure path.
drop.assert_not_called()
create.assert_not_called()
@@ -79,26 +84,28 @@ def test_upgrade_rejects_duplicate_tenant_name_pairs():
def test_upgrade_creates_composite_unique_index():
"""No conflicts + old UNIQUE(name) exists → swap to composite index."""
import importlib
mig = importlib.import_module(MIGRATION_MOD)
conn = _build_conn(
duplicate_rows=[],
index_exists={
'name': True, # legacy UNIQUE(name)
'uk_llm_server_tenant_name': False, # composite not yet present
},
)
conn = _build_conn(duplicate_rows=[])
indexes = {
"name": True, # legacy UNIQUE(name)
"uk_llm_server_tenant_name": False, # composite not yet present
}
with patch.object(mig.op, 'get_bind', return_value=conn), \
patch.object(mig.op, 'drop_index') as drop, \
patch.object(mig.op, 'create_index') as create:
with (
patch.object(mig.op, "get_bind", return_value=conn),
patch.object(mig, "index_exists", side_effect=_index_exists(indexes)),
patch.object(mig.op, "drop_index") as drop,
patch.object(mig.op, "create_index") as create,
):
mig.upgrade()
drop.assert_called_once_with('name', table_name='llm_server')
drop.assert_called_once_with("name", table_name="llm_server")
create.assert_called_once_with(
'uk_llm_server_tenant_name',
'llm_server',
['tenant_id', 'name'],
"uk_llm_server_tenant_name",
"llm_server",
["tenant_id", "name"],
unique=True,
)
@@ -106,22 +113,24 @@ def test_upgrade_creates_composite_unique_index():
def test_upgrade_skips_drop_when_legacy_index_absent():
"""Fresh v2.5.1 install (no legacy UNIQUE(name)) → just create composite."""
import importlib
mig = importlib.import_module(MIGRATION_MOD)
conn = _build_conn(
duplicate_rows=[],
index_exists={'name': False, 'uk_llm_server_tenant_name': False},
)
conn = _build_conn(duplicate_rows=[])
indexes = {"name": False, "uk_llm_server_tenant_name": False}
with patch.object(mig.op, 'get_bind', return_value=conn), \
patch.object(mig.op, 'drop_index') as drop, \
patch.object(mig.op, 'create_index') as create:
with (
patch.object(mig.op, "get_bind", return_value=conn),
patch.object(mig, "index_exists", side_effect=_index_exists(indexes)),
patch.object(mig.op, "drop_index") as drop,
patch.object(mig.op, "create_index") as create,
):
mig.upgrade()
drop.assert_not_called()
create.assert_called_once_with(
'uk_llm_server_tenant_name',
'llm_server',
['tenant_id', 'name'],
"uk_llm_server_tenant_name",
"llm_server",
["tenant_id", "name"],
unique=True,
)
@@ -15,16 +15,17 @@ from bisheng.database.models.department import Department
from bisheng.database.models.tenant import Tenant
@pytest.fixture(scope='module')
@pytest.fixture(scope="module")
def init_engine():
"""SQLite engine with tenant + department tables."""
engine = create_engine(
'sqlite://',
connect_args={'check_same_thread': False},
"sqlite://",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
with engine.begin() as conn:
conn.execute(text("""
conn.execute(
text("""
CREATE TABLE IF NOT EXISTS tenant (
id INTEGER PRIMARY KEY AUTOINCREMENT,
tenant_code VARCHAR(64) NOT NULL UNIQUE,
@@ -43,8 +44,10 @@ def init_engine():
create_time DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
update_time DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL
)
"""))
conn.execute(text("""
""")
)
conn.execute(
text("""
CREATE TABLE IF NOT EXISTS department (
id INTEGER PRIMARY KEY AUTOINCREMENT,
dept_id VARCHAR(64) NOT NULL UNIQUE,
@@ -70,7 +73,8 @@ def init_engine():
update_time DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
UNIQUE(source, external_id)
)
"""))
""")
)
yield engine
engine.dispose()
@@ -89,15 +93,14 @@ def session(init_engine):
class TestInitDefaultRootDepartment:
def test_init_creates_root_department(self, session):
"""AC-17: Default tenant (id=1) gets a root department on startup."""
# Setup: create default tenant with no root_dept_id
tenant = Tenant(
id=1,
tenant_code='default',
tenant_name='Default Tenant',
status='active',
tenant_code="default",
tenant_name="Default Tenant",
status="active",
)
session.add(tenant)
session.commit()
@@ -106,19 +109,20 @@ class TestInitDefaultRootDepartment:
# Simulate what _init_default_root_department does (sync version)
dept = Department(
dept_id='BS@root',
name='默认组织',
dept_id="BS@root",
name="默认组织",
parent_id=None,
tenant_id=1,
path='',
source='local',
status='active',
path="",
source="local",
external_id="BS@root",
status="active",
)
session.add(dept)
session.flush()
session.refresh(dept)
dept.path = f'/{dept.id}/'
dept.path = f"/{dept.id}/"
tenant.root_dept_id = dept.id
session.commit()
session.refresh(tenant)
@@ -127,8 +131,9 @@ class TestInitDefaultRootDepartment:
# Verify
assert dept.id is not None
assert dept.parent_id is None
assert dept.path == f'/{dept.id}/'
assert dept.name == '默认组织'
assert dept.path == f"/{dept.id}/"
assert dept.name == "默认组织"
assert dept.external_id == "BS@root"
assert tenant.root_dept_id == dept.id
def test_init_idempotent(self, session):
@@ -136,26 +141,27 @@ class TestInitDefaultRootDepartment:
# Setup: create tenant + root department
tenant = Tenant(
id=1,
tenant_code='default',
tenant_name='Default Tenant',
status='active',
tenant_code="default",
tenant_name="Default Tenant",
status="active",
)
session.add(tenant)
session.commit()
dept = Department(
dept_id='BS@root',
name='默认组织',
dept_id="BS@root",
name="默认组织",
parent_id=None,
tenant_id=1,
path='',
source='local',
status='active',
path="",
source="local",
external_id="BS@root",
status="active",
)
session.add(dept)
session.flush()
session.refresh(dept)
dept.path = f'/{dept.id}/'
dept.path = f"/{dept.id}/"
tenant.root_dept_id = dept.id
session.commit()
session.refresh(tenant)
@@ -170,8 +176,131 @@ class TestInitDefaultRootDepartment:
select(Department).where(
Department.parent_id.is_(None),
Department.tenant_id == 1,
Department.status == 'active',
Department.status == "active",
)
).all()
assert len(roots) == 1
assert roots[0].id == first_dept_id
class _FirstResult:
def __init__(self, value):
self._value = value
def first(self):
return self._value
class _FakeAsyncSession:
def __init__(self, responses):
self._responses = iter(responses)
self.added = []
self.commit_snapshots = []
self._next_department_id = 100
async def exec(self, _statement):
return _FirstResult(next(self._responses))
def add(self, value):
self.added.append(value)
async def flush(self):
for value in self.added:
if isinstance(value, Department) and value.id is None:
value.id = self._next_department_id
self._next_department_id += 1
async def refresh(self, _value):
return None
async def commit(self):
snapshot = {value.dept_id: value.external_id for value in self.added if isinstance(value, Department)}
self.commit_snapshots.append(snapshot)
async def test_default_departments_use_distinct_external_ids():
from bisheng.common.init_data import _init_default_root_department
tenant = Tenant(
id=1,
tenant_code="default",
tenant_name="Default Tenant",
status="active",
)
session = _FakeAsyncSession([tenant, None])
await _init_default_root_department(session)
departments = {value.dept_id: value for value in session.added if isinstance(value, Department)}
assert departments["BS@root"].external_id == "BS@root"
assert departments["BS@guest"].external_id == "BS@guest"
async def test_existing_null_root_identity_is_committed_before_guest_insert():
from bisheng.common.init_data import _init_default_root_department
tenant = Tenant(
id=1,
tenant_code="default",
tenant_name="Default Tenant",
root_dept_id=24,
status="active",
)
root = Department(
id=24,
dept_id="BS@root",
name="默认组织",
parent_id=None,
tenant_id=1,
path="/24/",
source="local",
external_id=None,
status="active",
)
session = _FakeAsyncSession([tenant, root, None])
await _init_default_root_department(session)
assert session.commit_snapshots[0] == {"BS@root": "BS@root"}
guest = next(value for value in session.added if isinstance(value, Department) and value.dept_id == "BS@guest")
assert guest.external_id == "BS@guest"
async def test_existing_null_guest_identity_is_repaired():
from bisheng.common.init_data import _init_default_root_department
tenant = Tenant(
id=1,
tenant_code="default",
tenant_name="Default Tenant",
root_dept_id=24,
status="active",
)
root = Department(
id=24,
dept_id="BS@root",
name="默认组织",
parent_id=None,
tenant_id=1,
path="/24/",
source="local",
external_id="BS@root",
status="active",
)
guest = Department(
id=25,
dept_id="BS@guest",
name="临时访客",
parent_id=24,
tenant_id=1,
path="/24/25/",
source="local",
external_id=None,
status="active",
)
session = _FakeAsyncSession([tenant, root, guest])
await _init_default_root_department(session)
assert guest.external_id == "BS@guest"
assert session.commit_snapshots[-1] == {"BS@guest": "BS@guest"}