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

locust_cloud/__init__.py

@@ -1,6 +1,8 @@
+import importlib.metadata
 import os
 
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
+__version__ = importlib.metadata.version("locust-cloud")
 
 import argparse
 import logging
@@ -13,6 +15,7 @@ from locust.argument_parser import LocustArgumentParser
 from locust_cloud.auth import register_auth
 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")
@@ -63,8 +66,17 @@ def create_connection_pool(
     pg_user: str, pg_host: str, pg_password: str, pg_database: str, pg_port: str | int
 ) -> ConnectionPool:
     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",
+        )
         return ConnectionPool(
-            conninfo=f"postgres://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_database}?sslmode=require",
+            conninfo,
             min_size=1,
             max_size=10,
             configure=set_autocommit,

locust_cloud/auth.py

@@ -7,9 +7,8 @@ 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)
 
@@ -69,6 +68,7 @@ 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:

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(
@@ -408,6 +407,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/timescale/exporter.py

@@ -28,6 +28,7 @@ class Exporter:
         self._run_id = None
         self._samples: list[dict] = []
         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()
@@ -97,6 +98,21 @@ class Exporter:
                     break
             gevent.sleep(0.5)
 
+    def _update_end_time(self):
+        # 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:
@@ -126,6 +142,7 @@ 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, "_user_count_logger", False):
             self._user_count_logger.kill()
         self.log_stop_test_run(exit_code)

locust_cloud/timescale/query.py

@@ -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']}"
                             )

{locust_cloud-1.5.11.dist-info → locust_cloud-1.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: locust-cloud
-Version: 1.5.11
+Version: 1.6.0
 Summary: Locust Cloud
 Project-URL: Homepage, https://locust.cloud
 Requires-Python: >=3.11

{locust_cloud-1.5.11.dist-info → locust_cloud-1.6.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
-locust_cloud/__init__.py,sha256=LBL_MzFOcQ92_Oyq9xn4RXkPnQLI3uiRqc3YwFgDWHI,3216
-locust_cloud/auth.py,sha256=PFox9bsu6m-jCgvXcdICWvx47bGe8fQ1OU4NSLEzuvg,2975
-locust_cloud/cloud.py,sha256=Q7L15mpceukU3YwcyS7Kln8QaNiOtdHN15R6RDuzwU8,16332
+locust_cloud/__init__.py,sha256=jryOnEQhJSi0zmb_NItb7ic_9e5K5Jf3cmmsuingFkw,3523
+locust_cloud/auth.py,sha256=aNyK2kiHa50UxNQ2NW2oomV8G6XIKCoZ1qcNMscnYiE,3062
+locust_cloud/cloud.py,sha256=B-5P-Yee5BEzz7g8n3sI_wDokMBiGSvBIjtZHpOADxo,16397
 locust_cloud/constants.py,sha256=75-M9cWMWOxB7RJrPccTviPz-WcMB04HatCYGC1wyRE,170
-locust_cloud/credential_manager.py,sha256=KFvoThBz7fMhgZiDqId0uQ7qIoOSnDACT4yXWgSy6X8,5235
-locust_cloud/timescale/exporter.py,sha256=76fvl0ylLuIBL_F_cXBBrEi6QdsPuXF8eTnJafo8OEs,11460
+locust_cloud/credential_manager.py,sha256=DYfog2doA2cfDZ7tgpVOdPegkShMJXUb_jkfV8NUGyE,5528
+locust_cloud/timescale/exporter.py,sha256=ZJLuGIEaCjLPtsqorpPBSkzkLrKZ4jMpeT-QW3CPwO8,12222
 locust_cloud/timescale/queries.py,sha256=vwORcodOB4z_ishj6_LxBD0i7XNSYzg4IeHFn0pBF1s,6296
-locust_cloud/timescale/query.py,sha256=NpihODXS4IvmfJL1AE7drDXOIunvsnwQVAnkcz1ypJg,2770
+locust_cloud/timescale/query.py,sha256=jLnKBEueTXdwisT3PmS_oOehR-maMs_VTRi7k46rzkE,2771
 locust_cloud/webui/.eslintrc,sha256=huuuTZNwbBGkhfelEOMhLQ-tgC9OYnnX0DRYgBn83Ig,1146
 locust_cloud/webui/.gitignore,sha256=i8EHIqDlVm1-Dkvf_GTZmP_Wu99GE7ABfbYzHXOU734,54
 locust_cloud/webui/.prettierrc,sha256=qZpzAJPlHUQ3--NHYK8WXgc_TlIqq-eoz4pbkTKK3nk,167
@@ -17,7 +17,7 @@ locust_cloud/webui/vite.config.ts,sha256=cqxPMkbwEA3H9mGGbuPulQUhIHCosUqm_1usxzs
 locust_cloud/webui/yarn.lock,sha256=YxZP905geusJV48VSr2lG2HrPyegUExE5CWmH8dS2Bg,227118
 locust_cloud/webui/dist/index.html,sha256=Z8x1932WXhLeHfvIVejfB3SxXI6wQP028GGl5VnkafM,664
 locust_cloud/webui/dist/assets/index-BrOP_HSY.js,sha256=GT8kkxuuVsbaVj-9KGFb6WTmoZ9hJm2eUP4LmZ7zUhU,2809361
-locust_cloud-1.5.11.dist-info/METADATA,sha256=5NQU9A_LKm5WQCUwhe-u4c59WyN-OSAAk9Sqgz7dCWw,1664
-locust_cloud-1.5.11.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-locust_cloud-1.5.11.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.5.11.dist-info/RECORD,,
+locust_cloud-1.6.0.dist-info/METADATA,sha256=tSYi56wjgFFvqxG9k8dTUeWRhr73v0AyJJLWtV8OjvI,1663
+locust_cloud-1.6.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+locust_cloud-1.6.0.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.6.0.dist-info/RECORD,,