How to Add Custom SQL Completion Rules to Chat2DB

Chat2DB exposes a ServiceLoader-based SPI (ISqlCompletionSlotRule) that lets you inject custom completion logic by implementing the interface, registering it in META-INF/services, and packaging it as a plugin module.

Chat2DB is an open-source database client that provides intelligent SQL editing capabilities. If you need to extend its autocomplete behavior with domain-specific keywords, table hints, or AI-driven suggestions, you can add custom SQL completion rules to Chat2DB without modifying the core codebase. The platform uses a plugin architecture centered on the chat2db-community-spi module, which discovers rules at runtime using Java's ServiceLoader mechanism.

Understanding the Completion Architecture

The SQL completion engine in Chat2DB is built on a decoupled, pipeline-based architecture. It discovers implementations at runtime, executes them against the current cursor context, and aggregates the results into ranked suggestions.

The Extension Point (ISqlCompletionSlotRule)

The core contract for custom rules is the ISqlCompletionSlotRule interface, located at:

chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/ISqlCompletionSlotRule.java

This interface defines a single method, apply(SqlCompletionPipelineState state, SqlCompletionSlot slot), which receives the current pipeline state and a slot object describing the cursor position. Implementations return a List<SqlCompletionCandidate> containing the suggested completions.

The Processing Pipeline

The SqlCompletionPipeline.java class orchestrates execution. It builds a SqlCompletionPipelineState object that carries the parsed SQL, cursor offset, and list of completion slots. The pipeline then iterates over every discovered ISqlCompletionSlotRule implementation, invoking apply() to gather candidates, evidence, and relevance scores before merging them into the final response.

Step-by-Step Implementation Guide

To add a custom rule, you must implement the interface, register it via the ServiceLoader mechanism, and ensure it is available on the classpath.

Step 1: Implement the ISqlCompletionSlotRule Interface

Create a new Java class that implements ai.chat2db.spi.ISqlCompletionSlotRule. The following example demonstrates a rule that suggests the LIMIT keyword when the cursor appears immediately after a SELECT clause:

package com.mycompany.chat2db.completion;

import ai.chat2db.community.domain.api.model.completion.slot.SqlCompletionSlot;
import ai.chat2db.community.domain.api.model.completion.result.SqlCompletionCandidate;
import ai.chat2db.spi.parser.completion.SqlCompletionPipelineState;
import ai.chat2db.spi.ISqlCompletionSlotRule;
import java.util.Collections;
import java.util.List;

/** Example rule that suggests the keyword “LIMIT” when the cursor is after a SELECT clause. */
public class LimitKeywordCompletionRule implements ISqlCompletionSlotRule {

    @Override
    public List<SqlCompletionCandidate> apply(SqlCompletionPipelineState state, SqlCompletionSlot slot) {
        // Heuristic: if the slot is positioned after a SELECT statement, return LIMIT candidate
        if (slot.getBeforeText().trim().toUpperCase().endsWith("SELECT")) {
            SqlCompletionCandidate candidate = new SqlCompletionCandidate();
            candidate.setContent(" LIMIT ");
            candidate.setDisplayText("LIMIT – limit the result set");
            candidate.setType(ai.chat2db.community.domain.api.enums.completion.SqlCompletionCandidateTypeEnum.KEYWORD);
            return Collections.singletonList(candidate);
        }
        return Collections.emptyList();
    }
}

The SqlCompletionSlot parameter provides access to the text before and after the cursor, while SqlCompletionPipelineState exposes the parsed SQL structure and database context.

Step 2: Register Your Implementation with ServiceLoader

Java’s ServiceLoader requires a provider-configuration file to locate implementations at runtime. Create a file named ai.chat2db.spi.ISqlCompletionSlotRule (the fully-qualified interface name) in:

src/main/resources/META-INF/services/

Populate this file with the fully-qualified class name of your implementation, one per line:

com.mycompany.chat2db.completion.LimitKeywordCompletionRule

At runtime, Chat2DB scans the classpath for these files and instantiates the listed classes automatically.

Step 3: Package as a Maven Module

Ensure your module declares a dependency on the SPI artifact so it can access the required interfaces:

<dependency>
    <groupId>ai.chat2db</groupId>
    <artifactId>chat2db-community-spi</artifactId>
    <version>${chat2db.version}</version>
</dependency>

Build the project with mvn clean package. The resulting JAR containing your META-INF/services registration and compiled rule class should be placed on Chat2DB’s classpath (either in the plugins directory or bundled into the main application classpath).

Runtime Execution Flow

When a user triggers autocomplete in the Chat2DB interface, the following sequence occurs:

1. Request ingress: The frontend sends a SqlCompletionRequest to the DbSqlParserController.java REST endpoint (chat2db-community-server/chat2db-community-web/src/main/java/ai/chat2db/community/web/api/controller/DbSqlParserController.java).
2. Pipeline initialization: The controller delegates to SqlCompletionPipeline, which constructs a SqlCompletionPipelineState from the request.
3. Rule discovery: The pipeline uses ServiceLoader.load(ISqlCompletionSlotRule.class) to discover all registered implementations.
4. Execution: Each rule’s apply() method is invoked with the current state and slot context.
5. Aggregation: Results are merged, scored, and returned as a SqlCompletionResponse to the client.

This architecture ensures that your custom logic executes alongside built-in rules without requiring changes to the core server code.

Testing Your Custom Rule

Unit test your implementation by mocking the pipeline state and asserting on the returned candidates:

@Test
public void shouldSuggestLimitKeyword() {
    DbSqlCompletionRequest request = new DbSqlCompletionRequest();
    request.setSql("SELECT * FROM users ");
    request.setCursorPosition(request.getSql().length()); // cursor after the space

    // Obtain the pipeline bean from Spring context or instantiate directly
    SqlCompletionResponse response = sqlCompletionPipeline.execute(request);

    // Verify that the response contains our custom candidate
    assertTrue(response.getCandidates()
               .stream()
               .anyMatch(c -> "LIMIT".equalsIgnoreCase(c.getContent().trim())));
}

Refer to existing tests such as DefaultSqlSyntaxHandlerCompletionTest.java for patterns on constructing request objects and asserting completion behavior.

Summary

• Implement ISqlCompletionSlotRule to define custom logic that inspects cursor context and returns SqlCompletionCandidate objects.
• Register your implementation by creating a file at META-INF/services/ai.chat2db.spi.ISqlCompletionSlotRule containing the fully-qualified class name.
• Package the rule as a Maven module depending on chat2db-community-spi and ensure it is present on the classpath at runtime.
• Test using the SqlCompletionPipeline directly with mocked requests to verify candidate generation.
• Deploy by rebuilding Chat2DB with your module included; the ServiceLoader discovers the rule automatically without code changes to the core platform.

Frequently Asked Questions

What interface must I implement to add custom SQL completion rules to Chat2DB?

You must implement ai.chat2db.spi.ISqlCompletionSlotRule, located in the chat2db-community-spi module. This interface defines the apply(SqlCompletionPipelineState, SqlCompletionSlot) method where you return a list of completion candidates based on the current cursor position and SQL context.

Where do I place the ServiceLoader configuration file?

Create a file named exactly ai.chat2db.spi.ISqlCompletionSlotRule inside src/main/resources/META-INF/services/. The file should contain the fully-qualified class name of your implementation, one per line. This follows the standard Java ServiceLoader pattern that Chat2DB uses for plugin discovery.

Can I add multiple custom completion rules to a single Chat2DB instance?

Yes. The ServiceLoader mechanism supports multiple providers. Simply list each fully-qualified implementation class on a separate line in the META-INF/services/ai.chat2db.spi.ISqlCompletionSlotRule file, or distribute them across multiple JARs each containing their own service descriptor. Chat2DB’s SqlCompletionPipeline will invoke every discovered rule and merge the results.

How do I debug why my custom completion rule is not appearing?

First, verify that your JAR contains the META-INF/services/ai.chat2db.spi.ISqlCompletionSlotRule file and that the class name is spelled correctly. Second, ensure the JAR is on the runtime classpath. Finally, add logging inside your apply() method or set a breakpoint in SqlCompletionPipeline.java where ServiceLoader.load(ISqlCompletionSlotRule.class) is called to confirm your implementation is being instantiated.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}