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

locust_cloud/__init__.py

@@ -3,8 +3,10 @@ import os
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
 
 import argparse
+import logging
 import sys
 
+import locust.env
 import psycopg
 from locust import events
 from locust.argument_parser import LocustArgumentParser
@@ -19,6 +21,7 @@ PG_PASSWORD = os.environ.get("PG_PASSWORD")
 PG_DATABASE = os.environ.get("PG_DATABASE")
 PG_PORT = os.environ.get("PG_PORT", 5432)
 GRAPH_VIEWER = os.environ.get("GRAPH_VIEWER")
+logger = logging.getLogger(__name__)
 
 
 @events.init_command_line_parser.add_listener
@@ -55,7 +58,9 @@ def set_autocommit(conn: psycopg.Connection):
     conn.autocommit = True
 
 
-def create_connection_pool(pg_user, pg_host, pg_password, pg_database, pg_port):
+def create_connection_pool(
+    pg_user: str, pg_host: str, pg_password: str, pg_database: str, pg_port: str | int
+) -> ConnectionPool:
     try:
         return ConnectionPool(
             conninfo=f"postgres://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_database}?sslmode=require",
@@ -69,19 +74,28 @@ def create_connection_pool(pg_user, pg_host, pg_password, pg_database, pg_port):
 
 
 @events.init.add_listener
-def on_locust_init(environment, **_args):
-    if not (PG_HOST or GRAPH_VIEWER):
+def on_locust_init(environment: locust.env.Environment, **_args):
+    if not (PG_HOST and PG_USER and PG_PASSWORD and PG_DATABASE and PG_PORT):
         return
 
-    pool = create_connection_pool(
-        pg_user=PG_USER,
-        pg_host=PG_HOST,
-        pg_password=PG_PASSWORD,
-        pg_database=PG_DATABASE,
-        pg_port=PG_PORT,
-    )
+    if GRAPH_VIEWER:
+        environment.web_ui.template_args["isGraphViewer"] = True
+
+    try:
+        pool = create_connection_pool(
+            pg_user=PG_USER,
+            pg_host=PG_HOST,
+            pg_password=PG_PASSWORD,
+            pg_database=PG_DATABASE,
+            pg_port=PG_PORT,
+        )
+        pool.wait()
+    except Exception as e:
+        logger.exception(e)
+        logger.error(f"{PG_HOST=}")
+        raise
 
-    if not GRAPH_VIEWER and environment.parsed_options.exporter:
+    if not GRAPH_VIEWER and environment.parsed_options and environment.parsed_options.exporter:
         Exporter(environment, pool)
 
     if environment.web_ui:

locust_cloud/auth.py

@@ -1,80 +1,75 @@
 import os
 from datetime import UTC, datetime, timedelta
+from typing import TypedDict
 
+import locust.env
 import requests
+import werkzeug
 from flask import redirect, request, url_for
-from flask_login import UserMixin, current_user, login_user
-from locust_cloud.constants import DEFAULT_LAMBDA_URL
+from flask_login import UserMixin, login_user
 
-LAMBDA = os.environ.get("LOCUST_API_BASE_URL", DEFAULT_LAMBDA_URL)
+
+class Credentials(TypedDict):
+    user_sub_id: str
+    refresh_token: str
 
 
 class AuthUser(UserMixin):
-    def __init__(self, user_sub_id):
+    def __init__(self, user_sub_id: str):
         self.user_sub_id = user_sub_id
 
     def get_id(self):
         return self.user_sub_id
 
 
-def set_credentials(credentials, response):
-    if not credentials.get("cognito_client_id_token"):
+def set_credentials(username: str, credentials: Credentials, response: werkzeug.wrappers.response.Response):
+    if not credentials.get("user_sub_id"):
         return response
 
-    id_token = credentials["cognito_client_id_token"]
     user_sub_id = credentials["user_sub_id"]
     refresh_token = credentials["refresh_token"]
 
-    response.set_cookie("cognito_token", id_token, expires=datetime.now(tz=UTC) + timedelta(days=1))
+    response.set_cookie("username", username, expires=datetime.now(tz=UTC) + timedelta(days=365))
     response.set_cookie("user_token", refresh_token, expires=datetime.now(tz=UTC) + timedelta(days=365))
     response.set_cookie("user_sub_id", user_sub_id, expires=datetime.now(tz=UTC) + timedelta(days=365))
 
     return response
 
 
-def load_user(user_sub_id):
-    refresh_token = request.cookies.get("user_token")
+def register_auth(environment: locust.env.Environment):
+    environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
+    environment.web_ui.app.debug = False
 
-    if refresh_token:
-        return AuthUser(user_sub_id)
+    def load_user(user_sub_id: str):
+        username = request.cookies.get("username")
+        refresh_token = request.cookies.get("user_token")
 
-    return None
+        if refresh_token:
+            environment.web_ui.template_args["username"] = username
+            return AuthUser(user_sub_id)
 
+        return None
 
-def register_auth(environment):
-    environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
-    environment.web_ui.app.debug = False
     environment.web_ui.login_manager.user_loader(load_user)
     environment.web_ui.auth_args = {
         "username_password_callback": "/authenticate",
     }
 
-    @environment.web_ui.app.after_request
-    def refresh_handler(response):
-        if request.path == "/" and current_user:
-            refresh_token = request.cookies.get("user_token")
-            user_sub_id = request.cookies.get("user_sub_id")
-            if user_sub_id and refresh_token:
-                auth_response = requests.post(
-                    f"{LAMBDA}/auth/login", json={"user_sub_id": user_sub_id, "refresh_token": refresh_token}
-                )
-                credentials = auth_response.json()
-                response = set_credentials(credentials, response)
-
-        return response
-
     @environment.web_ui.app.route("/authenticate", methods=["POST"])
     def login_submit():
-        username = request.form.get("username")
+        username = request.form.get("username", "")
         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()
                 response = redirect(url_for("index"))
-                response = set_credentials(credentials, response)
+                response = set_credentials(username, credentials, response)
                 login_user(AuthUser(credentials["user_sub_id"]))
 
                 return response

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/exporter.py

@@ -235,7 +235,7 @@ class Exporter:
                         """
 UPDATE testruns
 SET (requests, resp_time_avg, rps_avg, fail_ratio) =
-(SELECT reqs, resp_time, reqs / GREATEST(duration, 1), fails / reqs) FROM
+(SELECT reqs, resp_time, reqs / GREATEST(duration, 1), fails / GREATEST(reqs, 1)) FROM
 (SELECT
  COUNT(*)::numeric AS reqs,
  AVG(response_time)::numeric as resp_time
@@ -248,7 +248,7 @@ SET (requests, resp_time_avg, rps_avg, fail_ratio) =
 WHERE id = %(run_id)s""",
                         {"run_id": self._run_id},
                     )
-                except psycopg.errors.DivisionByZero:
+                except psycopg.errors.DivisionByZero:  # remove this except later, because it shouldnt happen any more
                     logging.info(
                         "Got DivisionByZero error when trying to update testruns, most likely because there were no requests logged"
                     )

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
@@ -255,6 +256,13 @@ JOIN avg_response_time_failed f ON a.time = f.time
 ORDER BY a.time
 """
 
+total_runtime = """
+SELECT
+  SUM((end_time - id) * num_users) AS "totalVuh"
+FROM testruns
+WHERE id >= date_trunc('month', NOW())
+"""
+
 queries: dict["str", LiteralString] = {
     "request-names": request_names,
     "requests": requests_query,
@@ -273,4 +281,5 @@ queries: dict["str", LiteralString] = {
     "testruns-table": testruns_table,
     "testruns-rps": testruns_rps,
     "testruns-response-time": testruns_response_time,
+    "total-runtime": total_runtime,
 }

locust_cloud/timescale/query.py

@@ -1,6 +1,7 @@
 import logging
 
 from flask import make_response, request
+from flask_login import login_required
 from locust_cloud.timescale.queries import queries
 
 logger = logging.getLogger(__name__)
@@ -12,6 +13,7 @@ def adapt_timestamp(result):
 
 def register_query(environment, pool):
     @environment.web_ui.app.route("/cloud-stats/<query>", methods=["POST"])
+    @login_required
     def query(query):
         results = []
         try:
@@ -21,6 +23,18 @@ def register_query(environment, pool):
                     # get_conn_time = (time.perf_counter() - start_time) * 1000
                     sql_params = request.get_json()
                     # start_time = time.perf_counter()
+                    from datetime import datetime, timedelta
+
+                    if "start" in sql_params:
+                        # protect the database against huge queries
+                        start_time = datetime.fromisoformat(sql_params["start"])
+                        end_time = datetime.fromisoformat(sql_params["end"])
+                        if end_time >= start_time + timedelta(hours=6):
+                            logger.warning(
+                                f"UI asked for too long time interval. Start was {sql_params['start']}, end was {sql_params['end']}"
+                            )
+                            return []
+
                     cursor = conn.execute(queries[query], sql_params)
                     # exec_time = (time.perf_counter() - start_time) * 1000
                     assert cursor
@@ -45,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)