locust-cloud 1.5.3__py3-none-any.whl → 1.5.4__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.runner.state = "STOPPED"
+
+    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,71 +1,66 @@
 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 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):
+    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:
@@ -74,7 +69,7 @@ def register_auth(environment):
             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/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

@@ -255,6 +255,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 +280,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