[KYUUBI #7379][2b/4] Data Agent Engine: agent runtime, middleware, and OpenAI provider

build/dist

@@ -288,6 +288,7 @@ mkdir -p "$DISTDIR/externals/engines/trino"
 mkdir -p "$DISTDIR/externals/engines/hive"
 mkdir -p "$DISTDIR/externals/engines/jdbc"
 mkdir -p "$DISTDIR/externals/engines/chat"
+mkdir -p "$DISTDIR/externals/engines/data-agent"
 echo "Kyuubi $VERSION $GITREVSTRING built for" > "$DISTDIR/RELEASE"
 echo "Java $JAVA_VERSION" >> "$DISTDIR/RELEASE"
 echo "Scala $SCALA_VERSION" >> "$DISTDIR/RELEASE"
@@ -363,6 +364,18 @@ for jar in $(ls "$DISTDIR/jars/"); do
   fi
 done
 
+# Copy data-agent engines
+cp "$KYUUBI_HOME/externals/kyuubi-data-agent-engine/target/kyuubi-data-agent-engine_${SCALA_VERSION}-${VERSION}.jar" "$DISTDIR/externals/engines/data-agent/"
+cp -r "$KYUUBI_HOME"/externals/kyuubi-data-agent-engine/target/scala-$SCALA_VERSION/jars/*.jar "$DISTDIR/externals/engines/data-agent/"
+
+# Share the jars w/ server to reduce binary size
+# shellcheck disable=SC2045
+for jar in $(ls "$DISTDIR/jars/"); do
+  if [[ -f "$DISTDIR/externals/engines/data-agent/$jar" ]]; then
+    (cd $DISTDIR/externals/engines/data-agent; ln -snf "../../../jars/$jar" "$DISTDIR/externals/engines/data-agent/$jar")
+  fi
+done
+
 # Copy Kyuubi Spark extension
 # shellcheck disable=SC2068
 for SPARK_EXTENSION_VERSION in ${SPARK_EXTENSION_VERSIONS[@]}; do

externals/kyuubi-data-agent-engine/pom.xml

@@ -0,0 +1,141 @@
+<?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.12.0-SNAPSHOT</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>kyuubi-data-agent-engine_${scala.binary.version}</artifactId>
+    <packaging>jar</packaging>
+    <name>Kyuubi Project Engine Data Agent</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.kyuubi</groupId>
+            <artifactId>${hive.jdbc.artifact}</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+
+        <!-- OpenAI official Java SDK -->
+        <dependency>
+            <groupId>com.openai</groupId>
+            <artifactId>openai-java</artifactId>
+            <version>${openai.sdk.version}</version>
+        </dependency>
+
+        <!-- JSON Schema generation from Jackson-annotated classes -->
+        <dependency>
+            <groupId>com.github.victools</groupId>
+            <artifactId>jsonschema-generator</artifactId>
+            <version>${victools.jsonschema.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.github.victools</groupId>
+            <artifactId>jsonschema-module-jackson</artifactId>
+            <version>${victools.jsonschema.version}</version>
+        </dependency>
+
+        <!-- SQLite JDBC driver -->
+        <dependency>
+            <groupId>org.xerial</groupId>
+            <artifactId>sqlite-jdbc</artifactId>
+            <version>${sqlite.version}</version>
+        </dependency>
+
+        <!-- MySQL JDBC driver (also works for StarRocks) -->
+        <dependency>
+            <groupId>com.mysql</groupId>
+            <artifactId>mysql-connector-j</artifactId>
+        </dependency>
+
+        <!-- Trino JDBC driver -->
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-jdbc</artifactId>
+        </dependency>
+
+        <!-- Connection pool -->
+        <dependency>
+            <groupId>com.zaxxer</groupId>
+            <artifactId>HikariCP</artifactId>
+        </dependency>
+
+        <!-- test dependencies -->
+        <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>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <configuration>
+                    <skipTests>${skipTests}</skipTests>
+                    <argLine>${extraJavaTestArgs}</argLine>
+                </configuration>
+            </plugin>
+            <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>

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/DataSourceFactory.java

@@ -0,0 +1,64 @@
+/*
+ * 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.dataagent.datasource;
+
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import javax.sql.DataSource;
+
+/** Factory for creating pooled DataSource instances from JDBC URLs. */
+public final class DataSourceFactory {
+
+  private static final int DEFAULT_MAX_POOL_SIZE = 5;
+
+  private DataSourceFactory() {}
+
+  /**
+   * Create a pooled DataSource from a JDBC URL. Supports any JDBC driver available on the
+   * classpath. For authenticated databases, credentials can be embedded in the URL.
+   *
+   * @param jdbcUrl the JDBC connection URL
+   * @return a HikariCP-backed DataSource
+   */
+  public static DataSource create(String jdbcUrl) {
+    return create(jdbcUrl, null);
+  }
+
+  /**
+   * Create a pooled DataSource from a JDBC URL with an explicit username. When the data-agent
+   * connects back to Kyuubi Server, the username determines the proxy user for the downstream
+   * engine (e.g. Spark). Without it, Kyuubi defaults to "anonymous" which typically fails Hadoop
+   * impersonation checks.
+   *
+   * @param jdbcUrl the JDBC connection URL
+   * @param user the username for the JDBC connection, may be null
+   * @return a HikariCP-backed DataSource
+   */
+  public static DataSource create(String jdbcUrl, String user) {
+    HikariConfig config = new HikariConfig();
+    config.setJdbcUrl(jdbcUrl);
+    if (user != null && !user.isEmpty()) {
+      config.setUsername(user);
+    }
+    config.setMaximumPoolSize(DEFAULT_MAX_POOL_SIZE);
+    config.setMinimumIdle(1);
+    config.setInitializationFailTimeout(-1);
+    config.setPoolName("kyuubi-data-agent");
+    return new HikariDataSource(config);
+  }
+}

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/JdbcDialect.java

@@ -0,0 +1,63 @@
+/*
+ * 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.dataagent.datasource;
+
+/**
+ * SQL dialect abstraction for datasource-specific SQL generation.
+ *
+ * <p>Each dialect maps to a datasource name used for prompt resource lookup ({@code
+ * prompts/datasource-{name}.md}).
+ */
+public interface JdbcDialect {
+
+  /** Datasource name for prompt resource lookup (e.g. "spark", "trino"). */
+  String datasourceName();
+
+  /**
+   * Quote an identifier (table/column/database name) using the dialect-appropriate quote character.
+   * Escapes any embedded quote characters by doubling them.
+   */
+  String quoteIdentifier(String identifier);
+
+  /**
+   * Infer the dialect from a JDBC URL.
+   *
+   * @param jdbcUrl the JDBC connection URL
+   * @return the matching dialect, or {@code null} if unrecognized
+   */
+  static JdbcDialect fromUrl(String jdbcUrl) {
+    if (jdbcUrl == null) {
+      return null;
+    }
+    String lower = jdbcUrl.toLowerCase();
+    if (lower.startsWith("jdbc:hive2:") || lower.startsWith("jdbc:spark:")) {
+      return SparkDialect.INSTANCE;
+    }
+    if (lower.startsWith("jdbc:trino:")) {
+      return TrinoDialect.INSTANCE;
+    }
+    if (lower.startsWith("jdbc:mysql:")) {
+      return MysqlDialect.INSTANCE;
+    }
+    if (lower.startsWith("jdbc:sqlite:")) {
+      return SqliteDialect.INSTANCE;
+    }
+    // MySQL as fallback for unrecognized JDBC URLs
+    return MysqlDialect.INSTANCE;
+  }
+}

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/MysqlDialect.java

@@ -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.engine.dataagent.datasource;
+
+/** MySQL dialect. Uses backtick quoting for identifiers. */
+public final class MysqlDialect implements JdbcDialect {
+
+  static final MysqlDialect INSTANCE = new MysqlDialect();
+
+  private MysqlDialect() {}
+
+  @Override
+  public String datasourceName() {
+    return "mysql";
+  }
+
+  @Override
+  public String quoteIdentifier(String identifier) {
+    String escaped = identifier.replace("`", "``");
+    return "`" + escaped + "`";
+  }
+}

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/SparkDialect.java

@@ -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.engine.dataagent.datasource;
+
+/** Spark SQL dialect. Uses backtick quoting for identifiers. */
+public final class SparkDialect implements JdbcDialect {
+
+  static final SparkDialect INSTANCE = new SparkDialect();
+
+  private SparkDialect() {}
+
+  @Override
+  public String datasourceName() {
+    return "spark";
+  }
+
+  @Override
+  public String quoteIdentifier(String identifier) {
+    String escaped = identifier.replace("`", "``");
+    return "`" + escaped + "`";
+  }
+}