From 72c0ad419fb24fbafac9b095b2dcbd2475b43ffe Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 26 Mar 2025 17:11:25 -0700
Subject: [PATCH] build(migration.sql): add migration file for new
 dailyusertable

documents prisma db changes
---
 ci_cd/run_migration.py                        | 96 +++++++++++++++++++
 .../migration.sql                             | 33 +++++++
 2 files changed, 129 insertions(+)
 create mode 100644 ci_cd/run_migration.py
 create mode 100644 deploy/migrations/20250326171002_add_daily_user_table/migration.sql

diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py
new file mode 100644
index 00000000000..69b4308e57d
--- /dev/null
+++ b/ci_cd/run_migration.py
@@ -0,0 +1,96 @@
+import os
+import subprocess
+from pathlib import Path
+from datetime import datetime
+import testing.postgresql
+import shutil
+
+
+def create_migration(migration_name: str = None):
+    """
+    Create a new migration SQL file in deploy/migrations directory by comparing
+    current database state with schema
+
+    Args:
+        migration_name (str): Name for the migration
+    """
+    try:
+        # Get paths
+        root_dir = Path(__file__).parent.parent
+        deploy_dir = root_dir / "deploy"
+        migrations_dir = deploy_dir / "migrations"
+        schema_path = root_dir / "schema.prisma"
+
+        # Create temporary PostgreSQL database
+        with testing.postgresql.Postgresql() as postgresql:
+            db_url = postgresql.url()
+
+            # Create temporary migrations directory next to schema.prisma
+            temp_migrations_dir = schema_path.parent / "migrations"
+
+            try:
+                # Copy existing migrations to temp directory
+                if temp_migrations_dir.exists():
+                    shutil.rmtree(temp_migrations_dir)
+                shutil.copytree(migrations_dir, temp_migrations_dir)
+
+                # Apply existing migrations to temp database
+                os.environ["DATABASE_URL"] = db_url
+                subprocess.run(
+                    ["prisma", "migrate", "deploy", "--schema", str(schema_path)],
+                    check=True,
+                )
+
+                # Generate diff between current database and schema
+                result = subprocess.run(
+                    [
+                        "prisma",
+                        "migrate",
+                        "diff",
+                        "--from-url",
+                        db_url,
+                        "--to-schema-datamodel",
+                        str(schema_path),
+                        "--script",
+                    ],
+                    capture_output=True,
+                    text=True,
+                    check=True,
+                )
+
+                if result.stdout.strip():
+                    # Generate timestamp and create migration directory
+                    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+                    migration_name = migration_name or "unnamed_migration"
+                    migration_dir = migrations_dir / f"{timestamp}_{migration_name}"
+                    migration_dir.mkdir(parents=True, exist_ok=True)
+
+                    # Write the SQL to migration.sql
+                    migration_file = migration_dir / "migration.sql"
+                    migration_file.write_text(result.stdout)
+
+                    print(f"Created migration in {migration_dir}")
+                    return True
+                else:
+                    print("No schema changes detected. Migration not needed.")
+                    return False
+
+            finally:
+                # Clean up: remove temporary migrations directory
+                if temp_migrations_dir.exists():
+                    shutil.rmtree(temp_migrations_dir)
+
+    except subprocess.CalledProcessError as e:
+        print(f"Error generating migration: {e.stderr}")
+        return False
+    except Exception as e:
+        print(f"Error creating migration: {str(e)}")
+        return False
+
+
+if __name__ == "__main__":
+    # If running directly, can optionally pass migration name as argument
+    import sys
+
+    migration_name = sys.argv[1] if len(sys.argv) > 1 else None
+    create_migration(migration_name)
diff --git a/deploy/migrations/20250326171002_add_daily_user_table/migration.sql b/deploy/migrations/20250326171002_add_daily_user_table/migration.sql
new file mode 100644
index 00000000000..3379d8e9fda
--- /dev/null
+++ b/deploy/migrations/20250326171002_add_daily_user_table/migration.sql
@@ -0,0 +1,33 @@
+-- CreateTable
+CREATE TABLE "LiteLLM_DailyUserSpend" (
+    "id" TEXT NOT NULL,
+    "user_id" TEXT NOT NULL,
+    "date" TEXT NOT NULL,
+    "api_key" TEXT NOT NULL,
+    "model" TEXT NOT NULL,
+    "model_group" TEXT,
+    "custom_llm_provider" TEXT,
+    "prompt_tokens" INTEGER NOT NULL DEFAULT 0,
+    "completion_tokens" INTEGER NOT NULL DEFAULT 0,
+    "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "LiteLLM_DailyUserSpend_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model");
+
+-- CreateIndex
+CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider");
+