diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLExplainButton.tsx b/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLExplainButton.tsx index cd9a1071c5f..e181470cb91 100644 --- a/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLExplainButton.tsx +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLExplainButton.tsx @@ -1,18 +1,20 @@ -import { useCallback } from 'react'; +import { useCallback, useMemo } from 'react'; import { t } from '@grafana/i18n'; import { GenAIButton } from '../../../../dashboard/components/GenAI/GenAIButton'; import { EventTrackingSrc } from '../../../../dashboard/components/GenAI/tracking'; import { Message, Role } from '../../../../dashboard/components/GenAI/utils'; +import { SQLSchemasResponse } from '../hooks/useSQLSchemas'; import { getSQLExplanationSystemPrompt, QueryUsageContext } from './sqlPromptConfig'; +import { formatSchemasForPrompt } from './utils/formatSchemas'; interface GenAISQLExplainButtonProps { currentQuery: string; onExplain: (explanation: string) => void; refIds: string[]; - schemas?: unknown; // Reserved for future schema implementation + schemas?: SQLSchemasResponse | null; queryContext?: QueryUsageContext; } @@ -33,20 +35,20 @@ Explain what this query does in simple terms.`; * * @param refIds - The list of RefIDs available in the current context * @param currentQuery - The current SQL query to explain - * @param schemas - Optional schema information (planned for future implementation) + * @param formattedSchemas - Pre-formatted schema information string * @param queryContext - Optional query usage context * @returns A list of messages to be sent to the LLM for explaining the SQL query */ const getSQLExplanationMessages = ( refIds: string[], currentQuery: string, - schemas?: unknown, + formattedSchemas?: string, queryContext?: QueryUsageContext ): Message[] => { const systemPrompt = getSQLExplanationSystemPrompt({ refIds: refIds.length > 0 ? refIds.join(', ') : 'A', currentQuery: currentQuery.trim() || 'No current query provided', - schemas, // Will be utilized once schema extraction is implemented + formattedSchemas, queryContext, }); @@ -69,11 +71,13 @@ export const GenAISQLExplainButton = ({ onExplain, queryContext, refIds, - schemas, // Future implementation will use this for enhanced context + schemas, }: GenAISQLExplainButtonProps) => { + const formattedSchemas = useMemo(() => formatSchemasForPrompt(schemas), [schemas]); + const messages = useCallback(() => { - return getSQLExplanationMessages(refIds, currentQuery, schemas, queryContext); - }, [refIds, currentQuery, schemas, queryContext]); + return getSQLExplanationMessages(refIds, currentQuery, formattedSchemas, queryContext); + }, [refIds, currentQuery, formattedSchemas, queryContext]); const hasQuery = currentQuery && currentQuery.trim() !== ''; diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLSuggestionsButton.tsx b/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLSuggestionsButton.tsx index bdc839739c3..d20905846b3 100644 --- a/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLSuggestionsButton.tsx +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/GenAISQLSuggestionsButton.tsx @@ -1,12 +1,14 @@ -import { useCallback } from 'react'; +import { useCallback, useMemo } from 'react'; import { t } from '@grafana/i18n'; import { GenAIButton } from '../../../../dashboard/components/GenAI/GenAIButton'; import { EventTrackingSrc } from '../../../../dashboard/components/GenAI/tracking'; import { Message, Role } from '../../../../dashboard/components/GenAI/utils'; +import { SQLSchemasResponse } from '../hooks/useSQLSchemas'; import { getSQLSuggestionSystemPrompt, QueryUsageContext } from './sqlPromptConfig'; +import { formatSchemasForPrompt } from './utils/formatSchemas'; interface GenAISQLSuggestionsButtonProps { currentQuery: string; @@ -14,7 +16,7 @@ interface GenAISQLSuggestionsButtonProps { onHistoryUpdate?: (history: string[]) => void; refIds: string[]; initialQuery: string; - schemas?: unknown; // Reserved for future schema implementation + schemas?: SQLSchemasResponse | null; errorContext?: string[]; queryContext?: QueryUsageContext; } @@ -41,15 +43,15 @@ const getContextualPrompts = (refIds: string[], currentQuery: string): string[] * * @param refIds - The list of RefIDs available in the current context * @param currentQuery - The current SQL query being edited - * @param schemas - Optional schema information (planned for future implementation) - * @param errorContext - Optional error context for targeted fixes (planned for future implementation) + * @param formattedSchemas - Pre-formatted schema information string + * @param errorContext - Optional error context for targeted fixes * @param queryContext - Optional query usage context * @returns A list of messages to be sent to the LLM for generating SQL suggestions */ const getSQLSuggestionMessages = ( refIds: string[], currentQuery: string, - schemas?: unknown, + formattedSchemas?: string, errorContext?: string[], queryContext?: QueryUsageContext ): Message[] => { @@ -62,7 +64,7 @@ const getSQLSuggestionMessages = ( refIds: refIds.length > 0 ? refIds.join(', ') : 'A', currentQuery: trimmedQuery || 'No current query provided', queryInstruction: queryInstruction, - schemas, // Will be utilized once schema extraction is implemented + formattedSchemas, errorContext, queryContext, }); @@ -88,13 +90,15 @@ export const GenAISQLSuggestionsButton = ({ onHistoryUpdate, refIds, initialQuery, - schemas, // Future implementation will use this for enhanced context + schemas, errorContext, queryContext, }: GenAISQLSuggestionsButtonProps) => { + const formattedSchemas = useMemo(() => formatSchemasForPrompt(schemas), [schemas]); + const messages = useCallback(() => { - return getSQLSuggestionMessages(refIds, currentQuery, schemas, errorContext, queryContext); - }, [refIds, currentQuery, schemas, errorContext, queryContext]); + return getSQLSuggestionMessages(refIds, currentQuery, formattedSchemas, errorContext, queryContext); + }, [refIds, currentQuery, formattedSchemas, errorContext, queryContext]); const text = !currentQuery || currentQuery === initialQuery ? 'Generate suggestion' : 'Improve query'; diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.test.ts b/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.test.ts new file mode 100644 index 00000000000..bb1dc66d3a4 --- /dev/null +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.test.ts @@ -0,0 +1,52 @@ +import { getSQLSuggestionSystemPrompt, QueryUsageContext } from './sqlPromptConfig'; + +describe('getSQLSuggestionSystemPrompt', () => { + it('includes all context fields in generated prompt', () => { + const queryContext: QueryUsageContext = { + panelId: 'timeseries', + alerting: false, + dashboardContext: { + dashboardTitle: 'Production Metrics', + panelName: 'CPU Usage', + }, + datasources: ['prometheus', 'postgres'], + totalRows: 5000, + requestTime: 250, + numberOfQueries: 3, + }; + + const result = getSQLSuggestionSystemPrompt({ + refIds: 'A, B', + currentQuery: 'SELECT * FROM A', + queryInstruction: 'Focus on fixing the current query', + formattedSchemas: 'RefID A:\n - time (TIMESTAMP, not null)\n - value (FLOAT, nullable)', + errorContext: ['Syntax error near WHERE'], + queryContext, + }); + + expect(result).toContain('A, B'); + expect(result).toContain('SELECT * FROM A'); + expect(result).toContain('Focus on fixing the current query'); + expect(result).toContain('time (TIMESTAMP, not null)'); + expect(result).toContain('Syntax error near WHERE'); + expect(result).toContain('Panel Type: timeseries'); + expect(result).toContain('Dashboard: Production Metrics, Panel: CPU Usage'); + expect(result).toContain('Datasources: prometheus, postgres'); + expect(result).toContain('Total rows in the query: 5000'); + expect(result).toContain('Request time: 250'); + expect(result).toContain('Number of queries: 3'); + }); + + it('formats multiple errors with newlines', () => { + const result = getSQLSuggestionSystemPrompt({ + refIds: 'A', + currentQuery: 'SELECT * FROM A', + queryInstruction: 'Fix errors', + errorContext: ['Error 1: Syntax error', 'Error 2: Column not found'], + }); + + expect(result).toContain('Error 1: Syntax error'); + expect(result).toContain('Error 2: Column not found'); + expect(result).toContain('Error 1: Syntax error\nError 2: Column not found'); + }); +}); diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.ts b/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.ts index 5c7d6238334..d8bc77f62a2 100644 --- a/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.ts +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/sqlPromptConfig.ts @@ -1,10 +1,7 @@ /** * Configuration file for SQL AI prompts used across expression components - * NOTE: Schema and error context information integration is planned for future implementation */ -import { DataQuery } from '@grafana/schema'; - // Common SQL context information shared across all prompts const COMMON_SQL_CONTEXT = { engineInfo: 'MySQL dialectic based on dolthub go-mysql-server. The tables are all in memory', @@ -17,15 +14,14 @@ const TEMPLATE_PLACEHOLDERS = { refIds: '{refIds}', currentQuery: '{currentQuery}', queryInstruction: '{queryInstruction}', - schemaInfo: '{schemaInfo}', // Note: Schema information will be implemented in future updates - errorContext: '{errorContext}', // Note: Error context will be implemented in future updates + schemaInfo: '{schemaInfo}', + errorContext: '{errorContext}', queryContext: '{queryContext}', } as const; export interface QueryUsageContext { panelId?: string; alerting?: boolean; - queries?: DataQuery[]; dashboardContext?: { dashboardTitle?: string; panelName?: string; @@ -34,7 +30,6 @@ export interface QueryUsageContext { totalRows?: number; requestTime?: number; numberOfQueries?: number; - seriesData?: unknown; } /** @@ -58,8 +53,6 @@ ${TEMPLATE_PLACEHOLDERS.queryContext} Query instruction: ${TEMPLATE_PLACEHOLDERS.queryInstruction} -You may be able to derive schema information from the series data in queryContext. - Given the above data, help users with their SQL query by: - **PRIORITY: If there are errors listed above, focus on fixing them first** - Fixing syntax errors using available field and data type information @@ -96,8 +89,7 @@ Explain SQL queries clearly and concisely, focusing on: - What data is being selected and from which RefIDs - How the data is being transformed or aggregated - The purpose and business meaning of the query using dashboard and panel name from query context if relevant -- Performance implications and optimization opportunities. Database columns can not be indexed in context of Grafana sql expressions. Don't focus on - performance unless the query context has a requestTime or totalRows that looks like it could benefit from it. +- Performance implications and optimization opportunities. Database columns cannot be indexed in context of Grafana sql expressions. Don't focus on performance unless the query context has a requestTime or totalRows that looks like it could benefit from it. - Time series specific patterns and their significance Provide a clear explanation of what this SQL query does:`; @@ -112,33 +104,19 @@ const generateQueryContext = (queryContext?: QueryUsageContext): string => { const contextParts = []; if (queryContext.panelId) { - contextParts.push( - `Panel Type: ${queryContext.panelId}. Please use this to generate suggestions that are relevant to the panel type.` - ); + contextParts.push(`Panel Type: ${queryContext.panelId}`); } if (queryContext.alerting) { - contextParts.push( - 'Context: Alerting rule (focus on boolean/threshold results). Please use this to generate suggestions that are relevant to the alerting rule.' - ); - } - if (queryContext.queries) { - const queriesText = Array.isArray(queryContext.queries) - ? JSON.stringify(queryContext.queries, null, 2) - : String(queryContext.queries); - contextParts.push(`Queries available to use in the SQL Expression: ${queriesText}`); + contextParts.push('Context: Alerting rule (focus on boolean/threshold results)'); } if (queryContext.dashboardContext) { - const dashboardText = - typeof queryContext.dashboardContext === 'object' - ? JSON.stringify(queryContext.dashboardContext, null, 2) - : String(queryContext.dashboardContext); - contextParts.push(`Dashboard context (dashboard title and panel name): ${dashboardText}`); + const { dashboardTitle, panelName } = queryContext.dashboardContext; + if (dashboardTitle || panelName) { + contextParts.push(`Dashboard: ${dashboardTitle || 'Unknown'}, Panel: ${panelName || 'Unknown'}`); + } } - if (queryContext.datasources) { - const datasourcesText = Array.isArray(queryContext.datasources) - ? JSON.stringify(queryContext.datasources, null, 2) - : String(queryContext.datasources); - contextParts.push(`Datasources available to use in the SQL Expression: ${datasourcesText}`); + if (queryContext.datasources && queryContext.datasources.length > 0) { + contextParts.push(`Datasources: ${queryContext.datasources.join(', ')}`); } if (queryContext.totalRows) { contextParts.push(`Total rows in the query: ${queryContext.totalRows}`); @@ -149,13 +127,6 @@ const generateQueryContext = (queryContext?: QueryUsageContext): string => { if (queryContext.numberOfQueries) { contextParts.push(`Number of queries: ${queryContext.numberOfQueries}`); } - if (queryContext.seriesData) { - const seriesDataText = - typeof queryContext.seriesData === 'object' - ? JSON.stringify(queryContext.seriesData, null, 2) - : String(queryContext.seriesData); - contextParts.push(`Series data: ${seriesDataText}`); - } return contextParts.length ? `Query Context: @@ -170,19 +141,17 @@ export interface SQLPromptVariables { refIds: string; currentQuery: string; queryInstruction: string; - schemas?: unknown; // Reserved for future schema implementation + formattedSchemas?: string; errorContext?: string[]; queryContext?: QueryUsageContext; } /** * Generate the complete system prompt for SQL suggestions with enhanced context - * - * Note: Schema information integration is planned for future implementation */ export const getSQLSuggestionSystemPrompt = (variables: SQLPromptVariables): string => { const queryContext = generateQueryContext(variables.queryContext); - const schemaInfo = ''; // Placeholder for future schema information + const schemaInfo = variables.formattedSchemas ?? 'No schema information available.'; const errorContext = variables.errorContext?.length ? variables.errorContext.join('\n') : 'No current errors detected.'; @@ -200,8 +169,7 @@ export const getSQLSuggestionSystemPrompt = (variables: SQLPromptVariables): str */ export const getSQLExplanationSystemPrompt = (variables: Omit): string => { const queryContext = generateQueryContext(variables.queryContext); - - const schemaInfo = ''; // Placeholder for future schema information + const schemaInfo = variables.formattedSchemas ?? 'No schema information available.'; return SQL_EXPLANATION_SYSTEM_PROMPT.replaceAll(TEMPLATE_PLACEHOLDERS.refIds, variables.refIds) .replaceAll(TEMPLATE_PLACEHOLDERS.schemaInfo, schemaInfo) diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.test.ts b/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.test.ts new file mode 100644 index 00000000000..7c203235005 --- /dev/null +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.test.ts @@ -0,0 +1,90 @@ +import { SQLSchemasResponse } from '../../hooks/useSQLSchemas'; + +import { formatSchemasForPrompt } from './formatSchemas'; + +describe('formatSchemasForPrompt', () => { + const createSchemaResponse = (sqlSchemas: SQLSchemasResponse['sqlSchemas']): SQLSchemasResponse => ({ + kind: 'SQLSchemaResponse', + apiVersion: 'query.grafana.app/v0alpha1', + sqlSchemas, + }); + + it('returns fallback message for null/undefined/empty schemas', () => { + expect(formatSchemasForPrompt(null)).toBe('No schema information available.'); + expect(formatSchemasForPrompt(undefined)).toBe('No schema information available.'); + expect(formatSchemasForPrompt(createSchemaResponse({}))).toBe('No schema information available.'); + }); + + it('includes schema errors instead of silently failing', () => { + const schemas = createSchemaResponse({ + A: { + columns: null, + sampleRows: null, + error: 'Connection timeout', + }, + }); + + const result = formatSchemasForPrompt(schemas); + + expect(result).toBe('RefID A: Error - Connection timeout'); + }); + + it('truncates columns beyond default limit to prevent token overflow', () => { + const maxColumns = 10; // Default limit + const totalColumns = 15; + const columns = Array.from({ length: totalColumns }, (_, i) => ({ + name: `column_${i + 1}`, + mysqlType: 'VARCHAR', + dataFrameFieldType: 'string', + nullable: false, + })); + + const schemas = createSchemaResponse({ + A: { + columns, + sampleRows: null, + error: undefined, + }, + }); + + const result = formatSchemasForPrompt(schemas); + + // Should include columns up to limit + expect(result).toContain('column_1 (VARCHAR, not null)'); + expect(result).toContain(`column_${maxColumns} (VARCHAR, not null)`); + + // Should truncate columns beyond limit + expect(result).not.toContain(`column_${maxColumns + 1}`); + expect(result).toContain(`... and ${totalColumns - maxColumns} more columns`); + }); + + it('formats multiple RefIDs with proper separation', () => { + const schemas = createSchemaResponse({ + A: { + columns: [{ name: 'time', mysqlType: 'TIMESTAMP', dataFrameFieldType: 'time', nullable: false }], + sampleRows: [[1234567890]], + error: undefined, + }, + B: { + columns: [{ name: 'value', mysqlType: 'FLOAT', dataFrameFieldType: 'number', nullable: true }], + sampleRows: [[42.5]], + error: undefined, + }, + C: { + columns: [{ name: 'label', mysqlType: 'VARCHAR', dataFrameFieldType: 'string', nullable: false }], + sampleRows: [['production']], + error: undefined, + }, + }); + + const result = formatSchemasForPrompt(schemas); + + expect(result).toContain('RefID A:'); + expect(result).toContain('time (TIMESTAMP, not null)'); + expect(result).toContain('RefID B:'); + expect(result).toContain('value (FLOAT, nullable)'); + expect(result).toContain('RefID C:'); + expect(result).toContain('label (VARCHAR, not null)'); + expect(result.split('\n\n').length).toBe(3); + }); +}); diff --git a/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.ts b/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.ts new file mode 100644 index 00000000000..1ab4a00d959 --- /dev/null +++ b/public/app/features/expressions/components/SqlExpressions/GenAI/utils/formatSchemas.ts @@ -0,0 +1,70 @@ +import { SQLSchemasResponse, SQLSchemaField } from '../../hooks/useSQLSchemas'; + +const DEFAULT_MAX_COLUMNS = 10; + +/** + * Format schemas into a readable string for LLM context with token budget management + * + * This function converts SQL schema information into a human-readable format suitable + * for LLM prompts. It includes column names, types, nullability, and sample data. + * + * Token budget management: + * - Limits columns per RefID to prevent excessive token usage + * - Respects original column order from schema + * - Provides summary for truncated columns + * + * @param schemas - The SQL schemas response from the API + * @param maxColumnsPerRefId - Maximum number of columns to include per RefID (default: 10) + * @returns A formatted string representation of the schemas + */ +export const formatSchemasForPrompt = ( + schemas?: SQLSchemasResponse | null, + maxColumnsPerRefId = DEFAULT_MAX_COLUMNS +): string => { + if (!schemas?.sqlSchemas || Object.keys(schemas.sqlSchemas).length === 0) { + return 'No schema information available.'; + } + + const schemaParts: string[] = []; + + for (const [refId, schemaData] of Object.entries(schemas.sqlSchemas)) { + if (schemaData.error) { + schemaParts.push(`RefID ${refId}: Error - ${schemaData.error}`); + continue; + } + + if (!schemaData.columns || schemaData.columns.length === 0) { + schemaParts.push(`RefID ${refId}: No columns available`); + continue; + } + + const columnsToShow = schemaData.columns.slice(0, maxColumnsPerRefId); + const remainingCount = schemaData.columns.length - columnsToShow.length; + + const columnDescriptions = columnsToShow.map(({ nullable, name, mysqlType }: SQLSchemaField) => { + const isNullable = nullable ? 'nullable' : 'not null'; + return ` - ${name} (${mysqlType}, ${isNullable})`; + }); + + // Add truncation notice if we hit the limit + if (remainingCount > 0) { + columnDescriptions.push(` ... and ${remainingCount} more column${remainingCount > 1 ? 's' : ''}`); + } + + // Build schema text parts + const textParts = [`RefID ${refId}:`, columnDescriptions.join('\n')]; + + // Add sample data if available (first row only) + if (schemaData.sampleRows && schemaData.sampleRows.length > 0) { + const sampleRow = schemaData.sampleRows[0]; + const sampleValues = columnsToShow + .map(({ name }: SQLSchemaField, idx: number) => `${name}=${JSON.stringify(sampleRow[idx])}`) + .join(', '); + textParts.push(` Sample: ${sampleValues}`); + } + + schemaParts.push(textParts.join('\n')); + } + + return schemaParts.join('\n\n'); +}; diff --git a/public/app/features/expressions/components/SqlExpressions/SqlExpr.tsx b/public/app/features/expressions/components/SqlExpressions/SqlExpr.tsx index 361d748b716..cc6b8aaab2f 100644 --- a/public/app/features/expressions/components/SqlExpressions/SqlExpr.tsx +++ b/public/app/features/expressions/components/SqlExpressions/SqlExpr.tsx @@ -121,7 +121,6 @@ LIMIT refetch: refetchSchemas, } = useSQLSchemas({ queries, - enabled: isSchemaInspectorOpen, timeRange: metadata?.range, }); @@ -129,7 +128,6 @@ LIMIT () => ({ alerting, panelId: metadata?.data?.request?.panelPluginId, - queries: metadata?.queries, dashboardContext: { dashboardTitle: metadata?.data?.request?.dashboardTitle ?? '', panelName: metadata?.data?.request?.panelName ?? '', @@ -140,7 +138,6 @@ LIMIT ? metadata?.data?.request?.endTime - metadata?.data?.request?.startTime : -1, numberOfQueries: metadata?.data?.request?.targets?.length ?? 0, - seriesData: metadata?.data?.series, }), [alerting, metadata] ); @@ -276,7 +273,7 @@ LIMIT onExplain={handleExplain} queryContext={queryContext} refIds={vars} - // schemas={schemas} // Will be added when schema extraction is implemented + schemas={schemas} /> )} @@ -288,8 +285,8 @@ LIMIT onHistoryUpdate={handleHistoryUpdate} queryContext={queryContext} refIds={vars} - errorContext={errorContext} // Will be added when error tracking is implemented - // schemas={schemas} // Will be added when schema extraction is implemented + errorContext={errorContext} + schemas={schemas} /> diff --git a/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.test.ts b/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.test.ts new file mode 100644 index 00000000000..88f71bebc08 --- /dev/null +++ b/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.test.ts @@ -0,0 +1,141 @@ +import { renderHook, waitFor } from '@testing-library/react'; +import { testWithFeatureToggles } from 'test/test-utils'; + +import { useSQLSchemas, SQLSchemasResponse } from './useSQLSchemas'; + +jest.mock('@grafana/runtime', () => ({ + ...jest.requireActual('@grafana/runtime'), + getBackendSrv: jest.fn(), +})); + +jest.mock('@grafana/api-clients', () => ({ + getAPINamespace: jest.fn(() => 'default'), +})); + +jest.mock('@grafana/data', () => ({ + ...jest.requireActual('@grafana/data'), + getDefaultTimeRange: jest.fn(() => ({ + from: { toISOString: () => '2024-01-01T00:00:00.000Z' }, + to: { toISOString: () => '2024-01-01T23:59:59.999Z' }, + })), +})); + +describe('useSQLSchemas', () => { + testWithFeatureToggles({ enable: ['queryService'] }); + + beforeEach(() => jest.clearAllMocks()); + + const mockQueries = [ + { + refId: 'A', + datasource: { type: 'prometheus', uid: 'prom-uid' }, + }, + ]; + + const mockSchemasResponse: SQLSchemasResponse = { + kind: 'SQLSchemaResponse', + apiVersion: 'query.grafana.app/v0alpha1', + sqlSchemas: { + A: { + columns: [ + { name: 'time', mysqlType: 'TIMESTAMP', dataFrameFieldType: 'time', nullable: false }, + { name: 'value', mysqlType: 'FLOAT', dataFrameFieldType: 'number', nullable: true }, + ], + sampleRows: [[1234567890, 42.5]], + error: undefined, + }, + }, + }; + + const setupMockBackendSrv = (mockImplementation: jest.Mock) => { + const { getBackendSrv } = require('@grafana/runtime'); + getBackendSrv.mockReturnValue({ post: mockImplementation }); + return mockImplementation; + }; + + it('fetches schemas successfully and updates state', async () => { + // Arrange + const mockPost = setupMockBackendSrv(jest.fn().mockResolvedValue(mockSchemasResponse)); + + // Act - Render the hook + const { result } = renderHook(() => useSQLSchemas({ queries: mockQueries })); + + // Assert - Initial state + expect(result.current.loading).toBe(true); + expect(result.current.schemas).toBe(null); + expect(result.current.error).toBe(null); + + // Wait for async fetch to complete + await waitFor(() => expect(result.current.loading).toBe(false)); + + // Assert - Final state + expect(result.current.schemas).toEqual(mockSchemasResponse); + expect(result.current.schemas?.sqlSchemas.A.columns).toHaveLength(2); + expect(result.current.schemas?.sqlSchemas.A.columns?.[0].name).toBe('time'); + expect(result.current.error).toBe(null); + + // Verify API was called correctly + expect(mockPost).toHaveBeenCalledWith( + '/apis/query.grafana.app/v0alpha1/namespaces/default/sqlschemas/name', + expect.objectContaining({ + queries: mockQueries, + from: '2024-01-01T00:00:00.000Z', + to: '2024-01-01T23:59:59.999Z', + }) + ); + }); + + it('handles API errors gracefully without crashing', async () => { + // Arrange + const mockError = new Error('Network error'); + setupMockBackendSrv(jest.fn().mockRejectedValue(mockError)); + + // Act + const { result } = renderHook(() => useSQLSchemas({ queries: mockQueries })); + + // Wait for error to be set + await waitFor(() => expect(result.current.loading).toBe(false)); + + // Assert - Error state is set, no crash + expect(result.current.error).toEqual(mockError); + expect(result.current.schemas).toBe(null); + expect(result.current.loading).toBe(false); + }); + + describe('when feature flag is disabled', () => { + testWithFeatureToggles({ disable: ['queryService', 'grafanaAPIServerWithExperimentalAPIs'] }); + + it('does not fetch schemas', async () => { + // Arrange + setupMockBackendSrv(jest.fn()); + + // Act + const { result } = renderHook(() => useSQLSchemas({ queries: mockQueries })); + + // Assert - No API call, feature disabled + expect(result.current.isFeatureEnabled).toBe(false); + expect(result.current.loading).toBe(false); + expect(result.current.schemas).toBe(null); + }); + }); + + it('returns empty schemas for empty queries array without calling API', async () => { + // Arrange + const mockPost = setupMockBackendSrv(jest.fn()); + + // Act - Empty queries array + const { result } = renderHook(() => useSQLSchemas({ queries: [] })); + + // Wait for state to settle + await waitFor(() => expect(result.current.loading).toBe(false)); + + // Assert - Returns empty schemas, no API call + expect(result.current.schemas).toEqual({ + kind: 'SQLSchemaResponse', + apiVersion: 'query.grafana.app/v0alpha1', + sqlSchemas: {}, + }); + expect(result.current.error).toBe(null); + expect(mockPost).not.toHaveBeenCalled(); // No API call for empty queries + }); +}); diff --git a/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.ts b/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.ts index fd9f5f41a1e..9ac6b7f2fba 100644 --- a/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.ts +++ b/public/app/features/expressions/components/SqlExpressions/hooks/useSQLSchemas.ts @@ -28,11 +28,10 @@ export interface SQLSchemasResponse { interface UseSQLSchemasOptions { queries?: DataQuery[]; - enabled: boolean; timeRange?: TimeRange; } -export function useSQLSchemas({ queries, enabled, timeRange }: UseSQLSchemasOptions) { +export function useSQLSchemas({ queries, timeRange }: UseSQLSchemasOptions) { const isFeatureEnabled = useMemo( () => config.featureToggles.queryService || config.featureToggles.grafanaAPIServerWithExperimentalAPIs || false, [] @@ -40,7 +39,7 @@ export function useSQLSchemas({ queries, enabled, timeRange }: UseSQLSchemasOpti // Start with loading=true if we're going to fetch on mount const [schemas, setSchemas] = useState(null); - const [loading, setLoading] = useState(enabled && isFeatureEnabled && Boolean(queries)); + const [loading, setLoading] = useState(isFeatureEnabled && Boolean(queries)); const [error, setError] = useState(null); // Store queries in ref so we can access current value without triggering effect @@ -48,7 +47,7 @@ export function useSQLSchemas({ queries, enabled, timeRange }: UseSQLSchemasOpti queriesRef.current = queries; const fetchSchemas = useCallback(async () => { - if (!enabled || !isFeatureEnabled) { + if (!isFeatureEnabled) { return; } @@ -85,7 +84,7 @@ export function useSQLSchemas({ queries, enabled, timeRange }: UseSQLSchemasOpti } finally { setLoading(false); } - }, [enabled, isFeatureEnabled, timeRange]); + }, [isFeatureEnabled, timeRange]); useEffect(() => { fetchSchemas();