Skip to content

Instantly share code, notes, and snippets.

@posulliv
Created December 15, 2022 19:51
Show Gist options
  • Save posulliv/e0f0165541e45a9acec2f97ff1db34b5 to your computer and use it in GitHub Desktop.
Save posulliv/e0f0165541e45a9acec2f97ff1db34b5 to your computer and use it in GitHub Desktop.
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()
locust==2.9.0
trino==0.320.0
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment