locust-cloud 1.5.11__py3-none-any.whl → 1.7.0__py3-none-any.whl

locust_cloud/__init__.py

@@ -1,18 +1,22 @@
+import importlib.metadata
 import os
 
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
+__version__ = importlib.metadata.version("locust-cloud")
 
 import argparse
 import logging
-import sys
 
+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
 
 PG_USER = os.environ.get("PG_USER")
@@ -53,39 +57,41 @@ def add_arguments(parser: LocustArgumentParser):
         default="",
         help="Description of the test being run",
     )
+    # 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,
+    )
 
 
 def set_autocommit(conn: psycopg.Connection):
     conn.autocommit = True
 
 
-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",
-            min_size=1,
-            max_size=10,
-            configure=set_autocommit,
-        )
-    except Exception:
-        sys.stderr.write(f"Could not connect to postgres ({pg_user}@{pg_host}:{pg_port}).")
-        sys.exit(1)
-
-
 @events.init.add_listener
 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
 
     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,
+        conninfo = make_conninfo(
+            dbname=PG_DATABASE,
+            user=PG_USER,
+            port=PG_PORT,
+            password=PG_PASSWORD,
+            host=PG_HOST,
+            sslmode="require",
+            # options="-c statement_timeout=55000",
+        )
+        pool = ConnectionPool(
+            conninfo,
+            min_size=1,
+            max_size=10,
+            configure=set_autocommit,
+            check=ConnectionPool.check_connection,
         )
         pool.wait()
     except Exception as e:
@@ -93,6 +99,9 @@ def on_locust_init(environment: locust.env.Environment, **_args):
         logger.error(f"{PG_HOST=}")
         raise
 
+    if not GRAPH_VIEWER:
+        IdleExit(environment)
+
     if not GRAPH_VIEWER and environment.parsed_options and environment.parsed_options.exporter:
         Exporter(environment, pool)
 

locust_cloud/auth.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from datetime import UTC, datetime, timedelta
 from typing import TypedDict
@@ -7,12 +8,13 @@ import requests
 import werkzeug
 from flask import redirect, request, url_for
 from flask_login import UserMixin, login_user
-from locust_cloud.constants import (
-    DEFAULT_DEPLOYER_URL,
-)
+from locust_cloud import __version__
+from locust_cloud.constants import DEFAULT_DEPLOYER_URL
 
 DEPLOYER_URL = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", DEFAULT_DEPLOYER_URL)
 
+logger = logging.getLogger(__name__)
+
 
 class Credentials(TypedDict):
     user_sub_id: str
@@ -69,23 +71,25 @@ def register_auth(environment: locust.env.Environment):
             auth_response = requests.post(
                 f"{DEPLOYER_URL}/auth/login",
                 json={"username": username, "password": password},
+                headers={"X-Client-Version": __version__},
             )
 
-            if auth_response.status_code == 200:
-                credentials = auth_response.json()
-                response = redirect(url_for("index"))
-                response = set_credentials(username, credentials, response)
-                login_user(AuthUser(credentials["user_sub_id"]))
+            auth_response.raise_for_status()
 
-                return response
+            credentials = auth_response.json()
+            response = redirect(url_for("index"))
+            response = set_credentials(username, credentials, response)
+            login_user(AuthUser(credentials["user_sub_id"]))
 
-            environment.web_ui.auth_args = {**environment.web_ui.auth_args, "error": "Invalid username or password"}
+            return response
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 401:
+                environment.web_ui.auth_args["error"] = "Invalid username or password"
+            else:
+                logger.error(f"Unknown response from auth: {e.response.status_code} {e.response.text}")
 
-            return redirect(url_for("login"))
-        except Exception:
-            environment.web_ui.auth_args = {
-                **environment.web_ui.auth_args,
-                "error": "An unknown error occured, please try again",
-            }
+                environment.web_ui.auth_args["error"] = (
+                    "Unknown error during authentication, check logs and/or contact support"
+                )
 
             return redirect(url_for("login"))

locust_cloud/cloud.py

@@ -1,4 +1,3 @@
-import importlib.metadata
 import json
 import logging
 import math
@@ -14,6 +13,7 @@ from typing import IO, Any
 import configargparse
 import requests
 from botocore.exceptions import ClientError
+from locust_cloud import __version__
 from locust_cloud.constants import (
     DEFAULT_CLUSTER_NAME,
     DEFAULT_DEPLOYER_URL,
@@ -23,8 +23,6 @@ from locust_cloud.constants import (
 )
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
-__version__ = importlib.metadata.version("locust-cloud")
-
 
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
     def parse(self, stream: IO[str]) -> OrderedDict[str, Any]:
@@ -198,7 +196,7 @@ logging.getLogger("urllib3").setLevel(logging.INFO)
 
 def main() -> None:
     s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
-    deployed_pods: list[Any] = []
+    deployments: 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 > 5000000:
@@ -305,6 +303,7 @@ def main() -> None:
             "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
@@ -314,7 +313,7 @@ def main() -> None:
             sys.exit(1)
 
         if response.status_code == 200:
-            deployed_pods = response.json().get("pods", [])
+            deployments = response.json().get("deployments", [])
         else:
             try:
                 logger.error(f"Error when deploying: {response.json()['Message']}")
@@ -331,7 +330,7 @@ def main() -> None:
         sys.exit(0)
 
     log_group_name = f"/eks/{options.kube_cluster_name}-{options.kube_namespace}"
-    master_pod_name = next((pod for pod in deployed_pods if "master" in pod), None)
+    master_pod_name = next((deployment for deployment in deployments if "master" in deployment), None)
 
     if not master_pod_name:
         logger.error(
@@ -373,6 +372,7 @@ def main() -> None:
                     startTime=timestamp,
                     startFromHead=True,
                 )
+                locust_shutdown = False
                 for event in response.get("events", []):
                     message = event.get("message", "")
                     event_timestamp = event.get("timestamp", timestamp) + 1
@@ -380,9 +380,17 @@ def main() -> None:
                         message_json = json.loads(message)
                         if "log" in message_json:
                             print(message_json["log"])
+
+                            if "Shutting down (exit code" in message_json["log"]:
+                                locust_shutdown = True
+
                     except json.JSONDecodeError:
                         print(message)
                     timestamp = event_timestamp
+
+                if locust_shutdown:
+                    break
+
                 time.sleep(5)
             except ClientError as e:
                 error_code = e.response.get("Error", {}).get("Code", "")
@@ -408,6 +416,7 @@ def delete(s3_bucket, credential_manager):
             "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")

locust_cloud/credential_manager.py

@@ -8,6 +8,7 @@ import jwt
 import requests
 from botocore.credentials import RefreshableCredentials
 from botocore.session import Session as BotocoreSession
+from locust_cloud import __version__
 
 logger = logging.getLogger(__name__)
 
@@ -67,7 +68,11 @@ class CredentialManager:
             raise CredentialError("Insufficient credentials to obtain AWS session.")
 
         try:
-            response = requests.post(f"{self.lambda_url}/auth/login", json=payload)
+            response = requests.post(
+                f"{self.lambda_url}/auth/login",
+                json=payload,
+                headers={"X-Client-Version": __version__},
+            )
             response.raise_for_status()
             data = response.json()
 
@@ -98,6 +103,9 @@ class CredentialManager:
             if response is not None and response.status_code == 401:
                 raise CredentialError("Incorrect username or password.") from http_err
             else:
+                if js := response.json():
+                    if message := js.get("Message"):
+                        raise CredentialError(message)
                 error_info = f"HTTP {response.status_code} {response.reason}" if response else "No response received."
                 raise CredentialError(f"HTTP error occurred while obtaining credentials: {error_info}") from http_err
         except requests.exceptions.RequestException as req_err:

locust_cloud/idle_exit.py

@@ -0,0 +1,38 @@
+import logging
+import sys
+
+import gevent
+import locust.env
+from locust import events
+
+logger = logging.getLogger(__name__)
+
+
+class IdleExit:
+    def __init__(self, environment: locust.env.Environment):
+        self.environment = environment
+        self._destroy_task: gevent.Greenlet | None = None
+        events.test_start.add_listener(self.on_locust_state_change)
+        events.test_stop.add_listener(self.on_test_stop)
+        events.quit.add_listener(self.on_locust_state_change)
+
+        if not self.environment.parsed_options.autostart:
+            self._destroy_task = gevent.spawn(self._destroy)
+
+    def _destroy(self):
+        gevent.sleep(1800)
+        logger.info("Locust was detected as idle (no test running) for more than 30 minutes")
+        self.environment.runner.quit()
+
+        if self.environment.web_ui:
+            self.environment.web_ui.greenlet.kill(timeout=5)
+
+            if self.environment.web_ui.greenlet.started:
+                sys.exit(1)
+
+    def on_test_stop(self, **_kwargs):
+        self._destroy_task = gevent.spawn(self._destroy)
+
+    def on_locust_state_change(self, **_kwargs):
+        if self._destroy_task:
+            self._destroy_task.kill()

locust_cloud/timescale/exporter.py

@@ -22,11 +22,19 @@ def safe_serialize(obj):
     return json.dumps(obj, default=default)
 
 
+def format_datetime(d: datetime):
+    return d.strftime("%Y-%m-%d, %H:%M:%S.%f")
+
+
+def parse_datetime(s: str):
+    return datetime.strptime(s, "%Y-%m-%d, %H:%M:%S.%f").replace(tzinfo=UTC)
+
+
 class Exporter:
     def __init__(self, environment: locust.env.Environment, pool):
         self.env = environment
         self._run_id = None
-        self._samples: list[dict] = []
+        self._samples: list[tuple] = []
         self._background = gevent.spawn(self._run)
         self._hostname = socket.gethostname()
         self._finished = False
@@ -42,13 +50,6 @@ class Exporter:
         events.spawning_complete.add_listener(self.spawning_complete)
         atexit.register(self.log_stop_test_run)
 
-        if self.env.runner is not None:
-            self.env.runner.register_message("run_id", self.set_run_id)
-
-    def set_run_id(self, environment, msg, **kwargs):
-        logging.debug(f"run id from master: {msg.data}")
-        self._run_id = datetime.strptime(msg.data, "%Y-%m-%d, %H:%M:%S.%f").replace(tzinfo=UTC)
-
     def on_cpu_warning(self, environment: locust.env.Environment, cpu_usage, message=None, timestamp=None, **kwargs):
         # passing a custom message & timestamp to the event is a haxx to allow using this event for reporting generic events
         if not timestamp:
@@ -63,12 +64,12 @@ class Exporter:
     def on_test_start(self, environment: locust.env.Environment):
         if not self.env.parsed_options or not self.env.parsed_options.worker:
             self._run_id = environment._run_id = datetime.now(UTC)  # type: ignore
-            msg = environment._run_id.strftime("%Y-%m-%d, %H:%M:%S.%f")  # type: ignore
-            if environment.runner is not None:
-                logging.debug(f"about to send run_id to workers: {msg}")
-                environment.runner.send_message("run_id", msg)
+            self.env.parsed_options.run_id = format_datetime(environment._run_id)  # type: ignore
             self.log_start_testrun()
             self._user_count_logger = gevent.spawn(self._log_user_count)
+            self._update_end_time_task = gevent.spawn(self._update_end_time)
+        if self.env.parsed_options.worker:
+            self._run_id = parse_datetime(self.env.parsed_options.run_id)
 
     def _log_user_count(self):
         while True:
@@ -97,22 +98,40 @@ class Exporter:
                     break
             gevent.sleep(0.5)
 
+    def _update_end_time(self):
+        # delay setting first end time
+        # so UI doesn't display temporary value
+        gevent.sleep(5)
+
+        # Regularly update endtime to prevent missing endtimes when a test crashes
+        while True:
+            current_end_time = datetime.now(UTC)
+            try:
+                with self.pool.connection() as conn:
+                    conn.execute(
+                        "UPDATE testruns SET end_time = %s WHERE id = %s",
+                        (current_end_time, self._run_id),
+                    )
+                gevent.sleep(60)
+            except psycopg.Error as error:
+                logging.error("Failed to update testruns table with end time: " + repr(error))
+                gevent.sleep(1)
+
     def write_samples_to_db(self, samples):
         try:
             with self.pool.connection() as conn:
                 conn: psycopg.connection.Connection
-                with conn.cursor() as cur:
-                    cur.executemany(
-                        """
-    INSERT INTO requests (time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,response_length,exception,pid,url,context)
-    VALUES (%(time)s, %(run_id)s, %(greenlet_id)s, %(loadgen)s, %(name)s, %(request_type)s, %(response_time)s, %(success)s, %(response_length)s, %(exception)s, %(pid)s, %(url)s, %(context)s)
-    """,
-                        samples,
-                    )
+                with conn.cursor().copy(
+                    "COPY requests (time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,response_length,exception,pid,url,context) FROM STDIN"
+                ) as copy:
+                    for sample in samples:
+                        copy.write_row(sample)
         except psycopg.Error as error:
             logging.error("Failed to write samples to Postgresql timescale database: " + repr(error))
 
     def on_test_stop(self, environment):
+        if getattr(self, "_update_end_time_task", False):
+            self._update_end_time_task.kill()
         if getattr(self, "_user_count_logger", False):
             self._user_count_logger.kill()
             with self.pool.connection() as conn:
@@ -126,6 +145,8 @@ class Exporter:
         self._finished = True
         atexit._clear()  # make sure we dont capture additional ctrl-c:s
         self._background.join(timeout=10)
+        if getattr(self, "_update_end_time_task", False):
+            self._update_end_time_task.kill()
         if getattr(self, "_user_count_logger", False):
             self._user_count_logger.kill()
         self.log_stop_test_run(exit_code)
@@ -142,6 +163,9 @@ class Exporter:
         url=None,
         **kwargs,
     ):
+        # handle if a worker connects after test_start
+        if not self._run_id:
+            self._run_id = parse_datetime(self.env.parsed_options.run_id)
         success = 0 if exception else 1
         if start_time:
             time = datetime.fromtimestamp(start_time, tz=UTC)
@@ -150,35 +174,38 @@ class Exporter:
             # (which will be horribly wrong if users spend a lot of time in a with/catch_response-block)
             time = datetime.now(UTC) - timedelta(milliseconds=response_time or 0)
         greenlet_id = getattr(greenlet.getcurrent(), "minimal_ident", 0)  # if we're debugging there is no greenlet
-        sample = {
-            "time": time,
-            "run_id": self._run_id,
-            "greenlet_id": greenlet_id,
-            "loadgen": self._hostname,
-            "name": name,
-            "request_type": request_type,
-            "response_time": response_time,
-            "success": success,
-            "url": url[0:255] if url else None,
-            "pid": self._pid,
-            "context": psycopg.types.json.Json(context, safe_serialize),
-        }
 
         if response_length >= 0:
-            sample["response_length"] = response_length
+            response_length = response_length
         else:
-            sample["response_length"] = None
+            response_length = None
 
         if exception:
             if isinstance(exception, CatchResponseError):
-                sample["exception"] = str(exception)
+                exception = str(exception)
             else:
                 try:
-                    sample["exception"] = repr(exception)
+                    exception = repr(exception)
                 except AttributeError:
-                    sample["exception"] = f"{exception.__class__} (and it has no string representation)"
+                    exception = f"{exception.__class__} (and it has no string representation)"
         else:
-            sample["exception"] = None
+            exception = None
+
+        sample = (
+            time,
+            self._run_id,
+            greenlet_id,
+            self._hostname,
+            name,
+            request_type,
+            response_time,
+            success,
+            response_length,
+            exception,
+            self._pid,
+            url[0:255] if url else None,
+            psycopg.types.json.Json(context, safe_serialize),
+        )
 
         self._samples.append(sample)
 

locust_cloud/timescale/query.py

@@ -21,7 +21,7 @@ def register_query(environment, pool):
                 # start_time = time.perf_counter()
                 with pool.connection() as conn:
                     # get_conn_time = (time.perf_counter() - start_time) * 1000
-                    sql_params = request.get_json()
+                    sql_params = request.get_json() if request.content_type == "application/json" else {}
                     # start_time = time.perf_counter()
                     from datetime import datetime, timedelta
 
@@ -29,7 +29,7 @@ def register_query(environment, pool):
                         # 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):
+                        if end_time >= start_time + timedelta(hours=48):
                             logger.warning(
                                 f"UI asked for too long time interval. Start was {sql_params['start']}, end was {sql_params['end']}"
                             )