fix(oceanbase-oracle): complete tables across schemas
This commit is contained in:
Mo-Moore authored and GitHub committed
1 parent
53184faf87
commit
a73bc4df69
10 files changed
+574
-22
No files matched your search
+332
@@ -2,6 +2,12 @@ package com.dbx.agent.oceanbaseoracle;
|
||||
|
||||
import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.AgentProtocol;
|
||||
import com.dbx.agent.CompletionAssistantCandidate;
|
||||
import com.dbx.agent.CompletionAssistantCandidateKind;
|
||||
import com.dbx.agent.CompletionAssistantMatchMode;
|
||||
import com.dbx.agent.CompletionAssistantObjectKind;
|
||||
import com.dbx.agent.CompletionAssistantRequest;
|
||||
import com.dbx.agent.CompletionAssistantResponse;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.DatabaseInfo;
|
||||
@@ -34,8 +40,10 @@ import java.util.ArrayDeque;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.Deque;
|
||||
import java.util.HashSet;
|
||||
import java.util.IdentityHashMap;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
@@ -388,6 +396,330 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent {
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public CompletionAssistantResponse completionAssistantSearch(CompletionAssistantRequest request) {
|
||||
if (hasTableLikeCompletionKind(request.getObject_kinds())) {
|
||||
return unchecked(() -> completionAssistantTables(request));
|
||||
}
|
||||
return super.completionAssistantSearch(request);
|
||||
}
|
||||
|
||||
private CompletionAssistantResponse completionAssistantTables(CompletionAssistantRequest request) throws SQLException {
|
||||
int limit = boundedCompletionLimit(request.getMax_results());
|
||||
int scanLimit = Math.min(1000, Math.max(limit * 3, limit + 1));
|
||||
String preferredSchema = preferredCompletionSchema(request);
|
||||
CompletionTablesQuery query = buildCompletionTablesQuery(request, preferredSchema, scanLimit + 1);
|
||||
List<CompletionTableRow> rows = new ArrayList<>();
|
||||
try (PreparedStatement stmt = requireConnection().prepareStatement(query.sql)) {
|
||||
bindCompletionArgs(stmt, query.args);
|
||||
try (ResultSet rs = stmt.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
rows.add(new CompletionTableRow(
|
||||
rs.getString(1),
|
||||
rs.getString(2),
|
||||
rs.getString(3),
|
||||
rs.getString(4),
|
||||
rs.getString(5)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Set<CompletionSynonymTarget> validTargets = validCompletionSynonymTargets(
|
||||
rows,
|
||||
completionTableObjectTypes(request.getObject_kinds())
|
||||
);
|
||||
List<CompletionAssistantCandidate> candidates = new ArrayList<>();
|
||||
for (CompletionTableRow row : rows) {
|
||||
if (row.name == null || row.name.isBlank()) {
|
||||
continue;
|
||||
}
|
||||
if ("SYNONYM".equalsIgnoreCase(row.objectType)
|
||||
&& (row.targetOwner == null || row.targetName == null
|
||||
|| !validTargets.contains(new CompletionSynonymTarget(row.targetOwner, row.targetName)))) {
|
||||
continue;
|
||||
}
|
||||
CompletionAssistantCandidateKind kind = "VIEW".equalsIgnoreCase(row.objectType)
|
||||
? CompletionAssistantCandidateKind.VIEW
|
||||
: CompletionAssistantCandidateKind.TABLE;
|
||||
candidates.add(new CompletionAssistantCandidate(
|
||||
row.name,
|
||||
kind,
|
||||
blankToNull(request.getDatabase()),
|
||||
row.owner,
|
||||
null,
|
||||
null,
|
||||
null,
|
||||
row.objectType
|
||||
));
|
||||
}
|
||||
boolean incomplete = candidates.size() > limit;
|
||||
if (incomplete) {
|
||||
candidates = new ArrayList<>(candidates.subList(0, limit));
|
||||
}
|
||||
return new CompletionAssistantResponse(candidates, incomplete, false);
|
||||
}
|
||||
|
||||
static CompletionTablesQuery buildCompletionTablesQuery(
|
||||
CompletionAssistantRequest request,
|
||||
String preferredSchema,
|
||||
int limit
|
||||
) {
|
||||
List<String> objectTypes = completionTableObjectTypes(request.getObject_kinds());
|
||||
boolean caseSensitive = request.getCase_sensitive();
|
||||
String pattern = completionLikePattern(request.getMask(), request.getMatch_mode());
|
||||
// OceanBase Oracle matches metadata the same way as listTables: fold the
|
||||
// bind value in Java and compare UPPER(column) LIKE ?. UPPER(?) on binds
|
||||
// is unreliable for fuzzy/lowercase masks in this driver.
|
||||
if (!caseSensitive) {
|
||||
pattern = pattern.toUpperCase(Locale.ROOT);
|
||||
}
|
||||
List<Object> args = new ArrayList<>();
|
||||
String namePredicate = completionNamePredicate("o.OBJECT_NAME", caseSensitive);
|
||||
String synonymNamePredicate = completionNamePredicate("s.SYNONYM_NAME", caseSensitive);
|
||||
args.add(pattern);
|
||||
|
||||
String ownerPredicate = "";
|
||||
String synonymOwnerPredicate = "";
|
||||
String owner = "";
|
||||
if (!request.getGlobal_search()) {
|
||||
owner = firstNonBlank(request.getParent_schema(), request.getSchema(), preferredSchema);
|
||||
if (owner == null) {
|
||||
owner = "";
|
||||
}
|
||||
owner = owner.toUpperCase(Locale.ROOT);
|
||||
args.add(owner);
|
||||
ownerPredicate = " AND UPPER(o.OWNER) = ?";
|
||||
}
|
||||
|
||||
args.add(pattern);
|
||||
if (!owner.isEmpty()) {
|
||||
args.add(owner);
|
||||
synonymOwnerPredicate = " AND UPPER(s.OWNER) = ?";
|
||||
}
|
||||
|
||||
String preferred = preferredSchema == null ? "" : preferredSchema.trim();
|
||||
String preferredFolded = caseSensitive ? preferred : preferred.toUpperCase(Locale.ROOT);
|
||||
String exactMask = request.getMask() == null ? "" : request.getMask().trim();
|
||||
String exactFolded = caseSensitive ? exactMask : exactMask.toUpperCase(Locale.ROOT);
|
||||
args.add(preferredFolded);
|
||||
args.add(exactFolded);
|
||||
args.add(limit);
|
||||
|
||||
String typeList = String.join(", ", objectTypes);
|
||||
String baseSql = """
|
||||
SELECT o.OWNER,
|
||||
o.OBJECT_NAME,
|
||||
o.OBJECT_TYPE,
|
||||
CAST(NULL AS VARCHAR2(128)) AS TARGET_OWNER,
|
||||
CAST(NULL AS VARCHAR2(128)) AS TARGET_NAME
|
||||
FROM ALL_OBJECTS o
|
||||
WHERE o.OBJECT_TYPE IN (%s)
|
||||
AND %s%s
|
||||
UNION ALL
|
||||
SELECT s.OWNER,
|
||||
s.SYNONYM_NAME AS OBJECT_NAME,
|
||||
'SYNONYM' AS OBJECT_TYPE,
|
||||
s.TABLE_OWNER AS TARGET_OWNER,
|
||||
s.TABLE_NAME AS TARGET_NAME
|
||||
FROM ALL_SYNONYMS s
|
||||
WHERE s.DB_LINK IS NULL
|
||||
AND %s%s
|
||||
""".formatted(typeList, namePredicate, ownerPredicate, synonymNamePredicate, synonymOwnerPredicate)
|
||||
.stripIndent()
|
||||
.trim();
|
||||
|
||||
String preferredOwnerPredicate = caseSensitive ? "OWNER = ?" : "UPPER(OWNER) = ?";
|
||||
String exactNamePredicate = caseSensitive
|
||||
? "OBJECT_NAME = ?"
|
||||
: "UPPER(OBJECT_NAME) = ?";
|
||||
String orderedSql = """
|
||||
SELECT OWNER, OBJECT_NAME, OBJECT_TYPE, TARGET_OWNER, TARGET_NAME
|
||||
FROM (
|
||||
%s
|
||||
)
|
||||
ORDER BY CASE
|
||||
WHEN %s THEN 0
|
||||
WHEN OWNER = 'PUBLIC' THEN 1
|
||||
WHEN OWNER IN ('SYS','SYSTEM','SYSMAN','DBSNMP','OUTLN','XDB','MDSYS','CTXSYS','WMSYS') THEN 3
|
||||
ELSE 2
|
||||
END,
|
||||
CASE WHEN %s THEN 0 ELSE 1 END,
|
||||
CASE OBJECT_TYPE WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'SYNONYM' THEN 3 ELSE 4 END,
|
||||
OBJECT_NAME,
|
||||
OWNER
|
||||
""".formatted(baseSql, preferredOwnerPredicate, exactNamePredicate).stripIndent().trim();
|
||||
|
||||
String sql = """
|
||||
SELECT OWNER, OBJECT_NAME, OBJECT_TYPE, TARGET_OWNER, TARGET_NAME
|
||||
FROM (
|
||||
%s
|
||||
)
|
||||
WHERE ROWNUM <= ?
|
||||
""".formatted(orderedSql).stripIndent().trim();
|
||||
|
||||
return new CompletionTablesQuery(sql, args);
|
||||
}
|
||||
|
||||
private static List<String> completionTableObjectTypes(List<CompletionAssistantObjectKind> kinds) {
|
||||
List<String> objectTypes = new ArrayList<>();
|
||||
boolean any = kinds == null || kinds.isEmpty();
|
||||
if (any || kinds.contains(CompletionAssistantObjectKind.TABLE)) {
|
||||
objectTypes.add("'TABLE'");
|
||||
}
|
||||
if (any || kinds.contains(CompletionAssistantObjectKind.VIEW)) {
|
||||
objectTypes.add("'VIEW'");
|
||||
}
|
||||
if (objectTypes.isEmpty()) {
|
||||
objectTypes.add("'TABLE'");
|
||||
objectTypes.add("'VIEW'");
|
||||
}
|
||||
return objectTypes;
|
||||
}
|
||||
|
||||
private static boolean hasTableLikeCompletionKind(List<CompletionAssistantObjectKind> kinds) {
|
||||
if (kinds == null || kinds.isEmpty()) {
|
||||
return true;
|
||||
}
|
||||
for (CompletionAssistantObjectKind kind : kinds) {
|
||||
if (kind == CompletionAssistantObjectKind.TABLE || kind == CompletionAssistantObjectKind.VIEW) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private String preferredCompletionSchema(CompletionAssistantRequest request) throws SQLException {
|
||||
String preferred = firstNonBlank(request.getSchema());
|
||||
if (preferred != null && !preferred.isBlank()) {
|
||||
return preferred;
|
||||
}
|
||||
String current = currentSchema();
|
||||
return current == null ? "" : current;
|
||||
}
|
||||
|
||||
private static String completionLikePattern(String mask, CompletionAssistantMatchMode matchMode) {
|
||||
String escaped = escapeLikePattern(mask == null ? "" : mask.trim());
|
||||
if (matchMode == CompletionAssistantMatchMode.CONTAINS) {
|
||||
return "%" + escaped + "%";
|
||||
}
|
||||
return escaped + "%";
|
||||
}
|
||||
|
||||
private static String escapeLikePattern(String mask) {
|
||||
return mask.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_");
|
||||
}
|
||||
|
||||
private static String completionNamePredicate(String column, boolean caseSensitive) {
|
||||
if (caseSensitive) {
|
||||
return column + " LIKE ? ESCAPE '\\'";
|
||||
}
|
||||
return "UPPER(" + column + ") LIKE ? ESCAPE '\\'";
|
||||
}
|
||||
|
||||
private static int boundedCompletionLimit(Integer requested) {
|
||||
if (requested == null || requested <= 0) {
|
||||
return 100;
|
||||
}
|
||||
return Math.min(requested, 1000);
|
||||
}
|
||||
|
||||
private static String firstNonBlank(String... values) {
|
||||
if (values == null) {
|
||||
return null;
|
||||
}
|
||||
for (String value : values) {
|
||||
if (value != null && !value.isBlank()) {
|
||||
return value.trim();
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private static String blankToNull(String value) {
|
||||
return value == null || value.isBlank() ? null : value;
|
||||
}
|
||||
|
||||
private static void bindCompletionArgs(PreparedStatement stmt, List<Object> args) throws SQLException {
|
||||
for (int i = 0; i < args.size(); i++) {
|
||||
Object arg = args.get(i);
|
||||
if (arg instanceof Integer integer) {
|
||||
stmt.setInt(i + 1, integer);
|
||||
} else {
|
||||
stmt.setString(i + 1, arg == null ? null : String.valueOf(arg));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private Set<CompletionSynonymTarget> validCompletionSynonymTargets(
|
||||
List<CompletionTableRow> rows,
|
||||
List<String> objectTypes
|
||||
) throws SQLException {
|
||||
Set<CompletionSynonymTarget> targets = new LinkedHashSet<>();
|
||||
for (CompletionTableRow row : rows) {
|
||||
if ("SYNONYM".equalsIgnoreCase(row.objectType) && row.targetOwner != null && row.targetName != null) {
|
||||
targets.add(new CompletionSynonymTarget(row.targetOwner, row.targetName));
|
||||
}
|
||||
}
|
||||
Set<CompletionSynonymTarget> valid = new HashSet<>();
|
||||
if (targets.isEmpty()) {
|
||||
return valid;
|
||||
}
|
||||
String typeList = String.join(", ", objectTypes);
|
||||
List<CompletionSynonymTarget> ordered = new ArrayList<>(targets);
|
||||
for (int start = 0; start < ordered.size(); start += 100) {
|
||||
List<CompletionSynonymTarget> batch = ordered.subList(start, Math.min(start + 100, ordered.size()));
|
||||
List<Object> args = new ArrayList<>();
|
||||
List<String> predicates = new ArrayList<>();
|
||||
for (CompletionSynonymTarget target : batch) {
|
||||
args.add(target.owner());
|
||||
args.add(target.name());
|
||||
predicates.add("(o.OWNER = ? AND o.OBJECT_NAME = ?)");
|
||||
}
|
||||
String sql = "SELECT DISTINCT o.OWNER, o.OBJECT_NAME"
|
||||
+ " FROM ALL_OBJECTS o"
|
||||
+ " WHERE o.OBJECT_TYPE IN (" + typeList + ")"
|
||||
+ " AND (" + String.join(" OR ", predicates) + ")";
|
||||
try (PreparedStatement stmt = requireConnection().prepareStatement(sql)) {
|
||||
bindCompletionArgs(stmt, args);
|
||||
try (ResultSet rs = stmt.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
valid.add(new CompletionSynonymTarget(rs.getString(1), rs.getString(2)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return valid;
|
||||
}
|
||||
|
||||
record CompletionSynonymTarget(String owner, String name) {
|
||||
}
|
||||
|
||||
static final class CompletionTableRow {
|
||||
final String owner;
|
||||
final String name;
|
||||
final String objectType;
|
||||
final String targetOwner;
|
||||
final String targetName;
|
||||
|
||||
CompletionTableRow(String owner, String name, String objectType, String targetOwner, String targetName) {
|
||||
this.owner = owner;
|
||||
this.name = name;
|
||||
this.objectType = objectType;
|
||||
this.targetOwner = targetOwner;
|
||||
this.targetName = targetName;
|
||||
}
|
||||
}
|
||||
|
||||
static final class CompletionTablesQuery {
|
||||
final String sql;
|
||||
final List<Object> args;
|
||||
|
||||
CompletionTablesQuery(String sql, List<Object> args) {
|
||||
this.sql = sql;
|
||||
this.args = List.copyOf(args);
|
||||
}
|
||||
}
|
||||
|
||||
private static MetadataSql oceanBaseMetadataSql(
|
||||
String baseSql,
|
||||
String selectList,
|
||||
|
||||
+89
@@ -2,6 +2,10 @@ package com.dbx.agent.oceanbaseoracle;
|
||||
|
||||
import com.dbx.agent.AgentProtocol;
|
||||
import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.CompletionAssistantMatchMode;
|
||||
import com.dbx.agent.CompletionAssistantObjectKind;
|
||||
import com.dbx.agent.CompletionAssistantRequest;
|
||||
import com.dbx.agent.CompletionAssistantResponse;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
@@ -460,6 +464,64 @@ class OceanBaseOracleAgentTest {
|
||||
Assertions.assertTrue(sql.get(0).contains("ROWNUM <= ?"), sql.get(0));
|
||||
}
|
||||
|
||||
@Test
|
||||
void globalCompletionTableSearchOmitsOwnerFilter() {
|
||||
CompletionAssistantRequest request = completionRequest("DWD", null, "STAG", true);
|
||||
OceanBaseOracleAgent.CompletionTablesQuery query = OceanBaseOracleAgent.buildCompletionTablesQuery(request, "DWD", 21);
|
||||
|
||||
Assertions.assertTrue(query.sql.contains("FROM ALL_OBJECTS"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("FROM ALL_SYNONYMS"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(o.OBJECT_NAME) LIKE ?"), query.sql);
|
||||
Assertions.assertFalse(query.sql.contains("UPPER(o.OWNER) = ?"), query.sql);
|
||||
Assertions.assertFalse(query.sql.contains("UPPER(s.OWNER) = ?"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("ROWNUM <= ?"), query.sql);
|
||||
Assertions.assertEquals(List.of("STAG%", "STAG%", "DWD", "STAG", 21), query.args);
|
||||
}
|
||||
|
||||
@Test
|
||||
void completionTableSearchFoldsLowercaseMasksForFuzzyMatch() {
|
||||
CompletionAssistantRequest request = completionRequest("dwd", null, "ord", true);
|
||||
setField(request, "match_mode", CompletionAssistantMatchMode.CONTAINS);
|
||||
OceanBaseOracleAgent.CompletionTablesQuery query = OceanBaseOracleAgent.buildCompletionTablesQuery(request, "dwd", 21);
|
||||
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(o.OBJECT_NAME) LIKE ?"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(OWNER) = ?"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(OBJECT_NAME) = ?"), query.sql);
|
||||
Assertions.assertFalse(query.sql.contains("LIKE UPPER(?)"), query.sql);
|
||||
Assertions.assertEquals(List.of("%ORD%", "%ORD%", "DWD", "ORD", 21), query.args);
|
||||
}
|
||||
|
||||
@Test
|
||||
void scopedCompletionTableSearchFiltersOwnerCaseInsensitively() {
|
||||
CompletionAssistantRequest request = completionRequest("dwd", "staging", "ord", false);
|
||||
OceanBaseOracleAgent.CompletionTablesQuery query = OceanBaseOracleAgent.buildCompletionTablesQuery(request, "dwd", 21);
|
||||
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(o.OWNER) = ?"), query.sql);
|
||||
Assertions.assertTrue(query.sql.contains("UPPER(s.OWNER) = ?"), query.sql);
|
||||
Assertions.assertEquals(List.of("ORD%", "STAGING", "ORD%", "STAGING", "DWD", "ORD", 21), query.args);
|
||||
}
|
||||
|
||||
@Test
|
||||
void completionAssistantSearchReturnsGlobalTableCandidates() {
|
||||
List<String> sql = new ArrayList<>();
|
||||
OceanBaseOracleAgent agent = new OceanBaseOracleAgent();
|
||||
TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet(
|
||||
new String[]{"OWNER", "OBJECT_NAME", "OBJECT_TYPE", "TARGET_OWNER", "TARGET_NAME"},
|
||||
new Object[][]{
|
||||
{"DWD", "ORDERS", "TABLE", null, null},
|
||||
{"STAGING", "ORDERS", "TABLE", null, null}
|
||||
}
|
||||
)));
|
||||
|
||||
CompletionAssistantResponse response = agent.completionAssistantSearch(completionRequest("DWD", null, "ORD", true));
|
||||
|
||||
Assertions.assertEquals(2, response.getCandidates().size());
|
||||
Assertions.assertEquals("ORDERS", response.getCandidates().get(0).getName());
|
||||
Assertions.assertEquals("DWD", response.getCandidates().get(0).getSchema());
|
||||
Assertions.assertEquals("STAGING", response.getCandidates().get(1).getSchema());
|
||||
Assertions.assertFalse(sql.get(0).contains("UPPER(o.OWNER) = ?"), sql.get(0));
|
||||
}
|
||||
|
||||
@Test
|
||||
void readsViewDdlWithDbmsMetadataForSchemaCompare() {
|
||||
List<String> sql = new ArrayList<>();
|
||||
@@ -1467,6 +1529,33 @@ class OceanBaseOracleAgentTest {
|
||||
return null;
|
||||
}
|
||||
|
||||
private static CompletionAssistantRequest completionRequest(
|
||||
String schema,
|
||||
String parentSchema,
|
||||
String mask,
|
||||
boolean globalSearch
|
||||
) {
|
||||
CompletionAssistantRequest request = new CompletionAssistantRequest();
|
||||
setField(request, "database", "OBORCL");
|
||||
setField(request, "schema", schema);
|
||||
setField(request, "parent_schema", parentSchema);
|
||||
setField(request, "mask", mask);
|
||||
setField(request, "global_search", globalSearch);
|
||||
setField(request, "max_results", 20);
|
||||
setField(request, "object_kinds", List.of(CompletionAssistantObjectKind.TABLE, CompletionAssistantObjectKind.VIEW));
|
||||
return request;
|
||||
}
|
||||
|
||||
private static void setField(Object target, String name, Object value) {
|
||||
try {
|
||||
var field = target.getClass().getDeclaredField(name);
|
||||
field.setAccessible(true);
|
||||
field.set(target, value);
|
||||
} catch (ReflectiveOperationException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private static <T> T proxy(Class<T> type, MethodHandler handler) {
|
||||
InvocationHandler invocationHandler = new InvocationHandler() {
|
||||
@Override
|
||||
|
||||
Reference in new issue
Block a user