Skip to content

Commit

Permalink
fix: Add container Trino (#642)
Browse files Browse the repository at this point in the history
Resolve #641
  • Loading branch information
grieve54706 committed Jul 15, 2024
1 parent 0ce4fec commit 49ce5a5
Show file tree
Hide file tree
Showing 5 changed files with 105 additions and 2 deletions.
2 changes: 2 additions & 0 deletions modules/trino/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.. autoclass:: testcontainers.trino.TrinoContainer
.. title:: testcontainers.trino.TrinoContainer
56 changes: 56 additions & 0 deletions modules/trino/testcontainers/trino/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import re

from testcontainers.core.config import testcontainers_config as c
from testcontainers.core.generic import DbContainer
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
from trino.dbapi import connect


class TrinoContainer(DbContainer):
def __init__(
self,
image="trinodb/trino:latest",
user: str = "test",
port: int = 8080,
**kwargs,
):
super().__init__(image=image, **kwargs)
self.user = user
self.port = port
self.with_exposed_ports(self.port)

@wait_container_is_ready()
def _connect(self) -> None:
wait_for_logs(
self,
re.compile(".*======== SERVER STARTED ========.*", re.MULTILINE).search,
c.max_tries,
c.sleep_time,
)
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()

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

def _configure(self):
pass
17 changes: 17 additions & 0 deletions modules/trino/tests/test_trino.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from testcontainers.trino import TrinoContainer
from trino.dbapi import connect


def test_docker_run_trino():
container = TrinoContainer("trinodb/trino:451")
with container as trino:
conn = connect(
host=trino.get_container_host_ip(),
port=trino.get_exposed_port(trino.port),
user="test",
)
cur = conn.cursor()
cur.execute("SELECT version()")
rows = cur.fetchall()
assert rows[0][0] == "451"
conn.close()
29 changes: 27 additions & 2 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ packages = [
{ include = "testcontainers", from = "modules/registry" },
{ include = "testcontainers", from = "modules/sftp" },
{ include = "testcontainers", from = "modules/selenium" },
{ include = "testcontainers", from = "modules/trino" },
{ include = "testcontainers", from = "modules/vault" },
{ include = "testcontainers", from = "modules/weaviate" },
]
Expand Down Expand Up @@ -111,6 +112,7 @@ bcrypt = { version = "*", optional = true }
httpx = { version = "*", optional = true }
azure-cosmos = { version = "*", optional = true }
cryptography = { version = "*", optional = true }
trino = { version = "*", optional = true }

[tool.poetry.extras]
arangodb = ["python-arango"]
Expand Down Expand Up @@ -153,6 +155,7 @@ sftp = ["cryptography"]
vault = []
weaviate = ["weaviate-client"]
chroma = ["chromadb-client"]
trino = ["trino"]

[tool.poetry.group.dev.dependencies]
mypy = "1.7.1"
Expand Down

0 comments on commit 49ce5a5

Please sign in to comment.