locust-cloud 1.7.0__py3-none-any.whl → 1.9.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/auth.py

@@ -1,17 +1,20 @@
 import logging
 import os
 from datetime import UTC, datetime, timedelta
-from typing import TypedDict
+from typing import Any, TypedDict, cast
 
 import locust.env
 import requests
 import werkzeug
-from flask import redirect, request, url_for
+from flask import Blueprint, redirect, request, session, url_for
 from flask_login import UserMixin, login_user
+from locust.html import render_template_from
 from locust_cloud import __version__
-from locust_cloud.constants import DEFAULT_DEPLOYER_URL
 
-DEPLOYER_URL = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", DEFAULT_DEPLOYER_URL)
+REGION = os.environ.get("AWS_DEFAULT_REGION")
+DEPLOYER_URL = f"https://api.{REGION}.locust.cloud/1"
+ALLOW_SIGNUP = os.environ.get("ALLOW_SIGNUP", True)
+
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +50,9 @@ def register_auth(environment: locust.env.Environment):
     environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
     environment.web_ui.app.debug = False
 
+    web_base_path = environment.parsed_options.web_base_path
+    auth_blueprint = Blueprint("locust_cloud_auth", __name__, url_prefix=web_base_path)
+
     def load_user(user_sub_id: str):
         username = request.cookies.get("username")
         refresh_token = request.cookies.get("user_token")
@@ -58,11 +64,19 @@ def register_auth(environment: locust.env.Environment):
         return None
 
     environment.web_ui.login_manager.user_loader(load_user)
-    environment.web_ui.auth_args = {
-        "username_password_callback": "/authenticate",
-    }
+    environment.web_ui.auth_args = cast(
+        Any,
+        {
+            "username_password_callback": f"{web_base_path}/authenticate",
+        },
+    )
+
+    if ALLOW_SIGNUP:
+        environment.web_ui.auth_args["auth_providers"] = [
+            {"label": "Sign Up", "callback_url": f"{web_base_path}/signup"}
+        ]
 
-    @environment.web_ui.app.route("/authenticate", methods=["POST"])
+    @auth_blueprint.route("/authenticate", methods=["POST"])
     def login_submit():
         username = request.form.get("username", "")
         password = request.form.get("password")
@@ -77,19 +91,168 @@ def register_auth(environment: locust.env.Environment):
             auth_response.raise_for_status()
 
             credentials = auth_response.json()
-            response = redirect(url_for("index"))
+            response = redirect(url_for("locust.index"))
             response = set_credentials(username, credentials, response)
             login_user(AuthUser(credentials["user_sub_id"]))
 
             return response
         except requests.exceptions.HTTPError as e:
             if e.response.status_code == 401:
-                environment.web_ui.auth_args["error"] = "Invalid username or password"
+                session["auth_error"] = "Invalid username or password"
             else:
                 logger.error(f"Unknown response from auth: {e.response.status_code} {e.response.text}")
 
-                environment.web_ui.auth_args["error"] = (
-                    "Unknown error during authentication, check logs and/or contact support"
-                )
+                session["auth_error"] = "Unknown error during authentication, check logs and/or contact support"
+
+            return redirect(url_for("locust.login"))
+
+    @auth_blueprint.route("/signup")
+    def signup():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        if session.get("username"):
+            sign_up_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "Confirmation Code",
+                            "name": "confirmation_code",
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/confirm-signup",
+                    "submit_button_text": "Confirm Email",
+                },
+            }
+        else:
+            sign_up_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "Username",
+                            "name": "username",
+                        },
+                        {
+                            "label": "Password",
+                            "name": "password",
+                            "is_secret": True,
+                        },
+                        {
+                            "label": "Access Code",
+                            "name": "access_code",
+                        },
+                        {
+                            "label": "I consent to:\n\n1. Only test your own website/service or our example example target\n\n2. Only use locust-cloud for its intended purpose: to load test other sites/services.\n\n3. Not attempt to circumvent your account limitations (e.g. max user count or max request count)\n\n4. Not use personal data (real names, addresses etc) in your tests.",
+                            "name": "consent",
+                            "default_value": False,
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/create-account",
+                    "submit_button_text": "Sign Up",
+                },
+            }
+
+        if session.get("auth_info"):
+            sign_up_args["info"] = session["auth_info"]
+        if session.get("auth_sign_up_error"):
+            sign_up_args["error"] = session["auth_sign_up_error"]
+
+        return render_template_from(
+            "auth.html",
+            auth_args=sign_up_args,
+        )
+
+    @auth_blueprint.route("/create-account", methods=["POST"])
+    def create_account():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        session["auth_sign_up_error"] = ""
+        session["auth_info"] = ""
+
+        username = request.form.get("username", "")
+        password = request.form.get("password")
+        access_code = request.form.get("access_code")
+        has_consented = request.form.get("consent")
+
+        if not has_consented:
+            session["auth_sign_up_error"] = "Please accept the terms and conditions to create an account."
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+        try:
+            auth_response = requests.post(
+                f"{DEPLOYER_URL}/auth/signup",
+                json={"username": username, "password": password, "access_code": access_code},
+            )
+
+            auth_response.raise_for_status()
+
+            session["user_sub"] = auth_response.json().get("user_sub")
+            session["username"] = username
+            session["auth_info"] = (
+                "Please check your email and enter the confirmation code. If you didn't get a code after one minute, you can [request a new one](/resend-code)"
+            )
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+    @auth_blueprint.route("/resend-code")
+    def resend_code():
+        try:
+            auth_response = requests.post(
+                "http://localhost:8000/1/auth/resend-confirmation",
+                json={"username": session["username"]},
+            )
+
+            auth_response.raise_for_status()
+
+            session["auth_sign_up_error"] = ""
+            session["auth_info"] = "Confirmation code sent, please check your email."
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+    @auth_blueprint.route("/confirm-signup", methods=["POST"])
+    def confirm_signup():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        session["auth_sign_up_error"] = ""
+        confirmation_code = request.form.get("confirmation_code")
+
+        try:
+            auth_response = requests.post(
+                f"{DEPLOYER_URL}/auth/confirm-signup",
+                json={
+                    "username": session.get("username"),
+                    "user_sub": session["user_sub"],
+                    "confirmation_code": confirmation_code,
+                },
+            )
+
+            auth_response.raise_for_status()
+
+            session["username"] = None
+            session["auth_info"] = "Account created successfully!"
+            session["auth_sign_up_error"] = ""
+
+            return redirect(url_for("locust.login"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
 
-            return redirect(url_for("login"))
+    environment.web_ui.app.register_blueprint(auth_blueprint)

locust_cloud/cloud.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import math
 import os
 import sys
 import time
@@ -14,13 +13,6 @@ 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,
-    DEFAULT_NAMESPACE,
-    DEFAULT_REGION_NAME,
-    USERS_PER_WORKER,
-)
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
 
@@ -110,29 +102,9 @@ advanced.add_argument(
 advanced.add_argument(
     "--region",
     type=str,
-    default=os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION_NAME),
+    default=os.environ.get("AWS_DEFAULT_REGION"),
     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.",
 )
-advanced.add_argument(
-    "--kube-cluster-name",
-    type=str,
-    default=DEFAULT_CLUSTER_NAME,
-    help=configargparse.SUPPRESS,
-    env_var="KUBE_CLUSTER_NAME",
-)
-advanced.add_argument(
-    "--kube-namespace",
-    type=str,
-    default=DEFAULT_NAMESPACE,
-    help=configargparse.SUPPRESS,
-    env_var="KUBE_NAMESPACE",
-)
-parser.add_argument(
-    "--deployer-url",
-    type=str,
-    default=DEFAULT_DEPLOYER_URL,
-    help=configargparse.SUPPRESS,
-)
 parser.add_argument(
     "--aws-access-key-id",
     type=str,
@@ -162,7 +134,7 @@ parser.add_argument(
 parser.add_argument(
     "--workers",
     type=int,
-    help=f"Number of workers to use for the deployment. Defaults to number of users divided by {USERS_PER_WORKER}.",
+    help="Number of workers to use for the deployment. Defaults to number of users divided by 500, but the default may be customized for your account.",
     default=None,
 )
 parser.add_argument(
@@ -194,30 +166,30 @@ logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
 
+api_url = f"https://api.{options.region}.locust.cloud/1"
+
+
 def main() -> None:
-    s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
-    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.")
+    if not options.region:
+        logger.error(
+            "Setting a region is required to use Locust Cloud. Please ensure the AWS_DEFAULT_REGION env variable or the --region flag is set."
+        )
         sys.exit(1)
-    if worker_count > 1000:
-        logger.error("You asked for more than 20 workers, that isn't allowed.")
+
+    s3_bucket = "dmdb-default" if options.region == "us-east-1" else "locust-default"
+    deployments: list[Any] = []
+
+    if not ((options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)):
+        logger.error(
+            "Authentication is required to use Locust Cloud. Please ensure the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
+        )
         sys.exit(1)
-    try:
-        if not (
-            (options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)
-        ):
-            logger.error(
-                "Authentication is required to use Locust Cloud. Please ensure the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
-            )
-            sys.exit(1)
 
-        logger.info(f"Authenticating ({os.environ['AWS_DEFAULT_REGION']}, v{__version__})")
-        logger.debug(f"Lambda url: {options.deployer_url}")
+    try:
+        logger.info(f"Authenticating ({options.region}, v{__version__})")
+        logger.debug(f"Lambda url: {api_url}")
         credential_manager = CredentialManager(
-            lambda_url=options.deployer_url,
+            lambda_url=api_url,
             access_key=options.aws_access_key_id,
             secret_key=options.aws_secret_access_key,
             username=options.username,
@@ -284,19 +256,21 @@ def main() -> None:
             ]
             and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{options.deployer_url}/{options.kube_cluster_name}"
+        deploy_endpoint = f"{api_url}/deploy"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_LOCUSTFILE", "value": locustfile_url},
                 {"name": "LOCUST_USERS", "value": str(options.users)},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
                 {"name": "LOCUSTCLOUD_REQUIREMENTS_URL", "value": requirements_url},
-                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": options.deployer_url},
+                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": api_url},
                 *locust_env_variables,
             ],
-            "worker_count": worker_count,
+            "user_count": options.users,
             "image_tag": options.image_tag,
         }
+        if options.workers is not None:
+            payload["worker_count"] = options.workers
         headers = {
             "Authorization": f"Bearer {cognito_client_id_token}",
             "Content-Type": "application/json",
@@ -329,7 +303,7 @@ def main() -> None:
         logger.debug("Interrupted by user")
         sys.exit(0)
 
-    log_group_name = f"/eks/{options.kube_cluster_name}-{options.kube_namespace}"
+    log_group_name = "/eks/dmdb-default" if options.region == "us-east-1" else "/eks/locust-default"
     master_pod_name = next((deployment for deployment in deployments if "master" in deployment), None)
 
     if not master_pod_name:
@@ -424,9 +398,8 @@ def delete(s3_bucket, credential_manager):
             headers["AWS_SESSION_TOKEN"] = token
 
         response = requests.delete(
-            f"{options.deployer_url}/{options.kube_cluster_name}",
+            f"{api_url}/teardown",
             headers=headers,
-            params={"namespace": options.kube_namespace} if options.kube_namespace else {},
         )
 
         if response.status_code != 200:

locust_cloud/credential_manager.py

@@ -100,13 +100,16 @@ class CredentialManager:
 
         except requests.exceptions.HTTPError as http_err:
             response = http_err.response
-            if response is not None and response.status_code == 401:
+            if response is None:
+                raise CredentialError("Response was None?!") from http_err
+
+            if 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."
+                error_info = f"HTTP {response.status_code} {response.reason}"
                 raise CredentialError(f"HTTP error occurred while obtaining credentials: {error_info}") from http_err
         except requests.exceptions.RequestException as req_err:
             raise CredentialError(f"Request exception occurred while obtaining credentials: {req_err}") from req_err

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)
+            )