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

locust_cloud/__init__.py

@@ -19,22 +19,16 @@ 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")
-PG_HOST = os.environ.get("PG_HOST")
-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")
-MAX_USER_COUNT = os.environ.get("MAX_USER_COUNT")
 logger = logging.getLogger(__name__)
 
 
 @events.init_command_line_parser.add_listener
 def add_arguments(parser: LocustArgumentParser):
-    if not (PG_HOST or GRAPH_VIEWER):
+    if not (os.environ.get("PGHOST") or GRAPH_VIEWER):
         parser.add_argument_group(
             "locust-cloud",
-            "locust-cloud disabled, because PG_HOST was not set - this is normal for local runs",
+            "locust-cloud disabled, because PGHOST was not set - this is normal for local runs",
         )
         return
 
@@ -73,30 +67,23 @@ def set_autocommit(conn: psycopg.Connection):
 
 @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):
+    if not (os.environ.get("PGHOST")):
         return
 
     try:
         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,
+            max_size=20,
             configure=set_autocommit,
             check=ConnectionPool.check_connection,
         )
         pool.wait()
     except Exception as e:
         logger.exception(e)
-        logger.error(f"{PG_HOST=}")
         raise
 
     if not GRAPH_VIEWER:
@@ -106,8 +93,6 @@ def on_locust_init(environment: locust.env.Environment, **_args):
         Exporter(environment, pool)
 
     if environment.web_ui:
-        environment.web_ui.template_args["maxUserCount"] = MAX_USER_COUNT
-
         if GRAPH_VIEWER:
             environment.web_ui.template_args["isGraphViewer"] = True
 

locust_cloud/cloud.py

@@ -199,12 +199,7 @@ def main() -> None:
     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:
-        logger.error("You asked for more than 5000000 Users, that isn't allowed.")
-        sys.exit(1)
-    if worker_count > 1000:
-        logger.error("You asked for more than 20 workers, that isn't allowed.")
-        sys.exit(1)
+
     try:
         if not (
             (options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)
@@ -295,6 +290,7 @@ def main() -> None:
                 *locust_env_variables,
             ],
             "worker_count": worker_count,
+            "user_count": options.users,
             "image_tag": options.image_tag,
         }
         headers = {

locust_cloud/idle_exit.py

@@ -30,9 +30,9 @@ class IdleExit:
             if self.environment.web_ui.greenlet.started:
                 sys.exit(1)
 
-    def on_test_stop(self, **_kwargs):
+    def on_test_stop(self, **kwargs):
         self._destroy_task = gevent.spawn(self._destroy)
 
-    def on_locust_state_change(self, **_kwargs):
+    def on_locust_state_change(self, **kwargs):
         if self._destroy_task:
             self._destroy_task.kill()

locust_cloud/timescale/exporter.py

@@ -38,6 +38,7 @@ class Exporter:
         self._background = gevent.spawn(self._run)
         self._hostname = socket.gethostname()
         self._finished = False
+        self._has_logged_test_stop = False
         self._pid = os.getpid()
         self.pool = pool
 
@@ -58,11 +59,13 @@ class Exporter:
             message = f"High CPU usage ({cpu_usage}%)"
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)", (timestamp, message, self._run_id)
+                "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
+                (timestamp, message, self._run_id),
             )
 
     def on_test_start(self, environment: locust.env.Environment):
         if not self.env.parsed_options or not self.env.parsed_options.worker:
+            self._has_logged_test_stop = False
             self._run_id = environment._run_id = datetime.now(UTC)  # type: ignore
             self.env.parsed_options.run_id = format_datetime(environment._run_id)  # type: ignore
             self.log_start_testrun()
@@ -78,7 +81,7 @@ class Exporter:
             try:
                 with self.pool.connection() as conn:
                     conn.execute(
-                        """INSERT INTO number_of_users(time, run_id, user_count) VALUES (%s, %s, %s)""",
+                        """INSERT INTO number_of_users(time, run_id, user_count, customer) VALUES (%s, %s, %s, current_user)""",
                         (datetime.now(UTC).isoformat(), self._run_id, self.env.runner.user_count),
                     )
             except psycopg.Error as error:
@@ -136,10 +139,11 @@ class Exporter:
             self._user_count_logger.kill()
             with self.pool.connection() as conn:
                 conn.execute(
-                    """INSERT INTO number_of_users(time, run_id, user_count) VALUES (%s, %s, %s)""",
+                    """INSERT INTO number_of_users(time, run_id, user_count, customer) VALUES (%s, %s, %s, current_user)""",
                     (datetime.now(UTC).isoformat(), self._run_id, 0),
                 )
         self.log_stop_test_run()
+        self._has_logged_test_stop = True
 
     def on_quit(self, exit_code, **kwargs):
         self._finished = True
@@ -149,7 +153,10 @@ class Exporter:
             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)
+        if not self._has_logged_test_stop:
+            self.log_stop_test_run()
+        if not self.env.parsed_options.worker:
+            self.log_exit_code(exit_code)
 
     def on_request(
         self,
@@ -175,11 +182,6 @@ class Exporter:
             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
 
-        if response_length >= 0:
-            response_length = response_length
-        else:
-            response_length = None
-
         if exception:
             if isinstance(exception, CatchResponseError):
                 exception = str(exception)
@@ -188,6 +190,8 @@ class Exporter:
                     exception = repr(exception)
                 except AttributeError:
                     exception = f"{exception.__class__} (and it has no string representation)"
+
+            exception = exception[:300]
         else:
             exception = None
 
@@ -213,7 +217,7 @@ class Exporter:
         cmd = sys.argv[1:]
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, description, arguments) VALUES (%s,%s,%s,%s,%s,%s,%s)",
+                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, description, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,%s,current_user)",
                 (
                     self._run_id,
                     self.env.runner.target_user_count if self.env.runner else 1,
@@ -230,7 +234,7 @@ class Exporter:
                 ),
             )
             conn.execute(
-                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
                 (datetime.now(UTC).isoformat(), "Test run started", self._run_id),
             )
 
@@ -240,7 +244,7 @@ class Exporter:
             try:
                 with self.pool.connection() as conn:
                     conn.execute(
-                        "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                        "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
                         (end_time, f"Rampup complete, {user_count} users spawned", self._run_id),
                     )
             except psycopg.Error as error:
@@ -248,7 +252,7 @@ class Exporter:
                     "Failed to insert rampup complete event time to Postgresql timescale database: " + repr(error)
                 )
 
-    def log_stop_test_run(self, exit_code=None):
+    def log_stop_test_run(self):
         logging.debug(f"Test run id {self._run_id} stopping")
         if self.env.parsed_options.worker:
             return  # only run on master or standalone
@@ -256,17 +260,14 @@ class Exporter:
         try:
             with self.pool.connection() as conn:
                 conn.execute(
-                    "UPDATE testruns SET end_time = %s, exit_code = %s where id = %s",
-                    (end_time, exit_code, self._run_id),
-                )
-                conn.execute(
-                    "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
-                    (end_time, f"Finished with exit code: {exit_code}", self._run_id),
+                    "UPDATE testruns SET end_time = %s WHERE id = %s",
+                    (end_time, self._run_id),
                 )
-                # The AND time > run_id clause in the following statements are there to help Timescale performance
-                # We dont use start_time / end_time to calculate RPS, instead we use the time between the actual first and last request
-                # (as this is a more accurate measurement of the actual test)
+
                 try:
+                    # The AND time > run_id clause in the following statements are there to help Timescale performance
+                    # We dont use start_time / end_time to calculate RPS, instead we use the time between the actual first and last request
+                    # (as this is a more accurate measurement of the actual test)
                     conn.execute(
                         """
 UPDATE testruns
@@ -275,12 +276,12 @@ SET (requests, resp_time_avg, rps_avg, fail_ratio) =
 (SELECT
  COUNT(*)::numeric AS reqs,
  AVG(response_time)::numeric as resp_time
- FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s) AS _,
+ FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s) AS _,
 (SELECT
- EXTRACT(epoch FROM (SELECT MAX(time)-MIN(time) FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s))::numeric AS duration) AS __,
+ EXTRACT(epoch FROM (SELECT MAX(time)-MIN(time) FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s))::numeric AS duration) AS __,
 (SELECT
  COUNT(*)::numeric AS fails
- FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s AND success = 0) AS ___
+ FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s AND success = 0) AS ___
 WHERE id = %(run_id)s""",
                         {"run_id": self._run_id},
                     )
@@ -293,3 +294,20 @@ WHERE id = %(run_id)s""",
                 "Failed to update testruns record (or events) with end time to Postgresql timescale database: "
                 + repr(error)
             )
+
+    def log_exit_code(self, exit_code=None):
+        try:
+            with self.pool.connection() as conn:
+                conn.execute(
+                    "UPDATE testruns SET exit_code = %s WHERE id = %s",
+                    (exit_code, self._run_id),
+                )
+                conn.execute(
+                    "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
+                    (datetime.now(UTC).isoformat(), f"Finished with exit code: {exit_code}", self._run_id),
+                )
+        except psycopg.Error as error:
+            logging.error(
+                "Failed to update testruns record (or events) with end time to Postgresql timescale database: "
+                + repr(error)
+            )

locust_cloud/timescale/queries.py

@@ -10,7 +10,7 @@ SELECT
 	MIN(min),
 	MAX(max),
 	SUM(failed_count) / SUM(count) * 100 as "errorPercentage"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY name, method
@@ -20,13 +20,13 @@ GROUP BY name, method
 failures_query = """
 SELECT
   name as name,
-  left(exception,300) as exception,
+  exception,
   count(*)
-FROM requests
+FROM requests_view
 WHERE time BETWEEN %(start)s AND %(end)s AND
- exception is not null
+ success = 0
 AND run_id = %(testrun)s
-GROUP BY "name",left(exception,300)
+GROUP BY "name",exception
 """
 
 
@@ -35,7 +35,7 @@ WITH request_count_agg AS (
   SELECT
     time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     COALESCE(SUM(count)/%(resolution)s, 0) as rps
-  FROM requests_summary
+  FROM requests_summary_view
   WHERE bucket BETWEEN %(start)s AND %(end)s
   AND run_id = %(testrun)s
   GROUP BY 1
@@ -55,7 +55,7 @@ errors_per_s_agg AS (
   SELECT
     time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     COALESCE(SUM(failed_count)/%(resolution)s, 0) as error_rate
-  FROM requests_summary
+  FROM requests_summary_view
   WHERE bucket BETWEEN %(start)s AND %(end)s
   AND run_id = %(testrun)s
   GROUP BY 1
@@ -76,7 +76,7 @@ ORDER BY r.time;
 total_requests = """
 SELECT
  SUM(count) as "totalRequests"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 """
@@ -85,7 +85,7 @@ AND run_id = %(testrun)s
 total_failed = """
 SELECT
  SUM(failed_count) as "totalFailures"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 """
@@ -94,7 +94,7 @@ AND run_id = %(testrun)s
 error_percentage = """
 SELECT
 	SUM(failed_count) / SUM(count) * 100 "errorPercentage"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 """
@@ -104,7 +104,7 @@ SELECT
     time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     name,
     COALESCE(SUM(count)/%(resolution)s, 0) as throughput
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name
@@ -117,7 +117,7 @@ SELECT
     time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) as time,
     name,
     avg(average) as "responseTime"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name
@@ -129,7 +129,7 @@ SELECT
     time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
     name,
     SUM(failed_count)/%(resolution)s as "errorRate"
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name
@@ -138,10 +138,10 @@ ORDER BY 1
 
 
 perc99_response_times = """
-SELECT time_bucket(%(resolution)s * interval '1 second', bucket) AS time,
+SELECT time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) AS time,
   name,
   MAX(perc99) as perc99
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name
@@ -151,10 +151,10 @@ ORDER BY 1
 
 response_length = """
 SELECT
-    time_bucket(%(resolution)s * interval '1 second', bucket) as time,
+    time_bucket_gapfill(%(resolution)s * interval '1 second', bucket) as time,
     AVG(response_length) as "responseLength",
     name
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 GROUP BY 1, name
@@ -164,7 +164,7 @@ ORDER BY 1
 
 request_names = """
 SELECT DISTINCT name
-FROM requests_summary
+FROM requests_summary_view
 WHERE bucket BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 """
@@ -174,7 +174,7 @@ SELECT
  time,
  name,
  response_time as "responseTime"
-FROM requests
+FROM requests_view
 WHERE time BETWEEN %(start)s AND %(end)s
 AND run_id = %(testrun)s
 ORDER BY 1,2
@@ -260,13 +260,23 @@ JOIN avg_response_time_failed f ON a.time = f.time
 ORDER BY a.time
 """
 
-total_runtime = """
+total_vuh = """
 SELECT
   SUM((end_time - id) * num_users) AS "totalVuh"
 FROM testruns
 WHERE id >= date_trunc('month', NOW()) AND NOT refund
 """
 
+customer = """
+SELECT
+  max_vuh as "maxVuh",
+  max_workers as "maxWorkers",
+  max_users as "maxUsers",
+  users_per_worker as "usersPerWorker"
+FROM customers
+WHERE customer = current_user
+"""
+
 queries: dict["str", LiteralString] = {
     "request-names": request_names,
     "requests": requests_query,
@@ -285,5 +295,6 @@ queries: dict["str", LiteralString] = {
     "testruns-table": testruns_table,
     "testruns-rps": testruns_rps,
     "testruns-response-time": testruns_response_time,
-    "total-runtime": total_runtime,
+    "total-vuh": total_vuh,
+    "customer": customer,
 }