fix: how to execute cloudzero sql

litellm/integrations/cloudzero/database.py

@@ -19,7 +19,7 @@
 """Database connection and data extraction for LiteLLM."""
 
 from datetime import datetime
-from typing import Any, Dict, Optional
+from typing import Any, Optional, List
 
 import polars as pl
 
@@ -46,19 +46,9 @@ class LiteLLMDatabase:
         """Retrieve usage data from LiteLLM daily user spend table."""
         client = self._ensure_prisma_client()
 
-        # Build WHERE clause for time filtering
-        where_conditions = []
-        if start_time_utc:
-            where_conditions.append(f"dus.updated_at >= '{start_time_utc.isoformat()}'")
-        if end_time_utc:
-            where_conditions.append(f"dus.updated_at <= '{end_time_utc.isoformat()}'")
-
-        where_clause = ""
-        if where_conditions:
-            where_clause = "WHERE " + " AND ".join(where_conditions)
-
-        # Query to get user spend data with team information
-        query = f"""
+        # Query to get user spend data with team information. Use parameter binding to
+        # avoid SQL injection from user-supplied timestamps or limits.
+        query = """
         SELECT
             dus.id,
             dus.date,
@@ -85,163 +75,27 @@ class LiteLLMDatabase:
         LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
         LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
         LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
-        {where_clause}
+        WHERE ($1::timestamptz IS NULL OR dus.updated_at >= $1::timestamptz)
+          AND ($2::timestamptz IS NULL OR dus.updated_at <= $2::timestamptz)
         ORDER BY dus.date DESC, dus.created_at DESC
         """
 
-        if limit:
-            query += f" LIMIT {limit}"
+        params: List[Any] = [
+            start_time_utc,
+            end_time_utc,
+        ]
+
+        if limit is not None:
+            try:
+                params.append(int(limit))
+            except (TypeError, ValueError):
+                raise ValueError("limit must be an integer")
+            query += " LIMIT $3"
 
         try:
-            db_response = await client.db.query_raw(query)
+            db_response = await client.db.query_raw(query, *params)
             # Convert the response to polars DataFrame with full schema inference
             # This prevents schema mismatch errors when data types vary across rows
             return pl.DataFrame(db_response, infer_schema_length=None)
         except Exception as e:
             raise Exception(f"Error retrieving usage data: {str(e)}")
-
-    async def get_table_info(self) -> Dict[str, Any]:
-        """Get information about the daily user spend table."""
-        client = self._ensure_prisma_client()
-
-        try:
-            # Get row count from user spend table
-            user_count = await self._get_table_row_count("LiteLLM_DailyUserSpend")
-
-            # Get column structure from user spend table
-            query = """
-            SELECT column_name, data_type, is_nullable
-            FROM information_schema.columns
-            WHERE table_name = 'LiteLLM_DailyUserSpend'
-            ORDER BY ordinal_position;
-            """
-            columns_response = await client.db.query_raw(query)
-
-            return {
-                "columns": columns_response,
-                "row_count": user_count,
-                "table_name": "LiteLLM_DailyUserSpend",
-            }
-        except Exception as e:
-            raise Exception(f"Error getting table info: {str(e)}")
-
-    async def _get_table_row_count(self, table_name: str) -> int:
-        """Get row count from specified table."""
-        client = self._ensure_prisma_client()
-
-        try:
-            query = f'SELECT COUNT(*) as count FROM "{table_name}"'
-            response = await client.db.query_raw(query)
-
-            if response and len(response) > 0:
-                return response[0].get("count", 0)
-            return 0
-        except Exception:
-            return 0
-
-    async def discover_all_tables(self) -> Dict[str, Any]:
-        """Discover all tables in the LiteLLM database and their schemas."""
-        client = self._ensure_prisma_client()
-
-        try:
-            # Get all LiteLLM tables
-            litellm_tables_query = """
-            SELECT table_name 
-            FROM information_schema.tables 
-            WHERE table_schema = 'public' 
-            AND table_name LIKE 'LiteLLM_%'
-            ORDER BY table_name;
-            """
-            tables_response = await client.db.query_raw(litellm_tables_query)
-            table_names = [row["table_name"] for row in tables_response]
-
-            # Get detailed schema for each table
-            tables_info = {}
-            for table_name in table_names:
-                # Get column information
-                columns_query = """
-                SELECT 
-                    column_name,
-                    data_type,
-                    is_nullable,
-                    column_default,
-                    character_maximum_length,
-                    numeric_precision,
-                    numeric_scale,
-                    ordinal_position
-                FROM information_schema.columns 
-                WHERE table_name = $1
-                AND table_schema = 'public'
-                ORDER BY ordinal_position;
-                """
-                columns_response = await client.db.query_raw(columns_query, table_name)
-
-                # Get primary key information
-                pk_query = """
-                SELECT a.attname
-                FROM pg_index i
-                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
-                WHERE i.indrelid = $1::regclass AND i.indisprimary;
-                """
-                pk_response = await client.db.query_raw(pk_query, f'"{table_name}"')
-                primary_keys = (
-                    [row["attname"] for row in pk_response] if pk_response else []
-                )
-
-                # Get foreign key information
-                fk_query = """
-                SELECT
-                    tc.constraint_name,
-                    kcu.column_name,
-                    ccu.table_name AS foreign_table_name,
-                    ccu.column_name AS foreign_column_name
-                FROM information_schema.table_constraints AS tc
-                JOIN information_schema.key_column_usage AS kcu
-                    ON tc.constraint_name = kcu.constraint_name
-                JOIN information_schema.constraint_column_usage AS ccu
-                    ON ccu.constraint_name = tc.constraint_name
-                WHERE tc.constraint_type = 'FOREIGN KEY'
-                AND tc.table_name = $1;
-                """
-                fk_response = await client.db.query_raw(fk_query, table_name)
-                foreign_keys = fk_response if fk_response else []
-
-                # Get indexes
-                indexes_query = """
-                SELECT
-                    i.relname AS index_name,
-                    array_agg(a.attname ORDER BY a.attnum) AS column_names,
-                    ix.indisunique AS is_unique
-                FROM pg_class t
-                JOIN pg_index ix ON t.oid = ix.indrelid
-                JOIN pg_class i ON i.oid = ix.indexrelid
-                JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
-                WHERE t.relname = $1
-                AND t.relkind = 'r'
-                GROUP BY i.relname, ix.indisunique
-                ORDER BY i.relname;
-                """
-                indexes_response = await client.db.query_raw(indexes_query, table_name)
-                indexes = indexes_response if indexes_response else []
-
-                # Get row count
-                try:
-                    row_count = await self._get_table_row_count(table_name)
-                except Exception:
-                    row_count = 0
-
-                tables_info[table_name] = {
-                    "columns": columns_response,
-                    "primary_keys": primary_keys,
-                    "foreign_keys": foreign_keys,
-                    "indexes": indexes,
-                    "row_count": row_count,
-                }
-
-            return {
-                "tables": tables_info,
-                "table_count": len(table_names),
-                "table_names": table_names,
-            }
-        except Exception as e:
-            raise Exception(f"Error discovering tables: {str(e)}")

tests/test_litellm/integrations/cloudzero/test_database.py

@@ -0,0 +1,57 @@
+"""Tests for LiteLLM CloudZero database helper."""
+
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from litellm.integrations.cloudzero.database import LiteLLMDatabase
+
+
+def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return):
+    """Return a database instance with prisma client mocked out."""
+    query_mock = AsyncMock(return_value=query_return)
+    mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock))
+    db = LiteLLMDatabase()
+    monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
+    return db, query_mock
+
+
+@pytest.mark.asyncio
+async def test_get_usage_data_parameterized(monkeypatch: pytest.MonkeyPatch):
+    """Start/end filters and limit should be parameterized via placeholders."""
+    start = datetime(2024, 5, 1, tzinfo=timezone.utc)
+    end = datetime(2024, 5, 2, tzinfo=timezone.utc)
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    await db.get_usage_data(limit=10, start_time_utc=start, end_time_utc=end)
+
+    query_text, *params = query_mock.await_args.args
+    assert "dus.updated_at >= $1::timestamptz" in query_text
+    assert "dus.updated_at <= $2::timestamptz" in query_text
+    assert "LIMIT $3" in query_text
+    assert params == [start, end, 10]
+
+
+@pytest.mark.asyncio
+async def test_get_usage_data_handles_missing_filters(monkeypatch: pytest.MonkeyPatch):
+    """When no filters provided the params should be None placeholders."""
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    await db.get_usage_data()
+
+    query_text, *params = query_mock.await_args.args
+    assert "LIMIT $3" not in query_text
+    assert params == [None, None]
+
+
+@pytest.mark.asyncio
+async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPatch):
+    """limit must coerce to int or raise ValueError before hitting the DB."""
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    with pytest.raises(ValueError):
+        await db.get_usage_data(limit="invalid")
+
+    assert query_mock.await_count == 0