You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2023/03/17 05:39:41 UTC
[kyuubi] branch master updated: [KYUUBI #4544] Initial implement Kyuubi Chat Engine
This is an automated email from the ASF dual-hosted git repository.
chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new f1eb449c4 [KYUUBI #4544] Initial implement Kyuubi Chat Engine
f1eb449c4 is described below
commit f1eb449c42b820b8ad3f4d388bb6320acf70bb1b
Author: Cheng Pan <ch...@apache.org>
AuthorDate: Fri Mar 17 13:39:30 2023 +0800
[KYUUBI #4544] Initial implement Kyuubi Chat Engine
### _Why are the changes needed?_
Introduce a brand new CHAT engine, it's supposed to support different backends, e.g. ChatGPT, 文心一言, etc.
This PR implements the following providers:
- ECHO, simply replies a welcome message.
- GPT: a.k.a ChatGPT, powered by OpenAI, which requires a API key for authentication. https://platform.openai.com/account/api-keys
Add the following configurations in `kyuubi-defaults.conf`
```
kyuubi.engine.chat.provider=[ECHO|GPT]
kyuubi.engine.chat.gpt.apiKey=<chat-gpt-api-key>
```
Open an ECHO beeline chat engine.
```
beeline -u 'jdbc:hive2://localhost:10009/?kyuubi.engine.type=CHAT;kyuubi.engine.chat.provider=ECHO'
```
```
Connecting to jdbc:hive2://localhost:10009/
Connected to: Kyuubi Chat Engine (version 1.8.0-SNAPSHOT)
Driver: Kyuubi Project Hive JDBC Client (version 1.7.0)
Beeline version 1.7.0 by Apache Kyuubi
0: jdbc:hive2://localhost:10009/> Hello, Kyuubi!;
+----------------------------------------+
| reply |
+----------------------------------------+
| This is ChatKyuubi, nice to meet you! |
+----------------------------------------+
1 row selected (0.397 seconds)
```
Open a ChatGPT beeline chat engine. (make sure your network can connect the open API and configure the API key)
```
beeline -u 'jdbc:hive2://localhost:10009/?kyuubi.engine.type=CHAT;kyuubi.engine.chat.provider=GPT'
```
<img width="1109" alt="image" src="https://user-images.githubusercontent.com/26535726/225813625-a002e6e2-3b0d-4194-b061-2e215d58ba94.png">
### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
- [x] Add screenshots for manual tests if appropriate
- [x] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request
Closes #4544 from pan3793/chatgpt.
Closes #4544
87bdebb6d [Cheng Pan] nit
f7dee18f3 [Cheng Pan] Update docs
9beb55162 [cxzl25] chat api (#1)
af38bdc7c [Cheng Pan] update docs
9aa6d83a6 [Cheng Pan] Initial implement Kyuubi Chat Engine
Lead-authored-by: Cheng Pan <ch...@apache.org>
Co-authored-by: cxzl25 <cx...@users.noreply.github.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
build/dist | 13 +++
docs/deployment/settings.md | 113 +++++++++---------
externals/kyuubi-chat-engine/pom.xml | 89 ++++++++++++++
.../kyuubi/engine/chat/ChatBackendService.scala | 15 +--
.../org/apache/kyuubi/engine/chat/ChatEngine.scala | 86 ++++++++++++++
.../engine/chat/ChatTBinaryFrontendService.scala | 21 ++--
.../engine/chat/operation/ChatOperation.scala | 100 ++++++++++++++++
.../chat/operation/ChatOperationManager.scala | 130 +++++++++++++++++++++
.../engine/chat/operation/ExecuteStatement.scala | 66 +++++++++++
.../engine/chat/provider/ChatGPTProvider.scala | 84 +++++++++++++
.../kyuubi/engine/chat/provider/ChatProvider.scala | 59 ++++++++++
.../kyuubi/engine/chat/provider/EchoProvider.scala | 15 +--
.../kyuubi/engine/chat/provider/Message.scala | 11 +-
.../apache/kyuubi/engine/chat/schema/RowSet.scala | 108 +++++++++++++++++
.../kyuubi/engine/chat/schema/SchemaHelper.scala | 56 +++++++++
.../engine/chat/session/ChatSessionImpl.scala | 68 +++++++++++
.../engine/chat/session/ChatSessionManager.scala | 71 +++++++++++
.../src/test/resources/log4j2-test.xml | 42 +++++++
.../apache/kyuubi/engine/chat/WithChatEngine.scala | 60 ++++++++++
.../engine/chat/operation/ChatOperationSuite.scala | 27 +++--
.../org/apache/kyuubi/config/KyuubiConf.scala | 46 ++++++++
.../org/apache/kyuubi/engine/EngineType.scala | 2 +-
.../scala/org/apache/kyuubi/engine/EngineRef.scala | 5 +-
.../kyuubi/engine/chat/ChatProcessBuilder.scala | 115 ++++++++++++++++++
pom.xml | 3 +-
25 files changed, 1311 insertions(+), 94 deletions(-)
diff --git a/build/dist b/build/dist
index c155dce3f..e0dae3479 100755
--- a/build/dist
+++ b/build/dist
@@ -256,6 +256,7 @@ mkdir -p "$DISTDIR/externals/engines/spark"
mkdir -p "$DISTDIR/externals/engines/trino"
mkdir -p "$DISTDIR/externals/engines/hive"
mkdir -p "$DISTDIR/externals/engines/jdbc"
+mkdir -p "$DISTDIR/externals/engines/chat"
echo "Kyuubi $VERSION $GITREVSTRING built for" > "$DISTDIR/RELEASE"
echo "Java $JAVA_VERSION" >> "$DISTDIR/RELEASE"
echo "Scala $SCALA_VERSION" >> "$DISTDIR/RELEASE"
@@ -313,6 +314,18 @@ for jar in $(ls "$DISTDIR/jars/"); do
fi
done
+# Copy chat engines
+cp "$KYUUBI_HOME/externals/kyuubi-chat-engine/target/kyuubi-chat-engine_${SCALA_VERSION}-${VERSION}.jar" "$DISTDIR/externals/engines/chat/"
+cp -r "$KYUUBI_HOME"/externals/kyuubi-chat-engine/target/scala-$SCALA_VERSION/jars/*.jar "$DISTDIR/externals/engines/chat/"
+
+# Share the jars w/ server to reduce binary size
+# shellcheck disable=SC2045
+for jar in $(ls "$DISTDIR/jars/"); do
+ if [[ -f "$DISTDIR/externals/engines/chat/$jar" ]]; then
+ (cd $DISTDIR/externals/engines/chat; ln -snf "../../../jars/$jar" "$DISTDIR/externals/engines/chat/$jar")
+ fi
+done
+
# Copy kyuubi tools
if [[ -f "$KYUUBI_HOME/tools/spark-block-cleaner/target/spark-block-cleaner_${SCALA_VERSION}-${VERSION}.jar" ]]; then
mkdir -p "$DISTDIR/tools/spark-block-cleaner/kubernetes"
diff --git a/docs/deployment/settings.md b/docs/deployment/settings.md
index a2474a506..9c5f76712 100644
--- a/docs/deployment/settings.md
+++ b/docs/deployment/settings.md
@@ -118,60 +118,65 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
### Engine
-| Key | Default | [...]
-|----------------------------------------------------------|---------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- [...]
-| kyuubi.engine.connection.url.use.hostname | true | (deprecated) When true, the engine registers with hostname to zookeeper. When Spark runs on K8s with cluster mode, set to false to ensure that server can connect to engine [...]
-| kyuubi.engine.deregister.exception.classes || A comma-separated list of exception classes. If there is any exception thrown, whose class matches the specified classes, the engine would deregister itself. [...]
-| kyuubi.engine.deregister.exception.messages || A comma-separated list of exception messages. If there is any exception thrown, whose message or stacktrace matches the specified message list, the engine would deregister itself. [...]
-| kyuubi.engine.deregister.exception.ttl | PT30M | Time to live(TTL) for exceptions pattern specified in kyuubi.engine.deregister.exception.classes and kyuubi.engine.deregister.exception.messages to deregister engines. Once the total error count hits the kyuubi.engine.deregister.job.max.failures within the TTL, an engine will deregister itself and wait for self-terminated. Otherwise, we suppose that the engine has recovered from temporary failures. [...]
-| kyuubi.engine.deregister.job.max.failures | 4 | Number of failures of job before deregistering the engine. [...]
-| kyuubi.engine.event.json.log.path | file:///tmp/kyuubi/events | The location where all the engine events go for the built-in JSON logger.<ul><li>Local Path: start with 'file://'</li><li>HDFS Path: start with 'hdfs://'</li></ul> [...]
-| kyuubi.engine.event.loggers | SPARK | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>SPARK: the events will be written to the Spark listener bus.</li> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: User-defined event handlers.</li></ul> Note that: Kyuubi supports custom event handlers with the Jav [...]
-| kyuubi.engine.flink.extra.classpath | <undefined> | The extra classpath for the Flink SQL engine, for configuring the location of hadoop client jars, etc [...]
-| kyuubi.engine.flink.java.options | <undefined> | The extra Java options for the Flink SQL engine [...]
-| kyuubi.engine.flink.memory | 1g | The heap memory for the Flink SQL engine [...]
-| kyuubi.engine.hive.event.loggers | JSON | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
-| kyuubi.engine.hive.extra.classpath | <undefined> | The extra classpath for the Hive query engine, for configuring location of the hadoop client jars and etc. [...]
-| kyuubi.engine.hive.java.options | <undefined> | The extra Java options for the Hive query engine [...]
-| kyuubi.engine.hive.memory | 1g | The heap memory for the Hive query engine [...]
-| kyuubi.engine.initialize.sql | SHOW DATABASES | SemiColon-separated list of SQL statements to be initialized in the newly created engine before queries. i.e. use `SHOW DATABASES` to eagerly active HiveClient. This configuration can not be used in JDBC url due to the limitation of Beeline/JDBC driver. [...]
-| kyuubi.engine.jdbc.connection.password | <undefined> | The password is used for connecting to server [...]
-| kyuubi.engine.jdbc.connection.properties || The additional properties are used for connecting to server [...]
-| kyuubi.engine.jdbc.connection.provider | <undefined> | The connection provider is used for getting a connection from the server [...]
-| kyuubi.engine.jdbc.connection.url | <undefined> | The server url that engine will connect to [...]
-| kyuubi.engine.jdbc.connection.user | <undefined> | The user is used for connecting to server [...]
-| kyuubi.engine.jdbc.driver.class | <undefined> | The driver class for JDBC engine connection [...]
-| kyuubi.engine.jdbc.extra.classpath | <undefined> | The extra classpath for the JDBC query engine, for configuring the location of the JDBC driver and etc. [...]
-| kyuubi.engine.jdbc.java.options | <undefined> | The extra Java options for the JDBC query engine [...]
-| kyuubi.engine.jdbc.memory | 1g | The heap memory for the JDBC query engine [...]
-| kyuubi.engine.jdbc.type | <undefined> | The short name of JDBC type [...]
-| kyuubi.engine.operation.convert.catalog.database.enabled | true | When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines [...]
-| kyuubi.engine.operation.log.dir.root | engine_operation_logs | Root directory for query operation log at engine-side. [...]
-| kyuubi.engine.pool.name | engine-pool | The name of the engine pool. [...]
-| kyuubi.engine.pool.selectPolicy | RANDOM | The select policy of an engine from the corresponding engine pool engine for a session. <ul><li>RANDOM - Randomly use the engine in the pool</li><li>POLLING - Polling use the engine in the pool</li></ul> [...]
-| kyuubi.engine.pool.size | -1 | The size of the engine pool. Note that, if the size is less than 1, the engine pool will not be enabled; otherwise, the size of the engine pool will be min(this, kyuubi.engine.pool.size.threshold). [...]
-| kyuubi.engine.pool.size.threshold | 9 | This parameter is introduced as a server-side parameter controlling the upper limit of the engine pool. [...]
-| kyuubi.engine.session.initialize.sql || SemiColon-separated list of SQL statements to be initialized in the newly created engine session before queries. This configuration can not be used in JDBC url due to the limitation of Beeline/JDBC driver. [...]
-| kyuubi.engine.share.level | USER | Engines will be shared in different levels, available configs are: <ul> <li>CONNECTION: engine will not be shared but only used by the current client connection</li> <li>USER: engine will be shared by all sessions created by a unique username, see also kyuubi.engine.share.level.subdomain</li> <li>GROUP: the engine will be shared by all sessions created by all users belong to the same primary group na [...]
-| kyuubi.engine.share.level.sub.domain | <undefined> | (deprecated) - Using kyuubi.engine.share.level.subdomain instead [...]
-| kyuubi.engine.share.level.subdomain | <undefined> | Allow end-users to create a subdomain for the share level of an engine. A subdomain is a case-insensitive string values that must be a valid zookeeper subpath. For example, for the `USER` share level, an end-user can share a certain engine within a subdomain, not for all of its clients. End-users are free to create multiple engines in the `USER` share level. When disable engine pool, use 'default' if [...]
-| kyuubi.engine.single.spark.session | false | When set to true, this engine is running in a single session mode. All the JDBC/ODBC connections share the temporary views, function registries, SQL configuration and the current database. [...]
-| kyuubi.engine.spark.event.loggers | SPARK | A comma-separated list of engine loggers, where engine/session/operation etc events go.<ul> <li>SPARK: the events will be written to the Spark listener bus.</li> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
-| kyuubi.engine.spark.python.env.archive | <undefined> | Portable Python env archive used for Spark engine Python language mode. [...]
-| kyuubi.engine.spark.python.env.archive.exec.path | bin/python | The Python exec path under the Python env archive. [...]
-| kyuubi.engine.spark.python.home.archive | <undefined> | Spark archive containing $SPARK_HOME/python directory, which is used to init session Python worker for Python language mode. [...]
-| kyuubi.engine.submit.timeout | PT30S | Period to tolerant Driver Pod ephemerally invisible after submitting. In some Resource Managers, e.g. K8s, the Driver Pod is not invisible immediately after `spark-submit` is returned. [...]
-| kyuubi.engine.trino.event.loggers | JSON | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
-| kyuubi.engine.trino.extra.classpath | <undefined> | The extra classpath for the Trino query engine, for configuring other libs which may need by the Trino engine [...]
-| kyuubi.engine.trino.java.options | <undefined> | The extra Java options for the Trino query engine [...]
-| kyuubi.engine.trino.memory | 1g | The heap memory for the Trino query engine [...]
-| kyuubi.engine.type | SPARK_SQL | Specify the detailed engine supported by Kyuubi. The engine type bindings to SESSION scope. This configuration is experimental. Currently, available configs are: <ul> <li>SPARK_SQL: specify this engine type will launch a Spark engine which can provide all the capacity of the Apache Spark. Note, it's a default engine type.</li> <li>FLINK_SQL: specify this engine type will launch a Flink engine which c [...]
-| kyuubi.engine.ui.retainedSessions | 200 | The number of SQL client sessions kept in the Kyuubi Query Engine web UI. [...]
-| kyuubi.engine.ui.retainedStatements | 200 | The number of statements kept in the Kyuubi Query Engine web UI. [...]
-| kyuubi.engine.ui.stop.enabled | true | When true, allows Kyuubi engine to be killed from the Spark Web UI. [...]
-| kyuubi.engine.user.isolated.spark.session | true | When set to false, if the engine is running in a group or server share level, all the JDBC/ODBC connections will be isolated against the user. Including the temporary views, function registries, SQL configuration, and the current database. Note that, it does not affect if the share level is connection or user. [...]
-| kyuubi.engine.user.isolated.spark.session.idle.interval | PT1M | The interval to check if the user-isolated Spark session is timeout. [...]
-| kyuubi.engine.user.isolated.spark.session.idle.timeout | PT6H | If kyuubi.engine.user.isolated.spark.session is false, we will release the Spark session if its corresponding user is inactive after this configured timeout. [...]
+| Key | Default | [...]
+|----------------------------------------------------------|---------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- [...]
+| kyuubi.engine.chat.extra.classpath | <undefined> | The extra classpath for the Chat engine, for configuring the location of the SDK and etc. [...]
+| kyuubi.engine.chat.gpt.apiKey | <undefined> | The key to access OpenAI open API, which could be got at https://platform.openai.com/account/api-keys [...]
+| kyuubi.engine.chat.java.options | <undefined> | The extra Java options for the Chat engine [...]
+| kyuubi.engine.chat.memory | 1g | The heap memory for the Chat engine [...]
+| kyuubi.engine.chat.provider | ECHO | The provider for the Chat engine. Candidates: <ul> <li>ECHO: simply replies a welcome message.</li> <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li></ul> [...]
+| kyuubi.engine.connection.url.use.hostname | true | (deprecated) When true, the engine registers with hostname to zookeeper. When Spark runs on K8s with cluster mode, set to false to ensure that server can connect to engine [...]
+| kyuubi.engine.deregister.exception.classes || A comma-separated list of exception classes. If there is any exception thrown, whose class matches the specified classes, the engine would deregister itself. [...]
+| kyuubi.engine.deregister.exception.messages || A comma-separated list of exception messages. If there is any exception thrown, whose message or stacktrace matches the specified message list, the engine would deregister itself. [...]
+| kyuubi.engine.deregister.exception.ttl | PT30M | Time to live(TTL) for exceptions pattern specified in kyuubi.engine.deregister.exception.classes and kyuubi.engine.deregister.exception.messages to deregister engines. Once the total error count hits the kyuubi.engine.deregister.job.max.failures within the TTL, an engine will deregister itself and wait for self-terminated. Otherwise, we suppose that the engine has recovered from temporary failures. [...]
+| kyuubi.engine.deregister.job.max.failures | 4 | Number of failures of job before deregistering the engine. [...]
+| kyuubi.engine.event.json.log.path | file:///tmp/kyuubi/events | The location where all the engine events go for the built-in JSON logger.<ul><li>Local Path: start with 'file://'</li><li>HDFS Path: start with 'hdfs://'</li></ul> [...]
+| kyuubi.engine.event.loggers | SPARK | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>SPARK: the events will be written to the Spark listener bus.</li> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: User-defined event handlers.</li></ul> Note that: Kyuubi supports custom event handlers with the Jav [...]
+| kyuubi.engine.flink.extra.classpath | <undefined> | The extra classpath for the Flink SQL engine, for configuring the location of hadoop client jars, etc [...]
+| kyuubi.engine.flink.java.options | <undefined> | The extra Java options for the Flink SQL engine [...]
+| kyuubi.engine.flink.memory | 1g | The heap memory for the Flink SQL engine [...]
+| kyuubi.engine.hive.event.loggers | JSON | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
+| kyuubi.engine.hive.extra.classpath | <undefined> | The extra classpath for the Hive query engine, for configuring location of the hadoop client jars and etc. [...]
+| kyuubi.engine.hive.java.options | <undefined> | The extra Java options for the Hive query engine [...]
+| kyuubi.engine.hive.memory | 1g | The heap memory for the Hive query engine [...]
+| kyuubi.engine.initialize.sql | SHOW DATABASES | SemiColon-separated list of SQL statements to be initialized in the newly created engine before queries. i.e. use `SHOW DATABASES` to eagerly active HiveClient. This configuration can not be used in JDBC url due to the limitation of Beeline/JDBC driver. [...]
+| kyuubi.engine.jdbc.connection.password | <undefined> | The password is used for connecting to server [...]
+| kyuubi.engine.jdbc.connection.properties || The additional properties are used for connecting to server [...]
+| kyuubi.engine.jdbc.connection.provider | <undefined> | The connection provider is used for getting a connection from the server [...]
+| kyuubi.engine.jdbc.connection.url | <undefined> | The server url that engine will connect to [...]
+| kyuubi.engine.jdbc.connection.user | <undefined> | The user is used for connecting to server [...]
+| kyuubi.engine.jdbc.driver.class | <undefined> | The driver class for JDBC engine connection [...]
+| kyuubi.engine.jdbc.extra.classpath | <undefined> | The extra classpath for the JDBC query engine, for configuring the location of the JDBC driver and etc. [...]
+| kyuubi.engine.jdbc.java.options | <undefined> | The extra Java options for the JDBC query engine [...]
+| kyuubi.engine.jdbc.memory | 1g | The heap memory for the JDBC query engine [...]
+| kyuubi.engine.jdbc.type | <undefined> | The short name of JDBC type [...]
+| kyuubi.engine.operation.convert.catalog.database.enabled | true | When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines [...]
+| kyuubi.engine.operation.log.dir.root | engine_operation_logs | Root directory for query operation log at engine-side. [...]
+| kyuubi.engine.pool.name | engine-pool | The name of the engine pool. [...]
+| kyuubi.engine.pool.selectPolicy | RANDOM | The select policy of an engine from the corresponding engine pool engine for a session. <ul><li>RANDOM - Randomly use the engine in the pool</li><li>POLLING - Polling use the engine in the pool</li></ul> [...]
+| kyuubi.engine.pool.size | -1 | The size of the engine pool. Note that, if the size is less than 1, the engine pool will not be enabled; otherwise, the size of the engine pool will be min(this, kyuubi.engine.pool.size.threshold). [...]
+| kyuubi.engine.pool.size.threshold | 9 | This parameter is introduced as a server-side parameter controlling the upper limit of the engine pool. [...]
+| kyuubi.engine.session.initialize.sql || SemiColon-separated list of SQL statements to be initialized in the newly created engine session before queries. This configuration can not be used in JDBC url due to the limitation of Beeline/JDBC driver. [...]
+| kyuubi.engine.share.level | USER | Engines will be shared in different levels, available configs are: <ul> <li>CONNECTION: engine will not be shared but only used by the current client connection</li> <li>USER: engine will be shared by all sessions created by a unique username, see also kyuubi.engine.share.level.subdomain</li> <li>GROUP: the engine will be shared by all sessions created by all users belong to the same primary group na [...]
+| kyuubi.engine.share.level.sub.domain | <undefined> | (deprecated) - Using kyuubi.engine.share.level.subdomain instead [...]
+| kyuubi.engine.share.level.subdomain | <undefined> | Allow end-users to create a subdomain for the share level of an engine. A subdomain is a case-insensitive string values that must be a valid zookeeper subpath. For example, for the `USER` share level, an end-user can share a certain engine within a subdomain, not for all of its clients. End-users are free to create multiple engines in the `USER` share level. When disable engine pool, use 'default' if [...]
+| kyuubi.engine.single.spark.session | false | When set to true, this engine is running in a single session mode. All the JDBC/ODBC connections share the temporary views, function registries, SQL configuration and the current database. [...]
+| kyuubi.engine.spark.event.loggers | SPARK | A comma-separated list of engine loggers, where engine/session/operation etc events go.<ul> <li>SPARK: the events will be written to the Spark listener bus.</li> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
+| kyuubi.engine.spark.python.env.archive | <undefined> | Portable Python env archive used for Spark engine Python language mode. [...]
+| kyuubi.engine.spark.python.env.archive.exec.path | bin/python | The Python exec path under the Python env archive. [...]
+| kyuubi.engine.spark.python.home.archive | <undefined> | Spark archive containing $SPARK_HOME/python directory, which is used to init session Python worker for Python language mode. [...]
+| kyuubi.engine.submit.timeout | PT30S | Period to tolerant Driver Pod ephemerally invisible after submitting. In some Resource Managers, e.g. K8s, the Driver Pod is not invisible immediately after `spark-submit` is returned. [...]
+| kyuubi.engine.trino.event.loggers | JSON | A comma-separated list of engine history loggers, where engine/session/operation etc events go.<ul> <li>JSON: the events will be written to the location of kyuubi.engine.event.json.log.path</li> <li>JDBC: to be done</li> <li>CUSTOM: to be done.</li></ul> [...]
+| kyuubi.engine.trino.extra.classpath | <undefined> | The extra classpath for the Trino query engine, for configuring other libs which may need by the Trino engine [...]
+| kyuubi.engine.trino.java.options | <undefined> | The extra Java options for the Trino query engine [...]
+| kyuubi.engine.trino.memory | 1g | The heap memory for the Trino query engine [...]
+| kyuubi.engine.type | SPARK_SQL | Specify the detailed engine supported by Kyuubi. The engine type bindings to SESSION scope. This configuration is experimental. Currently, available configs are: <ul> <li>SPARK_SQL: specify this engine type will launch a Spark engine which can provide all the capacity of the Apache Spark. Note, it's a default engine type.</li> <li>FLINK_SQL: specify this engine type will launch a Flink engine which c [...]
+| kyuubi.engine.ui.retainedSessions | 200 | The number of SQL client sessions kept in the Kyuubi Query Engine web UI. [...]
+| kyuubi.engine.ui.retainedStatements | 200 | The number of statements kept in the Kyuubi Query Engine web UI. [...]
+| kyuubi.engine.ui.stop.enabled | true | When true, allows Kyuubi engine to be killed from the Spark Web UI. [...]
+| kyuubi.engine.user.isolated.spark.session | true | When set to false, if the engine is running in a group or server share level, all the JDBC/ODBC connections will be isolated against the user. Including the temporary views, function registries, SQL configuration, and the current database. Note that, it does not affect if the share level is connection or user. [...]
+| kyuubi.engine.user.isolated.spark.session.idle.interval | PT1M | The interval to check if the user-isolated Spark session is timeout. [...]
+| kyuubi.engine.user.isolated.spark.session.idle.timeout | PT6H | If kyuubi.engine.user.isolated.spark.session is false, we will release the Spark session if its corresponding user is inactive after this configured timeout. [...]
### Event
diff --git a/externals/kyuubi-chat-engine/pom.xml b/externals/kyuubi-chat-engine/pom.xml
new file mode 100644
index 000000000..7e2178918
--- /dev/null
+++ b/externals/kyuubi-chat-engine/pom.xml
@@ -0,0 +1,89 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+ -->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+ <parent>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-parent</artifactId>
+ <version>1.8.0-SNAPSHOT</version>
+ <relativePath>../../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>kyuubi-chat-engine_2.12</artifactId>
+ <packaging>jar</packaging>
+ <name>Kyuubi Project Engine Chat</name>
+ <url>https://kyuubi.apache.org/</url>
+
+ <dependencies>
+ <!-- kyuubi dependency -->
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-ha_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.httpcomponents</groupId>
+ <artifactId>httpclient</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>kyuubi-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.kyuubi</groupId>
+ <artifactId>${hive.jdbc.artifact}</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
+
+ </dependencies>
+
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-jar-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>prepare-test-jar</id>
+ <goals>
+ <goal>test-jar</goal>
+ </goals>
+ <phase>test-compile</phase>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
+ <testOutputDirectory>target/scala-${scala.binary.version}/test-classes</testOutputDirectory>
+ </build>
+
+</project>
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatBackendService.scala
similarity index 69%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
copy to externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatBackendService.scala
index 88680a8c7..fdc710e2c 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatBackendService.scala
@@ -14,14 +14,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
+package org.apache.kyuubi.engine.chat
-package org.apache.kyuubi.engine
+import org.apache.kyuubi.engine.chat.session.ChatSessionManager
+import org.apache.kyuubi.service.AbstractBackendService
+import org.apache.kyuubi.session.SessionManager
-/**
- * Defines different engine types supported by Kyuubi.
- */
-object EngineType extends Enumeration {
- type EngineType = Value
+class ChatBackendService
+ extends AbstractBackendService("ChatBackendService") {
+
+ override val sessionManager: SessionManager = new ChatSessionManager()
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatEngine.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatEngine.scala
new file mode 100644
index 000000000..c1fdea953
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatEngine.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.engine.chat
+
+import ChatEngine.currentEngine
+
+import org.apache.kyuubi.{Logging, Utils}
+import org.apache.kyuubi.Utils.{addShutdownHook, JDBC_ENGINE_SHUTDOWN_PRIORITY}
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.ha.HighAvailabilityConf.HA_ZK_CONN_RETRY_POLICY
+import org.apache.kyuubi.ha.client.RetryPolicies
+import org.apache.kyuubi.service.Serverable
+import org.apache.kyuubi.util.SignalRegister
+
+class ChatEngine extends Serverable("ChatEngine") {
+
+ override val backendService = new ChatBackendService()
+ override val frontendServices = Seq(new ChatTBinaryFrontendService(this))
+
+ override def start(): Unit = {
+ super.start()
+ // Start engine self-terminating checker after all services are ready and it can be reached by
+ // all servers in engine spaces.
+ backendService.sessionManager.startTerminatingChecker(() => {
+ currentEngine.foreach(_.stop())
+ })
+ }
+
+ override protected def stopServer(): Unit = {}
+}
+
+object ChatEngine extends Logging {
+
+ val kyuubiConf: KyuubiConf = KyuubiConf()
+
+ var currentEngine: Option[ChatEngine] = None
+
+ def startEngine(): Unit = {
+ currentEngine = Some(new ChatEngine())
+ currentEngine.foreach { engine =>
+ engine.initialize(kyuubiConf)
+ engine.start()
+ addShutdownHook(
+ () => {
+ engine.stop()
+ },
+ JDBC_ENGINE_SHUTDOWN_PRIORITY + 1)
+ }
+ }
+
+ def main(args: Array[String]): Unit = {
+ SignalRegister.registerLogger(logger)
+
+ try {
+ Utils.fromCommandLineArgs(args, kyuubiConf)
+ kyuubiConf.setIfMissing(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)
+ kyuubiConf.setIfMissing(HA_ZK_CONN_RETRY_POLICY, RetryPolicies.N_TIME.toString)
+
+ startEngine()
+ } catch {
+ case t: Throwable if currentEngine.isDefined =>
+ currentEngine.foreach { engine =>
+ engine.stop()
+ }
+ error("Failed to create Chat Engine", t)
+ throw t
+ case t: Throwable =>
+ error("Failed to create Chat Engine.", t)
+ throw t
+ }
+ }
+}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatTBinaryFrontendService.scala
similarity index 57%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
copy to externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatTBinaryFrontendService.scala
index 88680a8c7..80702c97c 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ChatTBinaryFrontendService.scala
@@ -14,14 +14,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
+package org.apache.kyuubi.engine.chat
-package org.apache.kyuubi.engine
+import org.apache.kyuubi.ha.client.{EngineServiceDiscovery, ServiceDiscovery}
+import org.apache.kyuubi.service.{Serverable, Service, TBinaryFrontendService}
-/**
- * Defines different engine types supported by Kyuubi.
- */
-object EngineType extends Enumeration {
- type EngineType = Value
+class ChatTBinaryFrontendService(override val serverable: Serverable)
+ extends TBinaryFrontendService("ChatTBinaryFrontend") {
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
+ /**
+ * An optional `ServiceDiscovery` for [[FrontendService]] to expose itself
+ */
+ override lazy val discoveryService: Option[Service] =
+ if (ServiceDiscovery.supportServiceDiscovery(conf)) {
+ Some(new EngineServiceDiscovery(this))
+ } else {
+ None
+ }
}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala
new file mode 100644
index 000000000..38527cbf1
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala
@@ -0,0 +1,100 @@
+/*
+ * 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.engine.chat.operation
+
+import org.apache.hive.service.rpc.thrift._
+
+import org.apache.kyuubi.{KyuubiSQLException, Utils}
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.engine.chat.schema.{RowSet, SchemaHelper}
+import org.apache.kyuubi.operation.{AbstractOperation, FetchIterator, OperationState}
+import org.apache.kyuubi.operation.FetchOrientation.{FETCH_FIRST, FETCH_NEXT, FETCH_PRIOR, FetchOrientation}
+import org.apache.kyuubi.session.Session
+
+abstract class ChatOperation(session: Session) extends AbstractOperation(session) {
+
+ protected var iter: FetchIterator[Array[String]] = _
+
+ protected lazy val conf: KyuubiConf = session.sessionManager.getConf
+
+ override def getNextRowSet(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ validateDefaultFetchOrientation(order)
+ assertState(OperationState.FINISHED)
+ setHasResultSet(true)
+ order match {
+ case FETCH_NEXT =>
+ iter.fetchNext()
+ case FETCH_PRIOR =>
+ iter.fetchPrior(rowSetSize)
+ case FETCH_FIRST =>
+ iter.fetchAbsolute(0)
+ }
+
+ val taken = iter.take(rowSetSize)
+ val resultRowSet = RowSet.toTRowSet(taken.toSeq, 1, getProtocolVersion)
+ resultRowSet.setStartRowOffset(iter.getPosition)
+ resultRowSet
+ }
+
+ override def cancel(): Unit = {
+ cleanup(OperationState.CANCELED)
+ }
+
+ override def close(): Unit = {
+ cleanup(OperationState.CLOSED)
+ }
+
+ protected def onError(cancel: Boolean = false): PartialFunction[Throwable, Unit] = {
+ // We should use Throwable instead of Exception since `java.lang.NoClassDefFoundError`
+ // could be thrown.
+ case e: Throwable =>
+ state.synchronized {
+ val errMsg = Utils.stringifyException(e)
+ if (state == OperationState.TIMEOUT) {
+ val ke = KyuubiSQLException(s"Timeout operating $opType: $errMsg")
+ setOperationException(ke)
+ throw ke
+ } else if (isTerminalState(state)) {
+ setOperationException(KyuubiSQLException(errMsg))
+ warn(s"Ignore exception in terminal state with $statementId: $errMsg")
+ } else {
+ error(s"Error operating $opType: $errMsg", e)
+ val ke = KyuubiSQLException(s"Error operating $opType: $errMsg", e)
+ setOperationException(ke)
+ setState(OperationState.ERROR)
+ throw ke
+ }
+ }
+ }
+
+ override protected def beforeRun(): Unit = {
+ setState(OperationState.PENDING)
+ setHasResultSet(true)
+ }
+
+ override protected def afterRun(): Unit = {}
+
+ override def getResultSetMetadata: TGetResultSetMetadataResp = {
+ val tTableSchema = SchemaHelper.stringTTableSchema("reply")
+ val resp = new TGetResultSetMetadataResp
+ resp.setSchema(tTableSchema)
+ resp.setStatus(OK_STATUS)
+ resp
+ }
+
+ override def shouldRunAsync: Boolean = false
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationManager.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationManager.scala
new file mode 100644
index 000000000..1e8916517
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationManager.scala
@@ -0,0 +1,130 @@
+/*
+ * 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.engine.chat.operation
+
+import java.util
+
+import org.apache.kyuubi.KyuubiSQLException
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.engine.chat.provider.ChatProvider
+import org.apache.kyuubi.operation.{Operation, OperationManager}
+import org.apache.kyuubi.session.Session
+
+class ChatOperationManager(
+ conf: KyuubiConf,
+ chatProvider: ChatProvider) extends OperationManager("ChatOperationManager") {
+
+ override def newExecuteStatementOperation(
+ session: Session,
+ statement: String,
+ confOverlay: Map[String, String],
+ runAsync: Boolean,
+ queryTimeout: Long): Operation = {
+ val executeStatement =
+ new ExecuteStatement(
+ session,
+ statement,
+ runAsync,
+ queryTimeout,
+ chatProvider)
+ addOperation(executeStatement)
+ }
+
+ override def newGetTypeInfoOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCatalogsOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetSchemasOperation(
+ session: Session,
+ catalog: String,
+ schema: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetTablesOperation(
+ session: Session,
+ catalogName: String,
+ schemaName: String,
+ tableName: String,
+ tableTypes: util.List[String]): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetTableTypesOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetColumnsOperation(
+ session: Session,
+ catalogName: String,
+ schemaName: String,
+ tableName: String,
+ columnName: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetFunctionsOperation(
+ session: Session,
+ catalogName: String,
+ schemaName: String,
+ functionName: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetPrimaryKeysOperation(
+ session: Session,
+ catalogName: String,
+ schemaName: String,
+ tableName: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCrossReferenceOperation(
+ session: Session,
+ primaryCatalog: String,
+ primarySchema: String,
+ primaryTable: String,
+ foreignCatalog: String,
+ foreignSchema: String,
+ foreignTable: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def getQueryId(operation: Operation): String = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ExecuteStatement.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ExecuteStatement.scala
new file mode 100644
index 000000000..754a51932
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ExecuteStatement.scala
@@ -0,0 +1,66 @@
+/*
+ * 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.engine.chat.operation
+
+import org.apache.kyuubi.Logging
+import org.apache.kyuubi.engine.chat.provider.ChatProvider
+import org.apache.kyuubi.operation.{ArrayFetchIterator, OperationState}
+import org.apache.kyuubi.operation.log.OperationLog
+import org.apache.kyuubi.session.Session
+
+class ExecuteStatement(
+ session: Session,
+ override val statement: String,
+ override val shouldRunAsync: Boolean,
+ queryTimeout: Long,
+ chatProvider: ChatProvider)
+ extends ChatOperation(session) with Logging {
+
+ private val operationLog: OperationLog = OperationLog.createOperationLog(session, getHandle)
+ override def getOperationLog: Option[OperationLog] = Option(operationLog)
+
+ override protected def runInternal(): Unit = {
+ addTimeoutMonitor(queryTimeout)
+ if (shouldRunAsync) {
+ val asyncOperation = new Runnable {
+ override def run(): Unit = {
+ executeStatement()
+ }
+ }
+ val chatSessionManager = session.sessionManager
+ val backgroundHandle = chatSessionManager.submitBackgroundOperation(asyncOperation)
+ setBackgroundHandle(backgroundHandle)
+ } else {
+ executeStatement()
+ }
+ }
+
+ private def executeStatement(): Unit = {
+ setState(OperationState.RUNNING)
+
+ try {
+ val reply = chatProvider.ask(session.handle.identifier.toString, statement)
+ iter = new ArrayFetchIterator(Array(Array(reply)))
+
+ setState(OperationState.FINISHED)
+ } catch {
+ onError(true)
+ } finally {
+ shutdownTimeoutMonitor()
+ }
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatGPTProvider.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatGPTProvider.scala
new file mode 100644
index 000000000..f948eb154
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatGPTProvider.scala
@@ -0,0 +1,84 @@
+/*
+ * 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.engine.chat.provider
+
+import java.util
+import java.util.concurrent.TimeUnit
+
+import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache}
+import org.apache.http.HttpStatus
+import org.apache.http.client.methods.HttpPost
+import org.apache.http.entity.StringEntity
+import org.apache.http.impl.client.{CloseableHttpClient, HttpClientBuilder}
+import org.apache.http.util.EntityUtils
+
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.engine.chat.provider.ChatProvider.mapper
+
+class ChatGPTProvider(conf: KyuubiConf) extends ChatProvider {
+
+ val token = conf.get(KyuubiConf.ENGINE_CHAT_GPT_API_KEY).get
+
+ val httpClient: CloseableHttpClient = HttpClientBuilder.create().build()
+
+ private val chatHistory: LoadingCache[String, util.ArrayDeque[Message]] =
+ CacheBuilder.newBuilder()
+ .expireAfterWrite(10, TimeUnit.MINUTES)
+ .build(new CacheLoader[String, util.ArrayDeque[Message]] {
+ override def load(sessionId: String): util.ArrayDeque[Message] =
+ new util.ArrayDeque[Message]
+ })
+
+ override def open(sessionId: String): Unit = {
+ chatHistory.getIfPresent(sessionId)
+ }
+
+ override def ask(sessionId: String, q: String): String = {
+ val messages = chatHistory.get(sessionId)
+ messages.addLast(Message("user", q))
+
+ val request = new HttpPost("https://api.openai.com/v1/chat/completions")
+ request.addHeader("Content-Type", "application/json")
+ request.addHeader("Authorization", "Bearer " + token)
+
+ val req = Map(
+ "messages" -> messages,
+ "model" -> "gpt-3.5-turbo",
+ "max_tokens" -> 200,
+ "temperature" -> 0.5,
+ "top_p" -> 1)
+
+ request.setEntity(new StringEntity(mapper.writeValueAsString(req)))
+ val responseEntity = httpClient.execute(request)
+ val respJson = mapper.readTree(EntityUtils.toString(responseEntity.getEntity))
+ val statusCode = responseEntity.getStatusLine.getStatusCode
+ if (responseEntity.getStatusLine.getStatusCode == HttpStatus.SC_OK) {
+ val replyMessage = mapper.treeToValue[Message](
+ respJson.get("choices").get(0).get("message"))
+ messages.addLast(replyMessage)
+ replyMessage.content
+ } else {
+ messages.removeLast()
+ s"Chat failed. Status: $statusCode. ${respJson.get("error").get("message").asText}"
+ }
+ }
+
+ override def close(sessionId: String): Unit = {
+ chatHistory.invalidate(sessionId)
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatProvider.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatProvider.scala
new file mode 100644
index 000000000..af1ba434b
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ChatProvider.scala
@@ -0,0 +1,59 @@
+/*
+ * 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.engine.chat.provider
+
+import scala.util.control.NonFatal
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.{ClassTagExtensions, DefaultScalaModule}
+
+import org.apache.kyuubi.{KyuubiException, Logging}
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.reflection.DynConstructors
+
+trait ChatProvider {
+
+ def open(sessionId: String): Unit
+
+ def ask(sessionId: String, q: String): String
+
+ def close(sessionId: String): Unit
+}
+
+object ChatProvider extends Logging {
+
+ val mapper: ObjectMapper with ClassTagExtensions =
+ new ObjectMapper().registerModule(DefaultScalaModule) :: ClassTagExtensions
+
+ def load(conf: KyuubiConf): ChatProvider = {
+ val groupProviderClass = conf.get(KyuubiConf.ENGINE_CHAT_PROVIDER)
+ try {
+ DynConstructors.builder(classOf[ChatProvider])
+ .impl(groupProviderClass, classOf[KyuubiConf])
+ .impl(groupProviderClass)
+ .buildChecked
+ .newInstanceChecked(conf)
+ } catch {
+ case _: ClassCastException =>
+ throw new KyuubiException(
+ s"Class $groupProviderClass is not a child of '${classOf[ChatProvider].getName}'.")
+ case NonFatal(e) =>
+ throw new IllegalArgumentException(s"Error while instantiating '$groupProviderClass': ", e)
+ }
+ }
+}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/EchoProvider.scala
similarity index 73%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
copy to externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/EchoProvider.scala
index 88680a8c7..31ad3b8e3 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/EchoProvider.scala
@@ -15,13 +15,14 @@
* limitations under the License.
*/
-package org.apache.kyuubi.engine
+package org.apache.kyuubi.engine.chat.provider
-/**
- * Defines different engine types supported by Kyuubi.
- */
-object EngineType extends Enumeration {
- type EngineType = Value
+class EchoProvider extends ChatProvider {
+
+ override def open(sessionId: String): Unit = {}
+
+ override def ask(sessionId: String, q: String): String =
+ "This is ChatKyuubi, nice to meet you!"
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
+ override def close(sessionId: String): Unit = {}
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/Message.scala
similarity index 78%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
copy to externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/Message.scala
index 88680a8c7..e2162be9f 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/Message.scala
@@ -15,13 +15,6 @@
* limitations under the License.
*/
-package org.apache.kyuubi.engine
+package org.apache.kyuubi.engine.chat.provider
-/**
- * Defines different engine types supported by Kyuubi.
- */
-object EngineType extends Enumeration {
- type EngineType = Value
-
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
-}
+case class Message(role: String, content: String)
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/RowSet.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/RowSet.scala
new file mode 100644
index 000000000..3bb4ba7df
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/RowSet.scala
@@ -0,0 +1,108 @@
+/*
+ * 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.engine.chat.schema
+
+import java.util
+
+import org.apache.hive.service.rpc.thrift._
+
+import org.apache.kyuubi.util.RowSetUtils._
+
+object RowSet {
+
+ def emptyTRowSet(): TRowSet = {
+ new TRowSet(0, new java.util.ArrayList[TRow](0))
+ }
+
+ def toTRowSet(
+ rows: Seq[Array[String]],
+ columnSize: Int,
+ protocolVersion: TProtocolVersion): TRowSet = {
+ if (protocolVersion.getValue < TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6.getValue) {
+ toRowBasedSet(rows, columnSize)
+ } else {
+ toColumnBasedSet(rows, columnSize)
+ }
+ }
+
+ def toRowBasedSet(rows: Seq[Array[String]], columnSize: Int): TRowSet = {
+ val rowSize = rows.length
+ val tRows = new java.util.ArrayList[TRow](rowSize)
+ var i = 0
+ while (i < rowSize) {
+ val row = rows(i)
+ val tRow = new TRow()
+ var j = 0
+ val columnSize = row.length
+ while (j < columnSize) {
+ val columnValue = stringTColumnValue(j, row)
+ tRow.addToColVals(columnValue)
+ j += 1
+ }
+ i += 1
+ tRows.add(tRow)
+ }
+ new TRowSet(0, tRows)
+ }
+
+ def toColumnBasedSet(rows: Seq[Array[String]], columnSize: Int): TRowSet = {
+ val rowSize = rows.length
+ val tRowSet = new TRowSet(0, new util.ArrayList[TRow](rowSize))
+ var i = 0
+ while (i < columnSize) {
+ val tColumn = toTColumn(rows, i)
+ tRowSet.addToColumns(tColumn)
+ i += 1
+ }
+ tRowSet
+ }
+
+ private def toTColumn(rows: Seq[Array[String]], ordinal: Int): TColumn = {
+ val nulls = new java.util.BitSet()
+ val values = getOrSetAsNull[String](rows, ordinal, nulls, "")
+ TColumn.stringVal(new TStringColumn(values, nulls))
+ }
+
+ private def getOrSetAsNull[String](
+ rows: Seq[Array[String]],
+ ordinal: Int,
+ nulls: util.BitSet,
+ defaultVal: String): util.List[String] = {
+ val size = rows.length
+ val ret = new util.ArrayList[String](size)
+ var idx = 0
+ while (idx < size) {
+ val row = rows(idx)
+ val isNull = row(ordinal) == null
+ if (isNull) {
+ nulls.set(idx, true)
+ ret.add(idx, defaultVal)
+ } else {
+ ret.add(idx, row(ordinal))
+ }
+ idx += 1
+ }
+ ret
+ }
+
+ private def stringTColumnValue(ordinal: Int, row: Array[String]): TColumnValue = {
+ val tStringValue = new TStringValue
+ if (row(ordinal) != null) tStringValue.setValue(row(ordinal))
+ TColumnValue.stringVal(tStringValue)
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/SchemaHelper.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/SchemaHelper.scala
new file mode 100644
index 000000000..8ccfdda2f
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/schema/SchemaHelper.scala
@@ -0,0 +1,56 @@
+/*
+ * 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.engine.chat.schema
+
+import java.util.Collections
+
+import org.apache.hive.service.rpc.thrift._
+
+object SchemaHelper {
+
+ def stringTTypeQualifiers: TTypeQualifiers = {
+ val ret = new TTypeQualifiers()
+ val qualifiers = Collections.emptyMap[String, TTypeQualifierValue]()
+ ret.setQualifiers(qualifiers)
+ ret
+ }
+
+ def stringTTypeDesc: TTypeDesc = {
+ val typeEntry = new TPrimitiveTypeEntry(TTypeId.STRING_TYPE)
+ typeEntry.setTypeQualifiers(stringTTypeQualifiers)
+ val tTypeDesc = new TTypeDesc()
+ tTypeDesc.addToTypes(TTypeEntry.primitiveEntry(typeEntry))
+ tTypeDesc
+ }
+
+ def stringTColumnDesc(fieldName: String, pos: Int): TColumnDesc = {
+ val tColumnDesc = new TColumnDesc()
+ tColumnDesc.setColumnName(fieldName)
+ tColumnDesc.setTypeDesc(stringTTypeDesc)
+ tColumnDesc.setPosition(pos)
+ tColumnDesc
+ }
+
+ def stringTTableSchema(fieldsName: String*): TTableSchema = {
+ val tTableSchema = new TTableSchema()
+ fieldsName.zipWithIndex.foreach { case (f, i) =>
+ tTableSchema.addToColumns(stringTColumnDesc(f, i))
+ }
+ tTableSchema
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionImpl.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionImpl.scala
new file mode 100644
index 000000000..29f420768
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionImpl.scala
@@ -0,0 +1,68 @@
+/*
+ * 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.engine.chat.session
+
+import org.apache.hive.service.rpc.thrift.{TGetInfoType, TGetInfoValue, TProtocolVersion}
+
+import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException}
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY
+import org.apache.kyuubi.session.{AbstractSession, SessionHandle, SessionManager}
+
+class ChatSessionImpl(
+ protocol: TProtocolVersion,
+ user: String,
+ password: String,
+ ipAddress: String,
+ conf: Map[String, String],
+ sessionManager: SessionManager)
+ extends AbstractSession(protocol, user, password, ipAddress, conf, sessionManager) {
+
+ override val handle: SessionHandle =
+ conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID).getOrElse(SessionHandle())
+
+ private val chatProvider = sessionManager.asInstanceOf[ChatSessionManager].chatProvider
+
+ override def open(): Unit = {
+ info(s"Starting to open chat session.")
+ chatProvider.open(handle.identifier.toString)
+ super.open()
+ info(s"The chat session is started.")
+ }
+
+ override def getInfo(infoType: TGetInfoType): TGetInfoValue = withAcquireRelease() {
+ infoType match {
+ case TGetInfoType.CLI_SERVER_NAME | TGetInfoType.CLI_DBMS_NAME =>
+ TGetInfoValue.stringValue("Kyuubi Chat Engine")
+ case TGetInfoType.CLI_DBMS_VER =>
+ TGetInfoValue.stringValue(KYUUBI_VERSION)
+ case TGetInfoType.CLI_ODBC_KEYWORDS => TGetInfoValue.stringValue("Unimplemented")
+ case TGetInfoType.CLI_MAX_COLUMN_NAME_LEN =>
+ TGetInfoValue.lenValue(128)
+ case TGetInfoType.CLI_MAX_SCHEMA_NAME_LEN =>
+ TGetInfoValue.lenValue(128)
+ case TGetInfoType.CLI_MAX_TABLE_NAME_LEN =>
+ TGetInfoValue.lenValue(128)
+ case _ => throw KyuubiSQLException(s"Unrecognized GetInfoType value: $infoType")
+ }
+ }
+
+ override def close(): Unit = {
+ chatProvider.close(handle.identifier.toString)
+ super.close()
+ }
+
+}
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionManager.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionManager.scala
new file mode 100644
index 000000000..33a9dd450
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/session/ChatSessionManager.scala
@@ -0,0 +1,71 @@
+/*
+ * 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.engine.chat.session
+
+import org.apache.hive.service.rpc.thrift.TProtocolVersion
+
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.ENGINE_SHARE_LEVEL
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY
+import org.apache.kyuubi.engine.ShareLevel
+import org.apache.kyuubi.engine.chat.ChatEngine
+import org.apache.kyuubi.engine.chat.operation.ChatOperationManager
+import org.apache.kyuubi.engine.chat.provider.ChatProvider
+import org.apache.kyuubi.operation.OperationManager
+import org.apache.kyuubi.session.{Session, SessionHandle, SessionManager}
+
+class ChatSessionManager(name: String)
+ extends SessionManager(name) {
+
+ def this() = this(classOf[ChatSessionManager].getSimpleName)
+
+ override protected def isServer: Boolean = false
+
+ lazy val chatProvider: ChatProvider = ChatProvider.load(conf)
+
+ override lazy val operationManager: OperationManager =
+ new ChatOperationManager(conf, chatProvider)
+
+ override def initialize(conf: KyuubiConf): Unit = {
+ this.conf = conf
+ super.initialize(conf)
+ }
+
+ override protected def createSession(
+ protocol: TProtocolVersion,
+ user: String,
+ password: String,
+ ipAddress: String,
+ conf: Map[String, String]): Session = {
+ conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID)
+ .flatMap(getSessionOption).getOrElse {
+ new ChatSessionImpl(protocol, user, password, ipAddress, conf, this)
+ }
+ }
+
+ override def closeSession(sessionHandle: SessionHandle): Unit = {
+ super.closeSession(sessionHandle)
+ if (conf.get(ENGINE_SHARE_LEVEL) == ShareLevel.CONNECTION.toString) {
+ info("Session stopped due to shared level is Connection.")
+ stopSession()
+ }
+ }
+
+ private def stopSession(): Unit = {
+ ChatEngine.currentEngine.foreach(_.stop())
+ }
+}
diff --git a/externals/kyuubi-chat-engine/src/test/resources/log4j2-test.xml b/externals/kyuubi-chat-engine/src/test/resources/log4j2-test.xml
new file mode 100644
index 000000000..585a12c6f
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/test/resources/log4j2-test.xml
@@ -0,0 +1,42 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ 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.
+ -->
+
+<!-- Extra logging related to initialization of Log4j.
+ Set to debug or trace if log4j initialization is failing. -->
+<Configuration status="WARN">
+ <Appenders>
+ <Console name="stdout" target="SYSTEM_OUT">
+ <PatternLayout pattern="%d{HH:mm:ss.SSS} %p %c: %m%n"/>
+ <Filters>
+ <ThresholdFilter level="FATAL"/>
+ <RegexFilter regex=".*Thrift error occurred during processing of message.*" onMatch="DENY" onMismatch="NEUTRAL"/>
+ </Filters>
+ </Console>
+ <File name="file" fileName="target/unit-tests.log">
+ <PatternLayout pattern="%d{HH:mm:ss.SSS} %t %p %c{1}: %m%n"/>
+ <Filters>
+ <RegexFilter regex=".*Thrift error occurred during processing of message.*" onMatch="DENY" onMismatch="NEUTRAL"/>
+ </Filters>
+ </File>
+ </Appenders>
+ <Loggers>
+ <Root level="INFO">
+ <AppenderRef ref="file"/>
+ </Root>
+ </Loggers>
+</Configuration>
diff --git a/externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/WithChatEngine.scala b/externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/WithChatEngine.scala
new file mode 100644
index 000000000..287fdde2f
--- /dev/null
+++ b/externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/WithChatEngine.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.engine.chat
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.config.KyuubiConf
+
+trait WithChatEngine extends KyuubiFunSuite {
+
+ protected var engine: ChatEngine = _
+ protected var connectionUrl: String = _
+
+ protected val kyuubiConf: KyuubiConf = ChatEngine.kyuubiConf
+
+ def withKyuubiConf: Map[String, String]
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ startChatEngine()
+ }
+
+ override def afterAll(): Unit = {
+ stopChatEngine()
+ super.afterAll()
+ }
+
+ def stopChatEngine(): Unit = {
+ if (engine != null) {
+ engine.stop()
+ engine = null
+ }
+ }
+
+ def startChatEngine(): Unit = {
+ withKyuubiConf.foreach { case (k, v) =>
+ System.setProperty(k, v)
+ kyuubiConf.set(k, v)
+ }
+ ChatEngine.startEngine()
+ engine = ChatEngine.currentEngine.get
+ connectionUrl = engine.frontendServices.head.connectionUrl
+ }
+
+ protected def jdbcConnectionUrl: String = s"jdbc:hive2://$connectionUrl/;"
+
+}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationSuite.scala
similarity index 52%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
copy to externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationSuite.scala
index 88680a8c7..b14407a26 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/externals/kyuubi-chat-engine/src/test/scala/org/apache/kyuubi/engine/chat/operation/ChatOperationSuite.scala
@@ -15,13 +15,26 @@
* limitations under the License.
*/
-package org.apache.kyuubi.engine
+package org.apache.kyuubi.engine.chat.operation
-/**
- * Defines different engine types supported by Kyuubi.
- */
-object EngineType extends Enumeration {
- type EngineType = Value
+import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.engine.chat.WithChatEngine
+import org.apache.kyuubi.operation.HiveJDBCTestHelper
+
+class ChatOperationSuite extends HiveJDBCTestHelper with WithChatEngine {
+
+ override def withKyuubiConf: Map[String, String] = Map(
+ ENGINE_CHAT_PROVIDER.key -> "echo")
+
+ override protected def jdbcUrl: String = jdbcConnectionUrl
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
+ test("test echo chat provider") {
+ withJdbcStatement() { stmt =>
+ val result = stmt.executeQuery("Hello, Kyuubi")
+ assert(result.next())
+ val expected = "This is ChatKyuubi, nice to meet you!"
+ assert(result.getString("reply") === expected)
+ assert(!result.next())
+ }
+ }
}
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 1b7737344..b39bee307 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
@@ -1817,6 +1817,7 @@ object KyuubiConf {
" all the capacity of the Hive Server2.</li>" +
" <li>JDBC: specify this engine type will launch a JDBC engine which can provide" +
" a MySQL protocol connector, for now we only support Doris dialect.</li>" +
+ " <li>CHAT: specify this engine type will launch a Chat engine.</li>" +
"</ul>")
.version("1.4.0")
.stringConf
@@ -2621,6 +2622,51 @@ object KyuubiConf {
Map(configs.map { cfg => cfg.key -> cfg }: _*)
}
+ val ENGINE_CHAT_MEMORY: ConfigEntry[String] =
+ buildConf("kyuubi.engine.chat.memory")
+ .doc("The heap memory for the Chat engine")
+ .version("1.8.0")
+ .stringConf
+ .createWithDefault("1g")
+
+ val ENGINE_CHAT_JAVA_OPTIONS: OptionalConfigEntry[String] =
+ buildConf("kyuubi.engine.chat.java.options")
+ .doc("The extra Java options for the Chat engine")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
+
+ val ENGINE_CHAT_PROVIDER: ConfigEntry[String] =
+ buildConf("kyuubi.engine.chat.provider")
+ .doc("The provider for the Chat engine. Candidates: <ul>" +
+ " <li>ECHO: simply replies a welcome message.</li>" +
+ " <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li>" +
+ "</ul>")
+ .version("1.8.0")
+ .stringConf
+ .transform {
+ case "ECHO" | "echo" => "org.apache.kyuubi.engine.chat.provider.EchoProvider"
+ case "GPT" | "gpt" | "ChatGPT" => "org.apache.kyuubi.engine.chat.provider.ChatGPTProvider"
+ case other => other
+ }
+ .createWithDefault("ECHO")
+
+ val ENGINE_CHAT_GPT_API_KEY: OptionalConfigEntry[String] =
+ buildConf("kyuubi.engine.chat.gpt.apiKey")
+ .doc("The key to access OpenAI open API, which could be got at " +
+ "https://platform.openai.com/account/api-keys")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
+
+ val ENGINE_CHAT_EXTRA_CLASSPATH: OptionalConfigEntry[String] =
+ buildConf("kyuubi.engine.chat.extra.classpath")
+ .doc("The extra classpath for the Chat engine, for configuring the location " +
+ "of the SDK and etc.")
+ .version("1.8.0")
+ .stringConf
+ .createOptional
+
val ENGINE_JDBC_MEMORY: ConfigEntry[String] =
buildConf("kyuubi.engine.jdbc.memory")
.doc("The heap memory for the JDBC query engine")
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
index 88680a8c7..3d850ba14 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/EngineType.scala
@@ -23,5 +23,5 @@ package org.apache.kyuubi.engine
object EngineType extends Enumeration {
type EngineType = Value
- val SPARK_SQL, FLINK_SQL, TRINO, HIVE_SQL, JDBC = Value
+ val SPARK_SQL, FLINK_SQL, CHAT, TRINO, HIVE_SQL, JDBC = Value
}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala
index e2ddb4221..5f3af28f8 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala
@@ -29,8 +29,9 @@ import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException, Logging, Utils}
import org.apache.kyuubi.config.KyuubiConf
import org.apache.kyuubi.config.KyuubiConf._
import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_ENGINE_SUBMIT_TIME_KEY
-import org.apache.kyuubi.engine.EngineType.{EngineType, FLINK_SQL, HIVE_SQL, JDBC, SPARK_SQL, TRINO}
+import org.apache.kyuubi.engine.EngineType._
import org.apache.kyuubi.engine.ShareLevel.{CONNECTION, GROUP, SERVER, ShareLevel}
+import org.apache.kyuubi.engine.chat.ChatProcessBuilder
import org.apache.kyuubi.engine.flink.FlinkProcessBuilder
import org.apache.kyuubi.engine.hive.HiveProcessBuilder
import org.apache.kyuubi.engine.jdbc.JdbcProcessBuilder
@@ -192,6 +193,8 @@ private[kyuubi] class EngineRef(
new HiveProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
case JDBC =>
new JdbcProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
+ case CHAT =>
+ new ChatProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
}
MetricsSystem.tracing(_.incCount(ENGINE_TOTAL))
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/chat/ChatProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/chat/ChatProcessBuilder.scala
new file mode 100644
index 000000000..532c5dddf
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/chat/ChatProcessBuilder.scala
@@ -0,0 +1,115 @@
+/*
+ * 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.engine.chat
+
+import java.io.File
+import java.nio.file.{Files, Paths}
+import java.util
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
+
+import com.google.common.annotations.VisibleForTesting
+
+import org.apache.kyuubi.{Logging, SCALA_COMPILE_VERSION, Utils}
+import org.apache.kyuubi.Utils.REDACTION_REPLACEMENT_TEXT
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
+import org.apache.kyuubi.engine.ProcBuilder
+import org.apache.kyuubi.operation.log.OperationLog
+
+class ChatProcessBuilder(
+ override val proxyUser: String,
+ override val conf: KyuubiConf,
+ val engineRefId: String,
+ val extraEngineLog: Option[OperationLog] = None)
+ extends ProcBuilder with Logging {
+
+ @VisibleForTesting
+ def this(proxyUser: String, conf: KyuubiConf) {
+ this(proxyUser, conf, "")
+ }
+
+ /**
+ * The short name of the engine process builder, we use this for form the engine jar paths now
+ * see `mainResource`
+ */
+ override def shortName: String = "chat"
+
+ override protected def module: String = "kyuubi-chat-engine"
+
+ /**
+ * The class containing the main method
+ */
+ override protected def mainClass: String = "org.apache.kyuubi.engine.chat.ChatEngine"
+
+ override protected val commands: Array[String] = {
+ val buffer = new ArrayBuffer[String]()
+ buffer += executable
+
+ val memory = conf.get(ENGINE_CHAT_MEMORY)
+ buffer += s"-Xmx$memory"
+
+ val javaOptions = conf.get(ENGINE_CHAT_JAVA_OPTIONS)
+ javaOptions.foreach(buffer += _)
+
+ buffer += "-cp"
+ val classpathEntries = new util.LinkedHashSet[String]
+ mainResource.foreach(classpathEntries.add)
+ mainResource.foreach { path =>
+ val parent = Paths.get(path).getParent
+ val chatDevDepDir = parent
+ .resolve(s"scala-$SCALA_COMPILE_VERSION")
+ .resolve("jars")
+ if (Files.exists(chatDevDepDir)) {
+ // add dev classpath
+ classpathEntries.add(s"$chatDevDepDir${File.separator}*")
+ } else {
+ // add prod classpath
+ classpathEntries.add(s"$parent${File.separator}*")
+ }
+ }
+
+ val extraCp = conf.get(ENGINE_CHAT_EXTRA_CLASSPATH)
+ extraCp.foreach(classpathEntries.add)
+ buffer += classpathEntries.asScala.mkString(File.pathSeparator)
+ buffer += mainClass
+
+ buffer += "--conf"
+ buffer += s"$KYUUBI_SESSION_USER_KEY=$proxyUser"
+
+ for ((k, v) <- conf.getAll) {
+ buffer += "--conf"
+ buffer += s"$k=$v"
+ }
+ buffer.toArray
+ }
+
+ override def toString: String = {
+ if (commands == null) {
+ super.toString()
+ } else {
+ Utils.redactCommandLineArgs(conf, commands).map {
+ case arg if arg.contains(ENGINE_CHAT_GPT_API_KEY.key) =>
+ s"${ENGINE_CHAT_GPT_API_KEY.key}=$REDACTION_REPLACEMENT_TEXT"
+ case arg => arg
+ }.mkString("\n")
+ }
+ }
+}
diff --git a/pom.xml b/pom.xml
index 2082499dc..a2de5d29d 100644
--- a/pom.xml
+++ b/pom.xml
@@ -60,10 +60,11 @@
<module>extensions/spark/kyuubi-spark-connector-tpcds</module>
<module>extensions/spark/kyuubi-spark-connector-tpch</module>
<module>extensions/spark/kyuubi-spark-lineage</module>
+ <module>externals/kyuubi-chat-engine</module>
<module>externals/kyuubi-download</module>
<module>externals/kyuubi-flink-sql-engine</module>
- <module>externals/kyuubi-jdbc-engine</module>
<module>externals/kyuubi-hive-sql-engine</module>
+ <module>externals/kyuubi-jdbc-engine</module>
<module>externals/kyuubi-spark-sql-engine</module>
<module>externals/kyuubi-trino-engine</module>
<module>integration-tests</module>