diff --git a/apps/studio/data/database-policies/database-policies-query.ts b/apps/studio/data/database-policies/database-policies-query.ts index 8a447605223..cbec86be7b9 100644 --- a/apps/studio/data/database-policies/database-policies-query.ts +++ b/apps/studio/data/database-policies/database-policies-query.ts @@ -16,7 +16,8 @@ type DatabasePoliciesVariables = { export async function getDatabasePolicies( { projectRef, connectionString, schemas }: DatabasePoliciesVariables, - signal?: AbortSignal + signal?: AbortSignal, + headersInit?: HeadersInit ) { if (!projectRef) throw new Error('projectRef is required') @@ -28,7 +29,8 @@ export async function getDatabasePolicies( sql, queryKey: ['policies', schemas], }, - signal + signal, + headersInit ) return result as PGPolicy[] diff --git a/apps/studio/lib/ai/tools/index.test.ts b/apps/studio/lib/ai/tools/index.test.ts index c0157b35069..a97afb8cdd3 100644 --- a/apps/studio/lib/ai/tools/index.test.ts +++ b/apps/studio/lib/ai/tools/index.test.ts @@ -2,6 +2,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { getTools } from './index' import { getMcpTools } from './mcp-tools' +import { getSchemaTools } from './schema-tools' vi.mock('common', () => ({ IS_PLATFORM: true })) @@ -47,6 +48,16 @@ describe('ai/tools getTools', () => { expect(tools).toHaveProperty('incident_tool') }) + it('passes authorization through to schema tools', async () => { + await getTools(BASE_PARAMS) + + expect(getSchemaTools).toHaveBeenCalledWith({ + projectRef: BASE_PARAMS.projectRef, + connectionString: BASE_PARAMS.connectionString, + authorization: BASE_PARAMS.authorization, + }) + }) + it('degrades gracefully to the remaining tools when remote MCP fetch fails', async () => { const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) vi.mocked(getMcpTools).mockRejectedValueOnce(new Error('remote MCP unreachable')) diff --git a/apps/studio/lib/ai/tools/index.ts b/apps/studio/lib/ai/tools/index.ts index b4b041a78c6..30fad97232f 100644 --- a/apps/studio/lib/ai/tools/index.ts +++ b/apps/studio/lib/ai/tools/index.ts @@ -75,6 +75,7 @@ export const getTools = async ({ ...getSchemaTools({ projectRef, connectionString, + authorization, }), ...getReportTools({ projectRef, authorization }), ...(isExplorerEnabled diff --git a/apps/studio/lib/ai/tools/schema-tools.test.ts b/apps/studio/lib/ai/tools/schema-tools.test.ts new file mode 100644 index 00000000000..ecdbc4af7a2 --- /dev/null +++ b/apps/studio/lib/ai/tools/schema-tools.test.ts @@ -0,0 +1,61 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { getSchemaTools } from './schema-tools' +import { executeSql } from '@/data/sql/execute-sql-mutation' + +vi.mock('@/data/sql/execute-sql-mutation', () => ({ + executeSql: vi.fn(), +})) + +const TOOL_CONTEXT = { toolCallId: 'test', messages: [], context: {} } + +describe('ai/tools/schema-tools list_policies', () => { + beforeEach(() => { + vi.mocked(executeSql).mockReset() + vi.mocked(executeSql).mockResolvedValue({ + result: [ + { + name: 'Users can read own rows', + action: 'PERMISSIVE', + roles: ['authenticated'], + command: 'SELECT', + definition: '(auth.uid() = user_id)', + check: null, + }, + ], + }) + }) + + it('forwards the Authorization header to executeSql', async () => { + const tools = getSchemaTools({ + projectRef: 'abcdefghijklmnopqrst', + connectionString: 'encrypted', + authorization: 'Bearer token', + }) + if (!tools.list_policies.execute) throw new Error('execute is undefined') + + const result = await tools.list_policies.execute({ schemas: ['public'] }, TOOL_CONTEXT) + + expect(executeSql).toHaveBeenCalledWith( + expect.objectContaining({ + projectRef: 'abcdefghijklmnopqrst', + connectionString: 'encrypted', + }), + undefined, + { Authorization: 'Bearer token' } + ) + expect(result).toContain('Policy Name: "Users can read own rows"') + }) + + it('omits headers when no authorization is provided', async () => { + const tools = getSchemaTools({ + projectRef: 'abcdefghijklmnopqrst', + connectionString: 'encrypted', + }) + if (!tools.list_policies.execute) throw new Error('execute is undefined') + + await tools.list_policies.execute({ schemas: ['public'] }, TOOL_CONTEXT) + + expect(executeSql).toHaveBeenCalledWith(expect.any(Object), undefined, undefined) + }) +}) diff --git a/apps/studio/lib/ai/tools/schema-tools.ts b/apps/studio/lib/ai/tools/schema-tools.ts index 7e92adccb35..749e3c3cd6f 100644 --- a/apps/studio/lib/ai/tools/schema-tools.ts +++ b/apps/studio/lib/ai/tools/schema-tools.ts @@ -6,9 +6,11 @@ import { getDatabasePolicies } from '@/data/database-policies/database-policies- export const getSchemaTools = ({ projectRef, connectionString, + authorization, }: { projectRef: string connectionString: string + authorization?: string }) => ({ list_policies: tool({ description: 'Get existing RLS policies for a given schema', @@ -22,13 +24,15 @@ export const getSchemaTools = ({ connectionString, schemas, }, - undefined + undefined, + authorization ? { Authorization: authorization } : undefined ) const formattedPolicies = data .map( (policy) => ` Policy Name: "${policy.name}" + Table: "${policy.table}" Action: ${policy.action} Roles: ${policy.roles.join(', ')} Command: ${policy.command}