diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d0a8aa7..b48adaad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## 1.5.0 [unreleased] +### Features + +1. [#289](https://github.com/InfluxCommunity/influxdb3-java/pull/289) Add the possibility to disable gRPC compression via the `disableGRPCCompression` parameter in the `ClientConfig`. + ### CI 1. [#283](https://github.com/InfluxCommunity/influxdb3-java/pull/283) Fix pipeline not downloading the correct java images. diff --git a/src/main/java/com/influxdb/v3/client/config/ClientConfig.java b/src/main/java/com/influxdb/v3/client/config/ClientConfig.java index 70fb4ba8..7e0e8ac7 100644 --- a/src/main/java/com/influxdb/v3/client/config/ClientConfig.java +++ b/src/main/java/com/influxdb/v3/client/config/ClientConfig.java @@ -67,6 +67,7 @@ *
  • authenticator - HTTP proxy authenticator
    • *
  • headers - headers to be added to requests
    • *
  • sslRootsFilePath - path to the stored certificates file in PEM format
    • + *
  • disableGRPCCompression - disables the default gRPC compression header
    • * *

    * If you want to create a client with custom configuration, you can use following code: @@ -113,6 +114,7 @@ public final class ClientConfig { private final Authenticator authenticator; private final Map headers; private final String sslRootsFilePath; + private final boolean disableGRPCCompression; /** * Deprecated use {@link #proxyUrl}. @@ -318,6 +320,15 @@ public Map getHeaders() { return headers; } + /** + * Is gRPC compression disabled. + * + * @return true if gRPC compression is disabled + */ + public boolean getDisableGRPCCompression() { + return disableGRPCCompression; + } + /** * Validates the configuration properties. */ @@ -354,7 +365,8 @@ public boolean equals(final Object o) { && Objects.equals(proxyUrl, that.proxyUrl) && Objects.equals(authenticator, that.authenticator) && Objects.equals(headers, that.headers) - && Objects.equals(sslRootsFilePath, that.sslRootsFilePath); + && Objects.equals(sslRootsFilePath, that.sslRootsFilePath) + && disableGRPCCompression == that.disableGRPCCompression; } @Override @@ -363,7 +375,7 @@ public int hashCode() { database, writePrecision, gzipThreshold, writeNoSync, timeout, writeTimeout, queryTimeout, allowHttpRedirects, disableServerCertificateValidation, proxy, proxyUrl, authenticator, headers, - defaultTags, sslRootsFilePath); + defaultTags, sslRootsFilePath, disableGRPCCompression); } @Override @@ -386,6 +398,7 @@ public String toString() { .add("headers=" + headers) .add("defaultTags=" + defaultTags) .add("sslRootsFilePath=" + sslRootsFilePath) + .add("disableGRPCCompression=" + disableGRPCCompression) .toString(); } @@ -415,6 +428,7 @@ public static final class Builder { private Authenticator authenticator; private Map headers; private String sslRootsFilePath; + private boolean disableGRPCCompression; /** * Sets the URL of the InfluxDB server. @@ -697,6 +711,18 @@ public Builder sslRootsFilePath(@Nullable final String sslRootsFilePath) { return this; } + /** + * Sets whether to disable gRPC compression. Default is 'false'. + * + * @param disableGRPCCompression disable gRPC compression + * @return this + */ + @Nonnull + public Builder disableGRPCCompression(final boolean disableGRPCCompression) { + this.disableGRPCCompression = disableGRPCCompression; + return this; + } + /** * Build an instance of {@code ClientConfig}. * @@ -745,6 +771,9 @@ public ClientConfig build(@Nonnull final String connectionString) throws Malform if (parameters.containsKey("writeNoSync")) { this.writeNoSync(Boolean.parseBoolean(parameters.get("writeNoSync"))); } + if (parameters.containsKey("disableGRPCCompression")) { + this.disableGRPCCompression(Boolean.parseBoolean(parameters.get("disableGRPCCompression"))); + } return new ClientConfig(this); } @@ -807,6 +836,11 @@ public ClientConfig build(@Nonnull final Map env, final Properti long to = Long.parseLong(queryTimeout); this.queryTimeout(Duration.ofSeconds(to)); } + final String disableGRPCCompression = get.apply("INFLUX_DISABLE_GRPC_COMPRESSION", + "influx.disableGRPCCompression"); + if (disableGRPCCompression != null) { + this.disableGRPCCompression(Boolean.parseBoolean(disableGRPCCompression)); + } return new ClientConfig(this); } @@ -862,5 +896,6 @@ private ClientConfig(@Nonnull final Builder builder) { authenticator = builder.authenticator; headers = builder.headers; sslRootsFilePath = builder.sslRootsFilePath; + disableGRPCCompression = builder.disableGRPCCompression; } } diff --git a/src/main/java/com/influxdb/v3/client/internal/FlightSqlClient.java b/src/main/java/com/influxdb/v3/client/internal/FlightSqlClient.java index 697693fb..4337f6b3 100644 --- a/src/main/java/com/influxdb/v3/client/internal/FlightSqlClient.java +++ b/src/main/java/com/influxdb/v3/client/internal/FlightSqlClient.java @@ -41,6 +41,8 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import io.grpc.Codec; +import io.grpc.DecompressorRegistry; import io.grpc.HttpConnectProxiedSocketAddress; import io.grpc.Metadata; import io.grpc.ProxyDetector; @@ -93,8 +95,6 @@ final class FlightSqlClient implements AutoCloseable { defaultHeaders.put("Authorization", "Bearer " + new String(config.getToken())); } - defaultHeaders.put("User-Agent", Identity.getUserAgent()); - if (config.getHeaders() != null) { defaultHeaders.putAll(config.getHeaders()); } @@ -148,6 +148,8 @@ private FlightClient createFlightClient(@Nonnull final ClientConfig config) { URI uri = createLocation(config).getUri(); final NettyChannelBuilder nettyChannelBuilder = NettyChannelBuilder.forAddress(uri.getHost(), uri.getPort()); + nettyChannelBuilder.userAgent(Identity.getUserAgent()); + if (LocationSchemes.GRPC_TLS.equals(uri.getScheme())) { nettyChannelBuilder.useTransportSecurity(); @@ -169,6 +171,11 @@ private FlightClient createFlightClient(@Nonnull final ClientConfig config) { nettyChannelBuilder.maxTraceEvents(0) .maxInboundMetadataSize(Integer.MAX_VALUE); + if (config.getDisableGRPCCompression()) { + nettyChannelBuilder.decompressorRegistry(DecompressorRegistry.emptyInstance() + .with(Codec.Identity.NONE, false)); + } + return FlightGrpcUtils.createFlightClient(new RootAllocator(Long.MAX_VALUE), nettyChannelBuilder.build()); } diff --git a/src/test/java/com/influxdb/v3/client/ITQueryWrite.java b/src/test/java/com/influxdb/v3/client/ITQueryWrite.java index 820f4672..49e9110c 100644 --- a/src/test/java/com/influxdb/v3/client/ITQueryWrite.java +++ b/src/test/java/com/influxdb/v3/client/ITQueryWrite.java @@ -439,6 +439,9 @@ public void queryTimeoutSuperceededByGrpcOptTest() { Assertions.assertThat(thrown.getMessage()).matches(".*deadline.*exceeded.*"); } + @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_URL", matches = ".*") + @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_TOKEN", matches = ".*") + @EnabledIfEnvironmentVariable(named = "TESTING_INFLUXDB_DATABASE", matches = ".*") @Test public void repeatQueryWithTimeoutTest() { long timeout = 1000; diff --git a/src/test/java/com/influxdb/v3/client/TestUtils.java b/src/test/java/com/influxdb/v3/client/TestUtils.java new file mode 100644 index 00000000..d7f2d1c9 --- /dev/null +++ b/src/test/java/com/influxdb/v3/client/TestUtils.java @@ -0,0 +1,92 @@ +/* + * The MIT License + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +package com.influxdb.v3.client; + +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import javax.annotation.Nonnull; + +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.NoOpFlightProducer; +import org.apache.arrow.flight.Ticket; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +public final class TestUtils { + + private TestUtils() { + throw new IllegalStateException("Utility class"); + } + + public static FlightServer simpleFlightServer(@Nonnull final URI uri, + @Nonnull final BufferAllocator allocator, + @Nonnull final NoOpFlightProducer producer) throws Exception { + Location location = Location.forGrpcInsecure(uri.getHost(), uri.getPort()); + return FlightServer.builder(allocator, location, producer).build(); + } + + public static NoOpFlightProducer simpleProducer(@Nonnull final VectorSchemaRoot vectorSchemaRoot) { + return new NoOpFlightProducer() { + @Override + public void getStream(final CallContext context, + final Ticket ticket, + final ServerStreamListener listener) { + listener.start(vectorSchemaRoot); + if (listener.isReady()) { + listener.putNext(); + } + listener.completed(); + } + }; + } + + public static VectorSchemaRoot generateVectorSchemaRoot(final int fieldCount, final int rowCount) { + List fields = new ArrayList<>(); + for (int i = 0; i < fieldCount; i++) { + Field field = new Field("field" + i, FieldType.nullable(new ArrowType.Utf8()), null); + fields.add(field); + } + + Schema schema = new Schema(fields); + VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, new RootAllocator(Long.MAX_VALUE)); + for (Field field : fields) { + VarCharVector vector = (VarCharVector) vectorSchemaRoot.getVector(field); + vector.allocateNew(rowCount); + for (int i = 0; i < rowCount; i++) { + vector.set(i, "Value".getBytes(StandardCharsets.UTF_8)); + } + } + vectorSchemaRoot.setRowCount(rowCount); + + return vectorSchemaRoot; + } +} + diff --git a/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java b/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java index b17443d6..f5ce2e8f 100644 --- a/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java +++ b/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java @@ -46,7 +46,8 @@ class ClientConfigTest { .queryTimeout(Duration.ofSeconds(120)) .allowHttpRedirects(true) .disableServerCertificateValidation(true) - .headers(Map.of("X-device", "ab-01")); + .headers(Map.of("X-device", "ab-01")) + .disableGRPCCompression(true); @Test void equalConfig() { @@ -81,6 +82,7 @@ void toStringConfig() { Assertions.assertThat(configString).contains("timeout=PT30S"); Assertions.assertThat(configString).contains("writeTimeout=PT35S"); Assertions.assertThat(configString).contains("queryTimeout=PT2M"); + Assertions.assertThat(configString).contains("disableGRPCCompression=true"); } @@ -131,10 +133,11 @@ void fromConnectionString() throws MalformedURLException { cfg = new ClientConfig.Builder() .build("http://localhost:9999/" - + "?token=my-token&authScheme=my-auth"); + + "?token=my-token&authScheme=my-auth&disableGRPCCompression=true"); Assertions.assertThat(cfg.getHost()).isEqualTo("http://localhost:9999/"); Assertions.assertThat(cfg.getToken()).isEqualTo("my-token".toCharArray()); Assertions.assertThat(cfg.getAuthScheme()).isEqualTo("my-auth"); + Assertions.assertThat(cfg.getDisableGRPCCompression()).isEqualTo(true); } @Test @@ -204,7 +207,9 @@ void fromEnv() { "INFLUX_DATABASE", "my-db", "INFLUX_PRECISION", "ms", "INFLUX_GZIP_THRESHOLD", "64", - "INFLUX_WRITE_NO_SYNC", "true" + "INFLUX_WRITE_NO_SYNC", "true", + "INFLUX_DISABLE_GRPC_COMPRESSION", "true" + ); cfg = new ClientConfig.Builder() .build(env, null); @@ -215,6 +220,7 @@ void fromEnv() { Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.MS); Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(64); Assertions.assertThat(cfg.getWriteNoSync()).isEqualTo(true); + Assertions.assertThat(cfg.getDisableGRPCCompression()).isTrue(); } @Test @@ -318,6 +324,7 @@ void fromSystemProperties() { properties.put("influx.precision", "ms"); properties.put("influx.gzipThreshold", "64"); properties.put("influx.writeNoSync", "true"); + properties.put("influx.disableGRPCCompression", "true"); cfg = new ClientConfig.Builder() .build(new HashMap<>(), properties); Assertions.assertThat(cfg.getHost()).isEqualTo("http://localhost:9999/"); @@ -327,6 +334,7 @@ void fromSystemProperties() { Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.MS); Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(64); Assertions.assertThat(cfg.getWriteNoSync()).isEqualTo(true); + Assertions.assertThat(cfg.getDisableGRPCCompression()).isTrue(); } @Test diff --git a/src/test/java/com/influxdb/v3/client/internal/FlightSqlClientTest.java b/src/test/java/com/influxdb/v3/client/internal/FlightSqlClientTest.java index 84a609ad..a6873c56 100644 --- a/src/test/java/com/influxdb/v3/client/internal/FlightSqlClientTest.java +++ b/src/test/java/com/influxdb/v3/client/internal/FlightSqlClientTest.java @@ -23,26 +23,32 @@ import java.net.InetSocketAddress; import java.net.URISyntaxException; +import java.util.HashMap; import java.util.Map; +import java.util.stream.IntStream; +import java.util.stream.Stream; import io.grpc.HttpConnectProxiedSocketAddress; import io.grpc.ProxyDetector; import io.grpc.internal.GrpcUtil; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; +import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.Location; -import org.apache.arrow.flight.NoOpFlightProducer; +import org.apache.arrow.flight.RequestContext; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import com.influxdb.v3.client.InfluxDBClient; +import com.influxdb.v3.client.TestUtils; import com.influxdb.v3.client.config.ClientConfig; import com.influxdb.v3.client.query.QueryOptions; import com.influxdb.v3.client.query.QueryType; @@ -50,300 +56,366 @@ public class FlightSqlClientTest { private static final String LOCALHOST = "localhost"; + private static final FlightServerMiddleware.Key HEADER_CAPTURE_KEY = + FlightServerMiddleware.Key.of("header-capture"); private final Location grpcLocation = Location.forGrpcInsecure(LOCALHOST, 0); - private final String serverLocation = String.format("http://%s:%d", LOCALHOST, grpcLocation.getUri().getPort()); + private final HeaderCaptureMiddlewareFactory headerFactory = new HeaderCaptureMiddlewareFactory(); + private final int rowCount = 10; + private final VectorSchemaRoot vectorSchemaRoot = TestUtils.generateVectorSchemaRoot(10, rowCount); - private final CallHeadersMiddleware callHeadersMiddleware = new CallHeadersMiddleware(); private RootAllocator allocator; private FlightServer server; - private FlightClient client; - - @BeforeEach - void reset() { - callHeadersMiddleware.headers = null; - } - - @BeforeEach - void setUp() throws Exception { - allocator = new RootAllocator(Long.MAX_VALUE); - server = FlightServer.builder(allocator, grpcLocation, new NoOpFlightProducer()).build().start(); - client = FlightClient.builder(allocator, server.getLocation()).intercept(callHeadersMiddleware).build(); - callHeadersMiddleware.headers = null; - } - - @AfterEach - void tearDown() throws Exception { - if (client != null) { - client.close(); - } - if (server != null) { - server.shutdown(); - server.awaitTermination(); - } - if (allocator != null) { - allocator.close(); - } - } - - @Test - void flightSqlClient() throws Exception { - String correctHost = "grpc+unix://tmp/dummy.sock"; - ClientConfig clientConfig = new ClientConfig.Builder() - .host(correctHost) - .token("Token".toCharArray()) - .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig)) { - Assertions.assertThat(flightSqlClient).isNotNull(); - } - - FlightClient.Builder builder = FlightClient.builder(allocator, server.getLocation()); - try (FlightClient flightClient = builder.build()) { - FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, flightClient); - Assertions.assertThat(flightSqlClient).isNotNull(); - } - - var inCorrectHost = "grpc+unix://///tmp/dummy.sock"; - ClientConfig clientConfig1 = new ClientConfig.Builder() - .host(inCorrectHost) - .token("Token".toCharArray()) - .build(); - Assertions.assertThatThrownBy(() -> new FlightSqlClient(clientConfig1)); - } @Test public void invalidHost() { ClientConfig clientConfig = new ClientConfig.Builder() - .host("xyz://a bc") - .token("my-token".toCharArray()) - .build(); + .host("xyz://a bc") + .token("my-token".toCharArray()) + .build(); Assertions.assertThatThrownBy(() -> { - try (FlightSqlClient ignored = new FlightSqlClient(clientConfig)) { - Assertions.fail("Should not be here"); - } - }) - .isInstanceOf(RuntimeException.class) - .hasCauseInstanceOf(URISyntaxException.class) - .hasMessageContaining("xyz://a bc"); + try (FlightSqlClient ignored = new FlightSqlClient(clientConfig)) { + Assertions.fail("Should not be here"); + } + }) + .isInstanceOf(RuntimeException.class) + .hasCauseInstanceOf(URISyntaxException.class) + .hasMessageContaining("xyz://a bc"); } @Test public void callHeaders() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("my-token".toCharArray()) - .build(); + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of())) { - flightSqlClient.execute("select * from cpu", "mydb", QueryType.SQL, Map.of(), Map.of()); + Assertions.assertThat(data.count()).isEqualTo(rowCount); - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "authorization", - "user-agent", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + Assertions.assertThat(receivedHeaders.keySet()).contains( + "authorization", + "user-agent", + GrpcUtil.MESSAGE_ACCEPT_ENCODING ); - Assertions.assertThat(incomingHeaders.get("authorization")).isEqualTo("Bearer my-token"); - Assertions.assertThat(incomingHeaders.get("user-agent")).isEqualTo(Identity.getUserAgent()); + Assertions.assertThat(receivedHeaders.get("authorization")).isEqualTo("Bearer my-token"); + Assertions.assertThat(receivedHeaders.get("user-agent")).startsWith(Identity.getUserAgent()); } } @Test public void callHeadersWithoutToken() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .build(); + .host(server.getLocation().getUri().toString()) + .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of())) { - flightSqlClient.execute("select * from cpu", "mydb", QueryType.SQL, Map.of(), Map.of()); + Assertions.assertThat(data.count()).isEqualTo(rowCount); - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "user-agent", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + Assertions.assertThat(receivedHeaders.keySet()).containsOnly( + "user-agent", + GrpcUtil.MESSAGE_ACCEPT_ENCODING, + "content-type" ); - Assertions.assertThat(incomingHeaders.get("authorization")).isNull(); + Assertions.assertThat(receivedHeaders.get("authorization")).isNull(); } } @Test public void callHeadersEmptyToken() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("".toCharArray()) - .build(); + .host(server.getLocation().getUri().toString()) + .token("".toCharArray()) + .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of())) { - flightSqlClient.execute("select * from cpu", "mydb", QueryType.SQL, Map.of(), Map.of()); + Assertions.assertThat(data.count()).isEqualTo(rowCount); - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; - - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "user-agent", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); + Assertions.assertThat(receivedHeaders.keySet()).containsOnly( + "user-agent", + GrpcUtil.MESSAGE_ACCEPT_ENCODING, + "content-type" ); - Assertions.assertThat(incomingHeaders.get("authorization")).isNull(); + Assertions.assertThat(receivedHeaders.get("authorization")).isNull(); } } @Test public void callHeadersCustomHeader() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("my-token".toCharArray()) - .headers(Map.of("X-Tracing-Id", "123")) - .build(); + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .headers(Map.of("X-Tracing-Id", "123")) + .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of())) { - flightSqlClient.execute("select * from cpu", "mydb", QueryType.SQL, Map.of(), Map.of()); + Assertions.assertThat(data.count()).isEqualTo(rowCount); - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "authorization", - "user-agent", - "x-tracing-id", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + Assertions.assertThat(receivedHeaders.keySet()).contains( + "authorization", + "user-agent", + "x-tracing-id", + GrpcUtil.MESSAGE_ACCEPT_ENCODING ); - Assertions.assertThat(incomingHeaders.get("X-Tracing-Id")).isEqualTo("123"); + Assertions.assertThat(receivedHeaders.get("x-tracing-id")).isEqualTo("123"); } } @Test public void customHeaderForRequest() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("my-token".toCharArray()) - .headers(Map.of("X-Tracing-Id", "123")) - .build(); - - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { - - flightSqlClient.execute( - "select * from cpu", - "mydb", - QueryType.SQL, - Map.of(), - Map.of("X-Invoice-Id", "456")); - - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; - - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "authorization", - "user-agent", - "x-tracing-id", - "x-invoice-id", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .headers(Map.of("X-Tracing-Id", "123")) + .build(); + + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of("X-Invoice-Id", "456"))) { + + Assertions.assertThat(data.count()).isEqualTo(rowCount); + + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); + + Assertions.assertThat(receivedHeaders.keySet()).contains( + "authorization", + "user-agent", + "x-tracing-id", + "x-invoice-id", + GrpcUtil.MESSAGE_ACCEPT_ENCODING ); - Assertions.assertThat(incomingHeaders.get("X-Tracing-Id")).isEqualTo("123"); + Assertions.assertThat(receivedHeaders.get("x-invoice-id")).isEqualTo("456"); } } @Test public void customHeaderForRequestOverrideConfig() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("my-token".toCharArray()) - .headers(Map.of("X-Tracing-Id", "123")) - .build(); - - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client)) { - - flightSqlClient.execute( - "select * from cpu", - "mydb", - QueryType.SQL, - Map.of(), - Map.of("X-Tracing-Id", "456")); - - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; - - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "authorization", - "user-agent", - "x-tracing-id", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .headers(Map.of("X-Tracing-Id", "123")) + .build(); + + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of("X-Tracing-Id", "456"))) { + + Assertions.assertThat(data.count()).isEqualTo(rowCount); + + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); + + Assertions.assertThat(receivedHeaders.keySet()).contains( + "authorization", + "user-agent", + "x-tracing-id", + GrpcUtil.MESSAGE_ACCEPT_ENCODING ); - Assertions.assertThat(incomingHeaders.get("X-Tracing-Id")).isEqualTo("456"); + Assertions.assertThat(receivedHeaders.get("x-tracing-id")).isEqualTo("456"); } } @Test public void useParamsFromQueryConfig() throws Exception { ClientConfig clientConfig = new ClientConfig.Builder() - .host(serverLocation) - .token("my-token".toCharArray()) - .database("mydb") - .build(); + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .database("mydb") + .build(); + + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + InfluxDBClient influxDBClient = new InfluxDBClientImpl(clientConfig, null, flightSqlClient); + Stream data = influxDBClient.query( + "select * from cpu", + new QueryOptions(Map.of("X-Tracing-Id", "987")))) { + + Assertions.assertThat(data.count()).isEqualTo(rowCount); + + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); + + Assertions.assertThat(receivedHeaders.keySet()).contains( + "authorization", + "x-tracing-id", + "user-agent", + GrpcUtil.MESSAGE_ACCEPT_ENCODING + ); + Assertions.assertThat(receivedHeaders.get("x-tracing-id")).isEqualTo("987"); + } + } + + @Test + public void disableGRPCCompression() throws Exception { + ClientConfig clientConfig = new ClientConfig.Builder() + .host(server.getLocation().getUri().toString()) + .token("my-token".toCharArray()) + .disableGRPCCompression(true) + .build(); - try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, client); - InfluxDBClient influxDBClient = new InfluxDBClientImpl(clientConfig, null, flightSqlClient)) { + var qopts = new GrpcCallOptions.Builder().withCompressorName("identity").build(); + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig); + var data = executeQuery(flightSqlClient, Map.of(), Map.of(), qopts.getCallOptions())) { - influxDBClient.query("select * from cpu", new QueryOptions(Map.of("X-Tracing-Id", "987"))); + Assertions.assertThat(data.count()).isEqualTo(rowCount); - final CallHeaders incomingHeaders = callHeadersMiddleware.headers; + final Map receivedHeaders = headerFactory.getLastInstance().getHeaders(); - Assertions.assertThat(incomingHeaders.keys()).containsOnly( - "authorization", - "x-tracing-id", - "user-agent", - GrpcUtil.MESSAGE_ACCEPT_ENCODING + Assertions.assertThat(receivedHeaders.keySet()).containsOnly( + "authorization", + "user-agent", + "content-type" ); - Assertions.assertThat(incomingHeaders.get("X-Tracing-Id")).isEqualTo("987"); + Assertions.assertThat(receivedHeaders.get(GrpcUtil.MESSAGE_ACCEPT_ENCODING)).isNull(); } } + private Stream> executeQuery(final FlightSqlClient flightSqlClient, + final Map queryParameters, + final Map headers, + final CallOption... callOptions) { + return flightSqlClient.execute( + "select * from cpu", + "mydb", + QueryType.SQL, + queryParameters, + headers, + callOptions) + .flatMap(vector -> IntStream.range(0, vector.getRowCount()) + .mapToObj(rowNumber -> + VectorSchemaRootConverter.INSTANCE + .getMapFromVectorSchemaRoot( + vector, + rowNumber + ))); + } + + @BeforeEach + void setUp() throws Exception { + allocator = new RootAllocator(Long.MAX_VALUE); + server = FlightServer.builder(allocator, grpcLocation, TestUtils.simpleProducer(vectorSchemaRoot)).middleware( + HEADER_CAPTURE_KEY, + headerFactory).build().start(); + } + + @AfterEach + void tearDown() throws Exception { + if (server != null) { + server.shutdown(); + server.awaitTermination(); + } + if (allocator != null) { + allocator.close(); + } + } + + @Test + void flightSqlClient() throws Exception { + String correctHost = "grpc+unix://tmp/dummy.sock"; + ClientConfig clientConfig = new ClientConfig.Builder() + .host(correctHost) + .token("Token".toCharArray()) + .build(); + try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig)) { + Assertions.assertThat(flightSqlClient).isNotNull(); + } + + FlightClient.Builder builder = FlightClient.builder(allocator, server.getLocation()); + try (FlightClient flightClient = builder.build()) { + FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig, flightClient); + Assertions.assertThat(flightSqlClient).isNotNull(); + } + + var inCorrectHost = "grpc+unix://///tmp/dummy.sock"; + ClientConfig clientConfig1 = new ClientConfig.Builder() + .host(inCorrectHost) + .token("Token".toCharArray()) + .build(); + Assertions.assertThatThrownBy(() -> new FlightSqlClient(clientConfig1)); + } + @Test void createProxyDetector() { String targetUrl = "https://localhost:80"; ClientConfig clientConfig = new ClientConfig.Builder() - .host(targetUrl) - .build(); + .host(targetUrl) + .build(); try (FlightSqlClient flightSqlClient = new FlightSqlClient(clientConfig)) { String proxyUrl = "http://localhost:10000"; ProxyDetector proxyDetector = flightSqlClient.createProxyDetector(targetUrl, proxyUrl); Assertions.assertThat(proxyDetector.proxyFor( - new InetSocketAddress("localhost", 80) + new InetSocketAddress("localhost", 80) )).isEqualTo(HttpConnectProxiedSocketAddress.newBuilder() - .setProxyAddress(new InetSocketAddress("localhost", 10000)) - .setTargetAddress(new InetSocketAddress("localhost", 80)) - .build()); + .setProxyAddress(new InetSocketAddress("localhost", 10000)) + .setTargetAddress(new InetSocketAddress("localhost", 80)) + .build()); // Return null case Assertions.assertThat(proxyDetector.proxyFor( - new InetSocketAddress("123.2.3.1", 80) + new InetSocketAddress("123.2.3.1", 80) )).isNull(); } catch (Exception e) { throw new RuntimeException(e); } } - static class CallHeadersMiddleware implements FlightClientMiddleware.Factory { - CallHeaders headers; + static class HeaderCaptureMiddleware implements FlightServerMiddleware { + + private final Map headers = new HashMap<>(); + + public HeaderCaptureMiddleware(final CallHeaders callHeaders) { + for (String key : callHeaders.keys()) { + headers.put(key, callHeaders.get(key)); + } + } + + public Map getHeaders() { + return headers; + } + @Override - public FlightClientMiddleware onCallStarted(final CallInfo info) { - return new FlightClientMiddleware() { - @Override - public void onBeforeSendingHeaders(final CallHeaders outgoingHeaders) { - headers = outgoingHeaders; - } + public void onBeforeSendingHeaders(final CallHeaders callHeaders) { - @Override - public void onHeadersReceived(final CallHeaders incomingHeaders) { - } + } - @Override - public void onCallCompleted(final CallStatus status) { - } - }; + @Override + public void onCallCompleted(final CallStatus callStatus) { + + } + + @Override + public void onCallErrored(final Throwable throwable) { + + } + } + + static class HeaderCaptureMiddlewareFactory implements FlightServerMiddleware.Factory { + + private HeaderCaptureMiddleware lastInstance; + + + public HeaderCaptureMiddleware getLastInstance() { + return lastInstance; + } + + @Override + public HeaderCaptureMiddleware onCallStarted(final CallInfo callInfo, + final CallHeaders callHeaders, + final RequestContext requestContext) { + lastInstance = new HeaderCaptureMiddleware(callHeaders); + return lastInstance; } } } diff --git a/src/test/java/com/influxdb/v3/client/query/QueryOptionsTest.java b/src/test/java/com/influxdb/v3/client/query/QueryOptionsTest.java index 845c9522..ecfd988c 100644 --- a/src/test/java/com/influxdb/v3/client/query/QueryOptionsTest.java +++ b/src/test/java/com/influxdb/v3/client/query/QueryOptionsTest.java @@ -24,14 +24,10 @@ import java.io.IOException; import java.net.ServerSocket; import java.net.URI; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.List; import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.stream.Stream; -import javax.annotation.Nonnull; import io.grpc.Deadline; import io.grpc.ManagedChannel; @@ -41,24 +37,17 @@ import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightServer; -import org.apache.arrow.flight.Location; -import org.apache.arrow.flight.NoOpFlightProducer; -import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.Schema; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import com.influxdb.v3.client.InfluxDBClient; import com.influxdb.v3.client.PointValues; +import com.influxdb.v3.client.TestUtils; import com.influxdb.v3.client.config.ClientConfig; import com.influxdb.v3.client.internal.GrpcCallOptions; @@ -121,9 +110,10 @@ void setInboundMessageSizeSmall() throws Exception { int freePort = findFreePort(); URI uri = URI.create("http://127.0.0.1:" + freePort); int rowCount = 100; - try (VectorSchemaRoot vectorSchemaRoot = generateVectorSchemaRoot(10, rowCount); + try (VectorSchemaRoot vectorSchemaRoot = TestUtils.generateVectorSchemaRoot(10, rowCount); BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); - FlightServer flightServer = simpleFlightServer(uri, allocator, simpleProducer(vectorSchemaRoot)) + FlightServer flightServer = TestUtils.simpleFlightServer(uri, allocator, + TestUtils.simpleProducer(vectorSchemaRoot)) ) { flightServer.start(); @@ -160,9 +150,10 @@ void setInboundMessageSizeLarge() throws Exception { int freePort = findFreePort(); URI uri = URI.create("http://127.0.0.1:" + freePort); int rowCount = 100; - try (VectorSchemaRoot vectorSchemaRoot = generateVectorSchemaRoot(10, rowCount); + try (VectorSchemaRoot vectorSchemaRoot = TestUtils.generateVectorSchemaRoot(10, rowCount); BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); - FlightServer flightServer = simpleFlightServer(uri, allocator, simpleProducer(vectorSchemaRoot)) + FlightServer flightServer = TestUtils.simpleFlightServer(uri, allocator, + TestUtils.simpleProducer(vectorSchemaRoot)) ) { flightServer.start(); @@ -256,46 +247,4 @@ void grpcCallOptions() { Assertions.assertThat(stubCallOptions.getDeadline()).isEqualTo(grpcCallOption.getDeadline()); } - private FlightServer simpleFlightServer(@Nonnull final URI uri, - @Nonnull final BufferAllocator allocator, - @Nonnull final NoOpFlightProducer producer) throws Exception { - Location location = Location.forGrpcInsecure(uri.getHost(), uri.getPort()); - return FlightServer.builder(allocator, location, producer).build(); - } - - private NoOpFlightProducer simpleProducer(@Nonnull final VectorSchemaRoot vectorSchemaRoot) { - return new NoOpFlightProducer() { - @Override - public void getStream(final CallContext context, - final Ticket ticket, - final ServerStreamListener listener) { - listener.start(vectorSchemaRoot); - if (listener.isReady()) { - listener.putNext(); - } - listener.completed(); - } - }; - } - - private VectorSchemaRoot generateVectorSchemaRoot(final int fieldCount, final int rowCount) { - List fields = new ArrayList<>(); - for (int i = 0; i < fieldCount; i++) { - Field field = new Field("field" + i, FieldType.nullable(new ArrowType.Utf8()), null); - fields.add(field); - } - - Schema schema = new Schema(fields); - VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.create(schema, new RootAllocator(Long.MAX_VALUE)); - for (Field field : fields) { - VarCharVector vector = (VarCharVector) vectorSchemaRoot.getVector(field); - vector.allocateNew(rowCount); - for (int i = 0; i < rowCount; i++) { - vector.set(i, "Value".getBytes(StandardCharsets.UTF_8)); - } - } - vectorSchemaRoot.setRowCount(rowCount); - - return vectorSchemaRoot; - } }