Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add PostgresUser to examples #2836

Merged
merged 16 commits into from
Aug 24, 2024
Merged
Show file tree
Hide file tree
Changes from 15 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions examples/postgres/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Overview

Read the instruction below for your specific database

## PostgreSQL

### How to run the test

- Prerequisites:

- `psycopg3` - https://www.psycopg.org/psycopg3/docs/basic/install.html

- Set your environment variables for:

- PGHOST
- PGPORT
- PGDATABASE
- PGUSER
- PGPASSWORD

- Run locust as usual, see https://docs.locust.io/en/stable/quickstart.html
36 changes: 36 additions & 0 deletions examples/postgres/locustfile.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
from locust import TaskSet, between, task
from locust.contrib.postgres import PostgresUser

import os
import random


class UserTasks(TaskSet):
@task
def run_select_query(self):
self.client.execute_query(
"SELECT * FROM loadtesting.invoice WHERE amount > 500",
)

@task(3)
def run_update_query(self):
random_amount = random.randint(1, 12)
self.client.execute_query(
f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10",
)


class PostgresLocust(PostgresUser):
tasks = [UserTasks]
min_wait = 0
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove the wait_time stuff and just leave it as the default, to simplify the example.

You can skip the taskset and put the tasks directly under the user too.

max_wait = 3
wait_time = between(min_wait, max_wait)

# Use environment variables or default values
PGHOST = os.getenv("PGHOST", "localhost")
PGPORT = os.getenv("PGPORT", "5432")
PGDATABASE = os.getenv("PGDATABASE", "postgres")
PGUSER = os.getenv("PGUSER", "postgres")
PGPASSWORD = os.getenv("PGPASSWORD", "postgres")

conn_string = f"postgresql://{PGUSER}:{PGPASSWORD}@{PGHOST}:{PGPORT}/{PGDATABASE}"
52 changes: 52 additions & 0 deletions locust/contrib/postgres.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
from locust import TaskSet, User, events

import time

import psycopg


def create_conn(conn_string):
Copy link
Collaborator

@cyberw cyberw Aug 24, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this method is too small to make sense! The code is easier to understand if you just call .connect directly.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good call, it looks a lot cleaner calling the psycopg.connect directly

return psycopg.connect(conn_string)


class PostgresClient:
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()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

psycopg3 can do connection.execute(), no need to have the intermediate step of creating a 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}")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I dont think we should print the exception :)


def close(self):
self.connection.close()


class PostgresUser(User):
abstract = True

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.client = PostgresClient(conn_string=self.conn_string)

def on_stop(self):
self.client.close()