locust-cloud 1.12.3__py3-none-any.whl → 1.13.0__py3-none-any.whl

locust_cloud/credential_manager.py

@@ -1,141 +0,0 @@
-import logging
-import time
-from datetime import UTC, datetime
-from typing import Any
-
-import boto3
-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__)
-
-
-class CredentialError(Exception):
-    """Custom exception for credential-related errors."""
-
-    pass
-
-
-class CredentialManager:
-    def __init__(
-        self,
-        lambda_url: str,
-        username: str | None = None,
-        password: str | None = None,
-        user_sub_id: str | None = None,
-        refresh_token: str | None = None,
-        access_key: str | None = None,
-        secret_key: str | None = None,
-    ) -> None:
-        self.lambda_url = lambda_url
-        self.username = username
-        self.password = password
-        self.user_sub_id = user_sub_id
-        self.refresh_token = refresh_token
-
-        self.credentials = {
-            "access_key": access_key,
-            "secret_key": secret_key,
-        }
-        self.cognito_client_id_token: str = ""
-        self.expiry_time: float = 0
-
-        self.obtain_credentials()
-
-        self.refreshable_credentials = RefreshableCredentials.create_from_metadata(
-            metadata=self.get_current_credentials(),
-            refresh_using=self.refresh_credentials,
-            method="custom-refresh",
-        )
-
-        botocore_session = BotocoreSession()
-        botocore_session._credentials = self.refreshable_credentials  # type: ignore
-        botocore_session.set_config_variable("signature_version", "v4")
-
-        self.session = boto3.Session(botocore_session=botocore_session)
-        logger.debug("Boto3 session created with RefreshableCredentials.")
-
-    def obtain_credentials(self) -> None:
-        payload = {}
-        if self.username and self.password:
-            payload = {"username": self.username, "password": self.password}
-        elif self.user_sub_id and self.refresh_token:
-            payload = {"user_sub_id": self.user_sub_id, "refresh_token": self.refresh_token}
-        else:
-            raise CredentialError("Insufficient credentials to obtain AWS session.")
-
-        try:
-            response = requests.post(
-                f"{self.lambda_url}/auth/login",
-                json=payload,
-                headers={"X-Client-Version": __version__},
-            )
-            response.raise_for_status()
-            data = response.json()
-
-            token_key = next(
-                (key for key in ["cognito_client_id_token", "id_token", "access_token"] if key in data), None
-            )
-
-            if not token_key:
-                raise CredentialError("No valid token found in authentication response.")
-
-            self.credentials = {
-                "access_key": data.get("aws_access_key_id"),
-                "secret_key": data.get("aws_secret_access_key"),
-                "token": data.get("aws_session_token"),
-            }
-
-            token = data.get(token_key)
-            if not token:
-                raise CredentialError(f"Token '{token_key}' is missing in the authentication response.")
-
-            decoded = jwt.decode(token, options={"verify_signature": False})
-            self.expiry_time = decoded.get("exp", time.time() + 3600) - 60  # Refresh 1 minute before expiry
-
-            self.cognito_client_id_token = token
-
-        except requests.exceptions.HTTPError as http_err:
-            response = http_err.response
-            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}"
-                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
-        except jwt.DecodeError as decode_err:
-            raise CredentialError(f"Failed to decode JWT token: {decode_err}") from decode_err
-        except KeyError as key_err:
-            raise CredentialError(f"Missing expected key in authentication response: {key_err}") from key_err
-
-    def refresh_credentials(self) -> dict[str, Any]:
-        logger.debug("Refreshing credentials using refresh_credentials method.")
-        self.obtain_credentials()
-        return {
-            "access_key": self.credentials.get("access_key"),
-            "secret_key": self.credentials.get("secret_key"),
-            "token": self.credentials.get("token"),
-            "expiry_time": datetime.fromtimestamp(self.expiry_time, tz=UTC).isoformat(),
-        }
-
-    def get_current_credentials(self) -> dict[str, Any]:
-        if not self.cognito_client_id_token:
-            raise CredentialError("cognito_client_id_token not set in CredentialManager.")
-
-        return {
-            "access_key": self.credentials.get("access_key"),
-            "secret_key": self.credentials.get("secret_key"),
-            "token": self.credentials.get("token"),
-            "expiry_time": datetime.fromtimestamp(self.expiry_time, tz=UTC).isoformat(),
-            "cognito_client_id_token": self.cognito_client_id_token,
-        }

locust_cloud/idle_exit.py

@@ -1,38 +0,0 @@
-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/socket_logging.py

@@ -1,127 +0,0 @@
-import atexit
-import logging
-import os
-import sys
-from collections import deque
-
-import flask
-import gevent
-import gevent.pywsgi
-import geventwebsocket.handler
-import socketio
-import socketio.exceptions
-
-
-def setup_socket_logging():
-    """
-    Set up a separate server listening for incomming websocket connections.
-    Because it listens on a separate port from the webui and the ALB only
-    accepts connections on port 443, this will be exposed on a different
-    path (/<customer-id>/socket-logs) with a separate target group.
-    The targetgroup will want to do health checks before exposing the server
-    so a small flask app exposing the health check endpoint is added to the
-    server in addition to the websocket stuff.
-    """
-
-    quiet_logger = logging.getLogger("be_quiet")
-    quiet_logger.propagate = False
-    quiet_logger.addHandler(logging.NullHandler())
-
-    # This app will use a logger with the same name as the application
-    # which means it will pick up the one set up above.
-    healthcheck_app = flask.Flask("be_quiet")
-
-    # /login is the health check endpoint currently configured for the
-    # ALB controller in kubernetes.
-    # See cloud-onboarder/kubernetes/alb-load-balancer-ingress.yaml
-    @healthcheck_app.route("/login")
-    def healthcheck():
-        return ""
-
-    logger = logging.getLogger(__name__)
-    socketio_path = f"/{os.environ['CUSTOMER_ID']}/socket-logs"
-    sio = socketio.Server(async_handlers=True, always_connect=True, async_mode="gevent", cors_allowed_origins="*")
-    sio_app = socketio.WSGIApp(sio, healthcheck_app, socketio_path=socketio_path)
-    message_queue = deque(maxlen=500)
-
-    @sio.event
-    def connect(sid, environ, auth):  # noqa: ARG001
-        logger.debug("Client connected to socketio server")
-
-    @sio.event
-    def disconnect(sid):  # noqa: ARG001
-        logger.debug("Client disconnected from socketio server")
-
-    class QueueCopyStream:
-        def __init__(self, name, original):
-            self.name = name
-            self.original = original
-
-        def write(self, message):
-            """
-            Writes to the queue first since when running in gevent the write to
-            stdout/stderr can yield and we want to ensure the same ordering in
-            the queue as for the written messages.
-            """
-            message_queue.append((self.name, message))
-            self.original.write(message)
-
-        def flush(self):
-            self.original.flush()
-
-    def connected_sid():
-        for rooms in sio.manager.rooms.values():
-            for clients in rooms.values():
-                for sid in clients.keys():
-                    return sid
-
-    def emitter():
-        while True:
-            while message_queue and (sid := connected_sid()):
-                name, message = message_queue[0]
-
-                try:
-                    if name == "shutdown":
-                        logger.debug("Sending websocket shutdown event")
-
-                    sio.call(name, message, to=sid, timeout=5)
-                    message_queue.popleft()
-
-                    if name == "shutdown":
-                        logger.debug("Websocket shutdown event acknowledged by client")
-                        return
-
-                except socketio.exceptions.TimeoutError:
-                    logger.debug("Timed out waiting for client to aknowledge websocket message")
-
-            gevent.sleep(1)
-
-    emitter_greenlet = gevent.spawn(emitter)
-
-    sys.stderr = QueueCopyStream("stderr", sys.stderr)
-    sys.stdout = QueueCopyStream("stdout", sys.stdout)
-
-    @atexit.register
-    def notify_shutdown(*args, **kwargs):
-        logger.debug("Adding shutdown event to websocket queue")
-        message_queue.append(("shutdown", ""))
-        emitter_greenlet.join(timeout=30)
-
-    class WebSocketHandlerWithoutLogging(geventwebsocket.handler.WebSocketHandler):
-        """
-        Subclassing WebSocketHandler so it doesn't set a logger on
-        the server that I've explicitly configured not to have one.
-        """
-
-        @property
-        def logger(self):
-            return quiet_logger
-
-    @gevent.spawn
-    def start_websocket_server():
-        logger.debug(f"Starting socketio server on port 1095 with path {socketio_path}")
-        server = gevent.pywsgi.WSGIServer(
-            ("", 1095), sio_app, log=None, error_log=None, handler_class=WebSocketHandlerWithoutLogging
-        )
-        server.serve_forever()
-        gevent.get_hub().join()

locust_cloud/timescale/exporter.py

@@ -1,313 +0,0 @@
-import atexit
-import json
-import logging
-import os
-import socket
-import sys
-from datetime import UTC, datetime, timedelta
-
-import gevent
-import greenlet
-import locust.env
-import psycopg
-import psycopg.types.json
-from locust.exception import CatchResponseError
-from locust.runners import MasterRunner
-
-
-def safe_serialize(obj):
-    def default(o):
-        return f"<<non-serializable: {type(o).__qualname__}>>"
-
-    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[tuple] = []
-        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
-
-        events = self.env.events
-        events.test_start.add_listener(self.on_test_start)
-        events.test_stop.add_listener(self.on_test_stop)
-        events.request.add_listener(self.on_request)
-        events.cpu_warning.add_listener(self.on_cpu_warning)
-        events.quit.add_listener(self.on_quit)
-        events.spawning_complete.add_listener(self.spawning_complete)
-        atexit.register(self.log_stop_test_run)
-
-    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:
-            timestamp = datetime.now(UTC).isoformat()
-        if not message:
-            message = f"High CPU usage ({cpu_usage}%)"
-        with self.pool.connection() as conn:
-            conn.execute(
-                "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()
-            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:
-            if self.env.runner is None:
-                return  # there is no runner, so nothing to log...
-            try:
-                with self.pool.connection() as conn:
-                    conn.execute(
-                        """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:
-                logging.error("Failed to write user count to Postgresql: " + repr(error))
-            gevent.sleep(2.0)
-
-    def _run(self):
-        while True:
-            if self._samples:
-                # Buffer samples, so that a locust greenlet will write to the new list
-                # instead of the one that has been sent into postgres client
-                samples_buffer = self._samples
-                self._samples = []
-                self.write_samples_to_db(samples_buffer)
-            else:
-                if self._finished:
-                    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().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:
-                conn.execute(
-                    """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
-        atexit.unregister(self.log_stop_test_run)  # 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()
-        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,
-        request_type,
-        name,
-        response_time,
-        response_length,
-        exception,
-        context,
-        start_time=None,
-        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)
-        else:
-            # some users may not send start_time, so we just make an educated guess
-            # (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
-
-        if exception:
-            if isinstance(exception, CatchResponseError):
-                exception = str(exception)
-            else:
-                try:
-                    exception = repr(exception)
-                except AttributeError:
-                    exception = f"{exception.__class__} (and it has no string representation)"
-
-            exception = exception[:300]
-        else:
-            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)
-
-    def log_start_testrun(self):
-        cmd = sys.argv[1:]
-        with self.pool.connection() as conn:
-            conn.execute(
-                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, profile, 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,
-                    len(self.env.runner.clients)
-                    if isinstance(
-                        self.env.runner,
-                        MasterRunner,
-                    )
-                    else 0,
-                    self.env.web_ui.template_args.get("username", "") if self.env.web_ui else "",
-                    self.env.parsed_locustfiles[0].split("/")[-1].split("__")[-1],
-                    self.env.parsed_options.profile,
-                    " ".join(cmd),
-                ),
-            )
-            conn.execute(
-                "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
-                (datetime.now(UTC).isoformat(), "Test run started", self._run_id),
-            )
-
-    def spawning_complete(self, user_count):
-        if not self.env.parsed_options.worker:  # only log for master/standalone
-            end_time = datetime.now(UTC)
-            try:
-                with self.pool.connection() as conn:
-                    conn.execute(
-                        "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:
-                logging.error(
-                    "Failed to insert rampup complete event time to Postgresql timescale database: " + repr(error)
-                )
-
-    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
-        end_time = datetime.now(UTC)
-        try:
-            with self.pool.connection() as conn:
-                conn.execute(
-                    "UPDATE testruns SET end_time = %s WHERE id = %s",
-                    (end_time, self._run_id),
-                )
-
-                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
-SET (requests, resp_time_avg, rps_avg, fail_ratio) =
-(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
- 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_view WHERE run_id = %(run_id)s AND time > %(run_id)s))::numeric AS duration) AS __,
-(SELECT
- COUNT(*)::numeric AS fails
- 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},
-                    )
-                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"
-                    )
-        except psycopg.Error as error:
-            logging.error(
-                "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)
-            )