locust-cloud 1.12.4__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,10 +458,9 @@ 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)
 
@@ -318,10 +471,10 @@ def main() -> None:
     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)
 
         @sio.event
         def connect():
@@ -364,38 +517,23 @@ def main() -> None:
 
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
-        delete(credential_manager)
+        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:
@@ -407,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.4.dist-info → locust_cloud-1.13.0.dist-info}/METADATA

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: locust-cloud
-Version: 1.12.4
+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,,

locust_cloud/__init__.py

@@ -1,133 +0,0 @@
-import importlib.metadata
-import os
-import sys
-
-os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
-
-from locust_cloud.socket_logging import setup_socket_logging
-
-if os.environ.get("LOCUST_MODE_MASTER") == "1":
-    major, minor, *rest = os.environ["LOCUSTCLOUD_CLIENT_VERSION"].split(".")
-
-    if int(major) > 1 or int(major) == 1 and int(minor) >= 12:
-        setup_socket_logging()
-
-__version__ = importlib.metadata.version("locust-cloud")
-
-import logging
-
-import configargparse
-import locust.env
-import psycopg
-from locust import events
-from locust.argument_parser import LocustArgumentParser
-from locust_cloud.auth import register_auth
-from locust_cloud.idle_exit import IdleExit
-from locust_cloud.timescale.exporter import Exporter
-from locust_cloud.timescale.query import register_query
-from psycopg.conninfo import make_conninfo
-from psycopg_pool import ConnectionPool
-
-logger = logging.getLogger(__name__)
-
-
-@events.init_command_line_parser.add_listener
-def add_arguments(parser: LocustArgumentParser):
-    if not (os.environ.get("PGHOST")):
-        parser.add_argument_group(
-            "locust-cloud",
-            "locust-cloud disabled, because PGHOST was not set - this is normal for local runs",
-        )
-        return
-
-    try:
-        REGION = os.environ["AWS_DEFAULT_REGION"]
-    except KeyError:
-        logger.fatal("Missing AWS_DEFAULT_REGION env var")
-        sys.exit(1)
-
-    os.environ["LOCUST_BUILD_PATH"] = os.path.join(os.path.dirname(__file__), "webui/dist")
-    locust_cloud = parser.add_argument_group(
-        "locust-cloud",
-        "Arguments for use with Locust cloud",
-    )
-    # do not set
-    # used for sending the run id from master to workers
-    locust_cloud.add_argument(
-        "--run-id",
-        type=str,
-        env_var="LOCUSTCLOUD_RUN_ID",
-        help=configargparse.SUPPRESS,
-    )
-    locust_cloud.add_argument(
-        "--allow-signup",
-        env_var="LOCUSTCLOUD_ALLOW_SIGNUP",
-        help=configargparse.SUPPRESS,
-        default=False,
-        action="store_true",
-    )
-    locust_cloud.add_argument(
-        "--allow-forgot-password",
-        env_var="LOCUSTCLOUD_FORGOT_PASSWORD",
-        help=configargparse.SUPPRESS,
-        default=False,
-        action="store_true",
-    )
-    locust_cloud.add_argument(
-        "--graph-viewer",
-        env_var="LOCUSTCLOUD_GRAPH_VIEWER",
-        help=configargparse.SUPPRESS,
-        default=False,
-        action="store_true",
-    )
-    locust_cloud.add_argument(
-        "--deployer-url",
-        type=str,
-        env_var="LOCUSTCLOUD_DEPLOYER_URL",
-        help=configargparse.SUPPRESS,
-        default=f"https://api.{REGION}.locust.cloud/1",
-    )
-    locust_cloud.add_argument(
-        "--profile",
-        type=str,
-        env_var="LOCUSTCLOUD_PROFILE",
-        help=configargparse.SUPPRESS,
-        default=None,
-    )
-
-
-def set_autocommit(conn: psycopg.Connection):
-    conn.autocommit = True
-
-
-@events.init.add_listener
-def on_locust_init(environment: locust.env.Environment, **_args):
-    if not (os.environ.get("PGHOST")):
-        return
-
-    conninfo = make_conninfo(
-        sslmode="require",
-    )
-    pool = ConnectionPool(
-        conninfo,
-        min_size=1,
-        max_size=20,
-        configure=set_autocommit,
-        check=ConnectionPool.check_connection,
-    )
-    pool.wait(timeout=10)
-
-    if not environment.parsed_options.graph_viewer:
-        IdleExit(environment)
-        Exporter(environment, pool)
-
-    if environment.web_ui:
-        environment.web_ui.template_args["locustVersion"] = locust.__version__
-        environment.web_ui.template_args["locustCloudVersion"] = __version__
-        environment.web_ui.template_args["webBasePath"] = environment.parsed_options.web_base_path
-
-        if environment.parsed_options.graph_viewer:
-            environment.web_ui.template_args["isGraphViewer"] = True
-
-        register_auth(environment)
-        register_query(environment, pool)