feat(auth): Support multiple ways of passing in auth

- `AMP_AUTH_TOKEN` env var, `auth_token` param, or locally stored
auth file (from interactive browser login)

Author: fordN

src/amp/admin/client.py

@@ -4,6 +4,7 @@
 with the Amp Admin API over HTTP.
 """
 
+import os
 from typing import Optional
 
 import httpx
@@ -19,15 +20,24 @@ class AdminClient:
 
     Args:
         base_url: Base URL for Admin API (e.g., 'http://localhost:8080')
-        auth_token: Optional Bearer token for authentication
+        auth_token: Optional Bearer token for authentication (highest priority)
         auth: If True, load auth token from ~/.amp-cli-config (shared with TS CLI)
 
+    Authentication Priority (highest to lowest):
+        1. Explicit auth_token parameter
+        2. AMP_AUTH_TOKEN environment variable
+        3. auth=True - reads from ~/.amp-cli-config/amp_cli_auth
+
     Example:
-        >>> # Use amp auth system
+        >>> # Use amp auth from file
         >>> client = AdminClient('http://localhost:8080', auth=True)
         >>>
-        >>> # Or use manual token
+        >>> # Use manual token
         >>> client = AdminClient('http://localhost:8080', auth_token='your-token')
+        >>>
+        >>> # Use environment variable
+        >>> # export AMP_AUTH_TOKEN="eyJhbGci..."
+        >>> client = AdminClient('http://localhost:8080')
     """
 
     def __init__(self, base_url: str, auth_token: Optional[str] = None, auth: bool = False):
@@ -46,17 +56,25 @@ def __init__(self, base_url: str, auth_token: Optional[str] = None, auth: bool =
 
         self.base_url = base_url.rstrip('/')
 
-        # Load token from amp auth system if requested
-        if auth:
+        # Resolve auth token with priority: explicit param > env var > auth file
+        resolved_token = None
+        if auth_token:
+            # Priority 1: Explicit auth_token parameter
+            resolved_token = auth_token
+        elif os.getenv('AMP_AUTH_TOKEN'):
+            # Priority 2: AMP_AUTH_TOKEN environment variable
+            resolved_token = os.getenv('AMP_AUTH_TOKEN')
+        elif auth:
+            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth
             from amp.auth import AuthService
 
             auth_service = AuthService()
-            auth_token = auth_service.get_token()
+            resolved_token = auth_service.get_token()
 
         # Build headers
         headers = {}
-        if auth_token:
-            headers['Authorization'] = f'Bearer {auth_token}'
+        if resolved_token:
+            headers['Authorization'] = f'Bearer {resolved_token}'
 
         # Create HTTP client
         self._http = httpx.Client(

src/amp/client.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Dict, Iterator, List, Optional, Union
 
 import pyarrow as pa
@@ -270,19 +271,28 @@ class Client:
         url: Flight SQL URL (for backward compatibility, treated as query_url)
         query_url: Query endpoint URL via Flight SQL (e.g., 'grpc://localhost:1602')
         admin_url: Optional Admin API URL (e.g., 'http://localhost:8080')
-        auth_token: Optional Bearer token for Admin API authentication
+        auth_token: Optional Bearer token for authentication (highest priority)
         auth: If True, load auth token from ~/.amp-cli-config (shared with TS CLI)
 
+    Authentication Priority (highest to lowest):
+        1. Explicit auth_token parameter
+        2. AMP_AUTH_TOKEN environment variable
+        3. auth=True - reads from ~/.amp-cli-config/amp_cli_auth
+
     Example:
         >>> # Query-only client (backward compatible)
         >>> client = Client(url='grpc://localhost:1602')
         >>>
-        >>> # Client with admin capabilities and amp auth
+        >>> # Client with amp auth from file
         >>> client = Client(
         ...     query_url='grpc://localhost:1602',
         ...     admin_url='http://localhost:8080',
         ...     auth=True
         ... )
+        >>>
+        >>> # Client with auth from environment variable
+        >>> # export AMP_AUTH_TOKEN="eyJhbGci..."
+        >>> client = Client(query_url='grpc://localhost:1602')
     """
 
     def __init__(
@@ -297,15 +307,20 @@ def __init__(
         if url and not query_url:
             query_url = url
 
-        # Get auth token if using amp auth system
+        # Resolve auth token with priority: explicit param > env var > auth file
         flight_auth_token = None
-        if auth and not auth_token:
+        if auth_token:
+            # Priority 1: Explicit auth_token parameter
+            flight_auth_token = auth_token
+        elif os.getenv('AMP_AUTH_TOKEN'):
+            # Priority 2: AMP_AUTH_TOKEN environment variable
+            flight_auth_token = os.getenv('AMP_AUTH_TOKEN')
+        elif auth:
+            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth
             from amp.auth import AuthService
 
             auth_service = AuthService()
             flight_auth_token = auth_service.get_token()
-        elif auth_token:
-            flight_auth_token = auth_token
 
         # Initialize Flight SQL client
         if query_url:
@@ -327,7 +342,8 @@ def __init__(
         if admin_url:
             from amp.admin.client import AdminClient
 
-            self._admin_client = AdminClient(admin_url, auth_token=auth_token, auth=auth)
+            # Pass resolved token to AdminClient (maintains same priority logic)
+            self._admin_client = AdminClient(admin_url, auth_token=flight_auth_token, auth=False)
         else:
             self._admin_client = None
 

tests/unit/test_client.py

@@ -7,7 +7,7 @@
 
 import json
 from pathlib import Path
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 
@@ -87,6 +87,151 @@ def test_client_requires_url_or_query_url(self):
             Client()
 
 
+@pytest.mark.unit
+class TestClientAuthPriority:
+    """Test Client authentication priority (explicit token > env var > auth file)"""
+
+    @patch('amp.client.os.getenv')
+    @patch('amp.client.flight.connect')
+    def test_explicit_token_highest_priority(self, mock_connect, mock_getenv):
+        """Test that explicit auth_token parameter has highest priority"""
+        mock_getenv.return_value = 'env-var-token'
+
+        client = Client(query_url='grpc://localhost:1602', auth_token='explicit-token')
+
+        # Verify that explicit token was used (not env var)
+        mock_connect.assert_called_once()
+        call_args = mock_connect.call_args
+        middleware = call_args[1].get('middleware', [])
+        assert len(middleware) == 1
+        assert middleware[0].token == 'explicit-token'
+
+    @patch('amp.client.os.getenv')
+    @patch('amp.client.flight.connect')
+    def test_env_var_second_priority(self, mock_connect, mock_getenv):
+        """Test that AMP_AUTH_TOKEN env var has second priority"""
+        # Return 'env-var-token' for AMP_AUTH_TOKEN, None for others
+        def getenv_side_effect(key, default=None):
+            if key == 'AMP_AUTH_TOKEN':
+                return 'env-var-token'
+            return default
+
+        mock_getenv.side_effect = getenv_side_effect
+
+        client = Client(query_url='grpc://localhost:1602')
+
+        # Verify env var was checked
+        calls = [str(call) for call in mock_getenv.call_args_list]
+        assert any('AMP_AUTH_TOKEN' in call for call in calls)
+        mock_connect.assert_called_once()
+        call_args = mock_connect.call_args
+        middleware = call_args[1].get('middleware', [])
+        assert len(middleware) == 1
+        assert middleware[0].token == 'env-var-token'
+
+    @patch('amp.auth.AuthService')
+    @patch('amp.client.os.getenv')
+    @patch('amp.client.flight.connect')
+    def test_auth_file_lowest_priority(self, mock_connect, mock_getenv, mock_auth_service):
+        """Test that auth=True has lowest priority"""
+        # Return None for all getenv calls
+        def getenv_side_effect(key, default=None):
+            return default
+
+        mock_getenv.side_effect = getenv_side_effect
+
+        mock_service_instance = Mock()
+        mock_service_instance.get_token.return_value = 'file-token'
+        mock_auth_service.return_value = mock_service_instance
+
+        client = Client(query_url='grpc://localhost:1602', auth=True)
+
+        # Verify auth file was used
+        mock_auth_service.assert_called_once()
+        mock_service_instance.get_token.assert_called_once()
+        mock_connect.assert_called_once()
+        call_args = mock_connect.call_args
+        middleware = call_args[1].get('middleware', [])
+        assert len(middleware) == 1
+        assert middleware[0].token == 'file-token'
+
+    @patch('amp.client.os.getenv')
+    @patch('amp.client.flight.connect')
+    def test_no_auth_when_nothing_provided(self, mock_connect, mock_getenv):
+        """Test that no auth middleware is added when no auth is provided"""
+        # Return None/default for all getenv calls
+        def getenv_side_effect(key, default=None):
+            return default
+
+        mock_getenv.side_effect = getenv_side_effect
+
+        client = Client(query_url='grpc://localhost:1602')
+
+        # Verify no middleware was added
+        mock_connect.assert_called_once()
+        call_args = mock_connect.call_args
+        middleware = call_args[1].get('middleware')
+        assert middleware is None or len(middleware) == 0
+
+
+@pytest.mark.unit
+class TestAdminClientAuthPriority:
+    """Test AdminClient authentication priority"""
+
+    @patch('amp.admin.client.os.getenv')
+    def test_admin_explicit_token_highest_priority(self, mock_getenv):
+        """Test that explicit auth_token parameter has highest priority for AdminClient"""
+        from amp.admin.client import AdminClient
+
+        mock_getenv.return_value = 'env-var-token'
+
+        client = AdminClient('http://localhost:8080', auth_token='explicit-token')
+
+        # Verify explicit token was used
+        assert client._http.headers.get('Authorization') == 'Bearer explicit-token'
+
+    @patch('amp.admin.client.os.getenv')
+    def test_admin_env_var_second_priority(self, mock_getenv):
+        """Test that AMP_AUTH_TOKEN env var has second priority for AdminClient"""
+        from amp.admin.client import AdminClient
+
+        mock_getenv.return_value = 'env-var-token'
+
+        client = AdminClient('http://localhost:8080')
+
+        # Verify env var was used
+        mock_getenv.assert_called_with('AMP_AUTH_TOKEN')
+        assert client._http.headers.get('Authorization') == 'Bearer env-var-token'
+
+    @patch('amp.auth.AuthService')
+    @patch('amp.admin.client.os.getenv')
+    def test_admin_auth_file_lowest_priority(self, mock_getenv, mock_auth_service):
+        """Test that auth=True has lowest priority for AdminClient"""
+        from amp.admin.client import AdminClient
+
+        mock_getenv.return_value = None
+        mock_service_instance = Mock()
+        mock_service_instance.get_token.return_value = 'file-token'
+        mock_auth_service.return_value = mock_service_instance
+
+        client = AdminClient('http://localhost:8080', auth=True)
+
+        # Verify auth file was used
+        assert client._http.headers.get('Authorization') == 'Bearer file-token'
+
+    @patch('amp.admin.client.os.getenv')
+    def test_admin_no_auth_when_nothing_provided(self, mock_getenv):
+        """Test that no auth header is added when no auth is provided"""
+        from amp.admin.client import AdminClient
+
+        mock_getenv.return_value = None
+
+        client = AdminClient('http://localhost:8080')
+
+        # Verify no auth header
+        assert 'Authorization' not in client._http.headers
+
+
 @pytest.mark.unit
 class TestQueryBuilderManifest:
     """Test QueryBuilder manifest generation"""