Examples with TestingTrinoServer io.trino.server.testing.TestingTrinoServer used on opensource projects

Example 21 with TestingTrinoServer

use of io.trino.server.testing.TestingTrinoServer in project trino by trinodb.

the class DistributedQueryRunner method createTestingTrinoServer.

private static TestingTrinoServer createTestingTrinoServer(URI discoveryUri, boolean coordinator, Map<String, String> extraProperties, String environment, Module additionalModule, Optional<Path> baseDataDir, List<SystemAccessControl> systemAccessControls, List<EventListener> eventListeners) {
    long start = System.nanoTime();
    ImmutableMap.Builder<String, String> propertiesBuilder = ImmutableMap.<String, String>builder().put("internal-communication.shared-secret", "test-secret").put("query.client.timeout", "10m").put("discovery.http-client.min-threads", // default 8
    "1").put("exchange.http-client.min-threads", // default 8
    "1").put("node-manager.http-client.min-threads", // default 8
    "1").put("exchange.page-buffer-client.max-callback-threads", // default 25
    "5").put("exchange.http-client.idle-timeout", "1h").put("task.max-index-memory", // causes index joins to fault load
    "16kB").put("distributed-index-joins-enabled", "true");
    if (coordinator) {
        propertiesBuilder.put("node-scheduler.include-coordinator", "true");
        propertiesBuilder.put("join-distribution-type", "PARTITIONED");
        // Use few threads in tests to preserve resources on CI
        // default 8
        propertiesBuilder.put("failure-detector.http-client.min-threads", "1");
        // default 8
        propertiesBuilder.put("memoryManager.http-client.min-threads", "1");
        // default 8
        propertiesBuilder.put("scheduler.http-client.min-threads", "1");
        // default 8
        propertiesBuilder.put("workerInfo.http-client.min-threads", "1");
    }
    HashMap<String, String> properties = new HashMap<>(propertiesBuilder.buildOrThrow());
    properties.putAll(extraProperties);
    TestingTrinoServer server = TestingTrinoServer.builder().setCoordinator(coordinator).setProperties(properties).setEnvironment(environment).setDiscoveryUri(discoveryUri).setAdditionalModule(additionalModule).setBaseDataDir(baseDataDir).setSystemAccessControls(systemAccessControls).setEventListeners(eventListeners).build();
    String nodeRole = coordinator ? "coordinator" : "worker";
    log.info("Created %s TestingTrinoServer in %s: %s", nodeRole, nanosSince(start).convertToMostSuccinctTimeUnit(), server.getBaseUrl());
    return server;
}

Example 22 with TestingTrinoServer

use of io.trino.server.testing.TestingTrinoServer in project trino by trinodb.

the class TestMemoryWorkerCrash method closeWorker.

private void closeWorker() throws Exception {
    int nodeCount = getNodeCount();
    TestingTrinoServer worker = getDistributedQueryRunner().getServers().stream().filter(server -> !server.isCoordinator()).findAny().orElseThrow(() -> new IllegalStateException("No worker nodes"));
    worker.close();
    waitForNodes(nodeCount - 1);
}

Example 23 with TestingTrinoServer

use of io.trino.server.testing.TestingTrinoServer in project trino by trinodb.

the class TestMemoryManager method testOutOfMemoryKiller.

@Test(timeOut = 240_000, expectedExceptions = ExecutionException.class, expectedExceptionsMessageRegExp = ".*Query killed because the cluster is out of memory. Please try again in a few minutes.")
public void testOutOfMemoryKiller() throws Exception {
    Map<String, String> properties = ImmutableMap.<String, String>builder().put("query.low-memory-killer.delay", "5s").put("query.low-memory-killer.policy", "total-reservation").buildOrThrow();
    try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) {
        // Reserve all the memory
        QueryId fakeQueryId = new QueryId("fake");
        for (TestingTrinoServer server : queryRunner.getServers()) {
            MemoryPool memoryPool = server.getLocalMemoryManager().getMemoryPool();
            assertTrue(memoryPool.tryReserve(fakeQueryId, "test", memoryPool.getMaxBytes()));
        }
        List<Future<?>> queryFutures = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            queryFutures.add(executor.submit(() -> queryRunner.execute("SELECT COUNT(*), clerk FROM orders GROUP BY clerk")));
        }
        // Wait for one of the queries to die
        waitForQueryToBeKilled(queryRunner);
        for (TestingTrinoServer server : queryRunner.getServers()) {
            MemoryPool pool = server.getLocalMemoryManager().getMemoryPool();
            assertTrue(pool.getReservedBytes() > 0);
            // Free up the entire pool
            pool.free(fakeQueryId, "test", pool.getMaxBytes());
            assertTrue(pool.getFreeBytes() > 0);
        }
        for (Future<?> query : queryFutures) {
            query.get();
        }
    }
}

Example 24 with TestingTrinoServer

use of io.trino.server.testing.TestingTrinoServer in project trino by trinodb.

the class TestMemoryManager method testNoLeak.

private void testNoLeak(@Language("SQL") String query) throws Exception {
    Map<String, String> properties = ImmutableMap.<String, String>builder().put("task.verbose-stats", "true").buildOrThrow();
    try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) {
        executor.submit(() -> queryRunner.execute(query)).get();
        for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) {
            assertEquals(info.getState(), FINISHED);
        }
        // Make sure we didn't leak any memory on the workers
        for (TestingTrinoServer worker : queryRunner.getServers()) {
            MemoryPool pool = worker.getLocalMemoryManager().getMemoryPool();
            assertEquals(pool.getMaxBytes(), pool.getFreeBytes());
        }
    }
}

Example 25 with TestingTrinoServer

use of io.trino.server.testing.TestingTrinoServer in project trino by trinodb.

the class TestResourceSecurity method testPasswordAuthenticatorUserMapping.

@Test
public void testPasswordAuthenticatorUserMapping() throws Exception {
    try (TestingTrinoServer server = TestingTrinoServer.builder().setProperties(ImmutableMap.<String, String>builder().putAll(SECURE_PROPERTIES).put("password-authenticator.config-files", passwordConfigDummy.toString()).put("http-server.authentication.type", "password").put("http-server.authentication.password.user-mapping.pattern", ALLOWED_USER_MAPPING_PATTERN).buildOrThrow()).setAdditionalModule(binder -> jaxrsBinder(binder).bind(TestResource.class)).build()) {
        server.getInstance(Key.get(PasswordAuthenticatorManager.class)).setAuthenticators(TestResourceSecurity::authenticate);
        server.getInstance(Key.get(AccessControlManager.class)).addSystemAccessControl(TestSystemAccessControl.WITH_IMPERSONATION);
        HttpServerInfo httpServerInfo = server.getInstance(Key.get(HttpServerInfo.class));
        // Test sets basic auth user and X-Trino-User, and the authenticator is performing user mapping.
        // Normally this would result in an impersonation check to the X-Trino-User, but the password
        // authenticator has a hack to clear X-Trino-User in this case.
        Request request = new Request.Builder().url(getLocation(httpServerInfo.getHttpsUri(), "/protocol/identity")).addHeader("Authorization", Credentials.basic(TEST_USER_LOGIN, TEST_PASSWORD)).addHeader("X-Trino-User", TEST_USER_LOGIN).build();
        try (Response response = client.newCall(request).execute()) {
            assertEquals(response.code(), SC_OK);
            assertEquals(response.header("user"), TEST_USER);
            assertEquals(response.header("principal"), TEST_USER_LOGIN);
        }
    }
}