[KYUUBI #7379][PR 2a/4] Data Agent Engine: tool system, data source, and prompt templates

Author: wangzhigang1999

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 + "`";
+  }
+}

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/SqliteDialect.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;
+
+/** SQLite dialect. Uses double-quote quoting for identifiers. */
+public final class SqliteDialect implements JdbcDialect {
+
+  static final SqliteDialect INSTANCE = new SqliteDialect();
+
+  private SqliteDialect() {}
+
+  @Override
+  public String datasourceName() {
+    return "sqlite";
+  }
+
+  @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/TrinoDialect.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;
+
+/** Trino SQL dialect. Uses double-quote quoting for identifiers. */
+public final class TrinoDialect implements JdbcDialect {
+
+  static final TrinoDialect INSTANCE = new TrinoDialect();
+
+  private TrinoDialect() {}
+
+  @Override
+  public String datasourceName() {
+    return "trino";
+  }
+
+  @Override
+  public String quoteIdentifier(String identifier) {
+    String escaped = identifier.replace("\"", "\"\"");
+    return "\"" + escaped + "\"";
+  }
+}