locust-cloud 1.5.1__py3-none-any.whl → 1.5.3__py3-none-any.whl

locust_cloud/cloud.py

@@ -173,6 +173,13 @@ parser.add_argument(
     action="store_true",
     help="Delete a running cluster. Useful if locust-cloud was killed/disconnected or if there was an error.",
 )
+parser.add_argument(
+    "--image-tag",
+    type=str,
+    env_var="LOCUST_CLOUD_IMAGE_TAG",
+    default="latest",
+    help=configargparse.SUPPRESS,  # overrides the locust-cloud docker image tag. for internal use
+)
 
 options, locust_options = parser.parse_known_args()
 options: Namespace
@@ -195,7 +202,12 @@ def main() -> None:
     s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
     deployed_pods: list[Any] = []
     worker_count: int = max(options.workers or math.ceil(options.users / USERS_PER_WORKER), 2)
-
+    if options.users > 10000:
+        logger.error("You asked for more than 10000 Users, that isn't allowed.")
+        sys.exit(1)
+    if worker_count > 20:
+        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)
@@ -289,6 +301,7 @@ def main() -> None:
                 *locust_env_variables,
             ],
             "worker_count": worker_count,
+            "image_tag": options.image_tag,
         }
         headers = {
             "Authorization": f"Bearer {cognito_client_id_token}",

locust_cloud/timescale/exporter.py

@@ -99,13 +99,15 @@ class Exporter:
     def write_samples_to_db(self, samples):
         try:
             with self.pool.connection() as conn:
-                conn.cursor().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,
-                )
+                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,
+                    )
         except psycopg.Error as error:
             logging.error("Failed to write samples to Postgresql timescale database: " + repr(error))