fix(proxy-extras): honor a raised command timeout for migrate deploy and name the right knob on db push timeouts

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -26,7 +26,10 @@ finish any sooner. Migrate deploy therefore runs under its own budget.
 All three budgets are overridable so an operator can widen them without a
 release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install,
 ``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and
-``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command.
+``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The
+per-command budget used to bound migrate deploy as well, so a deployment that
+raised it above the deploy default keeps that larger budget for deploy unless
+the deploy override says otherwise.
 """
 
 import math
@@ -102,9 +105,11 @@ def prisma_bootstrap_timeout() -> float:
 
 def prisma_migrate_deploy_timeout() -> float:
     """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending."""
-    return _timeout_from_env(
-        PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT
-    )
+    if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None:
+        return _timeout_from_env(
+            PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT
+        )
+    return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout())
 
 
 def nodeenv_cache_dir() -> Optional[Path]:

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -15,6 +15,7 @@ from litellm_proxy_extras.replica_identity import (
     apply_replica_identity_full,
 )
 from litellm_proxy_extras.prisma_toolchain import (
+    PRISMA_COMMAND_TIMEOUT_ENV_VAR,
     PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
     ensure_prisma_toolchain,
     prisma_command_timeout,
@@ -1135,9 +1136,9 @@ class ProxyExtrasDBManager:
                     return True
             except subprocess.TimeoutExpired:
                 logger.warning(
-                    "Attempt %s timed out. Raise %s if this database needs longer to apply its pending migrations.",
+                    "Attempt %s timed out. Raise %s if this database needs longer to apply its schema.",
                     attempt + 1,
-                    PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
+                    PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR,
                 )
                 time.sleep(random.randrange(5, 15))
             except subprocess.CalledProcessError as e:

tests/proxy_migration_tests/test_prisma_toolchain.py

@@ -16,11 +16,13 @@ the proxy gave up after four identical timeouts.
 
 import ast
 import json
+import logging
 import os
 import sys
 import time
 from collections.abc import Callable
 from pathlib import Path
+from typing import Optional
 
 import pytest
 
@@ -52,10 +54,10 @@ import time
 args = sys.argv[1:]
 cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"]
 log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"])
-earlier_deploys = sum(
+earlier_same_command = sum(
     1
     for line in (log_path.read_text().splitlines() if log_path.exists() else [])
-    if json.loads(line)["args"][:2] == ["migrate", "deploy"]
+    if json.loads(line)["args"][:2] == args[:2]
 )
 with log_path.open("a") as log:
     log.write(
@@ -64,12 +66,14 @@ with log_path.open("a") as log:
     )
 time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0")))
 if args[:2] == ["migrate", "deploy"]:
-    if earlier_deploys == 0:
+    if earlier_same_command == 0:
         time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0")))
     elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"):
         print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr)
         sys.exit(1)
     print("No pending migrations to apply")
+if args[:2] == ["db", "push"] and earlier_same_command == 0:
+    time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_PUSH_SLEEP", "0")))
 sys.exit(0)
 """
 
@@ -269,6 +273,47 @@ def test_migrate_deploy_stops_at_its_own_timeout(
     assert elapsed < 30
 
 
+def test_db_push_timeout_hint_names_the_per_command_budget(
+    toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    """``db push`` keeps the per-command budget, so its timeout hint has to name that variable."""
+    _, log_path = toolchain_env
+    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
+    monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1")
+    monkeypatch.setenv("FAKE_PRISMA_FIRST_PUSH_SLEEP", "3")
+
+    with caplog.at_level(logging.WARNING, logger="litellm_proxy_extras"):
+        assert ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=False) is True
+
+    assert [call["args"][:2] for call in _fake_prisma_calls(log_path)].count(["db", "push"]) == 2
+    assert [record.getMessage() for record in caplog.records if "timed out" in record.getMessage()] == [
+        f"Attempt 1 timed out. Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer to apply its schema."
+    ]
+
+
+@pytest.mark.parametrize(
+    ("command_timeout", "deploy_timeout", "expected"),
+    [
+        ("900", None, 900.0),
+        ("12", None, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT),
+        ("900", "1200", 1200.0),
+        ("900", "300", 300.0),
+    ],
+    ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"],
+)
+def test_migrate_deploy_budget_keeps_a_raised_command_budget(
+    command_timeout: str, deploy_timeout: Optional[str], expected: float, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy."""
+    monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout)
+    if deploy_timeout is None:
+        monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False)
+    else:
+        monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, deploy_timeout)
+
+    assert prisma_migrate_deploy_timeout() == expected
+
+
 @pytest.mark.parametrize(
     "raw",
     ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"],