Examples with ClientSession io.prestosql.client.ClientSession used on opensource projects

Example 6 with ClientSession

use of io.prestosql.client.ClientSession in project hetu-core by openlookeng.

the class CubeConsole method processCubeInitialQuery.

private static boolean processCubeInitialQuery(QueryRunner queryRunner, String sql, ClientOptions.OutputFormat outputFormat, Runnable schemaChanged, boolean usePager, boolean showProgress, Terminal terminal, PrintStream out, PrintStream errorChannel) {
    String finalSql;
    try {
        finalSql = preprocessQuery(Optional.ofNullable(queryRunner.getSession().getCatalog()), Optional.ofNullable(queryRunner.getSession().getSchema()), sql);
    } catch (QueryPreprocessorException e) {
        System.err.println(e.getMessage());
        if (queryRunner.isDebug()) {
            e.printStackTrace(System.err);
        }
        return false;
    }
    try (CubeQuery query = queryRunner.startCubeQuery(finalSql)) {
        boolean success = query.renderCubeOutput(terminal, out, errorChannel, outputFormat, usePager, showProgress);
        if (success) {
            setResultInitCubeQuery(query.getcubeInitQueryResult());
        }
        ClientSession session = queryRunner.getSession();
        // update catalog and schema if present
        if (query.getSetCatalog().isPresent() || query.getSetSchema().isPresent()) {
            session = ClientSession.builder(session).withCatalog(query.getSetCatalog().orElse(session.getCatalog())).withSchema(query.getSetSchema().orElse(session.getSchema())).build();
        }
        // update transaction ID if necessary
        if (query.isClearTransactionId()) {
            session = stripTransactionId(session);
        }
        ClientSession.Builder builder = ClientSession.builder(session);
        if (query.getStartedTransactionId() != null) {
            builder = builder.withTransactionId(query.getStartedTransactionId());
        }
        // update path if present
        if (query.getSetPath().isPresent()) {
            builder = builder.withPath(query.getSetPath().get());
        }
        // update session properties if present
        if (!query.getSetSessionProperties().isEmpty() || !query.getResetSessionProperties().isEmpty()) {
            Map<String, String> sessionProperties = new HashMap<>(session.getProperties());
            sessionProperties.putAll(query.getSetSessionProperties());
            sessionProperties.keySet().removeAll(query.getResetSessionProperties());
            builder = builder.withProperties(sessionProperties);
        }
        // update session roles
        if (!query.getSetRoles().isEmpty()) {
            Map<String, ClientSelectedRole> roles = new HashMap<>(session.getRoles());
            roles.putAll(query.getSetRoles());
            builder = builder.withRoles(roles);
        }
        // update prepared statements if present
        if (!query.getAddedPreparedStatements().isEmpty() || !query.getDeallocatedPreparedStatements().isEmpty()) {
            Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());
            preparedStatements.putAll(query.getAddedPreparedStatements());
            preparedStatements.keySet().removeAll(query.getDeallocatedPreparedStatements());
            builder = builder.withPreparedStatements(preparedStatements);
        }
        session = builder.build();
        queryRunner.setSession(session);
        if (query.getSetCatalog().isPresent() || query.getSetSchema().isPresent()) {
            schemaChanged.run();
        }
        return success;
    } catch (RuntimeException e) {
        System.err.println("Error running command: " + e.getMessage());
        if (queryRunner.isDebug()) {
            e.printStackTrace(System.err);
        }
        return false;
    }
}

Example 7 with ClientSession

use of io.prestosql.client.ClientSession in project hetu-core by openlookeng.

the class PrestoBenchmarkDriver method run.

protected void run(String[] args) throws Exception {
    PrestoBenchmarkDriver prestoBenchmarkDriver = singleCommand(PrestoBenchmarkDriver.class).parse(args);
    if (prestoBenchmarkDriver.helpOption.showHelpIfRequested()) {
        return;
    }
    BenchmarkDriverOptions driverOptions = prestoBenchmarkDriver.benchmarkDriverOptions;
    initializeLogging(driverOptions.debug);
    // select suites
    List<Suite> suites = Suite.readSuites(new File(driverOptions.suiteConfigFile));
    if (!driverOptions.suites.isEmpty()) {
        suites = suites.stream().filter(suite -> driverOptions.suites.contains(suite.getName())).collect(Collectors.toList());
    }
    suites = ImmutableList.copyOf(suites);
    // load queries
    File queriesDir = new File(driverOptions.sqlTemplateDir);
    List<BenchmarkQuery> allQueries = readQueries(queriesDir);
    // select queries to run
    Set<BenchmarkQuery> queries;
    if (driverOptions.queries.isEmpty()) {
        queries = suites.stream().map(suite -> suite.selectQueries(allQueries)).flatMap(List::stream).collect(Collectors.toSet());
    } else {
        queries = driverOptions.queries.stream().map(Pattern::compile).map(pattern -> allQueries.stream().filter(query -> pattern.matcher(query.getName()).matches())).flatMap(identity()).collect(Collectors.toSet());
    }
    // create results store
    BenchmarkResultsStore resultsStore = getResultsStore(suites, queries);
    // create session
    ClientSession session = driverOptions.getClientSession();
    try (BenchmarkDriver benchmarkDriver = new BenchmarkDriver(resultsStore, session, queries, driverOptions.warm, driverOptions.runs, driverOptions.debug, driverOptions.maxFailures, Optional.ofNullable(driverOptions.socksProxy))) {
        for (Suite suite : suites) {
            benchmarkDriver.run(suite);
        }
    }
}

Example 8 with ClientSession

use of io.prestosql.client.ClientSession in project hetu-core by openlookeng.

the class PrestoConnection method startQuery.

StatementClient startQuery(String sql, Map<String, String> sessionPropertiesOverride) {
    String source = "presto-jdbc";
    String applicationName = clientInfo.get("ApplicationName");
    if (applicationNamePrefix.isPresent()) {
        source = applicationNamePrefix.get();
        if (applicationName != null) {
            source += applicationName;
        }
    } else if (applicationName != null) {
        source = applicationName;
    }
    Optional<String> traceToken = Optional.ofNullable(clientInfo.get("TraceToken"));
    Iterable<String> clientTags = Splitter.on(',').trimResults().omitEmptyStrings().split(nullToEmpty(clientInfo.get("ClientTags")));
    Map<String, String> allProperties = new HashMap<>(sessionProperties);
    allProperties.putAll(sessionPropertiesOverride);
    // zero means no timeout, so use a huge value that is effectively unlimited
    int millis = networkTimeoutMillis.get();
    Duration timeout = (millis > 0) ? new Duration(millis, MILLISECONDS) : new Duration(999, DAYS);
    ClientSession session = new ClientSession(httpUri, user, source, traceToken, ImmutableSet.copyOf(clientTags), clientInfo.get("ClientInfo"), catalog.get(), schema.get(), path.get(), timeZoneId.get(), locale.get(), ImmutableMap.of(), ImmutableMap.copyOf(allProperties), ImmutableMap.copyOf(preparedStatements), ImmutableMap.copyOf(roles), extraCredentials, transactionId.get(), timeout);
    return queryExecutor.startQuery(session, sql);
}

Example 9 with ClientSession

use of io.prestosql.client.ClientSession in project hetu-core by openlookeng.

the class ExecutionClient method scheduleExecution.

private UUID scheduleExecution(Duration timeout, QueryRunner queryRunner, QueryExecutionAuthorizer authorizer, BlockingQueue<Job> jobs, URI requestUri) {
    final Job job;
    try {
        job = jobs.take();
    } catch (InterruptedException e) {
        return null;
    }
    final Execution execution = new Execution(job, queryRunner, queryInfoClient, authorizer, timeout, outputBuilderFactory, persistorFactory, requestUri);
    executionMap.put(job.getUuid(), execution);
    activeJobsStore.jobStarted(job);
    ListenableFuture<Job> result = executor.submit(execution);
    Futures.addCallback(result, new FutureCallback<Job>() {

        @Override
        public void onSuccess(@Nullable Job result) {
            if (result != null) {
                result.setState(JobState.FINISHED);
            }
            // Add Active Job
            if (jobs.peek() != null) {
                QueryRunner nextQueryRunner = getNextQueryRunner();
                scheduleExecution(timeout, nextQueryRunner, authorizer, jobs, requestUri);
            }
            jobFinished(result);
        }

        // Re-Use session level fields among multi statement queries.
        private QueryRunner getNextQueryRunner() {
            StatementClient client = queryRunner.getCurrentClient();
            ClientSession session = queryRunner.getSession();
            ClientSession.Builder builder = ClientSession.builder(session).withoutTransactionId();
            if (client.getSetCatalog().isPresent()) {
                builder.withCatalog(client.getSetCatalog().get());
            }
            if (client.getSetSchema().isPresent()) {
                builder.withSchema(client.getSetSchema().get());
            }
            if (client.getStartedTransactionId() != null) {
                builder = builder.withTransactionId(client.getStartedTransactionId());
            }
            if (client.getSetPath().isPresent()) {
                builder = builder.withPath(client.getSetPath().get());
            }
            if (!client.getSetSessionProperties().isEmpty() || !client.getResetSessionProperties().isEmpty()) {
                Map<String, String> sessionProperties = new HashMap<>(session.getProperties());
                sessionProperties.putAll(client.getSetSessionProperties());
                sessionProperties.keySet().removeAll(client.getResetSessionProperties());
                builder = builder.withProperties(sessionProperties);
            }
            if (!client.getSetRoles().isEmpty()) {
                Map<String, ClientSelectedRole> roles = new HashMap<>(session.getRoles());
                roles.putAll(client.getSetRoles());
                builder = builder.withRoles(roles);
            }
            if (!client.getAddedPreparedStatements().isEmpty() || !client.getDeallocatedPreparedStatements().isEmpty()) {
                Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());
                preparedStatements.putAll(client.getAddedPreparedStatements());
                preparedStatements.keySet().removeAll(client.getDeallocatedPreparedStatements());
                builder = builder.withPreparedStatements(preparedStatements);
            }
            return queryRunnerFactory.create(builder.build());
        }

        @Override
        public void onFailure(@NotNull Throwable t) {
            job.setState(JobState.FAILED);
            if (job.getError() == null) {
                job.setError(new QueryError(t.getMessage(), null, -1, null, null, null, null, null));
            }
            jobFinished(job);
        }
    }, MoreExecutors.directExecutor());
    return job.getUuid();
}

Example 10 with ClientSession

use of io.prestosql.client.ClientSession in project hetu-core by openlookeng.

the class BenchmarkDriver method run.

public void run(Suite suite) {
    // select queries to run
    List<BenchmarkQuery> benchmarkQueries = suite.selectQueries(this.queries);
    if (benchmarkQueries.isEmpty()) {
        return;
    }
    Map<String, String> properties = new HashMap<>();
    properties.putAll(clientSession.getProperties());
    properties.putAll(suite.getSessionProperties());
    ClientSession session = ClientSession.builder(clientSession).withProperties(properties).build();
    // select schemas to use
    List<BenchmarkSchema> benchmarkSchemas;
    if (!suite.getSchemaNameTemplates().isEmpty()) {
        List<String> schemas = queryRunner.getSchemas(session);
        benchmarkSchemas = suite.selectSchemas(schemas);
    } else {
        benchmarkSchemas = ImmutableList.of(new BenchmarkSchema(session.getSchema()));
    }
    if (benchmarkSchemas.isEmpty()) {
        return;
    }
    for (BenchmarkSchema benchmarkSchema : benchmarkSchemas) {
        for (BenchmarkQuery benchmarkQuery : benchmarkQueries) {
            session = ClientSession.builder(session).withCatalog(session.getCatalog()).withSchema(benchmarkSchema.getName()).build();
            BenchmarkQueryResult result = queryRunner.execute(suite, session, benchmarkQuery);
            resultsStore.store(benchmarkSchema, result);
        }
    }
}