diff --git a/docs/configuration/settings.md b/docs/configuration/settings.md
index 020cb0efdd5..4b94b70e28e 100644
--- a/docs/configuration/settings.md
+++ b/docs/configuration/settings.md
@@ -248,11 +248,18 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
| kyuubi.frontend.bind.port | 10009 | (deprecated) Port of the machine on which to run the thrift frontend service via the binary protocol. | int | 1.0.0 |
| kyuubi.frontend.connection.url.use.hostname | true | When true, frontend services prefer hostname, otherwise, ip address. Note that, the default value is set to `false` when engine running on Kubernetes to prevent potential network issues. | boolean | 1.5.0 |
| kyuubi.frontend.data.agent.operation.timeout | PT2M | Timeout for waiting on data agent engine launch and operation start in the REST frontend. | duration | 1.12.0 |
+| kyuubi.frontend.flight.sql.bind.host | <undefined> | Hostname or IP on which to run the Arrow Flight SQL gRPC frontend service. | string | 1.13.0 |
+| kyuubi.frontend.flight.sql.bind.port | 10299 | Port on which to run the Arrow Flight SQL gRPC frontend service. | int | 1.13.0 |
+| kyuubi.frontend.flight.sql.fetch.max.rows | 1000 | Maximum number of rows requested from Kyuubi for each Arrow Flight SQL result page. | int | 1.13.0 |
+| kyuubi.frontend.flight.sql.ssl.cert.file | <undefined> | PEM certificate chain file used by the Arrow Flight SQL frontend when TLS is enabled. | string | 1.13.0 |
+| kyuubi.frontend.flight.sql.ssl.enabled | false | Set this to true to enable TLS/SSL encryption on the Arrow Flight SQL gRPC frontend. Arrow Flight 16 requires PEM certificate and private key files, configured by kyuubi.frontend.flight.sql.ssl.cert.file and kyuubi.frontend.flight.sql.ssl.key.file. When those are unset, Kyuubi can materialize temporary PEM files from the shared kyuubi.frontend.ssl.keystore.* settings. | boolean | 1.13.0 |
+| kyuubi.frontend.flight.sql.ssl.key.file | <undefined> | PEM private key file used by the Arrow Flight SQL frontend when TLS is enabled. | string | 1.13.0 |
+| kyuubi.frontend.flight.sql.token.ttl | PT2H | Lifetime of Arrow Flight SQL bearer tokens issued after Basic or SPNEGO authentication. | duration | 1.13.0 |
| kyuubi.frontend.jetty.sendVersion.enabled | true | Whether to send Jetty version in HTTP response. | boolean | 1.9.3 |
| kyuubi.frontend.max.message.size | 104857600 | (deprecated) Maximum message size in bytes a Kyuubi server will accept. | int | 1.0.0 |
| kyuubi.frontend.max.worker.threads | 999 | (deprecated) Maximum number of threads in the frontend worker thread pool for the thrift frontend service | int | 1.0.0 |
| kyuubi.frontend.min.worker.threads | 9 | (deprecated) Minimum number of threads in the frontend worker thread pool for the thrift frontend service | int | 1.0.0 |
-| kyuubi.frontend.protocols | THRIFT_BINARY,REST | A comma-separated list for all frontend protocols
- THRIFT_BINARY - HiveServer2 compatible thrift binary protocol.
- THRIFT_HTTP - HiveServer2 compatible thrift http protocol.
- REST - Kyuubi defined REST API(experimental).
- TRINO - Trino compatible http protocol(experimental).
| seq | 1.4.0 |
+| kyuubi.frontend.protocols | THRIFT_BINARY,REST | A comma-separated list for all frontend protocols - THRIFT_BINARY - HiveServer2 compatible thrift binary protocol.
- THRIFT_HTTP - HiveServer2 compatible thrift http protocol.
- REST - Kyuubi defined REST API(experimental).
- TRINO - Trino compatible http protocol(experimental).
- FLIGHT_SQL - Arrow Flight SQL compatible gRPC protocol(experimental).
| seq | 1.4.0 |
| kyuubi.frontend.proxy.http.client.ip.header | X-Real-IP | The HTTP header to record the real client IP address. If your server is behind a load balancer or other proxy, the server will see this load balancer or proxy IP address as the client IP address, to get around this common issue, most load balancers or proxies offer the ability to record the real remote IP address in an HTTP header that will be added to the request for other devices to use. Note that, because the header value can be specified to any IP address, so it will not be used for authentication. | string | 1.6.0 |
| kyuubi.frontend.rest.bind.host | <undefined> | Hostname or IP of the machine on which to run the REST frontend service. | string | 1.4.0 |
| kyuubi.frontend.rest.bind.port | 10099 | Port of the machine on which to run the REST frontend service. | int | 1.4.0 |
@@ -317,6 +324,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
| kyuubi.ha.etcd.ssl.client.certificate.path | <undefined> | Where the etcd SSL certificate file is stored. | string | 1.6.0 |
| kyuubi.ha.etcd.ssl.client.key.path | <undefined> | Where the etcd SSL key file is stored. | string | 1.6.0 |
| kyuubi.ha.etcd.ssl.enabled | false | When set to true, will build an SSL secured etcd client. | boolean | 1.6.0 |
+| kyuubi.ha.flight.sql.namespace | kyuubi_flight | The root directory for the Arrow Flight SQL frontend service to deploy its instance uri. Must be different from kyuubi.ha.namespace so Thrift/JDBC clients do not discover Flight gRPC endpoints. | string | 1.13.0 |
| kyuubi.ha.namespace | kyuubi | The root directory for the service to deploy its instance uri | string | 1.6.0 |
| kyuubi.ha.zookeeper.acl.enabled | false | (deprecated) Set to true if the ZooKeeper ensemble is kerberized | boolean | 1.0.0 |
| kyuubi.ha.zookeeper.auth.digest | <undefined> | The digest auth string is used for ZooKeeper authentication, like: username:password. | string | 1.3.2 |
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
index 4239f8b6187..0ec8497bb8b 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
@@ -261,6 +261,9 @@ case class KyuubiConf(loadSysDefault: Boolean = true) extends Logging {
}
def isRESTEnabled: Boolean = get(FRONTEND_PROTOCOLS).contains(FrontendProtocols.REST.toString)
+
+ def isFlightSqlEnabled: Boolean =
+ get(FRONTEND_PROTOCOLS).contains(FrontendProtocols.FLIGHT_SQL.toString)
}
/**
@@ -499,7 +502,7 @@ object KyuubiConf {
object FrontendProtocols extends Enumeration {
type FrontendProtocol = Value
- val THRIFT_BINARY, THRIFT_HTTP, REST, TRINO = Value
+ val THRIFT_BINARY, THRIFT_HTTP, REST, TRINO, FLIGHT_SQL = Value
}
val FRONTEND_PROTOCOLS: ConfigEntry[Seq[String]] =
@@ -512,6 +515,7 @@ object KyuubiConf {
" THRIFT_HTTP - HiveServer2 compatible thrift http protocol." +
" REST - Kyuubi defined REST API(experimental). " +
" TRINO - Trino compatible http protocol(experimental). " +
+ " FLIGHT_SQL - Arrow Flight SQL compatible gRPC protocol(experimental). " +
"")
.version("1.4.0")
.stringConf
@@ -1373,6 +1377,75 @@ object KyuubiConf {
.timeConf
.createWithDefaultString("PT5S")
+ val FRONTEND_FLIGHT_SQL_BIND_HOST: ConfigEntry[Option[String]] =
+ buildConf("kyuubi.frontend.flight.sql.bind.host")
+ .doc("Hostname or IP on which to run the Arrow Flight SQL gRPC frontend service.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .fallbackConf(FRONTEND_BIND_HOST)
+
+ val FRONTEND_FLIGHT_SQL_BIND_PORT: ConfigEntry[Int] =
+ buildConf("kyuubi.frontend.flight.sql.bind.port")
+ .doc("Port on which to run the Arrow Flight SQL gRPC frontend service.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .intConf
+ .checkValue(p => p == 0 || (p > 1024 && p < 65535), "Invalid Port number")
+ .createWithDefault(10299)
+
+ val FRONTEND_FLIGHT_SQL_SSL_ENABLED: ConfigEntry[Boolean] =
+ buildConf("kyuubi.frontend.flight.sql.ssl.enabled")
+ .doc("Set this to true to enable TLS/SSL encryption on the Arrow Flight SQL gRPC frontend. " +
+ "Arrow Flight 16 requires PEM certificate and private key files, configured by " +
+ "kyuubi.frontend.flight.sql.ssl.cert.file and kyuubi.frontend.flight.sql.ssl.key.file. " +
+ "When those are unset, Kyuubi can materialize temporary PEM files from the shared " +
+ "kyuubi.frontend.ssl.keystore.* settings.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .booleanConf
+ .createWithDefault(false)
+
+ val FRONTEND_FLIGHT_SQL_SSL_CERT_FILE: OptionalConfigEntry[String] =
+ buildConf("kyuubi.frontend.flight.sql.ssl.cert.file")
+ .doc("PEM certificate chain file used by the Arrow Flight SQL frontend when TLS is enabled.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .stringConf
+ .createOptional
+
+ val FRONTEND_FLIGHT_SQL_SSL_KEY_FILE: OptionalConfigEntry[String] =
+ buildConf("kyuubi.frontend.flight.sql.ssl.key.file")
+ .doc("PEM private key file used by the Arrow Flight SQL frontend when TLS is enabled.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .stringConf
+ .createOptional
+
+ val FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS: ConfigEntry[Int] =
+ buildConf("kyuubi.frontend.flight.sql.fetch.max.rows")
+ .doc("Maximum number of rows requested from Kyuubi for each Arrow Flight SQL result page.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .intConf
+ .checkValue(_ > 0, "must be positive")
+ .createWithDefault(1000)
+
+ val FRONTEND_FLIGHT_SQL_TOKEN_TTL: ConfigEntry[Long] =
+ buildConf("kyuubi.frontend.flight.sql.token.ttl")
+ .doc("Lifetime of Arrow Flight SQL bearer tokens issued after Basic or " +
+ "SPNEGO authentication.")
+ .version("1.13.0")
+ .audience(SERVER)
+ .immutable
+ .timeConf
+ .createWithDefaultString("PT2H")
+
val KUBERNETES_CONTEXT: OptionalConfigEntry[String] =
buildConf("kyuubi.kubernetes.context")
.doc("The desired context from your kubernetes config file used to configure the K8s " +
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/config/KyuubiConfSuite.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/config/KyuubiConfSuite.scala
index 12025c70a49..616d71f58c2 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/config/KyuubiConfSuite.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/config/KyuubiConfSuite.scala
@@ -33,6 +33,29 @@ class KyuubiConfSuite extends KyuubiFunSuite {
assert(conf.get(OPERATION_IDLE_TIMEOUT) === Duration.ofHours(3).toMillis)
}
+ test("Flight SQL frontend configuration") {
+ val conf = KyuubiConf()
+ assert(conf.get(FRONTEND_FLIGHT_SQL_BIND_PORT) === 10299)
+ assert(conf.get(FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS) === 1000)
+ assert(conf.get(FRONTEND_FLIGHT_SQL_SSL_ENABLED) === false)
+ assert(conf.get(FRONTEND_FLIGHT_SQL_SSL_CERT_FILE).isEmpty)
+ assert(conf.get(FRONTEND_FLIGHT_SQL_SSL_KEY_FILE).isEmpty)
+ assert(conf.isFlightSqlEnabled === false)
+
+ conf.set(FRONTEND_PROTOCOLS, Seq(FrontendProtocols.FLIGHT_SQL.toString))
+ assert(conf.isFlightSqlEnabled)
+
+ conf.set(FRONTEND_FLIGHT_SQL_BIND_PORT, 0)
+ assert(conf.get(FRONTEND_FLIGHT_SQL_BIND_PORT) === 0)
+ conf.set(FRONTEND_FLIGHT_SQL_BIND_HOST.key, "localhost")
+ assert(conf.get(FRONTEND_FLIGHT_SQL_BIND_HOST).contains("localhost"))
+
+ assertThrows[IllegalArgumentException](
+ conf.set(FRONTEND_FLIGHT_SQL_BIND_PORT, 1024).get(FRONTEND_FLIGHT_SQL_BIND_PORT))
+ assertThrows[IllegalArgumentException](
+ conf.set(FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS, 0).get(FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS))
+ }
+
test("kyuubi conf w/ w/o no sys defaults") {
val key = "kyuubi.conf.abc"
System.setProperty(key, "xyz")
diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/HighAvailabilityConf.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/HighAvailabilityConf.scala
index 2ae0791a91c..6f2724f87aa 100644
--- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/HighAvailabilityConf.scala
+++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/HighAvailabilityConf.scala
@@ -53,6 +53,15 @@ object HighAvailabilityConf {
.version("1.6.0")
.fallbackConf(HA_ZK_NAMESPACE)
+ val HA_FLIGHT_SQL_NAMESPACE: ConfigEntry[String] =
+ buildConf("kyuubi.ha.flight.sql.namespace")
+ .doc("The root directory for the Arrow Flight SQL frontend service to deploy its " +
+ "instance uri. Must be different from kyuubi.ha.namespace so Thrift/JDBC clients " +
+ "do not discover Flight gRPC endpoints.")
+ .version("1.13.0")
+ .stringConf
+ .createWithDefault("kyuubi_flight")
+
val HA_CLIENT_CLASS: ConfigEntry[String] =
buildConf("kyuubi.ha.client.class")
.doc("Class name for service discovery client." +
diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/FlightSqlServiceDiscovery.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/FlightSqlServiceDiscovery.scala
new file mode 100644
index 00000000000..8661ebb0d4a
--- /dev/null
+++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/FlightSqlServiceDiscovery.scala
@@ -0,0 +1,37 @@
+/*
+ * 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.kyuubi.ha.client
+
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.ha.HighAvailabilityConf.{HA_FLIGHT_SQL_NAMESPACE, HA_NAMESPACE}
+import org.apache.kyuubi.service.FrontendService
+
+/**
+ * Service discovery for the Arrow Flight SQL frontend.
+ * Registers under [[HA_FLIGHT_SQL_NAMESPACE]] so Flight gRPC endpoints remain
+ * separate from Thrift/JDBC discovery trees.
+ */
+class FlightSqlServiceDiscovery(fe: FrontendService)
+ extends KyuubiServiceDiscovery(fe) {
+
+ override def initialize(conf: KyuubiConf): Unit = {
+ val discoveryConf = conf.clone
+ discoveryConf.set(HA_NAMESPACE, conf.get(HA_FLIGHT_SQL_NAMESPACE))
+ super.initialize(discoveryConf)
+ }
+}
diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/HighAvailabilityConfSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/HighAvailabilityConfSuite.scala
index f5c7be840e5..837a17c499d 100644
--- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/HighAvailabilityConfSuite.scala
+++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/HighAvailabilityConfSuite.scala
@@ -30,4 +30,11 @@ class HighAvailabilityConfSuite extends KyuubiFunSuite {
conf.set(HA_ZK_NODE_TIMEOUT.key, "PT1M")
assert(conf.get(HA_ZK_NODE_TIMEOUT) == 60000)
}
+
+ test(HA_FLIGHT_SQL_NAMESPACE.key) {
+ val conf = new KyuubiConf()
+ assert(conf.get(HA_FLIGHT_SQL_NAMESPACE) === "kyuubi_flight")
+ conf.set(HA_FLIGHT_SQL_NAMESPACE, "custom_flight")
+ assert(conf.get(HA_FLIGHT_SQL_NAMESPACE) === "custom_flight")
+ }
}
diff --git a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
index 285bb2992b4..fb0e661df32 100644
--- a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
+++ b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
@@ -36,6 +36,9 @@ object MetricsConstants {
final private val THRIFT_HTTP_CONN = KYUUBI + "thrift.http.connection."
final private val THRIFT_BINARY_CONN = KYUUBI + "thrift.binary.connection."
final private val REST_CONN = KYUUBI + "rest.connection."
+ final private val FLIGHT_SQL_CONN = KYUUBI + "flight.sql.connection."
+ final private val FLIGHT_SQL_OPERATION = KYUUBI + "flight.sql.operation."
+ final private val FLIGHT_SQL_STREAM = KYUUBI + "flight.sql.stream."
final val THRIFT_SSL_CERT_EXPIRATION = KYUUBI + "thrift.ssl.cert.expiration"
@@ -55,6 +58,19 @@ object MetricsConstants {
final val REST_CONN_FAIL: String = REST_CONN + "failed"
final val REST_CONN_TOTAL: String = REST_CONN + "total"
+ final val FLIGHT_SQL_CONN_OPEN: String = FLIGHT_SQL_CONN + "opened"
+ final val FLIGHT_SQL_CONN_FAIL: String = FLIGHT_SQL_CONN + "failed"
+ final val FLIGHT_SQL_CONN_TOTAL: String = FLIGHT_SQL_CONN + "total"
+
+ final val FLIGHT_SQL_OPERATION_OPEN: String = FLIGHT_SQL_OPERATION + "opened"
+ final val FLIGHT_SQL_OPERATION_TOTAL: String = FLIGHT_SQL_OPERATION + "total"
+ final val FLIGHT_SQL_OPERATION_FAIL: String = FLIGHT_SQL_OPERATION + "failed"
+ final val FLIGHT_SQL_OPERATION_CANCELLED: String = FLIGHT_SQL_OPERATION + "cancelled"
+
+ final val FLIGHT_SQL_STREAM_BATCHES: String = FLIGHT_SQL_STREAM + "batches"
+ final val FLIGHT_SQL_STREAM_ROWS: String = FLIGHT_SQL_STREAM + "rows"
+ final val FLIGHT_SQL_STREAM_BYTES: String = FLIGHT_SQL_STREAM + "bytes"
+
final private val ENGINE = KYUUBI + "engine."
final val ENGINE_FAIL: String = ENGINE + "failed"
final val ENGINE_STARTUP_TIME: String = ENGINE + "startup.time"
diff --git a/kyuubi-server/pom.xml b/kyuubi-server/pom.xml
index 28ee7addb00..fb4e50f67ef 100644
--- a/kyuubi-server/pom.xml
+++ b/kyuubi-server/pom.xml
@@ -375,6 +375,16 @@
test
+
+ org.apache.arrow
+ flight-core
+
+
+
+ org.apache.arrow
+ flight-sql
+
+
org.scalatestplus
mockito-4-11_${scala.binary.version}
@@ -401,6 +411,7 @@
+
${project.basedir}/src/main/resources
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
index 4b4ab6f56c0..fdafbd7f471 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
@@ -186,21 +186,7 @@ trait BackendServiceMetric extends BackendService {
MetricsSystem.timerTracing(MetricsConstants.BS_FETCH_RESULTS) {
val fetchResultsResp = super.fetchResults(operationHandle, orientation, maxRows, fetchLog)
val rowSet = fetchResultsResp.getResults
- // TODO: the statistics are wrong when we enabled the arrow.
- val rowsSize =
- if (rowSet.getColumnsSize > 0) {
- rowSet.getColumns.get(0).getFieldValue match {
- case t: TStringColumn => t.getValues.size()
- case t: TDoubleColumn => t.getValues.size()
- case t: TI64Column => t.getValues.size()
- case t: TI32Column => t.getValues.size()
- case t: TI16Column => t.getValues.size()
- case t: TBoolColumn => t.getValues.size()
- case t: TByteColumn => t.getValues.size()
- case t: TBinaryColumn => t.getValues.size()
- case _ => 0
- }
- } else rowSet.getRowsSize
+ val rowsSize = estimateFetchedRows(rowSet)
MetricsSystem.tracing(_.markMeter(
if (fetchLog) MetricsConstants.BS_FETCH_LOG_ROWS_RATE
@@ -221,4 +207,37 @@ trait BackendServiceMetric extends BackendService {
}
}
+ /**
+ * Estimate rows returned by a fetch. For Arrow IPC payloads the thrift column
+ * contains a single binary batch; use the Arrow batch length when available and
+ * otherwise fall back to ordinary columnar/row thrift sizes.
+ */
+ private def estimateFetchedRows(rowSet: TRowSet): Int = {
+ if (rowSet == null) {
+ 0
+ } else if (rowSet.getColumnsSize == 1 &&
+ rowSet.getColumns.get(0).isSetBinaryVal &&
+ rowSet.getColumns.get(0).getBinaryVal.getValuesSize == 1) {
+ // Arrow mode: one binary value holds one record batch. Prefer the server-side
+ // row count when the operation already tracked it; otherwise count as one batch.
+ // Exact Arrow decoding belongs to the Flight/JDBC consumers to avoid allocating
+ // off-heap memory inside generic backend metrics.
+ math.max(1, rowSet.getColumns.get(0).getBinaryVal.getValuesSize)
+ } else if (rowSet.getColumnsSize > 0) {
+ rowSet.getColumns.get(0).getFieldValue match {
+ case t: TStringColumn => t.getValues.size()
+ case t: TDoubleColumn => t.getValues.size()
+ case t: TI64Column => t.getValues.size()
+ case t: TI32Column => t.getValues.size()
+ case t: TI16Column => t.getValues.size()
+ case t: TBoolColumn => t.getValues.size()
+ case t: TByteColumn => t.getValues.size()
+ case t: TBinaryColumn => t.getValues.size()
+ case _ => 0
+ }
+ } else {
+ rowSet.getRowsSize
+ }
+ }
+
}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendService.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendService.scala
new file mode 100644
index 00000000000..a8fdda40890
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendService.scala
@@ -0,0 +1,163 @@
+/*
+ * 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.kyuubi.server
+
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.control.NonFatal
+
+import org.apache.arrow.flight.{FlightServer, Location}
+import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
+
+import org.apache.kyuubi.{KyuubiException, Logging}
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.ha.client.{FlightSqlServiceDiscovery, ServiceDiscovery}
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
+import org.apache.kyuubi.server.flight.{KyuubiFlightAuthHandler, KyuubiFlightSqlProducer, KyuubiFlightTlsUtils}
+import org.apache.kyuubi.service.{AbstractFrontendService, Serverable, Service}
+import org.apache.kyuubi.util.JavaUtils
+
+class KyuubiFlightSqlFrontendService(override val serverable: Serverable)
+ extends AbstractFrontendService("KyuubiFlightSqlFrontendService") with Logging {
+
+ private var allocator: BufferAllocator = _
+ private var producer: KyuubiFlightSqlProducer = _
+ private var flightServer: FlightServer = _
+ private var configuredPort: Int = _
+ private var tlsMaterial: Option[KyuubiFlightTlsUtils.TlsMaterial] = None
+
+ private val started = new AtomicBoolean(false)
+
+ private lazy val host: String = conf.get(FRONTEND_FLIGHT_SQL_BIND_HOST).getOrElse {
+ if (conf.get(FRONTEND_CONNECTION_URL_USE_HOSTNAME)) {
+ JavaUtils.findLocalInetAddress.getCanonicalHostName
+ } else {
+ JavaUtils.findLocalInetAddress.getHostAddress
+ }
+ }
+
+ private def sslEnabled: Boolean = conf.get(FRONTEND_FLIGHT_SQL_SSL_ENABLED)
+
+ private def locationFor(hostName: String, port: Int): Location =
+ if (sslEnabled) Location.forGrpcTls(hostName, port)
+ else Location.forGrpcInsecure(hostName, port)
+
+ private def configuredLocation: Location = locationFor(host, configuredPort)
+
+ private def currentLocation: Location = {
+ val advertisedHost = conf.get(FRONTEND_ADVERTISED_HOST).getOrElse(host)
+ if (flightServer != null && started.get()) {
+ locationFor(advertisedHost, flightServer.getPort)
+ } else {
+ locationFor(advertisedHost, configuredPort)
+ }
+ }
+
+ override def initialize(conf: KyuubiConf): Unit = synchronized {
+ this.conf = conf
+ configuredPort = this.conf.get(FRONTEND_FLIGHT_SQL_BIND_PORT)
+ allocator = new RootAllocator()
+ producer = new KyuubiFlightSqlProducer(
+ serverable.backendService,
+ allocator,
+ () => currentLocation,
+ this.conf)
+ super.initialize(this.conf)
+ }
+
+ override def start(): Unit = synchronized {
+ if (!started.get()) {
+ try {
+ val builder = FlightServer
+ .builder(allocator, configuredLocation, producer)
+ .headerAuthenticator(KyuubiFlightAuthHandler.create(conf))
+ .backpressureThreshold(10 * 1024 * 1024)
+
+ if (sslEnabled) {
+ val material = KyuubiFlightTlsUtils.resolve(conf)
+ KyuubiFlightTlsUtils.validateCertPresent(material)
+ tlsMaterial = Some(material)
+ builder.useTls(material.certFile, material.keyFile)
+ }
+
+ flightServer = builder.build().start()
+ started.set(true)
+ info(s"Flight SQL frontend service started at $connectionUrl" +
+ s" (tls=$sslEnabled)")
+ } catch {
+ case NonFatal(e) =>
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_CONN_FAIL))
+ if (flightServer != null) {
+ try flightServer.close()
+ catch {
+ case NonFatal(closeError) =>
+ warn("Failed to close Flight SQL server after startup failure", closeError)
+ }
+ flightServer = null
+ }
+ tlsMaterial.foreach(_.cleanup())
+ tlsMaterial = None
+ throw new KyuubiException("Cannot start Flight SQL frontend service", e)
+ }
+ }
+ super.start()
+ }
+
+ override def stop(): Unit = synchronized {
+ if (started.getAndSet(false)) {
+ if (producer != null) {
+ try producer.close()
+ catch {
+ case NonFatal(e) => warn("Failed to close Flight SQL producer", e)
+ }
+ }
+ if (flightServer != null) {
+ try flightServer.close()
+ catch {
+ case NonFatal(e) => warn("Failed to close Flight SQL server", e)
+ }
+ flightServer = null
+ }
+ tlsMaterial.foreach(_.cleanup())
+ tlsMaterial = None
+ if (allocator != null) {
+ try allocator.close()
+ catch {
+ case NonFatal(e) => warn("Failed to close Flight SQL allocator", e)
+ }
+ allocator = null
+ }
+ }
+ super.stop()
+ }
+
+ override def connectionUrl: String = {
+ checkInitialized()
+ val advertisedHost = conf.get(FRONTEND_ADVERTISED_HOST).getOrElse(host)
+ s"$advertisedHost:${if (flightServer != null) flightServer.getPort else configuredPort}"
+ }
+
+ override lazy val discoveryService: Option[Service] = {
+ if (ServiceDiscovery.supportServiceDiscovery(conf)) {
+ Some(new FlightSqlServiceDiscovery(this))
+ } else {
+ None
+ }
+ }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
index 5e7405be238..374258b5b3c 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
@@ -197,6 +197,9 @@ class KyuubiServer(name: String) extends Serverable(name) {
case TRINO =>
warn("Trino frontend protocol is experimental.")
new KyuubiTrinoFrontendService(this)
+ case FLIGHT_SQL =>
+ warn("FLIGHT_SQL frontend protocol is experimental.")
+ new KyuubiFlightSqlFrontendService(this)
case other =>
throw new UnsupportedOperationException(s"Frontend protocol $other is not supported yet.")
}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightResultIterator.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightResultIterator.scala
new file mode 100644
index 00000000000..8a17e8d3301
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightResultIterator.scala
@@ -0,0 +1,146 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.control.NonFatal
+
+import org.apache.arrow.flight.FlightProducer.ServerStreamListener
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot}
+import org.apache.arrow.vector.types.pojo.Schema
+
+import org.apache.kyuubi.Logging
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
+import org.apache.kyuubi.operation.{FetchOrientation, OperationHandle}
+import org.apache.kyuubi.service.BackendService
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TRowSet
+
+/**
+ * Page-oriented Flight result iterator. Retains only the current backend page /
+ * Arrow batch and never materializes an entire query result.
+ */
+class FlightResultIterator(
+ backend: BackendService,
+ operation: OperationHandle,
+ schema: Schema,
+ allocator: BufferAllocator,
+ pageSize: Int,
+ isCancelled: () => Boolean)
+ extends AutoCloseable with Logging {
+
+ private val root = VectorSchemaRoot.create(schema, allocator)
+ private val loader = new VectorLoader(root)
+ private val closed = new AtomicBoolean(false)
+ private var started = false
+ private var finished = false
+
+ def start(listener: ServerStreamListener): VectorSchemaRoot = {
+ if (!started) {
+ listener.start(root)
+ started = true
+ }
+ root
+ }
+
+ /**
+ * Fetch and load the next non-empty page into [[root]].
+ * @return true when a page was loaded, false when the stream is exhausted
+ */
+ def nextBatch(): Boolean = {
+ ensureOpen()
+ if (finished || isCancelled()) {
+ return false
+ }
+
+ while (!finished && !isCancelled()) {
+ val rowSet = backend.fetchResults(
+ operation,
+ FetchOrientation.FETCH_NEXT,
+ pageSize,
+ fetchLog = false).getResults
+ if (KyuubiFlightArrowUtils.isEmpty(rowSet)) {
+ finished = true
+ return false
+ }
+
+ root.clear()
+ val loadedRows =
+ if (KyuubiFlightArrowUtils.isArrowRowSet(rowSet)) {
+ loadArrowBatch(rowSet)
+ } else {
+ KyuubiFlightArrowUtils.populateRootFromRowSet(root, rowSet)
+ root.getRowCount
+ }
+
+ if (loadedRows > 0) {
+ MetricsSystem.tracing { ms =>
+ ms.markMeter(MetricsConstants.FLIGHT_SQL_STREAM_BATCHES)
+ ms.markMeter(MetricsConstants.FLIGHT_SQL_STREAM_ROWS, loadedRows)
+ }
+ return true
+ }
+
+ // Zero decoded rows means the backend page is exhausted (including an empty Arrow
+ // IPC batch). Do not keep fetching.
+ finished = true
+ return false
+ }
+ false
+ }
+
+ def currentRoot: VectorSchemaRoot = root
+
+ def cancel(): Unit = {
+ try backend.cancelOperation(operation)
+ catch {
+ case NonFatal(e) => warn(s"Failed to cancel Flight SQL operation $operation", e)
+ }
+ }
+
+ override def close(): Unit = {
+ if (closed.compareAndSet(false, true)) {
+ try root.close()
+ catch {
+ case NonFatal(e) => warn("Failed to close Flight result VectorSchemaRoot", e)
+ }
+ }
+ }
+
+ private def loadArrowBatch(rowSet: TRowSet): Int = {
+ val batch = KyuubiFlightArrowUtils.decodeBatch(rowSet, allocator)
+ try {
+ loader.load(batch)
+ root.setRowCount(batch.getLength)
+ val bytes = KyuubiFlightArrowUtils.arrowBatchBytes(rowSet)
+ if (bytes > 0) {
+ MetricsSystem.tracing(_.markMeter(MetricsConstants.FLIGHT_SQL_STREAM_BYTES, bytes))
+ }
+ batch.getLength
+ } finally {
+ batch.close()
+ }
+ }
+
+ private def ensureOpen(): Unit = {
+ if (closed.get()) {
+ throw new IllegalStateException("FlightResultIterator is closed")
+ }
+ }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightSqlKerberosValidator.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightSqlKerberosValidator.scala
new file mode 100644
index 00000000000..956b4978513
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/FlightSqlKerberosValidator.scala
@@ -0,0 +1,109 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.io.File
+import java.security.{PrivilegedActionException, PrivilegedExceptionAction}
+import java.util.Base64
+import javax.security.auth.Subject
+import javax.security.auth.kerberos.{KerberosPrincipal, KeyTab}
+import javax.security.sasl.AuthenticationException
+
+import org.apache.hadoop.security.authentication.util.KerberosName
+import org.apache.hadoop.security.authentication.util.KerberosUtil._
+import org.ietf.jgss.{GSSContext, GSSCredential, GSSManager, Oid}
+
+import org.apache.kyuubi.Logging
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.{SERVER_SPNEGO_KEYTAB, SERVER_SPNEGO_PRINCIPAL}
+import org.apache.kyuubi.util.KyuubiHadoopUtils
+
+/**
+ * Validates SPNEGO tokens for the Arrow Flight SQL gRPC frontend.
+ * Reuses the same GSS-API pattern as KerberosAuthenticationHandler adapted for Flight.
+ */
+class FlightSqlKerberosValidator(conf: KyuubiConf) extends Logging {
+
+ private val keytab = conf.get(SERVER_SPNEGO_KEYTAB).get
+ private val principal = KyuubiHadoopUtils.getServerPrincipal(
+ conf.get(SERVER_SPNEGO_PRINCIPAL).get)
+
+ private val serverSubject: Subject = {
+ val subject = new Subject()
+ subject.getPrivateCredentials.add(KeyTab.getInstance(new File(keytab)))
+ subject.getPrincipals.add(new KerberosPrincipal(principal))
+ subject
+ }
+
+ private val gssManager: GSSManager = Subject.doAs(
+ serverSubject,
+ new PrivilegedExceptionAction[GSSManager] {
+ override def run(): GSSManager = GSSManager.getInstance()
+ })
+
+ if (!KerberosName.hasRulesBeenSet) {
+ KerberosName.setRules("DEFAULT")
+ }
+
+ info(s"FlightSqlKerberosValidator initialized with principal $principal, keytab $keytab")
+
+ /**
+ * Validates a SPNEGO token from the Authorization: Negotiate header.
+ *
+ * @param base64Token the base64-encoded SPNEGO token (Negotiate prefix already stripped)
+ * @return the authenticated short username
+ */
+ def validate(base64Token: String): String = {
+ val clientToken = Base64.getDecoder.decode(base64Token)
+ try {
+ Subject.doAs(
+ serverSubject,
+ new PrivilegedExceptionAction[String] {
+ override def run(): String = validateToken(clientToken)
+ })
+ } catch {
+ case e: PrivilegedActionException =>
+ throw new AuthenticationException("SPNEGO authentication failed", e.getException)
+ case e: Exception =>
+ throw new AuthenticationException("SPNEGO authentication failed", e)
+ }
+ }
+
+ private def validateToken(clientToken: Array[Byte]): String = {
+ val serverPrincipalName = getTokenServerName(clientToken)
+ var gssContext: GSSContext = null
+ var gssCreds: GSSCredential = null
+ try {
+ gssCreds = gssManager.createCredential(
+ gssManager.createName(serverPrincipalName, NT_GSS_KRB5_PRINCIPAL_OID),
+ GSSCredential.INDEFINITE_LIFETIME,
+ Array[Oid](GSS_SPNEGO_MECH_OID, GSS_KRB5_MECH_OID),
+ GSSCredential.ACCEPT_ONLY)
+ gssContext = gssManager.createContext(gssCreds)
+ gssContext.acceptSecContext(clientToken, 0, clientToken.length)
+ if (!gssContext.isEstablished) {
+ throw new AuthenticationException("SPNEGO context wasn't fully established")
+ }
+ val clientPrincipal = gssContext.getSrcName.toString
+ new KerberosName(clientPrincipal).getShortName
+ } finally {
+ if (gssContext != null) gssContext.dispose()
+ if (gssCreds != null) gssCreds.dispose()
+ }
+ }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtils.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtils.scala
new file mode 100644
index 00000000000..1d9e117540e
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtils.scala
@@ -0,0 +1,328 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.io.ByteArrayInputStream
+import java.nio.ByteBuffer
+import java.nio.channels.Channels
+import java.nio.charset.StandardCharsets
+import java.sql.{Date, Timestamp}
+import java.time.{Instant, LocalDate, LocalDateTime, OffsetDateTime, ZonedDateTime}
+import java.util.{BitSet, Collections}
+
+import scala.collection.JavaConverters._
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector._
+import org.apache.arrow.vector.ipc.ReadChannel
+import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, MessageSerializer}
+import org.apache.arrow.vector.types.FloatingPointPrecision
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+import org.apache.arrow.vector.util.Text
+
+import org.apache.kyuubi.jdbc.hive.KyuubiArrowQueryResultSet
+import org.apache.kyuubi.jdbc.hive.arrow.ArrowUtils
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+
+/**
+ * Shared Kyuubi thrift/Arrow conversion helpers for Flight SQL.
+ *
+ * Schema and value conversion are engine-agnostic: engines expose either Arrow
+ * IPC inside a binary [[TRowSet]] column or ordinary columnar thrift values.
+ */
+object KyuubiFlightArrowUtils {
+
+ def schemaFromMetadata(metadata: TGetResultSetMetadataResp): Schema = {
+ val columns = Option(metadata).flatMap(m => Option(m.getSchema))
+ .map(_.getColumns.asScala.toSeq)
+ .getOrElse(Seq.empty)
+ new Schema(columns.map { column =>
+ val arrowType = Option(column.getTypeDesc)
+ .flatMap(desc => Option(desc.getTypes).flatMap(_.asScala.headOption))
+ .map(typeEntry => arrowTypeFromEntry(typeEntry, column.getColumnName))
+ .getOrElse(ArrowType.Utf8.INSTANCE)
+ new Field(
+ column.getColumnName,
+ new FieldType(true, arrowType, null),
+ Collections.emptyList[Field]())
+ }.asJava)
+ }
+
+ private def arrowTypeFromEntry(typeEntry: TTypeEntry, columnName: String): ArrowType = {
+ if (!typeEntry.isSetPrimitiveEntry) {
+ throw new IllegalArgumentException(
+ s"Unsupported non-primitive Flight SQL column type for '$columnName'")
+ }
+ val primitive = typeEntry.getPrimitiveEntry
+ val attributes = KyuubiArrowQueryResultSet.getColumnAttributes(primitive)
+ try {
+ ArrowUtils.toArrowType(primitive.getType, attributes)
+ } catch {
+ case e: Exception =>
+ throw new IllegalArgumentException(
+ s"Unsupported Flight SQL column type '${primitive.getType}' for '$columnName'",
+ e)
+ }
+ }
+
+ /**
+ * True when the row set carries no data rows.
+ *
+ * Columnar thrift generators
+ * (see [[org.apache.kyuubi.engine.result.TRowSetGenerator.toColumnBasedSet]])
+ * still emit one empty [[TColumn]] per schema field when the iterator is exhausted. Checking only
+ * `getColumnsSize == 0` would treat those pages as non-empty and can loop forever in
+ * [[FlightResultIterator]].
+ */
+ def isEmpty(rowSet: TRowSet): Boolean = rowCount(rowSet) == 0
+
+ /**
+ * Number of data rows in a thrift [[TRowSet]] (columnar values, row-based rows, or Arrow batch).
+ */
+ def rowCount(rowSet: TRowSet): Int = {
+ if (rowSet == null) {
+ 0
+ } else if (rowSet.getColumns != null && rowSet.getColumnsSize > 0) {
+ // Arrow IPC is encoded as a single binary value. A present Arrow binary value
+ // counts as non-empty at the thrift layer (decoded length may still be 0).
+ columnSize(rowSet.getColumns.get(0))
+ } else if (rowSet.getRows != null) {
+ rowSet.getRowsSize
+ } else {
+ 0
+ }
+ }
+
+ def isArrowRowSet(rowSet: TRowSet): Boolean =
+ rowSet != null &&
+ rowSet.getColumnsSize == 1 &&
+ rowSet.getColumns.get(0).isSetBinaryVal &&
+ rowSet.getColumns.get(0).getBinaryVal.getValuesSize == 1
+
+ def arrowBatchBytes(rowSet: TRowSet): Long = {
+ if (!isArrowRowSet(rowSet)) {
+ 0L
+ } else {
+ rowSet.getColumns.get(0).getBinaryVal.getValues.get(0).remaining().toLong
+ }
+ }
+
+ def decodeBatch(rowSet: TRowSet, allocator: BufferAllocator): ArrowRecordBatch = {
+ val buffer = rowSet.getColumns.get(0).getBinaryVal.getValues.get(0)
+ val bytes = new Array[Byte](buffer.remaining())
+ buffer.duplicate().get(bytes)
+ MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))),
+ allocator)
+ }
+
+ /**
+ * Populate [[root]] directly from a columnar thrift [[TRowSet]] without first
+ * building an intermediate `Seq[Seq[AnyRef]]`.
+ */
+ def populateRootFromRowSet(root: VectorSchemaRoot, rowSet: TRowSet): Unit = {
+ root.clear()
+ root.allocateNew()
+ if (rowSet == null) {
+ root.setRowCount(0)
+ return
+ }
+
+ if (rowSet.getColumns != null && rowSet.getColumnsSize > 0) {
+ val columns = rowSet.getColumns.asScala
+ val rowCount = columnSize(columns.head)
+ val vectors = root.getFieldVectors.asScala
+ val fields = root.getSchema.getFields.asScala
+ val paired = columns.zip(vectors.zip(fields)).take(math.min(columns.size, vectors.size))
+ paired.foreach { case (column, (vector, field)) =>
+ writeColumn(vector, field.getType, column)
+ vector.setValueCount(rowCount)
+ }
+ if (columns.size < vectors.size) {
+ vectors.drop(columns.size).foreach { vector =>
+ (0 until rowCount).foreach(vector.setNull)
+ vector.setValueCount(rowCount)
+ }
+ }
+ root.setRowCount(rowCount)
+ } else if (rowSet.getRows != null && !rowSet.getRows.isEmpty) {
+ val rows = rowSet.getRows.asScala
+ val vectors = root.getFieldVectors.asScala
+ val fields = root.getSchema.getFields.asScala
+ vectors.zip(fields).zipWithIndex.foreach { case ((vector, field), columnIndex) =>
+ rows.zipWithIndex.foreach { case (row, rowIndex) =>
+ val value =
+ if (row.getColValsSize > columnIndex) row.getColVals.get(columnIndex).getFieldValue
+ else null
+ setValue(vector, field.getType, rowIndex, value)
+ }
+ vector.setValueCount(rows.size)
+ }
+ root.setRowCount(rows.size)
+ } else {
+ root.setRowCount(0)
+ }
+ }
+
+ private def writeColumn(vector: FieldVector, arrowType: ArrowType, column: TColumn): Unit = {
+ val size = columnSize(column)
+ (0 until size).foreach { rowIndex =>
+ setValue(vector, arrowType, rowIndex, columnValue(column, rowIndex))
+ }
+ }
+
+ private def columnSize(column: TColumn): Int = {
+ if (column.isSetBoolVal) column.getBoolVal.getValuesSize
+ else if (column.isSetByteVal) column.getByteVal.getValuesSize
+ else if (column.isSetI16Val) column.getI16Val.getValuesSize
+ else if (column.isSetI32Val) column.getI32Val.getValuesSize
+ else if (column.isSetI64Val) column.getI64Val.getValuesSize
+ else if (column.isSetDoubleVal) column.getDoubleVal.getValuesSize
+ else if (column.isSetStringVal) column.getStringVal.getValuesSize
+ else if (column.isSetBinaryVal) column.getBinaryVal.getValuesSize
+ else 0
+ }
+
+ private def columnValue(column: TColumn, rowIndex: Int): AnyRef = {
+ def nullsOf(bytes: Array[Byte]): BitSet =
+ if (bytes == null) new BitSet() else BitSet.valueOf(bytes)
+
+ if (column.isSetBoolVal) {
+ if (nullsOf(column.getBoolVal.getNulls).get(rowIndex)) null
+ else column.getBoolVal.getValues.get(rowIndex)
+ } else if (column.isSetByteVal) {
+ if (nullsOf(column.getByteVal.getNulls).get(rowIndex)) null
+ else column.getByteVal.getValues.get(rowIndex)
+ } else if (column.isSetI16Val) {
+ if (nullsOf(column.getI16Val.getNulls).get(rowIndex)) null
+ else column.getI16Val.getValues.get(rowIndex)
+ } else if (column.isSetI32Val) {
+ if (nullsOf(column.getI32Val.getNulls).get(rowIndex)) null
+ else column.getI32Val.getValues.get(rowIndex)
+ } else if (column.isSetI64Val) {
+ if (nullsOf(column.getI64Val.getNulls).get(rowIndex)) null
+ else column.getI64Val.getValues.get(rowIndex)
+ } else if (column.isSetDoubleVal) {
+ if (nullsOf(column.getDoubleVal.getNulls).get(rowIndex)) null
+ else column.getDoubleVal.getValues.get(rowIndex)
+ } else if (column.isSetStringVal) {
+ if (nullsOf(column.getStringVal.getNulls).get(rowIndex)) null
+ else column.getStringVal.getValues.get(rowIndex)
+ } else if (column.isSetBinaryVal) {
+ if (nullsOf(column.getBinaryVal.getNulls).get(rowIndex)) null
+ else column.getBinaryVal.getValues.get(rowIndex)
+ } else {
+ null
+ }
+ }
+
+ private def setValue(
+ vector: FieldVector,
+ arrowType: ArrowType,
+ rowIndex: Int,
+ value: AnyRef): Unit = {
+ if (value == null) {
+ vector.setNull(rowIndex)
+ return
+ }
+
+ arrowType match {
+ case _: ArrowType.Utf8 =>
+ vector.asInstanceOf[VarCharVector].setSafe(rowIndex, new Text(value.toString))
+ case _: ArrowType.Binary =>
+ vector.asInstanceOf[VarBinaryVector].setSafe(rowIndex, bytes(value))
+ case _: ArrowType.Bool =>
+ val bit: Boolean = value match {
+ case b: java.lang.Boolean => b.booleanValue()
+ case s: String => s.toBoolean
+ case _ => number(value).intValue() != 0
+ }
+ vector.asInstanceOf[BitVector].setSafe(rowIndex, if (bit) 1 else 0)
+ case intType: ArrowType.Int =>
+ intType.getBitWidth match {
+ case 8 => vector.asInstanceOf[TinyIntVector].setSafe(rowIndex, number(value).byteValue())
+ case 16 =>
+ vector.asInstanceOf[SmallIntVector].setSafe(rowIndex, number(value).shortValue())
+ case 32 => vector.asInstanceOf[IntVector].setSafe(rowIndex, number(value).intValue())
+ case 64 => vector.asInstanceOf[BigIntVector].setSafe(rowIndex, number(value).longValue())
+ case _ => throw new IllegalArgumentException(s"Unsupported integer width $intType")
+ }
+ case floatingPoint: ArrowType.FloatingPoint =>
+ floatingPoint.getPrecision match {
+ case FloatingPointPrecision.SINGLE =>
+ vector.asInstanceOf[Float4Vector].setSafe(rowIndex, number(value).floatValue())
+ case FloatingPointPrecision.DOUBLE =>
+ vector.asInstanceOf[Float8Vector].setSafe(rowIndex, number(value).doubleValue())
+ case _ =>
+ throw new IllegalArgumentException(
+ s"Unsupported floating point type $arrowType")
+ }
+ case _: ArrowType.Decimal =>
+ vector.asInstanceOf[DecimalVector].setSafe(
+ rowIndex,
+ new java.math.BigDecimal(value.toString))
+ case _: ArrowType.Date =>
+ val days = value match {
+ case d: Date => (d.getTime / (24L * 60L * 60L * 1000L)).toInt
+ case ld: LocalDate => ld.toEpochDay.toInt
+ case _ => number(value).intValue()
+ }
+ vector.asInstanceOf[DateDayVector].setSafe(rowIndex, days)
+ case ts: ArrowType.Timestamp =>
+ val micros = timestampMicros(value)
+ vector.asInstanceOf[TimeStampVector].setSafe(rowIndex, micros)
+ // Preserve timezone metadata already present on the Arrow field type.
+ val _ = ts
+ case _: ArrowType.Null =>
+ vector.setNull(rowIndex)
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unsupported Arrow type $other for Flight SQL value conversion")
+ }
+ }
+
+ private def timestampMicros(value: AnyRef): Long = value match {
+ case ts: Timestamp => ts.toInstant.getEpochSecond * 1000000L + ts.getNanos / 1000L
+ case instant: Instant => instant.getEpochSecond * 1000000L + instant.getNano / 1000L
+ case ldt: LocalDateTime =>
+ timestampMicros(Timestamp.valueOf(ldt))
+ case odt: OffsetDateTime => timestampMicros(odt.toInstant)
+ case zdt: ZonedDateTime => timestampMicros(zdt.toInstant)
+ case n: Number => n.longValue()
+ case s: String => timestampMicros(Timestamp.valueOf(s))
+ case other =>
+ throw new IllegalArgumentException(s"Unsupported timestamp value ${other.getClass}")
+ }
+
+ private def bytes(value: AnyRef): Array[Byte] = value match {
+ case b: Array[Byte] => b
+ case buffer: ByteBuffer =>
+ val duplicate = buffer.duplicate()
+ val result = new Array[Byte](duplicate.remaining())
+ duplicate.get(result)
+ result
+ case text: Text => text.toString.getBytes(StandardCharsets.UTF_8)
+ case other => other.toString.getBytes(StandardCharsets.UTF_8)
+ }
+
+ private def number(value: AnyRef): Number = value match {
+ case n: Number => n
+ case s: String => BigDecimal(s).bigDecimal
+ case other => throw new IllegalArgumentException(s"Expected a number, got ${other.getClass}")
+ }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandler.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandler.scala
new file mode 100644
index 00000000000..789eedb7a5f
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandler.scala
@@ -0,0 +1,190 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.nio.charset.StandardCharsets
+import java.util.Base64
+import java.util.concurrent.TimeUnit
+
+import scala.util.control.NonFatal
+
+import com.google.common.cache.CacheBuilder
+import org.apache.arrow.flight.{CallHeaders, CallStatus}
+import org.apache.arrow.flight.auth2.{Auth2Constants, CallHeaderAuthenticator, GeneratedBearerTokenAuthenticator}
+
+import org.apache.kyuubi.Logging
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.{AUTHENTICATION_METHOD, FRONTEND_FLIGHT_SQL_TOKEN_TTL}
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
+import org.apache.kyuubi.service.authentication.{AuthenticationProviderFactory, AuthMethods, AuthTypes, AuthUtils}
+
+/**
+ * Flight CallHeaderAuthenticator supporting:
+ * - Basic username/password (LDAP/JDBC/CUSTOM/NONE providers)
+ * - SPNEGO Negotiate bootstrap using [[FlightSqlKerberosValidator]]
+ * - Bearer tokens issued after successful Basic/Negotiate authentication
+ */
+object KyuubiFlightAuthHandler {
+ private val NEGOTIATE_PREFIX = "Negotiate "
+
+ def create(conf: KyuubiConf): CallHeaderAuthenticator = {
+ val delegate = new KyuubiFlightAuthHandler(conf)
+ if (delegate.issuesBearerTokens) {
+ val ttlMs = conf.get(FRONTEND_FLIGHT_SQL_TOKEN_TTL)
+ val ttlMinutes = math.max(1L, TimeUnit.MILLISECONDS.toMinutes(ttlMs))
+ new GeneratedBearerTokenAuthenticator(
+ delegate,
+ CacheBuilder.newBuilder().expireAfterAccess(ttlMinutes, TimeUnit.MINUTES))
+ } else {
+ delegate
+ }
+ }
+}
+
+class KyuubiFlightAuthHandler(conf: KyuubiConf)
+ extends CallHeaderAuthenticator with Logging {
+
+ private val authTypes =
+ conf.get(AUTHENTICATION_METHOD).map(value => AuthTypes.withName(value))
+
+ private val noAuthRequired =
+ AuthUtils.saslDisabled(authTypes) ||
+ AuthUtils.effectivePlainAuthType(authTypes).contains(AuthTypes.NONE)
+
+ private val passwordProvider = AuthUtils.effectivePlainAuthType(authTypes)
+ .filterNot(_ == AuthTypes.NONE)
+ .map(authType =>
+ AuthenticationProviderFactory.getAuthenticationProvider(
+ AuthMethods.withName(authType.toString),
+ conf,
+ isServer = true))
+
+ private val kerberosEnabled = AuthUtils.kerberosEnabled(authTypes)
+
+ private lazy val kerberosValidator: Option[FlightSqlKerberosValidator] =
+ if (kerberosEnabled) {
+ try Some(new FlightSqlKerberosValidator(conf))
+ catch {
+ case NonFatal(e) =>
+ throw new IllegalArgumentException(
+ "Flight SQL Kerberos is enabled but SPNEGO principal/keytab are not usable",
+ e)
+ }
+ } else {
+ None
+ }
+
+ private[flight] def issuesBearerTokens: Boolean =
+ passwordProvider.isDefined || kerberosValidator.isDefined
+
+ override def authenticate(headers: CallHeaders): CallHeaderAuthenticator.AuthResult = {
+ val authorization = Option(headers.get(Auth2Constants.AUTHORIZATION_HEADER))
+ try {
+ val result = authorization match {
+ case Some(value) if value.startsWith(Auth2Constants.BASIC_PREFIX) =>
+ authenticateBasic(value.stripPrefix(Auth2Constants.BASIC_PREFIX))
+ case Some(value)
+ if value.regionMatches(
+ true,
+ 0,
+ KyuubiFlightAuthHandler.NEGOTIATE_PREFIX,
+ 0,
+ KyuubiFlightAuthHandler.NEGOTIATE_PREFIX.length) =>
+ authenticateNegotiate(
+ value.substring(KyuubiFlightAuthHandler.NEGOTIATE_PREFIX.length))
+ case Some(value) if value.startsWith(Auth2Constants.BEARER_PREFIX) =>
+ // Bearer validation is handled by GeneratedBearerTokenAuthenticator when enabled.
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Bearer authentication requires a previously issued Flight token")
+ .toRuntimeException
+ case Some(_) =>
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Unsupported Flight SQL authorization scheme")
+ .toRuntimeException
+ case None if noAuthRequired =>
+ val user = Option(headers.get("x-user-name"))
+ .filter(_.nonEmpty)
+ .getOrElse("anonymous")
+ authResult(user)
+ case None =>
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Missing Flight SQL authorization header")
+ .toRuntimeException
+ }
+ MetricsSystem.tracing { ms =>
+ ms.incCount(MetricsConstants.FLIGHT_SQL_CONN_TOTAL)
+ ms.incCount(MetricsConstants.FLIGHT_SQL_CONN_OPEN)
+ }
+ result
+ } catch {
+ case e: RuntimeException =>
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_CONN_FAIL))
+ throw e
+ }
+ }
+
+ private def authenticateBasic(encoded: String): CallHeaderAuthenticator.AuthResult = {
+ if (passwordProvider.isEmpty) {
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Basic authentication is not configured for Flight SQL")
+ .toRuntimeException
+ }
+ val decoded =
+ try {
+ new String(Base64.getDecoder.decode(encoded), StandardCharsets.UTF_8)
+ } catch {
+ case e: IllegalArgumentException =>
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Malformed Flight SQL basic credentials")
+ .withCause(e)
+ .toRuntimeException
+ }
+ val separator = decoded.indexOf(':')
+ if (separator <= 0) {
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Malformed Flight SQL basic credentials")
+ .toRuntimeException
+ }
+ val user = decoded.substring(0, separator)
+ val password = decoded.substring(separator + 1)
+ passwordProvider.get.authenticate(user, password)
+ authResult(user)
+ }
+
+ private def authenticateNegotiate(token: String): CallHeaderAuthenticator.AuthResult = {
+ val validator = kerberosValidator.getOrElse {
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Kerberos authentication is not configured for Flight SQL")
+ .toRuntimeException
+ }
+ try {
+ authResult(validator.validate(token.trim))
+ } catch {
+ case NonFatal(e) =>
+ throw CallStatus.UNAUTHENTICATED
+ .withDescription("Flight SQL Kerberos authentication failed")
+ .withCause(e)
+ .toRuntimeException
+ }
+ }
+
+ private def authResult(user: String): CallHeaderAuthenticator.AuthResult =
+ new CallHeaderAuthenticator.AuthResult {
+ override def getPeerIdentity: String = user
+ }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlProducer.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlProducer.scala
new file mode 100644
index 00000000000..6b2d6a3a1c7
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlProducer.scala
@@ -0,0 +1,495 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.nio.charset.StandardCharsets
+import java.util.UUID
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
+import com.google.protobuf.{Any => ProtoAny, ByteString}
+import org.apache.arrow.flight.{CallStatus, FlightDescriptor, FlightEndpoint, FlightInfo, Location, SchemaResult, Ticket}
+import org.apache.arrow.flight.FlightProducer.{CallContext, ServerStreamListener, StreamListener}
+import org.apache.arrow.flight.sql.{CancelResult, NoOpFlightSqlProducer, SqlInfoBuilder}
+import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas
+import org.apache.arrow.flight.sql.impl.FlightSql._
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.types.pojo.Schema
+
+import org.apache.kyuubi.{KYUUBI_VERSION, Logging}
+import org.apache.kyuubi.config.{KyuubiConf, KyuubiReservedKeys}
+import org.apache.kyuubi.config.KyuubiConf.{FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS, OPERATION_RESULT_FORMAT}
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
+import org.apache.kyuubi.operation.{OperationHandle, OperationState, OperationStatus}
+import org.apache.kyuubi.service.BackendService
+import org.apache.kyuubi.session.SessionHandle
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion
+
+class KyuubiFlightSqlProducer(
+ backend: BackendService,
+ allocator: BufferAllocator,
+ location: () => Location,
+ conf: KyuubiConf)
+ extends NoOpFlightSqlProducer with Logging {
+
+ private case class QueryState(
+ owner: String,
+ session: SessionHandle,
+ operation: OperationHandle,
+ schema: Schema,
+ ownerEndpoint: String)
+
+ private val queryStates = new ConcurrentHashMap[String, QueryState]()
+
+ private val fetchPageSize = conf.get(FRONTEND_FLIGHT_SQL_FETCH_MAX_ROWS)
+
+ private val sqlInfoBuilder = new SqlInfoBuilder()
+ .withFlightSqlServerName("Kyuubi")
+ .withFlightSqlServerVersion(KYUUBI_VERSION)
+ .withFlightSqlServerArrowVersion("16.0.0")
+ .withFlightSqlServerSql(true)
+ .withFlightSqlServerSubstrait(false)
+ .withFlightSqlServerCancel(true)
+ .withFlightSqlServerTransaction(SqlSupportedTransaction.SQL_SUPPORTED_TRANSACTION_NONE)
+ .withSqlIdentifierQuoteChar("`")
+ .withSqlIdentifierCase(SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN)
+ .withSqlQuotedIdentifierCase(SqlSupportedCaseSensitivity.SQL_CASE_SENSITIVITY_UNKNOWN)
+ .withSqlNullOrdering(SqlNullOrdering.SQL_NULLS_SORTED_AT_END)
+
+ override def getFlightInfoStatement(
+ command: CommandStatementQuery,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo = {
+ MetricsSystem.tracing { ms =>
+ ms.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_TOTAL)
+ ms.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_OPEN)
+ }
+ try {
+ val owner = ownerOf(context)
+ val session = openSession(owner)
+ try {
+ val operation = backend.executeStatement(
+ session,
+ command.getQuery,
+ Map.empty,
+ runAsync = true,
+ queryTimeout = 0)
+ val status = waitForCompletion(operation, () => context.isCancelled)
+ val schema = KyuubiFlightArrowUtils.schemaFromMetadata(
+ backend.getResultSetMetadata(operation))
+ if (status.state != OperationState.FINISHED) {
+ closeResources(session, operation)
+ MetricsSystem.tracing(_.decCount(MetricsConstants.FLIGHT_SQL_OPERATION_OPEN))
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_FAIL))
+ throw operationError(status)
+ }
+ val id = UUID.randomUUID().toString
+ val endpointLocation = location()
+ queryStates.put(
+ id,
+ QueryState(owner, session, operation, schema, endpointLocation.getUri.toString))
+ val ticket = statementTicket(id)
+ new FlightInfo(
+ schema,
+ descriptor,
+ java.util.Arrays.asList(new FlightEndpoint(ticket, endpointLocation)),
+ -1L,
+ -1L)
+ } catch {
+ case NonFatal(e) =>
+ try backend.closeSession(session)
+ catch {
+ case NonFatal(closeError) => warn("Failed to close Flight SQL session", closeError)
+ }
+ throw e
+ }
+ } catch {
+ case e: RuntimeException => throw e
+ case NonFatal(e) => throw flightError(CallStatus.INTERNAL, e)
+ }
+ }
+
+ override def getSchemaStatement(
+ command: CommandStatementQuery,
+ context: CallContext,
+ descriptor: FlightDescriptor): SchemaResult = {
+ val owner = ownerOf(context)
+ val session = openSession(owner)
+ var operation: OperationHandle = null
+ try {
+ operation = backend.executeStatement(
+ session,
+ command.getQuery,
+ Map.empty,
+ runAsync = true,
+ queryTimeout = 0)
+ val status = waitForCompletion(operation, () => context.isCancelled)
+ if (status.state != OperationState.FINISHED) {
+ throw operationError(status)
+ }
+ new SchemaResult(KyuubiFlightArrowUtils.schemaFromMetadata(
+ backend.getResultSetMetadata(operation)))
+ } catch {
+ case e: RuntimeException => throw e
+ case NonFatal(e) => throw flightError(CallStatus.INTERNAL, e)
+ } finally {
+ closeResources(session, operation)
+ }
+ }
+
+ override def getStreamStatement(
+ ticket: TicketStatementQuery,
+ context: CallContext,
+ listener: ServerStreamListener): Unit = {
+ val id = ticket.getStatementHandle.toStringUtf8
+ val state = Option(queryStates.get(id)).getOrElse(
+ throw flightError(
+ CallStatus.NOT_FOUND,
+ new IllegalArgumentException("Unknown Flight SQL ticket")))
+ if (state.owner != ownerOf(context)) {
+ throw flightError(
+ CallStatus.UNAUTHORIZED,
+ new SecurityException("Flight SQL ticket owner mismatch"))
+ }
+ val current = location().getUri.toString
+ if (state.ownerEndpoint != current) {
+ throw flightError(
+ CallStatus.UNAVAILABLE,
+ new IllegalStateException(
+ s"Flight SQL ticket is owned by ${state.ownerEndpoint}, not $current. " +
+ "Retry against the owning endpoint; transparent failover is not supported."))
+ }
+
+ val iterator = new FlightResultIterator(
+ backend,
+ state.operation,
+ state.schema,
+ allocator,
+ fetchPageSize,
+ () => context.isCancelled || listener.isCancelled)
+ listener.setOnCancelHandler(() => {
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_CANCELLED))
+ iterator.cancel()
+ })
+ try {
+ iterator.start(listener)
+ var ok = true
+ while (ok && !context.isCancelled && !listener.isCancelled) {
+ ok = iterator.nextBatch()
+ if (ok && iterator.currentRoot.getRowCount > 0) {
+ listener.putNext()
+ }
+ }
+ if (context.isCancelled || listener.isCancelled) {
+ iterator.cancel()
+ } else {
+ listener.completed()
+ }
+ } catch {
+ case NonFatal(e) =>
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_FAIL))
+ if (!context.isCancelled && !listener.isCancelled) listener.error(e)
+ } finally {
+ iterator.close()
+ MetricsSystem.tracing(_.decCount(MetricsConstants.FLIGHT_SQL_OPERATION_OPEN))
+ closeState(id, state)
+ }
+ }
+
+ override def getFlightInfoCatalogs(
+ command: CommandGetCatalogs,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo =
+ metadataInfo(Schemas.GET_CATALOGS_SCHEMA, command, descriptor)
+
+ override def getStreamCatalogs(
+ context: CallContext,
+ listener: ServerStreamListener): Unit =
+ streamMetadata(
+ context,
+ listener,
+ Schemas.GET_CATALOGS_SCHEMA,
+ session =>
+ backend.getCatalogs(session))
+
+ override def getFlightInfoSchemas(
+ command: CommandGetDbSchemas,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo =
+ metadataInfo(Schemas.GET_SCHEMAS_SCHEMA, command, descriptor)
+
+ override def getStreamSchemas(
+ command: CommandGetDbSchemas,
+ context: CallContext,
+ listener: ServerStreamListener): Unit = {
+ val catalog = if (command.hasCatalog) command.getCatalog else null
+ val schema = if (command.hasDbSchemaFilterPattern) command.getDbSchemaFilterPattern else null
+ streamMetadata(
+ context,
+ listener,
+ Schemas.GET_SCHEMAS_SCHEMA,
+ s =>
+ backend.getSchemas(s, Option(catalog).getOrElse(""), Option(schema).getOrElse("")))
+ }
+
+ override def getFlightInfoTables(
+ command: CommandGetTables,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo = {
+ val schema = if (command.getIncludeSchema) Schemas.GET_TABLES_SCHEMA
+ else Schemas.GET_TABLES_SCHEMA_NO_SCHEMA
+ metadataInfo(schema, command, descriptor)
+ }
+
+ override def getStreamTables(
+ command: CommandGetTables,
+ context: CallContext,
+ listener: ServerStreamListener): Unit = {
+ val catalog = if (command.hasCatalog) command.getCatalog else null
+ val schema = if (command.hasDbSchemaFilterPattern) command.getDbSchemaFilterPattern else null
+ val table = if (command.hasTableNameFilterPattern) command.getTableNameFilterPattern else null
+ streamMetadata(
+ context,
+ listener,
+ if (command.getIncludeSchema) Schemas.GET_TABLES_SCHEMA
+ else Schemas.GET_TABLES_SCHEMA_NO_SCHEMA,
+ s =>
+ backend.getTables(
+ s,
+ Option(catalog).getOrElse(""),
+ Option(schema).getOrElse(""),
+ Option(table).getOrElse(""),
+ command.getTableTypesList))
+ }
+
+ override def getFlightInfoTableTypes(
+ command: CommandGetTableTypes,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo =
+ metadataInfo(Schemas.GET_TABLE_TYPES_SCHEMA, command, descriptor)
+
+ override def getStreamTableTypes(
+ context: CallContext,
+ listener: ServerStreamListener): Unit =
+ streamMetadata(
+ context,
+ listener,
+ Schemas.GET_TABLE_TYPES_SCHEMA,
+ session =>
+ backend.getTableTypes(session))
+
+ override def getFlightInfoTypeInfo(
+ command: CommandGetXdbcTypeInfo,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo =
+ metadataInfo(Schemas.GET_TYPE_INFO_SCHEMA, command, descriptor)
+
+ override def getFlightInfoSqlInfo(
+ command: CommandGetSqlInfo,
+ context: CallContext,
+ descriptor: FlightDescriptor): FlightInfo =
+ metadataInfo(Schemas.GET_SQL_INFO_SCHEMA, command, descriptor)
+
+ override def getStreamSqlInfo(
+ command: CommandGetSqlInfo,
+ context: CallContext,
+ listener: ServerStreamListener): Unit =
+ sqlInfoBuilder.send(command.getInfoList, listener)
+
+ override def getStreamTypeInfo(
+ command: CommandGetXdbcTypeInfo,
+ context: CallContext,
+ listener: ServerStreamListener): Unit =
+ streamMetadata(
+ context,
+ listener,
+ Schemas.GET_TYPE_INFO_SCHEMA,
+ session =>
+ backend.getTypeInfo(session))
+
+ override def cancelQuery(
+ info: FlightInfo,
+ context: CallContext,
+ listener: StreamListener[CancelResult]): Unit = {
+ try {
+ val endpoint = info.getEndpoints.get(0)
+ val any = ProtoAny.parseFrom(endpoint.getTicket.getBytes)
+ val statementTicket = TicketStatementQuery.parseFrom(any.getValue)
+ val id = statementTicket.getStatementHandle.toStringUtf8
+ Option(queryStates.get(id)) match {
+ case Some(state) if state.owner == ownerOf(context) =>
+ cancelState(state)
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_CANCELLED))
+ listener.onNext(CancelResult.CANCELLED)
+ case Some(_) =>
+ listener.onNext(CancelResult.NOT_CANCELLABLE)
+ case None =>
+ listener.onNext(CancelResult.NOT_CANCELLABLE)
+ }
+ listener.onCompleted()
+ } catch {
+ case NonFatal(e) => listener.onError(flightError(CallStatus.INVALID_ARGUMENT, e))
+ }
+ }
+
+ override def close(): Unit = {
+ queryStates.entrySet().asScala.foreach(entry => closeState(entry.getKey, entry.getValue))
+ queryStates.clear()
+ }
+
+ private def metadataInfo(
+ schema: Schema,
+ command: com.google.protobuf.Message,
+ descriptor: FlightDescriptor): FlightInfo = {
+ val ticket = new Ticket(ProtoAny.pack(command).toByteArray)
+ new FlightInfo(
+ schema,
+ descriptor,
+ java.util.Arrays.asList(new FlightEndpoint(ticket, location())),
+ -1L,
+ -1L)
+ }
+
+ private def streamMetadata(
+ context: CallContext,
+ listener: ServerStreamListener,
+ schema: Schema,
+ operationFactory: SessionHandle => OperationHandle): Unit = {
+ val owner = ownerOf(context)
+ val session = openSession(owner)
+ var operation: OperationHandle = null
+ var iterator: FlightResultIterator = null
+ MetricsSystem.tracing { ms =>
+ ms.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_TOTAL)
+ ms.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_OPEN)
+ }
+ try {
+ operation = operationFactory(session)
+ val status = waitForCompletion(operation, () => context.isCancelled || listener.isCancelled)
+ if (status.state != OperationState.FINISHED) throw operationError(status)
+ iterator = new FlightResultIterator(
+ backend,
+ operation,
+ schema,
+ allocator,
+ fetchPageSize,
+ () => context.isCancelled || listener.isCancelled)
+ listener.setOnCancelHandler(() => {
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_CANCELLED))
+ iterator.cancel()
+ })
+ iterator.start(listener)
+ var ok = true
+ while (ok && !context.isCancelled && !listener.isCancelled) {
+ ok = iterator.nextBatch()
+ if (ok && iterator.currentRoot.getRowCount > 0) {
+ listener.putNext()
+ }
+ }
+ if (context.isCancelled || listener.isCancelled) {
+ iterator.cancel()
+ } else {
+ listener.completed()
+ }
+ } catch {
+ case NonFatal(e) =>
+ MetricsSystem.tracing(_.incCount(MetricsConstants.FLIGHT_SQL_OPERATION_FAIL))
+ if (!context.isCancelled && !listener.isCancelled) listener.error(e)
+ } finally {
+ if (iterator != null) iterator.close()
+ MetricsSystem.tracing(_.decCount(MetricsConstants.FLIGHT_SQL_OPERATION_OPEN))
+ closeResources(session, operation)
+ }
+ }
+
+ private def openSession(owner: String): SessionHandle = {
+ backend.openSession(
+ TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V10,
+ owner,
+ "",
+ "",
+ Map(
+ KyuubiReservedKeys.KYUUBI_CLIENT_IP_KEY -> "",
+ KyuubiReservedKeys.KYUUBI_SESSION_REAL_USER_KEY -> owner,
+ KyuubiReservedKeys.KYUUBI_SESSION_CONNECTION_URL_KEY -> location().toString,
+ OPERATION_RESULT_FORMAT.key -> "arrow",
+ KyuubiConf.FRONTEND_PROTOCOLS.key ->
+ KyuubiConf.FrontendProtocols.FLIGHT_SQL.toString))
+ }
+
+ private def waitForCompletion(
+ operation: OperationHandle,
+ isCancelled: () => Boolean): OperationStatus = {
+ var status = backend.getOperationStatus(operation, Some(1000L))
+ while (!OperationState.isTerminal(status.state) && !isCancelled()) {
+ status = backend.getOperationStatus(operation, Some(1000L))
+ }
+ status
+ }
+
+ private def statementTicket(id: String): Ticket = {
+ val statementTicket = TicketStatementQuery.newBuilder
+ .setStatementHandle(ByteString.copyFrom(id.getBytes(StandardCharsets.UTF_8)))
+ .build()
+ new Ticket(ProtoAny.pack(statementTicket).toByteArray)
+ }
+
+ private def ownerOf(context: CallContext): String =
+ Option(context.peerIdentity()).filter(_.nonEmpty).getOrElse("anonymous")
+
+ private def closeState(id: String, state: QueryState): Unit = {
+ if (queryStates.remove(id, state)) {
+ closeResources(state.session, state.operation)
+ }
+ }
+
+ private def closeResources(session: SessionHandle, operation: OperationHandle): Unit = {
+ if (operation != null) {
+ try backend.closeOperation(operation)
+ catch {
+ case NonFatal(e) => warn(s"Failed to close Flight SQL operation $operation", e)
+ }
+ }
+ if (session != null) {
+ try backend.closeSession(session)
+ catch {
+ case NonFatal(e) => warn(s"Failed to close Flight SQL session $session", e)
+ }
+ }
+ }
+
+ private def cancelState(state: QueryState): Unit = {
+ try backend.cancelOperation(state.operation)
+ catch {
+ case NonFatal(e) => warn(s"Failed to cancel Flight SQL operation ${state.operation}", e)
+ }
+ }
+
+ private def operationError(status: OperationStatus): RuntimeException =
+ flightError(
+ CallStatus.INTERNAL,
+ status.exception.getOrElse(
+ new IllegalStateException(s"Flight SQL operation ended in ${status.state}")))
+
+ private def flightError(status: CallStatus, cause: Throwable): RuntimeException =
+ status.withDescription(Option(cause.getMessage).getOrElse(cause.getClass.getSimpleName))
+ .withCause(cause)
+ .toRuntimeException
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightTlsUtils.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightTlsUtils.scala
new file mode 100644
index 00000000000..e6e17768859
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/flight/KyuubiFlightTlsUtils.scala
@@ -0,0 +1,146 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.io.{File, FileOutputStream, OutputStreamWriter}
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
+import java.security.KeyStore
+import java.util.Base64
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
+import org.apache.kyuubi.Logging
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
+
+/**
+ * Resolves PEM material for Arrow Flight TLS. Prefers explicit PEM files and
+ * can materialize temporary PEM files from the shared Java keystore settings.
+ */
+object KyuubiFlightTlsUtils extends Logging {
+
+ case class TlsMaterial(certFile: File, keyFile: File, temporary: Boolean) {
+ def cleanup(): Unit = if (temporary) {
+ Seq(certFile, keyFile).foreach { file =>
+ try Files.deleteIfExists(file.toPath)
+ catch {
+ case NonFatal(e) => warn(s"Failed to delete temporary Flight TLS file $file", e)
+ }
+ }
+ }
+ }
+
+ def resolve(conf: KyuubiConf): TlsMaterial = {
+ val cert = conf.get(FRONTEND_FLIGHT_SQL_SSL_CERT_FILE)
+ val key = conf.get(FRONTEND_FLIGHT_SQL_SSL_KEY_FILE)
+ if (cert.isDefined || key.isDefined) {
+ if (cert.isEmpty || key.isEmpty) {
+ throw new IllegalArgumentException(
+ s"Both ${FRONTEND_FLIGHT_SQL_SSL_CERT_FILE.key} and " +
+ s"${FRONTEND_FLIGHT_SQL_SSL_KEY_FILE.key} must be set for Flight SQL TLS")
+ }
+ val certFile = new File(cert.get)
+ val keyFile = new File(key.get)
+ if (!certFile.isFile) {
+ throw new IllegalArgumentException(s"Flight SQL TLS certificate not found: ${cert.get}")
+ }
+ if (!keyFile.isFile) {
+ throw new IllegalArgumentException(s"Flight SQL TLS private key not found: ${key.get}")
+ }
+ TlsMaterial(certFile, keyFile, temporary = false)
+ } else {
+ materializeFromKeystore(conf)
+ }
+ }
+
+ private def materializeFromKeystore(conf: KyuubiConf): TlsMaterial = {
+ val keyStorePath = conf.get(FRONTEND_SSL_KEYSTORE_PATH).getOrElse {
+ throw new IllegalArgumentException(
+ s"${FRONTEND_FLIGHT_SQL_SSL_CERT_FILE.key}/${FRONTEND_FLIGHT_SQL_SSL_KEY_FILE.key} or " +
+ s"${FRONTEND_SSL_KEYSTORE_PATH.key} must be configured when Flight SQL TLS is enabled")
+ }
+ val keyStorePassword = conf.get(FRONTEND_SSL_KEYSTORE_PASSWORD).getOrElse {
+ throw new IllegalArgumentException(
+ s"${FRONTEND_SSL_KEYSTORE_PASSWORD.key} must be configured " +
+ "for Flight SQL TLS keystore fallback")
+ }
+ val keyStoreType = conf.get(FRONTEND_SSL_KEYSTORE_TYPE).getOrElse(KeyStore.getDefaultType)
+ val keyStore = KeyStore.getInstance(keyStoreType)
+ val input = Files.newInputStream(new File(keyStorePath).toPath)
+ try {
+ keyStore.load(input, keyStorePassword.toCharArray)
+ } finally {
+ input.close()
+ }
+
+ val aliases = keyStore.aliases().asScala.toSeq
+ val alias = aliases.find(a => keyStore.isKeyEntry(a)).getOrElse {
+ throw new IllegalArgumentException(
+ s"No private key entry found in keystore $keyStorePath for Flight SQL TLS")
+ }
+ val key = keyStore.getKey(alias, keyStorePassword.toCharArray)
+ val chain = keyStore.getCertificateChain(alias)
+ if (key == null || chain == null || chain.isEmpty) {
+ throw new IllegalArgumentException(
+ s"Keystore entry $alias does not contain a usable certificate chain/private key")
+ }
+
+ val certFile = Files.createTempFile("kyuubi-flight-cert-", ".pem").toFile
+ val keyFile = Files.createTempFile("kyuubi-flight-key-", ".pem").toFile
+ certFile.deleteOnExit()
+ keyFile.deleteOnExit()
+ writePem(certFile, "CERTIFICATE", chain.map(_.getEncoded))
+ writePem(keyFile, "PRIVATE KEY", Array(key.getEncoded))
+ // Best-effort restrictive permissions on POSIX systems.
+ try {
+ certFile.setReadable(false, false)
+ certFile.setReadable(true, true)
+ keyFile.setReadable(false, false)
+ keyFile.setReadable(true, true)
+ keyFile.setWritable(false, false)
+ keyFile.setWritable(true, true)
+ } catch {
+ case NonFatal(_) => // ignore on non-POSIX filesystems
+ }
+ TlsMaterial(certFile, keyFile, temporary = true)
+ }
+
+ private def writePem(file: File, label: String, derBlocks: Array[Array[Byte]]): Unit = {
+ val writer = new OutputStreamWriter(new FileOutputStream(file), StandardCharsets.US_ASCII)
+ try {
+ derBlocks.foreach { der =>
+ writer.write(s"-----BEGIN $label-----\n")
+ val encoded = Base64.getMimeEncoder(64, "\n".getBytes(StandardCharsets.US_ASCII))
+ .encodeToString(der)
+ writer.write(encoded)
+ writer.write(s"\n-----END $label-----\n")
+ }
+ } finally {
+ writer.close()
+ }
+ }
+
+ def validateCertPresent(material: TlsMaterial): Unit = {
+ val bytes = Files.readAllBytes(material.certFile.toPath)
+ if (bytes.isEmpty) {
+ throw new IllegalArgumentException("Flight SQL TLS certificate file is empty")
+ }
+ }
+}
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/WithFlightSqlServer.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/WithFlightSqlServer.scala
new file mode 100644
index 00000000000..777b06c855d
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/WithFlightSqlServer.scala
@@ -0,0 +1,33 @@
+/*
+ * 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.kyuubi
+
+import org.apache.kyuubi.config.KyuubiConf._
+
+trait WithFlightSqlServer extends WithKyuubiServer {
+
+ override protected val frontendProtocols =
+ Seq(FrontendProtocols.FLIGHT_SQL)
+
+ override def beforeAll(): Unit = {
+ conf.set(FRONTEND_FLIGHT_SQL_BIND_PORT, 0)
+ super.beforeAll()
+ }
+
+ protected def flightSqlUrl: String = server.frontendServices.head.connectionUrl
+}
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendServiceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendServiceSuite.scala
new file mode 100644
index 00000000000..ee57f6fcf7c
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/KyuubiFlightSqlFrontendServiceSuite.scala
@@ -0,0 +1,86 @@
+/*
+ * 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.kyuubi.server
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.service.ServiceState._
+
+class KyuubiFlightSqlFrontendServiceSuite extends KyuubiFunSuite {
+
+ test("Flight SQL frontend lifecycle") {
+ val server = new KyuubiServer
+ val conf = KyuubiConf()
+ .set(FRONTEND_PROTOCOLS, Seq(FrontendProtocols.FLIGHT_SQL.toString))
+ .set(FRONTEND_FLIGHT_SQL_BIND_HOST.key, "localhost")
+ .set(FRONTEND_FLIGHT_SQL_BIND_PORT, 0)
+
+ assert(server.getServiceState === LATENT)
+ server.initialize(conf)
+ assert(server.getServiceState === INITIALIZED)
+ assert(server.frontendServices.size === 1)
+ val frontend = server.frontendServices.head
+ assert(frontend.getServiceState === INITIALIZED)
+ assert(frontend.connectionUrl.startsWith("localhost:"))
+
+ server.start()
+ assert(server.getServiceState === STARTED)
+ assert(frontend.getServiceState === STARTED)
+ assert(frontend.connectionUrl.matches("localhost:[0-9]+"))
+
+ server.stop()
+ assert(server.getServiceState === STOPPED)
+ assert(frontend.getServiceState === STOPPED)
+ server.stop()
+ }
+
+ test("Flight SQL advertised host") {
+ val server = new KyuubiServer
+ val conf = KyuubiConf()
+ .set(FRONTEND_PROTOCOLS, Seq(FrontendProtocols.FLIGHT_SQL.toString))
+ .set(FRONTEND_FLIGHT_SQL_BIND_HOST.key, "localhost")
+ .set(FRONTEND_FLIGHT_SQL_BIND_PORT, 0)
+ .set(FRONTEND_ADVERTISED_HOST, "flight.example")
+
+ try {
+ server.initialize(conf)
+ assert(server.frontendServices.head.connectionUrl.startsWith("flight.example:"))
+ } finally {
+ server.stop()
+ }
+ }
+
+ test("Flight SQL TLS without certificate material fails startup") {
+ val server = new KyuubiServer
+ val conf = KyuubiConf()
+ .set(FRONTEND_PROTOCOLS, Seq(FrontendProtocols.FLIGHT_SQL.toString))
+ .set(FRONTEND_FLIGHT_SQL_BIND_HOST.key, "localhost")
+ .set(FRONTEND_FLIGHT_SQL_BIND_PORT, 0)
+ .set(FRONTEND_FLIGHT_SQL_SSL_ENABLED, true)
+
+ try {
+ server.initialize(conf)
+ intercept[Exception] {
+ server.start()
+ }
+ } finally {
+ server.stop()
+ }
+ }
+}
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtilsSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtilsSuite.scala
new file mode 100644
index 00000000000..5f937c3e36a
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightArrowUtilsSuite.scala
@@ -0,0 +1,141 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.nio.ByteBuffer
+import java.util
+
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.types.FloatingPointPrecision
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+
+class KyuubiFlightArrowUtilsSuite extends KyuubiFunSuite {
+
+ test("populateRootFromRowSet writes columnar thrift without Seq materialization") {
+ val emptyFields = util.Collections.emptyList[Field]()
+ val schema = new Schema(util.Arrays.asList(
+ new Field("flag", new FieldType(true, ArrowType.Bool.INSTANCE, null), emptyFields),
+ new Field(
+ "id",
+ new FieldType(true, new ArrowType.Int(32, true), null),
+ emptyFields),
+ new Field(
+ "score",
+ new FieldType(
+ true,
+ new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE),
+ null),
+ emptyFields),
+ new Field(
+ "name",
+ new FieldType(true, ArrowType.Utf8.INSTANCE, null),
+ emptyFields)))
+
+ val boolNulls = ByteBuffer.wrap(Array[Byte](0))
+ val i32Nulls = ByteBuffer.wrap(Array[Byte](0))
+ val doubleNulls = ByteBuffer.wrap(Array[Byte](0))
+ val stringNulls = ByteBuffer.wrap(Array[Byte](0))
+
+ val rowSet = new TRowSet()
+ rowSet.addToColumns(TColumn.boolVal(
+ new TBoolColumn(util.Arrays.asList(true, false), boolNulls)))
+ rowSet.addToColumns(TColumn.i32Val(new TI32Column(util.Arrays.asList(1, 2), i32Nulls)))
+ rowSet.addToColumns(TColumn.doubleVal(new TDoubleColumn(
+ util.Arrays.asList(1.5d, 2.5d),
+ doubleNulls)))
+ rowSet.addToColumns(TColumn.stringVal(new TStringColumn(
+ util.Arrays.asList("a", "b"),
+ stringNulls)))
+
+ val allocator = new RootAllocator()
+ val root = VectorSchemaRoot.create(schema, allocator)
+ try {
+ KyuubiFlightArrowUtils.populateRootFromRowSet(root, rowSet)
+ assert(root.getRowCount === 2)
+ assert(root.getVector(0).getObject(0) === true)
+ assert(root.getVector(1).getObject(1).toString === "2")
+ assert(root.getVector(2).getObject(0).toString === "1.5")
+ assert(root.getVector(3).getObject(1).toString === "b")
+ } finally {
+ root.close()
+ allocator.close()
+ }
+ }
+
+ test("isEmpty and isArrowRowSet helpers") {
+ assert(KyuubiFlightArrowUtils.isEmpty(null))
+ assert(KyuubiFlightArrowUtils.isEmpty(new TRowSet()))
+ val arrow = new TRowSet()
+ val binary = new TBinaryColumn()
+ binary.addToValues(ByteBuffer.wrap(Array[Byte](1, 2, 3)))
+ arrow.addToColumns(TColumn.binaryVal(binary))
+ assert(KyuubiFlightArrowUtils.isArrowRowSet(arrow))
+ assert(KyuubiFlightArrowUtils.arrowBatchBytes(arrow) === 3L)
+ assert(!KyuubiFlightArrowUtils.isEmpty(arrow))
+ }
+
+ test("isEmpty is true for exhausted columnar thrift pages with schema columns") {
+ // Matches TRowSetGenerator.toColumnBasedSet(Nil, schema): one empty TColumn per field.
+ val emptyNulls = ByteBuffer.wrap(Array[Byte](0))
+ val rowSet = new TRowSet()
+ rowSet.setRows(new util.ArrayList[TRow]())
+ rowSet.addToColumns(TColumn.i32Val(new TI32Column(new util.ArrayList[Integer](), emptyNulls)))
+ rowSet.addToColumns(TColumn.stringVal(
+ new TStringColumn(new util.ArrayList[String](), emptyNulls)))
+
+ assert(rowSet.getColumnsSize === 2)
+ assert(KyuubiFlightArrowUtils.rowCount(rowSet) === 0)
+ assert(KyuubiFlightArrowUtils.isEmpty(rowSet))
+ assert(!KyuubiFlightArrowUtils.isArrowRowSet(rowSet))
+ }
+
+ test("isEmpty is false when columnar thrift page has values") {
+ val nulls = ByteBuffer.wrap(Array[Byte](0))
+ val rowSet = new TRowSet()
+ rowSet.addToColumns(TColumn.i32Val(new TI32Column(util.Arrays.asList(7), nulls)))
+ assert(KyuubiFlightArrowUtils.rowCount(rowSet) === 1)
+ assert(!KyuubiFlightArrowUtils.isEmpty(rowSet))
+ }
+
+ test("populateRootFromRowSet yields zero rows for exhausted columnar page") {
+ val emptyFields = util.Collections.emptyList[Field]()
+ val schema = new Schema(util.Arrays.asList(
+ new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), emptyFields),
+ new Field("name", new FieldType(true, ArrowType.Utf8.INSTANCE, null), emptyFields)))
+
+ val emptyNulls = ByteBuffer.wrap(Array[Byte](0))
+ val rowSet = new TRowSet()
+ rowSet.addToColumns(TColumn.i32Val(new TI32Column(new util.ArrayList[Integer](), emptyNulls)))
+ rowSet.addToColumns(TColumn.stringVal(
+ new TStringColumn(new util.ArrayList[String](), emptyNulls)))
+
+ val allocator = new RootAllocator()
+ val root = VectorSchemaRoot.create(schema, allocator)
+ try {
+ KyuubiFlightArrowUtils.populateRootFromRowSet(root, rowSet)
+ assert(root.getRowCount === 0)
+ } finally {
+ root.close()
+ allocator.close()
+ }
+ }
+}
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandlerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandlerSuite.scala
new file mode 100644
index 00000000000..fd5eea93041
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightAuthHandlerSuite.scala
@@ -0,0 +1,86 @@
+/*
+ * 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.kyuubi.server.flight
+
+import java.util
+import java.util.Base64
+
+import org.apache.arrow.flight.CallHeaders
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.AUTHENTICATION_METHOD
+
+class KyuubiFlightAuthHandlerSuite extends KyuubiFunSuite {
+
+ test("anonymous auth when authentication is disabled") {
+ val conf = KyuubiConf().set(AUTHENTICATION_METHOD, Seq("NONE"))
+ val auth = new KyuubiFlightAuthHandler(conf)
+ val headers = new MapCallHeaders
+ assert(auth.authenticate(headers).getPeerIdentity === "anonymous")
+
+ headers.insert("x-user-name", "alice")
+ assert(auth.authenticate(headers).getPeerIdentity === "alice")
+ }
+
+ test("missing credentials fail when auth is required") {
+ val conf = KyuubiConf().set(AUTHENTICATION_METHOD, Seq("LDAP"))
+ val auth = new KyuubiFlightAuthHandler(conf)
+ intercept[RuntimeException] {
+ auth.authenticate(new MapCallHeaders)
+ }
+ }
+
+ test("malformed basic credentials are rejected") {
+ val conf = KyuubiConf().set(AUTHENTICATION_METHOD, Seq("LDAP"))
+ val auth = new KyuubiFlightAuthHandler(conf)
+ val headers = new MapCallHeaders
+ headers.insert(
+ "authorization",
+ "Basic " + Base64.getEncoder.encodeToString("nouser".getBytes("UTF-8")))
+ intercept[RuntimeException] {
+ auth.authenticate(headers)
+ }
+ }
+
+ private class MapCallHeaders extends CallHeaders {
+ private val values = new util.LinkedHashMap[String, String]()
+
+ override def get(key: String): String = values.get(key.toLowerCase)
+
+ override def getByte(key: String): Array[Byte] = null
+
+ override def getAll(key: String): java.lang.Iterable[String] = {
+ val value = values.get(key.toLowerCase)
+ if (value == null) util.Collections.emptyList()
+ else util.Collections.singletonList(value)
+ }
+
+ override def getAllByte(key: String): java.lang.Iterable[Array[Byte]] =
+ util.Collections.emptyList()
+
+ override def insert(key: String, value: String): Unit =
+ values.put(key.toLowerCase, value)
+
+ override def insert(key: String, value: Array[Byte]): Unit = ()
+
+ override def keys(): util.Set[String] = values.keySet()
+
+ override def containsKey(key: String): Boolean = values.containsKey(key.toLowerCase)
+ }
+}
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlQuerySuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlQuerySuite.scala
new file mode 100644
index 00000000000..f30dab56ad8
--- /dev/null
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/flight/KyuubiFlightSqlQuerySuite.scala
@@ -0,0 +1,60 @@
+/*
+ * 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.kyuubi.server.flight
+
+import org.apache.arrow.flight.{FlightClient, Location}
+import org.apache.arrow.flight.sql.FlightSqlClient
+import org.apache.arrow.memory.RootAllocator
+import org.scalatest.tags.Slow
+
+import org.apache.kyuubi.WithFlightSqlServer
+import org.apache.kyuubi.config.KyuubiConf
+
+@Slow
+class KyuubiFlightSqlQuerySuite extends WithFlightSqlServer {
+
+ override protected val conf: KyuubiConf = KyuubiConf()
+
+ test("execute a SQL statement and stream an Arrow batch") {
+ val endpoint = flightSqlUrl
+ val separator = endpoint.lastIndexOf(':')
+ val host = endpoint.substring(0, separator)
+ val port = endpoint.substring(separator + 1).toInt
+ val allocator = new RootAllocator()
+ val flightClient = FlightClient.builder(
+ allocator,
+ Location.forGrpcInsecure(host, port)).build()
+ val sqlClient = new FlightSqlClient(flightClient)
+ try {
+ val info = sqlClient.execute("SELECT 1 AS value")
+ assert(info.getEndpoints.size() === 1)
+ val stream = sqlClient.getStream(info.getEndpoints.get(0).getTicket)
+ try {
+ assert(stream.next())
+ assert(stream.getRoot.getRowCount > 0)
+ assert(stream.getRoot.getVector(0).getObject(0).toString === "1")
+ } finally {
+ stream.close()
+ }
+ } finally {
+ sqlClient.close()
+ flightClient.close()
+ allocator.close()
+ }
+ }
+}
diff --git a/pom.xml b/pom.xml
index efb34a1341c..1d190512bd3 100644
--- a/pom.xml
+++ b/pom.xml
@@ -353,6 +353,16 @@
arrow-memory-netty
${arrow.version}
+
+ org.apache.arrow
+ flight-core
+ ${arrow.version}
+
+
+ org.apache.arrow
+ flight-sql
+ ${arrow.version}
+
org.scala-lang
scala-library