Skip to content

Commit

Permalink
feat(trino): use waiting tpch instead of sleep delay time
Browse files Browse the repository at this point in the history
  • Loading branch information
grieve54706 committed Nov 19, 2024
1 parent 5b1ac8c commit 622d00a
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions modules/trino/testcontainers/trino/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,12 @@ def __init__(
image="trinodb/trino:latest",
user: str = "test",
port: int = 8080,
delay: int = 5,
**kwargs,
):
super().__init__(image=image, **kwargs)
self.user = user
self.port = port
self.with_exposed_ports(self.port)
self.delay = delay

@wait_container_is_ready()
def _connect(self) -> None:
Expand All @@ -42,17 +40,21 @@ def _connect(self) -> None:
c.max_tries,
c.sleep_time,
)
# To avoid `TrinoQueryError(type=INTERNAL_ERROR, name=GENERIC_INTERNAL_ERROR, message="nodes is empty")`
time.sleep(self.delay)
conn = connect(
host=self.get_container_host_ip(),
port=self.get_exposed_port(self.port),
user=self.user,
)
cur = conn.cursor()
cur.execute("SELECT 1")
cur.fetchall()
conn.close()
deadline = time.time() + c.max_tries
while time.time() < deadline:
try:
cur = conn.cursor()
cur.execute("SELECT * FROM tpch.tiny.nation LIMIT 1")
cur.fetchall()
return
except Exception:
time.sleep(c.sleep_time)
raise TimeoutError(f"Trino did not start within {c.max_tries:.3f} seconds")

def get_connection_url(self):
return f"trino://{self.user}@{self.get_container_host_ip()}:{self.port}"
Expand Down

0 comments on commit 622d00a

Please sign in to comment.