mirror of
https://github.com/dataelement/bisheng.git
synced 2026-08-30 17:58:00 +08:00
fix(database): harden DM fresh deployment startup
This commit is contained in:
@@ -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"}
|
||||
|
||||
Reference in New Issue
Block a user