client, admin-client: Keep auth token fresh

fordN · fordN · commit cc9de479ee42 · 2025-11-20T11:37:52.000-03:00
diff --git a/src/amp/admin/client.py b/src/amp/admin/client.py
@@ -56,30 +56,25 @@ def __init__(self, base_url: str, auth_token: Optional[str] = None, auth: bool =
 
         self.base_url = base_url.rstrip('/')
 
-        # Resolve auth token with priority: explicit param > env var > auth file
-        resolved_token = None
+        # Resolve auth token provider with priority: explicit param > env var > auth file
+        self._get_token = None
         if auth_token:
-            # Priority 1: Explicit auth_token parameter
-            resolved_token = auth_token
+            # Priority 1: Explicit auth_token parameter (static token)
+            self._get_token = lambda: auth_token
         elif os.getenv('AMP_AUTH_TOKEN'):
-            # Priority 2: AMP_AUTH_TOKEN environment variable
-            resolved_token = os.getenv('AMP_AUTH_TOKEN')
+            # Priority 2: AMP_AUTH_TOKEN environment variable (static token)
+            env_token = os.getenv('AMP_AUTH_TOKEN')
+            self._get_token = lambda: env_token
         elif auth:
-            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth
+            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth (auto-refreshing)
             from amp.auth import AuthService
 
             auth_service = AuthService()
-            resolved_token = auth_service.get_token()
+            self._get_token = auth_service.get_token  # Callable that auto-refreshes
 
-        # Build headers
-        headers = {}
-        if resolved_token:
-            headers['Authorization'] = f'Bearer {resolved_token}'
-
-        # Create HTTP client
+        # Create HTTP client (no auth header yet - will be added per-request)
         self._http = httpx.Client(
             base_url=self.base_url,
-            headers=headers,
             timeout=30.0,
             follow_redirects=True,
         )
@@ -102,6 +97,12 @@ def _request(
         Raises:
             AdminAPIError: If the API returns an error response
         """
+        # Add auth header dynamically (auto-refreshes if needed)
+        headers = kwargs.get('headers', {})
+        if self._get_token:
+            headers['Authorization'] = f'Bearer {self._get_token()}'
+            kwargs['headers'] = headers
+
         response = self._http.request(method, path, json=json, params=params, **kwargs)
 
         # Handle error responses
diff --git a/src/amp/client.py b/src/amp/client.py
@@ -24,33 +24,33 @@
 class AuthMiddleware(ClientMiddleware):
     """Flight middleware to add Bearer token authentication header."""
 
-    def __init__(self, token: str):
+    def __init__(self, get_token):
         """Initialize auth middleware.
 
         Args:
-            token: Bearer token to add to requests
+            get_token: Callable that returns the current access token
         """
-        self.token = token
+        self.get_token = get_token
 
     def sending_headers(self):
         """Add Authorization header to outgoing requests."""
-        return {'authorization': f'Bearer {self.token}'}
+        return {'authorization': f'Bearer {self.get_token()}'}
 
 
 class AuthMiddlewareFactory(ClientMiddlewareFactory):
     """Factory for creating auth middleware instances."""
 
-    def __init__(self, token: str):
+    def __init__(self, get_token):
         """Initialize auth middleware factory.
 
         Args:
-            token: Bearer token to use for authentication
+            get_token: Callable that returns the current access token
         """
-        self.token = token
+        self.get_token = get_token
 
     def start_call(self, info):
         """Create auth middleware for each call."""
-        return AuthMiddleware(self.token)
+        return AuthMiddleware(self.get_token)
 
 
 class QueryBuilder:
@@ -307,26 +307,30 @@ def __init__(
         if url and not query_url:
             query_url = url
 
-        # Resolve auth token with priority: explicit param > env var > auth file
-        flight_auth_token = None
+        # Resolve auth token provider with priority: explicit param > env var > auth file
+        get_token = None
         if auth_token:
-            # Priority 1: Explicit auth_token parameter
-            flight_auth_token = auth_token
+            # Priority 1: Explicit auth_token parameter (static token)
+            def get_token():
+                return auth_token
         elif os.getenv('AMP_AUTH_TOKEN'):
-            # Priority 2: AMP_AUTH_TOKEN environment variable
-            flight_auth_token = os.getenv('AMP_AUTH_TOKEN')
+            # Priority 2: AMP_AUTH_TOKEN environment variable (static token)
+            env_token = os.getenv('AMP_AUTH_TOKEN')
+
+            def get_token():
+                return env_token
         elif auth:
-            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth
+            # Priority 3: Load from ~/.amp-cli-config/amp_cli_auth (auto-refreshing)
             from amp.auth import AuthService
 
             auth_service = AuthService()
-            flight_auth_token = auth_service.get_token()
+            get_token = auth_service.get_token  # Callable that auto-refreshes
 
         # Initialize Flight SQL client
         if query_url:
-            # Add auth middleware if token is provided
-            if flight_auth_token:
-                middleware = [AuthMiddlewareFactory(flight_auth_token)]
+            # Add auth middleware if token provider exists
+            if get_token:
+                middleware = [AuthMiddlewareFactory(get_token)]
                 self.conn = flight.connect(query_url, middleware=middleware)
             else:
                 self.conn = flight.connect(query_url)
@@ -342,8 +346,18 @@ def __init__(
         if admin_url:
             from amp.admin.client import AdminClient
 
-            # Pass resolved token to AdminClient (maintains same priority logic)
-            self._admin_client = AdminClient(admin_url, auth_token=flight_auth_token, auth=False)
+            # Pass auth=True if we have a get_token callable from auth file
+            # Otherwise pass the static token if available
+            if auth:
+                # Use auth file (auto-refreshing)
+                self._admin_client = AdminClient(admin_url, auth=True)
+            elif auth_token or os.getenv('AMP_AUTH_TOKEN'):
+                # Use static token
+                static_token = auth_token or os.getenv('AMP_AUTH_TOKEN')
+                self._admin_client = AdminClient(admin_url, auth_token=static_token)
+            else:
+                # No auth
+                self._admin_client = AdminClient(admin_url)
         else:
             self._admin_client = None
 
diff --git a/tests/integration/admin/test_admin_client.py b/tests/integration/admin/test_admin_client.py
@@ -25,8 +25,7 @@ def test_admin_client_with_auth_token(self):
         """Test AdminClient with authentication token."""
         client = AdminClient('http://localhost:8080', auth_token='test-token')
 
-        assert 'Authorization' in client._http.headers
-        assert client._http.headers['Authorization'] == 'Bearer test-token'
+        assert client._get_token() == 'test-token'
 
     @respx.mock
     def test_request_success(self):
diff --git a/tests/integration/test_snowflake_loader.py b/tests/integration/test_snowflake_loader.py
@@ -67,7 +67,7 @@ def wait_for_snowpipe_data(loader, table_name, expected_count, max_wait=30, poll
 
 
 # Skip all Snowflake tests
-# pytestmark = pytest.mark.skip(reason='Requires active Snowflake account - see module docstring for details')
+pytestmark = pytest.mark.skip(reason='Requires active Snowflake account - see module docstring for details')
 
 
 @pytest.fixture
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -104,7 +104,7 @@ def test_explicit_token_highest_priority(self, mock_connect, mock_getenv):
         call_args = mock_connect.call_args
         middleware = call_args[1].get('middleware', [])
         assert len(middleware) == 1
-        assert middleware[0].token == 'explicit-token'
+        assert middleware[0].get_token() == 'explicit-token'
 
     @patch('amp.client.os.getenv')
     @patch('amp.client.flight.connect')
@@ -128,7 +128,7 @@ def getenv_side_effect(key, default=None):
         call_args = mock_connect.call_args
         middleware = call_args[1].get('middleware', [])
         assert len(middleware) == 1
-        assert middleware[0].token == 'env-var-token'
+        assert middleware[0].get_token() == 'env-var-token'
 
     @patch('amp.auth.AuthService')
     @patch('amp.client.os.getenv')
@@ -150,12 +150,12 @@ def getenv_side_effect(key, default=None):
 
         # 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'
+        # The middleware should use the auth service's get_token method directly
+        assert middleware[0].get_token == mock_service_instance.get_token
 
     @patch('amp.client.os.getenv')
     @patch('amp.client.flight.connect')
@@ -190,8 +190,8 @@ def test_admin_explicit_token_highest_priority(self, mock_getenv):
 
         client = AdminClient('http://localhost:8080', auth_token='explicit-token')
 
-        # Verify explicit token was used
-        assert client._http.headers.get('Authorization') == 'Bearer explicit-token'
+        # Verify explicit token was used (check get_token callable)
+        assert client._get_token() == 'explicit-token'
 
     @patch('amp.admin.client.os.getenv')
     def test_admin_env_var_second_priority(self, mock_getenv):
@@ -204,7 +204,7 @@ def test_admin_env_var_second_priority(self, mock_getenv):
 
         # Verify env var was used
         mock_getenv.assert_called_with('AMP_AUTH_TOKEN')
-        assert client._http.headers.get('Authorization') == 'Bearer env-var-token'
+        assert client._get_token() == 'env-var-token'
 
     @patch('amp.auth.AuthService')
     @patch('amp.admin.client.os.getenv')
@@ -220,7 +220,7 @@ def test_admin_auth_file_lowest_priority(self, mock_getenv, mock_auth_service):
         client = AdminClient('http://localhost:8080', auth=True)
 
         # Verify auth file was used
-        assert client._http.headers.get('Authorization') == 'Bearer file-token'
+        assert client._get_token == mock_service_instance.get_token
 
     @patch('amp.admin.client.os.getenv')
     def test_admin_no_auth_when_nothing_provided(self, mock_getenv):
