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

locust_cloud/__init__.py

@@ -6,13 +6,14 @@ __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
@@ -56,48 +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:
+@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:
         conninfo = make_conninfo(
-            dbname=pg_database,
-            user=pg_user,
-            port=pg_port,
-            password=pg_password,
-            host=pg_host,
+            dbname=PG_DATABASE,
+            user=PG_USER,
+            port=PG_PORT,
+            password=PG_PASSWORD,
+            host=PG_HOST,
             sslmode="require",
-            options="-c statement_timeout=55000",
+            # options="-c statement_timeout=55000",
         )
-        return ConnectionPool(
+        pool = ConnectionPool(
             conninfo,
             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,
+            check=ConnectionPool.check_connection,
         )
         pool.wait()
     except Exception as e:
@@ -105,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
@@ -12,6 +13,8 @@ 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
@@ -71,21 +74,22 @@ def register_auth(environment: locust.env.Environment):
                 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

@@ -372,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
@@ -379,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", "")

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,13 +22,20 @@ 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._update_end_time_task = gevent.spawn(self._update_end_time)
         self._hostname = socket.gethostname()
         self._finished = False
         self._pid = os.getpid()
@@ -43,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:
@@ -64,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:
@@ -99,6 +99,10 @@ class Exporter:
             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)
@@ -117,18 +121,17 @@ class Exporter:
         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:
@@ -142,7 +145,8 @@ class Exporter:
         self._finished = True
         atexit._clear()  # make sure we dont capture additional ctrl-c:s
         self._background.join(timeout=10)
-        self._update_end_time_task.kill()
+        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)
@@ -159,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)
@@ -167,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