diff --git a/app/api/generate-query/route.ts b/app/api/generate-query/route.ts index 51f289fd..740ff680 100644 --- a/app/api/generate-query/route.ts +++ b/app/api/generate-query/route.ts @@ -1,19 +1,17 @@ import { getDatabaseSchema } from '~/lib/schema'; import { NextRequest, NextResponse } from 'next/server'; -import { generateSqlQueries, detectDatabaseTypeFromPrompt, type Table } from '~/lib/.server/llm/database-source'; +import { generateSqlQueries } from '~/lib/.server/llm/database-source'; import { createScopedLogger } from '~/utils/logger'; import { z } from 'zod'; import { prisma } from '~/lib/prisma'; import { requireUserId } from '~/auth/session'; -import { getConnectionProtocol } from '@liblab/data-access/utils/connection'; -import { DataAccessor } from '@liblab/data-access/dataAccessor'; const logger = createScopedLogger('generate-sql'); const requestSchema = z.object({ prompt: z.string(), existingQuery: z.string().optional(), - dataSourceId: z.string().optional(), + dataSourceId: z.string(), suggestedDatabaseType: z.string().optional(), }); @@ -22,54 +20,19 @@ export async function POST(request: NextRequest) { try { const body = await request.json(); - const { prompt, existingQuery, dataSourceId, suggestedDatabaseType } = requestSchema.parse(body); + const { prompt, existingQuery, dataSourceId } = requestSchema.parse(body); const existingQueries = existingQuery ? [existingQuery] : []; - let schema: Table[]; - let type: string; + const schema = await getDatabaseSchema(dataSourceId, userId); - // If dataSourceId is provided, use the existing data source - if (dataSourceId) { - schema = await getDatabaseSchema(dataSourceId, userId); - - const dataSource = await prisma.dataSource.findUniqueOrThrow({ - where: { id: dataSourceId, createdById: userId }, - }); - - type = getConnectionProtocol(dataSource.connectionString); - } else { - // Use AI to determine database type from prompt - const availableTypes = DataAccessor.getAvailableDatabaseTypes(); - const detectedType = suggestedDatabaseType || (await detectDatabaseTypeFromPrompt(prompt, availableTypes)); - - if (!detectedType) { - return NextResponse.json( - { error: 'Could not determine database type from prompt. Please specify a data source.' }, - { status: 400 }, - ); - } - - type = detectedType; - - logger.warn( - `⚠ī¸ WARNING: No dataSourceId provided. Using sample schema for ${type}. Create a data source to query real data.`, - ); - - // For demonstration purposes, create a sample schema based on common patterns - // In a real implementation, you might want to ask the user to specify their schema - const sampleSchema = DataAccessor.getSampleSchema(type); - - if (!sampleSchema) { - return NextResponse.json({ error: 'Unsupported database type' }, { status: 400 }); - } - - schema = sampleSchema; - } + const dataSource = await prisma.dataSource.findUniqueOrThrow({ + where: { id: dataSourceId, createdById: userId }, + }); const queries = await generateSqlQueries({ schema, userPrompt: prompt, - databaseType: type, + connectionString: dataSource.connectionString, existingQueries, }); diff --git a/app/components/@settings/tabs/data/forms/AddDataSourceForm.tsx b/app/components/@settings/tabs/data/forms/AddDataSourceForm.tsx index 9234c4cc..d1e25d4e 100644 --- a/app/components/@settings/tabs/data/forms/AddDataSourceForm.tsx +++ b/app/components/@settings/tabs/data/forms/AddDataSourceForm.tsx @@ -1,7 +1,7 @@ import { classNames } from '~/utils/classNames'; import { useState } from 'react'; import { toast } from 'sonner'; -import { XCircle, CheckCircle, Loader2, Plug, Save } from 'lucide-react'; +import { CheckCircle, Loader2, Plug, Save, XCircle } from 'lucide-react'; import type { TestConnectionResponse } from '~/components/@settings/tabs/data/DataTab'; import { z } from 'zod'; import { BaseSelect } from '~/components/ui/Select'; @@ -181,6 +181,7 @@ export default function AddDataSourceForm({ isSubmitting, setIsSubmitting, onSuc
{ setDbType(value as DataSourceOption); diff --git a/app/components/sidebar/Menu.client.tsx b/app/components/sidebar/Menu.client.tsx index 3c9b1aa8..2fb8836c 100644 --- a/app/components/sidebar/Menu.client.tsx +++ b/app/components/sidebar/Menu.client.tsx @@ -150,6 +150,7 @@ export const Menu = () => { animate={open ? 'open' : 'closed'} variants={menuVariants} style={{ width: '340px' }} + data-testid="menu" className={classNames( 'flex selection-accent flex-col side-menu fixed top-0 h-full', 'bg-white dark:bg-gray-950 border-r border-gray-100 dark:border-gray-800/50', diff --git a/app/components/ui/IconButton.tsx b/app/components/ui/IconButton.tsx index 251d17c4..4c50feee 100644 --- a/app/components/ui/IconButton.tsx +++ b/app/components/ui/IconButton.tsx @@ -11,6 +11,7 @@ interface BaseIconButtonProps { title?: string; disabled?: boolean; onClick?: (event: React.MouseEvent) => void; + dataTestId?: string; } type IconButtonWithoutChildrenProps = { @@ -37,6 +38,7 @@ export const IconButton = forwardRef( title, onClick, children, + dataTestId, }: IconButtonProps, ref: ForwardedRef, ) => { @@ -52,6 +54,7 @@ export const IconButton = forwardRef( )} title={title} disabled={disabled} + data-testid={dataTestId} onClick={(event) => { if (disabled) { return; diff --git a/app/components/ui/Select.tsx b/app/components/ui/Select.tsx index c16aeb9d..a9125acf 100644 --- a/app/components/ui/Select.tsx +++ b/app/components/ui/Select.tsx @@ -64,6 +64,7 @@ interface SelectProps { components?: any; styles?: Partial>; controlIcon?: React.ReactNode; + dataTestId?: string; } const createDefaultStyles = (): StylesConfig => ({ @@ -207,6 +208,7 @@ export const BaseSelect = ({ components: customComponents, styles: customStyles, controlIcon, + dataTestId, }: SelectProps) => { const defaultStyles = createDefaultStyles(); const mergedStyles = { ...defaultStyles, ...customStyles }; @@ -219,7 +221,7 @@ export const BaseSelect = ({ return ( -
+
value={value} onChange={onChange} diff --git a/app/components/ui/SettingsButton.tsx b/app/components/ui/SettingsButton.tsx index c3810506..a7d5714a 100644 --- a/app/components/ui/SettingsButton.tsx +++ b/app/components/ui/SettingsButton.tsx @@ -13,6 +13,7 @@ export const SettingsButton = memo(({ onClick }: SettingsButtonProps) => { size="xl" title="Settings" className="text-[#666] hover:text-primary hover:bg-depth-3/10 transition-colors" + dataTestId="settings-button" > diff --git a/app/layout.tsx b/app/layout.tsx index 72701535..b2c057be 100644 --- a/app/layout.tsx +++ b/app/layout.tsx @@ -32,7 +32,6 @@ async function getRootData() { let user = null; let dataSources: any[] = []; - let pluginAccess = FREE_PLUGIN_ACCESS; let dataSourceTypes: any[] = []; if (session?.user) { @@ -42,14 +41,15 @@ async function getRootData() { // Get data sources for the user const userAbility = await getUserAbility(session.user.id); dataSources = await getDataSources(userAbility); + } - // Initialize plugin manager - await PluginManager.getInstance().initialize(); - pluginAccess = PluginManager.getInstance().getAccessMap(); + // Initialize plugin manager + await PluginManager.getInstance().initialize(); - // Get available data source types - dataSourceTypes = DataSourcePluginManager.getAvailableDatabaseTypes(); - } + const pluginAccess = PluginManager.getInstance().getAccessMap(); + + // Get available data source types + dataSourceTypes = DataSourcePluginManager.getAvailableDatabaseTypes(); return { user, diff --git a/app/lib/.server/llm/database-source.ts b/app/lib/.server/llm/database-source.ts index 50a190b2..e5cd6dd8 100644 --- a/app/lib/.server/llm/database-source.ts +++ b/app/lib/.server/llm/database-source.ts @@ -3,6 +3,7 @@ import { generateObject } from 'ai'; import { z } from 'zod'; import { DataAccessor } from '@liblab/data-access/dataAccessor'; import { getLlm } from './get-llm'; +import { getConnectionProtocol } from '@liblab/data-access/utils/connection'; const queryDecisionSchema = z.object({ shouldUpdateSql: z.boolean(), @@ -60,7 +61,7 @@ export interface Table { export type GenerateSqlQueriesOptions = { schema: Table[]; userPrompt: string; - databaseType: string; + connectionString: string; implementationPlan?: string; existingQueries?: string[]; }; @@ -68,19 +69,20 @@ export type GenerateSqlQueriesOptions = { export async function generateSqlQueries({ schema, userPrompt, - databaseType, + connectionString, existingQueries, }: GenerateSqlQueriesOptions): Promise { const dbSchema = formatDbSchemaForLLM(schema); // Get the appropriate accessor for this database type - const accessor = DataAccessor.getByDatabaseType(databaseType); + const accessor = DataAccessor.getAccessor(connectionString); if (!accessor) { - throw new Error(`No accessor found for database type: ${databaseType}`); + const protocol = getConnectionProtocol(connectionString); + throw new Error(`No accessor found for database type: ${protocol}`); } - const systemPrompt = accessor.generateSystemPrompt(databaseType, dbSchema, existingQueries, userPrompt); + const systemPrompt = accessor.generateSystemPrompt(accessor.label, dbSchema, existingQueries, userPrompt); try { const llm = await getLlm(); diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts index fa605321..19d459a4 100644 --- a/app/lib/.server/llm/stream-text.ts +++ b/app/lib/.server/llm/stream-text.ts @@ -136,13 +136,10 @@ ${props.summary} where: { id: currentDataSourceId, createdById: userId }, }); - const connectionDetails = new URL(dataSource.connectionString); - const type = connectionDetails.protocol.replace(':', ''); - const sqlQueries = await generateSqlQueries({ schema, userPrompt: lastUserMessage, - databaseType: type, + connectionString: dataSource.connectionString, implementationPlan, existingQueries, }); diff --git a/tests/e2e/tests/add-postgres-datasource-flow.spec.ts b/tests/e2e/tests/add-postgres-datasource-flow.spec.ts new file mode 100644 index 00000000..4cf9ef3b --- /dev/null +++ b/tests/e2e/tests/add-postgres-datasource-flow.spec.ts @@ -0,0 +1,255 @@ +import { type ConsoleMessage, type Page, test } from '@playwright/test'; + +test.describe('Add PostgreSQL Data Source Flow', () => { + test('Create PostgreSQL data source when connection string starts with postgresql', async ({ + page, + }: { + page: Page; + }) => { + test.setTimeout(60000); // 1 minute for this specific test + + // Enable browser console logging for debugging + page.on('console', (msg: ConsoleMessage) => console.log('đŸ–Ĩī¸ Browser console:', msg.text())); + page.on('pageerror', (error: Error) => console.log('đŸ–Ĩī¸ Browser error:', error.message)); + + console.log('Starting PostgreSQL data source creation test...'); + + console.log('🧭 Navigating to application...'); + await page.goto('/'); + + // Wait for the page to load + await page.waitForLoadState('networkidle'); + console.log('✅ Page loaded successfully'); + + try { + console.log('🔍 Checking for telemetry consent page...'); + + const telemetryHeading = page.locator('h1:has-text("Help us improve liblab ai")'); + await telemetryHeading.waitFor({ state: 'visible', timeout: 5000 }); + console.log('📋 Found telemetry consent page, clicking Decline...'); + + const declineButton = page.locator('button:has-text("Decline")'); + await declineButton.waitFor({ state: 'visible' }); + await declineButton.click(); + + await page.waitForLoadState('networkidle'); + console.log('✅ Declined telemetry, waiting for redirect...'); + } catch { + console.warn('ℹī¸ No telemetry consent page found, continuing...'); + } + + try { + console.log('🔍 Checking for data source connection page...'); + + const dataSourceHeading = page.locator('h1:has-text("Let\'s connect your data source")'); + await dataSourceHeading.waitFor({ state: 'visible', timeout: 5000 }); + console.log('💾 Found data source connection page, connecting to sample database...'); + + const connectButton = page.locator('button:has-text("Connect")'); + await connectButton.waitFor({ state: 'visible', timeout: 10000 }); + console.log('🔗 Found Connect button, clicking...'); + await connectButton.click(); + + await page.waitForLoadState('networkidle'); + console.log('✅ Connected to sample database, waiting for redirect...'); + } catch { + console.warn('ℹī¸ No data source connection page found, continuing...'); + } + + console.log('🔍 Looking for settings cog icon...'); + + await page.mouse.move(0, 500); + await page.mouse.down(); // Move mouse to ensure the menu is visible + + const menu = page.locator('[data-testid="menu"]'); + await menu.waitFor({ state: 'attached', timeout: 10000 }); + await menu.hover(); + + const settingsButton = page.locator('[data-testid="settings-button"]'); + await settingsButton.waitFor({ state: 'attached', timeout: 10000 }); + await settingsButton.click(); + + await page.getByRole('button', { name: 'Add Data Source' }).click(); + + // Look for the database type selector + console.log('🔍 Looking for database type selector...'); + + // help here + const dbTypeSelector = page.locator('select, [data-testid="add-data-source-select"]'); + await dbTypeSelector.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found database type selector, selecting PostgreSQL...'); + + // Click to open the dropdown and select PostgreSQL + await dbTypeSelector.click(); + + await page.waitForLoadState('domcontentloaded'); + + const postgresOption = page.locator('[id="postgres"]'); + await postgresOption.waitFor({ state: 'visible', timeout: 5000 }); + await postgresOption.click(); + console.log('✅ Selected PostgreSQL database type'); + + console.log('🔍 Looking for database name input...'); + + const dbNameInput = page.locator( + 'input[placeholder*="database name"], input[placeholder*="Database Name"], input[name*="dbName"], input[name*="name"]', + ); + await dbNameInput.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found database name input, filling "foobar"...'); + await dbNameInput.fill('foobar'); + + console.log('🔍 Looking for connection string input...'); + + const connStrInput = page.locator('input[placeholder*="postgres(ql)://username:password@host:port/database"]'); + await connStrInput.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found connection string input, filling connection string...'); + await connStrInput.fill('postgresql://user:pass@localhost:5432/db'); + + console.log('🔍 Looking for save/create button...'); + + const saveButton = page.locator('button:has-text("Create")'); + await saveButton.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found save button, clicking...'); + await saveButton.click(); + + console.log('💾 Waiting for data source creation to complete...'); + await page.waitForLoadState('networkidle'); + + // Look for success message or redirect + try { + const successMessage = page.locator('text=successfully, text=Success, text=created, text=added'); + await successMessage.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Data source created successfully!'); + } catch { + console.log('ℹī¸ No explicit success message found, but form submission completed'); + } + + console.log('🎉 PostgreSQL data source creation test completed successfully!'); + }); + + test('Create PostgreSQL data source when connection string starts with postgres', async ({ + page, + }: { + page: Page; + }) => { + test.setTimeout(60000); // 1 minute for this specific test + + // Enable browser console logging for debugging + page.on('console', (msg: ConsoleMessage) => console.log('đŸ–Ĩī¸ Browser console:', msg.text())); + page.on('pageerror', (error: Error) => console.log('đŸ–Ĩī¸ Browser error:', error.message)); + + console.log('Starting PostgreSQL data source creation test...'); + + console.log('🧭 Navigating to application...'); + await page.goto('/'); + + // Wait for the page to load + await page.waitForLoadState('networkidle'); + console.log('✅ Page loaded successfully'); + + try { + console.log('🔍 Checking for telemetry consent page...'); + + const telemetryHeading = page.locator('h1:has-text("Help us improve liblab ai")'); + await telemetryHeading.waitFor({ state: 'visible', timeout: 5000 }); + console.log('📋 Found telemetry consent page, clicking Decline...'); + + const declineButton = page.locator('button:has-text("Decline")'); + await declineButton.waitFor({ state: 'visible' }); + await declineButton.click(); + + await page.waitForLoadState('networkidle'); + console.log('✅ Declined telemetry, waiting for redirect...'); + } catch { + console.warn('ℹī¸ No telemetry consent page found, continuing...'); + } + + try { + console.log('🔍 Checking for data source connection page...'); + + const dataSourceHeading = page.locator('h1:has-text("Let\'s connect your data source")'); + await dataSourceHeading.waitFor({ state: 'visible', timeout: 5000 }); + console.log('💾 Found data source connection page, connecting to sample database...'); + + const connectButton = page.locator('button:has-text("Connect")'); + await connectButton.waitFor({ state: 'visible', timeout: 10000 }); + console.log('🔗 Found Connect button, clicking...'); + await connectButton.click(); + + await page.waitForLoadState('networkidle'); + console.log('✅ Connected to sample database, waiting for redirect...'); + } catch { + console.warn('ℹī¸ No data source connection page found, continuing...'); + } + + console.log('🔍 Looking for settings cog icon...'); + + await page.mouse.move(0, 500); + await page.mouse.down(); // Move mouse to ensure the menu is visible + + const menu = page.locator('[data-testid="menu"]'); + await menu.waitFor({ state: 'attached', timeout: 10000 }); + await menu.hover(); + + const settingsButton = page.locator('[data-testid="settings-button"]'); + await settingsButton.waitFor({ state: 'attached', timeout: 10000 }); + await settingsButton.click(); + + await page.getByRole('button', { name: 'Add Data Source' }).click(); + + // Look for the database type selector + console.log('🔍 Looking for database type selector...'); + + // help here + const dbTypeSelector = page.locator('select, [data-testid="add-data-source-select"]'); + await dbTypeSelector.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found database type selector, selecting PostgreSQL...'); + + // Click to open the dropdown and select PostgreSQL + await dbTypeSelector.click(); + + await page.waitForLoadState('domcontentloaded'); + + const postgresOption = page.locator('[id="postgres"]'); + await postgresOption.waitFor({ state: 'visible', timeout: 5000 }); + await postgresOption.click(); + console.log('✅ Selected PostgreSQL database type'); + + console.log('🔍 Looking for database name input...'); + + const dbNameInput = page.locator( + 'input[placeholder*="database name"], input[placeholder*="Database Name"], input[name*="dbName"], input[name*="name"]', + ); + await dbNameInput.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found database name input, filling "foobar"...'); + await dbNameInput.fill('baz'); + + console.log('🔍 Looking for connection string input...'); + + const connStrInput = page.locator('input[placeholder*="postgres(ql)://username:password@host:port/database"]'); + await connStrInput.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found connection string input, filling connection string...'); + await connStrInput.fill('postgres://user:pass@localhost:5432/db'); + + console.log('🔍 Looking for save/create button...'); + + const saveButton = page.locator('button:has-text("Create")'); + await saveButton.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Found save button, clicking...'); + await saveButton.click(); + + console.log('💾 Waiting for data source creation to complete...'); + await page.waitForLoadState('networkidle'); + + // Look for success message or redirect + try { + const successMessage = page.locator('text=successfully, text=Success, text=created, text=added'); + await successMessage.waitFor({ state: 'visible', timeout: 10000 }); + console.log('✅ Data source created successfully!'); + } catch { + console.log('ℹī¸ No explicit success message found, but form submission completed'); + } + + console.log('🎉 PostgreSQL data source creation test completed successfully!'); + }); +});