Add "keys import" command to CLI (#12620)

* Add "keys import" command to CLI

E.g.:

```
litellm-proxy keys import \
  --source-base-url=https://old-litellm.company.com \
  --source-api-key=$LITELLM_KEY \
  --dry-run
```

* Add --created-since option

* Add tests

* Fix lint errors

* Fix lint issues

* Fix lint errors

* Fix response.raise_for_status not being a thing

* Fix a mypy error

litellm/proxy/client/cli/commands/keys.py

@@ -1,5 +1,6 @@
 import json
-from typing import Literal, Optional
+from datetime import datetime
+from typing import Literal, Optional, List, Dict, Any
 
 import click
 import rich
@@ -156,3 +157,201 @@ def delete(ctx: click.Context, keys: Optional[str], key_aliases: Optional[str]):
         except json.JSONDecodeError:
             click.echo(e.response.text, err=True)
         raise click.Abort()
+
+
+def _parse_created_since_filter(created_since: Optional[str]) -> Optional[datetime]:
+    """Parse and validate the created_since date filter."""
+    if not created_since:
+        return None
+
+    try:
+        # Support formats: YYYY-MM-DD_HH:MM or YYYY-MM-DD
+        if "_" in created_since:
+            return datetime.strptime(created_since, "%Y-%m-%d_%H:%M")
+        else:
+            return datetime.strptime(created_since, "%Y-%m-%d")
+    except ValueError:
+        click.echo(f"Error: Invalid date format '{created_since}'. Use YYYY-MM-DD_HH:MM or YYYY-MM-DD", err=True)
+        raise click.Abort()
+
+
+def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_base_url: str) -> List[Dict[str, Any]]:
+    """Fetch all keys from source instance using pagination."""
+    click.echo(f"Fetching keys from source server: {source_base_url}")
+    source_keys = []
+    page = 1
+    page_size = 100  # Use a larger page size to minimize API calls
+
+    while True:
+        source_response = source_client.list(return_full_object=True, page=page, size=page_size)
+        # source_client.list() returns Dict[str, Any] when return_request is False (default)
+        assert isinstance(source_response, dict), "Expected dict response from list API"
+        page_keys = source_response.get("keys", [])
+
+        if not page_keys:
+            break
+
+        source_keys.extend(page_keys)
+        click.echo(f"Fetched page {page}: {len(page_keys)} keys")
+
+        # Check if we got fewer keys than the page size, indicating last page
+        if len(page_keys) < page_size:
+            break
+
+        page += 1
+
+    return source_keys
+
+
+def _filter_keys_by_created_since(
+    source_keys: List[Dict[str, Any]], created_since_dt: Optional[datetime], created_since: str
+) -> List[Dict[str, Any]]:
+    """Filter keys by created_since date if specified."""
+    if not created_since_dt:
+        return source_keys
+
+    filtered_keys = []
+    for key in source_keys:
+        key_created_at = key.get("created_at")
+        if key_created_at:
+            # Parse the key's created_at timestamp
+            if isinstance(key_created_at, str):
+                if "T" in key_created_at:
+                    key_dt = datetime.fromisoformat(key_created_at.replace("Z", "+00:00"))
+                else:
+                    key_dt = datetime.fromisoformat(key_created_at)
+
+                # Convert to naive datetime for comparison (assuming UTC)
+                if key_dt.tzinfo:
+                    key_dt = key_dt.replace(tzinfo=None)
+
+                if key_dt >= created_since_dt:
+                    filtered_keys.append(key)
+
+    click.echo(f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}")
+    return filtered_keys
+
+
+def _display_dry_run_table(source_keys: List[Dict[str, Any]]) -> None:
+    """Display a table of keys that would be imported in dry-run mode."""
+    click.echo("\n--- DRY RUN MODE ---")
+    table = Table(title="Keys that would be imported")
+    table.add_column("Key Alias", style="green")
+    table.add_column("User ID", style="magenta")
+    table.add_column("Created", style="cyan")
+
+    for key in source_keys:
+        created_at = key.get("created_at", "")
+        # Format the timestamp if it exists
+        if created_at:
+            # Try to parse and format the timestamp for better readability
+            if isinstance(created_at, str):
+                # Handle common timestamp formats
+                if "T" in created_at:
+                    dt = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
+                    created_at = dt.strftime("%Y-%m-%d %H:%M")
+
+        table.add_row(str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at))
+    rich.print(table)
+
+
+def _prepare_key_import_data(key: Dict[str, Any]) -> Dict[str, Any]:
+    """Prepare key data for import by extracting relevant fields."""
+    import_data = {}
+
+    # Copy relevant fields if they exist
+    for field in ["models", "aliases", "spend", "key_alias", "team_id", "user_id", "budget_id", "config"]:
+        if key.get(field):
+            import_data[field] = key[field]
+
+    return import_data
+
+
+def _import_keys_to_destination(
+    source_keys: List[Dict[str, Any]], dest_client: KeysManagementClient
+) -> tuple[int, int]:
+    """Import each key to the destination instance and return counts."""
+    imported_count = 0
+    failed_count = 0
+
+    for key in source_keys:
+        try:
+            # Prepare key data for import
+            import_data = _prepare_key_import_data(key)
+
+            # Generate the key in destination instance
+            response = dest_client.generate(**import_data)
+            click.echo(f"Generated key: {response}")
+            # The generate method returns JSON data directly, not a Response object
+            imported_count += 1
+
+            key_alias = key.get("key_alias", "N/A")
+            click.echo(f"✓ Imported key: {key_alias}")
+
+        except Exception as e:
+            failed_count += 1
+            key_alias = key.get("key_alias", "N/A")
+            click.echo(f"✗ Failed to import key {key_alias}: {str(e)}", err=True)
+
+    return imported_count, failed_count
+
+
+@keys.command(name="import")
+@click.option(
+    "--source-base-url", required=True, help="Base URL of the source LiteLLM proxy server to import keys from"
+)
+@click.option("--source-api-key", help="API key for authentication to the source server")
+@click.option("--dry-run", is_flag=True, help="Show what would be imported without actually importing")
+@click.option(
+    "--created-since", help="Only import keys created after this date/time (format: YYYY-MM-DD_HH:MM or YYYY-MM-DD)"
+)
+@click.pass_context
+def import_keys(
+    ctx: click.Context, source_base_url: str, source_api_key: Optional[str], dry_run: bool, created_since: Optional[str]
+):
+    """Import API keys from another LiteLLM instance"""
+    # Parse created_since filter if provided
+    created_since_dt = _parse_created_since_filter(created_since)
+
+    # Create clients for both source and destination
+    source_client = KeysManagementClient(source_base_url, source_api_key)
+    dest_client = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
+
+    try:
+        # Get all keys from source instance with pagination
+        source_keys = _fetch_all_keys_with_pagination(source_client, source_base_url)
+
+        # Filter keys by created_since if specified
+        if created_since:
+            source_keys = _filter_keys_by_created_since(source_keys, created_since_dt, created_since)
+
+        if not source_keys:
+            click.echo("No keys found in source instance.")
+            return
+
+        click.echo(f"Found {len(source_keys)} keys in source instance.")
+
+        if dry_run:
+            _display_dry_run_table(source_keys)
+            return
+
+        # Import each key
+        imported_count, failed_count = _import_keys_to_destination(source_keys, dest_client)
+
+        # Summary
+        click.echo("\nImport completed:")
+        click.echo(f"  Successfully imported: {imported_count}")
+        click.echo(f"  Failed to import: {failed_count}")
+        click.echo(f"  Total keys processed: {len(source_keys)}")
+
+    except requests.exceptions.HTTPError as e:
+        click.echo(f"Error: HTTP {e.response.status_code}", err=True)
+        try:
+            error_body = e.response.json()
+            rich.print_json(data=error_body)
+        except json.JSONDecodeError:
+            click.echo(e.response.text, err=True)
+        raise click.Abort()
+    except Exception as e:
+        click.echo(f"Error: {str(e)}", err=True)
+        raise click.Abort()

tests/test_litellm/proxy/client/cli/test_keys_commands.py

@@ -3,6 +3,8 @@ import os
 import sys
 from unittest.mock import patch
 
+import requests
+
 sys.path.insert(
     0, os.path.abspath("../../..")
 )  # Adds the parent directory to the system path
@@ -34,9 +36,7 @@ def mock_env():
 
 @pytest.fixture
 def mock_keys_client():
-    with patch(
-        "litellm.proxy.client.cli.commands.keys.KeysManagementClient"
-    ) as MockClient:
+    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
         yield MockClient
 
 
@@ -88,9 +88,7 @@ def test_async_keys_generate_success(mock_keys_client, cli_runner):
         "key": "new-key",
         "spend": 100.0,
     }
-    result = cli_runner.invoke(
-        cli, ["keys", "generate", "--models", "gpt-4", "--spend", "100"]
-    )
+    result = cli_runner.invoke(cli, ["keys", "generate", "--models", "gpt-4", "--spend", "100"])
     assert result.exit_code == 0
     assert "new-key" in result.output
     mock_keys_client.return_value.generate.assert_called_once()
@@ -158,3 +156,326 @@ def test_async_keys_delete_http_error_handling(mock_keys_client, cli_runner):
     assert result.exit_code != 0
     # HTTPError should be caught and converted to click.Abort
     assert isinstance(result.exception, SystemExit)  # click.Abort raises SystemExit
+
+
+# Tests for keys import command
+def test_keys_import_dry_run_success(mock_keys_client, cli_runner):
+    """Test successful dry-run import showing table of keys that would be imported"""
+    # Mock source client response (paginated)
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.side_effect = [
+        {
+            "keys": [
+                {
+                    "key_alias": "test-key-1",
+                    "user_id": "user1@example.com",
+                    "created_at": "2024-01-15T10:30:00Z",
+                    "models": ["gpt-4"],
+                    "spend": 10.0,
+                },
+                {
+                    "key_alias": "test-key-2", 
+                    "user_id": "user2@example.com",
+                    "created_at": "2024-01-16T11:45:00Z",
+                    "models": [],
+                    "spend": 5.0,
+                }
+            ]
+        },
+        {"keys": []}  # Empty second page
+    ]
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com",
+        "--source-api-key", "sk-source-123",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code == 0
+    assert "Found 2 keys in source instance" in result.output
+    assert "DRY RUN MODE" in result.output
+    assert "test-key-1" in result.output
+    assert "user1@example.com" in result.output
+    assert "test-key-2" in result.output
+    assert "user2@example.com" in result.output
+    
+    # Verify source client was called (pagination stops early when fewer keys than page_size)
+    assert mock_source_instance.list.call_count >= 1
+    mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100)
+
+
+def test_keys_import_actual_import_success(mock_keys_client, cli_runner):
+    """Test successful actual import of keys"""
+    # Create separate mock instances for source and destination
+    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
+        mock_source_instance = MockClient.return_value
+        mock_dest_instance = MockClient.return_value
+        
+        # Configure source client
+        mock_source_instance.list.side_effect = [
+            {
+                "keys": [
+                    {
+                        "key_alias": "import-key-1",
+                        "user_id": "user1@example.com",
+                        "models": ["gpt-4"],
+                        "spend": 100.0,
+                        "team_id": "team-1"
+                    }
+                ]
+            },
+            {"keys": []}  # Empty second page
+        ]
+        
+        # Configure destination client
+        mock_dest_instance.generate.return_value = {
+            "key": "sk-new-generated-key",
+            "status": "success"
+        }
+        
+        result = cli_runner.invoke(cli, [
+            "keys", "import",
+            "--source-base-url", "https://source.example.com",
+            "--source-api-key", "sk-source-123"
+        ])
+        
+        assert result.exit_code == 0
+        assert "Found 1 keys in source instance" in result.output
+        assert "✓ Imported key: import-key-1" in result.output
+        assert "Successfully imported: 1" in result.output
+        assert "Failed to import: 0" in result.output
+        
+        # Verify generate was called with correct parameters
+        mock_dest_instance.generate.assert_called_once_with(
+            models=["gpt-4"],
+            spend=100.0,
+            key_alias="import-key-1",
+            team_id="team-1",
+            user_id="user1@example.com"
+        )
+
+
+def test_keys_import_pagination_handling(mock_keys_client, cli_runner):
+    """Test that import correctly handles pagination to get all keys"""
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.side_effect = [
+        {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100)]},  # Page 1: 100 keys
+        {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100, 150)]},  # Page 2: 50 keys
+        {"keys": []}  # Page 3: Empty
+    ]
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import", 
+        "--source-base-url", "https://source.example.com",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code == 0
+    assert "Fetched page 1: 100 keys" in result.output
+    assert "Fetched page 2: 50 keys" in result.output
+    assert "Found 150 keys in source instance" in result.output
+    
+    # Verify pagination calls (stops early when fewer keys than page_size)
+    assert mock_source_instance.list.call_count >= 2
+    mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100)
+    mock_source_instance.list.assert_any_call(return_full_object=True, page=2, size=100)
+
+
+def test_keys_import_created_since_filter(mock_keys_client, cli_runner):
+    """Test that --created-since filter works correctly"""
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.side_effect = [
+        {
+            "keys": [
+                {
+                    "key_alias": "old-key",
+                    "user_id": "user1@example.com", 
+                    "created_at": "2024-01-01T10:00:00Z"  # Before filter
+                },
+                {
+                    "key_alias": "new-key",
+                    "user_id": "user2@example.com",
+                    "created_at": "2024-07-08T10:00:00Z"  # After filter
+                }
+            ]
+        },
+        {"keys": []}
+    ]
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com", 
+        "--created-since", "2024-07-07_18:19",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code == 0
+    assert "Filtered 2 keys to 1 keys created since 2024-07-07_18:19" in result.output
+    assert "Found 1 keys in source instance" in result.output
+    assert "new-key" in result.output
+    assert "old-key" not in result.output
+
+
+def test_keys_import_created_since_date_only_format(mock_keys_client, cli_runner):
+    """Test --created-since with date-only format (YYYY-MM-DD)"""
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.side_effect = [
+        {
+            "keys": [
+                {
+                    "key_alias": "test-key",
+                    "user_id": "user@example.com",
+                    "created_at": "2024-07-08T10:00:00Z"
+                }
+            ]
+        },
+        {"keys": []}
+    ]
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com",
+        "--created-since", "2024-07-07",  # Date only format
+        "--dry-run"
+    ])
+    
+    assert result.exit_code == 0
+    assert "Filtered 1 keys to 1 keys created since 2024-07-07" in result.output
+
+
+def test_keys_import_no_keys_found(mock_keys_client, cli_runner):
+    """Test handling when no keys are found in source instance"""
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.return_value = {"keys": []}
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code == 0
+    assert "No keys found in source instance" in result.output
+
+
+def test_keys_import_invalid_date_format(cli_runner):
+    """Test error handling for invalid --created-since date format"""
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com",
+        "--created-since", "invalid-date",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code != 0
+    assert "Invalid date format" in result.output
+    assert "Use YYYY-MM-DD_HH:MM or YYYY-MM-DD" in result.output
+
+
+def test_keys_import_source_api_error(mock_keys_client, cli_runner):
+    """Test error handling when source API returns an error"""
+    mock_source_instance = mock_keys_client.return_value
+    mock_source_instance.list.side_effect = Exception("Source API Error")
+    
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--source-base-url", "https://source.example.com",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code != 0
+    assert "Source API Error" in result.output
+
+
+def test_keys_import_partial_failure(mock_keys_client, cli_runner):
+    """Test handling when some keys fail to import"""
+    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
+        mock_source_instance = MockClient.return_value
+        mock_dest_instance = MockClient.return_value
+        
+        # Source returns 2 keys
+        mock_source_instance.list.side_effect = [
+            {
+                "keys": [
+                    {"key_alias": "success-key", "user_id": "user1@example.com"},
+                    {"key_alias": "fail-key", "user_id": "user2@example.com"}
+                ]
+            },
+            {"keys": []}
+        ]
+        
+        # Destination: first succeeds, second fails
+        mock_dest_instance.generate.side_effect = [
+            {"key": "sk-new-key", "status": "success"},
+            Exception("Import failed for this key")
+        ]
+        
+        result = cli_runner.invoke(cli, [
+            "keys", "import",
+            "--source-base-url", "https://source.example.com"
+        ])
+        
+        assert result.exit_code == 0  # Command completes even with partial failures
+        assert "✓ Imported key: success-key" in result.output
+        assert "✗ Failed to import key fail-key" in result.output
+        assert "Successfully imported: 1" in result.output
+        assert "Failed to import: 1" in result.output
+        assert "Total keys processed: 2" in result.output
+
+
+def test_keys_import_missing_required_source_url(cli_runner):
+    """Test error when required --source-base-url is missing"""
+    result = cli_runner.invoke(cli, [
+        "keys", "import",
+        "--dry-run"
+    ])
+    
+    assert result.exit_code != 0
+    assert "Missing option" in result.output or "required" in result.output.lower()
+
+
+def test_keys_import_with_all_key_properties(mock_keys_client, cli_runner):
+    """Test import preserves all key properties (models, aliases, config, etc.)"""
+    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
+        mock_source_instance = MockClient.return_value
+        mock_dest_instance = MockClient.return_value
+        
+        mock_source_instance.list.side_effect = [
+            {
+                "keys": [
+                    {
+                        "key_alias": "full-key",
+                        "user_id": "user@example.com",
+                        "team_id": "team-123",
+                        "budget_id": "budget-456", 
+                        "models": ["gpt-4", "gpt-3.5-turbo"],
+                        "aliases": {"custom-model": "gpt-4"},
+                        "spend": 50.0,
+                        "config": {"max_tokens": 1000}
+                    }
+                ]
+            },
+            {"keys": []}
+        ]
+        
+        mock_dest_instance.generate.return_value = {"key": "sk-imported", "status": "success"}
+        
+        result = cli_runner.invoke(cli, [
+            "keys", "import",
+            "--source-base-url", "https://source.example.com"
+        ])
+        
+        assert result.exit_code == 0
+        
+        # Verify all properties were passed to generate
+        mock_dest_instance.generate.assert_called_once_with(
+            models=["gpt-4", "gpt-3.5-turbo"],
+            aliases={"custom-model": "gpt-4"},
+            spend=50.0,
+            key_alias="full-key",
+            team_id="team-123",
+            user_id="user@example.com",
+            budget_id="budget-456",
+            config={"max_tokens": 1000}
+        )