locustfile.py

import time

from locust import User, TaskSet, task, between, events
from trino.dbapi import connect
from trino import auth
from configparser import ConfigParser

config = ConfigParser()
config.read('config.ini')
trino_config = config['trino']

conn = connect(
    host = trino_config['host'],
    port = trino_config['port'],
    user = trino_config['user'],
    catalog = trino_config['catalog'],
    schema = trino_config['schema'],
    http_scheme = trino_config['http_scheme'],
    auth = auth.BasicAuthentication(trino_config['user'], trino_config['password']),
    verify = trino_config['verify_certs'].lower() == 'true'
)

query = """explain analyze SELECT sum(extendedprice * discount) AS revenue
FROM
  lineitem
WHERE
  shipdate >= DATE '1994-01-01'
  AND shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
  AND discount BETWEEN decimal '0.06' - decimal '0.01' AND decimal '0.06' + decimal '0.01'
  AND quantity < 24"""

def execute_query():

    cur = conn.cursor()
    cur.execute(query)
    rows = cur.fetchall()
    cur.close()

    return rows

class TrinoClient:

    def __getattr__(self, name):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            try:
                res = execute_query()
                events.request_success.fire(
                    request_type="trino",
                    name=name,
                    response_time=int((time.time() - start_time) * 1000),
                    response_length=len(res)
                )
            except Exception as e:
                events.request_failure.fire(
                    request_type="trino",
                    name=name,
                    response_time=int((time.time() - start_time) * 1000),
                    exception=e
                )

        return wrapper


class TrinoTaskSet(TaskSet):
    @task
    def execute_query(self):
        self.client.execute_query()


class TrinoUser(User):
    host = trino_config['host']

    wait_time = between(0.1, 1)

    def __init__(self, environment):
        super(TrinoUser, self).__init__(environment)
        self.client = TrinoClient()

    @task
    def execute_query(self):
        self.client.execute_query()

requirements.txt

locust==2.9.0
trino==0.320.0