diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java index 179e0c416e..5e94a39a24 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java @@ -16,9 +16,14 @@ */ package org.apache.arrow.adbc.driver.flightsql; +import java.net.URI; import java.net.URISyntaxException; +import java.net.URLDecoder; +import java.nio.charset.StandardCharsets; +import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.function.Supplier; import org.apache.arrow.adbc.core.AdbcDatabase; import org.apache.arrow.adbc.core.AdbcDriver; import org.apache.arrow.adbc.core.AdbcException; @@ -27,8 +32,19 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; -/** An ADBC driver wrapping Arrow Flight SQL. */ +/** + * An ADBC driver wrapping Arrow Flight SQL. + * + *

The "uri" option accepts the {@code flightsql://} scheme: secure TLS by default, or add {@code + * ?transport=tcp} for plaintext, or {@code ?transport=unix} with a socket path for a Unix domain + * socket. The {@code transport} value is matched case-insensitively, and an unrecognized value is + * rejected rather than silently falling back to a default. The legacy {@code grpc://}, {@code + * grpc+tcp://}, {@code grpc+tls://}, and {@code grpc+unix://} schemes are also still accepted. + */ public class FlightSqlDriver implements AdbcDriver { + private static final String FLIGHTSQL_SCHEME = "flightsql"; + private static final String TRANSPORT_PARAM = "transport"; + private final BufferAllocator allocator; public FlightSqlDriver(BufferAllocator allocator) { @@ -47,14 +63,8 @@ public AdbcDatabase open(Map parameters) throws AdbcException { uri = (String) target; } - Location location; - try { - location = new Location(uri); - } catch (URISyntaxException e) { - throw AdbcException.invalidArgument( - String.format("[Flight SQL] Location %s is invalid: %s", uri, e)) - .withCause(e); - } + Location location = parseLocation(uri); + Object quirks = parameters.get(PARAM_SQL_QUIRKS); if (quirks != null) { Preconditions.checkArgument( @@ -67,4 +77,129 @@ public AdbcDatabase open(Map parameters) throws AdbcException { } return new FlightSqlDatabase(allocator, location, (SqlQuirks) quirks, parameters); } + + /** + * Parses the "uri" option into a {@link Location}, translating a {@code flightsql://} scheme (and + * its {@code transport} query parameter) into the legacy {@code grpc*://} scheme that {@link + * Location} understands natively. Any other scheme is passed through unchanged. + */ + static Location parseLocation(String uri) throws AdbcException { + final URI parsed; + try { + parsed = new URI(uri); + } catch (URISyntaxException e) { + throw AdbcException.invalidArgument( + String.format("[Flight SQL] Location %s is invalid: %s", uri, e)) + .withCause(e); + } + + if (!FLIGHTSQL_SCHEME.equalsIgnoreCase(parsed.getScheme())) { + return new Location(parsed); + } + + if (parsed.getUserInfo() != null || parsed.getFragment() != null) { + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': userinfo and fragment are not supported in %s://" + + " URIs", + uri, FLIGHTSQL_SCHEME)); + } + + final String transport = transportOf(uri, parsed); + final String host = parsed.getHost(); + final String path = parsed.getPath(); + switch (transport) { + case "": + case "tls": + requireNoPath(uri, path); + requireHost(uri, host); + return buildLocation(uri, () -> Location.forGrpcTls(host, parsed.getPort())); + case "tcp": + requireNoPath(uri, path); + requireHost(uri, host); + return buildLocation(uri, () -> Location.forGrpcInsecure(host, parsed.getPort())); + case "unix": + if (host != null && !host.isEmpty()) { + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': a host is not valid with transport=unix", uri)); + } + if (path == null || path.isEmpty()) { + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': transport=unix requires a socket path", uri)); + } + return buildLocation(uri, () -> Location.forGrpcDomainSocket(path)); + default: + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': unrecognized transport '%s' (expected 'tls'," + + " 'tcp', or 'unix')", + uri, transport)); + } + } + + private static void requireNoPath(String uri, String path) throws AdbcException { + if (path != null && !path.isEmpty()) { + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': a socket path is only valid with transport=unix", + uri)); + } + } + + private static void requireHost(String uri, String host) throws AdbcException { + if (host == null || host.isEmpty()) { + throw AdbcException.invalidArgument( + String.format("[Flight SQL] Invalid URI '%s': a host is required", uri)); + } + } + + /** + * Builds a {@link Location} via the given factory, wrapping the {@link IllegalArgumentException} + * that {@link Location}'s {@code forGrpc*} factories document throwing on an invalid URI into an + * {@link AdbcException} instead of letting it escape as an unchecked exception. + */ + private static Location buildLocation(String uri, Supplier factory) + throws AdbcException { + try { + return factory.get(); + } catch (IllegalArgumentException e) { + throw AdbcException.invalidArgument( + String.format("[Flight SQL] Location %s is invalid: %s", uri, e)) + .withCause(e); + } + } + + /** Extracts and lowercases the {@code transport} query parameter, defaulting to "". */ + private static String transportOf(String uri, URI parsed) throws AdbcException { + final String query = parsed.getRawQuery(); + if (query == null) { + return ""; + } + int start = 0; + while (start <= query.length()) { + int amp = query.indexOf('&', start); + int end = amp >= 0 ? amp : query.length(); + String pair = query.substring(start, end); + int eq = pair.indexOf('='); + String key = eq >= 0 ? pair.substring(0, eq) : pair; + if (TRANSPORT_PARAM.equalsIgnoreCase(key)) { + String value = eq >= 0 ? pair.substring(eq + 1) : ""; + try { + return URLDecoder.decode(value, StandardCharsets.UTF_8).toLowerCase(Locale.ROOT); + } catch (IllegalArgumentException e) { + throw AdbcException.invalidArgument( + String.format( + "[Flight SQL] Invalid URI '%s': malformed transport value: %s", uri, e)) + .withCause(e); + } + } + if (amp < 0) { + break; + } + start = amp + 1; + } + return ""; + } } diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriverUriTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriverUriTest.java new file mode 100644 index 0000000000..54c3be4c17 --- /dev/null +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriverUriTest.java @@ -0,0 +1,214 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adbc.driver.flightsql; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.HashMap; +import java.util.Map; +import org.apache.arrow.adbc.core.AdbcConnection; +import org.apache.arrow.adbc.core.AdbcDatabase; +import org.apache.arrow.adbc.core.AdbcDriver; +import org.apache.arrow.adbc.core.AdbcException; +import org.apache.arrow.adbc.core.AdbcStatusCode; +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.sql.NoOpFlightSqlProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** Tests for the {@code flightsql://} URI scheme handled by {@link FlightSqlDriver}. */ +class FlightSqlDriverUriTest { + static BufferAllocator allocator; + static FlightServer server; + static FlightSqlDriver driver; + + @BeforeAll + static void beforeAll() throws Exception { + allocator = new RootAllocator(); + server = + FlightServer.builder() + .allocator(allocator) + .producer(new NoOpFlightSqlProducer()) + .location(Location.forGrpcInsecure("localhost", 0)) + .build(); + server.start(); + driver = new FlightSqlDriver(allocator); + } + + @AfterAll + static void afterAll() throws Exception { + AutoCloseables.close(server, allocator); + } + + @Test + void defaultTransportIsTls() throws Exception { + Location location = FlightSqlDriver.parseLocation("flightsql://localhost:1234"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tls"); + assertThat(location.getUri().getHost()).isEqualTo("localhost"); + assertThat(location.getUri().getPort()).isEqualTo(1234); + } + + @Test + void explicitTlsTransport() throws Exception { + Location location = FlightSqlDriver.parseLocation("flightsql://localhost:1234?transport=tls"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tls"); + } + + @Test + void tcpTransportIsCaseInsensitive() throws Exception { + Location location = FlightSqlDriver.parseLocation("flightsql://localhost:1234?transport=TCP"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tcp"); + assertThat(location.getUri().getPort()).isEqualTo(1234); + } + + @Test + void unixTransport() throws Exception { + Location location = FlightSqlDriver.parseLocation("flightsql:///tmp/adbc.sock?transport=unix"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+unix"); + assertThat(location.getUri().getPath()).isEqualTo("/tmp/adbc.sock"); + } + + @Test + void schemeIsCaseInsensitive() throws Exception { + Location location = FlightSqlDriver.parseLocation("FLIGHTSQL://localhost:1234"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tls"); + } + + @Test + void legacySchemesStillWork() throws Exception { + Location location = FlightSqlDriver.parseLocation("grpc+tls://localhost:1234"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tls"); + } + + @Test + void connectsOverTcpTransport() throws Exception { + Map parameters = new HashMap<>(); + AdbcDriver.PARAM_URI.set( + parameters, "flightsql://localhost:" + server.getPort() + "?transport=tcp"); + + try (AdbcDatabase database = driver.open(parameters); + AdbcConnection connection = database.connect()) { + assertThat(connection).isNotNull(); + } + } + + @Test + void rejectsUnrecognizedTransport() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://localhost:1234?transport=bogus")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsHostWithUnixTransport() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> + FlightSqlDriver.parseLocation( + "flightsql://localhost:1234/socket?transport=unix")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsPathWithTcpTransport() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> + FlightSqlDriver.parseLocation("flightsql://localhost:1234/socket?transport=tcp")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsPathWithDefaultTransport() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://localhost:1234/socket")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsUnixTransportWithoutPath() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://?transport=unix")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsHostWithUnderscoreAsAdbcException() { + // java.net.URI treats an underscore as invalid in a "server-based" authority and silently + // parses the host/port as absent instead of throwing, so this must be caught explicitly + // rather than letting Location.forGrpcTls's IllegalArgumentException escape uncaught. + AdbcException ex = + assertThrows( + AdbcException.class, () -> FlightSqlDriver.parseLocation("flightsql://my_host:1234")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsMissingHostAsAdbcException() { + AdbcException ex = + assertThrows(AdbcException.class, () -> FlightSqlDriver.parseLocation("flightsql://:1234")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void transportKeyIsCaseInsensitive() throws Exception { + Location location = FlightSqlDriver.parseLocation("flightsql://localhost:1234?TRANSPORT=tcp"); + assertThat(location.getUri().getScheme()).isEqualTo("grpc+tcp"); + } + + @Test + void rejectsMalformedTransportPercentEncoding() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://localhost:1234?transport=%zz")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsUserInfo() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://alice:secret@localhost:1234")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + + @Test + void rejectsFragment() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> FlightSqlDriver.parseLocation("flightsql://localhost:1234#frag")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.INVALID_ARGUMENT); + } + +} diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/TlsTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/TlsTest.java index bcfff64dd0..60a8d2505e 100644 --- a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/TlsTest.java +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/TlsTest.java @@ -170,10 +170,37 @@ public void testClientTlsOnGoodHostnameOverride() throws Exception { } } + @Test + public void testClientFlightSqlSchemeDefaultTls() throws Exception { + params.put(AdbcDriver.PARAM_URI.getKey(), getFlightSqlUri(null)); + params.put(FlightSqlConnectionProperties.TLS_SKIP_VERIFY.getKey(), true); + try (AdbcDatabase db = + AdbcDriverManager.getInstance() + .connect(FlightSqlDriverFactory.class.getCanonicalName(), allocator, params); + AdbcConnection conn = db.connect()) {} + } + + @Test + public void testClientFlightSqlSchemeExplicitTlsTransport() throws Exception { + params.put(AdbcDriver.PARAM_URI.getKey(), getFlightSqlUri("tls")); + params.put(FlightSqlConnectionProperties.TLS_SKIP_VERIFY.getKey(), true); + try (AdbcDatabase db = + AdbcDriverManager.getInstance() + .connect(FlightSqlDriverFactory.class.getCanonicalName(), allocator, params); + AdbcConnection conn = db.connect()) {} + } + private String getUri(boolean withTls) { String protocol = String.format("grpc%s", withTls ? "+tls" : "+tcp"); return String.format( "%s://%s:%d", protocol, FLIGHT_SERVER_TEST_EXTENSION.getHost(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); } + + private String getFlightSqlUri(String transport) { + String query = transport == null ? "" : "?transport=" + transport; + return String.format( + "flightsql://%s:%d%s", + FLIGHT_SERVER_TEST_EXTENSION.getHost(), FLIGHT_SERVER_TEST_EXTENSION.getPort(), query); + } }