mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge pull request #19409 from BerriAI/revert-19261-feat/redis-migration-lock-safe
Revert "feat: Add Redis-based migration lock with bug fixes"
This commit is contained in:
Sameer Kankute
GitHub
commit
12c556e485
3 changed files with 25 additions and 452 deletions
|
|
@ -11,11 +11,6 @@ from typing import Optional
|
|||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
||||
try:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
except ImportError:
|
||||
RedisCache = None # type: ignore
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
if value is None:
|
||||
|
|
@ -23,154 +18,6 @@ def str_to_bool(value: Optional[str]) -> bool:
|
|||
return value.lower() in ("true", "1", "t", "y", "yes")
|
||||
|
||||
|
||||
class MigrationLockManager:
|
||||
"""Redis-based lock manager for database migrations.
|
||||
|
||||
Prevents concurrent Prisma migrations in multi-pod deployments by using
|
||||
a distributed lock. Only one pod can hold the lock and run migrations at a time.
|
||||
"""
|
||||
|
||||
MIGRATION_LOCK_KEY = "migration_lock"
|
||||
LOCK_TTL_SECONDS = 300 # 5 minutes TTL
|
||||
|
||||
def __init__(self, redis_cache: Optional["RedisCache"] = None):
|
||||
"""Initialize the migration lock manager.
|
||||
|
||||
Args:
|
||||
redis_cache: Optional RedisCache instance for distributed locking.
|
||||
If None, migrations run without lock protection (single instance mode).
|
||||
"""
|
||||
self.redis_cache = redis_cache
|
||||
self.lock_acquired = False
|
||||
self.pod_id = f"pod_{os.getpid()}_{int(time.time())}"
|
||||
|
||||
def _get_redis_lock_key(self) -> str:
|
||||
"""Get Redis lock key for migration."""
|
||||
return f"migration_lock:{self.MIGRATION_LOCK_KEY}"
|
||||
|
||||
def acquire_lock(self) -> bool:
|
||||
"""Acquire migration lock using Redis SET NX.
|
||||
|
||||
Returns:
|
||||
bool: True if lock acquired, False otherwise.
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
logger.warning(
|
||||
"Redis cache is not available, running migration without lock protection"
|
||||
)
|
||||
self.lock_acquired = True
|
||||
return True
|
||||
|
||||
try:
|
||||
lock_key = self._get_redis_lock_key()
|
||||
|
||||
# FIX: Use native Redis SET NX instead of RedisCache.set_cache()
|
||||
# Original bug: set_cache() doesn't support nx parameter and returns None
|
||||
# Fixed: Use redis_client.set() directly which returns True/False
|
||||
acquired = self.redis_cache.redis_client.set(
|
||||
name=lock_key,
|
||||
value=self.pod_id,
|
||||
nx=True, # Only set if key doesn't exist
|
||||
ex=self.LOCK_TTL_SECONDS, # Set expiration time
|
||||
)
|
||||
|
||||
if acquired:
|
||||
self.lock_acquired = True
|
||||
logger.info(f"Migration lock acquired by pod {self.pod_id}")
|
||||
return True
|
||||
else:
|
||||
logger.info("Migration lock is already held by another pod")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to acquire migration lock: {e}")
|
||||
return False
|
||||
|
||||
def wait_for_lock_release(
|
||||
self, check_interval: int = 5, max_wait: int = 300
|
||||
) -> bool:
|
||||
"""Wait for another process to release the lock.
|
||||
|
||||
Args:
|
||||
check_interval: Seconds to wait between lock acquisition attempts.
|
||||
max_wait: Maximum seconds to wait for lock release.
|
||||
|
||||
Returns:
|
||||
bool: True if lock acquired after waiting, False if timeout.
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
logger.warning("Redis cache is not available, cannot wait for lock")
|
||||
return False
|
||||
|
||||
logger.info(f"Waiting for migration lock to be released (max {max_wait}s)...")
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < max_wait:
|
||||
# Try to acquire lock using the public acquire_lock method
|
||||
if self.acquire_lock():
|
||||
logger.info(
|
||||
f"Migration lock acquired after waiting by pod {self.pod_id}"
|
||||
)
|
||||
return True
|
||||
|
||||
time.sleep(check_interval)
|
||||
|
||||
logger.warning(f"Failed to acquire migration lock within {max_wait} seconds")
|
||||
return False
|
||||
|
||||
def release_lock(self):
|
||||
"""Release migration lock atomically using Lua script.
|
||||
|
||||
FIX: Use Lua script for atomic compare-and-delete to prevent race conditions.
|
||||
Original bug: Non-atomic GET then DELETE allows another pod to acquire lock
|
||||
between the GET and DELETE operations.
|
||||
"""
|
||||
if not self.lock_acquired or self.redis_cache is None:
|
||||
return
|
||||
|
||||
try:
|
||||
lock_key = self._get_redis_lock_key()
|
||||
|
||||
# FIX: Use Lua script for atomic compare-and-delete
|
||||
# This prevents race condition where:
|
||||
# 1. Pod A reads lock value (sees its own pod_id)
|
||||
# 2. Lock TTL expires
|
||||
# 3. Pod B acquires lock
|
||||
# 4. Pod A deletes lock (deletes Pod B's lock!)
|
||||
lua_script = """
|
||||
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||||
return redis.call("DEL", KEYS[1])
|
||||
else
|
||||
return 0
|
||||
end
|
||||
"""
|
||||
|
||||
result = self.redis_cache.redis_client.eval(
|
||||
lua_script,
|
||||
1, # Number of keys
|
||||
lock_key, # KEYS[1]
|
||||
self.pod_id, # ARGV[1]
|
||||
)
|
||||
|
||||
if result == 1:
|
||||
logger.info(f"Migration lock released by pod {self.pod_id}")
|
||||
else:
|
||||
logger.warning(f"Pod {self.pod_id} cannot release lock (not owner)")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to release migration lock: {e}")
|
||||
finally:
|
||||
self.lock_acquired = False
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - acquire lock when entering with statement."""
|
||||
self.acquire_lock()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - release lock when exiting with statement."""
|
||||
self.release_lock()
|
||||
|
||||
|
||||
def _get_prisma_env() -> dict:
|
||||
"""Get environment variables for Prisma, handling offline mode if configured."""
|
||||
|
|
@ -178,9 +25,7 @@ def _get_prisma_env() -> dict:
|
|||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
# These env vars prevent Prisma from attempting downloads
|
||||
prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true"
|
||||
prisma_env["NPM_CONFIG_CACHE"] = os.getenv(
|
||||
"NPM_CONFIG_CACHE", "/app/.cache/npm"
|
||||
)
|
||||
prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm")
|
||||
return prisma_env
|
||||
|
||||
|
||||
|
|
@ -189,28 +34,29 @@ def _get_prisma_command() -> str:
|
|||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
# Primary location where Prisma Python package installs the CLI
|
||||
default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma"
|
||||
|
||||
|
||||
# Check if custom path is provided (for flexibility)
|
||||
custom_cli_path = os.getenv("PRISMA_CLI_PATH")
|
||||
if custom_cli_path and os.path.exists(custom_cli_path):
|
||||
logger.info(f"Using custom Prisma CLI at {custom_cli_path}")
|
||||
return custom_cli_path
|
||||
|
||||
|
||||
# Check the default location
|
||||
if os.path.exists(default_cli_path):
|
||||
logger.info(f"Using cached Prisma CLI at {default_cli_path}")
|
||||
return default_cli_path
|
||||
|
||||
|
||||
# If not found, log warning and fall back
|
||||
logger.warning(
|
||||
f"Prisma CLI not found at {default_cli_path}. "
|
||||
"Falling back to Python wrapper (may attempt downloads)"
|
||||
)
|
||||
|
||||
|
||||
# Fall back to the Python wrapper (will work in online mode)
|
||||
return "prisma"
|
||||
|
||||
|
||||
|
||||
class ProxyExtrasDBManager:
|
||||
@staticmethod
|
||||
def _get_prisma_dir() -> str:
|
||||
|
|
@ -273,7 +119,7 @@ class ProxyExtrasDBManager:
|
|||
stdout=open(migration_file, "w"),
|
||||
check=True,
|
||||
timeout=30,
|
||||
env=prisma_env,
|
||||
env=prisma_env
|
||||
)
|
||||
|
||||
# 3. Mark the migration as applied since it represents current state
|
||||
|
|
@ -288,7 +134,7 @@ class ProxyExtrasDBManager:
|
|||
],
|
||||
check=True,
|
||||
timeout=30,
|
||||
env=prisma_env,
|
||||
env=prisma_env
|
||||
)
|
||||
|
||||
return True
|
||||
|
|
@ -313,20 +159,14 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _roll_back_migration(migration_name: str):
|
||||
"""Mark a specific migration as rolled back"""
|
||||
# Set up environment for offline mode if configured
|
||||
# Set up environment for offline mode if configured
|
||||
prisma_env = _get_prisma_env()
|
||||
subprocess.run(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
"resolve",
|
||||
"--rolled-back",
|
||||
migration_name,
|
||||
],
|
||||
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
|
||||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
env=prisma_env
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -338,7 +178,7 @@ class ProxyExtrasDBManager:
|
|||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
env=prisma_env
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -408,7 +248,7 @@ class ProxyExtrasDBManager:
|
|||
if not database_url:
|
||||
logger.error("DATABASE_URL not set")
|
||||
return
|
||||
|
||||
|
||||
diff_dir = (
|
||||
Path(migrations_dir)
|
||||
/ "migrations"
|
||||
|
|
@ -443,7 +283,7 @@ class ProxyExtrasDBManager:
|
|||
check=True,
|
||||
timeout=60,
|
||||
stdout=f,
|
||||
env=_get_prisma_env(),
|
||||
env=_get_prisma_env()
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(f"Failed to generate migration diff: {e.stderr}")
|
||||
|
|
@ -473,7 +313,7 @@ class ProxyExtrasDBManager:
|
|||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
env=_get_prisma_env()
|
||||
)
|
||||
logger.info(f"prisma db execute stdout: {result.stdout}")
|
||||
logger.info("✅ Migration diff applied successfully")
|
||||
|
|
@ -491,18 +331,12 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
logger.info(f"Resolving migration: {migration_name}")
|
||||
subprocess.run(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
"resolve",
|
||||
"--applied",
|
||||
migration_name,
|
||||
],
|
||||
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
|
||||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
env=_get_prisma_env()
|
||||
)
|
||||
logger.debug(f"Resolved migration: {migration_name}")
|
||||
except subprocess.CalledProcessError as e:
|
||||
|
|
@ -512,57 +346,19 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(
|
||||
use_migrate: bool = False, redis_cache: Optional["RedisCache"] = None
|
||||
) -> bool:
|
||||
def setup_database(use_migrate: bool = False) -> bool:
|
||||
"""
|
||||
Set up the database using either prisma migrate or prisma db push.
|
||||
Uses migrations from litellm-proxy-extras package.
|
||||
In multi-instance environment, use redis lock to prevent concurrent execution.
|
||||
Set up the database using either prisma migrate or prisma db push
|
||||
Uses migrations from litellm-proxy-extras package
|
||||
|
||||
Args:
|
||||
schema_path (str): Path to the Prisma schema file
|
||||
use_migrate (bool): Whether to use prisma migrate instead of db push
|
||||
redis_cache: Redis cache instance for distributed locking
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
logger.error("DATABASE_URL environment variable is not set")
|
||||
return False
|
||||
|
||||
# Use MigrationLockManager to prevent concurrent migration execution
|
||||
with MigrationLockManager(redis_cache) as lock_manager:
|
||||
# Lock is already acquired in __enter__, check if it was successful
|
||||
if not lock_manager.lock_acquired:
|
||||
# Cannot acquire lock, another process is running migration
|
||||
logger.info(
|
||||
"Another pod is running migration, waiting for completion..."
|
||||
)
|
||||
|
||||
# Wait for other process to complete migration
|
||||
if not lock_manager.wait_for_lock_release():
|
||||
logger.error("Failed to acquire migration lock after waiting")
|
||||
return False
|
||||
|
||||
# Successfully acquired lock, proceed with migration
|
||||
logger.info("Acquired migration lock, proceeding with migration")
|
||||
return ProxyExtrasDBManager._execute_migration(use_migrate, schema_path)
|
||||
|
||||
@staticmethod
|
||||
def _execute_migration(use_migrate: bool, schema_path: str) -> bool:
|
||||
"""Execute the actual migration.
|
||||
|
||||
Args:
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
schema_path: Path to the Prisma schema file
|
||||
|
||||
Returns:
|
||||
bool: True if migration was successful, False otherwise
|
||||
"""
|
||||
for attempt in range(4):
|
||||
original_dir = os.getcwd()
|
||||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
|
|
@ -579,7 +375,7 @@ class ProxyExtrasDBManager:
|
|||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
env=_get_prisma_env()
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
|
||||
|
|
@ -617,7 +413,7 @@ class ProxyExtrasDBManager:
|
|||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
env=_get_prisma_env()
|
||||
)
|
||||
logger.info(
|
||||
f"✅ Migration {failed_migration} marked as rolled back... retrying"
|
||||
|
|
|
|||
|
|
@ -383,18 +383,7 @@ class PrismaManager:
|
|||
|
||||
prisma_dir = PrismaManager._get_prisma_dir()
|
||||
|
||||
# Import redis_usage_cache for distributed locking
|
||||
try:
|
||||
from litellm.proxy.proxy_server import redis_usage_cache
|
||||
except ImportError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not import redis_usage_cache, migrations will run without distributed locking"
|
||||
)
|
||||
redis_usage_cache = None
|
||||
|
||||
return ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=use_migrate, redis_cache=redis_usage_cache
|
||||
)
|
||||
return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate)
|
||||
else:
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
|
|
|
|||
|
|
@ -1,13 +1,11 @@
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager, MigrationLockManager
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
|
||||
|
||||
def test_custom_prisma_dir(monkeypatch):
|
||||
|
|
@ -127,213 +125,3 @@ class TestErrorClassificationPriority:
|
|||
error_message = "connection timeout"
|
||||
assert ProxyExtrasDBManager._is_permission_error(error_message) is False
|
||||
assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False
|
||||
|
||||
|
||||
class TestMigrationLockManager:
|
||||
"""Test cases for MigrationLockManager"""
|
||||
|
||||
def test_acquire_lock_success(self):
|
||||
"""Test successful lock acquisition using Redis SET NX"""
|
||||
mock_redis = Mock()
|
||||
# FIX VERIFICATION: redis_client.set() returns True for successful SET NX
|
||||
mock_redis.redis_client.set.return_value = True
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
result = lock_manager.acquire_lock()
|
||||
|
||||
assert result is True
|
||||
assert lock_manager.lock_acquired is True
|
||||
# Verify SET NX was called with correct parameters
|
||||
mock_redis.redis_client.set.assert_called_once()
|
||||
call_args = mock_redis.redis_client.set.call_args
|
||||
assert call_args[1]["nx"] is True # SET NX parameter
|
||||
assert call_args[1]["ex"] == 300 # TTL
|
||||
|
||||
def test_acquire_lock_failure(self):
|
||||
"""Test lock acquisition failure when lock is already held"""
|
||||
mock_redis = Mock()
|
||||
# FIX VERIFICATION: redis_client.set() returns False when key exists
|
||||
mock_redis.redis_client.set.return_value = False
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
result = lock_manager.acquire_lock()
|
||||
|
||||
assert result is False
|
||||
assert lock_manager.lock_acquired is False
|
||||
|
||||
def test_acquire_lock_without_redis(self):
|
||||
"""Test lock acquisition without Redis (graceful fallback)"""
|
||||
lock_manager = MigrationLockManager(redis_cache=None)
|
||||
result = lock_manager.acquire_lock()
|
||||
|
||||
assert result is True
|
||||
assert lock_manager.lock_acquired is True
|
||||
|
||||
def test_release_lock_success(self):
|
||||
"""Test successful lock release using Lua script"""
|
||||
mock_redis = Mock()
|
||||
# FIX VERIFICATION: Lua script returns 1 for successful delete
|
||||
mock_redis.redis_client.eval.return_value = 1
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
lock_manager.pod_id = "pod_123_456"
|
||||
lock_manager.lock_acquired = True
|
||||
|
||||
lock_manager.release_lock()
|
||||
|
||||
# Verify Lua script was called for atomic compare-and-delete
|
||||
mock_redis.redis_client.eval.assert_called_once()
|
||||
call_args = mock_redis.redis_client.eval.call_args
|
||||
# Verify Lua script contains atomic compare-and-delete logic
|
||||
lua_script = call_args[0][0]
|
||||
assert "GET" in lua_script
|
||||
assert "DEL" in lua_script
|
||||
assert lock_manager.lock_acquired is False
|
||||
|
||||
def test_release_lock_wrong_owner(self):
|
||||
"""Test releasing lock when not the owner"""
|
||||
mock_redis = Mock()
|
||||
# FIX VERIFICATION: Lua script returns 0 when ownership check fails
|
||||
mock_redis.redis_client.eval.return_value = 0
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
lock_manager.pod_id = "pod_123_456"
|
||||
lock_manager.lock_acquired = True
|
||||
|
||||
lock_manager.release_lock()
|
||||
|
||||
mock_redis.redis_client.eval.assert_called_once()
|
||||
assert lock_manager.lock_acquired is False
|
||||
|
||||
def test_context_manager(self):
|
||||
"""Test MigrationLockManager as context manager"""
|
||||
mock_redis = Mock()
|
||||
mock_redis.redis_client.set.return_value = True
|
||||
mock_redis.redis_client.eval.return_value = 1
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
lock_manager.pod_id = "pod_123_456"
|
||||
|
||||
with lock_manager:
|
||||
assert lock_manager.lock_acquired is True
|
||||
|
||||
# Should call release_lock when exiting context
|
||||
mock_redis.redis_client.eval.assert_called_once()
|
||||
|
||||
def test_wait_for_lock_release_success(self):
|
||||
"""Test waiting for lock release and acquiring it"""
|
||||
mock_redis = Mock()
|
||||
# First call fails, second call succeeds
|
||||
mock_redis.redis_client.set.side_effect = [False, True]
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=1)
|
||||
|
||||
assert result is True
|
||||
assert lock_manager.lock_acquired is True
|
||||
assert mock_redis.redis_client.set.call_count == 2
|
||||
|
||||
def test_wait_for_lock_release_timeout(self):
|
||||
"""Test timeout when waiting for lock release"""
|
||||
mock_redis = Mock()
|
||||
# Always fails to acquire lock
|
||||
mock_redis.redis_client.set.return_value = False
|
||||
|
||||
lock_manager = MigrationLockManager(mock_redis)
|
||||
start_time = time.time()
|
||||
result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=0.5)
|
||||
end_time = time.time()
|
||||
|
||||
assert result is False
|
||||
assert lock_manager.lock_acquired is False
|
||||
assert end_time - start_time >= 0.5 # Should wait at least max_wait time
|
||||
assert mock_redis.redis_client.set.call_count > 1 # Multiple attempts
|
||||
|
||||
|
||||
class TestProxyExtrasDBManagerMigrationLock:
|
||||
"""Test cases for ProxyExtrasDBManager with migration locking"""
|
||||
|
||||
@patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration")
|
||||
def test_setup_database_with_redis_lock_success(
|
||||
self, mock_execute_migration, monkeypatch
|
||||
):
|
||||
"""Test successful database setup with Redis lock"""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
|
||||
mock_execute_migration.return_value = True
|
||||
|
||||
# Mock Redis cache
|
||||
mock_redis = Mock()
|
||||
mock_redis.redis_client.set.return_value = True # Lock acquired
|
||||
mock_redis.redis_client.eval.return_value = 1 # Lock released
|
||||
|
||||
result = ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, redis_cache=mock_redis
|
||||
)
|
||||
|
||||
assert result is True
|
||||
mock_execute_migration.assert_called_once()
|
||||
# Verify lock was acquired
|
||||
mock_redis.redis_client.set.assert_called_once()
|
||||
|
||||
@patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration")
|
||||
def test_setup_database_with_redis_lock_wait_and_acquire(
|
||||
self, mock_execute_migration, monkeypatch
|
||||
):
|
||||
"""Test database setup when lock is held, then acquired after waiting"""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
|
||||
mock_execute_migration.return_value = True
|
||||
|
||||
# Mock Redis cache - first call fails, second call succeeds
|
||||
mock_redis = Mock()
|
||||
mock_redis.redis_client.set.side_effect = [False, True]
|
||||
mock_redis.redis_client.eval.return_value = 1
|
||||
|
||||
result = ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, redis_cache=mock_redis
|
||||
)
|
||||
|
||||
assert result is True
|
||||
mock_execute_migration.assert_called_once()
|
||||
# Should have tried to acquire lock twice
|
||||
assert mock_redis.redis_client.set.call_count == 2
|
||||
|
||||
@patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration")
|
||||
def test_setup_database_without_redis(self, mock_execute_migration, monkeypatch):
|
||||
"""Test database setup without Redis cache (single instance mode)"""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
|
||||
mock_execute_migration.return_value = True
|
||||
|
||||
result = ProxyExtrasDBManager.setup_database(use_migrate=True, redis_cache=None)
|
||||
|
||||
assert result is True
|
||||
mock_execute_migration.assert_called_once()
|
||||
|
||||
def test_setup_database_no_database_url(self):
|
||||
"""Test database setup without DATABASE_URL"""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
result = ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, redis_cache=None
|
||||
)
|
||||
|
||||
assert result is False
|
||||
|
||||
@patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration")
|
||||
def test_setup_database_lock_timeout(self, mock_execute_migration, monkeypatch):
|
||||
"""Test database setup when lock acquisition times out"""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
|
||||
|
||||
# Mock Redis cache - always fails to acquire lock
|
||||
mock_redis = Mock()
|
||||
mock_redis.redis_client.set.return_value = False
|
||||
|
||||
# Patch wait_for_lock_release to simulate timeout
|
||||
with patch.object(MigrationLockManager, "wait_for_lock_release") as mock_wait:
|
||||
mock_wait.return_value = False # Simulate timeout
|
||||
|
||||
result = ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, redis_cache=mock_redis
|
||||
)
|
||||
|
||||
assert result is False
|
||||
mock_execute_migration.assert_not_called() # Should not run migration
|
||||
mock_wait.assert_called_once()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue