Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,11 @@
*/
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 org.apache.arrow.adbc.core.AdbcDatabase;
Expand All @@ -27,8 +31,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.
*
* <p>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) {
Expand All @@ -47,14 +62,8 @@ public AdbcDatabase open(Map<String, Object> 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);
Comment thread
unikdahal marked this conversation as resolved.

Object quirks = parameters.get(PARAM_SQL_QUIRKS);
if (quirks != null) {
Preconditions.checkArgument(
Expand All @@ -67,4 +76,103 @@ public AdbcDatabase open(Map<String, Object> 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();

if (transport.isEmpty() || "tls".equals(transport) || "tcp".equals(transport)) {
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));
}
if (host == null || host.isEmpty()) {
throw AdbcException.invalidArgument(
String.format("[Flight SQL] Invalid URI '%s': a valid host is required", uri));
}
if ("tcp".equals(transport)) {
return Location.forGrpcInsecure(host, parsed.getPort());
} else {
return Location.forGrpcTls(host, parsed.getPort());
}
} else if ("unix".equals(transport)) {
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 Location.forGrpcDomainSocket(path);
}
throw AdbcException.invalidArgument(
String.format(
"[Flight SQL] Invalid URI '%s': unrecognized transport '%s' (expected 'tls',"
+ " 'tcp', or 'unix')",
uri, transport));
}

/** 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 "";
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,211 @@
/*
* 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<String, Object> 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);
}
}
Loading
Loading