Merge pull request #901 from vespa-engine/thomasht86/redo-authclients

Remove `get_valid_auth_clients`

.github/workflows/notebooks-cloud.yml

@@ -59,6 +59,7 @@ jobs:
       - name: Run notebooks tests
         env:
           VESPA_TEAM_API_KEY: ${{ secrets.VESPA_TEAM_API_KEY }}
+          VESPA_CLOUD_SECRET_TOKEN: ${{ secrets.VESPA_CLOUD_SECRET_TOKEN }}
           CO_API_KEY: ${{ secrets.CO_API_KEY }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         run: |

tests/integration/test_integration_vespa_cloud_token.py

@@ -55,9 +55,15 @@ def setUp(self) -> None:
         )
         self.disk_folder = os.path.join(os.getcwd(), "sample_application")
         self.instance_name = "token"
-        self.app: Vespa = self.vespa_cloud.deploy(
+        self.mtls_app: Vespa = self.vespa_cloud.deploy(
             instance=self.instance_name, disk_folder=self.disk_folder
         )
+        self.app = self.vespa_cloud.get_application(
+            instance=self.instance_name,
+            environment="dev",
+            endpoint_type="token",
+            vespa_cloud_secret_token=os.getenv("VESPA_CLOUD_SECRET_TOKEN"),
+        )
         print("Endpoint used " + self.app.url)
 
     def test_right_endpoint_used_with_token(self):
@@ -111,9 +117,15 @@ def setUp(self) -> None:
         )
         self.disk_folder = os.path.join(os.getcwd(), "sample_application")
         self.instance_name = "token"
-        self.app = self.vespa_cloud.deploy(
+        self.mtls_app = self.vespa_cloud.deploy(
             instance=self.instance_name, disk_folder=self.disk_folder
         )
+        self.app = self.vespa_cloud.get_application(
+            instance=self.instance_name,
+            environment="dev",
+            endpoint_type="token",
+            vespa_cloud_secret_token=os.getenv("VESPA_CLOUD_SECRET_TOKEN"),
+        )
         print("Endpoint used " + self.app.url)
         self.fields_to_send = [
             {
@@ -243,7 +255,10 @@ def test_application_status(self):
         if not success:
             self.fail("Deployment failed")
         self.app = self.vespa_cloud.get_application(
-            instance=self.instance_name, environment="prod"
+            instance=self.instance_name,
+            environment="prod",
+            endpoint_type="token",
+            vespa_cloud_secret_token=os.getenv("VESPA_CLOUD_SECRET_TOKEN"),
         )
         self.app.wait_for_application_up(max_wait=APP_INIT_TIMEOUT)
 

tests/unit/test_application.py

@@ -224,7 +224,10 @@ def test_query_token_from_env(self):
 
         os.environ["VESPA_CLOUD_SECRET_TOKEN"] = "vespa_cloud_str_secret"
         self.assertEqual(
-            Vespa(url="https://cord19.vespa.ai").vespa_cloud_secret_token,
+            Vespa(
+                url="https://cord19.vespa.ai",
+                vespa_cloud_secret_token=os.getenv("VESPA_CLOUD_SECRET_TOKEN"),
+            ).vespa_cloud_secret_token,
             "vespa_cloud_str_secret",
         )
 

vespa/application.py

@@ -8,7 +8,6 @@
 from concurrent.futures import ThreadPoolExecutor, Future, as_completed
 from queue import Queue, Empty
 import threading
-import requests
 from requests import Session
 from requests.models import Response
 from requests.exceptions import ConnectionError, HTTPError, JSONDecodeError
@@ -26,7 +25,6 @@
     RetryCallState,
 )
 from time import sleep
-from os import environ
 from urllib.parse import quote
 import random
 import time
@@ -121,11 +119,13 @@ def __init__(
         else:
             self.end_point = str(url).rstrip("/") + ":" + str(port)
         self.search_end_point = self.end_point + "/search/"
-        if vespa_cloud_secret_token is None:
-            token = environ.get(VESPA_CLOUD_SECRET_TOKEN, None)
-            if token is not None:
-                self.vespa_cloud_secret_token = token
-        self.auth_method = None
+        if self.vespa_cloud_secret_token is not None:
+            self.auth_method = "token"
+            self.base_headers.update(
+                {"Authorization": f"Bearer {self.vespa_cloud_secret_token}"}
+            )
+        else:
+            self.auth_method = "mtls"
 
     def asyncio(
         self, connections: Optional[int] = 8, total_timeout: int = 10
@@ -250,62 +250,6 @@ def wait_for_application_up(self, max_wait: int = 300) -> None:
                 )
             )
 
-    def _get_valid_auth_method(self) -> Optional[str]:
-        """
-        Get auth method for Vespa connection.
-
-        :return: Auth method used for Vespa connection. Either 'token','mtls_key_cert','mtls_cert' or 'http'. None if not able to authenticate.
-        """
-        endpoint = f"{self.end_point}/ApplicationStatus"
-
-        if self.auth_method:
-            return self.auth_method
-
-        # Plain HTTP
-        response = requests.get(endpoint, headers=self.base_headers)
-        if response.status_code == 200:
-            print(
-                f"Using plain HTTP to connect to Vespa endpoint {self.end_point}",
-                file=self.output_file,
-            )
-            return "http"
-
-        # Vespa Cloud Secret Token
-        if self.vespa_cloud_secret_token is not None:
-            headers = {"Authorization": f"Bearer {self.vespa_cloud_secret_token}"}
-            response = requests.get(endpoint, headers={**self.base_headers, **headers})
-            if response.status_code == 200:
-                print(
-                    f"Using Vespa Cloud Secret Token to connect to Vespa endpoint {self.end_point}",
-                    file=self.output_file,
-                )
-                return "token"
-
-        # Mutual TLS with key and cert
-        if self.key and self.cert:
-            response = requests.get(
-                endpoint, headers=self.base_headers, cert=(self.cert, self.key)
-            )
-            if response.status_code == 200:
-                print(
-                    f"Using Mutual TLS with key and cert to connect to Vespa endpoint {self.end_point}",
-                    file=self.output_file,
-                )
-                return "mtls_key_cert"
-
-        # Mutual TLS with cert
-        if self.cert:
-            response = requests.get(endpoint, headers=self.base_headers, cert=self.cert)
-            if response.status_code == 200:
-                print(
-                    f"Using Mutual TLS with cert to connect to Vespa endpoint {self.end_point}",
-                    file=self.output_file,
-                )
-                return "mtls_cert"
-
-        # There may be some cases where ApplicationStatus is not available, such as http://api.cord19.vespa.ai
-        return None
-
     def get_application_status(self) -> Optional[Response]:
         """
         Get application status (/ApplicationStatus)
@@ -1054,7 +998,6 @@ def __init__(
             self.cert = (self.app.cert, self.app.key)
         else:
             self.cert = self.app.cert
-        self.app.auth_method = self.app._get_valid_auth_method()
         self.headers = self.app.base_headers.copy()
         if self.app.auth_method == "token" and self.app.vespa_cloud_secret_token:
             # Bearer and user-agent
@@ -1472,7 +1415,6 @@ def __init__(
         self.httpx_client = None
         self.connections = connections
         self.total_timeout = total_timeout
-        self.app.auth_method = self.app._get_valid_auth_method()
         self.headers = self.app.base_headers.copy()
         if self.app.auth_method == "token" and self.app.vespa_cloud_secret_token:
             # Bearer and user-agent

vespa/deployment.py

@@ -28,7 +28,7 @@
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import ec
 
-from vespa.application import Vespa, VESPA_CLOUD_SECRET_TOKEN
+from vespa.application import Vespa
 from vespa.package import ApplicationPackage
 from vespa.utils.notebook import is_jupyter_notebook
 import vespa
@@ -619,7 +619,7 @@ def deploy(
             folder within user's current working directory.
         :param max_wait: Seconds to wait for the deployment.
 
-        :return: a Vespa connection instance.
+        :return: a Vespa connection instance. Returns a connection to the mtls endpoint. To connect to the token endpoint, use :func:`VespaCloud.get_application(endpoint_type="token")`.
         """
         if not disk_folder:
             disk_folder = os.path.join(os.getcwd(), self.application)
@@ -629,33 +629,9 @@ def deploy(
         job = "dev-" + region
         run = self._start_deployment(instance, job, disk_folder, None)
         self._follow_deployment(instance, job, run)
-
-        mtls_endpoint = self.get_mtls_endpoint(
-            instance=instance,
-            region=region,
-            environment="dev",
-        )
-        if self.auth_client_token_id is not None:
-            try:  # May have client_token_id set but the deployed app was not configured to use it
-                token_endpoint = self.get_token_endpoint(
-                    instance=instance,
-                    region=region,
-                    environment="dev",
-                )
-            except Exception as _:
-                token_endpoint = None
-        else:
-            token_endpoint = None
-        print(f"Connecting to {token_endpoint or mtls_endpoint}", file=self.output)
-        app = Vespa(
-            url=token_endpoint or mtls_endpoint,
-            cert=self.data_cert_path,
-            key=self.data_key_path or None,
-            application_package=self.application_package,
-            vespa_cloud_secret_token=os.environ.get(VESPA_CLOUD_SECRET_TOKEN),
+        app: Vespa = self.get_application(
+            instance=instance, environment="dev", endpoint_type="mtls"
         )
-        app.wait_for_application_up(max_wait=max_wait)
-        print("Finished deployment.", file=self.output)
         return app
 
     def deploy_to_prod(
@@ -728,6 +704,8 @@ def get_application(
         self,
         instance: str = "default",
         environment: str = "dev",
+        endpoint_type: str = "mtls",
+        vespa_cloud_secret_token: Optional[str] = None,
         region: Optional[str] = None,
         max_wait: int = 60,
     ) -> Vespa:
@@ -743,11 +721,15 @@ def get_application(
 
         :param instance: Name of this instance of the application, in the Vespa Cloud. Default is "default".
         :param environment: Environment of the application. Default is "dev". Options are "dev" or "prod".
+        :param endpoint_type: Type of endpoint to connect to. Default is "mtls". Options are "mtls" or "token".
+        :param vespa_cloud_secret_token: Vespa Cloud Secret Token. Only required if endpoint_type is "token".
         :param region: Region of the application in Vespa cloud, eg "aws-us-east-1c". If not provided, the first region from the environment will be used.
         :param max_wait: Seconds to wait for the application to be up. Default is 60 seconds.
 
         :return: Vespa application instance.
         """
+        if endpoint_type not in ["mtls", "token"]:
+            raise ValueError("Endpoint type must be 'mtls' or 'token'.")
         if environment == "dev":
             region = self.get_dev_region()
             print(
@@ -764,31 +746,35 @@ def get_application(
                 region = valid_regions[0]
         else:
             raise ValueError("Environment must be 'dev' or 'prod'.")
-
-        mtls_endpoint = self.get_mtls_endpoint(
-            instance=instance, region=region, environment=environment
-        )
-        if self.auth_client_token_id is not None:
+        if endpoint_type == "mtls":
+            mtls_endpoint = self.get_mtls_endpoint(
+                instance=instance, region=region, environment=environment
+            )
+            app: Vespa = Vespa(
+                url=mtls_endpoint,
+                cert=self.data_cert_path,
+                key=self.data_key_path or None,
+                application_package=self.application_package,
+            )
+        elif endpoint_type == "token":
             try:  # May have client_token_id set but the deployed app was not configured to use it
                 token_endpoint = self.get_token_endpoint(
                     instance=instance, region=region, environment=environment
                 )
             except Exception as _:
-                token_endpoint = None
-        else:
-            token_endpoint = None
-        if token_endpoint is None and mtls_endpoint is None:
-            raise ValueError(
-                "No token endpoint or mtls endpoint found. Please check your deployment."
+                raise ValueError(
+                    "No token endpoint found. Make sure the application is configured with a token endpoint."
+                )
+            if vespa_cloud_secret_token is None:
+                raise ValueError(
+                    "Vespa Cloud Secret Token must be provided for token based authentication."
+                )
+            app: Vespa = Vespa(
+                url=token_endpoint,
+                application_package=self.application_package,
+                vespa_cloud_secret_token=vespa_cloud_secret_token,
             )
-        print(f"Connecting to {token_endpoint or mtls_endpoint}", file=self.output)
-        app: Vespa = Vespa(
-            url=token_endpoint or mtls_endpoint,
-            cert=self.data_cert_path,
-            key=self.data_key_path,
-            application_package=self.application_package,
-            vespa_cloud_secret_token=os.environ.get(VESPA_CLOUD_SECRET_TOKEN),
-        )
+        app.wait_for_application_up(max_wait=max_wait)
         return app
 
     def check_production_build_status(self, build_no: Optional[int]) -> dict:
@@ -878,7 +864,7 @@ def deploy_from_disk(
         :param instance: Name of the instance where the application is to be run
         :param application_root: Application package directory root
         :param max_wait: Seconds to wait for the deployment.
-        :return: a Vespa connection instance.
+        :return: a Vespa connection instance.  Returns a connection to the mtls endpoint. To connect to the token endpoint, use :func:`VespaCloud.get_application(endpoint_type="token")`.
         """
         data = BytesIO(self.read_app_package_from_disk(application_root))
 
@@ -890,26 +876,11 @@ def deploy_from_disk(
             instance, job, disk_folder, application_zip_bytes=data
         )
         self._follow_deployment(instance, job, run)
-        mtls_endpoint = self.get_mtls_endpoint(instance=instance, region=region)
-        if self.auth_client_token_id is not None:
-            try:  # May have client_token_id set but the deployed app was not configured to use it
-                token_endpoint = self.get_token_endpoint(
-                    instance=instance, region=region
-                )
-            except Exception as _:
-                token_endpoint = None
-        else:
-            token_endpoint = None
-        app = Vespa(
-            url=token_endpoint or mtls_endpoint,
-            cert=self.data_cert_path,
-            key=self.data_key_path,
-            application_package=self.application_package,
-            vespa_cloud_secret_token=os.environ.get(VESPA_CLOUD_SECRET_TOKEN),
+        run = self._start_deployment(instance, job, disk_folder, None)
+        self._follow_deployment(instance, job, run)
+        app: Vespa = self.get_application(
+            instance=instance, environment="dev", endpoint_type="mtls"
         )
-        app.wait_for_application_up(max_wait=max_wait)
-        print("Finished deployment.", file=self.output)
-
         return app
 
     def delete(self, instance: Optional[str] = "default") -> None: