LLMs are quite good at producing the source code when they are provided enough context information and asked very clear questions. The same is true for generating SQL queries when they are provided the table/column information through DDL statements, and clear instructions of output format of the generated SQL statements.
Let us see a quick example of generating SQL statements when we use Spring AI to interact with LLM. The flow of sequence is:
- The user creates a prompt with instructions for LLM, DDL statement, and user question.
- LLM generates the SQL query.
- The application executes the SQL query into the database. In this demo, only ‘SELECT‘ queries are allowed.
- The query result is returned to the API user.
- Any create/insert/update query request is returned as a non-supported operation.
1. Project Setup
Start with creating a Spring boot application. The application uses OpenAI’s GPT model for generating SQL queries, so we add the spring-ai-openai-spring-boot-starter module as a dependency. Read more for detailed instructions on getting started with Spring AI.
For executing database queries, we will use JdbcTemplate and H2 database so add these dependencies as well.
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-jdbc</artifactId>
</dependency>
<dependency>
<groupId>com.h2database</groupId>
<artifactId>h2</artifactId>
<scope>runtime</scope>
</dependency>Do not forget to add the OpenAI API key in the properties file or environment variable.
spring.ai.openai.api-key=${OPENAI_API_KEY}2. DDL Statements and Initial Data
For the demo purpose, we are using the H2 database which is automatically configured by Spring Boot autoconfiguration. Spring boot will automatically execute the schema.sql and data.sql scripts when they are placed in the ‘/src/main/resources’ directory.
create table TBL_USER (
id int not null auto_increment,
username varchar(255) not null,
email varchar(255) not null,
password varchar(255) not null,
primary key (id)
);
create table TBL_ACCOUNT (
id int not null auto_increment,
accountNumber varchar(255) not null,
user_id int not null,
balance decimal(10, 2) not null,
openDate date not null,
primary key (id),
foreign key (user_id) references TBL_USER(id)
);INSERT INTO TBL_USER (username, email, password)
VALUES
('user1', '[email protected]', 'password1'),
('user2', '[email protected]', 'password2'),
('user3', '[email protected]', 'password3'),
('user4', '[email protected]', 'password4'),
('user5', '[email protected]', 'password5'),
('user6', '[email protected]', 'password6'),
('user7', '[email protected]', 'password7'),
('user8', '[email protected]', 'password8'),
('user9', '[email protected]', 'password9'),
('user10', '[email protected]', 'password10');
INSERT INTO TBL_ACCOUNT (accountNumber, user_id, balance, openDate)
VALUES
('ACC001', 1, 1000.00, '2024-07-09'),
('ACC002', 1, 500.00, '2024-07-10'),
('ACC003', 2, 1500.00, '2024-07-09'),
('ACC004', 2, 200.00, '2024-07-10'),
('ACC005', 3, 800.00, '2024-07-09'),
('ACC006', 4, 3000.00, '2024-07-09'),
('ACC007', 4, 100.00, '2024-07-10'),
('ACC008', 5, 250.00, '2024-07-09'),
('ACC009', 6, 1800.00, '2024-07-09'),
('ACC010', 6, 700.00, '2024-07-10'),
('ACC011', 7, 500.00, '2024-07-09'),
('ACC012', 8, 1200.00, '2024-07-09'),
('ACC013', 9, 900.00, '2024-07-09'),
('ACC014', 9, 300.00, '2024-07-10'),
('ACC015', 10, 2000.00, '2024-07-09');3. The Prompt
Writing a detailed and clear prompt is very important for generating the correct SQL statements that can be executed as it is by the application.
The following prompt requests for an SQL SELECT statement that can produce the request data from the database. The other create/update/delete operations are not supported for execution. You may choose to allow these statements but do not execute them and only return the generated query as API output.
Given the DDL in the DDL section, write an SQL query that answers the asked question in the QUESTION section.
Only produce select queries. Do not append any text or markup in the start or end of response.
Remove the markups such as ``` , sql , \n as well.
If the question would result in an insert, update, or delete, or if the query would alter the DDL in any way, say that the operation isn't supported.
If the question can't be answered, say that the DDL doesn't support answering that question.
QUESTION
{question}
DDL
{ddl}4. Sql Controller
Let’s write the API that uses ChatClient API to send the input prompt to LLM and receive the generated SQL query. If a valid query is generated, it uses the JdbcTemplate to execute that query and return the response.
import java.io.IOException;
import java.nio.charset.Charset;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
@RestController
public class SqlController {
@Value("classpath:/schema.sql")
private Resource ddlResource;
@Value("classpath:/sql-prompt-template.st")
private Resource sqlPromptTemplateResource;
private final ChatClient aiClient;
private final JdbcTemplate jdbcTemplate;
public SqlController(ChatClient.Builder aiClientBuilder, JdbcTemplate jdbcTemplate) {
this.aiClient = aiClientBuilder.build();
this.jdbcTemplate = jdbcTemplate;
}
@PostMapping(path = "/sql")
public AiResponse sql(@RequestBody AiRequest request) throws IOException {
String schema = ddlResource.getContentAsString(Charset.defaultCharset());
String query = aiClient.prompt()
.advisors(new SimpleLoggerAdvisor())
.user(userSpec -> userSpec
.text(sqlPromptTemplateResource)
.param("question", request.text())
.param("ddl", schema)
)
.call()
.content();
if (query.toLowerCase().startsWith("select")) {
return new AiResponse(query, jdbcTemplate.queryForList(query));
}
throw new AiException(query);
}
}
public record AiRequest(String text) { }
public record AiResponse(String sqlQuery, List<Map<String, Object>> results) { }5. Exception Handling
Not all user requests will be valid. Some users will try to trick the system and generate queries that alter the database schema or stored data. For such queries, we have used the Spring ProblemDetail API to return the error message in a standard format.
public class AiException extends RuntimeException {
public AiException(String response) {
super(response);
}
}@RestControllerAdvice
public class CustomExceptionHandler {
@ExceptionHandler(AiException.class)
public ProblemDetail handle(AiException ex) {
return ProblemDetail.forStatusAndDetail(HttpStatus.EXPECTATION_FAILED, ex.getMessage());
}
}6. Demo
Start the application in the embedded server which, by default, listens to port ‘8080‘. Now, send some valid and invalid SQL generation requests and the application will respond accordingly.
User request: “Find the count of accounts.”
API response:
{
"sqlQuery": "select count(*) as account_count from TBL_ACCOUNT;",
"results": [
{
"ACCOUNT_COUNT": 15
}
]
}
User request: “Sum of all accounts for all users.”
API response:
{
"sqlQuery": "select sum(balance) as total_balance\r\nfrom TBL_ACCOUNT;",
"results": [
{
"TOTAL_BALANCE": 14750.00
}
]
}
Invalid User request: “Empty the account balance of user1.”
API response:
{
"type": "about:blank",
"title": "Expectation Failed",
"status": 417,
"detail": "The operation isn't supported.",
"instance": "/sql"
}
7. Summary
In this short tutorial, we discussed how to generate the SQL statements from an LLM using Spring AI, and execute the SQL statements using Spring JDBC support. This SQL generation capability is highly useful for developers writing basic as well as complex SQL statements involving joins.
Do not forget to validate the generated SQL for all types of security and audit-related concerns. Sometimes, LLM-generated SQL can be quite wrong as well, so verify the SQL statements before using them in production.
Happy Learning !!
Comments