locust-cloud 1.5.4__py3-none-any.whl → 1.5.5__py3-none-any.whl

locust_cloud/__init__.py

@@ -79,7 +79,7 @@ def on_locust_init(environment: locust.env.Environment, **_args):
         return
 
     if GRAPH_VIEWER:
-        environment.runner.state = "STOPPED"
+        environment.web_ui.template_args["isGraphViewer"] = True
 
     try:
         pool = create_connection_pool(

locust_cloud/auth.py

@@ -7,9 +7,6 @@ import requests
 import werkzeug
 from flask import redirect, request, url_for
 from flask_login import UserMixin, login_user
-from locust_cloud.constants import DEFAULT_LAMBDA_URL
-
-LAMBDA = os.environ.get("LOCUST_API_BASE_URL", DEFAULT_LAMBDA_URL)
 
 
 class Credentials(TypedDict):
@@ -64,7 +61,10 @@ def register_auth(environment: locust.env.Environment):
         password = request.form.get("password")
 
         try:
-            auth_response = requests.post(f"{LAMBDA}/auth/login", json={"username": username, "password": password})
+            auth_response = requests.post(
+                f"{environment.parsed_options.deployer_url}/auth/login",
+                json={"username": username, "password": password},
+            )
 
             if auth_response.status_code == 200:
                 credentials = auth_response.json()

locust_cloud/cloud.py

@@ -16,14 +16,13 @@ import requests
 from botocore.exceptions import ClientError
 from locust_cloud.constants import (
     DEFAULT_CLUSTER_NAME,
-    DEFAULT_LAMBDA_URL,
+    DEFAULT_DEPLOYER_URL,
     DEFAULT_NAMESPACE,
     DEFAULT_REGION_NAME,
     USERS_PER_WORKER,
 )
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
-LOCUST_ENV_VARIABLE_IGNORE_LIST = ["LOCUST_BUILD_PATH", "LOCUST_SKIP_MONKEY_PATCH"]
 __version__ = importlib.metadata.version("locust-cloud")
 
 
@@ -57,7 +56,7 @@ parser = configargparse.ArgumentParser(
         "~/.cloud.conf",
         "cloud.conf",
     ],
-    auto_env_var_prefix="LOCUST_",
+    auto_env_var_prefix="LOCUSTCLOUD_",
     formatter_class=configargparse.RawDescriptionHelpFormatter,
     config_file_parser_class=configargparse.CompositeConfigParser(
         [
@@ -67,7 +66,7 @@ parser = configargparse.ArgumentParser(
     ),
     description="""Launches distributed Locust runs on locust.cloud infrastructure.
 
-Example: locust-cloud -f my_locustfile.py --aws-region-name us-east-1 --users 1000""",
+Example: locust-cloud -f my_locustfile.py --region us-east-1 --users 2000""",
     epilog="""Any parameters not listed here are forwarded to locust master unmodified, so go ahead and use things like --users, --host, --run-time, ...
 Locust config can also be set using config file (~/.locust.conf, locust.conf, pyproject.toml, ~/.cloud.conf or cloud.conf).
 Parameters specified on command line override env vars, which in turn override config files.""",
@@ -88,21 +87,19 @@ parser.add_argument(
     "--users",
     type=int,
     default=1,
-    help="Number of users to launch. This is the same as the regular Locust argument, but also decides how many workers to launch by default.",
+    help="Number of users to launch. This is the same as the regular Locust argument, but also affects how many workers to launch.",
     env_var="LOCUST_USERS",
 )
 parser.add_argument(
     "--requirements",
     type=str,
     help="Optional requirements.txt file that contains your external libraries.",
-    env_var="LOCUST_REQUIREMENTS",
 )
 parser.add_argument(
-    "--aws-region-name",
+    "--region",
     type=str,
-    default=DEFAULT_REGION_NAME,
-    help="Sets the region to use for the deployed cluster",
-    env_var="AWS_REGION_NAME",
+    default=os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION_NAME),
+    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(
     "--kube-cluster-name",
@@ -119,11 +116,10 @@ parser.add_argument(
     env_var="KUBE_NAMESPACE",
 )
 parser.add_argument(
-    "--lambda-url",
+    "--deployer-url",
     type=str,
-    default=DEFAULT_LAMBDA_URL,
-    help="Sets the namespace for scoping the deployed cluster",
-    env_var="LOCUST_API_BASE_URL",
+    default=DEFAULT_DEPLOYER_URL,
+    help=configargparse.SUPPRESS,
 )
 parser.add_argument(
     "--aws-access-key-id",
@@ -143,29 +139,25 @@ parser.add_argument(
     "--username",
     type=str,
     help="Authentication for deploying with Locust Cloud",
-    env_var="LOCUST_CLOUD_USERNAME",
-    default=None,
+    default=os.getenv("LOCUST_CLOUD_USERNAME", None),  # backwards compatitibility for dmdb
 )
 parser.add_argument(
     "--password",
     type=str,
     help="Authentication for deploying with Locust Cloud",
-    env_var="LOCUST_CLOUD_PASSWORD",
-    default=None,
+    default=os.getenv("LOCUST_CLOUD_PASSWORD", None),  # backwards compatitibility for dmdb
 )
 parser.add_argument(
     "--loglevel",
     "-L",
     type=str,
     help="Log level",
-    env_var="LOCUST_CLOUD_LOGLEVEL",
     default="INFO",
 )
 parser.add_argument(
     "--workers",
     type=int,
-    help="Number of workers to use for the deployment",
-    env_var="LOCUST_CLOUD_WORKERS",
+    help=f"Number of workers to use for the deployment. Defaults to number of users divided by {USERS_PER_WORKER}",
     default=None,
 )
 parser.add_argument(
@@ -176,7 +168,6 @@ parser.add_argument(
 parser.add_argument(
     "--image-tag",
     type=str,
-    env_var="LOCUST_CLOUD_IMAGE_TAG",
     default="latest",
     help=configargparse.SUPPRESS,  # overrides the locust-cloud docker image tag. for internal use
 )
@@ -202,6 +193,7 @@ def main() -> None:
     s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
     deployed_pods: list[Any] = []
     worker_count: int = max(options.workers or math.ceil(options.users / USERS_PER_WORKER), 2)
+    os.environ["AWS_DEFAULT_REGION"] = options.region
     if options.users > 10000:
         logger.error("You asked for more than 10000 Users, that isn't allowed.")
         sys.exit(1)
@@ -217,15 +209,14 @@ def main() -> None:
             )
             sys.exit(1)
 
-        logger.info(f"Authenticating ({options.aws_region_name}, v{__version__})")
-        logger.debug(f"Lambda url: {options.lambda_url}")
+        logger.info(f"Authenticating ({os.environ['AWS_DEFAULT_REGION']}, v{__version__})")
+        logger.debug(f"Lambda url: {options.deployer_url}")
         credential_manager = CredentialManager(
-            lambda_url=options.lambda_url,
+            lambda_url=options.deployer_url,
             access_key=options.aws_access_key_id,
             secret_key=options.aws_secret_access_key,
             username=options.username,
             password=options.password,
-            region_name=options.aws_region_name,
         )
 
         credentials = credential_manager.get_current_credentials()
@@ -278,26 +269,23 @@ def main() -> None:
             {"name": env_variable, "value": str(os.environ[env_variable])}
             for env_variable in os.environ
             if env_variable.startswith("LOCUST_")
-            and not env_variable.startswith("LOCUST_CLOUD")
             and not env_variable
             in [
                 "LOCUST_LOCUSTFILE",
                 "LOCUST_USERS",
                 "LOCUST_WEB_HOST_DISPLAY_NAME",
-                "LOCUST_API_BASE_URL",
                 "LOCUST_SKIP_MONKEY_PATCH",
-                "LOCUST_REQUIREMENTS_URL",
             ]
             and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{options.lambda_url}/{options.kube_cluster_name}"
+        deploy_endpoint = f"{options.deployer_url}/{options.kube_cluster_name}"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_LOCUSTFILE", "value": locustfile_url},
                 {"name": "LOCUST_USERS", "value": str(options.users)},
-                {"name": "LOCUST_REQUIREMENTS_URL", "value": requirements_url},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
-                {"name": "LOCUST_API_BASE_URL", "value": DEFAULT_LAMBDA_URL},
+                {"name": "LOCUSTCLOUD_REQUIREMENTS_URL", "value": requirements_url},
+                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": options.deployer_url},
                 *locust_env_variables,
             ],
             "worker_count": worker_count,
@@ -416,7 +404,7 @@ def delete(s3_bucket, credential_manager):
             headers["AWS_SESSION_TOKEN"] = token
 
         response = requests.delete(
-            f"{options.lambda_url}/{options.kube_cluster_name}",
+            f"{options.deployer_url}/{options.kube_cluster_name}",
             headers=headers,
             params={"namespace": options.kube_namespace} if options.kube_namespace else {},
         )
@@ -426,7 +414,7 @@ def delete(s3_bucket, credential_manager):
                 f"Could not automatically tear down Locust Cloud: HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
             )
     except Exception as e:
-        logger.error(f"Could not automatically tear down Locust Cloud: {e}")
+        logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
     try:
         logger.debug("Cleaning up locustfiles")

locust_cloud/constants.py

@@ -1,5 +1,5 @@
 DEFAULT_REGION_NAME = "us-east-1"
 DEFAULT_CLUSTER_NAME = "dmdb"
 DEFAULT_NAMESPACE = "default"
-DEFAULT_LAMBDA_URL = "https://api.locust.cloud/1"
+DEFAULT_DEPLOYER_URL = "https://api.locust.cloud/1"
 USERS_PER_WORKER = 500

locust_cloud/credential_manager.py

@@ -8,7 +8,6 @@ import jwt
 import requests
 from botocore.credentials import RefreshableCredentials
 from botocore.session import Session as BotocoreSession
-from locust_cloud.constants import DEFAULT_REGION_NAME
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +26,6 @@ class CredentialManager:
         password: str | None = None,
         user_sub_id: str | None = None,
         refresh_token: str | None = None,
-        region_name: str = DEFAULT_REGION_NAME,
         access_key: str | None = None,
         secret_key: str | None = None,
     ) -> None:
@@ -36,7 +34,6 @@ class CredentialManager:
         self.password = password
         self.user_sub_id = user_sub_id
         self.refresh_token = refresh_token
-        self.region_name = region_name
 
         self.credentials = {
             "access_key": access_key,
@@ -56,12 +53,8 @@ class CredentialManager:
         botocore_session = BotocoreSession()
         botocore_session._credentials = self.refreshable_credentials  # type: ignore
         botocore_session.set_config_variable("signature_version", "v4")
-        botocore_session.set_config_variable("region", self.region_name)
 
-        self.session = boto3.Session(
-            botocore_session=botocore_session,
-            region_name=self.region_name,
-        )
+        self.session = boto3.Session(botocore_session=botocore_session)
         logger.debug("Boto3 session created with RefreshableCredentials.")
 
     def obtain_credentials(self) -> None:

locust_cloud/timescale/queries.py

@@ -33,7 +33,7 @@ GROUP BY "name",left(exception,300)
 requests_per_second = """
 WITH request_count_agg AS (
   SELECT
-    time_bucket_gapfill('%(resolution)ss', bucket) AS time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     COALESCE(SUM(count)/%(resolution)s, 0) as rps
   FROM requests_summary
   WHERE bucket BETWEEN %(start)s AND %(end)s
@@ -43,7 +43,7 @@ WITH request_count_agg AS (
 ),
 user_count_agg AS (
   SELECT
-    time_bucket_gapfill('%(resolution)ss', time) AS time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', time) AS time,
     COALESCE(avg(user_count), 0) as users
   FROM number_of_users
   WHERE time BETWEEN %(start)s AND %(end)s
@@ -53,7 +53,7 @@ user_count_agg AS (
 ),
 errors_per_s_agg AS (
   SELECT
-    time_bucket_gapfill('%(resolution)ss', bucket) AS time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     COALESCE(SUM(failed_count)/%(resolution)s, 0) as error_rate
   FROM requests_summary
   WHERE bucket BETWEEN %(start)s AND %(end)s
@@ -101,7 +101,7 @@ AND run_id = %(testrun)s
 
 rps_per_request = """
 SELECT
-    time_bucket_gapfill('%(resolution)ss', bucket) AS time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     name,
     COALESCE(SUM(count)/%(resolution)s, 0) as throughput
 FROM requests_summary
@@ -114,7 +114,7 @@ ORDER BY 1,2
 
 avg_response_times = """
 SELECT
-    time_bucket_gapfill('%(resolution)ss', bucket) as time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) as time,
     name,
     avg(average) as "responseTime"
 FROM requests_summary
@@ -126,7 +126,7 @@ ORDER BY 1, 2
 
 errors_per_request = """
 SELECT
-    time_bucket_gapfill('%(resolution)ss', bucket) AS time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     name,
     SUM(failed_count)/%(resolution)s as "errorRate"
 FROM requests_summary
@@ -138,19 +138,20 @@ ORDER BY 1
 
 
 perc99_response_times = """
-SELECT time_bucket('%(resolution)ss', bucket) AS time,
+SELECT time_bucket(%(resolution)s * interval '1 second', bucket) AS time,
   name,
   perc99
 FROM requests_summary
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name, perc99
+ORDER BY 1
 """
 
 
 response_length = """
 SELECT
-    time_bucket('%(resolution)ss', bucket) as time,
+    time_bucket(%(resolution)s * interval '1 second', bucket) as time,
     response_length as "responseLength",
     name
 FROM requests_summary

locust_cloud/timescale/query.py

@@ -59,6 +59,5 @@ def register_query(environment, pool):
                 logger.warning(f"Received invalid query key: '{query}'")
                 return make_response("Invalid query key", 401)
         except Exception as e:
-            print(e)
             logger.error(f"Error executing query '{query}': {e}", exc_info=True)
             return make_response("Error executing query", 401)