Skip to content

Commit bb646e9

Browse files
authored
fix(trino): Remove deprecated class and decorator from Trino container (#895)
1 parent 2e4d80a commit bb646e9

File tree

1 file changed

+8
-21
lines changed

1 file changed

+8
-21
lines changed

modules/trino/testcontainers/trino/__init__.py

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -13,41 +13,28 @@
1313
import re
1414

1515
from testcontainers.core.config import testcontainers_config as c
16-
from testcontainers.core.generic import DbContainer
17-
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
18-
from trino.dbapi import connect
16+
from testcontainers.core.generic import DockerContainer
17+
from testcontainers.core.wait_strategies import LogMessageWaitStrategy
1918

2019

21-
class TrinoContainer(DbContainer):
20+
class TrinoContainer(DockerContainer):
2221
def __init__(
2322
self,
2423
image="trinodb/trino:latest",
2524
user: str = "test",
2625
port: int = 8080,
26+
container_start_timeout: int = 30,
2727
**kwargs,
2828
):
2929
super().__init__(image=image, **kwargs)
3030
self.user = user
3131
self.port = port
3232
self.with_exposed_ports(self.port)
33-
34-
@wait_container_is_ready()
35-
def _connect(self) -> None:
36-
wait_for_logs(
37-
self,
38-
re.compile(".*======== SERVER STARTED ========.*", re.MULTILINE).search,
39-
c.max_tries,
40-
c.sleep_time,
41-
)
42-
conn = connect(
43-
host=self.get_container_host_ip(),
44-
port=self.get_exposed_port(self.port),
45-
user=self.user,
33+
self.waiting_for(
34+
LogMessageWaitStrategy(re.compile(".*======== SERVER STARTED ========.*", re.MULTILINE))
35+
.with_poll_interval(c.sleep_time)
36+
.with_startup_timeout(container_start_timeout)
4637
)
47-
cur = conn.cursor()
48-
cur.execute("SELECT 1")
49-
cur.fetchall()
50-
conn.close()
5138

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

0 commit comments

Comments
 (0)