Add development server

mdesmet · mdesmet · commit 40c3a4e8f122 · 2024-11-30T10:30:56.000+08:00
diff --git a/etc/catalog/jmx.properties b/etc/catalog/jmx.properties
@@ -0,0 +1 @@
+connector.name=jmx
diff --git a/etc/catalog/memory.properties b/etc/catalog/memory.properties
@@ -0,0 +1 @@
+connector.name=memory
diff --git a/etc/catalog/tpcds.properties b/etc/catalog/tpcds.properties
@@ -0,0 +1 @@
+connector.name=tpcds
diff --git a/etc/catalog/tpch.properties b/etc/catalog/tpch.properties
@@ -0,0 +1,2 @@
+connector.name=tpch
+tpch.splits-per-node=4
diff --git a/etc/config-pre-466.properties b/etc/config-pre-466.properties
@@ -0,0 +1,11 @@
+node.id=coordinator
+node.environment=test
+
+coordinator=true
+node-scheduler.include-coordinator=true
+http-server.http.port=8080
+query.max-memory=1GB
+discovery.uri=http://localhost:8080
+
+# Disable http request log
+http-server.log.enabled=false
diff --git a/etc/config.properties b/etc/config.properties
@@ -0,0 +1,22 @@
+node.id=coordinator
+node.environment=test
+
+coordinator=true
+experimental.concurrent-startup=true
+node-scheduler.include-coordinator=true
+http-server.http.port=8080
+query.max-memory=1GB
+discovery.uri=http://localhost:8080
+
+# Use task.min-writer-count > 1, as this allows to expose writer-concurrency related bugs.
+task.min-writer-count=2
+task.concurrency=2
+task.max-writer-count=2
+
+# Experimental protocol spooling settings
+protocol.spooling.enabled=true
+protocol.spooling.shared-secret-key=jxTKysfCBuMZtFqUf8UJDQ1w9ez8rynEJsJqgJf66u0=
+protocol.spooling.retrieval-mode=coordinator_proxy
+
+# Disable http request log
+http-server.log.enabled=false
diff --git a/etc/jvm-pre-466.config b/etc/jvm-pre-466.config
@@ -0,0 +1,16 @@
+-server
+-Xmx2G
+-XX:G1HeapRegionSize=32M
+-XX:+ExplicitGCInvokesConcurrent
+-XX:+ExitOnOutOfMemoryError
+-XX:+HeapDumpOnOutOfMemoryError
+-XX:-OmitStackTraceInFastThrow
+-XX:ReservedCodeCacheSize=150M
+-XX:PerMethodRecompilationCutoff=10000
+-XX:PerBytecodeRecompilationCutoff=10000
+-Djdk.attach.allowAttachSelf=true
+# jdk.nio.maxCachedBufferSize controls what buffers can be allocated in per-thread "temporary buffer cache" (sun.nio.ch.Util). Value of 0 disables the cache.
+-Djdk.nio.maxCachedBufferSize=0
+# Allow loading dynamic agent used by JOL
+-XX:+EnableDynamicAgentLoading
+-XX:+UnlockDiagnosticVMOptions
diff --git a/etc/jvm.config b/etc/jvm.config
@@ -0,0 +1,17 @@
+-server
+-Xmx2G
+-XX:G1HeapRegionSize=32M
+-XX:+ExplicitGCInvokesConcurrent
+-XX:+ExitOnOutOfMemoryError
+-XX:+HeapDumpOnOutOfMemoryError
+-XX:-OmitStackTraceInFastThrow
+-XX:ReservedCodeCacheSize=150M
+-XX:PerMethodRecompilationCutoff=10000
+-XX:PerBytecodeRecompilationCutoff=10000
+-Djdk.attach.allowAttachSelf=true
+# jdk.nio.maxCachedBufferSize controls what buffers can be allocated in per-thread "temporary buffer cache" (sun.nio.ch.Util). Value of 0 disables the cache.
+-Djdk.nio.maxCachedBufferSize=0
+# Allow loading dynamic agent used by JOL
+-XX:+EnableDynamicAgentLoading
+-XX:+UnlockDiagnosticVMOptions
+--enable-native-access=ALL-UNNAMED
diff --git a/etc/spooling-manager.properties b/etc/spooling-manager.properties
@@ -0,0 +1,13 @@
+spooling-manager.name=filesystem
+fs.s3.enabled=true
+fs.location=s3://spooling/
+s3.endpoint=http://172.17.0.1:9000/
+s3.region=us-east-1
+s3.aws-access-key=minio-access-key
+s3.aws-secret-key=minio-secret-key
+s3.path-style-access=true
+fs.segment.ttl=5m
+fs.segment.pruning.interval=15s
+fs.segment.pruning.batch-size=250
+# Disable as we don't support SSE-C while writing/reading from S3
+fs.segment.encryption=false
diff --git a/setup.py b/setup.py
@@ -46,7 +46,9 @@
     "pre-commit",
     "black",
     "isort",
-    "keyring"
+    "keyring",
+    "testcontainers",
+    "minio"
 ]
 
 setup(
diff --git a/tests/development_server.py b/tests/development_server.py
@@ -0,0 +1,118 @@
+import os
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+from minio import Minio, S3Error
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+from trino.constants import DEFAULT_PORT
+
+MINIO_ROOT_USER = "minio-access-key"
+MINIO_ROOT_PASSWORD = "minio-secret-key"
+
+TRINO_VERSION = os.environ.get("TRINO_VERSION") or "latest"
+TRINO_HOST = "localhost"
+
+
+def create_bucket():
+    client = Minio(
+        "localhost:9000",
+        access_key=MINIO_ROOT_USER,
+        secret_key=MINIO_ROOT_PASSWORD,
+        secure=False
+    )
+    bucket_name = "spooling"
+    try:
+        # Check if the bucket exists
+        if not client.bucket_exists(bucket_name):
+            # Create the bucket if it does not exist
+            client.make_bucket(bucket_name)
+            print(f"Bucket {bucket_name} created successfully.")
+        else:
+            print(f"Bucket {bucket_name} already exists.")
+    except S3Error as e:
+        print(f"Error occurred: {e}")
+
+
+@contextmanager
+def start_development_server(port=None, trino_version=TRINO_VERSION):
+    minio = None
+    trino = None
+
+    try:
+        if TRINO_VERSION >= "466":
+            minio = DockerContainer("minio/minio:latest") \
+                .with_name("minio") \
+                .with_command(f"server --address '0.0.0.0:{9000}' --console-address '0.0.0.0:{9001}' -- /data") \
+                .with_env("MINIO_ROOT_USER", "minio-access-key") \
+                .with_env("MINIO_ROOT_PASSWORD", "minio-secret-key") \
+                .with_exposed_ports(9000, 9001) \
+                .with_bind_ports(9000, 9000) \
+                .with_bind_ports(9001, 9001)
+
+            # Start the container
+            print("Starting MinIO container...")
+            minio.start()
+
+            # Wait for logs indicating MinIO has started
+            wait_for_logs(minio, "API: http://", timeout=30)
+
+            # create spooling bucket
+            create_bucket()
+
+        trino = DockerContainer(f"trinodb/trino:{trino_version}") \
+            .with_name("trino") \
+            .with_env("TRINO_CONFIG_DIR", "/etc/trino") \
+            .with_bind_ports(DEFAULT_PORT, port)
+
+        root = Path(__file__).parent.parent
+
+        trino = trino \
+            .with_volume_mapping(str(root / "etc/catalog"), "/etc/trino/catalog")
+
+        # Enable spooling config
+        if TRINO_VERSION >= "466":
+            trino \
+                .with_volume_mapping(
+                    str(root / "etc/spooling-manager.properties"),
+                    "/etc/trino/spooling-manager.properties", "rw") \
+                .with_volume_mapping(str(root / "etc/jvm.config"), "/etc/trino/jvm.config") \
+                .with_volume_mapping(str(root / "etc/config.properties"), "/etc/trino/config.properties")
+        else:
+            trino \
+                .with_volume_mapping(str(root / "etc/jvm-pre-466.config"), "/etc/trino/jvm.config") \
+                .with_volume_mapping(str(root / "etc/config-pre-466.properties"), "/etc/trino/config.properties")
+
+        print("Starting Trino container...")
+        trino.start()
+
+        # Wait for logs indicating the service has started
+        wait_for_logs(trino, "SERVER STARTED", timeout=60)
+
+        # Otherwise some tests fail with No nodes available
+        time.sleep(2)
+
+        yield minio, trino
+    finally:
+        # Stop containers when exiting the context
+        if trino:
+            print("Stopping Trino container...")
+            trino.stop()
+        if minio:
+            print("Stopping MinIO container...")
+            minio.stop()
+
+
+def main():
+    """Run Trino setup independently from pytest."""
+    with start_development_server(port=DEFAULT_PORT):
+        print(f"Trino started at {TRINO_HOST}:{DEFAULT_PORT}")
+
+        # Keep the process running so that the containers stay up
+        input("Press Enter to stop containers...")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -12,31 +12,24 @@
 
 import os
 import socket
-import subprocess
-import time
 from contextlib import closing
-from uuid import uuid4
 
 import pytest
 
 import trino.logging
-from trino.client import ClientSession, TrinoQuery, TrinoRequest
+from tests.development_server import TRINO_HOST, TRINO_VERSION, start_development_server
 from trino.constants import DEFAULT_PORT
 
 logger = trino.logging.get_logger(__name__)
 
 
-TRINO_VERSION = os.environ.get("TRINO_VERSION") or "latest"
-TRINO_HOST = "127.0.0.1"
-TRINO_PORT = 8080
-
-
-def is_trino_available():
+def is_trino_available(host, port):
     with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
         sock.settimeout(2)
-        result = sock.connect_ex((TRINO_HOST, DEFAULT_PORT))
+        result = sock.connect_ex((host, port))
         if result == 0:
             return True
+    return False
 
 
 def get_local_port():
@@ -45,115 +38,21 @@ def get_local_port():
         return s.getsockname()[1]
 
 
-def get_default_trino_image_tag():
-    return "trinodb/trino:" + TRINO_VERSION
-
-
-def start_trino(image_tag=None):
-    if not image_tag:
-        image_tag = get_default_trino_image_tag()
-
-    container_id = "trino-python-client-tests-" + uuid4().hex[:7]
-    local_port = get_local_port()
-    logger.info("starting Docker container")
-    docker_run = [
-        "docker",
-        "run",
-        "--rm",
-        "-p",
-        "{host_port}:{cont_port}".format(host_port=local_port, cont_port=TRINO_PORT),
-        "--name",
-        container_id,
-        image_tag,
-    ]
-    run = subprocess.Popen(docker_run, universal_newlines=True, stderr=subprocess.PIPE)
-    return (container_id, run, "localhost", local_port)
-
-
-def wait_for_trino_workers(host, port, timeout=180):
-    request = TrinoRequest(
-        host=host,
-        port=port,
-        client_session=ClientSession(
-            user="test_fixture"
-        )
-    )
-    sql = "SELECT state FROM system.runtime.nodes"
-    t0 = time.time()
-    while True:
-        query = TrinoQuery(request, sql)
-        rows = list(query.execute())
-        if any(row[0] == "active" for row in rows):
-            break
-        if time.time() - t0 > timeout:
-            raise TimeoutError
-        time.sleep(1)
-
-
-def wait_for_trino_coordinator(stream, timeout=180):
-    started_tag = "======== SERVER STARTED ========"
-    t0 = time.time()
-    for line in iter(stream.readline, b""):
-        if line:
-            print(line)
-        if started_tag in line:
-            time.sleep(5)
-            return True
-        if time.time() - t0 > timeout:
-            logger.error("coordinator took longer than %s to start", timeout)
-            raise TimeoutError
-    return False
-
-
-def start_local_trino_server(image_tag):
-    container_id, proc, host, port = start_trino(image_tag)
-    print("trino.server.state starting")
-    trino_ready = wait_for_trino_coordinator(proc.stderr)
-    if not trino_ready:
-        raise Exception("Trino server did not start")
-    wait_for_trino_workers(host, port)
-    print("trino.server.state ready")
-    return container_id, proc, host, port
-
-
-def start_trino_and_wait(image_tag=None):
-    container_id = None
-    proc = None
-    host = os.environ.get("TRINO_RUNNING_HOST", None)
-    if host:
-        port = os.environ.get("TRINO_RUNNING_PORT", DEFAULT_PORT)
-    else:
-        container_id, proc, host, port = start_local_trino_server(
-            image_tag
-        )
-
-    print("trino.server.hostname {}".format(host))
-    print("trino.server.port {}".format(port))
-    if proc:
-        print("trino.server.pid {}".format(proc.pid))
-    if container_id:
-        print("trino.server.contained_id {}".format(container_id))
-    return container_id, proc, host, port
-
-
-def stop_trino(container_id, proc):
-    subprocess.check_call(["docker", "kill", container_id])
-
-
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def run_trino():
-    if is_trino_available():
-        yield None, TRINO_HOST, DEFAULT_PORT
-        return
+    host = os.environ.get("TRINO_RUNNING_HOST", TRINO_HOST)
+    port = os.environ.get("TRINO_RUNNING_PORT", DEFAULT_PORT)
 
-    image_tag = os.environ.get("TRINO_IMAGE")
-    if not image_tag:
-        image_tag = get_default_trino_image_tag()
+    # Is there any local Trino available
+    if is_trino_available(host, port):
+        yield host, port
+        return
 
-    container_id, proc, host, port = start_trino_and_wait(image_tag)
-    yield proc, host, port
-    if container_id or proc:
-        stop_trino(container_id, proc)
+    # Start Trino and MinIO server
+    print(f"Could not connect to Trino at {host}:{port}, starting server...")
+    local_port = get_local_port()
+    with start_development_server(port=local_port):
+        yield TRINO_HOST, local_port
 
 
 def trino_version():
diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py
diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py
diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+connector.name=tpch`
	`2`	`+tpch.splits-per-node=4`