public ClientSession build() { return new ClientSession( server, user, source, traceToken, clientTags, clientInfo, catalog, schema, path, timeZone.getId(), locale, resourceEstimates, properties, preparedStatements, transactionId, clientRequestTimeout); } }
private Builder(ClientSession clientSession) { requireNonNull(clientSession, "clientSession is null"); server = clientSession.getServer(); user = clientSession.getUser(); source = clientSession.getSource(); traceToken = clientSession.getTraceToken(); clientTags = clientSession.getClientTags(); clientInfo = clientSession.getClientInfo(); catalog = clientSession.getCatalog(); schema = clientSession.getSchema(); path = clientSession.getPath(); timeZone = clientSession.getTimeZone(); locale = clientSession.getLocale(); resourceEstimates = clientSession.getResourceEstimates(); properties = clientSession.getProperties(); preparedStatements = clientSession.getPreparedStatements(); transactionId = clientSession.getTransactionId(); clientRequestTimeout = clientSession.getClientRequestTimeout(); }
try { finalSql = preprocessQuery( Optional.ofNullable(queryRunner.getSession().getCatalog()), Optional.ofNullable(queryRunner.getSession().getSchema()), sql); session = ClientSession.builder(session) .withCatalog(query.getSetCatalog().orElse(session.getCatalog())) .withSchema(query.getSetSchema().orElse(session.getSchema())) .build(); schemaChanged.run(); session = stripTransactionId(session); ClientSession.Builder builder = ClientSession.builder(session); Map<String, String> sessionProperties = new HashMap<>(session.getProperties()); sessionProperties.putAll(query.getSetSessionProperties()); sessionProperties.keySet().removeAll(query.getResetSessionProperties()); Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements()); preparedStatements.putAll(query.getAddedPreparedStatements()); preparedStatements.keySet().removeAll(query.getDeallocatedPreparedStatements());
public StatementClientV1(OkHttpClient httpClient, ClientSession session, String query) { requireNonNull(httpClient, "httpClient is null"); requireNonNull(session, "session is null"); requireNonNull(query, "query is null"); this.httpClient = httpClient; this.timeZone = session.getTimeZone(); this.query = query; this.requestTimeoutNanos = session.getClientRequestTimeout(); this.user = session.getUser(); this.clientCapabilities = Joiner.on(",").join(ClientCapabilities.values()); Request request = buildQueryRequest(session, query); JsonResponse<QueryResults> response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request); if ((response.getStatusCode() != HTTP_OK) || !response.hasValue()) { state.compareAndSet(State.RUNNING, State.CLIENT_ERROR); throw requestFailedException("starting query", request, response); } processResponse(response.getHeaders(), response.getValue()); }
public static ClientSession stripTransactionId(ClientSession session) { return new ClientSession( session.getServer(), session.getUser(), session.getSource(), session.getCatalog(), session.getSchema(), session.getTimeZoneId(), session.getLocale(), session.getProperties(), null, session.isDebug(), session.getClientRequestTimeout()); }
private static Request buildQueryRequest(ClientSession session, String query) { Request.Builder builder = preparePost() .setUri(uriBuilderFrom(session.getServer()).replacePath("/v1/statement").build()) .setBodyGenerator(createStaticBodyGenerator(query, UTF_8)); if (session.getUser() != null) { builder.setHeader(PrestoHeaders.PRESTO_USER, session.getUser()); } if (session.getSource() != null) { builder.setHeader(PrestoHeaders.PRESTO_SOURCE, session.getSource()); } if (session.getCatalog() != null) { builder.setHeader(PrestoHeaders.PRESTO_CATALOG, session.getCatalog()); } if (session.getSchema() != null) { builder.setHeader(PrestoHeaders.PRESTO_SCHEMA, session.getSchema()); } builder.setHeader(PrestoHeaders.PRESTO_TIME_ZONE, session.getTimeZoneId()); builder.setHeader(PrestoHeaders.PRESTO_LANGUAGE, session.getLocale().toLanguageTag()); builder.setHeader(USER_AGENT, USER_AGENT_VALUE); Map<String, String> property = session.getProperties(); for (Entry<String, String> entry : property.entrySet()) { builder.addHeader(PrestoHeaders.PRESTO_SESSION, entry.getKey() + "=" + entry.getValue()); } builder.setHeader(PrestoHeaders.PRESTO_TRANSACTION_ID, session.getTransactionId() == null ? "NONE" : session.getTransactionId()); return builder.build(); }
public static ClientSession withCatalogAndSchema(ClientSession session, String catalog, String schema) { return new ClientSession( session.getServer(), session.getUser(), session.getSource(), catalog, schema, session.getTimeZoneId(), session.getLocale(), session.getProperties(), session.getTransactionId(), session.isDebug(), session.getClientRequestTimeout()); }
private static Request buildQueryRequest(ClientSession session, String query) { Request.Builder builder = preparePost() .setUri(uriBuilderFrom(session.getServer()).replacePath("/v1/execute").build()) .setBodyGenerator(createStaticBodyGenerator(query, UTF_8)); if (session.getUser() != null) { builder.setHeader(PrestoHeaders.PRESTO_USER, session.getUser()); } if (session.getSource() != null) { builder.setHeader(PrestoHeaders.PRESTO_SOURCE, session.getSource()); } if (session.getCatalog() != null) { builder.setHeader(PrestoHeaders.PRESTO_CATALOG, session.getCatalog()); } if (session.getSchema() != null) { builder.setHeader(PrestoHeaders.PRESTO_SCHEMA, session.getSchema()); } builder.setHeader(PrestoHeaders.PRESTO_TIME_ZONE, session.getTimeZoneId()); builder.setHeader(USER_AGENT, USER_AGENT_VALUE); return builder.build(); }
properties.putAll(clientSession.getProperties()); properties.putAll(suite.getSessionProperties()); ClientSession session = ClientSession.builder(clientSession) .withProperties(properties) .build(); benchmarkSchemas = ImmutableList.of(new BenchmarkSchema(session.getSchema())); session = ClientSession.builder(session) .withCatalog(session.getCatalog()) .withSchema(benchmarkSchema.getName()) .build();
private static void setupBasicAuth( OkHttpClient.Builder clientBuilder, ClientSession session, Optional<String> user, Optional<String> password) { if (user.isPresent() && password.isPresent()) { checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), "Authentication using username/password requires HTTPS to be enabled"); clientBuilder.addInterceptor(basicAuth(user.get(), password.get())); } }
public void populateCache() { String schemaName = queryRunner.getSession().getSchema(); if (schemaName != null) { executor.execute(() -> { functionCache.refresh(schemaName); tableCache.refresh(schemaName); }); } }
private static void process(QueryRunner queryRunner, String sql, OutputFormat outputFormat, boolean interactive) { try (Query query = queryRunner.startQuery(sql)) { query.renderOutput(System.out, outputFormat, interactive); ClientSession session = queryRunner.getSession(); // 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()); session = withProperties(session, sessionProperties); } // update transaction ID if necessary if (query.isClearTransactionId()) { session = stripTransactionId(session); } if (query.getStartedTransactionId() != null) { session = withTransactionId(session, query.getStartedTransactionId()); } queryRunner.setSession(session); } catch (RuntimeException e) { System.err.println("Error running command: " + e.getMessage()); if (queryRunner.getSession().isDebug()) { e.printStackTrace(); } } }
if (session.getSchema() != null) { prompt += ":" + session.getSchema(); Optional<Object> statement = getParsedStatement(split.statement()); if (statement.isPresent() && isSessionParameterChange(statement.get())) { Map<String, String> properties = queryRunner.getSession().getProperties(); session = processSessionParameterChange(statement.get(), session, properties); queryRunner.setSession(session);
static ClientSession processSessionParameterChange(Object parsedStatement, ClientSession session, Map<String, String> existingProperties) { if (parsedStatement instanceof Use) { Use use = (Use) parsedStatement; session = withCatalogAndSchema(session, use.getCatalog().orElse(session.getCatalog()), use.getSchema()); session = withSessionProperties(session, existingProperties); } return session; }
@Test public void testSource() { ClientOptions options = new ClientOptions(); options.source = "test"; ClientSession session = options.toClientSession(); assertEquals(session.getSource(), "test"); }
public static ClientSession stripTransactionId(ClientSession session) { return ClientSession.builder(session) .withoutTransactionId() .build(); }
public static ClientSession withTransactionId(ClientSession session, String transactionId) { return new ClientSession( session.getServer(), session.getUser(), session.getSource(), session.getCatalog(), session.getSchema(), session.getTimeZoneId(), session.getLocale(), session.getProperties(), transactionId, session.isDebug(), session.getClientRequestTimeout()); }
private static void setupTokenAuth( OkHttpClient.Builder clientBuilder, ClientSession session, Optional<String> accessToken) { if (accessToken.isPresent()) { checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), "Authentication using an access token requires HTTPS to be enabled"); clientBuilder.addInterceptor(tokenAuth(accessToken.get())); } } }
@Override public int complete(String buffer, int cursor, List<CharSequence> candidates) { if (cursor <= 0) { return cursor; } int blankPos = findLastBlank(buffer.substring(0, cursor)); String prefix = buffer.substring(blankPos + 1, cursor); String schemaName = queryRunner.getSession().getSchema(); if (schemaName != null) { List<String> functionNames = functionCache.getIfPresent(schemaName); List<String> tableNames = tableCache.getIfPresent(schemaName); SortedSet<String> sortedCandidates = new TreeSet<>(); if (functionNames != null) { sortedCandidates.addAll(filterResults(functionNames, prefix)); } if (tableNames != null) { sortedCandidates.addAll(filterResults(tableNames, prefix)); } candidates.addAll(sortedCandidates); } return blankPos + 1; }