locust-cloud 1.12.3__py3-none-any.whl → 1.13.0__py3-none-any.whl

locust_cloud/cloud.py

@@ -1,21 +1,29 @@
 import base64
 import gzip
+import importlib.metadata
+import json
 import logging
 import os
+import pathlib
 import sys
+import threading
 import time
 import tomllib
 import urllib.parse
+import webbrowser
 from argparse import Namespace
 from collections import OrderedDict
+from dataclasses import dataclass
 from typing import IO, Any
 
 import configargparse
+import jwt
+import platformdirs
 import requests
 import socketio
 import socketio.exceptions
-from locust_cloud import __version__
-from locust_cloud.credential_manager import CredentialError, CredentialManager
+
+__version__ = importlib.metadata.version("locust-cloud")
 
 
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
@@ -49,7 +57,7 @@ parser = configargparse.ArgumentParser(
         "cloud.conf",
     ],
     auto_env_var_prefix="LOCUSTCLOUD_",
-    formatter_class=configargparse.RawDescriptionHelpFormatter,
+    formatter_class=configargparse.RawTextHelpFormatter,
     config_file_parser_class=configargparse.CompositeConfigParser(
         [
             LocustTomlConfigParser(["tool.locust"]),
@@ -108,36 +116,16 @@ advanced.add_argument(
     help="Optional requirements.txt file that contains your external libraries.",
 )
 advanced.add_argument(
-    "--region",
-    type=str,
-    default=os.environ.get("AWS_DEFAULT_REGION"),
-    help="Sets the AWS region to use for the deployed cluster, e.g. us-east-1. It defaults to use AWS_DEFAULT_REGION env var, like AWS tools.",
-)
-parser.add_argument(
-    "--aws-access-key-id",
-    type=str,
-    help=configargparse.SUPPRESS,
-    env_var="AWS_ACCESS_KEY_ID",
-    default=None,
-)
-parser.add_argument(
-    "--aws-secret-access-key",
-    type=str,
-    help=configargparse.SUPPRESS,
-    env_var="AWS_SECRET_ACCESS_KEY",
-    default=None,
-)
-parser.add_argument(
-    "--username",
-    type=str,
+    "--login",
+    action="store_true",
+    default=False,
     help=configargparse.SUPPRESS,
-    default=os.getenv("LOCUST_CLOUD_USERNAME", None),  # backwards compatitibility for dmdb
 )
-parser.add_argument(
-    "--password",
-    type=str,
-    help=configargparse.SUPPRESS,
-    default=os.getenv("LOCUST_CLOUD_PASSWORD", None),  # backwards compatitibility for dmdb
+advanced.add_argument(
+    "--non-interactive",
+    action="store_true",
+    default=False,
+    help="This can be set when, for example, running in a CI/CD environment to ensure no interactive steps while executing.\nRequires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set.",
 )
 parser.add_argument(
     "--workers",
@@ -153,7 +141,7 @@ parser.add_argument(
 parser.add_argument(
     "--image-tag",
     type=str,
-    default="latest",
+    default=None,
     help=configargparse.SUPPRESS,  # overrides the locust-cloud docker image tag. for internal use
 )
 parser.add_argument(
@@ -169,6 +157,7 @@ parser.add_argument(
 )
 
 options, locust_options = parser.parse_known_args()
+
 options: Namespace
 locust_options: list
 
@@ -183,51 +172,222 @@ logging.getLogger("boto3").setLevel(logging.INFO)
 logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
+cloud_conf_file = pathlib.Path(platformdirs.user_config_dir(appname="locust-cloud")) / "config"
+valid_regions = ["us-east-1", "eu-north-1"]
 
-api_url = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{options.region}.locust.cloud/1")
 
+def get_api_url(region):
+    return os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{region}.locust.cloud/1")
 
-def main() -> None:
-    if options.version:
-        print(f"locust-cloud version {__version__}")
-        sys.exit(0)
 
-    if not options.region:
-        logger.error(
-            "Setting a region is required to use Locust Cloud. Please ensure the AWS_DEFAULT_REGION env variable or the --region flag is set."
-        )
+@dataclass
+class CloudConfig:
+    id_token: str | None = None
+    refresh_token: str | None = None
+    refresh_token_expires: int = 0
+    region: str | None = None
+
+
+def read_cloud_config() -> CloudConfig:
+    if cloud_conf_file.exists():
+        with open(cloud_conf_file) as f:
+            return CloudConfig(**json.load(f))
+
+    return CloudConfig()
+
+
+def write_cloud_config(config: CloudConfig) -> None:
+    cloud_conf_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(cloud_conf_file, "w") as f:
+        json.dump(config.__dict__, f)
+
+
+def web_login() -> None:
+    print("Enter the number for the region to authenticate against")
+    print()
+    for i, valid_region in enumerate(valid_regions, start=1):
+        print(f"  {i}. {valid_region}")
+    print()
+    choice = input("> ")
+    try:
+        region_index = int(choice) - 1
+        assert 0 <= region_index < len(valid_regions)
+    except (ValueError, AssertionError):
+        print(f"Not a valid choice: '{choice}'")
         sys.exit(1)
-    if options.region:
-        os.environ["AWS_DEFAULT_REGION"] = options.region
 
-    if not ((options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)):
-        logger.error(
-            "Authentication is required to use Locust Cloud. Please ensure the LOCUSTCLOUD_USERNAME and LOCUSTCLOUD_PASSWORD environment variables are set."
-        )
+    region = valid_regions[region_index]
+
+    try:
+        response = requests.post(f"{get_api_url(region)}/cli-auth")
+        response.raise_for_status()
+        response_data = response.json()
+        authentication_url = response_data["authentication_url"]
+        result_url = response_data["result_url"]
+    except Exception as e:
+        print("Something went wrong trying to authorize the locust-cloud CLI:", str(e))
         sys.exit(1)
+
+    message = f"""
+Attempting to automatically open the SSO authorization page in your default browser.
+If the browser does not open or you wish to use a different device to authorize this request, open the following URL:
+
+{authentication_url}
+    """.strip()
+    print()
+    print(message)
+
+    webbrowser.open_new_tab(authentication_url)
+
+    while True:  # Should there be some kind of timeout?
+        response = requests.get(result_url)
+
+        if not response.ok:
+            print("Oh no!")
+            print(response.text)
+            sys.exit(1)
+
+        data = response.json()
+
+        if data["state"] == "pending":
+            time.sleep(1)
+            continue
+        elif data["state"] == "failed":
+            print(f"\nFailed to authorize CLI: {data['reason']}")
+            sys.exit(1)
+        elif data["state"] == "authorized":
+            print("\nAuthorization succeded")
+            break
+        else:
+            print("\nGot unexpected response when authorizing CLI")
+            sys.exit(1)
+
+    config = CloudConfig(
+        id_token=data["id_token"],
+        refresh_token=data["refresh_token"],
+        refresh_token_expires=data["refresh_token_expires"],
+        region=region,
+    )
+    write_cloud_config(config)
+
+
+class ApiSession(requests.Session):
+    def __init__(self) -> None:
+        super().__init__()
+
+        if options.non_interactive:
+            username = os.getenv("LOCUSTCLOUD_USERNAME")
+            password = os.getenv("LOCUSTCLOUD_PASSWORD")
+            region = os.getenv("LOCUSTCLOUD_REGION")
+
+            if not all([username, password, region]):
+                print(
+                    "Running with --non-interaction requires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set."
+                )
+                sys.exit(1)
+
+            if region not in valid_regions:
+                print("Environment variable LOCUSTCLOUD_REGION needs to be set to one of", ", ".join(valid_regions))
+                sys.exit(1)
+
+            self.__configure_for_region(region)
+            response = requests.post(
+                self.__login_url,
+                json={"username": username, "password": password},
+                headers={"X-Client-Version": __version__},
+            )
+            if not response.ok:
+                print(f"Authentication failed: {response.text}")
+                sys.exit(1)
+
+            self.__refresh_token = response.json()["refresh_token"]
+            id_token = response.json()["cognito_client_id_token"]
+
+        else:
+            config = read_cloud_config()
+
+            if config.refresh_token_expires < time.time() + 24 * 60 * 60:
+                message = "You need to authenticate before proceeding. Please run:\n    locust-cloud --login"
+                print(message)
+                sys.exit(1)
+
+            assert config.region
+            self.__configure_for_region(config.region)
+            self.__refresh_token = config.refresh_token
+            id_token = config.id_token
+
+        assert id_token
+
+        decoded = jwt.decode(id_token, options={"verify_signature": False})
+        self.__expiry_time = decoded["exp"] - 60  # Refresh 1 minute before expiry
+        self.headers["Authorization"] = f"Bearer {id_token}"
+
+        self.__sub = decoded["sub"]
+        self.headers["X-Client-Version"] = __version__
+
+    def __configure_for_region(self, region: str) -> None:
+        self.__region = region
+        self.api_url = get_api_url(region)
+        self.__login_url = f"{self.api_url}/auth/login"
+
+        logger.debug(f"Lambda url: {self.api_url}")
+
+    def __ensure_valid_authorization_header(self) -> None:
+        if self.__expiry_time > time.time():
+            return
+
+        logger.info(f"Authenticating ({self.__region}, v{__version__})")
+
+        response = requests.post(
+            self.__login_url,
+            json={"user_sub_id": self.__sub, "refresh_token": self.__refresh_token},
+            headers={"X-Client-Version": __version__},
+        )
+
+        if not response.ok:
+            logger.error(f"Authentication failed: {response.text}")
+            sys.exit(1)
+
+        # TODO: Technically the /login endpoint can return a challenge for you
+        #       to change your password. Don't know how we should handle that
+        #       in the cli.
+
+        id_token = response.json()["cognito_client_id_token"]
+        decoded = jwt.decode(id_token, options={"verify_signature": False})
+        self.__expiry_time = decoded["exp"] - 60  # Refresh 1 minute before expiry
+        self.headers["Authorization"] = f"Bearer {id_token}"
+
+        if not options.non_interactive:
+            config = read_cloud_config()
+            config.id_token = id_token
+            write_cloud_config(config)
+
+    def request(self, method, url, *args, **kwargs) -> requests.Response:
+        self.__ensure_valid_authorization_header()
+        return super().request(method, f"{self.api_url}{url}", *args, **kwargs)
+
+
+def main() -> None:
+    if options.version:
+        print(f"locust-cloud version {__version__}")
+        sys.exit(0)
     if not options.locustfile:
         logger.error("A locustfile is required to run a test.")
         sys.exit(1)
 
-    try:
-        logger.info(f"Authenticating ({options.region}, v{__version__})")
-        logger.debug(f"Lambda url: {api_url}")
-        credential_manager = CredentialManager(
-            lambda_url=api_url,
-            access_key=options.aws_access_key_id,
-            secret_key=options.aws_secret_access_key,
-            username=options.username,
-            password=options.password,
-        )
+    if options.login:
+        try:
+            web_login()
+        except KeyboardInterrupt:
+            pass
+        sys.exit()
 
-        credentials = credential_manager.get_current_credentials()
-        cognito_client_id_token = credentials["cognito_client_id_token"]
-        aws_access_key_id = credentials.get("access_key")
-        aws_secret_access_key = credentials.get("secret_key")
-        aws_session_token = credentials.get("token", "")
+    session = ApiSession()
 
+    try:
         if options.delete:
-            delete(credential_manager)
+            delete(session)
             return
 
         try:
@@ -249,47 +409,41 @@ def main() -> None:
 
         logger.info("Deploying load generators")
         locust_env_variables = [
-            {"name": env_variable, "value": str(os.environ[env_variable])}
+            {"name": env_variable, "value": os.environ[env_variable]}
             for env_variable in os.environ
             if env_variable.startswith("LOCUST_")
-            and not env_variable
-            in [
+            and env_variable
+            not in [
                 "LOCUST_LOCUSTFILE",
                 "LOCUST_USERS",
                 "LOCUST_WEB_HOST_DISPLAY_NAME",
                 "LOCUST_SKIP_MONKEY_PATCH",
             ]
-            and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{api_url}/deploy"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_USERS", "value": str(options.users)},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
-                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": api_url},
+                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": session.api_url},
                 {"name": "LOCUSTCLOUD_PROFILE", "value": options.profile},
                 *locust_env_variables,
             ],
             "locustfile": {"filename": options.locustfile, "data": locustfile_data},
             "user_count": options.users,
-            "image_tag": options.image_tag,
             "mock_server": options.mock_server,
         }
+
+        if options.image_tag is not None:
+            payload["image_tag"] = options.image_tag
+
         if options.workers is not None:
             payload["worker_count"] = options.workers
+
         if options.requirements:
             payload["requirements"] = {"filename": options.requirements, "data": requirements_data}
-        headers = {
-            "Authorization": f"Bearer {cognito_client_id_token}",
-            "Content-Type": "application/json",
-            "AWS_ACCESS_KEY_ID": aws_access_key_id,
-            "AWS_SECRET_ACCESS_KEY": aws_secret_access_key,
-            "AWS_SESSION_TOKEN": aws_session_token,
-            "X-Client-Version": __version__,
-        }
+
         try:
-            # logger.info(payload) # might be useful when debugging sometimes
-            response = requests.post(deploy_endpoint, json=payload, headers=headers)
+            response = session.post("/deploy", json=payload)
         except requests.exceptions.RequestException as e:
             logger.error(f"Failed to deploy the load generators: {e}")
             sys.exit(1)
@@ -304,32 +458,28 @@ def main() -> None:
                     f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
                 )
             sys.exit(1)
-    except CredentialError as ce:
-        logger.error(f"Credential error: {ce}")
-        sys.exit(1)
+
     except KeyboardInterrupt:
+        # TODO: This would potentially leave a deployment running, combine with try-catch below?
         logger.debug("Interrupted by user")
         sys.exit(0)
 
     logger.debug("Load generators deployed successfully!")
     logger.info("Waiting for pods to be ready...")
 
+    shutdown_allowed = threading.Event()
+    shutdown_allowed.set()
+    reconnect_aborted = threading.Event()
+    connect_timeout = threading.Timer(2 * 60, reconnect_aborted.set)
+    sio = socketio.Client(handle_sigint=False)
+
     try:
         ws_connection_info = urllib.parse.urlparse(log_ws_url)
-        sio = socketio.Client(handle_sigint=False)
-
-        run = True
-
-        def wait():
-            logger.debug("Waiting for shutdown event")
-            while run:
-                time.sleep(0.1)
-
-            logger.debug("Shutting down websocket connection")
-            sio.shutdown()
 
         @sio.event
         def connect():
+            shutdown_allowed.clear()
+            connect_timeout.cancel()
             logger.debug("Websocket connection established, switching to Locust logs")
 
         @sio.event
@@ -350,55 +500,40 @@ def main() -> None:
             if message:
                 print(message)
 
-            nonlocal run
-            run = False
-
-        for _ in range(5 * 60):  # try for 5 minutes
-            try:
-                sio.connect(
-                    f"{ws_connection_info.scheme}://{ws_connection_info.netloc}", socketio_path=ws_connection_info.path
-                )
-                break
-            except socketio.exceptions.ConnectionError:
-                time.sleep(1)
-
-        else:  # no break
-            raise Exception("Failed to obtain socket connection")
-
-        wait()
+            shutdown_allowed.set()
+
+        # The _reconnect_abort value on the socketio client will be populated with a newly created threading.Event if it's not already set.
+        # There is no way to set this by passing it in the constructor.
+        # This event is the only way to interupt the retry logic when the connection is attempted.
+        sio._reconnect_abort = reconnect_aborted
+        connect_timeout.start()
+        sio.connect(
+            f"{ws_connection_info.scheme}://{ws_connection_info.netloc}",
+            socketio_path=ws_connection_info.path,
+            retry=True,
+        )
+        logger.debug("Waiting for shutdown")
+        shutdown_allowed.wait()
 
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
-        delete(credential_manager)
-        wait()
+        delete(session)
+        shutdown_allowed.wait(timeout=90)
     except Exception as e:
         logger.exception(e)
-        delete(credential_manager)
+        delete(session)
         sys.exit(1)
     else:
-        delete(credential_manager)
+        delete(session)
+    finally:
+        sio.shutdown()
 
 
-def delete(credential_manager):
+def delete(session):
     try:
         logger.info("Tearing down Locust cloud...")
-        credential_manager.refresh_credentials()
-        refreshed_credentials = credential_manager.get_current_credentials()
-
-        headers = {
-            "AWS_ACCESS_KEY_ID": refreshed_credentials.get("access_key", ""),
-            "AWS_SECRET_ACCESS_KEY": refreshed_credentials.get("secret_key", ""),
-            "Authorization": f"Bearer {refreshed_credentials.get('cognito_client_id_token', '')}",
-            "X-Client-Version": __version__,
-        }
-
-        token = refreshed_credentials.get("token")
-        if token:
-            headers["AWS_SESSION_TOKEN"] = token
-
-        response = requests.delete(
-            f"{api_url}/teardown",
-            headers=headers,
+        response = session.delete(
+            "/teardown",
         )
 
         if response.status_code == 200:
@@ -410,7 +545,7 @@ def delete(credential_manager):
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
-    logger.info("Done! ✨")
+    logger.info("Done! ✨")  # FIXME: Should probably not say it's done since at this point it could still be running
 
 
 if __name__ == "__main__":

{locust_cloud-1.12.3.dist-info → locust_cloud-1.13.0.dist-info}/METADATA

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: locust-cloud
-Version: 1.12.3
+Version: 1.13.0
 Summary: Locust Cloud
 Project-URL: Homepage, https://locust.cloud
 Requires-Python: >=3.11
 Requires-Dist: boto3==1.34.125
-Requires-Dist: gevent-websocket==0.10.1
-Requires-Dist: locust>=2.32.5.dev10
-Requires-Dist: psycopg[binary,pool]>=3.2.1
+Requires-Dist: configargparse==1.7
+Requires-Dist: platformdirs>=4.3.6
 Requires-Dist: pyjwt>=2.0
 Requires-Dist: python-socketio[client]==5.11.4
 Description-Content-Type: text/markdown

locust_cloud-1.13.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+locust_cloud/cloud.py,sha256=r0lWoP1QiT1Klkrecsr6LyPD_l1GMW2BFGRCRVkLowM,18124
+locust_cloud-1.13.0.dist-info/METADATA,sha256=STWxyZOw0qE003r-Pf-8RU3KPFK3yjH_NsSo5ffoHh0,496
+locust_cloud-1.13.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust_cloud-1.13.0.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.13.0.dist-info/RECORD,,