change how the connection is opened

examples/postgres/locustfile.py

@@ -9,15 +9,13 @@ class UserTasks(TaskSet):
     @task
     def run_select_query(self):
         self.client.execute_query(
-            self.user.conn_string,
             "SELECT * FROM loadtesting.invoice WHERE amount > 500",
         )
 
     @task(3)
     def run_update_query(self):
         random_amount = random.randint(1, 12)
         self.client.execute_query(
-            self.user.conn_string,
             f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10",
         )
 
@@ -31,7 +29,7 @@ class PostgresLocust(PostgresUser):
     # Use environment variables or default values
     PGHOST = os.getenv("PGHOST", "localhost")
     PGPORT = os.getenv("PGPORT", "5432")
-    PGDATABASE = os.getenv("PGDATABASE", "loadtesting_db")
+    PGDATABASE = os.getenv("PGDATABASE", "postgres")
     PGUSER = os.getenv("PGUSER", "postgres")
     PGPASSWORD = os.getenv("PGPASSWORD", "postgres")
 

locust/contrib/postgres.py

@@ -9,41 +9,44 @@ def create_conn(conn_string):
     return psycopg.connect(conn_string)
 
 
-def execute_query(conn_string, query):
-    db_conn = create_conn(conn_string)
-    return db_conn.cursor().execute(query)
-
-
 class PostgresClient:
-    def __getattr__(self, name):
-        def request_handler(*args, **kwargs):
-            start_time = time.time()
-            try:
-                execute_query(*args, **kwargs)
-                response_time = int((time.time() - start_time) * 1000)
-                events.request.fire(
-                    request_type="postgres_success",
-                    name=name,
-                    response_time=response_time,
-                    response_length=0,
-                )
-            except Exception as e:
-                response_time = int((time.time() - start_time) * 1000)
-                events.request.fire(
-                    request_type="postgres_failure",
-                    name=name,
-                    response_time=response_time,
-                    response_length=0,
-                    exception=e,
-                )
-                print(f"error: {e}")
-
-        return request_handler
+    def __init__(self, conn_string):
+        self.conn_string = conn_string
+        self.connection = create_conn(conn_string)
+
+    def execute_query(self, query):
+        start_time = time.time()
+        try:
+            cursor = self.connection.cursor()
+            cursor.execute(query)
+            response_time = int((time.time() - start_time) * 1000)
+            events.request.fire(
+                request_type="postgres_success",
+                name="execute_query",
+                response_time=response_time,
+                response_length=0,
+            )
+        except Exception as e:
+            response_time = int((time.time() - start_time) * 1000)
+            events.request.fire(
+                request_type="postgres_failure",
+                name="execute_query",
+                response_time=response_time,
+                response_length=0,
+                exception=e,
+            )
+            print(f"error: {e}")
+
+    def close(self):
+        self.connection.close()
 
 
 class PostgresUser(User):
     abstract = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.client = PostgresClient()
+        self.client = PostgresClient(conn_string=self.conn_string)
+
+    def on_stop(self):
+        self.client.close()