diff --git a/apps/dbagent/src/app/api/sql/route.ts b/apps/dbagent/src/app/api/sql/route.ts new file mode 100644 index 00000000..24c08dc3 --- /dev/null +++ b/apps/dbagent/src/app/api/sql/route.ts @@ -0,0 +1,49 @@ +import { NextResponse } from 'next/server'; +import { auth } from '~/auth'; // Assuming auth is used for protecting API routes + +export async function POST(req: Request) { + const session = await auth(); + if (!session) { + return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }); + } + + try { + const body = await req.json(); + const { query } = body; + + if (!query || typeof query !== 'string') { + return NextResponse.json({ error: 'Query is required and must be a string' }, { status: 400 }); + } + + // --- Placeholder for actual query execution --- + console.log('Received SQL query:', query); + // Simulate database execution + // In a real scenario, you would connect to the database and run the query here. + // For example: + // const dbClient = await getDbClient(); // Function to get a database client + // const results = await dbClient.query(query); + // --- End of placeholder --- + + // Simulate successful execution with dummy results + const dummyResults = [ + { id: 1, name: 'Dummy Result 1' }, + { id: 2, name: 'Dummy Result 2' } + ]; + + // Simulate an error for demonstration purposes if query contains "ERROR" + if (query.toUpperCase().includes('ERROR')) { + console.error('Simulated error executing query:', query); + return NextResponse.json({ error: 'Simulated error executing query' }, { status: 500 }); + } + + console.log('Simulated query execution successful.'); + return NextResponse.json({ results: dummyResults }); + } catch (error) { + console.error('Error in /api/sql/route.ts:', error); + let errorMessage = 'Internal Server Error'; + if (error instanceof Error) { + errorMessage = error.message; + } + return NextResponse.json({ error: errorMessage }, { status: 500 }); + } +} diff --git a/apps/dbagent/src/components/chat/artifacts/artifact.tsx b/apps/dbagent/src/components/chat/artifacts/artifact.tsx index e46707d3..8d0b78f0 100644 --- a/apps/dbagent/src/components/chat/artifacts/artifact.tsx +++ b/apps/dbagent/src/components/chat/artifacts/artifact.tsx @@ -14,12 +14,13 @@ import { ArtifactActions } from './artifact-actions'; import { ArtifactCloseButton } from './artifact-close-button'; import { ArtifactMessages } from './artifact-messages'; import { sheetArtifact } from './sheet/client'; +import { sqlArtifact } from './sql/client'; import { textArtifact } from './text/client'; import { Toolbar } from './toolbar'; import { useArtifact } from './use-artifact'; import { VersionFooter } from './version-footer'; -export const artifactDefinitions = [textArtifact, sheetArtifact]; +export const artifactDefinitions = [textArtifact, sheetArtifact, sqlArtifact]; export type ArtifactKind = (typeof artifactDefinitions)[number]['kind']; export interface UIArtifact { diff --git a/apps/dbagent/src/components/chat/artifacts/sql/client.test.tsx b/apps/dbagent/src/components/chat/artifacts/sql/client.test.tsx new file mode 100644 index 00000000..7698724d --- /dev/null +++ b/apps/dbagent/src/components/chat/artifacts/sql/client.test.tsx @@ -0,0 +1,285 @@ +import '@testing-library/jest-dom'; +import { render, screen } from '@testing-library/react'; +import React from 'react'; +import { sqlArtifact } from './client'; // Adjust path as necessary + +// Mock toast +const mockToast = { + info: jest.fn(), + success: jest.fn(), + error: jest.fn() +}; +jest.mock('@xata.io/components', () => ({ + ...jest.requireActual('@xata.io/components'), // Import and retain default exports + toast: mockToast // Mock the toast export +})); + +// Mock dependencies and props +const mockSetMetadata = jest.fn(); +const mockOnSaveContent = jest.fn(); +const mockGetDocumentContentById = jest.fn(); + +const defaultProps: React.ComponentProps = { + title: 'Test SQL Query', + content: 'SELECT * FROM users;', + mode: 'edit', + isCurrentVersion: true, + currentVersionIndex: 0, + status: 'idle', + suggestions: [], + onSaveContent: mockOnSaveContent, + isInline: false, + getDocumentContentById: mockGetDocumentContentById, + isLoading: false, + metadata: {}, + setMetadata: mockSetMetadata +}; + +describe('SqlArtifact Content', () => { + it('should render the SQL query content', () => { + render(React.createElement(sqlArtifact.content, defaultProps)); + + // Check if the SQL query is displayed + const queryElement = screen.getByText((content, element) => { + // Allow matching part of the text content if it's inside a
 or similar
+      const hasText = (node: Element | null) => node?.textContent === defaultProps.content;
+      const elementHasText = hasText(element);
+      const childrenDontHaveText = Array.from(element?.children || []).every((child) => !hasText(child));
+      return elementHasText && childrenDontHaveText;
+    });
+    expect(queryElement).toBeInTheDocument();
+
+    // Check if it's in a 
 tag for formatting
+    expect(queryElement.tagName).toBe('PRE');
+  });
+
+  // More tests will be added here for "Run Query" and "View Results"
+
+  describe('Run Query Action', () => {
+    // Mock fetch globally for these tests
+    global.fetch = jest.fn();
+
+    beforeEach(() => {
+      // Reset mocks before each test
+      (global.fetch as jest.Mock).mockClear();
+      mockToast.info.mockClear();
+      mockToast.success.mockClear();
+      mockToast.error.mockClear();
+      mockSetMetadata.mockClear(); // Assuming mockSetMetadata is available from outer scope
+    });
+
+    const runQueryAction = sqlArtifact.actions.find((a) => a.description === 'Run Query');
+
+    if (!runQueryAction) {
+      throw new Error('Run Query action not found in sqlArtifact.actions');
+    }
+
+    it('should call /api/sql with the query and show success toast on successful execution', async () => {
+      (global.fetch as jest.Mock).mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({ results: [{ id: 1, name: 'Test' }] })
+      });
+
+      const actionContext = {
+        content: 'SELECT * FROM test_table;',
+        handleVersionChange: jest.fn(),
+        currentVersionIndex: 0,
+        isCurrentVersion: true,
+        mode: 'edit' as 'edit' | 'diff',
+        metadata: {},
+        setMetadata: mockSetMetadata
+      };
+
+      await runQueryAction.onClick(actionContext);
+
+      expect(global.fetch).toHaveBeenCalledWith('/api/sql', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ query: actionContext.content })
+      });
+      expect(mockToast.info).toHaveBeenCalledWith('Running query...');
+      expect(mockToast.success).toHaveBeenCalledWith('Query executed successfully!');
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        1,
+        expect.objectContaining({ isRunningQuery: true, error: null })
+      );
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        2,
+        expect.objectContaining({ results: [{ id: 1, name: 'Test' }], error: null, isRunningQuery: false })
+      );
+    });
+
+    it('should show error toast if query is empty', async () => {
+      const actionContext = {
+        content: ' ', // Empty query
+        handleVersionChange: jest.fn(),
+        currentVersionIndex: 0,
+        isCurrentVersion: true,
+        mode: 'edit' as 'edit' | 'diff',
+        metadata: {},
+        setMetadata: mockSetMetadata
+      };
+
+      await runQueryAction.onClick(actionContext);
+
+      expect(global.fetch).not.toHaveBeenCalled();
+      expect(mockToast.error).toHaveBeenCalledWith('Query is empty.');
+      expect(mockSetMetadata).not.toHaveBeenCalled();
+    });
+
+    it('should show error toast on API failure', async () => {
+      (global.fetch as jest.Mock).mockResolvedValueOnce({
+        ok: false,
+        status: 500,
+        json: async () => ({ error: 'Internal Server Error' })
+      });
+
+      const actionContext = {
+        content: 'SELECT * FROM error_table;',
+        handleVersionChange: jest.fn(),
+        currentVersionIndex: 0,
+        isCurrentVersion: true,
+        mode: 'edit' as 'edit' | 'diff',
+        metadata: {},
+        setMetadata: mockSetMetadata
+      };
+
+      await runQueryAction.onClick(actionContext);
+
+      expect(global.fetch).toHaveBeenCalledWith('/api/sql', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ query: actionContext.content })
+      });
+      expect(mockToast.info).toHaveBeenCalledWith('Running query...');
+      expect(mockToast.error).toHaveBeenCalledWith('Failed to run query: Internal Server Error');
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        1,
+        expect.objectContaining({ isRunningQuery: true, error: null })
+      );
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        2,
+        expect.objectContaining({ results: null, error: 'Internal Server Error', isRunningQuery: false })
+      );
+    });
+    it('should show error toast on network failure', async () => {
+      (global.fetch as jest.Mock).mockRejectedValueOnce(new Error('Network failed'));
+
+      const actionContext = {
+        content: 'SELECT * FROM network_failure;',
+        handleVersionChange: jest.fn(),
+        currentVersionIndex: 0,
+        isCurrentVersion: true,
+        mode: 'edit' as 'edit' | 'diff',
+        metadata: {},
+        setMetadata: mockSetMetadata
+      };
+
+      await runQueryAction.onClick(actionContext);
+
+      expect(mockToast.info).toHaveBeenCalledWith('Running query...');
+      expect(mockToast.error).toHaveBeenCalledWith('Failed to run query: Network failed');
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        1,
+        expect.objectContaining({ isRunningQuery: true, error: null })
+      );
+      expect(mockSetMetadata).toHaveBeenNthCalledWith(
+        2,
+        expect.objectContaining({ results: null, error: 'Network failed', isRunningQuery: false })
+      );
+    });
+  });
+
+  describe('View Results Action', () => {
+    const viewResultsAction = sqlArtifact.actions.find((a) => a.description === 'View Results');
+
+    if (!viewResultsAction) {
+      throw new Error('View Results action not found in sqlArtifact.actions');
+    }
+
+    // Spy on window.alert and toast
+    const alertSpy = jest.spyOn(window, 'alert').mockImplementation(() => {});
+    // mockToast is already defined in the outer scope and is the one used by the component due to jest.mock at the top
+
+    beforeEach(() => {
+      alertSpy.mockClear();
+      mockToast.info.mockClear();
+      mockToast.success.mockClear();
+      mockToast.error.mockClear();
+      mockSetMetadata.mockClear();
+    });
+
+    afterAll(() => {
+      alertSpy.mockRestore();
+    });
+
+    it('should be disabled if query is running', () => {
+      const actionContext = {
+        metadata: { isRunningQuery: true }
+        // other context properties are not relevant for isDisabled here
+      } as any; // Cast to any to simplify context for isDisabled
+      expect(viewResultsAction.isDisabled?.(actionContext)).toBe(true);
+    });
+
+    it('should be enabled if query is not running', () => {
+      const actionContext = {
+        metadata: { isRunningQuery: false }
+      } as any;
+      expect(viewResultsAction.isDisabled?.(actionContext)).toBe(false);
+    });
+
+    it('should show info toast if query is running when onClick is called', () => {
+      const actionContext = {
+        metadata: { isRunningQuery: true }
+      } as any;
+      viewResultsAction.onClick(actionContext);
+      expect(mockToast.info).toHaveBeenCalledWith('Query is currently running.');
+      expect(alertSpy).not.toHaveBeenCalled();
+    });
+
+    it('should show results via alert and toast if results are present', () => {
+      const mockResults = [{ id: 1, data: 'some data' }];
+      const actionContext = {
+        metadata: { results: mockResults, isRunningQuery: false }
+      } as any;
+      viewResultsAction.onClick(actionContext);
+      expect(mockToast.success).toHaveBeenCalledWith('Displaying results (see alert/console).');
+      expect(alertSpy).toHaveBeenCalledWith(`Results:\n${JSON.stringify(mockResults, null, 2)}`);
+    });
+
+    it('should show error toast if error is present in metadata', () => {
+      const mockError = 'Failed query';
+      const actionContext = {
+        metadata: { error: mockError, isRunningQuery: false }
+      } as any;
+      viewResultsAction.onClick(actionContext);
+      expect(mockToast.error).toHaveBeenCalledWith(`Error from previous query run: ${mockError}`);
+      expect(alertSpy).not.toHaveBeenCalled();
+    });
+
+    it('should show info toast if no results or error are present', () => {
+      const actionContext = {
+        metadata: { isRunningQuery: false } // No results, no error
+      } as any;
+      viewResultsAction.onClick(actionContext);
+      expect(mockToast.info).toHaveBeenCalledWith('No results to display. Run a query first.');
+      expect(alertSpy).not.toHaveBeenCalled();
+    });
+  });
+});
+
+// Basic test for the artifact definition itself
+describe('SqlArtifact Definition', () => {
+  it('should have the correct kind and description', () => {
+    expect(sqlArtifact.kind).toBe('sql');
+    expect(sqlArtifact.description).toBe('Useful for SQL queries, allowing execution and viewing results.');
+  });
+
+  it('should have actions defined', () => {
+    expect(sqlArtifact.actions).toBeInstanceOf(Array);
+    expect(sqlArtifact.actions.length).toBeGreaterThan(0);
+    // Check for specific actions by description
+    expect(sqlArtifact.actions.find((action) => action.description === 'Run Query')).toBeDefined();
+    expect(sqlArtifact.actions.find((action) => action.description === 'View Results')).toBeDefined();
+  });
+});
diff --git a/apps/dbagent/src/components/chat/artifacts/sql/client.tsx b/apps/dbagent/src/components/chat/artifacts/sql/client.tsx
new file mode 100644
index 00000000..08137dfc
--- /dev/null
+++ b/apps/dbagent/src/components/chat/artifacts/sql/client.tsx
@@ -0,0 +1,128 @@
+import { toast } from '@xata.io/components';
+import { Artifact } from '../create-artifact';
+import { DocumentSkeleton } from '../document-skeleton';
+
+interface SqlArtifactMetadata {
+  results?: any[]; // Or a more specific type for query results
+  error?: string | null;
+  // Potentially add a loading state for when the query is running
+  isRunningQuery?: boolean;
+}
+
+export const sqlArtifact = new Artifact<'sql', SqlArtifactMetadata>({
+  kind: 'sql',
+  description: 'Useful for SQL queries, allowing execution and viewing results.',
+  initialize: async ({ documentId, setMetadata }) => {
+    // TODO: Initialization logic if needed
+  },
+  onStreamPart: ({ streamPart, setMetadata, setArtifact }) => {
+    // TODO: Handle stream parts if the SQL query can be streamed
+    if (streamPart.type === 'sql-delta') {
+      // Or a more generic type if applicable
+      setArtifact((draftArtifact) => {
+        return {
+          ...draftArtifact,
+          content: draftArtifact.content + (streamPart.content as string),
+          isVisible:
+            draftArtifact.status === 'streaming' &&
+            draftArtifact.content.length > 50 && // Adjust visibility condition as needed
+            draftArtifact.content.length < 100
+              ? true
+              : draftArtifact.isVisible,
+          status: 'streaming'
+        };
+      });
+    }
+  },
+  content: ({
+    mode,
+    status,
+    content,
+    isCurrentVersion,
+    currentVersionIndex,
+    onSaveContent,
+    getDocumentContentById,
+    isLoading,
+    metadata
+  }) => {
+    if (isLoading) {
+      return ;
+    }
+
+    // TODO: Implement diff view if necessary
+    // if (mode === 'diff') {
+    //   const oldContent = getDocumentContentById(currentVersionIndex - 1);
+    //   const newContent = getDocumentContentById(currentVersionIndex);
+    //   return ;
+    // }
+
+    return (
+      
+
{content}
+
+ ); + }, + actions: [ + { + icon: ▶️, // Placeholder icon + description: 'Run Query', + onClick: async ({ content, metadata, setMetadata }) => { + if (!content.trim()) { + toast.error('Query is empty.'); + return; + } + setMetadata((prev) => ({ ...prev, isRunningQuery: true, error: null })); + try { + toast.info('Running query...'); + const response = await fetch('/api/sql', { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ query: content }) + }); + + const result = await response.json(); + + if (!response.ok) { + throw new Error(result.error || `HTTP error! status: ${response.status}`); + } + + toast.success('Query executed successfully!'); + console.log('Query results:', result.results); + setMetadata((prev) => ({ ...prev, results: result.results, error: null, isRunningQuery: false })); + } catch (error: any) { + console.error('Failed to run query:', error); + toast.error(`Failed to run query: ${error.message}`); + setMetadata((prev) => ({ ...prev, results: null, error: error.message, isRunningQuery: false })); + } + } + }, + { + icon: 📄, // Placeholder icon + description: 'View Results', + onClick: ({ metadata }) => { + if (metadata?.isRunningQuery) { + toast.info('Query is currently running.'); + return; + } + if (metadata?.results) { + // TODO: Implement a proper modal or display area for results + // For now, using alert and console.log + toast.success('Displaying results (see alert/console).'); + console.log('Viewing results:', metadata.results); + alert(`Results: +${JSON.stringify(metadata.results, null, 2)}`); + } else if (metadata?.error) { + toast.error(`Error from previous query run: ${metadata.error}`); + } else { + toast.info('No results to display. Run a query first.'); + } + }, + isDisabled: ({ metadata }) => !!metadata?.isRunningQuery + } + ], + toolbar: [ + // TODO: Add toolbar actions if needed + ] +}); diff --git a/apps/dbagent/src/components/chat/artifacts/sql/server.ts b/apps/dbagent/src/components/chat/artifacts/sql/server.ts new file mode 100644 index 00000000..0ef8ac62 --- /dev/null +++ b/apps/dbagent/src/components/chat/artifacts/sql/server.ts @@ -0,0 +1,43 @@ +import { createDocumentHandler } from '../server'; + +export const sqlDocumentHandler = createDocumentHandler<'sql'>({ + kind: 'sql', + onCreateDocument: async ({ title, dataStream }) => { + // For SQL, the initial content is likely the query itself. + // We might not need to call an AI model to generate it unless specified. + // For now, we'll assume the 'title' or a dedicated field in `onCreateDocument` options + // will contain the SQL query. + const sqlQuery = title; // Or from another property if available + + // Stream the SQL query back. + // If the query is short and doesn't need streaming, this can be simplified. + dataStream.writeData({ + type: 'sql-delta', // Or a generic 'text-delta' if the client handles it + content: sqlQuery + }); + + // Return the full query as the document content. + return sqlQuery; + }, + onUpdateDocument: async ({ document, description, dataStream }) => { + // This function would be called if we want to update/modify an existing SQL query, + // possibly using an AI model based on a 'description'. + // For now, let's assume updates are direct or not yet implemented for SQL artifacts. + + // Example: if we were to stream back changes or a new query + // const updatedSqlQuery = `/* Updated based on: ${description} */ +${document.content}`; + // dataStream.writeData({ + // type: 'sql-delta', + // content: updatedSqlQuery + // }); + // return updatedSqlQuery; + + // For now, just return the existing content if no update logic is in place. + dataStream.writeData({ + type: 'sql-delta', + content: document.content + }); + return document.content; + } +});