diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 750c1a934aa..ccb9f2def30 100644 --- a/db_scripts/partition_spend_logs.sql +++ b/db_scripts/partition_spend_logs.sql @@ -18,6 +18,8 @@ -- * Test on a staging copy first and take a backup. -- * Postgres cannot convert a populated table to partitioned in place, so this -- renames the old table aside and creates a fresh partitioned table. +-- * Dependent views are repointed to the new table within the transaction. +-- Drop directly dependent materialized views first and recreate afterwards. -- * The partition key ("startTime") must be part of the primary key, so the -- PK becomes composite ("request_id", "startTime"). LiteLLM's write path uses -- INSERT ... ON CONFLICT DO NOTHING, which is compatible with this. @@ -90,12 +92,38 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault" PARTITION OF "LiteLLM_SpendLogs" DEFAULT; +-- Preserve view OIDs so views of these views keep working without recreation. +-- Materialized views require an explicit rebuild by the operator instead. +DO $$ +DECLARE + v record; +BEGIN + FOR v IN + SELECT DISTINCT view_class.oid::regclass AS view_name, view_class.relkind, view_class.reloptions + FROM pg_depend dep + JOIN pg_rewrite rule ON rule.oid = dep.objid + JOIN pg_class view_class ON view_class.oid = rule.ev_class + WHERE dep.refobjid = '"LiteLLM_SpendLogs_legacy"'::regclass + AND view_class.relkind IN ('v', 'm') + LOOP + IF v.relkind = 'm' THEN + RAISE EXCEPTION 'Materialized view % depends on LiteLLM_SpendLogs_legacy. Drop it before running this script and recreate it afterwards.', v.view_name; + END IF; + EXECUTE format('CREATE OR REPLACE VIEW %s %s AS %s', v.view_name, + CASE WHEN v.reloptions IS NULL THEN '' + ELSE format('WITH (%s)', array_to_string(v.reloptions, ', ')) + END, + replace(pg_get_viewdef(v.view_name), '"LiteLLM_SpendLogs_legacy"', '"LiteLLM_SpendLogs"')); + END LOOP; +END $$; + COMMIT; -- Backfill (optional). Rows route to the correct partition automatically. -- For large legacy tables, copy in time-bounded batches during a low-traffic -- window instead of one statement, or simply keep "LiteLLM_SpendLogs_legacy" -- read-only until its data ages past your retention, then DROP it. +-- Views now read the new table, so legacy rows appear only after backfilling. -- -- Backfilled rows land in the DEFAULT partition until explicit partitions -- cover their dates. Postgres refuses to create a partition whose range diff --git a/db_scripts/unpartition_spend_logs.sql b/db_scripts/unpartition_spend_logs.sql index 2555eca212b..ae19b4f058c 100644 --- a/db_scripts/unpartition_spend_logs.sql +++ b/db_scripts/unpartition_spend_logs.sql @@ -11,6 +11,8 @@ -- * Test on a staging copy first and take a backup. -- * Postgres cannot convert a partitioned table back in place, so this -- renames the partitioned table aside and creates a fresh plain table. +-- * Dependent views are repointed before dropping the partitioned table. +-- Drop directly dependent materialized views first and recreate afterwards. -- * The composite PK could in principle hold the same "request_id" in more -- than one partition, so rows are copied with ON CONFLICT DO NOTHING to -- restore the single-column PK without failing on such duplicates. @@ -69,6 +71,31 @@ INSERT INTO "LiteLLM_SpendLogs" SELECT * FROM "LiteLLM_SpendLogs_partitioned" ON CONFLICT ("request_id") DO NOTHING; +-- Preserve view OIDs so views of these views keep working without recreation. +-- Materialized views require an explicit rebuild by the operator instead. +DO $$ +DECLARE + v record; +BEGIN + FOR v IN + SELECT DISTINCT view_class.oid::regclass AS view_name, view_class.relkind, view_class.reloptions + FROM pg_depend dep + JOIN pg_rewrite rule ON rule.oid = dep.objid + JOIN pg_class view_class ON view_class.oid = rule.ev_class + WHERE dep.refobjid = '"LiteLLM_SpendLogs_partitioned"'::regclass + AND view_class.relkind IN ('v', 'm') + LOOP + IF v.relkind = 'm' THEN + RAISE EXCEPTION 'Materialized view % depends on LiteLLM_SpendLogs_partitioned. Drop it before running this script and recreate it afterwards.', v.view_name; + END IF; + EXECUTE format('CREATE OR REPLACE VIEW %s %s AS %s', v.view_name, + CASE WHEN v.reloptions IS NULL THEN '' + ELSE format('WITH (%s)', array_to_string(v.reloptions, ', ')) + END, + replace(pg_get_viewdef(v.view_name), '"LiteLLM_SpendLogs_partitioned"', '"LiteLLM_SpendLogs"')); + END LOOP; +END $$; + DROP TABLE "LiteLLM_SpendLogs_partitioned"; COMMIT; diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 5741c85470b..a47ec86990a 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -539,6 +539,59 @@ def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_atta _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx") +@requires_db +@pytest.mark.parametrize("partition_first", (False, True), ids=("partition", "unpartition")) +def test_spend_runbooks_repoint_views_and_allow_old_table_drop(scratch_database: str, partition_first: bool) -> None: + with psycopg.connect(scratch_database, autocommit=True) as conn: + if partition_first: + conn.execute(PARTITION_SCRIPT.read_bytes()) + conn.execute("CREATE SCHEMA reports") + conn.execute( + "CREATE VIEW reports.direct WITH (security_barrier=true) AS " + 'SELECT request_id FROM public."LiteLLM_SpendLogs" WITH LOCAL CHECK OPTION' + ) + conn.execute("CREATE VIEW reports.nested AS SELECT * FROM reports.direct") + _insert_spend_log(conn, "before", "2026-08-01") + script: Final = ( + PARTITION_SCRIPT.with_name("unpartition_spend_logs.sql") if partition_first else PARTITION_SCRIPT + ) + conn.execute(script.read_bytes()) + assert conn.execute("SELECT reloptions FROM pg_class WHERE oid = 'reports.direct'::regclass").fetchone() == ( + ["security_barrier=true", "check_option=local"], + ) + _insert_spend_log(conn, "after", "2026-08-01") + expected: Final = [("after",), ("before",)] if partition_first else [("after",)] + assert conn.execute("SELECT * FROM reports.direct ORDER BY request_id").fetchall() == expected + assert conn.execute("SELECT * FROM reports.nested ORDER BY request_id").fetchall() == expected + if not partition_first: + conn.execute('DROP TABLE "LiteLLM_SpendLogs_legacy"') + assert conn.execute( + "SELECT relkind FROM pg_class WHERE oid = '\"LiteLLM_SpendLogs\"'::regclass" + ).fetchone() == ("r" if partition_first else "p",) + + +@requires_db +@pytest.mark.parametrize("partition_first", (False, True), ids=("partition", "unpartition")) +def test_spend_runbooks_reject_materialized_views_without_changing_the_table( + scratch_database: str, partition_first: bool +) -> None: + with psycopg.connect(scratch_database, autocommit=True) as conn: + if partition_first: + conn.execute(PARTITION_SCRIPT.read_bytes()) + conn.execute('CREATE MATERIALIZED VIEW spend_snapshot AS SELECT request_id FROM "LiteLLM_SpendLogs"') + original: Final = conn.execute("SELECT '\"LiteLLM_SpendLogs\"'::regclass::oid").fetchone() + script: Final = ( + PARTITION_SCRIPT.with_name("unpartition_spend_logs.sql") if partition_first else PARTITION_SCRIPT + ) + with pytest.raises(psycopg.errors.RaiseException, match="Materialized view spend_snapshot"): + conn.execute(script.read_bytes()) + conn.execute("ROLLBACK") + assert conn.execute("SELECT '\"LiteLLM_SpendLogs\"'::regclass::oid").fetchone() == original + _insert_spend_log(conn, "still-writable", "2026-08-01") + conn.execute("REFRESH MATERIALIZED VIEW spend_snapshot") + assert conn.execute("SELECT * FROM spend_snapshot").fetchall() == [("still-writable",)] + + @requires_db def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_index_aside( scratch_database: str,