Core LLM integration infrastructure to allow pgAdmin to connect to AI providers. #9641

* Core infrastructure for LLM integration. * Add support for a number of different AI generated reports on security, performance, and schema design on servers, databases, and schemas, as appropriate. * Add a Natural Language AI assistant to the Query Tool. * Add an AI Insights panel to the EXPLAIN tool in the Query Tool, to analyse and report on issues in query plans.
2026-08-19 01:15:04 -05:00 · 2026-02-17 17:16:06 +05:30
parent 2715932464
commit f2756a3dcf
71 changed files with 14706 additions and 324 deletions
@@ -0,0 +1,181 @@
+/////////////////////////////////////////////////////////////
+//
+// pgAdmin 4 - PostgreSQL Tools
+//
+// Copyright (C) 2013 - 2025, The pgAdmin Development Team
+// This software is released under the PostgreSQL Licence
+//
+//////////////////////////////////////////////////////////////
+
+// Mock url_for
+jest.mock('sources/url_for', () => ({
+  __esModule: true,
+  default: jest.fn((endpoint) => `/mock/${endpoint}`),
+}));
+
+// Mock preferences store
+jest.mock('../../../pgadmin/preferences/static/js/store', () => ({
+  __esModule: true,
+  default: jest.fn(() => ({
+    getPreferencesForModule: jest.fn(() => ({})),
+  })),
+}));
+
+// Mock the QueryToolComponent to avoid importing all its dependencies
+jest.mock('../../../pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx', () => {
+  const React = require('react');
+  return {
+    QueryToolContext: React.createContext(null),
+    QueryToolEventsContext: React.createContext(null),
+  };
+});
+
+// Mock CodeMirror
+jest.mock('../../../pgadmin/static/js/components/ReactCodeMirror', () => ({
+  __esModule: true,
+  default: ({ value }) => <pre data-testid="codemirror">{value}</pre>,
+}));
+
+// Mock EmptyPanelMessage
+jest.mock('../../../pgadmin/static/js/components/EmptyPanelMessage', () => ({
+  __esModule: true,
+  default: ({ text }) => <div data-testid="empty-message">{text}</div>,
+}));
+
+// Mock Loader
+jest.mock('sources/components/Loader', () => ({
+  __esModule: true,
+  default: () => <div data-testid="loader">Loading...</div>,
+}));
+
+import { render, screen, fireEvent } from '@testing-library/react';
+import '@testing-library/jest-dom';
+import { withTheme } from '../fake_theme';
+import { NLQChatPanel } from '../../../pgadmin/tools/sqleditor/static/js/components/sections/NLQChatPanel.jsx';
+import {
+  QueryToolContext,
+  QueryToolEventsContext,
+} from '../../../pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx';
+
+// Mock the EventBus
+const createMockEventBus = () => ({
+  fireEvent: jest.fn(),
+  registerListener: jest.fn(),
+});
+
+// Mock the QueryToolContext
+const createMockQueryToolCtx = (isQueryTool = true) => ({
+  params: {
+    trans_id: 12345,
+    is_query_tool: isQueryTool,
+  },
+  api: {
+    post: jest.fn(),
+    get: jest.fn(),
+  },
+});
+
+// Helper to render with contexts
+const renderWithContexts = (component, { queryToolCtx, eventBus } = {}) => {
+  const mockEventBus = eventBus || createMockEventBus();
+  const mockQueryToolCtx = queryToolCtx || createMockQueryToolCtx();
+
+  return render(
+    <QueryToolContext.Provider value={mockQueryToolCtx}>
+      <QueryToolEventsContext.Provider value={mockEventBus}>
+        {component}
+      </QueryToolEventsContext.Provider>
+    </QueryToolContext.Provider>
+  );
+};
+
+describe('NLQChatPanel Component', () => {
+  let ThemedNLQChatPanel;
+
+  beforeAll(() => {
+    ThemedNLQChatPanel = withTheme(NLQChatPanel);
+
+    // Mock fetch for SSE
+    global.fetch = jest.fn();
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should render without crashing', () => {
+    const { container } = renderWithContexts(<ThemedNLQChatPanel />);
+    expect(container).toBeInTheDocument();
+  });
+
+  it('should show AI Assistant header', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    expect(screen.getByText('AI Assistant')).toBeInTheDocument();
+  });
+
+  it('should show empty state message when no messages', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    expect(
+      screen.getByText(/Describe what SQL you need/i)
+    ).toBeInTheDocument();
+  });
+
+  it('should have input field for typing queries', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    const input = screen.getByPlaceholderText(/Describe the SQL you need/i);
+    expect(input).toBeInTheDocument();
+  });
+
+  it('should have send button', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    const sendButton = screen.getByLabelText('Send');
+    expect(sendButton).toBeInTheDocument();
+  });
+
+  it('should have clear conversation button', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    const clearButton = screen.getByText('Clear');
+    expect(clearButton).toBeInTheDocument();
+  });
+
+  it('should disable send button when input is empty', () => {
+    const { container } = renderWithContexts(<ThemedNLQChatPanel />);
+    const sendButton = container.querySelector('button[data-label="Send"]');
+    expect(sendButton).toBeDisabled();
+  });
+
+  it('should enable send button when input has text', () => {
+    const { container } = renderWithContexts(<ThemedNLQChatPanel />);
+    const input = screen.getByPlaceholderText(/Describe the SQL you need/i);
+
+    fireEvent.change(input, { target: { value: 'Find all users' } });
+
+    const sendButton = container.querySelector('button[data-label="Send"]');
+    expect(sendButton).not.toBeDisabled();
+  });
+
+  it('should show message when not in query tool mode', () => {
+    const mockQueryToolCtx = createMockQueryToolCtx(false);
+    renderWithContexts(<ThemedNLQChatPanel />, {
+      queryToolCtx: mockQueryToolCtx,
+    });
+
+    expect(
+      screen.getByText(/AI Assistant is only available in Query Tool mode/i)
+    ).toBeInTheDocument();
+  });
+
+  it('should clear input after typing and clicking clear', () => {
+    renderWithContexts(<ThemedNLQChatPanel />);
+    const input = screen.getByPlaceholderText(/Describe the SQL you need/i);
+
+    fireEvent.change(input, { target: { value: 'Find all users' } });
+    expect(input.value).toBe('Find all users');
+
+    const clearButton = screen.getByText('Clear');
+    fireEvent.click(clearButton);
+
+    // Input should still have text (clear only clears messages)
+    expect(input.value).toBe('Find all users');
+  });
+});