Java 类com.facebook.presto.client.ClientSession 实例源码

public ClientSession toClientSession(URI server, boolean debug, Duration clientRequestTimeout)
{
    ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
    properties.putAll(systemProperties);
    for (Entry<String, Map<String, String>> catalogProperties : this.catalogProperties.entrySet()) {
        String catalog = catalogProperties.getKey();
        for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
            properties.put(catalog + "." + entry.getKey(), entry.getValue());
        }
    }

    return new ClientSession(
            requireNonNull(server, "server is null"),
            identity.getUser(),
            source.orElse(null),
            catalog.orElse(null),
            schema.orElse(null),
            timeZoneKey.getId(),
            locale,
            properties.build(),
            transactionId.map(TransactionId::toString).orElse(null),
            debug,
            clientRequestTimeout);
}

@POST
@Produces(MediaType.APPLICATION_JSON)
public Response createQuery(String query, @Context HttpServletRequest servletRequest)
{
    assertRequest(!isNullOrEmpty(query), "SQL query is empty");

    Session session = createSessionForRequest(servletRequest, accessControl, sessionPropertyManager, queryIdGenerator.createNextQueryId());
    ClientSession clientSession = session.toClientSession(serverUri(), false, new Duration(2, MINUTES));

    StatementClient client = new StatementClient(httpClient, queryResultsCodec, clientSession, query);

    List<Column> columns = getColumns(client);
    Iterator<List<Object>> iterator = flatten(new ResultsPageIterator(client));
    SimpleQueryResults results = new SimpleQueryResults(columns, iterator);

    return Response.ok(results, MediaType.APPLICATION_JSON_TYPE).build();
}

StatementClient startQuery(String sql)
{
    URI uri = createHttpUri(address);

    String source = firstNonNull(clientInfo.get("ApplicationName"), "presto-jdbc");

    ClientSession session = new ClientSession(
            uri,
            user,
            source,
            catalog.get(),
            schema.get(),
            timeZoneId.get(),
            locale.get(),
            ImmutableMap.copyOf(sessionProperties),
            transactionId.get(),
            false,
            new Duration(2, MINUTES));

    return queryExecutor.startQuery(session, sql);
}

public ParallelQueryRunner(int maxParallelism, URI server, String catalog, String schema, boolean debug, int timeout, Duration clientRequestTimeout)
{
    executor = listeningDecorator(newCachedThreadPool(daemonThreadsNamed("query-runner-%s")));

    ImmutableList.Builder<QueryRunner> runners = ImmutableList.builder();
    for (int i = 0; i < maxParallelism; i++) {
        ClientSession session = new ClientSession(
                server,
                "test-" + i,
                "presto-perf",
                catalog,
                schema,
                TimeZone.getDefault().getID(),
                Locale.getDefault(),
                ImmutableMap.<String, String>of(),
                null,
                debug,
                clientRequestTimeout);
        runners.add(new QueryRunner(session, executor, timeout));
    }
    this.runners = runners.build();
}

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();
}

public QueryRunner(
        ClientSession session,
        JsonCodec<QueryResults> queryResultsCodec,
        Optional<HostAndPort> socksProxy,
        Optional<String> keystorePath,
        Optional<String> keystorePassword,
        Optional<String> kerberosPrincipal,
        Optional<String> kerberosRemoteServiceName,
        boolean authenticationEnabled,
        KerberosConfig kerberosConfig)
{
    this.session = new AtomicReference<>(requireNonNull(session, "session is null"));
    this.queryResultsCodec = requireNonNull(queryResultsCodec, "queryResultsCodec is null");
    this.httpClient = new JettyHttpClient(
            getHttpClientConfig(socksProxy, keystorePath, keystorePassword, kerberosPrincipal, kerberosRemoteServiceName, authenticationEnabled),
            kerberosConfig,
            Optional.<JettyIoPool>empty(),
            ImmutableList.<HttpRequestFilter>of());
}

public static QueryRunner create(
        ClientSession session,
        Optional<HostAndPort> socksProxy,
        Optional<String> keystorePath,
        Optional<String> keystorePassword,
        Optional<String> kerberosPrincipal,
        Optional<String> kerberosRemoteServiceName,
        boolean authenticationEnabled,
        KerberosConfig kerberosConfig)
{
    return new QueryRunner(
            session,
            jsonCodec(QueryResults.class),
            socksProxy,
            keystorePath,
            keystorePassword,
            kerberosPrincipal,
            kerberosRemoteServiceName,
            authenticationEnabled,
            kerberosConfig);
}

private StatementStats execute(ClientSession session, String name, String query)
{
    // start query
    StatementClient client = new StatementClient(httpClient, queryResultsCodec, session, query);

    // read query output
    while (client.isValid() && client.advance()) {
        // we do not process the output
    }

    // verify final state
    if (client.isClosed()) {
        throw new IllegalStateException("Query aborted by user");
    }

    if (client.isGone()) {
        throw new IllegalStateException("Query is gone (server restarted?)");
    }

    QueryError resultsError = client.finalResults().getError();
    if (resultsError != null) {
        RuntimeException cause = null;
        if (resultsError.getFailureInfo() != null) {
            cause = resultsError.getFailureInfo().toException();
        }

        throw new BenchmarkDriverExecutionException(format("Query %s failed: %s", name, resultsError.getMessage()), cause);
    }

    return client.finalResults().getStats();
}

public BenchmarkDriver(BenchmarkResultsStore resultsStore,
        ClientSession clientSession,
        Iterable<BenchmarkQuery> queries,
        int warm,
        int runs,
        boolean debug,
        int maxFailures,
        Optional<HostAndPort> socksProxy)
{
    this.resultsStore = requireNonNull(resultsStore, "resultsStore is null");
    this.clientSession = requireNonNull(clientSession, "clientSession is null");
    this.queries = ImmutableList.copyOf(requireNonNull(queries, "queries is null"));

    queryRunner = new BenchmarkQueryRunner(warm, runs, debug, maxFailures, clientSession.getServer(), socksProxy);
}

public void run(Suite suite)
        throws Exception
{
    // select queries to run
    List<BenchmarkQuery> queries = suite.selectQueries(this.queries);
    if (queries.isEmpty()) {
        return;
    }

    ClientSession session = ClientSession.withSessionProperties(clientSession, suite.getSessionProperties());

    // 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 : queries) {
            session = ClientSession.withCatalogAndSchema(session, session.getCatalog(), benchmarkSchema.getName());
            BenchmarkQueryResult result = queryRunner.execute(suite, session, benchmarkQuery);

            resultsStore.store(benchmarkSchema, result);
        }
    }
}

public ClientSession getClientSession()
{
    return new ClientSession(
            parseServer(server),
            user,
            "presto-benchmark",
            catalog,
            schema,
            TimeZone.getDefault().getID(),
            Locale.getDefault(),
            toProperties(this.sessionProperties),
            null,
            debug,
            clientRequestTimeout);
}

public QueryRunner(ClientSession session, ListeningExecutorService executor, int timeout)
{
    this.session = session;
    this.executor = executor;

    HttpClientConfig clientConfig = new HttpClientConfig();
    clientConfig.setConnectTimeout(new Duration(10, TimeUnit.SECONDS));
    clientConfig.setIdleTimeout(new Duration(timeout, TimeUnit.SECONDS));
    httpClient = new JettyHttpClient(clientConfig);
}

@Override
public void run()
{
    ClientSession session = clientOptions.toClientSession();
    KerberosConfig kerberosConfig = clientOptions.toKerberosConfig();
    boolean hasQuery = !Strings.isNullOrEmpty(clientOptions.execute);
    boolean isFromFile = !Strings.isNullOrEmpty(clientOptions.file);

    if (!hasQuery || !isFromFile) {
        AnsiConsole.systemInstall();
    }

    initializeLogging(clientOptions.logLevelsFile);

    String query = clientOptions.execute;
    if (hasQuery) {
        query += ";";
    }

    if (isFromFile) {
        if (hasQuery) {
            throw new RuntimeException("both --execute and --file specified");
        }
        try {
            query = Files.toString(new File(clientOptions.file), UTF_8);
            hasQuery = true;
        }
        catch (IOException e) {
            throw new RuntimeException(format("Error reading from file %s: %s", clientOptions.file, e.getMessage()));
        }
    }

    try (QueryRunner queryRunner = QueryRunner.create(
            session,
            Optional.ofNullable(clientOptions.socksProxy),
            Optional.ofNullable(clientOptions.keystorePath),
            Optional.ofNullable(clientOptions.keystorePassword),
            Optional.ofNullable(clientOptions.krb5Principal),
            Optional.ofNullable(clientOptions.krb5RemoteServiceName),
            clientOptions.authenticationEnabled,
            kerberosConfig)) {
        if (hasQuery) {
            executeCommand(queryRunner, query, clientOptions.outputFormat);
        }
        else {
            runConsole(queryRunner, 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;
}

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();
        }
    }
}

public ClientSession toClientSession()
{
    return new ClientSession(
            parseServer(server),
            user,
            source,
            catalog,
            schema,
            TimeZone.getDefault().getID(),
            Locale.getDefault(),
            toProperties(sessionProperties),
            null,
            debug,
            clientRequestTimeout);
}

@Test
public void testDefault()
{
    ClientSession session = new ClientOptions().toClientSession();
    assertEquals(session.getServer().toString(), "http://localhost:8080");
    assertEquals(session.getSource(), "presto-cli");
}

@Test
public void testSource()
{
    ClientOptions options = new ClientOptions();
    options.source = "test";
    ClientSession session = options.toClientSession();
    assertEquals(session.getSource(), "test");
}

@Test
public void testServerHostOnly()
{
    ClientOptions options = new ClientOptions();
    options.server = "localhost";
    ClientSession session = options.toClientSession();
    assertEquals(session.getServer().toString(), "http://localhost:80");
}

@Test
public void testServerHostPort()
{
    ClientOptions options = new ClientOptions();
    options.server = "localhost:8888";
    ClientSession session = options.toClientSession();
    assertEquals(session.getServer().toString(), "http://localhost:8888");
}

@Test
public void testServerHttpUri()
{
    ClientOptions options = new ClientOptions();
    options.server = "http://localhost/foo";
    ClientSession session = options.toClientSession();
    assertEquals(session.getServer().toString(), "http://localhost/foo");
}

@Test
public void testServerHttpsUri()
{
    ClientOptions options = new ClientOptions();
    options.server = "https://localhost/foo";
    ClientSession session = options.toClientSession();
    assertEquals(session.getServer().toString(), "https://localhost/foo");
}

@Test
public void testUpdateSessionParameters()
        throws Exception
{
    ClientOptions options = new ClientOptions();
    ClientSession session = options.toClientSession();
    SqlParser sqlParser = new SqlParser();

    ImmutableMap<String, String> existingProperties = ImmutableMap.of("query_max_memory", "10GB", "distributed_join", "true");
    session = Console.processSessionParameterChange(sqlParser.createStatement("USE test_catalog.test_schema"), session, existingProperties);
    assertEquals(session.getCatalog(), "test_catalog");
    assertEquals(session.getSchema(), "test_schema");
    assertEquals(session.getProperties().get("query_max_memory"), "10GB");
    assertEquals(session.getProperties().get("distributed_join"), "true");

    session = Console.processSessionParameterChange(sqlParser.createStatement("USE test_schema_b"), session, existingProperties);
    assertEquals(session.getCatalog(), "test_catalog");
    assertEquals(session.getSchema(), "test_schema_b");
    assertEquals(session.getProperties().get("query_max_memory"), "10GB");
    assertEquals(session.getProperties().get("distributed_join"), "true");

    session = Console.processSessionParameterChange(sqlParser.createStatement("USE test_catalog_2.test_schema"), session, existingProperties);
    assertEquals(session.getCatalog(), "test_catalog_2");
    assertEquals(session.getSchema(), "test_schema");
    assertEquals(session.getProperties().get("query_max_memory"), "10GB");
    assertEquals(session.getProperties().get("distributed_join"), "true");
}

@Test
public void testAutoCompleteWithoutSchema()
{
    ClientSession session = new ClientOptions().toClientSession();
    QueryRunner runner = QueryRunner.create(session,
            Optional.<HostAndPort>empty(),
            Optional.<String>empty(),
            Optional.<String>empty(),
            Optional.<String>empty(),
            Optional.<String>empty(),
            false,
            null);
    TableNameCompleter completer = new TableNameCompleter(runner);
    assertEquals(completer.complete("SELECT is_infi", 14, ImmutableList.of()), 7);
}

@Inject
public PrestoRakamRaptorMetastore(
        @Named("presto.metastore.jdbc") JDBCPoolDataSource prestoMetastoreDataSource,
        EventBus eventBus,
        ProjectConfig projectConfig,
        PrestoConfig prestoConfig) {
    super(eventBus);
    dbi = new DBI(prestoMetastoreDataSource);
    dbi.registerMapper(new TableColumn.Mapper(new SignatureReferenceTypeManager()));
    this.dao = onDemandDao(dbi, MetadataDao.class);
    this.projectConfig = projectConfig;
    this.prestoConfig = prestoConfig;
    defaultSession = new ClientSession(
            prestoConfig.getAddress(),
            "rakam",
            "api-server",
            ImmutableSet.of(), null,
            prestoConfig.getColdStorageConnector(),
            "default",
            TimeZone.getTimeZone(ZoneOffset.UTC).getID(),
            ENGLISH,
            ImmutableMap.of(),
            ImmutableMap.of(),
            null, false, Duration.succinctDuration(1, MINUTES));

    activeNodeCount = Suppliers.memoizeWithExpiration(() -> {
        String getNodeCount = "select count(*) from system.runtime.nodes where state = 'active'";
        QueryResult queryResult = new PrestoQueryExecution(defaultSession, getNodeCount, false).getResult().join();
        if (queryResult.isFailed()) {
            throw new RakamException(queryResult.getError().message, SERVICE_UNAVAILABLE);
        }

        return Ints.checkedCast((long) queryResult.getResult().get(0).get(0));
    }, 5, MINUTES);
}

public ClientSession createSession(String catalog, ZoneId timezone, Map<String, String> sessionProperties, String user) {
    return new ClientSession(
            prestoConfig.getAddress(),
            user == null ? "rakam" : user,
            "rakam",
            ImmutableSet.of(), null,
            catalog == null ? "default" : catalog,
            "default",
            TimeZone.getTimeZone(timezone == null ? ZoneOffset.UTC : timezone).getID(),
            Locale.ENGLISH,
            sessionProperties,
            ImmutableMap.of(), null, false, new Duration(1, TimeUnit.MINUTES));
}

public ClientSession create(String user, String schema)
{
    return new ClientSession(server.get(),
            user,
            source,
            catalog,
            schema,
            timeZoneId,
            locale,
            ImmutableMap.<String, String>of(),
            null,
            debug,
            clientSessionTimeout
    );
}

public ClientSession create(String schema)
{
    return new ClientSession(server.get(),
            user,
            source,
            catalog,
            schema,
            timeZoneId,
            locale,
            ImmutableMap.<String, String>of(),
            null,
            debug,
            clientSessionTimeout
    );
}

public ClientSession create()
{
    return new ClientSession(server.get(),
            user,
            source,
            catalog,
            defaultSchema,
            timeZoneId,
            locale,
            ImmutableMap.<String, String>of(),
            null,
            debug,
            clientSessionTimeout
    );
}

public List<String> getSchemas(ClientSession session)
{
    failures = 0;
    while (true) {
        // start query
        StatementClient client = new StatementClient(httpClient, queryResultsCodec, session, "show schemas");

        // read query output
        ImmutableList.Builder<String> schemas = ImmutableList.builder();
        while (client.isValid() && client.advance()) {
            // we do not process the output
            Iterable<List<Object>> data = client.current().getData();
            if (data != null) {
                for (List<Object> objects : data) {
                    schemas.add(objects.get(0).toString());
                }
            }
        }

        // verify final state
        if (client.isClosed()) {
            throw new IllegalStateException("Query aborted by user");
        }

        if (client.isGone()) {
            throw new IllegalStateException("Query is gone (server restarted?)");
        }

        QueryError resultsError = client.finalResults().getError();
        if (resultsError != null) {
            RuntimeException cause = null;
            if (resultsError.getFailureInfo() != null) {
                cause = resultsError.getFailureInfo().toException();
            }
            handleFailure(cause);

            continue;
        }

        return schemas.build();
    }
}

public StatementClient startQuery(ClientSession session, String query)
{
    return new StatementClient(httpClient, queryInfoCodec, session, query);
}

public ClientSession getSession()
{
    return session.get();
}

public void setSession(ClientSession session)
{
    this.session.set(requireNonNull(session, "session is null"));
}

public PrestoQueryExecution internalExecuteRawQuery(RequestContext context, String query, ClientSession clientSession, boolean update) {
    return new PrestoQueryExecution(clientSession, query, update);
}

protected QueryRunner(ClientSession session, JsonCodec<QueryResults> queryResultsCodec, HttpClient httpClient)
{
    this.session = checkNotNull(session, "session is null");
    this.queryResultsCodec = checkNotNull(queryResultsCodec, "queryResultsCodec is null");
    this.httpClient = httpClient;
}