File size: 11,573 Bytes
979853c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 | """
Tests for workspace isolation during PostgreSQL migration.
This test module verifies that setup_table() properly filters migration data
by workspace, preventing cross-workspace data leakage during legacy table migration.
Critical Bug: Migration copied ALL records from legacy table regardless of workspace,
causing workspace A to receive workspace B's data, violating multi-tenant isolation.
"""
import pytest
from unittest.mock import AsyncMock
from lightrag.kg.postgres_impl import PGVectorStorage
class TestWorkspaceMigrationIsolation:
"""Test suite for workspace-scoped migration in PostgreSQL."""
async def test_migration_filters_by_workspace(self):
"""
Test that migration only copies data from the specified workspace.
Scenario: Legacy table contains data from multiple workspaces.
Migrate only workspace_a's data to new table.
Expected: New table contains only workspace_a data, workspace_b data excluded.
"""
db = AsyncMock()
# Configure mock return values to avoid unawaited coroutine warnings
db._create_vector_index.return_value = None
# Track state for new table count (starts at 0, increases after migration)
new_table_record_count = {"count": 0}
# Mock table existence checks
async def table_exists_side_effect(db_instance, name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False # New table doesn't exist initially
return False
# Mock data for workspace_a
mock_records_a = [
{
"id": "a1",
"workspace": "workspace_a",
"content": "content_a1",
"content_vector": [0.1] * 1536,
},
{
"id": "a2",
"workspace": "workspace_a",
"content": "content_a2",
"content_vector": [0.2] * 1536,
},
]
# Mock query responses
async def query_side_effect(sql, params, **kwargs):
multirows = kwargs.get("multirows", False)
sql_upper = sql.upper()
# Count query for new table workspace data (verification before migration)
if (
"COUNT(*)" in sql_upper
and "MODEL_1536D" in sql_upper
and "WHERE WORKSPACE" in sql_upper
):
return new_table_record_count # Initially 0
# Count query with workspace filter (legacy table) - for workspace count
elif "COUNT(*)" in sql_upper and "WHERE WORKSPACE" in sql_upper:
if params and params[0] == "workspace_a":
return {"count": 2} # workspace_a has 2 records
elif params and params[0] == "workspace_b":
return {"count": 3} # workspace_b has 3 records
return {"count": 0}
# Count query for legacy table (total, no workspace filter)
elif (
"COUNT(*)" in sql_upper
and "LIGHTRAG" in sql_upper
and "WHERE WORKSPACE" not in sql_upper
):
return {"count": 5} # Total records in legacy
# SELECT with workspace filter for migration (multirows)
elif "SELECT" in sql_upper and "FROM" in sql_upper and multirows:
workspace = params[0] if params else None
if workspace == "workspace_a":
# Handle keyset pagination: check for "id >" pattern
if "id >" in sql.lower():
# Keyset pagination: params = [workspace, last_id, limit]
last_id = params[1] if len(params) > 1 else None
# Find records after last_id
found_idx = -1
for i, rec in enumerate(mock_records_a):
if rec["id"] == last_id:
found_idx = i
break
if found_idx >= 0:
return mock_records_a[found_idx + 1 :]
return []
else:
# First batch: params = [workspace, limit]
return mock_records_a
return [] # No data for other workspaces
return {}
db.query.side_effect = query_side_effect
db.execute = AsyncMock()
# Mock check_table_exists on db
async def check_table_exists_side_effect(name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False # New table doesn't exist initially
return False
db.check_table_exists = AsyncMock(side_effect=check_table_exists_side_effect)
# Track migration through _run_with_retry calls
migration_executed = []
async def mock_run_with_retry(operation, *args, **kwargs):
migration_executed.append(True)
new_table_record_count["count"] = 2 # Simulate 2 records migrated
return None
db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
# Migrate for workspace_a only - correct parameter order
await PGVectorStorage.setup_table(
db,
"LIGHTRAG_DOC_CHUNKS_model_1536d",
workspace="workspace_a", # CRITICAL: Only migrate workspace_a
embedding_dim=1536,
legacy_table_name="LIGHTRAG_DOC_CHUNKS",
base_table="LIGHTRAG_DOC_CHUNKS",
)
# Verify the migration was triggered
assert (
len(migration_executed) > 0
), "Migration should have been executed for workspace_a"
async def test_migration_without_workspace_raises_error(self):
"""
Test that migration without workspace parameter raises ValueError.
Scenario: setup_table called without workspace parameter.
Expected: ValueError is raised because workspace is required.
"""
db = AsyncMock()
# workspace is now a required parameter - calling with None should raise ValueError
with pytest.raises(ValueError, match="workspace must be provided"):
await PGVectorStorage.setup_table(
db,
"lightrag_doc_chunks_model_1536d",
workspace=None, # No workspace - should raise ValueError
embedding_dim=1536,
legacy_table_name="lightrag_doc_chunks",
base_table="lightrag_doc_chunks",
)
async def test_no_cross_workspace_contamination(self):
"""
Test that workspace B's migration doesn't include workspace A's data.
Scenario: Migration for workspace_b only.
Expected: Only workspace_b data is queried, workspace_a data excluded.
"""
db = AsyncMock()
# Configure mock return values to avoid unawaited coroutine warnings
db._create_vector_index.return_value = None
# Track which workspace is being queried
queried_workspace = None
new_table_count = {"count": 0}
# Mock data for workspace_b
mock_records_b = [
{
"id": "b1",
"workspace": "workspace_b",
"content": "content_b1",
"content_vector": [0.3] * 1536,
},
]
async def table_exists_side_effect(db_instance, name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False
return False
async def query_side_effect(sql, params, **kwargs):
nonlocal queried_workspace
multirows = kwargs.get("multirows", False)
sql_upper = sql.upper()
# Count query for new table workspace data (should be 0 initially)
if (
"COUNT(*)" in sql_upper
and "MODEL_1536D" in sql_upper
and "WHERE WORKSPACE" in sql_upper
):
return new_table_count
# Count query with workspace filter (legacy table)
elif "COUNT(*)" in sql_upper and "WHERE WORKSPACE" in sql_upper:
queried_workspace = params[0] if params else None
return {"count": 1} # 1 record for the queried workspace
# Count query for legacy table total (no workspace filter)
elif (
"COUNT(*)" in sql_upper
and "LIGHTRAG" in sql_upper
and "WHERE WORKSPACE" not in sql_upper
):
return {"count": 3} # 3 total records in legacy
# SELECT with workspace filter for migration (multirows)
elif "SELECT" in sql_upper and "FROM" in sql_upper and multirows:
workspace = params[0] if params else None
if workspace == "workspace_b":
# Handle keyset pagination: check for "id >" pattern
if "id >" in sql.lower():
# Keyset pagination: params = [workspace, last_id, limit]
last_id = params[1] if len(params) > 1 else None
# Find records after last_id
found_idx = -1
for i, rec in enumerate(mock_records_b):
if rec["id"] == last_id:
found_idx = i
break
if found_idx >= 0:
return mock_records_b[found_idx + 1 :]
return []
else:
# First batch: params = [workspace, limit]
return mock_records_b
return [] # No data for other workspaces
return {}
db.query.side_effect = query_side_effect
db.execute = AsyncMock()
# Mock check_table_exists on db
async def check_table_exists_side_effect(name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False
return False
db.check_table_exists = AsyncMock(side_effect=check_table_exists_side_effect)
# Track migration through _run_with_retry calls
migration_executed = []
async def mock_run_with_retry(operation, *args, **kwargs):
migration_executed.append(True)
new_table_count["count"] = 1 # Simulate migration
return None
db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
# Migrate workspace_b - correct parameter order
await PGVectorStorage.setup_table(
db,
"LIGHTRAG_DOC_CHUNKS_model_1536d",
workspace="workspace_b", # Only migrate workspace_b
embedding_dim=1536,
legacy_table_name="LIGHTRAG_DOC_CHUNKS",
base_table="LIGHTRAG_DOC_CHUNKS",
)
# Verify only workspace_b was queried
assert queried_workspace == "workspace_b", "Should only query workspace_b"
|