diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index e93deb54cd..6483ebd2fb 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -88,7 +88,15 @@ import { driverProfileHasCompletionCandidates } from "@/lib/database/driverProfi import { sqlCompletionContextFromSemantic, sqlSemanticSelectStarIsOnlyProjection, sqlSemanticSelectStarQualifierSql, sqlSemanticSelectStarTableSources } from "@/lib/sql/semantic/completion"; import { buildSqlSemanticModel } from "@/lib/sql/semantic/model"; import { mergeSqlSemanticReferenceAnalysis, resolveSqlSemanticNavigationTarget } from "@/lib/sql/semantic/references"; -import { buildElasticsearchCompletionItemsFromContext, getElasticsearchCompletionContext, getElasticsearchCompletionResultValidFor, shouldAutoOpenElasticsearchCompletion, type ElasticsearchCompletionItem } from "@/lib/elasticsearch/elasticsearchCompletion"; +import { + buildElasticsearchCompletionItemsFromContext, + elasticsearchCompletionNeedsFields, + getElasticsearchCompletionContext, + getElasticsearchCompletionResultValidFor, + shouldAutoOpenElasticsearchCompletion, + type ElasticsearchCompletionField, + type ElasticsearchCompletionItem, +} from "@/lib/elasticsearch/elasticsearchCompletion"; import { buildMongoCompletionItemsFromContext, getMongoCompletionContext, getMongoCompletionResultValidFor, mongoCompletionNeedsCollections, mongoCompletionNeedsFields, shouldAutoOpenMongoCompletion, type MongoCompletionItem } from "@/lib/mongo/mongoCompletion"; import { buildSqlServerUseDatabaseCompletionItems, @@ -4526,9 +4534,11 @@ function shouldApplyCompletionAsSnippet(item: QueryCompletionItem): boolean { function completionOptionForItem(item: QueryCompletionItem | BatchColumnSelectionActionItem) { const filterText = "filterText" in item && typeof item.filterText === "string" ? item.filterText : undefined; const labelPresentation = completionLabelPresentation(item.label, filterText); + const sortText = "sortText" in item && typeof item.sortText === "string" ? item.sortText : labelPresentation.sortText; if (isBatchColumnSelectionAction(item)) { return { ...labelPresentation, + ...(sortText ? { sortText } : {}), dbxBatchColumnSelectionAction: { sessionKey: item.sessionKey }, type: item.type, detail: item.detail, @@ -4545,6 +4555,7 @@ function completionOptionForItem(item: QueryCompletionItem | BatchColumnSelectio if (shouldApplyCompletionAsSnippet(item) && item.apply) { const completion = codeMirrorSnippetCompletion(item.apply, { ...labelPresentation, + ...(sortText ? { sortText } : {}), type: item.type, detail: item.detail, info: item.info, @@ -4578,6 +4589,7 @@ function completionOptionForItem(item: QueryCompletionItem | BatchColumnSelectio } return cacheBatchColumnSelectionOption(batchColumnSelection, { ...labelPresentation, + ...(sortText ? { sortText } : {}), ...(batchColumnSelection ? { dbxBatchColumnSelection: batchColumnSelection } : {}), type: item.type, detail: item.detail, @@ -4612,6 +4624,7 @@ async function provideElasticsearchCompletions(currentState: import("@codemirror const completionContext = getElasticsearchCompletionContext(fullDoc, position); let indices: string[] = []; + let fields: ElasticsearchCompletionField[] = []; if (props.database != null && completionContext.mode === "path") { try { indices = await connectionStore.listElasticsearchCompletionIndices(props.connectionId, props.database); @@ -4619,9 +4632,16 @@ async function provideElasticsearchCompletions(currentState: import("@codemirror indices = []; } } + if (elasticsearchCompletionNeedsFields(completionContext) && completionContext.index) { + try { + fields = await connectionStore.listElasticsearchCompletionFields(props.connectionId, completionContext.index); + } catch { + fields = []; + } + } if (epoch !== completionEpoch) return null; - const items = buildElasticsearchCompletionItemsFromContext(completionContext, { indices }); + const items = buildElasticsearchCompletionItemsFromContext(completionContext, { indices, fields }); return buildCompletionResult(items, completionContext.from, getElasticsearchCompletionResultValidFor()); } diff --git a/apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts b/apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts index 3ecfec07df..c37569e80e 100644 --- a/apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts +++ b/apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts @@ -4,11 +4,12 @@ export type ElasticsearchJsonSlot = "key" | "value"; export interface ElasticsearchCompletionItem { label: string; - type: "keyword" | "property" | "text" | "snippet"; + type: "keyword" | "property" | "text" | "snippet" | "column"; detail?: string; info?: string; apply?: string; filterText?: string; + sortText?: string; // Set when `apply` carries snippet fields so the editor expands it instead of // inserting the raw `${}` markers. applyAsSnippet?: boolean; @@ -17,6 +18,11 @@ export interface ElasticsearchCompletionItem { boost: number; } +export interface ElasticsearchCompletionField { + name: string; + dataType?: string; +} + export interface ElasticsearchCompletionContext { mode: ElasticsearchCompletionMode; prefix: string; @@ -26,12 +32,15 @@ export interface ElasticsearchCompletionContext { segmentIndex?: number; replaceClosingQuote?: '"'; jsonSlot?: ElasticsearchJsonSlot; + jsonFieldSlot?: boolean; hasKeySeparator?: boolean; jsonIndent?: string; + index?: string; } export interface ElasticsearchCompletionInput { indices?: string[]; + fields?: ElasticsearchCompletionField[]; } const HTTP_METHODS = ["GET", "POST", "PUT", "DELETE", "HEAD"] as const; @@ -76,6 +85,20 @@ const JSON_KEYWORDS = ["query", "bool", "must", "should", "must_not", "filter", const JSON_ARRAY_KEYS = new Set(["must", "should", "must_not", "filter", "sort", "_source", "fields"]); const JSON_SCALAR_KEYS = new Set(["size", "from", "track_total_hits"]); +// Query-clause objects whose keys are mapping field names (`"term": { "status": ... }`). +const FIELD_NAME_OBJECT_KEYS = new Set(["term", "terms", "match", "match_phrase", "match_phrase_prefix", "match_bool_prefix", "range", "wildcard", "prefix", "fuzzy", "regexp", "span_term"]); + +// `terms`/`range` are also aggregation types whose keys are `field`/`ranges`, not mapping names. +const AGGREGATION_TYPE_FIELD_OBJECT_KEYS = new Set(["terms", "range"]); +const AGGREGATION_KEYS = new Set(["aggs", "aggregations"]); +const QUERY_CLAUSE_KEYS = new Set(["query", "bool", "must", "should", "filter", "must_not", "nested", "constant_score"]); + +// Arrays whose values (or object-element keys) are mapping field names. +const FIELD_NAME_ARRAY_KEYS = new Set(["sort", "_source", "fields"]); + +// Object keys whose JSON value is a mapping field name (`"exists": { "field": "status" }`). +const FIELD_VALUE_KEYS = new Set(["field", "path"]); + // Value scaffold inserted after a completed key so the matching brackets come // along with it, with the cursor parked where the value goes. function jsonKeyValueTemplate(key: string): string { @@ -170,14 +193,17 @@ export function getElasticsearchCompletionContext(text: string, cursor: number): if (lineStart > 0 || safeCursor > firstLineLimit || looksLikeJsonBody(text, safeCursor)) { const jsonPrefix = readJsonPrefix(text, safeCursor); const replaceClosingQuote = jsonPrefix.prefix.startsWith('"') && text[safeCursor] === '"' ? ('"' as const) : undefined; + const inspection = inspectJsonBody(text, jsonPrefix.from); return { mode: "json", prefix: jsonPrefix.prefix, from: jsonPrefix.from, replaceClosingQuote, - jsonSlot: readJsonSlot(text, jsonPrefix.from), + jsonSlot: inspection.slot, + jsonFieldSlot: isJsonFieldNameSlot(inspection), hasKeySeparator: hasJsonKeySeparator(text, safeCursor), jsonIndent: beforeCursorOnLine.match(/^\s*/)?.[0] ?? "", + index: elasticsearchRequestIndex(text, safeCursor), }; } @@ -192,20 +218,81 @@ export function buildElasticsearchCompletionItems(text: string, cursor: number, export function buildElasticsearchCompletionItemsFromContext(context: ElasticsearchCompletionContext, input: ElasticsearchCompletionInput = {}): ElasticsearchCompletionItem[] { if (context.mode === "method") return methodItems(context.prefix); - if (context.mode === "json") return jsonItems(context); + if (context.mode === "json") return jsonItems(context, input.fields ?? []); return pathItems(context, input.indices ?? []); } +export function elasticsearchCompletionNeedsFields(context: ElasticsearchCompletionContext): boolean { + return context.mode === "json" && context.jsonFieldSlot === true; +} + +const MAX_MAPPING_FIELDS = 500; + +/** Flatten `GET /{index}/_mapping` into completion fields. Iterative on purpose: a recursive walk is the wrong model for this JSON. */ +export function flattenElasticsearchMappingFields(mapping: unknown): ElasticsearchCompletionField[] { + const properties = elasticsearchMappingProperties(mapping); + if (!properties) return []; + + const fields: ElasticsearchCompletionField[] = []; + const seenNames = new Set(); + const seenObjects = new Set>(); + const stack: Array<{ prefix: string; properties: Record }> = [{ prefix: "", properties }]; + + while (stack.length > 0 && fields.length < MAX_MAPPING_FIELDS) { + const current = stack.pop(); + if (!current || seenObjects.has(current.properties)) continue; + seenObjects.add(current.properties); + for (const [name, definition] of Object.entries(current.properties)) { + if (fields.length >= MAX_MAPPING_FIELDS) break; + if (!isJsonObject(definition)) continue; + const fieldName = current.prefix ? `${current.prefix}.${name}` : name; + const fieldType = typeof definition.type === "string" ? definition.type : isJsonObject(definition.properties) ? "object" : undefined; + if (fieldType && !seenNames.has(fieldName)) { + seenNames.add(fieldName); + fields.push({ name: fieldName, dataType: fieldType }); + } + // Only mapping subtrees: `properties` (object/nested) and `fields` (multi-fields). + // Never walk `index_options`, `path`, analyzer settings, etc. + if (isJsonObject(definition.properties)) stack.push({ prefix: fieldName, properties: definition.properties }); + if (isJsonObject(definition.fields)) stack.push({ prefix: fieldName, properties: definition.fields }); + } + } + return fields; +} + +function isJsonObject(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function elasticsearchMappingProperties(mapping: unknown): Record | undefined { + if (!isJsonObject(mapping)) return undefined; + const collected: Record = {}; + const consider = (node: unknown) => { + if (!isJsonObject(node)) return; + const mappings = isJsonObject(node.mappings) ? node.mappings : node; + if (isJsonObject(mappings.properties)) { + Object.assign(collected, mappings.properties); + return; + } + for (const typed of Object.values(mappings)) { + if (isJsonObject(typed) && isJsonObject(typed.properties)) Object.assign(collected, typed.properties); + } + }; + if (isJsonObject(mapping.mappings)) consider(mapping); + else for (const value of Object.values(mapping)) consider(value); + return Object.keys(collected).length > 0 ? collected : undefined; +} + export function shouldAutoOpenElasticsearchCompletion(text: string, cursor: number): boolean { const previousChar = text[cursor - 1]; if (!previousChar) return false; if (/[{,}\]\n\r]/.test(previousChar)) return false; - if (/[\w/_."]/.test(previousChar)) return true; + if (/[\w/_.".-]/.test(previousChar)) return true; return false; } export function getElasticsearchCompletionResultValidFor(): RegExp { - return /[\w/_."]*$/; + return /[\w/_."-]*$/; } function methodItems(prefix: string): ElasticsearchCompletionItem[] { @@ -274,7 +361,8 @@ function indexEndpointItems(prefix: string): ElasticsearchCompletionItem[] { })); } -function jsonItems(context: ElasticsearchCompletionContext): ElasticsearchCompletionItem[] { +function jsonItems(context: ElasticsearchCompletionContext, fields: ElasticsearchCompletionField[]): ElasticsearchCompletionItem[] { + if (context.jsonFieldSlot) return fieldItems(context, fields); const quoted = context.prefix.startsWith('"'); const normalizedPrefix = context.prefix.replace(/^"/, ""); // Only an object key without its own `:` yet may pull the value scaffold in; @@ -306,6 +394,32 @@ function jsonItems(context: ElasticsearchCompletionContext): ElasticsearchComple return dedupeAndSort([...snippetItems, ...keyItems]); } +function fieldItems(context: ElasticsearchCompletionContext, fields: ElasticsearchCompletionField[]): ElasticsearchCompletionItem[] { + const quoted = context.prefix.startsWith('"'); + const normalizedPrefix = context.prefix.replace(/^"/, ""); + return fields + .filter((field) => matchesElasticsearchFieldPrefix(field.name, normalizedPrefix)) + .map((field) => { + const depth = field.name.split(".").length; + return { + label: `"${field.name}"`, + type: "column" as const, + detail: field.dataType || "field", + apply: `"${field.name}"`, + // Quoted labels cannot prefix-match an unquoted typed token (`ci` vs `"city"`), + // so expose the bare field name as filterText for CodeMirror's matcher. + filterText: quoted ? `"${field.name}"` : field.name, + replaceClosingQuote: context.replaceClosingQuote, + // Shallower paths first. Default editor sort uses localeCompare on labels, + // which puts `"nested_comments"` *after* its children; sortText avoids that. + sortText: `${String(depth).padStart(2, "0")}\t${field.name}`, + boost: 200 - depth * 20, + }; + }) + .sort((left, right) => right.boost - left.boost || (left.sortText ?? left.label).localeCompare(right.sortText ?? right.label)) + .slice(0, 100); +} + function indentJsonSnippet(snippet: string, baseIndent: string): string { if (!baseIndent) return snippet; return snippet.replace(/\n/g, `\n${baseIndent}`); @@ -328,25 +442,169 @@ function hasOpenJsonContainer(text: string, cursor: number): boolean { return stack.length > 0; } -// Tracks open brackets so an array element or a value after `:` is never -// mistaken for an object key. -function readJsonSlot(text: string, from: number): ElasticsearchJsonSlot { - const stack: string[] = []; +interface JsonFrame { + kind: "{" | "["; + key: string | null; +} + +interface JsonInspection { + slot: ElasticsearchJsonSlot; + parentKey: string | null; + parentKind: "{" | "[" | null; + nearestNamedKey: string | null; + currentKey: string | null; + ancestorKeys: string[]; +} + +// Walks the current REST request body while keeping object keys, so a field +// name inside `"term": { | }` is not treated the same as a DSL key inside `"bool"`. +function inspectJsonBody(text: string, from: number): JsonInspection { + const before = text.slice(0, from); + const requestLine = /[\s\S]*(?:^|\n)[ \t]*(?:GET|POST|PUT|DELETE|HEAD)[ \t][^\n]*/i.exec(before); + const bodyStart = requestLine?.[0].length ?? 0; + const stack: JsonFrame[] = []; + let pendingKey: string | null = null; let last = ""; - for (const char of readRequestBody(text, from)) { - if (/\s/.test(char)) continue; - if (char === "{" || char === "[") stack.push(char); - else if (char === "}" || char === "]") stack.pop(); + let index = bodyStart; + + while (index < from) { + const char = before[index] ?? ""; + if (char === "#" || (char === "/" && before[index + 1] === "/")) { + const newline = before.indexOf("\n", index); + index = newline < 0 ? from : newline; + continue; + } + + if (char === '"') { + const string = readJsonString(before, index, from); + if (string.end < 0) break; + index = string.end; + let lookahead = index; + while (lookahead < from && (before[lookahead] === " " || before[lookahead] === "\t")) lookahead++; + if (before[lookahead] === ":") { + pendingKey = string.value; + continue; + } + last = '"'; + continue; + } + + if (/\s/.test(char)) { + index++; + continue; + } + + if (char === "{" || char === "[") { + stack.push({ kind: char, key: pendingKey }); + pendingKey = null; + last = char; + index++; + continue; + } + + if (char === "}" || char === "]") { + stack.pop(); + pendingKey = null; + last = char; + index++; + continue; + } + + if (char === ",") { + pendingKey = null; + last = char; + index++; + continue; + } + last = char; + index++; + } + + const top = stack[stack.length - 1]; + const slot: ElasticsearchJsonSlot = top?.kind === "{" && (last === "{" || last === ",") ? "key" : "value"; + let nearestNamedKey: string | null = null; + for (let depth = stack.length - 1; depth >= 0; depth--) { + const key = stack[depth]?.key; + if (key) { + nearestNamedKey = key; + break; + } + } + + return { + slot, + parentKey: top?.key ?? null, + parentKind: top?.kind ?? null, + nearestNamedKey, + currentKey: slot === "value" ? pendingKey : null, + ancestorKeys: stack.map((frame) => frame.key).filter((key): key is string => !!key), + }; +} + +function readJsonString(text: string, start: number, limit: number): { end: number; value: string } { + let index = start + 1; + let escaped = false; + let value = ""; + while (index < limit) { + const char = text[index] ?? ""; + if (escaped) { + value += char; + escaped = false; + index++; + continue; + } + if (char === "\\") { + escaped = true; + index++; + continue; + } + if (char === '"') return { end: index + 1, value }; + value += char; + index++; + } + return { end: -1, value }; +} + +function isJsonFieldNameSlot(inspection: JsonInspection): boolean { + if (inspection.slot === "key") { + if (inspection.parentKey && FIELD_NAME_OBJECT_KEYS.has(inspection.parentKey)) { + return !isAggregationTypeClause(inspection); + } + // Object-form sort (`"sort": { "createdAt": "desc" }`) uses field names as keys. + // Array-form sort objects (`"sort": [{ "createdAt": ... }]`) have no key of their + // own; `_source`/`fields` objects do not (`includes`/`field`/`format`). + if (inspection.parentKey === "sort") return true; + return inspection.parentKind === "{" && !inspection.parentKey && inspection.nearestNamedKey === "sort"; + } + if (inspection.currentKey && FIELD_VALUE_KEYS.has(inspection.currentKey)) return true; + return inspection.parentKind === "[" && !!inspection.parentKey && FIELD_NAME_ARRAY_KEYS.has(inspection.parentKey); +} + +function isAggregationTypeClause(inspection: JsonInspection): boolean { + if (!inspection.parentKey || !AGGREGATION_TYPE_FIELD_OBJECT_KEYS.has(inspection.parentKey)) return false; + if (!inspection.ancestorKeys.some((key) => AGGREGATION_KEYS.has(key))) return false; + return !inspection.ancestorKeys.some((key) => QUERY_CLAUSE_KEYS.has(key)); +} + +function elasticsearchRequestIndex(text: string, cursor: number): string | undefined { + const match = /[\s\S]*(?:^|\n)[ \t]*(?:GET|POST|PUT|DELETE|HEAD)[ \t]+(\S+)/i.exec(text.slice(0, cursor)); + const path = match?.[1]; + if (!path) return undefined; + const first = (path.split("?")[0] ?? "").split("/").filter(Boolean)[0]; + if (!first || first.startsWith("_")) return undefined; + try { + return decodeURIComponent(first); + } catch { + return first; } - return stack[stack.length - 1] === "{" && (last === "{" || last === ",") ? "key" : "value"; } function hasJsonKeySeparator(text: string, cursor: number): boolean { let index = cursor; // The caret may sit inside the key, so step over the rest of the token and // its closing quote before looking for the separator. - while (index < text.length && /[\w_]/.test(text[index] ?? "")) index++; + while (index < text.length && /[\w_.-]/.test(text[index] ?? "")) index++; if (text[index] === '"') index++; while (text[index] === " " || text[index] === "\t") index++; return text[index] === ":"; @@ -359,7 +617,7 @@ function looksLikeJsonBody(text: string, cursor: number): boolean { function readJsonPrefix(text: string, cursor: number): { prefix: string; from: number } { let from = cursor; - while (from > 0 && /[\w_"]/.test(text[from - 1] ?? "")) from--; + while (from > 0 && /[\w_.\-"]/.test(text[from - 1] ?? "")) from--; return { prefix: text.slice(from, cursor), from }; } @@ -391,6 +649,15 @@ function matchesFuzzyPrefix(value: string, prefix: string): boolean { return !normalizedPrefix || normalizedValue.includes(normalizedPrefix); } +// Kibana Console: flatten mapping to dotted paths, then prefix-filter the full name. +// `"ne"` matches `nested_comments` and `nested_comments.user`; `"nested_comments."` keeps only children. +// Do not use substring match (`rf_hotness` must not appear for `"ne"`). +function matchesElasticsearchFieldPrefix(fieldName: string, prefix: string): boolean { + const name = fieldName.toLowerCase(); + const typed = prefix.toLowerCase(); + return !typed || name.startsWith(typed); +} + function dedupeAndSort(items: ElasticsearchCompletionItem[]): ElasticsearchCompletionItem[] { const seen = new Set(); const deduped: ElasticsearchCompletionItem[] = []; diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index c724e6d1f4..cbec431135 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -138,6 +138,7 @@ import { useSavedSqlStore } from "@/stores/savedSqlStore"; import { decorateDatabaseSavedSqlTreeNodes, indexSavedSqlFilesByDatabase, stripDatabaseSavedSqlTreeNodes, withDatabaseSavedSqlRoot } from "@/lib/savedSql/savedSqlDatabaseTree"; import { encodeSqlServerLinkedSchema, parseSqlServerLinkedSchema } from "@/lib/database/sqlServerLinkedServers"; import { inferMongoCompletionFields, type MongoCompletionField } from "@/lib/mongo/mongoCompletion"; +import { flattenElasticsearchMappingFields, type ElasticsearchCompletionField } from "@/lib/elasticsearch/elasticsearchCompletion"; import { isMongoLegacyDriverProfile } from "@/lib/mongo/mongoCapabilities"; import { mongoCollectionKindFromNode, toMongoCollectionKind, visibleMongoCollections } from "@/lib/sidebar/mongoCollectionMutation"; import { completionSchemasFromTree, completionTablesFromTree } from "@/lib/metadata/completionTreeIndex"; @@ -486,6 +487,7 @@ export const useConnectionStore = defineStore("connection", () => { const primaryVisibleObjectNames = ref>({}); const sqlServerCompletionContextCache = ref>({}); const elasticsearchCompletionIndicesCache = ref>({}); + const elasticsearchCompletionFieldsCache = ref>({}); const redisCompletionKeysCache = ref>({}); const redisCommandDocsCache = ref>({}); const redisCommandDocsCacheGeneration = new Map(); @@ -3511,6 +3513,9 @@ export const useConnectionStore = defineStore("connection", () => { for (const key of Object.keys(elasticsearchCompletionIndicesCache.value)) { if (key === exactCacheKey || key.startsWith(cachePrefix)) delete elasticsearchCompletionIndicesCache.value[key]; } + for (const key of Object.keys(elasticsearchCompletionFieldsCache.value)) { + if (key === exactCacheKey || key.startsWith(cachePrefix)) delete elasticsearchCompletionFieldsCache.value[key]; + } for (const key of Object.keys(redisCompletionKeysCache.value)) { if (key === exactCacheKey || key.startsWith(cachePrefix)) delete redisCompletionKeysCache.value[key]; } @@ -7975,6 +7980,21 @@ export const useConnectionStore = defineStore("connection", () => { return elasticsearchCompletionIndicesCache.value[cacheKey]; } + async function listElasticsearchCompletionFields(connectionId: string, index: string): Promise { + if (!index) return []; + const cacheKey = `${connectionId}:${index}`; + const cached = elasticsearchCompletionFieldsCache.value[cacheKey]; + if (cached) return cached; + return withCompletionInFlight(`${cacheKey}:es-fields`, async () => { + await ensureConnected(connectionId); + const mapping = await api.elasticsearchGetIndexMetadata(connectionId, index, "mapping"); + const fields = flattenElasticsearchMappingFields(mapping); + elasticsearchCompletionFieldsCache.value[cacheKey] = fields; + evictOldestCacheEntries(elasticsearchCompletionFieldsCache.value, COMPLETION_CACHE_MAX); + return fields; + }); + } + // Upper bound on cached key names per db, to keep completion memory bounded // (Redis can hold far more keys than we ever want resident for autocomplete). const REDIS_COMPLETION_KEYS_MAX = 1000; @@ -9388,6 +9408,7 @@ export const useConnectionStore = defineStore("connection", () => { refreshCompletionSchemas, refreshCompletionDatabases, listElasticsearchCompletionIndices, + listElasticsearchCompletionFields, listRedisCompletionKeys, listRedisCompletionCommandDocs, listMongoCompletionCollections, diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs index 54acfce871..2f2ed2c33d 100644 --- a/crates/dbx-core/src/db/elasticsearch_driver.rs +++ b/crates/dbx-core/src/db/elasticsearch_driver.rs @@ -39,6 +39,9 @@ const KIBANA_PROXY_STATUS_HEADER: &str = "x-console-proxy-status-code"; const ELASTICSEARCH_REST_TABLE_MAX_BODY_BYTES: usize = 8 * 1024 * 1024; const ELASTICSEARCH_REST_TABLE_MAX_ROWS: usize = 2_000; const ELASTICSEARCH_REST_TABLE_MAX_CELLS: usize = 200_000; +const ELASTICSEARCH_MAPPING_MAX_DEPTH: u8 = 32; +const ELASTICSEARCH_MAPPING_MAX_COLUMNS: usize = 2_000; +const ELASTICSEARCH_MAPPING_THREAD_STACK: usize = 16 * 1024 * 1024; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ElasticsearchTransportMode { @@ -711,20 +714,12 @@ pub async fn get_columns(client: &EsClient, index: &str) -> Result Option<&serde_json::Map> { @@ -739,32 +734,85 @@ fn mapping_properties(mapping: &Value) -> Option<&serde_json::Map .find_map(|typed_mapping| typed_mapping.get("properties").and_then(Value::as_object)) } +fn mapping_bytes_to_columns_on_large_stack(bytes: Vec) -> Result, String> { + let (tx, rx) = std::sync::mpsc::sync_channel(1); + std::thread::Builder::new() + .name("es-mapping-columns".into()) + .stack_size(ELASTICSEARCH_MAPPING_THREAD_STACK) + .spawn(move || { + let _ = tx.send(mapping_bytes_to_columns(&bytes)); + }) + .map_err(|error| format!("Elasticsearch mapping thread failed: {error}"))?; + rx.recv().map_err(|error| format!("Elasticsearch mapping thread failed: {error}"))? +} + +fn mapping_bytes_to_columns(bytes: &[u8]) -> Result, String> { + let body: Value = serde_json::from_slice(bytes).map_err(|e| format!("Elasticsearch parse error: {e}"))?; + let mut seen = HashSet::new(); + let mut columns = Vec::new(); + if let Some(indices) = body.as_object() { + for index_mapping in indices.values() { + if let Some(properties) = mapping_properties(index_mapping) { + collect_mapping_columns("", properties, &mut seen, &mut columns); + } + } + } + columns.sort_by(|left, right| left.name.cmp(&right.name)); + drop_json_iteratively(body); + Ok(columns) +} + fn collect_mapping_columns( prefix: &str, properties: &serde_json::Map, seen: &mut HashSet, columns: &mut Vec, ) { - for (name, definition) in properties { - let field_name = if prefix.is_empty() { name.clone() } else { format!("{prefix}.{name}") }; - let field_type = definition.get("type").and_then(Value::as_str); - let children = definition.get("properties").and_then(Value::as_object); - - if let Some(data_type) = field_type { - push_mapping_column(&field_name, data_type, seen, columns); - } else if children.is_some() { - // Elasticsearch infers `object` for a field with `properties` when - // the mapping omits an explicit type. Return the parent too because - // the document grid renders top-level object values as one column. - push_mapping_column(&field_name, "object", seen, columns); + // Walk nested `properties` / multi-`fields` on the heap with hard caps. + // Recursive walks overflow tokio worker stacks on indexes like `dbx_all_types`. + let mut stack = vec![(prefix.to_string(), properties, 0_u8)]; + while let Some((prefix, properties, depth)) = stack.pop() { + if columns.len() >= ELASTICSEARCH_MAPPING_MAX_COLUMNS { + break; } + for (name, definition) in properties { + if columns.len() >= ELASTICSEARCH_MAPPING_MAX_COLUMNS { + break; + } + let field_name = if prefix.is_empty() { name.clone() } else { format!("{prefix}.{name}") }; + let field_type = definition.get("type").and_then(Value::as_str); + let nested_properties = definition.get("properties").and_then(Value::as_object); + let nested_fields = definition.get("fields").and_then(Value::as_object); - if let Some(fields) = definition.get("fields").and_then(Value::as_object) { - collect_mapping_columns(&field_name, fields, seen, columns); + if let Some(data_type) = field_type { + push_mapping_column(&field_name, data_type, seen, columns); + } else if nested_properties.is_some() { + // Elasticsearch infers `object` for a field with `properties` when + // the mapping omits an explicit type. Return the parent too because + // the document grid renders top-level object values as one column. + push_mapping_column(&field_name, "object", seen, columns); + } + + if depth + 1 >= ELASTICSEARCH_MAPPING_MAX_DEPTH { + continue; + } + if let Some(properties) = nested_properties { + stack.push((field_name.clone(), properties, depth + 1)); + } + if let Some(fields) = nested_fields { + stack.push((field_name, fields, depth + 1)); + } } + } +} - if let Some(children) = children { - collect_mapping_columns(&field_name, children, seen, columns); +fn drop_json_iteratively(value: Value) { + let mut stack = vec![value]; + while let Some(node) = stack.pop() { + match node { + Value::Array(items) => stack.extend(items), + Value::Object(map) => stack.extend(map.into_iter().map(|(_, child)| child)), + _ => {} } } } @@ -3217,6 +3265,23 @@ mod tests { assert!(columns.iter().any(|column| column.name == "profile.name" && column.data_type == "keyword")); } + #[test] + fn mapping_columns_do_not_overflow_on_deep_nested_objects() { + let mut current = json!({ "type": "keyword" }); + for depth in 0..512 { + current = json!({ "properties": { format!("n{depth}"): current } }); + } + let properties = current.get("properties").and_then(|value| value.as_object()).unwrap(); + let mut seen = std::collections::HashSet::new(); + let mut columns = Vec::new(); + + super::collect_mapping_columns("", properties, &mut seen, &mut columns); + + assert!(columns.iter().any(|column| column.name == "n511" && column.data_type == "object")); + assert_eq!(columns.len(), super::ELASTICSEARCH_MAPPING_MAX_DEPTH as usize); + super::drop_json_iteratively(current); + } + #[test] fn connectivity_check_disabled_parses_bool_and_string() { use super::elasticsearch_connectivity_check_disabled as disabled; diff --git a/packages/app-tests/elasticsearchCompletion.test.ts b/packages/app-tests/elasticsearchCompletion.test.ts index abeb822efb..6a7a0d6c12 100644 --- a/packages/app-tests/elasticsearchCompletion.test.ts +++ b/packages/app-tests/elasticsearchCompletion.test.ts @@ -1,13 +1,19 @@ import { strict as assert } from "node:assert"; import { test } from "vitest"; -import { buildElasticsearchCompletionItems, getElasticsearchCompletionContext, shouldAutoOpenElasticsearchCompletion } from "../../apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts"; +import { buildElasticsearchCompletionItems, elasticsearchCompletionNeedsFields, flattenElasticsearchMappingFields, getElasticsearchCompletionContext, shouldAutoOpenElasticsearchCompletion } from "../../apps/desktop/src/lib/elasticsearch/elasticsearchCompletion.ts"; import { buildSqlCompletionItems } from "../../apps/desktop/src/lib/sql/sqlCompletion.ts"; const indices = ["orders", "order_items", "users"]; +const fields = [ + { name: "city", dataType: "keyword" }, + { name: "status", dataType: "keyword" }, + { name: "createdAt", dataType: "date" }, + { name: "user.name", dataType: "text" }, +]; -function applyCompletion(text: string, cursor: number, label: string): string { +function applyCompletion(text: string, cursor: number, label: string, extra: { fields?: typeof fields } = {}): string { const context = getElasticsearchCompletionContext(text, cursor); - const item = buildElasticsearchCompletionItems(text, cursor, { indices }).find((candidate) => candidate.label === label); + const item = buildElasticsearchCompletionItems(text, cursor, { indices, ...extra }).find((candidate) => candidate.label === label); assert.ok(item, `Expected completion item ${label}`); // Mirrors the editor: the auto-closed quote at the cursor is part of the // replaced range, and snippet fields expand to their placeholder text. @@ -115,16 +121,13 @@ test("Elasticsearch JSON keys bring their value brackets along", () => { }); test("Elasticsearch query snippets provide term field and value placeholders", () => { - const text = 'GET /orders/_search\n{\n te'; + const text = "GET /orders/_search\n{\n te"; assert.equal(getElasticsearchCompletionContext(text, text.length).mode, "json"); const term = buildElasticsearchCompletionItems(text, text.length).find((item) => item.label === "term"); assert.ok(term); assert.equal(term?.type, "snippet"); - assert.equal( - applyCompletion(text, text.length, "term"), - 'GET /orders/_search\n{\n "term": {\n "field": "value"\n }', - ); + assert.equal(applyCompletion(text, text.length, "term"), 'GET /orders/_search\n{\n "term": {\n "field": "value"\n }'); assert.match(term?.apply ?? "", /\$\{field\}/); assert.match(term?.apply ?? "", /\$\{value\}/); }); @@ -150,7 +153,7 @@ test("Elasticsearch JSON slot survives comments and earlier requests", () => { }); test("Elasticsearch JSON value slots are not completed as keys", () => { - const text = 'GET /orders/_search\n{\n "sort": [\n "fi'; + const text = 'GET /orders/_search\n{\n "query": {\n "bool": {\n "must": [\n "fi'; const items = buildElasticsearchCompletionItems(text, text.length); assert.equal(getElasticsearchCompletionContext(text, text.length).jsonSlot, "value"); @@ -177,6 +180,200 @@ test("Elasticsearch completion auto trigger ignores structural JSON punctuation" assert.equal(shouldAutoOpenElasticsearchCompletion('GET /orders/_search\n{"query": {},', 'GET /orders/_search\n{"query": {},'.length), false); }); +test("Elasticsearch term field slot suggests index mapping fields instead of DSL keywords", () => { + const text = 'GET /orders/_search\n{\n "query": {\n "term": {\n "'; + const context = getElasticsearchCompletionContext(text, text.length); + + assert.equal(context.mode, "json"); + assert.equal(context.jsonFieldSlot, true); + assert.equal(context.index, "orders"); + assert.equal(elasticsearchCompletionNeedsFields(context), true); + + const items = buildElasticsearchCompletionItems(text, text.length, { fields }); + assert.equal(items.find((item) => item.label === '"city"')?.detail, "keyword"); + assert.equal(items.find((item) => item.label === '"city"')?.type, "column"); + assert.equal( + items.some((item) => item.label === '"should"' || item.label === '"size"' || item.label === '"sort"'), + false, + ); +}); + +test("Elasticsearch bool and query slots still suggest DSL keywords", () => { + const queryText = 'GET /orders/_search\n{\n "query": {\n "'; + assert.equal(getElasticsearchCompletionContext(queryText, queryText.length).jsonFieldSlot, false); + assert.ok(buildElasticsearchCompletionItems(queryText, queryText.length, { fields }).find((item) => item.label === '"bool"')); + assert.equal( + buildElasticsearchCompletionItems(queryText, queryText.length, { fields }).some((item) => item.label === '"city"'), + false, + ); + + const boolText = 'GET /orders/_search\n{\n "query": {\n "bool": {\n "'; + assert.equal(getElasticsearchCompletionContext(boolText, boolText.length).jsonFieldSlot, false); + assert.ok(buildElasticsearchCompletionItems(boolText, boolText.length, { fields }).find((item) => item.label === '"must"')); + assert.equal( + buildElasticsearchCompletionItems(boolText, boolText.length, { fields }).some((item) => item.label === '"city"'), + false, + ); +}); + +test("Elasticsearch exists field value and sort slots suggest mapping fields", () => { + const existsText = 'GET /orders/_search\n{\n "query": {\n "exists": {\n "field": "'; + assert.equal(getElasticsearchCompletionContext(existsText, existsText.length).jsonFieldSlot, true); + assert.ok(buildElasticsearchCompletionItems(existsText, existsText.length, { fields }).find((item) => item.label === '"status"')); + + const sortArray = 'GET /orders/_search\n{\n "sort": [\n "'; + assert.equal(getElasticsearchCompletionContext(sortArray, sortArray.length).jsonFieldSlot, true); + assert.ok(buildElasticsearchCompletionItems(sortArray, sortArray.length, { fields }).find((item) => item.label === '"createdAt"')); + + const sortObject = 'GET /orders/_search\n{\n "sort": [\n {\n "'; + assert.equal(getElasticsearchCompletionContext(sortObject, sortObject.length).jsonFieldSlot, true); + assert.ok(buildElasticsearchCompletionItems(sortObject, sortObject.length, { fields }).find((item) => item.label === '"createdAt"')); + + const sortMap = 'GET /orders/_search\n{\n "sort": {\n "'; + assert.equal(getElasticsearchCompletionContext(sortMap, sortMap.length).jsonFieldSlot, true); +}); + +test("Elasticsearch _source and fields objects are not mapping-field key slots", () => { + const sourceObject = 'GET /orders/_search\n{\n "_source": {\n "'; + assert.equal(getElasticsearchCompletionContext(sourceObject, sourceObject.length).jsonFieldSlot, false); + assert.ok(buildElasticsearchCompletionItems(sourceObject, sourceObject.length, { fields }).find((item) => item.label === '"fields"')); + + const fieldsObject = 'GET /orders/_search\n{\n "fields": [\n {\n "'; + assert.equal(getElasticsearchCompletionContext(fieldsObject, fieldsObject.length).jsonFieldSlot, false); + + const sourceArray = 'GET /orders/_search\n{\n "_source": [\n "'; + assert.equal(getElasticsearchCompletionContext(sourceArray, sourceArray.length).jsonFieldSlot, true); + assert.ok(buildElasticsearchCompletionItems(sourceArray, sourceArray.length, { fields }).find((item) => item.label === '"city"')); +}); + +test("Elasticsearch field completion ranks prefix matches before the 100-item cap", () => { + const text = 'GET /orders/_search\n{\n "query": {\n "term": {\n "st'; + const manyFields = [...Array.from({ length: 120 }, (_, index) => ({ name: `status_${index}`, dataType: "keyword" })), { name: "status", dataType: "keyword" }]; + const items = buildElasticsearchCompletionItems(text, text.length, { fields: manyFields }); + assert.ok(items.find((item) => item.label === '"status"')); + assert.equal(items.length, 100); + assert.ok(items.every((item) => item.label.startsWith('"st'))); +}); + +test("Elasticsearch field completion follows Kibana-style dotted-path prefix matching", () => { + const nestedFields = [ + { name: "nested_comments", dataType: "nested" }, + { name: "nested_comments.message", dataType: "text" }, + { name: "nested_comments.user", dataType: "keyword" }, + { name: "rf_hotness", dataType: "rank_feature" }, + ]; + const top = 'GET /orders/_search\n{\n "query": {\n "term": {\n "ne'; + const topItems = buildElasticsearchCompletionItems(top, top.length, { fields: nestedFields }); + assert.deepEqual( + topItems.map((item) => item.label), + ['"nested_comments"', '"nested_comments.message"', '"nested_comments.user"'], + ); + assert.ok((topItems[0]?.boost ?? 0) > (topItems[1]?.boost ?? 0)); + assert.equal( + [...topItems] + .sort((left, right) => (left.sortText ?? left.label).localeCompare(right.sortText ?? right.label)) + .map((item) => item.label) + .join(","), + topItems.map((item) => item.label).join(","), + ); + + const child = 'GET /orders/_search\n{\n "query": {\n "term": {\n "nested_comments.'; + const childItems = buildElasticsearchCompletionItems(child, child.length, { fields: nestedFields }); + assert.deepEqual(childItems.map((item) => item.label).sort(), ['"nested_comments.message"', '"nested_comments.user"']); +}); + +test("Elasticsearch aggregation terms and range keep DSL keys, while query terms stay field slots", () => { + const aggTerms = 'GET /orders/_search\n{\n "aggs": {\n "by_status": {\n "terms": {\n "'; + assert.equal(getElasticsearchCompletionContext(aggTerms, aggTerms.length).jsonFieldSlot, false); + assert.equal( + buildElasticsearchCompletionItems(aggTerms, aggTerms.length, { fields }).some((item) => item.label === '"status"'), + false, + ); + + const aggField = 'GET /orders/_search\n{\n "aggs": {\n "by_status": {\n "terms": {\n "field": "'; + assert.equal(getElasticsearchCompletionContext(aggField, aggField.length).jsonFieldSlot, true); + assert.ok(buildElasticsearchCompletionItems(aggField, aggField.length, { fields }).find((item) => item.label === '"status"')); + + const queryTerms = 'GET /orders/_search\n{\n "query": {\n "terms": {\n "'; + assert.equal(getElasticsearchCompletionContext(queryTerms, queryTerms.length).jsonFieldSlot, true); + + const filterAgg = 'GET /orders/_search\n{\n "aggs": {\n "active": {\n "filter": {\n "term": {\n "'; + assert.equal(getElasticsearchCompletionContext(filterAgg, filterAgg.length).jsonFieldSlot, true); +}); + +test("flattens a shallow all-types mapping without recursion", () => { + const mapping = { + dbx_all_types: { + mappings: { + properties: { + ali_title: { type: "alias", path: "title_text" }, + b_is_active: { type: "boolean" }, + nested_comments: { type: "nested", properties: { message: { type: "text" }, user: { type: "keyword" } } }, + obj_author: { properties: { age: { type: "integer" }, name: { type: "text" } } }, + title_text: { type: "text", fields: { keyword: { type: "keyword" } } }, + dv_vector: { type: "dense_vector", dims: 3, index_options: { type: "int8_hnsw", m: 16 } }, + }, + }, + }, + }; + const flattened = flattenElasticsearchMappingFields(mapping); + assert.equal(flattened.find((field) => field.name === "ali_title")?.dataType, "alias"); + assert.equal(flattened.find((field) => field.name === "nested_comments.message")?.dataType, "text"); + assert.equal(flattened.find((field) => field.name === "obj_author.name")?.dataType, "text"); + assert.equal(flattened.find((field) => field.name === "title_text.keyword")?.dataType, "keyword"); + assert.equal(flattened.find((field) => field.name === "dv_vector")?.dataType, "dense_vector"); + assert.equal( + flattened.some((field) => field.name.includes("index_options")), + false, + ); + assert.ok(flattened.length < 20); +}); + +test("flattening mapping fields ignores object identity cycles", () => { + const properties: Record = { + title_text: { type: "text" }, + }; + (properties.title_text as Record).properties = properties; + const flattened = flattenElasticsearchMappingFields({ dbx: { mappings: { properties } } }); + assert.equal(flattened.find((field) => field.name === "title_text")?.dataType, "text"); + assert.ok(flattened.length < 10); +}); + +test("Elasticsearch unquoted field items expose a bare filterText", () => { + const text = 'GET /orders/_search\n{\n "query": {\n "term": {\n ci'; + const item = buildElasticsearchCompletionItems(text, text.length, { fields }).find((candidate) => candidate.label === '"city"'); + assert.equal(item?.filterText, "city"); + assert.equal(item?.apply, '"city"'); +}); + +test("Elasticsearch root _search field slots do not expose an index", () => { + const text = 'GET /_search\n{\n "query": {\n "term": {\n "'; + const context = getElasticsearchCompletionContext(text, text.length); + assert.equal(context.jsonFieldSlot, true); + assert.equal(context.index, undefined); + assert.equal(elasticsearchCompletionNeedsFields(context), true); +}); + +test("Elasticsearch request index keeps wildcard and comma patterns", () => { + const wildcard = 'GET /logs-*/_search\n{\n "query": {\n "term": {\n "'; + assert.equal(getElasticsearchCompletionContext(wildcard, wildcard.length).index, "logs-*"); + + const multi = 'GET /orders,users/_search\n{\n "query": {\n "term": {\n "'; + assert.equal(getElasticsearchCompletionContext(multi, multi.length).index, "orders,users"); +}); + +test("Elasticsearch term field completion replaces the empty key in a value pair", () => { + const text = 'GET /orders/_search\n{\n "query": {\n "bool": {\n "must": [\n {\n "term": {\n "": "31"'; + const cursor = text.indexOf('""') + 1; + const context = getElasticsearchCompletionContext(text, cursor); + + assert.equal(context.jsonFieldSlot, true); + assert.equal(context.hasKeySeparator, true); + assert.equal(context.index, "orders"); + assert.equal(context.replaceClosingQuote, '"'); + assert.equal(applyCompletion(text, cursor, '"city"', { fields }), text.replace('""', '"city"')); +}); + test("SQL completion does not include Elasticsearch endpoints", () => { const items = buildSqlCompletionItems("select", "select".length, { tables: [],