diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.test.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.test.ts index d1e8f9ebaed..cf2bf247779 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.test.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.test.ts @@ -460,6 +460,18 @@ describe('DatabricksVectorStore', () => { ).rejects.toThrow('Failed primary keys: row-7'); }); + it('rejects the primary key as the content column without a request', async () => { + const store = await createStore(directDescribe, { contentColumn: 'id' }); + + await expect(store.addDocuments([{ pageContent: 'hello', metadata: {} }])).rejects.toThrow( + 'primary key', + ); + expect(fetchMock).not.toHaveBeenCalledWith( + expect.stringContaining('upsert-data'), + expect.anything(), + ); + }); + it('rejects a managed Delta Sync index without a request', async () => { const store = await managedStore(); diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.ts index 3554ad36023..897f9ebc6cb 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/DatabricksVectorStore.ts @@ -311,6 +311,12 @@ export class DatabricksVectorStore extends VectorStore { ): Promise { const vectorColumn = this.assertDirectAccess(); const { primaryKey, name } = this.index; + // The row below would write the document text over its own primary key + if (this.contentColumn === primaryKey) { + throw new UserError( + `Column ${primaryKey} is the primary key of ${name}. Select another content column`, + ); + } const schemaColumns = new Set(this.index.schemaColumns ?? []); const ids = documents.map((doc, i) => options?.ids?.[i] ?? doc.id ?? randomUUID()); diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.test.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.test.ts new file mode 100644 index 00000000000..3d64053c23d --- /dev/null +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.test.ts @@ -0,0 +1,474 @@ +/* eslint-disable n8n-nodes-base/node-param-display-name-miscased */ +/* eslint-disable n8n-nodes-base/node-param-display-name-miscased-id */ +import type { Embeddings } from '@langchain/core/embeddings'; +import { proxyFetch } from '@n8n/ai-utilities'; +import { DATABRICKS_PARTNER_USER_AGENT } from 'n8n-nodes-base/dist/nodes/Databricks/constants'; +import { createMockExecuteFunction } from 'n8n-nodes-base/test/nodes/Helpers'; +import type { + IDataObject, + IExecuteFunctions, + ILoadOptionsFunctions, + INode, + INodeProperties, + INodeType, + ISupplyDataFunctions, +} from 'n8n-workflow'; +import type { Mocked } from 'vitest'; +import { mock } from 'vitest-mock-extended'; + +import { getDatabricksTokenProvider } from '@utils/databricks/token-provider'; +import { DatabricksVectorStore } from './DatabricksVectorStore'; +import { VectorStoreDatabricks } from './VectorStoreDatabricks.node'; + +vi.mock('./DatabricksVectorStore', () => ({ + DatabricksVectorStore: { fromExistingIndex: vi.fn(), describeIndex: vi.fn() }, +})); +vi.mock('@utils/databricks/token-provider'); +vi.mock('@n8n/ai-utilities', async (importActual) => ({ + ...(await importActual()), + proxyFetch: vi.fn(), +})); + +const mockedFromExistingIndex = vi.mocked(DatabricksVectorStore.fromExistingIndex); +const mockedDescribeIndex = vi.mocked(DatabricksVectorStore.describeIndex); + +const mockCredential = { + host: 'https://ws.example.com/', + grantType: 'clientCredentials', + clientId: 'test-client-id', + clientSecret: 'test-client-secret', + allowedHttpRequestDomains: 'all', + allowedDomains: '', +}; + +const nodeDef: INode = { + id: '1', + name: 'Databricks Vector Store', + typeVersion: 1.3, + type: '@n8n/n8n-nodes-langchain.vectorStoreDatabricks', + position: [0, 0], + parameters: {}, +}; + +const baseParams = { + databricksIndex: { mode: 'list', value: 'cat.sch.idx' }, + contentColumn: 'text', + options: {}, +}; + +const indexInfo = { + name: 'cat.sch.idx', + primaryKey: 'id', + indexType: 'DELTA_SYNC' as const, + embedding: { kind: 'managed' as const, sourceColumn: 'text' }, + schemaColumns: ['id', 'text', 'source'], +}; + +describe('VectorStoreDatabricks', () => { + let node: VectorStoreDatabricks; + const methods = new VectorStoreDatabricks().methods as Required< + NonNullable + >; + let embeddings: ReturnType>; + + const setupContext = ( + params: IDataObject, + credentialOverrides: Partial = {}, + ) => { + const ctx = createMockExecuteFunction(params, nodeDef) as Mocked; + ctx.getCredentials = vi.fn().mockResolvedValue({ ...mockCredential, ...credentialOverrides }); + ctx.getInputData = vi.fn().mockReturnValue([{ json: {} }]); + ctx.getInputConnectionData = vi.fn().mockResolvedValue(embeddings); + ctx.logAiEvent = vi.fn(); + return ctx; + }; + + beforeEach(() => { + vi.resetAllMocks(); + node = new VectorStoreDatabricks(); + embeddings = mock(); + mockedDescribeIndex.mockResolvedValue(indexInfo); + vi.mocked(getDatabricksTokenProvider).mockReturnValue({ + getToken: vi.fn(async () => 'test-token'), + expiredStatus: 403, + }); + }); + + describe('description', () => { + it('is a hidden vector store node with the four default modes', () => { + expect(node.description).toMatchObject({ + name: 'vectorStoreDatabricks', + hidden: true, + credentials: [{ name: 'databricksOAuth2Api', required: true }], + }); + const mode = node.description.properties.find((p) => p.name === 'mode'); + expect(mode?.options?.map((option) => 'value' in option && option.value)).toEqual([ + 'load', + 'insert', + 'retrieve', + 'retrieve-as-tool', + ]); + }); + + it('declares the shared fields', () => { + const { properties } = node.description; + const index = properties.find((p) => p.name === 'databricksIndex'); + expect(index?.type).toBe('resourceLocator'); + expect(index?.modes?.[0].typeOptions?.searchListMethod).toBe('searchIndexes'); + + const contentColumn = properties.find((p) => p.name === 'contentColumn'); + expect(contentColumn?.typeOptions).toEqual({ + loadOptionsMethod: 'getIndexColumns', + loadOptionsDependsOn: ['databricksIndex.value'], + }); + + const options = properties.find((p) => p.name === 'options'); + const metadataColumns = (options?.options as INodeProperties[]).find( + (o) => o.name === 'metadataColumns', + ); + expect(metadataColumns?.type).toBe('multiOptions'); + expect(metadataColumns?.typeOptions?.loadOptionsMethod).toBe('getIndexColumns'); + expect(options?.displayOptions?.show?.mode).toEqual(['load', 'retrieve', 'retrieve-as-tool']); + }); + }); + + describe('supplyData in retrieve mode', () => { + it('creates the store from the credential host and the shared fields', async () => { + const store = {}; + mockedFromExistingIndex.mockResolvedValue(store as DatabricksVectorStore); + mockedDescribeIndex.mockResolvedValue(indexInfo); + const ctx = setupContext({ + ...baseParams, + mode: 'retrieve', + options: { + metadataColumns: ['source'], + metadata: { metadataValues: [{ name: 'source', value: 'hr' }] }, + }, + }); + + const result = await node.supplyData.call(ctx, 0); + + expect(mockedFromExistingIndex).toHaveBeenCalledWith(embeddings, { + fetch: expect.any(Function), + host: 'https://ws.example.com', + indexName: 'cat.sch.idx', + contentColumn: 'text', + metadataColumns: ['source'], + filter: { source: 'hr' }, + index: indexInfo, + }); + expect(mockedDescribeIndex).toHaveBeenCalledWith( + expect.any(Function), + 'https://ws.example.com', + 'cat.sch.idx', + ); + expect(result.response).toBeDefined(); + }); + + it('describes the index again after a failed describe in the same run', async () => { + mockedFromExistingIndex.mockResolvedValue({} as DatabricksVectorStore); + mockedDescribeIndex + .mockRejectedValueOnce(new Error('503 Service Unavailable')) + .mockResolvedValueOnce(indexInfo); + const ctx = setupContext({ ...baseParams, mode: 'retrieve' }); + + await expect(node.supplyData.call(ctx, 0)).rejects.toThrow('503'); + await node.supplyData.call(ctx, 0); + + expect(mockedDescribeIndex).toHaveBeenCalledTimes(2); + expect(mockedFromExistingIndex).toHaveBeenCalledTimes(1); + expect(mockedFromExistingIndex).toHaveBeenCalledWith( + embeddings, + expect.objectContaining({ index: indexInfo }), + ); + }); + + it('opens the run again after a failed credential lookup in the same run', async () => { + mockedFromExistingIndex.mockResolvedValue({} as DatabricksVectorStore); + mockedDescribeIndex.mockResolvedValue(indexInfo); + const ctx = setupContext({ ...baseParams, mode: 'retrieve' }); + ctx.getCredentials = vi + .fn() + .mockRejectedValueOnce(new Error('credential store unavailable')) + .mockResolvedValue(mockCredential); + + await expect(node.supplyData.call(ctx, 0)).rejects.toThrow('credential store unavailable'); + await node.supplyData.call(ctx, 0); + + expect(ctx.getCredentials).toHaveBeenCalledTimes(2); + expect(mockedFromExistingIndex).toHaveBeenCalledTimes(1); + }); + + it('rejects an http host before describing the index', async () => { + const ctx = setupContext( + { ...baseParams, mode: 'retrieve' }, + { host: 'http://ws.example.com' }, + ); + + await expect(node.supplyData.call(ctx, 0)).rejects.toThrow('must use https'); + expect(mockedFromExistingIndex).not.toHaveBeenCalled(); + }); + + it('sends the bearer, the partner User-Agent and a cancellable signal through the shared pool', async () => { + mockedFromExistingIndex.mockResolvedValue({} as DatabricksVectorStore); + vi.mocked(proxyFetch).mockResolvedValue(new Response('{}')); + const cancel = new AbortController(); + const ctx = setupContext({ ...baseParams, mode: 'retrieve' }); + ctx.getExecutionCancelSignal = vi.fn().mockReturnValue(cancel.signal); + await node.supplyData.call(ctx, 0); + + const { fetch } = mockedFromExistingIndex.mock.calls[0][1]; + await fetch('https://ws.example.com/x'); + + const [{ input, init, timeoutOptions }] = vi.mocked(proxyFetch).mock.calls[0]; + expect(String(input)).toBe('https://ws.example.com/x'); + const headers = new Headers(init?.headers); + expect(headers.get('authorization')).toBe('Bearer test-token'); + expect(headers.get('user-agent')).toBe(DATABRICKS_PARTNER_USER_AGENT); + // No timeoutOptions: they would force a fresh undici Agent for every request + expect(timeoutOptions).toBeUndefined(); + expect(init?.signal).toBeInstanceOf(AbortSignal); + expect(init?.signal?.aborted).toBe(false); + cancel.abort(); + expect(init?.signal?.aborted).toBe(true); + }); + + it('aborts on a signal the caller already set', async () => { + mockedFromExistingIndex.mockResolvedValue({} as DatabricksVectorStore); + vi.mocked(proxyFetch).mockResolvedValue(new Response('{}')); + const ctx = setupContext({ ...baseParams, mode: 'retrieve' }); + ctx.getExecutionCancelSignal = vi.fn().mockReturnValue(new AbortController().signal); + await node.supplyData.call(ctx, 0); + + const { fetch } = mockedFromExistingIndex.mock.calls[0][1]; + const caller = new AbortController(); + await fetch('https://ws.example.com/x', { signal: caller.signal }); + + const signal = vi.mocked(proxyFetch).mock.calls[0][0].init?.signal; + expect(signal?.aborted).toBe(false); + caller.abort(); + expect(signal?.aborted).toBe(true); + }); + }); + + describe('execute in load mode', () => { + it('searches by text and never embeds the prompt', async () => { + const doc = { pageContent: 'hello', metadata: { source: 'hr' } }; + const store = { similaritySearchWithScore: vi.fn().mockResolvedValue([[doc, 0.9]]) }; + mockedFromExistingIndex.mockResolvedValue(store as unknown as DatabricksVectorStore); + const ctx = setupContext({ + ...baseParams, + mode: 'load', + prompt: 'what is up', + topK: 2, + }); + + const result = await node.execute.call(ctx); + + expect(store.similaritySearchWithScore).toHaveBeenCalledWith('what is up', 2, undefined); + expect(embeddings.embedQuery).not.toHaveBeenCalled(); + expect(result).toEqual([[{ json: { document: doc, score: 0.9 }, pairedItem: { item: 0 } }]]); + }); + + it('mints the token and describes the index once for every item of the run', async () => { + const store = { similaritySearchWithScore: vi.fn().mockResolvedValue([]) }; + mockedFromExistingIndex.mockResolvedValue(store as unknown as DatabricksVectorStore); + mockedDescribeIndex.mockResolvedValue(indexInfo); + const ctx = setupContext({ + ...baseParams, + mode: 'load', + prompt: 'what is up', + topK: 2, + }); + ctx.getInputData = vi.fn().mockReturnValue([{ json: {} }, { json: {} }]); + + await node.execute.call(ctx); + + expect(store.similaritySearchWithScore).toHaveBeenCalledTimes(2); + expect(mockedFromExistingIndex).toHaveBeenCalledTimes(2); + expect(mockedDescribeIndex).toHaveBeenCalledTimes(1); + expect(getDatabricksTokenProvider).toHaveBeenCalledTimes(1); + expect(ctx.getCredentials).toHaveBeenCalledTimes(1); + }); + }); + + describe('execute in insert mode', () => { + it('adds the loaded documents through the store', async () => { + const documents = [ + { pageContent: 'hello', metadata: { source: 'hr' } }, + { pageContent: 'world', metadata: {} }, + ]; + const store = { + addDocuments: vi.fn(async (docs: unknown[]) => docs.map((_, i) => String(i))), + }; + mockedFromExistingIndex.mockResolvedValue(store as unknown as DatabricksVectorStore); + const ctx = setupContext({ ...baseParams, mode: 'insert' }); + ctx.getInputConnectionData = vi + .fn() + .mockResolvedValueOnce(embeddings) + .mockResolvedValueOnce(documents); + ctx.getExecutionCancelSignal = vi.fn().mockReturnValue(undefined); + + const result = await node.execute.call(ctx); + + expect(mockedFromExistingIndex).toHaveBeenCalledWith( + embeddings, + expect.objectContaining({ indexName: 'cat.sch.idx', filter: undefined }), + ); + expect(store.addDocuments).toHaveBeenCalledWith(documents); + expect(result).toEqual([documents.map((doc) => ({ json: doc, pairedItem: { item: 0 } }))]); + }); + }); + + describe('searchIndexes', () => { + const pages: Record = { + endpoints: { endpoints: [{ name: 'ep1' }, { name: 'ep2' }] }, + 'ep1:': { vector_indexes: [{ name: 'cat.sch.zeta', index_type: 'DIRECT_ACCESS' }] }, + 'ep2:': { + vector_indexes: [{ name: 'cat.sch.beta', index_type: 'DELTA_SYNC' }], + next_page_token: 'p2', + }, + 'ep2:p2': { vector_indexes: [{ name: 'cat.sch.alpha', index_type: 'DELTA_SYNC' }] }, + }; + + let httpRequestWithAuthentication: ReturnType; + let ctx: ILoadOptionsFunctions; + + const setupSearchContext = (host: string) => { + httpRequestWithAuthentication = vi.fn( + async (_type: string, options: { url: string; qs: Record }) => + options.url.endsWith('/endpoints') + ? pages.endpoints + : pages[`${options.qs.endpoint_name}:${options.qs.page_token ?? ''}`], + ); + ctx = { + getCredentials: vi.fn().mockResolvedValue({ ...mockCredential, host }), + getNode: vi.fn().mockReturnValue(nodeDef), + helpers: { httpRequestWithAuthentication }, + } as unknown as ILoadOptionsFunctions; + }; + + it('lists every index of every endpoint, sorted, with the type and endpoint', async () => { + setupSearchContext('https://ws.example.com/'); + + const result = await methods.listSearch.searchIndexes.call(ctx); + + expect(result.results).toEqual([ + { name: 'cat.sch.alpha', value: 'cat.sch.alpha', description: 'DELTA_SYNC - ep2' }, + { name: 'cat.sch.beta', value: 'cat.sch.beta', description: 'DELTA_SYNC - ep2' }, + { name: 'cat.sch.zeta', value: 'cat.sch.zeta', description: 'DIRECT_ACCESS - ep1' }, + ]); + expect(httpRequestWithAuthentication).toHaveBeenCalledTimes(4); + for (const [type, options] of httpRequestWithAuthentication.mock.calls) { + expect(type).toBe('databricksOAuth2Api'); + expect(options.url).toMatch(/^https:\/\/ws\.example\.com\/api\/2\.0\/vector-search\//); + expect(options.headers).toMatchObject({ 'User-Agent': DATABRICKS_PARTNER_USER_AGENT }); + } + }); + + it('does not let the bearer follow a cross-origin redirect', async () => { + setupSearchContext('https://ws.example.com'); + + await methods.listSearch.searchIndexes.call(ctx); + + for (const [, options] of httpRequestWithAuthentication.mock.calls) { + expect(options.sendCredentialsOnCrossOriginRedirect).toBe(false); + } + }); + + it('applies the substring filter to the name', async () => { + setupSearchContext('https://ws.example.com'); + + const result = await methods.listSearch.searchIndexes.call(ctx, 'ZETA'); + + expect(result.results.map((r) => r.value)).toEqual(['cat.sch.zeta']); + }); + + it('stops after 50 pages when a page echoes its own token', async () => { + setupSearchContext('https://ws.example.com'); + httpRequestWithAuthentication.mockResolvedValue({ endpoints: [], next_page_token: 'loop' }); + + await expect(methods.listSearch.searchIndexes.call(ctx)).rejects.toThrow('50 pages'); + expect(httpRequestWithAuthentication).toHaveBeenCalledTimes(50); + }); + + it('rejects an http host before any request', async () => { + setupSearchContext('http://ws.example.com'); + + await expect(methods.listSearch.searchIndexes.call(ctx)).rejects.toThrow('must use https'); + expect(httpRequestWithAuthentication).not.toHaveBeenCalled(); + }); + }); + + describe('getIndexColumns', () => { + const setupLoadOptionsContext = (indexName: string) => { + const ctx = createMockExecuteFunction( + {}, + nodeDef, + ) as Mocked; + ctx.getCredentials = vi.fn().mockResolvedValue(mockCredential); + ctx.getCurrentNodeParameter = vi.fn().mockReturnValue(indexName); + return ctx; + }; + + it('lists the embedding source column first and hides the primary key', async () => { + mockedDescribeIndex.mockResolvedValue({ + name: 'cat.sch.idx', + primaryKey: 'id', + indexType: 'DELTA_SYNC', + embedding: { kind: 'managed', sourceColumn: 'text' }, + schemaColumns: ['id', 'text', 'source'], + }); + const ctx = setupLoadOptionsContext('cat.sch.idx'); + + const result = await methods.loadOptions.getIndexColumns.call(ctx); + + expect(mockedDescribeIndex).toHaveBeenCalledWith( + expect.any(Function), + 'https://ws.example.com', + 'cat.sch.idx', + ); + expect(ctx.getCurrentNodeParameter).toHaveBeenCalledWith('databricksIndex', { + extractValue: true, + }); + expect(result).toEqual([ + { name: 'text', value: 'text', description: 'Embedding source column' }, + { name: 'source', value: 'source' }, + ]); + }); + + it('hides the vector column of a Direct Access index', async () => { + mockedDescribeIndex.mockResolvedValue({ + name: 'cat.sch.idx', + primaryKey: 'id', + indexType: 'DIRECT_ACCESS', + embedding: { kind: 'self', vectorColumn: 'embedding' }, + schemaColumns: ['id', 'text', 'embedding'], + }); + const ctx = setupLoadOptionsContext('cat.sch.idx'); + + await expect(methods.loadOptions.getIndexColumns.call(ctx)).resolves.toEqual([ + { name: 'text', value: 'text' }, + ]); + }); + + it('returns no columns when the index schema is unknown', async () => { + mockedDescribeIndex.mockResolvedValue({ + name: 'cat.sch.idx', + primaryKey: 'id', + indexType: 'DELTA_SYNC', + embedding: { kind: 'self', vectorColumn: 'embedding' }, + }); + const ctx = setupLoadOptionsContext('cat.sch.idx'); + + await expect(methods.loadOptions.getIndexColumns.call(ctx)).resolves.toEqual([]); + }); + + it('returns no columns before an index is chosen', async () => { + const ctx = setupLoadOptionsContext(''); + + await expect(methods.loadOptions.getIndexColumns.call(ctx)).resolves.toEqual([]); + expect(mockedDescribeIndex).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.ts new file mode 100644 index 00000000000..22ebe399f71 --- /dev/null +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.ts @@ -0,0 +1,296 @@ +import type { Embeddings } from '@langchain/core/embeddings'; +import { createVectorStoreNode, proxyFetch } from '@n8n/ai-utilities'; +import { + assertParamIsArray, + assertParamIsString, + type IExecuteFunctions, + type ILoadOptionsFunctions, + type INodeListSearchResult, + type INodeProperties, + type INodePropertyOptions, + type ISupplyDataFunctions, +} from 'n8n-workflow'; + +import { + createDatabricksAuthFetch, + type DatabricksFetchContext, +} from '@utils/databricks/auth-fetch'; +import { assertHttpsHost } from '@utils/databricks/constants'; +import { listDatabricksPages } from '@utils/databricks/list-pages'; +import { + DATABRICKS_CREDENTIAL_TYPE, + type DatabricksOAuth2Credential, +} from '@utils/databricks/token-provider'; +import { DatabricksVectorStore, type DatabricksIndexInfo } from './DatabricksVectorStore'; + +const REQUEST_TIMEOUT_MS = 60_000; + +const databricksIndexRLC: INodeProperties = { + displayName: 'Index', + name: 'databricksIndex', + type: 'resourceLocator', + default: { mode: 'list', value: '' }, + required: true, + modes: [ + { + displayName: 'From List', + name: 'list', + type: 'list', + placeholder: 'Select an index...', + typeOptions: { + searchListMethod: 'searchIndexes', + searchable: true, + }, + }, + { + displayName: 'Name', + name: 'id', + type: 'string', + placeholder: 'catalog.schema.index', + validation: [ + { + type: 'regex', + properties: { + regex: '^[^.\\s]+\\.[^.\\s]+\\.[^.\\s]+$', + errorMessage: 'Use the full name: catalog.schema.index', + }, + }, + ], + }, + ], +}; + +const columnTypeOptions = { + loadOptionsMethod: 'getIndexColumns', + loadOptionsDependsOn: ['databricksIndex.value'], +}; + +const sharedFields: INodeProperties[] = [ + databricksIndexRLC, + { + displayName: 'Content Column', + name: 'contentColumn', + type: 'options', + default: '', + typeOptions: columnTypeOptions, + description: + 'Column that becomes the document text. Leave empty on a managed-embedding index to use its embedding source column. Choose from the list, or specify an ID using an expression.', + }, + { + displayName: 'Options', + name: 'options', + type: 'collection', + placeholder: 'Add Option', + default: {}, + // Metadata Columns shapes the documents a read returns; the write path ignores it + displayOptions: { show: { mode: ['load', 'retrieve', 'retrieve-as-tool'] } }, + options: [ + { + displayName: 'Metadata Columns', + name: 'metadataColumns', + type: 'multiOptions', + default: [], + typeOptions: columnTypeOptions, + description: + 'Columns to return in document metadata. Defaults to all columns except the key, content and vector columns. Choose from the list, or specify IDs using an expression.', + }, + ], + }, +]; + +interface Run { + fetch: typeof fetch; + host: string; + /** Describe results of this run, keyed by index name */ + indexes: Map>; +} + +// The class owns its HTTP calls, so the runtime and the column dropdown share +// the token-refreshing fetch; the transport goes through the proxy-aware fetch +async function openRun(ctx: DatabricksFetchContext): Promise { + const credential = await ctx.getCredentials( + DATABRICKS_CREDENTIAL_TYPE, + ); + assertHttpsHost(ctx, credential.host); + const host = credential.host.replace(/\/$/, ''); + const egressFilter = ctx.helpers.getSecureEgressFilter(); + // ILoadOptionsFunctions has no cancel signal + const cancelSignal = + 'getExecutionCancelSignal' in ctx ? ctx.getExecutionCancelSignal() : undefined; + const { fetch } = createDatabricksAuthFetch(ctx, credential, { + endpointUrl: host, + egressFilter, + baseFetch: async (input, init) => { + // A whole-request deadline on the signal keeps the shared undici pool; per-request + // timeoutOptions would build a new Agent for every hop. The responses are small JSON + // bodies, so the deadline covers headers and body alike + const signals = [AbortSignal.timeout(REQUEST_TIMEOUT_MS)]; + if (init?.signal) signals.push(init.signal); + if (cancelSignal) signals.push(cancelSignal); + return await proxyFetch({ + input, + init: { ...init, signal: AbortSignal.any(signals) }, + egressFilter, + }); + }, + }); + return { fetch, host, indexes: new Map() }; +} + +// One context object serves every item of an execute call and every tool call of a +// supplyData closure, so the token mint and the describe GETs happen once per run +const runs = new WeakMap>(); + +async function databricksFetch(ctx: DatabricksFetchContext): Promise { + let run = runs.get(ctx); + if (!run) { + // Evict a rejected run so a failed credential lookup does not poison the rest of the run + run = openRun(ctx).catch((error) => { + runs.delete(ctx); + throw error; + }); + runs.set(ctx, run); + } + return await run; +} + +async function createStore( + ctx: IExecuteFunctions | ISupplyDataFunctions, + embeddings: Embeddings, + itemIndex: number, + filter?: Record, +): Promise { + const { fetch, host, indexes } = await databricksFetch(ctx); + const node = ctx.getNode(); + const indexName = ctx.getNodeParameter('databricksIndex', itemIndex, '', { extractValue: true }); + assertParamIsString('databricksIndex', indexName, node); + const contentColumn = ctx.getNodeParameter('contentColumn', itemIndex, ''); + assertParamIsString('contentColumn', contentColumn, node); + const metadataColumns = ctx.getNodeParameter('options.metadataColumns', itemIndex, []); + assertParamIsArray( + 'metadataColumns', + metadataColumns, + (value): value is string => typeof value === 'string', + node, + ); + + let index = indexes.get(indexName); + if (!index) { + // Evict a rejected describe so one transient failure does not poison the rest of the run + index = DatabricksVectorStore.describeIndex(fetch, host, indexName).catch((error) => { + indexes.delete(indexName); + throw error; + }); + indexes.set(indexName, index); + } + + return await DatabricksVectorStore.fromExistingIndex(embeddings, { + fetch, + host, + indexName, + contentColumn, + metadataColumns, + filter, + index: await index, + }); +} + +interface ListPage { + endpoints?: Array<{ name: string }>; + vector_indexes?: Array<{ name: string; index_type?: string }>; + next_page_token?: string; +} + +async function searchIndexes( + this: ILoadOptionsFunctions, + filter?: string, +): Promise { + const credentials = await this.getCredentials( + DATABRICKS_CREDENTIAL_TYPE, + ); + assertHttpsHost(this, credentials.host); + const host = credentials.host.replace(/\/$/, ''); + + const endpoints = await listDatabricksPages( + this, + `${host}/api/2.0/vector-search/endpoints`, + {}, + (page: ListPage) => page.endpoints, + ); + const results = ( + await Promise.all( + endpoints.map(async (endpoint) => { + const indexes = await listDatabricksPages( + this, + `${host}/api/2.0/vector-search/indexes`, + { endpoint_name: endpoint.name }, + (page: ListPage) => page.vector_indexes, + ); + // The type tells the user which indexes accept inserts + return indexes.map((index) => ({ + name: index.name, + value: index.name, + description: `${index.index_type} - ${endpoint.name}`, + })); + }), + ) + ) + .flat() + .sort((a, b) => a.name.localeCompare(b.name)); + + const filterLower = filter?.toLowerCase(); + return { + results: filterLower + ? results.filter((result) => result.name.toLowerCase().includes(filterLower)) + : results, + }; +} + +async function getIndexColumns(this: ILoadOptionsFunctions): Promise { + const indexName = this.getCurrentNodeParameter('databricksIndex', { extractValue: true }); + if (typeof indexName !== 'string' || indexName === '') return []; + + const { fetch, host } = await databricksFetch(this); + const info = await DatabricksVectorStore.describeIndex(fetch, host, indexName); + const { embedding } = info; + const source = embedding.kind === 'managed' ? embedding.sourceColumn : undefined; + const vector = embedding.kind === 'self' ? embedding.vectorColumn : undefined; + const columns = (info.schemaColumns ?? []) + .filter((column) => column !== info.primaryKey && column !== source && column !== vector) + .map((column) => ({ name: column, value: column })); + return source + ? [{ name: source, value: source, description: 'Embedding source column' }, ...columns] + : columns; +} + +export class VectorStoreDatabricks extends createVectorStoreNode({ + meta: { + displayName: 'Databricks Vector Store', + name: 'vectorStoreDatabricks', + description: 'Work with your data in Databricks Vector Search', + icon: { + light: 'file:../../shared/icons/databricks.svg', + dark: 'file:../../shared/icons/databricks.dark.svg', + }, + docsUrl: + 'https://docs.n8n.io/integrations/builtin/cluster-nodes/root-nodes/n8n-nodes-langchain.vectorstoredatabricks/', + credentials: [ + { + name: 'databricksOAuth2Api', + required: true, + }, + ], + }, + hidden: true, + // The index decides between query_text and query_vector inside the class + searchByText: true, + methods: { listSearch: { searchIndexes }, loadOptions: { getIndexColumns } }, + sharedFields, + async getVectorStoreClient(context, filter, embeddings, itemIndex) { + return await createStore(context, embeddings, itemIndex, filter); + }, + async populateVectorStore(context, embeddings, documents, itemIndex) { + const store = await createStore(context, embeddings, itemIndex); + await store.addDocuments(documents); + }, +}) {} diff --git a/packages/@n8n/nodes-langchain/package.json b/packages/@n8n/nodes-langchain/package.json index 472ad885a21..f6195124164 100644 --- a/packages/@n8n/nodes-langchain/package.json +++ b/packages/@n8n/nodes-langchain/package.json @@ -183,6 +183,7 @@ "dist/nodes/trigger/ChatTrigger/ChatTrigger.node.js", "dist/nodes/trigger/ChatTrigger/Chat.node.js", "dist/nodes/vector_store/VectorStoreAzureAISearch/VectorStoreAzureAISearch.node.js", + "dist/nodes/vector_store/VectorStoreDatabricks/VectorStoreDatabricks.node.js", "dist/nodes/vector_store/VectorStoreInMemory/VectorStoreInMemory.node.js", "dist/nodes/vector_store/VectorStoreInMemoryInsert/VectorStoreInMemoryInsert.node.js", "dist/nodes/vector_store/VectorStoreInMemoryLoad/VectorStoreInMemoryLoad.node.js", diff --git a/packages/@n8n/nodes-langchain/utils/databricks/auth-fetch.ts b/packages/@n8n/nodes-langchain/utils/databricks/auth-fetch.ts new file mode 100644 index 00000000000..4b7a92d7f6d --- /dev/null +++ b/packages/@n8n/nodes-langchain/utils/databricks/auth-fetch.ts @@ -0,0 +1,66 @@ +import { createRefreshingAuthFetch } from '@n8n/ai-utilities'; +import { + assertUrlAllowed, + getCredentialAllowedDomains, + type IExecuteFunctions, + type ILoadOptionsFunctions, + type ISupplyDataFunctions, + type NodeEgressFilter, +} from 'n8n-workflow'; + +import type { RefreshingTokenSource } from '../oauth2-token-provider'; +import { databricksAuthHeaders } from './constants'; +import type { DatabricksOAuth2Credential } from './token-provider'; +import { getDatabricksTokenProvider } from './token-provider'; + +export type DatabricksFetchContext = + | IExecuteFunctions + | ISupplyDataFunctions + | ILoadOptionsFunctions; + +/** + * The token-refreshing fetch every Databricks AI node sends its requests through. + * The caller owns the transport (`baseFetch`) and the egress filter; the token + * source comes back so the caller can read `expiredStatus`. + */ +export function createDatabricksAuthFetch( + ctx: DatabricksFetchContext, + credential: DatabricksOAuth2Credential, + options: { endpointUrl: string; egressFilter: NodeEgressFilter; baseFetch: typeof fetch }, +): { fetch: typeof fetch; tokenSource: RefreshingTokenSource } { + const { endpointUrl, egressFilter } = options; + const node = ctx.getNode(); + // endpointUrl derives from the credential's own host, so credentialOwnedSurface joins it to the allowlist + const allowedDomains = getCredentialAllowedDomains({ + node, + credentialData: credential, + credentialOwnedSurface: true, + nodeEndpointUrl: endpointUrl, + }); + + const tokenSource = getDatabricksTokenProvider(ctx, credential, egressFilter); + const { refreshAfterRejection } = tokenSource; + // The caller builds its own transport, so it never reaches the + // request helpers: `resolveHeaders` runs the expiry clock before every + // request, and `refreshHeaders` covers the rejection the clock missed - + // revoked server-side, or clock skew + const authFetch = createRefreshingAuthFetch({ + baseFetch: options.baseFetch, + expiredStatus: tokenSource.expiredStatus, + resolveHeaders: async () => databricksAuthHeaders(await tokenSource.getToken()), + ...(refreshAfterRejection && { + refreshHeaders: async () => { + const refreshed = await refreshAfterRejection(); + return refreshed ? databricksAuthHeaders(refreshed) : null; + }, + }), + assertAllowedUrl: async (hopUrl) => { + assertUrlAllowed({ url: hopUrl, allowedDomains, node }); + if (!egressFilter) return; + const result = await egressFilter.validateUrl(hopUrl); + if (!result.ok) throw result.error; + }, + }); + + return { fetch: authFetch, tokenSource }; +} diff --git a/packages/@n8n/nodes-langchain/utils/databricks/constants.ts b/packages/@n8n/nodes-langchain/utils/databricks/constants.ts index d99fa4bfa85..8de196ee060 100644 --- a/packages/@n8n/nodes-langchain/utils/databricks/constants.ts +++ b/packages/@n8n/nodes-langchain/utils/databricks/constants.ts @@ -1,6 +1,5 @@ import { DATABRICKS_PARTNER_USER_AGENT } from 'n8n-nodes-base/dist/nodes/Databricks/constants'; -import type { ILoadOptionsFunctions, ISupplyDataFunctions } from 'n8n-workflow'; -import { NodeOperationError } from 'n8n-workflow'; +import { NodeOperationError, type INode } from 'n8n-workflow'; /** * The bearer plus the partner User-Agent, set on every request and every retry. @@ -16,10 +15,7 @@ export const databricksAuthHeaders = (token: string): HeadersInit => ({ * Every request carries a secret (bearer token, or the client secret on the * mint path), so an http host would ship it in cleartext. */ -export function assertHttpsHost( - ctx: ILoadOptionsFunctions | ISupplyDataFunctions, - host: string, -): void { +export function assertHttpsHost(ctx: { getNode(): INode }, host: string): void { if (!URL.canParse(host) || new URL(host).protocol !== 'https:') { throw new NodeOperationError(ctx.getNode(), 'Databricks host must use https'); } diff --git a/packages/@n8n/nodes-langchain/utils/databricks/gateway-config.ts b/packages/@n8n/nodes-langchain/utils/databricks/gateway-config.ts index ae14d23c4e8..fb90606b195 100644 --- a/packages/@n8n/nodes-langchain/utils/databricks/gateway-config.ts +++ b/packages/@n8n/nodes-langchain/utils/databricks/gateway-config.ts @@ -1,15 +1,12 @@ import type { ClientOptions } from '@langchain/openai'; -import { createRefreshingAuthFetch, getProxyAgent } from '@n8n/ai-utilities'; -import { - assertUrlAllowed, - getCredentialAllowedDomains, - type ISupplyDataFunctions, -} from 'n8n-workflow'; +import { getProxyAgent } from '@n8n/ai-utilities'; +import type { ISupplyDataFunctions } from 'n8n-workflow'; -import { assertHttpsHost, databricksAuthHeaders, gatewayBaseUrl } from './constants'; -import { wrapDatabricksErrorFetch } from './error-handling'; import type { RefreshingTokenSource } from '../oauth2-token-provider'; -import { getDatabricksTokenProvider, type DatabricksOAuth2Credential } from './token-provider'; +import { createDatabricksAuthFetch } from './auth-fetch'; +import { assertHttpsHost, gatewayBaseUrl } from './constants'; +import { wrapDatabricksErrorFetch } from './error-handling'; +import type { DatabricksOAuth2Credential } from './token-provider'; /** * Builds the OpenAI-client transport for the Unity AI Gateway: the gateway @@ -28,44 +25,16 @@ export function createDatabricksGatewayConfig( assertHttpsHost(ctx, credential.host); const baseURL = gatewayBaseUrl(credential.host); - const node = ctx.getNode(); - // baseURL derives from the credential's own host, so credentialOwnedSurface joins it to the allowlist - const allowedDomains = getCredentialAllowedDomains({ - node, - credentialData: credential, - credentialOwnedSurface: true, - nodeEndpointUrl: baseURL, - }); - const egressFilter = ctx.helpers.getSecureEgressFilter(); - const tokenSource = getDatabricksTokenProvider(ctx, credential, egressFilter); - const { refreshAfterRejection } = tokenSource; + const { fetch: authFetch, tokenSource } = createDatabricksAuthFetch(ctx, credential, { + endpointUrl: baseURL, + egressFilter, + baseFetch: fetch, + }); const configuration: ClientOptions = { baseURL, - // The model client builds its own transport, so it never reaches the - // request helpers: `resolveHeaders` runs the expiry clock before every - // request, and `refreshHeaders` covers the rejection the clock missed - - // revoked server-side, or clock skew - fetch: wrapDatabricksErrorFetch( - createRefreshingAuthFetch({ - baseFetch: fetch, - expiredStatus: tokenSource.expiredStatus, - resolveHeaders: async () => databricksAuthHeaders(await tokenSource.getToken()), - ...(refreshAfterRejection && { - refreshHeaders: async () => { - const refreshed = await refreshAfterRejection(); - return refreshed ? databricksAuthHeaders(refreshed) : null; - }, - }), - assertAllowedUrl: async (hopUrl) => { - assertUrlAllowed({ url: hopUrl, allowedDomains, node }); - if (!egressFilter) return; - const result = await egressFilter.validateUrl(hopUrl); - if (!result.ok) throw result.error; - }, - }), - ), + fetch: wrapDatabricksErrorFetch(authFetch), fetchOptions: { dispatcher: getProxyAgent( baseURL, diff --git a/packages/@n8n/nodes-langchain/utils/databricks/list-pages.ts b/packages/@n8n/nodes-langchain/utils/databricks/list-pages.ts new file mode 100644 index 00000000000..87ad1a62d74 --- /dev/null +++ b/packages/@n8n/nodes-langchain/utils/databricks/list-pages.ts @@ -0,0 +1,48 @@ +import { DATABRICKS_PARTNER_USER_AGENT } from 'n8n-nodes-base/dist/nodes/Databricks/constants'; +import { NodeOperationError, type ILoadOptionsFunctions } from 'n8n-workflow'; + +import { DATABRICKS_CREDENTIAL_TYPE } from './token-provider'; + +const MAX_PAGES = 50; + +interface Page { + next_page_token?: string; +} + +/** + * Follows `next_page_token` through a Databricks list endpoint and concatenates + * what `pick` keeps from each page. Every picker pages through this one function + * so the request shape, and the redirect guard, cannot drift between them. + */ +export async function listDatabricksPages

( + ctx: ILoadOptionsFunctions, + url: string, + qs: Record, + pick: (page: P) => T[] | undefined, +): Promise { + let items: T[] = []; + let pageToken: string | undefined; + let pages = 0; + do { + // Guard against a host or proxy that echoes the same next_page_token back + if (++pages > MAX_PAGES) { + throw new NodeOperationError(ctx.getNode(), `Databricks list exceeded ${MAX_PAGES} pages`); + } + const page: P = await ctx.helpers.httpRequestWithAuthentication.call( + ctx, + DATABRICKS_CREDENTIAL_TYPE, + { + method: 'GET', + url, + qs: { ...qs, page_token: pageToken }, + headers: { Accept: 'application/json', 'User-Agent': DATABRICKS_PARTNER_USER_AGENT }, + json: true, + // The bearer must not follow a redirect off the workspace host + sendCredentialsOnCrossOriginRedirect: false, + }, + ); + items = items.concat(pick(page) ?? []); + pageToken = page.next_page_token; + } while (pageToken); + return items; +} diff --git a/packages/@n8n/nodes-langchain/utils/databricks/model-services.ts b/packages/@n8n/nodes-langchain/utils/databricks/model-services.ts index 669b5077e69..77e9ed93cb7 100644 --- a/packages/@n8n/nodes-langchain/utils/databricks/model-services.ts +++ b/packages/@n8n/nodes-langchain/utils/databricks/model-services.ts @@ -1,4 +1,3 @@ -import { DATABRICKS_PARTNER_USER_AGENT } from 'n8n-nodes-base/dist/nodes/Databricks/constants'; import { NodeApiError, NodeOperationError, @@ -7,6 +6,7 @@ import { } from 'n8n-workflow'; import { assertHttpsHost } from './constants'; +import { listDatabricksPages } from './list-pages'; import { DATABRICKS_CREDENTIAL_TYPE, type DatabricksOAuth2Credential } from './token-provider'; interface ModelService { @@ -57,34 +57,14 @@ export function makeModelSearch(capability: ModelCapability) { assertHttpsHost(this, credentials.host); const host = credentials.host.replace(/\/$/, ''); - const listModelServices = async (parent?: string): Promise => { - let services: ModelService[] = []; - let pageToken: string | undefined; - let pages = 0; - do { - // Guard against a host or proxy that echoes the same next_page_token back - if (++pages > 50) { - throw new NodeOperationError(this.getNode(), 'Model service list exceeded 50 pages'); - } - const page: ModelServicesResponse = await this.helpers.httpRequestWithAuthentication.call( - this, - DATABRICKS_CREDENTIAL_TYPE, - { - method: 'GET', - url: `${host}/api/2.1/unity-catalog/model-services`, - // FULL view is needed for supported_api_types - qs: { view: 'FULL', parent, page_token: pageToken }, - headers: { Accept: 'application/json', 'User-Agent': DATABRICKS_PARTNER_USER_AGENT }, - json: true, - // The bearer must not follow a redirect off the workspace host - sendCredentialsOnCrossOriginRedirect: false, - }, - ); - services = services.concat(page.model_services ?? []); - pageToken = page.next_page_token; - } while (pageToken); - return services; - }; + const listModelServices = async (parent?: string): Promise => + await listDatabricksPages( + this, + `${host}/api/2.1/unity-catalog/model-services`, + // FULL view is needed for supported_api_types + { view: 'FULL', parent }, + (page) => page.model_services, + ); let services: ModelService[]; try { diff --git a/packages/@n8n/nodes-langchain/utils/databricks/token-provider.ts b/packages/@n8n/nodes-langchain/utils/databricks/token-provider.ts index 8bb43fd87fe..e40dc9b6062 100644 --- a/packages/@n8n/nodes-langchain/utils/databricks/token-provider.ts +++ b/packages/@n8n/nodes-langchain/utils/databricks/token-provider.ts @@ -1,6 +1,7 @@ import { ClientOAuth2 } from '@n8n/client-oauth2'; import { DATABRICKS_PARTNER_USER_AGENT } from 'n8n-nodes-base/dist/nodes/Databricks/constants'; import type { + IExecuteFunctions, INode, ISupplyDataFunctions, ILoadOptionsFunctions, @@ -29,7 +30,7 @@ export type DatabricksOAuth2Credential = { * one-time-use refresh token, so it has to go through core to persist the rotation. */ export function getDatabricksTokenProvider( - ctx: ISupplyDataFunctions | ILoadOptionsFunctions, + ctx: IExecuteFunctions | ISupplyDataFunctions | ILoadOptionsFunctions, credential: DatabricksOAuth2Credential, egressFilter?: NodeEgressFilter, ): RefreshingTokenSource {