fix(migration_data_checks): size btree entries by the bytes Postgres stores (LIT-10244) (#48911)

* fix(migration_data_checks): size btree entries by the bytes Postgres stores (LIT-10244)

* test(migration_data_checks): flag exactly the compressed keys Postgres rejects at the entry limit (LIT-10244)

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py

@@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
 
 
 def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
+    """The bytes a value takes in the index entry: its length, or its stored size when
+    Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
     from psycopg import sql
 
-    if column.type_name in _OCTET_LENGTH_TYPES:
-        return sql.SQL("octet_length({})").format(identifier)
-    return sql.SQL("octet_length({}::text)").format(identifier)
+    length: Final = (
+        sql.SQL("octet_length({})").format(identifier)
+        if column.type_name in _OCTET_LENGTH_TYPES
+        else sql.SQL("octet_length({}::text)").format(identifier)
+    )
+    return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
 
 
 def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":

tests/integration/database/migrations/test_request_log_indexes.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 import queue
 import secrets
@@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key
         _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
         _insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
         _insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
+        _insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000)
     ledger_before: Final = _ledger(scratch_database)
 
     with pytest.raises(RuntimeError) as refused:
@@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
     assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})"
 
 
+@requires_db
+def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is(
+    scratch_database: str,
+) -> None:
+    """Postgres compresses a long value before it checks the entry size, so 3000 repeated
+    bytes index fine while 3000 random hex bytes do not. The check must draw the same line."""
+    create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")')
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)')
+        conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500)))
+        findings: Final = migration_data_checks.find_blockers(
+            conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
+        )
+        with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
+            conn.execute(create_index)
+        conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,))
+        conn.execute(create_index)
+        stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone()
+
+    assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored
+    (finding,) = findings
+    assert isinstance(finding, migration_data_checks.BlockedChange)
+    assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other"
+    assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)"
+
+
+def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool:
+    conn.execute('DROP TABLE IF EXISTS "one_key"')
+    conn.execute('CREATE TABLE "one_key" ("k" TEXT)')
+    conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,))
+    try:
+        conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")')
+    except psycopg.errors.ProgramLimitExceeded:
+        return True
+    return False
+
+
+@requires_db
+def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None:
+    """Keys that compress to a stored size just under and just over the entry limit go through the
+    check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it."""
+    noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40))
+    keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500))
+    create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")')
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)')
+        conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),))
+        compressed: Final = conn.execute(
+            'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704'
+        ).fetchone()
+        findings: Final = migration_data_checks.find_blockers(
+            conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
+        )
+        rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
+
+    assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed"
+    assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}"
+    (finding,) = findings
+    assert isinstance(finding, migration_data_checks.BlockedChange)
+    assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other"
+
+
 @requires_db
 def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
     with psycopg.connect(scratch_database, autocommit=True) as conn:

tests/unit/litellm_proxy_extras/test_migration_data_checks.py

@@ -7,7 +7,9 @@ from litellm_proxy_extras.migration_data_checks import (
     MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR,
     BlockedChange,
     Check,
+    ColumnType,
     CouldNotVerify,
+    _count_query,
     btree_max_entry_bytes,
     checks_for_index,
     checks_for_migration,
@@ -113,6 +115,19 @@ def test_the_btree_entry_limit_on_the_default_page_is_the_number_in_postgres_err
     assert btree_max_entry_bytes(8192) == 2704
 
 
+def test_the_entry_size_query_sizes_a_text_column_by_its_stored_bytes_when_postgres_compressed_it() -> None:
+    """Postgres compresses a long value before it checks the entry size, so a repetitive
+    3000-byte key indexes fine. The query has to size it as stored, never by its length alone."""
+    (check,) = checks_for_index(
+        "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
+    )
+    types: Final = (ColumnType("text", -1, "i", -1), ColumnType("timestamp", 8, "d", 3))
+    query: Final = _count_query(check, "public", types, 2704).as_string(None)
+    stored: Final = 'least(octet_length("api_key"), pg_column_size("api_key"))'
+    assert stored in query
+    assert 'octet_length("api_key")' not in query.replace(stored, ""), query
+
+
 def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_was_applied() -> None:
     (check,) = checks_for_index(
         "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")