feat: add login command

pandasai/cli/main.py

@@ -12,6 +12,12 @@
 from pandasai.helpers.path import find_project_root, get_validated_dataset_path
 
 
+def validate_api_key(api_key: str) -> bool:
+    """Validate PandaBI API key format."""
+    pattern = r"^PAI-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
+    return bool(re.match(pattern, api_key))
+
+
 @click.group()
 def cli():
     """🐼 PandaAI CLI - Manage your datasets with ease"""
@@ -93,6 +99,37 @@ def create():
     click.echo(f"\n✨ Dataset created successfully at: {dataset_directory}")
 
 
+@cli.command()
+@click.argument("api_key")
+def login(api_key: str):
+    """🔑 Authenticate with your PandaBI API key"""
+    if not validate_api_key(api_key):
+        click.echo(
+            "❌ Invalid API key format. Expected format: PAI-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
+        )
+        return
+
+    env_path = os.path.join(find_project_root(), ".env")
+    env_content = ""
+    new_line = f"PANDABI_API_KEY={api_key}\n"
+
+    # Read existing .env if it exists
+    if os.path.exists(env_path):
+        with open(env_path, "r") as f:
+            lines = f.readlines()
+            # Filter out existing PANDABI_API_KEY line if present
+            lines = [line for line in lines if not line.startswith("PANDABI_API_KEY=")]
+            env_content = "".join(lines)
+            if env_content and not env_content.endswith("\n"):
+                env_content += "\n"
+
+    # Write updated content
+    with open(env_path, "w") as f:
+        f.write(env_content + new_line)
+
+    click.echo("✅ Successfully authenticated with PandaBI!")
+
+
 @cli.command()
 @click.argument("dataset_path")
 def pull(dataset_path):

tests/unit/cli/test_main.py → tests/test_cli.py

@@ -1,9 +1,77 @@
+import os
 from unittest.mock import MagicMock, patch
 
 import pytest
 from click.testing import CliRunner
 
-from pandasai.cli.main import cli, get_validated_dataset_path
+from pandasai.cli.main import cli, get_validated_dataset_path, validate_api_key
+
+
+def test_validate_api_key():
+    # Valid API key
+    assert validate_api_key("PAI-59ca2c4a-7998-4195-81d1-5c597f998867") == True
+
+    # Invalid API keys
+    assert validate_api_key("PAI-59ca2c4a-7998-4195-81d1") == False  # Too short
+    assert (
+        validate_api_key("XXX-59ca2c4a-7998-4195-81d1-5c597f998867") == False
+    )  # Wrong prefix
+    assert (
+        validate_api_key("PAI-59ca2c4a-7998-4195-81d1-5c597f99886") == False
+    )  # Wrong length
+    assert (
+        validate_api_key("PAI-59ca2c4a7998419581d15c597f998867") == False
+    )  # Missing hyphens
+    assert (
+        validate_api_key("PAI-XXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX") == False
+    )  # Invalid characters
+
+
+def test_login_command(tmp_path):
+    runner = CliRunner()
+
+    with runner.isolated_filesystem(temp_dir=tmp_path) as td:
+        # Test with valid API key
+        result = runner.invoke(
+            cli, ["login", "PAI-59ca2c4a-7998-4195-81d1-5c597f998867"]
+        )
+        assert result.exit_code == 0
+        assert "Successfully authenticated with PandaBI!" in result.output
+
+        # Verify .env file content
+        with open(os.path.join(td, ".env")) as f:
+            content = f.read()
+            assert "PANDABI_API_KEY=PAI-59ca2c4a-7998-4195-81d1-5c597f998867" in content
+
+        # Test with invalid API key
+        result = runner.invoke(cli, ["login", "invalid-key"])
+        assert result.exit_code == 0  # Click returns 0 for validation errors by default
+        assert "Invalid API key format" in result.output
+
+
+def test_login_command_preserves_existing_env(tmp_path):
+    runner = CliRunner()
+
+    with runner.isolated_filesystem(temp_dir=tmp_path) as td:
+        # Create .env with existing variables
+        with open(os.path.join(td, ".env"), "w") as f:
+            f.write("EXISTING_VAR=value\n")
+            f.write("PANDABI_API_KEY=PAI-old-key-that-should-be-replaced\n")
+            f.write("ANOTHER_VAR=another_value\n")
+
+        # Update API key
+        result = runner.invoke(
+            cli, ["login", "PAI-59ca2c4a-7998-4195-81d1-5c597f998867"]
+        )
+        assert result.exit_code == 0
+
+        # Verify .env file content
+        with open(os.path.join(td, ".env")) as f:
+            content = f.read().splitlines()
+            assert "EXISTING_VAR=value" in content
+            assert "ANOTHER_VAR=another_value" in content
+            assert "PANDABI_API_KEY=PAI-59ca2c4a-7998-4195-81d1-5c597f998867" in content
+            assert "PANDABI_API_KEY=PAI-old-key-that-should-be-replaced" not in content
 
 
 def test_get_validated_dataset_path_valid():
@@ -16,22 +84,26 @@ def test_get_validated_dataset_path_valid():
 def test_get_validated_dataset_path_invalid_format():
     """Test get_validated_dataset_path with invalid format"""
     with pytest.raises(
-        ValueError, match="Path must be in format: organization/dataset"
+        ValueError, match="Path must be in format 'organization/dataset'"
     ):
         get_validated_dataset_path("invalid-path")
 
 
 def test_get_validated_dataset_path_invalid_org():
     """Test get_validated_dataset_path with invalid organization name"""
     with pytest.raises(
-        ValueError, match="Organization name must be lowercase with hyphens"
+        ValueError,
+        match="Organization name must be lowercase and use hyphens instead of spaces",
     ):
         get_validated_dataset_path("INVALID_ORG/dataset")
 
 
 def test_get_validated_dataset_path_invalid_dataset():
     """Test get_validated_dataset_path with invalid dataset name"""
-    with pytest.raises(ValueError, match="Dataset name must be lowercase with hyphens"):
+    with pytest.raises(
+        ValueError,
+        match="Dataset name must be lowercase and use hyphens instead of spaces",
+    ):
         get_validated_dataset_path("my-org/INVALID_DATASET")