mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-18 00:10:59 +00:00
chore: merge master into dev/4.11.x
This commit is contained in:
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_membership_source_migration_backfills_existing_rows_as_local_and_enforces_constraint(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "membership-source.db"}')
|
||||
try:
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sa.text(
|
||||
"""
|
||||
CREATE TABLE workspace_memberships (
|
||||
uuid VARCHAR(36) PRIMARY KEY,
|
||||
workspace_uuid VARCHAR(36) NOT NULL,
|
||||
account_uuid VARCHAR(36) NOT NULL,
|
||||
role VARCHAR(32) NOT NULL,
|
||||
status VARCHAR(32) NOT NULL,
|
||||
projection_revision BIGINT NOT NULL DEFAULT 0,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
await connection.execute(
|
||||
sa.text(
|
||||
"""
|
||||
INSERT INTO workspace_memberships
|
||||
(uuid, workspace_uuid, account_uuid, role, status, projection_revision)
|
||||
VALUES
|
||||
('00000000-0000-4000-8000-000000000001', 'workspace', 'local-account',
|
||||
'viewer', 'active', 0),
|
||||
('00000000-0000-4000-8000-000000000002', 'workspace', 'cloud-account',
|
||||
'viewer', 'active', 0)
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
await run_alembic_stamp(engine, '0019_single_workspace_owner')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
async with engine.connect() as connection:
|
||||
rows = (
|
||||
await connection.execute(sa.text('SELECT uuid, source FROM workspace_memberships ORDER BY uuid'))
|
||||
).all()
|
||||
columns = await connection.run_sync(
|
||||
lambda sync_connection: {
|
||||
column['name']: column
|
||||
for column in sa.inspect(sync_connection).get_columns('workspace_memberships')
|
||||
}
|
||||
)
|
||||
assert rows == [
|
||||
('00000000-0000-4000-8000-000000000001', 'local'),
|
||||
('00000000-0000-4000-8000-000000000002', 'local'),
|
||||
]
|
||||
assert columns['source']['nullable'] is False
|
||||
|
||||
with pytest.raises(sa.exc.IntegrityError):
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sa.text("UPDATE workspace_memberships SET source = 'guessed-from-user-source'")
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@@ -9,8 +9,11 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
@@ -105,7 +108,7 @@ class TestSQLiteMigrationUpgrade:
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
|
||||
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||
assert _get_script_head() == '0020_merge_agent_cloud_heads'
|
||||
assert _get_script_head() == '0022_merge_agent_reasoning_heads'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upgrade_from_development_workspace_head_to_merged_head(self, sqlite_engine):
|
||||
@@ -117,7 +120,18 @@ class TestSQLiteMigrationUpgrade:
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
|
||||
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||
assert _get_script_head() == '0020_merge_agent_cloud_heads'
|
||||
assert _get_script_head() == '0022_merge_agent_reasoning_heads'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upgrade_from_reasoning_config_head_to_merged_head(self, sqlite_engine):
|
||||
"""A database that already ran the feature migration must remain upgradable."""
|
||||
async with sqlite_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
await run_alembic_stamp(sqlite_engine, '0018_llm_reasoning_config')
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
|
||||
assert await get_alembic_current(sqlite_engine) == '0022_merge_agent_reasoning_heads'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
|
||||
@@ -214,6 +228,66 @@ class TestSQLiteMigrationUpgrade:
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reasoning_config_migrates_existing_models(self, sqlite_engine):
|
||||
"""Upgrade from 0017 backfills reasoning config and keeps a database default."""
|
||||
async with sqlite_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
"""
|
||||
CREATE TABLE llm_models (
|
||||
uuid VARCHAR(255) PRIMARY KEY,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
provider_uuid VARCHAR(255) NOT NULL,
|
||||
abilities JSON NOT NULL,
|
||||
context_length INTEGER,
|
||||
extra_args JSON NOT NULL,
|
||||
prefered_ranking INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO llm_models (
|
||||
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
||||
) VALUES (
|
||||
'existing-model', 'Existing Model', 'provider', '[]', '{}', 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
await run_alembic_stamp(sqlite_engine, '0017_oss_workspace_identity')
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
|
||||
async with sqlite_engine.begin() as conn:
|
||||
columns = await conn.run_sync(lambda sync_conn: sqlalchemy.inspect(sync_conn).get_columns('llm_models'))
|
||||
reasoning_column = next(column for column in columns if column['name'] == 'reasoning_config')
|
||||
assert reasoning_column['nullable'] is False
|
||||
|
||||
existing_value = (
|
||||
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'existing-model'"))
|
||||
).scalar_one()
|
||||
assert json.loads(existing_value) == {'level': 'provider_default'}
|
||||
|
||||
await conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO llm_models (
|
||||
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
||||
) VALUES (
|
||||
'new-model', 'New Model', 'provider', '[]', '{}', 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
new_value = (
|
||||
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'new-model'"))
|
||||
).scalar_one()
|
||||
assert json.loads(new_value) == {'level': 'provider_default'}
|
||||
|
||||
|
||||
class TestSQLiteMigrationFreshDatabase:
|
||||
"""Tests for fresh database workflow."""
|
||||
|
||||
@@ -85,6 +85,32 @@ async def clean_database(postgres_engine: AsyncEngine):
|
||||
await clean()
|
||||
|
||||
|
||||
async def test_upgrade_adds_3072_dimension_index_and_constraint(
|
||||
postgres_engine: AsyncEngine,
|
||||
clean_database,
|
||||
) -> None:
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
await run_alembic_stamp(postgres_engine, '0010_scope_resources')
|
||||
await run_alembic_upgrade(postgres_engine, 'head')
|
||||
|
||||
async with postgres_engine.connect() as conn:
|
||||
constraint = await conn.scalar(
|
||||
text(
|
||||
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
|
||||
"WHERE conrelid = 'langbot_vectors'::regclass "
|
||||
"AND conname = 'ck_langbot_vectors_embedding_dimension_enabled'"
|
||||
)
|
||||
)
|
||||
assert '3072' in constraint
|
||||
index_definition = await conn.scalar(
|
||||
text("SELECT indexdef FROM pg_indexes WHERE indexname = 'ix_langbot_vectors_hnsw_cosine_3072'")
|
||||
)
|
||||
assert 'halfvec(3072)' in index_definition
|
||||
assert 'halfvec_cosine_ops' in index_definition
|
||||
|
||||
|
||||
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
|
||||
postgres_url: str,
|
||||
postgres_engine: AsyncEngine,
|
||||
|
||||
@@ -92,7 +92,7 @@ def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_
|
||||
'use': 'pgvector',
|
||||
'pgvector': {
|
||||
'use_business_database': True,
|
||||
'allowed_dimensions': [384, 512, 768, 1024, 1536],
|
||||
'allowed_dimensions': [384, 512, 768, 1024, 1536, 3072],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user