feat(completion): add scoped metadata assistant

agents/common/src/main/java/com/dbx/agent/AgentProtocol.java

@@ -15,6 +15,7 @@ public final class AgentProtocol {
     public static final String METHOD_LIST_SCHEMAS = "list_schemas";
     public static final String METHOD_LIST_TABLES = "list_tables";
     public static final String METHOD_LIST_OBJECTS = "list_objects";
+    public static final String METHOD_COMPLETION_ASSISTANT_SEARCH_V1 = "completion_assistant_search_v1";
     public static final String METHOD_GET_OBJECT_SOURCE = "get_object_source";
     public static final String METHOD_GET_TABLE_DDL = "get_table_ddl";
     public static final String METHOD_GET_COLUMNS = "get_columns";
@@ -81,6 +82,7 @@ public final class AgentProtocol {
         METHOD_LIST_SCHEMAS,
         METHOD_LIST_TABLES,
         METHOD_LIST_OBJECTS,
+        METHOD_COMPLETION_ASSISTANT_SEARCH_V1,
         METHOD_GET_OBJECT_SOURCE,
         METHOD_GET_TABLE_DDL,
         METHOD_GET_COLUMNS,

agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java

@@ -20,6 +20,11 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent {
         throw new UnsupportedOperationException("Object source is not supported");
     }
 
+    @Override
+    public CompletionAssistantResponse completionAssistantSearch(CompletionAssistantRequest request) {
+        throw new UnsupportedOperationException("Completion assistant search is not supported by this agent");
+    }
+
     @Override
     public String getTableDdl(String schema, String table) {
         List<IndexInfo> indexes;

agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidate.java

@@ -0,0 +1,41 @@
+package com.dbx.agent;
+
+public final class CompletionAssistantCandidate {
+    private final String name;
+    private final CompletionAssistantCandidateKind kind;
+    private final String database;
+    private final String schema;
+    private final String parent_schema;
+    private final String parent_name;
+    private final String comment;
+    private final String data_type;
+
+    public CompletionAssistantCandidate(
+        String name,
+        CompletionAssistantCandidateKind kind,
+        String database,
+        String schema,
+        String parentSchema,
+        String parentName,
+        String comment,
+        String dataType
+    ) {
+        this.name = name;
+        this.kind = kind;
+        this.database = database;
+        this.schema = schema;
+        this.parent_schema = parentSchema;
+        this.parent_name = parentName;
+        this.comment = comment;
+        this.data_type = dataType;
+    }
+
+    public String getName() { return name; }
+    public CompletionAssistantCandidateKind getKind() { return kind; }
+    public String getDatabase() { return database; }
+    public String getSchema() { return schema; }
+    public String getParent_schema() { return parent_schema; }
+    public String getParent_name() { return parent_name; }
+    public String getComment() { return comment; }
+    public String getData_type() { return data_type; }
+}

agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidateKind.java

@@ -0,0 +1,22 @@
+package com.dbx.agent;
+
+import com.google.gson.annotations.SerializedName;
+
+public enum CompletionAssistantCandidateKind {
+    @SerializedName("database")
+    DATABASE,
+    @SerializedName("schema")
+    SCHEMA,
+    @SerializedName("table")
+    TABLE,
+    @SerializedName("view")
+    VIEW,
+    @SerializedName("procedure")
+    PROCEDURE,
+    @SerializedName("function")
+    FUNCTION,
+    @SerializedName("column")
+    COLUMN,
+    @SerializedName("object")
+    OBJECT
+}

agents/common/src/main/java/com/dbx/agent/CompletionAssistantMatchMode.java

@@ -0,0 +1,10 @@
+package com.dbx.agent;
+
+import com.google.gson.annotations.SerializedName;
+
+public enum CompletionAssistantMatchMode {
+    @SerializedName("prefix")
+    PREFIX,
+    @SerializedName("contains")
+    CONTAINS
+}

agents/common/src/main/java/com/dbx/agent/CompletionAssistantObjectKind.java

@@ -0,0 +1,22 @@
+package com.dbx.agent;
+
+import com.google.gson.annotations.SerializedName;
+
+public enum CompletionAssistantObjectKind {
+    @SerializedName("database")
+    DATABASE,
+    @SerializedName("schema")
+    SCHEMA,
+    @SerializedName("table")
+    TABLE,
+    @SerializedName("view")
+    VIEW,
+    @SerializedName("routine")
+    ROUTINE,
+    @SerializedName("procedure")
+    PROCEDURE,
+    @SerializedName("function")
+    FUNCTION,
+    @SerializedName("column")
+    COLUMN
+}

agents/common/src/main/java/com/dbx/agent/CompletionAssistantRequest.java

@@ -0,0 +1,34 @@
+package com.dbx.agent;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public final class CompletionAssistantRequest {
+    private String connection_id;
+    private String database;
+    private String schema;
+    private List<CompletionAssistantObjectKind> object_kinds = new ArrayList<>();
+    private String mask = "";
+    private boolean case_sensitive;
+    private boolean global_search;
+    private Integer max_results;
+    private boolean search_in_comments;
+    private boolean search_in_definitions;
+    private String parent_schema;
+    private String parent_name;
+    private CompletionAssistantMatchMode match_mode;
+
+    public String getConnection_id() { return connection_id; }
+    public String getDatabase() { return database == null ? "" : database; }
+    public String getSchema() { return schema; }
+    public List<CompletionAssistantObjectKind> getObject_kinds() { return object_kinds == null ? new ArrayList<>() : object_kinds; }
+    public String getMask() { return mask == null ? "" : mask; }
+    public boolean getCase_sensitive() { return case_sensitive; }
+    public boolean getGlobal_search() { return global_search; }
+    public Integer getMax_results() { return max_results; }
+    public boolean getSearch_in_comments() { return search_in_comments; }
+    public boolean getSearch_in_definitions() { return search_in_definitions; }
+    public String getParent_schema() { return parent_schema; }
+    public String getParent_name() { return parent_name; }
+    public CompletionAssistantMatchMode getMatch_mode() { return match_mode; }
+}

agents/common/src/main/java/com/dbx/agent/CompletionAssistantResponse.java

@@ -0,0 +1,19 @@
+package com.dbx.agent;
+
+import java.util.List;
+
+public final class CompletionAssistantResponse {
+    private final List<CompletionAssistantCandidate> candidates;
+    private final boolean incomplete;
+    private final boolean fallback_used;
+
+    public CompletionAssistantResponse(List<CompletionAssistantCandidate> candidates, boolean incomplete, boolean fallbackUsed) {
+        this.candidates = candidates;
+        this.incomplete = incomplete;
+        this.fallback_used = fallbackUsed;
+    }
+
+    public List<CompletionAssistantCandidate> getCandidates() { return candidates; }
+    public boolean getIncomplete() { return incomplete; }
+    public boolean getFallback_used() { return fallback_used; }
+}

agents/common/src/main/java/com/dbx/agent/ConfiguredJdbcAgent.java

@@ -60,6 +60,11 @@ public abstract class ConfiguredJdbcAgent extends AbstractJdbcAgent {
         return StandardJdbcMetadata.INSTANCE.listObjects(listTables(schema), schema);
     }
 
+    @Override
+    public CompletionAssistantResponse completionAssistantSearch(CompletionAssistantRequest request) {
+        return StandardJdbcMetadata.INSTANCE.completionAssistantSearch(requireConnection(), profile, configuredDatabase, request);
+    }
+
     @Override
     public ObjectSource getObjectSource(String schema, String name, String objectType) {
         throw new UnsupportedOperationException("Object source is not supported");

agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java

@@ -24,6 +24,10 @@ public interface DatabaseAgent {
         return result;
     }
 
+    default CompletionAssistantResponse completionAssistantSearch(CompletionAssistantRequest request) {
+        throw new UnsupportedOperationException("Completion assistant search is not supported by this agent");
+    }
+
     List<ColumnInfo> getColumns(String schema, String table);
 
     default ObjectSource getObjectSource(String schema, String name, String objectType) {

agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java

@@ -115,6 +115,10 @@ public final class JsonRpcServer {
             switchCatalog(params);
             return agent.listObjects(params.get("schema").getAsString());
         }
+        if (AgentProtocol.METHOD_COMPLETION_ASSISTANT_SEARCH_V1.equals(method)) {
+            switchCatalog(params);
+            return agent.completionAssistantSearch(gson.fromJson(params, CompletionAssistantRequest.class));
+        }
         if (AgentProtocol.METHOD_GET_OBJECT_SOURCE.equals(method)) {
             switchCatalog(params);
             return agent.getObjectSource(

agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java

@@ -90,6 +90,40 @@ public final class StandardJdbcMetadata {
         return result;
     }
 
+    public CompletionAssistantResponse completionAssistantSearch(
+        Connection conn,
+        JdbcAgentProfile profile,
+        String configuredDatabase,
+        CompletionAssistantRequest request
+    ) {
+        return unchecked(() -> {
+            int limit = boundedLimit(request.getMax_results());
+            List<CompletionAssistantCandidate> candidates = new ArrayList<>();
+            List<CompletionAssistantObjectKind> kinds = normalizedObjectKinds(request);
+            DatabaseMetaData meta = conn.getMetaData();
+            String catalog = blankToNull(request.getDatabase());
+            String schema = completionSchema(request);
+
+            if (containsKind(kinds, CompletionAssistantObjectKind.SCHEMA)) {
+                appendCompletionSchemas(candidates, meta, profile, request, limit);
+            }
+            if (candidates.size() < limit && containsTableLike(kinds)) {
+                appendCompletionTables(candidates, meta, profile, catalog, schema, request, kinds, limit);
+                if (candidates.isEmpty() && profile.getCatalogFallbackEnabled() && configuredDatabase != null && !configuredDatabase.trim().isEmpty()) {
+                    appendCompletionTables(candidates, meta, profile, configuredDatabase, schema, request, kinds, limit);
+                }
+            }
+            if (candidates.size() < limit && containsKind(kinds, CompletionAssistantObjectKind.COLUMN)) {
+                appendCompletionColumns(candidates, meta, catalog, schema, request, limit);
+                if (candidates.isEmpty() && profile.getCatalogFallbackEnabled() && configuredDatabase != null && !configuredDatabase.trim().isEmpty()) {
+                    appendCompletionColumns(candidates, meta, configuredDatabase, schema, request, limit);
+                }
+            }
+
+            return new CompletionAssistantResponse(candidates, candidates.size() >= limit, false);
+        });
+    }
+
     public List<ColumnInfo> getColumns(Connection conn, JdbcAgentProfile profile, String configuredDatabase, String schema, String table) {
         return unchecked(() -> {
             DatabaseMetaData meta = conn.getMetaData();
@@ -229,6 +263,188 @@ public final class StandardJdbcMetadata {
         }
     }
 
+    private void appendCompletionSchemas(
+        List<CompletionAssistantCandidate> result,
+        DatabaseMetaData meta,
+        JdbcAgentProfile profile,
+        CompletionAssistantRequest request,
+        int limit
+    ) throws Exception {
+        Set<String> names = new LinkedHashSet<>();
+        try {
+            appendSchemas(names, meta.getSchemas(null, null));
+        } catch (Exception | AbstractMethodError first) {
+            try {
+                appendSchemas(names, meta.getSchemas());
+            } catch (Exception | AbstractMethodError ignored) {
+            }
+        }
+        List<String> sorted = new ArrayList<>(names);
+        Collections.sort(sorted);
+        for (String name : sorted) {
+            if (result.size() >= limit) {
+                return;
+            }
+            if (profile.getExcludedSchemas().contains(name.toUpperCase(Locale.ROOT)) || !completionNameMatches(name, request)) {
+                continue;
+            }
+            result.add(new CompletionAssistantCandidate(
+                name,
+                CompletionAssistantCandidateKind.SCHEMA,
+                request.getDatabase(),
+                name,
+                null,
+                null,
+                null,
+                null
+            ));
+        }
+    }
+
+    private void appendCompletionTables(
+        List<CompletionAssistantCandidate> result,
+        DatabaseMetaData meta,
+        JdbcAgentProfile profile,
+        String catalog,
+        String schema,
+        CompletionAssistantRequest request,
+        List<CompletionAssistantObjectKind> kinds,
+        int limit
+    ) throws Exception {
+        String[] tableTypes = getDriverTableTypes(meta, profile);
+        try (ResultSet rs = meta.getTables(catalog, blankToNull(schema), completionPattern(request), tableTypes)) {
+            while (rs.next() && result.size() < limit) {
+                String name = rs.getString("TABLE_NAME");
+                String type = normalizeTableType(rs.getString("TABLE_TYPE"));
+                CompletionAssistantCandidateKind kind = completionTableKind(type);
+                if (!completionTableKindAllowed(kind, kinds) || !completionNameMatches(name, request)) {
+                    continue;
+                }
+                result.add(new CompletionAssistantCandidate(
+                    name,
+                    kind,
+                    request.getDatabase(),
+                    schema,
+                    null,
+                    null,
+                    rs.getString("REMARKS"),
+                    null
+                ));
+            }
+        }
+        result.sort(Comparator.comparing(CompletionAssistantCandidate::getName));
+    }
+
+    private void appendCompletionColumns(
+        List<CompletionAssistantCandidate> result,
+        DatabaseMetaData meta,
+        String catalog,
+        String schema,
+        CompletionAssistantRequest request,
+        int limit
+    ) throws Exception {
+        String table = request.getParent_name();
+        if (table == null || table.trim().isEmpty()) {
+            return;
+        }
+        try (ResultSet rs = meta.getColumns(catalog, blankToNull(schema), table, completionPattern(request))) {
+            while (rs.next() && result.size() < limit) {
+                String name = rs.getString("COLUMN_NAME");
+                if (!completionNameMatches(name, request)) {
+                    continue;
+                }
+                result.add(new CompletionAssistantCandidate(
+                    name,
+                    CompletionAssistantCandidateKind.COLUMN,
+                    request.getDatabase(),
+                    schema,
+                    schema,
+                    table,
+                    rs.getString("REMARKS"),
+                    rs.getString("TYPE_NAME")
+                ));
+            }
+        }
+    }
+
+    private static List<CompletionAssistantObjectKind> normalizedObjectKinds(CompletionAssistantRequest request) {
+        List<CompletionAssistantObjectKind> kinds = request.getObject_kinds();
+        if (kinds.isEmpty()) {
+            kinds = new ArrayList<>();
+            kinds.add(CompletionAssistantObjectKind.TABLE);
+            kinds.add(CompletionAssistantObjectKind.VIEW);
+        }
+        return kinds;
+    }
+
+    private static String completionSchema(CompletionAssistantRequest request) {
+        String parentSchema = request.getParent_schema();
+        if (parentSchema != null && !parentSchema.trim().isEmpty()) {
+            return parentSchema;
+        }
+        return request.getSchema();
+    }
+
+    private static String completionPattern(CompletionAssistantRequest request) {
+        String mask = request.getMask();
+        if (mask.trim().isEmpty()) {
+            return "%";
+        }
+        String escaped = mask.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_");
+        if (request.getMatch_mode() == CompletionAssistantMatchMode.CONTAINS) {
+            return "%" + escaped + "%";
+        }
+        return escaped + "%";
+    }
+
+    private static boolean completionNameMatches(String name, CompletionAssistantRequest request) {
+        if (name == null) {
+            return false;
+        }
+        String mask = request.getMask();
+        if (mask.trim().isEmpty()) {
+            return true;
+        }
+        String candidate = request.getCase_sensitive() ? name : name.toLowerCase(Locale.ROOT);
+        String expected = request.getCase_sensitive() ? mask : mask.toLowerCase(Locale.ROOT);
+        if (request.getMatch_mode() == CompletionAssistantMatchMode.CONTAINS) {
+            return candidate.contains(expected);
+        }
+        return candidate.startsWith(expected);
+    }
+
+    private static int boundedLimit(Integer requested) {
+        if (requested == null) {
+            return 100;
+        }
+        return Math.max(1, Math.min(1000, requested));
+    }
+
+    private static boolean containsKind(List<CompletionAssistantObjectKind> kinds, CompletionAssistantObjectKind kind) {
+        return kinds.contains(kind);
+    }
+
+    private static boolean containsTableLike(List<CompletionAssistantObjectKind> kinds) {
+        return kinds.contains(CompletionAssistantObjectKind.TABLE) || kinds.contains(CompletionAssistantObjectKind.VIEW);
+    }
+
+    private static CompletionAssistantCandidateKind completionTableKind(String type) {
+        if ("VIEW".equalsIgnoreCase(type)) {
+            return CompletionAssistantCandidateKind.VIEW;
+        }
+        return CompletionAssistantCandidateKind.TABLE;
+    }
+
+    private static boolean completionTableKindAllowed(
+        CompletionAssistantCandidateKind kind,
+        List<CompletionAssistantObjectKind> requestedKinds
+    ) {
+        if (kind == CompletionAssistantCandidateKind.VIEW) {
+            return requestedKinds.contains(CompletionAssistantObjectKind.VIEW);
+        }
+        return requestedKinds.contains(CompletionAssistantObjectKind.TABLE);
+    }
+
     private Set<String> primaryKeys(DatabaseMetaData meta, String catalog, String schema, String table) {
         Set<String> keys = new LinkedHashSet<>();
         try (ResultSet rs = meta.getPrimaryKeys(catalog, blankToNull(schema), table)) {

agents/common/src/main/resources/agent-protocol-v1.json

@@ -43,6 +43,7 @@
     "list_schemas",
     "list_tables",
     "list_objects",
+    "completion_assistant_search_v1",
     "get_object_source",
     "get_table_ddl",
     "get_columns",