locust-cloud 1.11.1__py3-none-any.whl → 1.11.2__py3-none-any.whl

locust_cloud/__init__.py

@@ -57,6 +57,13 @@ def add_arguments(parser: LocustArgumentParser):
         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",
@@ -71,6 +78,12 @@ def add_arguments(parser: LocustArgumentParser):
         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,
+    )
 
 
 def set_autocommit(conn: psycopg.Connection):
@@ -103,6 +116,9 @@ def on_locust_init(environment: locust.env.Environment, **_args):
         Exporter(environment, pool)
 
     if environment.web_ui:
+        environment.web_ui.template_args["locustVersion"] = locust.__version__
+        environment.web_ui.template_args["locustCloudVersion"] = __version__
+
         if environment.parsed_options.graph_viewer:
             environment.web_ui.template_args["isGraphViewer"] = True
 

locust_cloud/auth.py

@@ -66,10 +66,15 @@ def register_auth(environment: locust.env.Environment):
         },
     )
 
+    environment.web_ui.auth_args["auth_providers"] = []
     if environment.parsed_options.allow_signup:
-        environment.web_ui.auth_args["auth_providers"] = [
+        environment.web_ui.auth_args["auth_providers"].append(
             {"label": "Sign Up", "callback_url": f"{web_base_path}/signup"}
-        ]
+        )
+    if environment.parsed_options.allow_forgot_password:
+        environment.web_ui.auth_args["auth_providers"].append(
+            {"label": "Forgot Password?", "callback_url": f"{web_base_path}/forgot-password"}
+        )
 
     @auth_blueprint.route("/authenticate", methods=["POST"])
     def login_submit():
@@ -87,6 +92,11 @@ def register_auth(environment: locust.env.Environment):
 
             credentials = auth_response.json()
 
+            if credentials.get("challenge_session"):
+                session["challenge_session"] = credentials.get("challenge_session")
+                session["username"] = username
+                return redirect(url_for("locust_cloud_auth.password_reset"))
+
             if os.getenv("CUSTOMER_ID", "") and credentials["customer_id"] != os.getenv("CUSTOMER_ID", ""):
                 session["auth_error"] = "Invalid login for this deployment"
                 return redirect(url_for("locust.login"))
@@ -263,4 +273,137 @@ def register_auth(environment: locust.env.Environment):
 
             return redirect(url_for("locust_cloud_auth.signup"))
 
+    @auth_blueprint.route("/forgot-password")
+    def forgot_password():
+        if not environment.parsed_options.allow_forgot_password:
+            return redirect(url_for("locust.login"))
+
+        forgot_password_args = {
+            "custom_form": {
+                "inputs": [
+                    {
+                        "label": "Username",
+                        "name": "username",
+                        "is_required": True,
+                        "type": "email",
+                    },
+                ],
+                "callback_url": f"{web_base_path}/send-forgot-password",
+                "submit_button_text": "Reset Password",
+            },
+            "info": "Enter your email and we will send a code to reset your password",
+        }
+
+        if session.get("auth_error"):
+            forgot_password_args["error"] = session["auth_error"]
+
+        return render_template_from("auth.html", auth_args=forgot_password_args)
+
+    @auth_blueprint.route("/send-forgot-password", methods=["POST"])
+    def send_forgot_password():
+        if not environment.parsed_options.allow_forgot_password:
+            return redirect(url_for("locust.login"))
+
+        try:
+            username = request.form.get("username", "")
+
+            auth_response = requests.post(
+                f"{environment.parsed_options.deployer_url}/auth/forgot-password",
+                json={"username": username},
+            )
+
+            auth_response.raise_for_status()
+
+            session["username"] = username
+
+            return redirect(url_for("locust_cloud_auth.password_reset"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.forgot_password"))
+
+    @auth_blueprint.route("/password-reset")
+    def password_reset():
+        if not environment.parsed_options.allow_forgot_password and not session.get("challenge_session"):
+            return redirect(url_for("locust.login"))
+
+        if session.get("challenge_session"):
+            reset_password_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "New Password",
+                            "name": "new_password",
+                            "is_required": True,
+                            "is_secret": True,
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/confirm-reset-password",
+                    "submit_button_text": "Reset Password",
+                },
+                "info": "You must set a new password",
+            }
+        else:
+            reset_password_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "Confirmation Code",
+                            "name": "confirmation_code",
+                            "is_required": True,
+                        },
+                        {
+                            "label": "New Password",
+                            "name": "new_password",
+                            "is_required": True,
+                            "is_secret": True,
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/confirm-reset-password",
+                    "submit_button_text": "Reset Password",
+                },
+                "info": "Enter your the confirmation code that was sent to your email",
+            }
+
+        if session.get("auth_error"):
+            reset_password_args["error"] = session["auth_error"]
+
+        return render_template_from("auth.html", auth_args=reset_password_args)
+
+    @auth_blueprint.route("/confirm-reset-password", methods=["POST"])
+    def confirm_reset_password():
+        if not environment.parsed_options.allow_forgot_password and not session.get("challenge_session"):
+            return redirect(url_for("locust.login"))
+
+        try:
+            username = session["username"]
+            confirmation_code = request.form.get("confirmation_code")
+            new_password = request.form.get("new_password")
+
+            auth_response = requests.post(
+                f"{environment.parsed_options.deployer_url}/auth/password-reset",
+                json={
+                    "username": username,
+                    "confirmation_code": confirmation_code,
+                    "new_password": new_password,
+                    "challenge_session": session.get("challenge_session"),
+                },
+            )
+
+            auth_response.raise_for_status()
+
+            session["username"] = ""
+            session["auth_info"] = "Password reset successfully! Please login"
+            session["auth_error"] = ""
+
+            return redirect(url_for("locust.login"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.password_reset"))
+
     environment.web_ui.app.register_blueprint(auth_blueprint)

locust_cloud/cloud.py

@@ -1,3 +1,5 @@
+import base64
+import gzip
 import json
 import logging
 import os
@@ -70,6 +72,12 @@ parser.add_argument(
     action="help",
     help=configargparse.SUPPRESS,
 )
+parser.add_argument(
+    "-V",
+    "--version",
+    action="store_true",
+    help=configargparse.SUPPRESS,
+)
 parser.add_argument(
     "-f",
     "--locustfile",
@@ -152,7 +160,12 @@ parser.add_argument(
     "--mock-server",
     action="store_true",
     default=False,
-    help=configargparse.SUPPRESS,
+    help="Start a demo mock service and set --host parameter to point Locust towards it",
+)
+parser.add_argument(
+    "--profile",
+    type=str,
+    help="Set a profile to group the testruns together",
 )
 
 options, locust_options = parser.parse_known_args()
@@ -167,7 +180,6 @@ logger = logging.getLogger(__name__)
 # Restore log level for other libs. Yes, this can be done more nicely
 logging.getLogger("botocore").setLevel(logging.INFO)
 logging.getLogger("boto3").setLevel(logging.INFO)
-logging.getLogger("s3transfer").setLevel(logging.INFO)
 logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
@@ -176,6 +188,10 @@ api_url = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{options.regi
 
 
 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."
@@ -184,14 +200,16 @@ def main() -> None:
     if options.region:
         os.environ["AWS_DEFAULT_REGION"] = options.region
 
-    s3_bucket = "dmdb-default" if options.region == "us-east-1" else "locust-default"
     deployments: list[Any] = []
 
     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 LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
+            "Authentication is required to use Locust Cloud. Please ensure the LOCUSTCLOUD_USERNAME and LOCUSTCLOUD_PASSWORD environment variables are set."
         )
         sys.exit(1)
+    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__})")
@@ -214,43 +232,22 @@ def main() -> None:
             delete(credential_manager)
             return
 
-        logger.info(f"Uploading {options.locustfile}")
-        logger.debug(f"... to {s3_bucket}")
-        s3 = credential_manager.session.client("s3")
         try:
-            filename = options.username + "__" + os.path.basename(options.locustfile)
-            s3.upload_file(options.locustfile, s3_bucket, filename)
-            locustfile_url = s3.generate_presigned_url(
-                ClientMethod="get_object",
-                Params={"Bucket": s3_bucket, "Key": filename},
-                ExpiresIn=3600,
-            )
-            logger.debug(f"Uploaded {options.locustfile} successfully")
+            with open(options.locustfile, "rb") as f:
+                locustfile_data = base64.b64encode(gzip.compress(f.read())).decode()
         except FileNotFoundError:
             logger.error(f"File not found: {options.locustfile}")
             sys.exit(1)
-        except ClientError as e:
-            logger.error(f"Failed to upload {options.locustfile}: {e}")
-            sys.exit(1)
 
-        requirements_url = ""
+        requirements_data = None
+
         if options.requirements:
-            logger.info(f"Uploading {options.requirements}")
             try:
-                filename = options.username + "__" + "requirements.txt"
-                s3.upload_file(options.requirements, s3_bucket, filename)
-                requirements_url = s3.generate_presigned_url(
-                    ClientMethod="get_object",
-                    Params={"Bucket": s3_bucket, "Key": filename},
-                    ExpiresIn=3600,
-                )
-                logger.debug(f"Uploaded {options.requirements} successfully")
+                with open(options.requirements, "rb") as f:
+                    requirements_data = base64.b64encode(gzip.compress(f.read())).decode()
             except FileNotFoundError:
                 logger.error(f"File not found: {options.requirements}")
                 sys.exit(1)
-            except ClientError as e:
-                logger.error(f"Failed to upload {options.requirements}: {e}")
-                sys.exit(1)
 
         logger.info("Deploying load generators")
         locust_env_variables = [
@@ -269,19 +266,21 @@ def main() -> None:
         deploy_endpoint = f"{api_url}/deploy"
         payload = {
             "locust_args": [
-                {"name": "LOCUST_LOCUSTFILE", "value": locustfile_url},
                 {"name": "LOCUST_USERS", "value": str(options.users)},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
-                {"name": "LOCUSTCLOUD_REQUIREMENTS_URL", "value": requirements_url},
                 {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": 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.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",
@@ -301,7 +300,7 @@ def main() -> None:
             deployments = response.json().get("deployments", [])
         else:
             try:
-                logger.error(f"Error when deploying: {response.json()['Message']}")
+                logger.error(f"{response.json()['Message']} (HTTP {response.status_code}/{response.reason})")
             except Exception:
                 logger.error(
                     f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
@@ -421,15 +420,6 @@ def delete(credential_manager):
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
-    try:
-        logger.debug("Cleaning up locustfiles")
-        # s3 = credential_manager.session.resource("s3")
-        # bucket = s3.Bucket(s3_bucket)
-        # bucket.objects.all().delete()
-    except ClientError as e:
-        logger.debug(f"Failed to clean up locust files: {e}")
-        # sys.exit(1)
-
     logger.info("Done! ✨")
 
 

locust_cloud/timescale/exporter.py

@@ -217,7 +217,7 @@ class Exporter:
         cmd = sys.argv[1:]
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,current_user)",
+                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, profile, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,%s,current_user)",
                 (
                     self._run_id,
                     self.env.runner.target_user_count if self.env.runner else 1,
@@ -229,6 +229,7 @@ class Exporter:
                     else 0,
                     self.env.web_ui.template_args.get("username", "") if self.env.web_ui else "",
                     self.env.parsed_locustfiles[0].split("/")[-1].split("__")[-1],
+                    self.env.parsed_options.profile,
                     " ".join(cmd),
                 ),
             )

locust_cloud/timescale/queries.py

@@ -183,7 +183,9 @@ ORDER BY 1,2
 testruns = """
 SELECT
   id as "runId",
-  end_time as "endTime"
+  end_time as "endTime",
+  locustfile,
+  profile
 FROM testruns
 ORDER BY id DESC
 """
@@ -277,6 +279,15 @@ FROM customers
 WHERE id = current_user
 """
 
+profiles = """
+SELECT DISTINCT
+CASE
+    WHEN profile IS NOT NULL THEN profile
+    ELSE locustfile
+END AS profile
+FROM testruns
+"""
+
 queries: dict["str", LiteralString] = {
     "request-names": request_names,
     "requests": requests_query,
@@ -297,4 +308,5 @@ queries: dict["str", LiteralString] = {
     "testruns-response-time": testruns_response_time,
     "total-vuh": total_vuh,
     "customer": customer,
+    "profiles": profiles,
 }