diff --git a/src/constants.ts b/src/constants.ts index 5c9c2c9..e6b9826 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -23,3 +23,5 @@ export const ANALYTICS_WRITE_KEY = export const SENTRY_DSN = process.env.SENTRY_DSN ?? 'https://b3564134667aa2dfeaa3992a12d9c12f@o1373725.ingest.us.sentry.io/4509328350380033'; + +export const NEON_CONSOLE_HOST = NEON_API_HOST.replace(/\/api\/v2$/, ''); diff --git a/src/tools/definitions.ts b/src/tools/definitions.ts index f5e6050..8e661c7 100644 --- a/src/tools/definitions.ts +++ b/src/tools/definitions.ts @@ -23,6 +23,8 @@ import { listOrganizationsInputSchema, listSharedProjectsInputSchema, resetFromParentInputSchema, + searchInputSchema, + fetchInputSchema, } from './toolsSchema.js'; export const NEON_TOOLS = [ @@ -608,4 +610,14 @@ export const NEON_TOOLS = [ description: 'Lists compute endpoints for a project or specific branch', inputSchema: listBranchComputesInputSchema, }, + { + name: 'search' as const, + description: `Searches across all user organizations, projects, and branches that match the query. Returns a list of objects with id, title, and url. This tool searches through all accessible resources and provides direct links to the Neon Console.`, + inputSchema: searchInputSchema, + }, + { + name: 'fetch' as const, + description: `Fetches detailed information about a specific organization, project, or branch using the ID returned by the search tool. This tool provides comprehensive information about Neon resources for detailed analysis and management.`, + inputSchema: fetchInputSchema, + }, ]; diff --git a/src/tools/handlers/connection-string.ts b/src/tools/handlers/connection-string.ts new file mode 100644 index 0000000..ec4417d --- /dev/null +++ b/src/tools/handlers/connection-string.ts @@ -0,0 +1,84 @@ +import { Api } from '@neondatabase/api-client'; +import { ToolHandlerExtraParams } from '../types.js'; +import { startSpan } from '@sentry/node'; +import { getDefaultDatabase } from '../utils.js'; +import { getDefaultBranch, getOnlyProject } from './utils.js'; + +export async function handleGetConnectionString( + { + projectId, + branchId, + computeId, + databaseName, + roleName, + }: { + projectId?: string; + branchId?: string; + computeId?: string; + databaseName?: string; + roleName?: string; + }, + neonClient: Api, + extra: ToolHandlerExtraParams, +) { + return await startSpan( + { + name: 'get_connection_string', + }, + async () => { + // If projectId is not provided, get the first project but only if there is only one project + if (!projectId) { + const project = await getOnlyProject(neonClient, extra); + projectId = project.id; + } + + if (!branchId) { + const defaultBranch = await getDefaultBranch(projectId, neonClient); + branchId = defaultBranch.id; + } + + // If databaseName is not provided, use default `neondb` or first database + let dbObject; + if (!databaseName) { + dbObject = await getDefaultDatabase( + { + projectId, + branchId, + databaseName, + }, + neonClient, + ); + databaseName = dbObject.name; + + if (!roleName) { + roleName = dbObject.owner_name; + } + } else if (!roleName) { + const { data } = await neonClient.getProjectBranchDatabase( + projectId, + branchId, + databaseName, + ); + roleName = data.database.owner_name; + } + + // Get connection URI with the provided parameters + const connectionString = await neonClient.getConnectionUri({ + projectId, + role_name: roleName, + database_name: databaseName, + branch_id: branchId, + endpoint_id: computeId, + }); + + return { + uri: connectionString.data.uri, + projectId, + branchId, + databaseName, + roleName, + computeId, + }; + }, + ); +} diff --git a/src/tools/handlers/decribe-project.ts b/src/tools/handlers/decribe-project.ts new file mode 100644 index 0000000..d786819 --- /dev/null +++ b/src/tools/handlers/decribe-project.ts @@ -0,0 +1,17 @@ +import { Api } from '@neondatabase/api-client'; + +async function handleDescribeProject( + projectId: string, + neonClient: Api, +) { + const { data: branchesData } = await neonClient.listProjectBranches({ + projectId: projectId, + }); + const { data: projectData } = await neonClient.getProject(projectId); + return { + branches: branchesData.branches, + project: projectData.project, + }; +} + +export { handleDescribeProject }; diff --git a/src/tools/handlers/describe-branch.ts b/src/tools/handlers/describe-branch.ts new file mode 100644 index 0000000..61480b9 --- /dev/null +++ b/src/tools/handlers/describe-branch.ts @@ -0,0 +1,96 @@ +import { Api, Branch } from '@neondatabase/api-client'; +import { ToolHandlerExtraParams } from '../types.js'; +import { handleGetConnectionString } from './connection-string.js'; +import { neon } from '@neondatabase/serverless'; +import { DESCRIBE_DATABASE_STATEMENTS } from '../utils.js'; +import { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import { CONSOLE_URLS, generateConsoleUrl } from './urls.js'; + +const branchInfo = (branch: Branch) => { + return `Branch Details: +Name: ${branch.name} +ID: ${branch.id} +Parent Branch: ${branch.parent_id} +Default: ${branch.default} +Protected: ${branch.protected ? 'Yes' : 'No'} + +${branch.created_by ? `Created By: ${branch.created_by.name}` : ''} +Created: ${new Date(branch.created_at).toLocaleDateString()} +Updated: ${new Date(branch.updated_at).toLocaleDateString()} + +Compute Usage: ${branch.compute_time_seconds} seconds +Written Data: ${branch.written_data_bytes} bytes +Data Transfer: ${branch.data_transfer_bytes} bytes + +Console Link: ${generateConsoleUrl(CONSOLE_URLS.PROJECT_BRANCH, { + projectId: branch.project_id, + branchId: branch.id, + })} +`; +}; + +export async function handleDescribeBranch( + { + projectId, + databaseName, + branchId, + }: { + projectId: string; + databaseName?: string; + branchId: string; + }, + neonClient: Api, + extra: ToolHandlerExtraParams, +): Promise { + const { data: branchData } = await neonClient.getProjectBranch( + projectId, + branchId, + ); + + const branch = branchData.branch; + + let response: Record[][]; + try { + const connectionString = await handleGetConnectionString( + { + projectId, + branchId: branch.id, + databaseName, + }, + neonClient, + extra, + ); + const runQuery = neon(connectionString.uri); + response = await runQuery.transaction( + DESCRIBE_DATABASE_STATEMENTS.map((sql) => runQuery.query(sql)), + ); + + return { + content: [ + { + type: 'text', + text: branchInfo(branch), + metadata: branch, + }, + { + type: 'text', + text: ['Database Structure:', JSON.stringify(response, null, 2)].join( + '\n', + ), + databasetree: response, + }, + ], + }; + } catch { + // Ignore database connection errors + } + + return { + content: [ + { + type: 'text', + text: branchInfo(branch), + }, + ], + }; +} diff --git a/src/tools/handlers/fetch.ts b/src/tools/handlers/fetch.ts new file mode 100644 index 0000000..6e77b80 --- /dev/null +++ b/src/tools/handlers/fetch.ts @@ -0,0 +1,208 @@ +import { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import { Api, MemberWithUser, ProjectListItem } from '@neondatabase/api-client'; +import { fetchInputSchema } from '../toolsSchema.js'; +import { z } from 'zod'; +import { handleDescribeProject } from './decribe-project.js'; +import { handleDescribeBranch } from './describe-branch.js'; +import { ToolHandlerExtraParams } from '../types.js'; +import { generateConsoleUrl, CONSOLE_URLS } from './urls.js'; + +type FetchProps = z.infer; + +export async function handleFetch( + { id }: FetchProps, + neonClient: Api, + extra: ToolHandlerExtraParams, +): Promise { + try { + // Parse the ID format + if (id.startsWith('org:') || id.startsWith('org-')) { + return await fetchOrganizationDetails(id.slice(4), neonClient); + } else if (id.startsWith('branch:')) { + const [projectId, branchId] = id.slice(7).split('/'); + return await fetchBranchDetails(projectId, branchId, neonClient, extra); + } else if (id.startsWith('project:')) { + return await fetchProjectDetails(id.slice(8), neonClient); + } else { + return { + isError: true, + content: [ + { + type: 'text', + text: `Invalid ID format: "${id}". Expected format: org:*, project:*, or branch:*`, + }, + ], + }; + } + } catch (error) { + return { + isError: true, + content: [ + { + type: 'text', + text: `Failed to fetch details: ${error instanceof Error ? error.message : 'Unknown error'}`, + }, + ], + }; + } +} + +async function fetchOrganizationDetails( + orgId: string, + neonClient: Api, +): Promise { + try { + const { data: orgData } = await neonClient.getOrganization(orgId); + + let members: MemberWithUser[] = []; + try { + const { data: membersData } = + await neonClient.getOrganizationMembers(orgId); + members = membersData.members || []; + } catch { + // Skip if we can't access members + } + + // Get projects count in this organization + let projects: ProjectListItem[] = []; + try { + const { data: projectsData } = await neonClient.listProjects({ + org_id: orgId, + }); + projects = projectsData.projects || []; + } catch { + // Skip if we can't access projects + } + + const details = { + organization: { + id: orgData.id, + name: orgData.name, + created_at: orgData.created_at, + updated_at: orgData.updated_at, + }, + console_url: generateConsoleUrl(CONSOLE_URLS.ORGANIZATION, { + orgId: orgData.id, + }), + }; + + return { + content: [ + { + type: 'text', + text: `**Organization Details** +**Basic Information:** +- Name: ${details.organization.name} +- ID: ${details.organization.id} +- Created: ${new Date(details.organization.created_at).toLocaleDateString()} + +**Statistics:** +${members.length > 0 ? `- Members: ${members.length}` : undefined} +- Projects: ${projects.length} +`, + metadata: { + org: orgData, + members: members, + projects: projects, + }, + }, + ], + }; + } catch (error) { + return { + isError: true, + content: [ + { + type: 'text', + text: `Failed to fetch organization details: ${error instanceof Error ? error.message : 'Unknown error'}`, + }, + ], + }; + } +} + +async function fetchProjectDetails( + projectId: string, + neonClient: Api, +): Promise { + try { + const { project, branches } = await handleDescribeProject( + projectId, + neonClient, + ); + + const defaultBranch = branches.find((branch) => branch.default); + return { + content: [ + { + type: 'text', + text: `**Project Details** + +**Basic Information:** +- Name: ${project.name} +- ID: ${project.id} +- Region: ${project.region_id} +- PostgreSQL Version: ${project.pg_version} +- Created: ${new Date(project.created_at).toLocaleDateString()} +- Last Updated: ${new Date(project.updated_at).toLocaleDateString()} + +**Statistics:** +- Branches: ${branches.length} +- Default Branch: ${defaultBranch?.name} (${defaultBranch?.id}) +`, + metadata: { + project: project, + branches: branches, + }, + }, + ], + }; + } catch (error) { + return { + isError: true, + content: [ + { + type: 'text', + text: `Failed to fetch project details: ${error instanceof Error ? error.message : 'Unknown error'}`, + }, + ], + }; + } +} + +async function fetchBranchDetails( + projectId: string, + branchId: string, + neonClient: Api, + extra: ToolHandlerExtraParams, +): Promise { + try { + const result = await handleDescribeBranch( + { + projectId, + branchId, + }, + neonClient, + extra, + ); + + return { + content: [ + { + type: 'text', + text: JSON.stringify(result, null, 2), + }, + ], + }; + } catch (error) { + return { + isError: true, + content: [ + { + type: 'text', + text: `Failed to fetch branch details: ${error instanceof Error ? error.message : 'Unknown error'}`, + }, + ], + }; + } +} diff --git a/src/tools/handlers/list-orgs.ts b/src/tools/handlers/list-orgs.ts new file mode 100644 index 0000000..a09d84f --- /dev/null +++ b/src/tools/handlers/list-orgs.ts @@ -0,0 +1,19 @@ +import { Api, Organization } from '@neondatabase/api-client'; +import { ToolHandlerExtraParams } from '../types.js'; +import { filterOrganizations } from '../utils.js'; + +export async function handleListOrganizations( + neonClient: Api, + account: ToolHandlerExtraParams['account'], + search?: string, +): Promise { + if (account.isOrg) { + const orgId = account.id; + const { data } = await neonClient.getOrganization(orgId); + return filterOrganizations([data], search); + } + + const { data: response } = await neonClient.getCurrentUserOrganizations(); + const organizations = response.organizations || []; + return filterOrganizations(organizations, search); +} diff --git a/src/tools/handlers/list-projects.ts b/src/tools/handlers/list-projects.ts new file mode 100644 index 0000000..26a09d3 --- /dev/null +++ b/src/tools/handlers/list-projects.ts @@ -0,0 +1,55 @@ +import { Api, ListProjectsParams } from '@neondatabase/api-client'; +import { ToolHandlerExtraParams } from '../types.js'; +import { getOrgByOrgIdOrDefault } from '../utils.js'; +import { handleListOrganizations } from './list-orgs.js'; + +export async function handleListProjects( + params: ListProjectsParams, + neonClient: Api, + extra: ToolHandlerExtraParams, +) { + const organization = await getOrgByOrgIdOrDefault(params, neonClient, extra); + + const response = await neonClient.listProjects({ + ...params, + org_id: organization?.id, + }); + if (response.status !== 200) { + throw new Error(`Failed to list projects: ${response.statusText}`); + } + + let projects = response.data.projects; + + // If search is provided and no org_id specified, and no projects found in personal account, + // search across all user organizations + if (params.search && !params.org_id && projects.length === 0) { + const organizations = await handleListOrganizations( + neonClient, + extra.account, + ); + + // Search projects across all organizations + const allProjects = []; + for (const org of organizations) { + // Skip the default organization + if (organization?.id === org.id) { + continue; + } + + const orgResponse = await neonClient.listProjects({ + ...params, + org_id: org.id, + }); + if (orgResponse.status === 200) { + allProjects.push(...orgResponse.data.projects); + } + } + + // If we found projects in other organizations, return them + if (allProjects.length > 0) { + projects = allProjects; + } + } + + return projects; +} diff --git a/src/tools/handlers/search.ts b/src/tools/handlers/search.ts new file mode 100644 index 0000000..ebfe9cd --- /dev/null +++ b/src/tools/handlers/search.ts @@ -0,0 +1,199 @@ +import { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import { Api, Organization, ProjectListItem } from '@neondatabase/api-client'; +import { Branch } from '@neondatabase/api-client'; +import { searchInputSchema } from '../toolsSchema.js'; +import { z } from 'zod'; +import { ToolHandlerExtraParams } from '../types.js'; +import { handleListProjects } from './list-projects.js'; +import { CONSOLE_URLS, generateConsoleUrl } from './urls.js'; + +type SearchProps = z.infer; + +type MCPOrgId = `org:${string}`; +type MCPProjectId = `project:${string}`; +type MCPBranchId = `branch:${string}/${string}`; // projectId/branchId +type SearchResultId = MCPOrgId | MCPProjectId | MCPBranchId; + +type SearchResult = { + id: SearchResultId; + title: string; + url: string; + type: 'organization' | 'project' | 'branch'; +}; + +const matches = ( + entity: Organization | ProjectListItem | Branch, + query: string, +) => { + return ( + entity.name.toLowerCase().includes(query) || + entity.id.toLowerCase().includes(query) + ); +}; + +export async function handleSearch( + { query }: SearchProps, + neonClient: Api, + extra: ToolHandlerExtraParams, +): Promise { + try { + const results: SearchResult[] = []; + const searchQuery = query.toLowerCase(); + // Search through all user's organizations + let organizations; + if (extra.account.isOrg) { + const orgId = extra.account.id; + const { data } = await neonClient.getOrganization(orgId); + organizations = [data]; + } else { + const { data: response } = await neonClient.getCurrentUserOrganizations(); + organizations = response.organizations || []; + } + + // If in personal account, search projects + if (!extra.account.isOrg) { + const projects = await handleListProjects( + { + limit: 400, + }, + neonClient, + extra, + ); + const searchResults = await searchProjectsAndBranches( + projects, + neonClient, + searchQuery, + ); + + results.push(...searchResults); + } + + // Search in all organizations + for (const org of organizations) { + // Check if organization matches the search query + if (matches(org, searchQuery)) { + results.push({ + id: `org:${org.id}`, + title: org.name, + url: generateConsoleUrl(CONSOLE_URLS.ORGANIZATION, { + orgId: org.id, + }), + type: 'organization', + }); + } + + try { + const projects = await handleListProjects( + { + org_id: org.id, + limit: 400, + }, + neonClient, + extra, + ); + + const searchResults = await searchProjectsAndBranches( + projects, + neonClient, + searchQuery, + ); + + results.push(...searchResults); + } catch { + // Skip projects if we can't access them + continue; + } + } + + // Also search shared projects + try { + const { data } = await neonClient.listSharedProjects({ + limit: 400, + }); + + const searchResults = await searchProjectsAndBranches( + data.projects, + neonClient, + searchQuery, + ); + results.push(...searchResults); + } catch { + // Skip shared projects if we can't access them + } + + return { + content: [ + { + type: 'text', + text: JSON.stringify(results, null, 2), + }, + ], + }; + } catch (error) { + return { + isError: true, + content: [ + { + type: 'text', + text: `Failed to search: ${error instanceof Error ? error.message : 'Unknown error'}`, + }, + ], + }; + } +} + +const searchProjectsAndBranches = async ( + projects: ProjectListItem[], + neonClient: Api, + query: string, +): Promise => { + const results: SearchResult[] = []; + projects.forEach((project) => { + if (matches(project, query)) { + results.push({ + id: `project:${project.id}`, + title: project.name, + url: generateConsoleUrl(CONSOLE_URLS.PROJECT, { + projectId: project.id, + }), + type: 'project', + }); + } + }); + + const branches = await Promise.all( + projects.map(async (project) => { + return searchBranches(project.id, neonClient, query); + }), + ); + results.push(...branches.flat()); + return results; +}; + +const searchBranches = async ( + projectId: string, + neonClient: Api, + query: string, +): Promise => { + try { + const { data } = await neonClient.listProjectBranches({ + projectId, + }); + const branches = data.branches; + + return branches + .filter((branch) => matches(branch, query)) + .map((branch) => ({ + id: `branch:${projectId}/${branch.id}`, + title: branch.name, + url: generateConsoleUrl(CONSOLE_URLS.PROJECT_BRANCH, { + projectId, + branchId: branch.id, + }), + type: 'branch', + })); + } catch { + // Ignore if we fail to fetch branches + return []; + } +}; diff --git a/src/tools/handlers/urls.ts b/src/tools/handlers/urls.ts new file mode 100644 index 0000000..f73f3c5 --- /dev/null +++ b/src/tools/handlers/urls.ts @@ -0,0 +1,28 @@ +import { NEON_CONSOLE_HOST } from '../../constants.js'; +import { NotFoundError } from '../../server/errors.js'; + +export enum CONSOLE_URLS { + ORGANIZATION = '/app/:orgId/projects', + PROJECT = '/app/projects/:projectId', + PROJECT_BRANCH = '/app/projects/:projectId/branches/:branchId', +} + +type ExtractPathParams = + T extends `${string}:${infer Param}/${infer Rest}` + ? { [k in Param | keyof ExtractPathParams<`/${Rest}`>]: string | number } + : T extends `${string}:${infer Param}` + ? Record + : Record; + +export function generateConsoleUrl( + url: T, + params: ExtractPathParams, +): string { + const link = url.replace(/:([a-zA-Z0-9_]+)/g, (_, key) => { + if ((params as any)[key] === undefined) { + throw new NotFoundError(`Missing parameter '${key}' for url '${url}'`); + } + return encodeURIComponent(String((params as any)[key])); + }); + return new URL(link, NEON_CONSOLE_HOST).toString(); +} diff --git a/src/tools/handlers/utils.ts b/src/tools/handlers/utils.ts new file mode 100644 index 0000000..b42440b --- /dev/null +++ b/src/tools/handlers/utils.ts @@ -0,0 +1,33 @@ +import { Api } from '@neondatabase/api-client'; +import { handleListProjects } from './list-projects.js'; +import { ToolHandlerExtraParams } from '../types.js'; +import { NotFoundError } from '../../server/errors.js'; + +export async function getOnlyProject( + neonClient: Api, + extra: ToolHandlerExtraParams, +) { + const projects = await handleListProjects({}, neonClient, extra); + if (projects.length === 1) { + return projects[0]; + } + throw new NotFoundError( + 'Please provide a project ID or ensure you have only one project in your account.', + ); +} + +export const getDefaultBranch = async ( + projectId: string, + neonClient: Api, +) => { + const branches = await neonClient.listProjectBranches({ + projectId, + }); + const defaultBranch = branches.data.branches.find((branch) => branch.default); + if (defaultBranch) { + return defaultBranch; + } + throw new NotFoundError( + 'No default branch found in this project. Please provide a branch ID.', + ); +}; diff --git a/src/tools/tools.ts b/src/tools/tools.ts index 864b441..4aa7bf8 100644 --- a/src/tools/tools.ts +++ b/src/tools/tools.ts @@ -2,9 +2,7 @@ import { Api, Branch, EndpointType, - ListProjectsParams, ListSharedProjectsParams, - Organization, ProjectCreateRequest, } from '@neondatabase/api-client'; import { neon } from '@neondatabase/serverless'; @@ -13,69 +11,23 @@ import { InvalidArgumentError, NotFoundError } from '../server/errors.js'; import { describeTable, formatTableDescription } from '../describeUtils.js'; import { handleProvisionNeonAuth } from './handlers/neon-auth.js'; +import { handleSearch } from './handlers/search.js'; +import { handleFetch } from './handlers/fetch.js'; import { getMigrationFromMemory, persistMigrationToMemory } from './state.js'; import { - DESCRIBE_DATABASE_STATEMENTS, getDefaultDatabase, splitSqlStatements, getOrgByOrgIdOrDefault, - filterOrganizations, resolveBranchId, } from './utils.js'; import { startSpan } from '@sentry/node'; import { ToolHandlerExtraParams, ToolHandlers } from './types.js'; - -async function handleListProjects( - params: ListProjectsParams, - neonClient: Api, - extra: ToolHandlerExtraParams, -) { - const organization = await getOrgByOrgIdOrDefault(params, neonClient, extra); - - const response = await neonClient.listProjects({ - ...params, - org_id: organization?.id, - }); - if (response.status !== 200) { - throw new Error(`Failed to list projects: ${response.statusText}`); - } - - let projects = response.data.projects; - - // If search is provided and no org_id specified, and no projects found in personal account, - // search across all user organizations - if (params.search && !params.org_id && projects.length === 0) { - const organizations = await handleListOrganizations( - neonClient, - extra.account, - ); - - // Search projects across all organizations - const allProjects = []; - for (const org of organizations) { - // Skip the default organization - if (organization?.id === org.id) { - continue; - } - - const orgResponse = await neonClient.listProjects({ - ...params, - org_id: org.id, - }); - if (orgResponse.status === 200) { - allProjects.push(...orgResponse.data.projects); - } - } - - // If we found projects in other organizations, return them - if (allProjects.length > 0) { - projects = allProjects; - } - } - - return projects; -} +import { handleListOrganizations } from './handlers/list-orgs.js'; +import { handleListProjects } from './handlers/list-projects.js'; +import { handleDescribeProject } from './handlers/decribe-project.js'; +import { handleGetConnectionString } from './handlers/connection-string.js'; +import { handleDescribeBranch } from './handlers/describe-branch.js'; async function handleCreateProject( params: ProjectCreateRequest, @@ -99,28 +51,6 @@ async function handleDeleteProject( return response.data; } -async function handleDescribeProject( - projectId: string, - neonClient: Api, -) { - const projectBranches = await neonClient.listProjectBranches({ - projectId: projectId, - }); - const projectDetails = await neonClient.getProject(projectId); - if (projectBranches.status !== 200) { - throw new Error( - `Failed to get project branches: ${projectBranches.statusText}`, - ); - } - if (projectDetails.status !== 200) { - throw new Error(`Failed to get project: ${projectDetails.statusText}`); - } - return { - branches: projectBranches.data, - project: projectDetails.data, - }; -} - async function handleRunSql( { sql, @@ -373,102 +303,6 @@ async function handleResetFromParent( }; } -async function handleGetConnectionString( - { - projectId, - branchId, - computeId, - databaseName, - roleName, - }: { - projectId?: string; - branchId?: string; - computeId?: string; - databaseName?: string; - roleName?: string; - }, - neonClient: Api, - extra: ToolHandlerExtraParams, -) { - return await startSpan( - { - name: 'get_connection_string', - }, - async () => { - // If projectId is not provided, get the first project but only if there is only one project - if (!projectId) { - const projects = await handleListProjects({}, neonClient, extra); - if (projects.length === 1) { - projectId = projects[0].id; - } else { - throw new NotFoundError( - 'Please provide a project ID or ensure you have only one project in your account.', - ); - } - } - - if (!branchId) { - const branches = await neonClient.listProjectBranches({ - projectId, - }); - const defaultBranch = branches.data.branches.find( - (branch) => branch.default, - ); - if (defaultBranch) { - branchId = defaultBranch.id; - } else { - throw new NotFoundError( - 'No default branch found in this project. Please provide a branch ID.', - ); - } - } - - // If databaseName is not provided, use default `neondb` or first database - let dbObject; - if (!databaseName) { - dbObject = await getDefaultDatabase( - { - projectId, - branchId, - databaseName, - }, - neonClient, - ); - databaseName = dbObject.name; - - if (!roleName) { - roleName = dbObject.owner_name; - } - } else if (!roleName) { - const { data } = await neonClient.getProjectBranchDatabase( - projectId, - branchId, - databaseName, - ); - roleName = data.database.owner_name; - } - - // Get connection URI with the provided parameters - const connectionString = await neonClient.getConnectionUri({ - projectId, - role_name: roleName, - database_name: databaseName, - branch_id: branchId, - endpoint_id: computeId, - }); - - return { - uri: connectionString.data.uri, - projectId, - branchId, - databaseName, - roleName, - computeId, - }; - }, - ); -} - async function handleSchemaMigration( { migrationSql, @@ -570,36 +404,6 @@ async function handleCommitMigration( }); } -async function handleDescribeBranch( - { - projectId, - databaseName, - branchId, - }: { - projectId: string; - databaseName?: string; - branchId?: string; - }, - neonClient: Api, - extra: ToolHandlerExtraParams, -) { - const connectionString = await handleGetConnectionString( - { - projectId, - branchId, - databaseName, - }, - neonClient, - extra, - ); - const runQuery = neon(connectionString.uri); - const response = await runQuery.transaction( - DESCRIBE_DATABASE_STATEMENTS.map((sql) => runQuery.query(sql)), - ); - - return response; -} - async function handleExplainSqlStatement( { params, @@ -1188,22 +992,6 @@ async function handleListBranchComputes( })); } -async function handleListOrganizations( - neonClient: Api, - account: ToolHandlerExtraParams['account'], - search?: string, -): Promise { - if (account.isOrg) { - const orgId = account.id; - const { data } = await neonClient.getOrganization(orgId); - return filterOrganizations([data], search); - } - - const { data: response } = await neonClient.getCurrentUserOrganizations(); - const organizations = response.organizations || []; - return filterOrganizations(organizations, search); -} - async function handleListSharedProjects( params: ListSharedProjectsParams, neonClient: Api, @@ -1334,7 +1122,7 @@ export const NEON_HANDLERS = { content: [ { type: 'text', - text: `This project is called ${result.project.project.name}.`, + text: `This project is called ${result.project.name}.`, }, { type: 'text', @@ -1503,7 +1291,7 @@ export const NEON_HANDLERS = { }, describe_branch: async ({ params }, neonClient, extra) => { - const result = await handleDescribeBranch( + return await handleDescribeBranch( { projectId: params.projectId, branchId: params.branchId, @@ -1512,16 +1300,6 @@ export const NEON_HANDLERS = { neonClient, extra, ); - return { - content: [ - { - type: 'text', - text: ['Database Structure:', JSON.stringify(result, null, 2)].join( - '\n', - ), - }, - ], - }; }, delete_branch: async ({ params }, neonClient) => { @@ -1774,4 +1552,12 @@ export const NEON_HANDLERS = { ], }; }, + + search: async ({ params }, neonClient, extra) => { + return await handleSearch(params, neonClient, extra); + }, + + fetch: async ({ params }, neonClient, extra) => { + return await handleFetch(params, neonClient, extra); + }, } satisfies ToolHandlers; diff --git a/src/tools/toolsSchema.ts b/src/tools/toolsSchema.ts index 855aaf0..9efa432 100644 --- a/src/tools/toolsSchema.ts +++ b/src/tools/toolsSchema.ts @@ -334,3 +334,21 @@ export const resetFromParentInputSchema = z.object({ 'Optional name to preserve the current state under a new branch before resetting', ), }); + +export const searchInputSchema = z.object({ + query: z + .string() + .min(3) + .describe( + 'The search query to find matching organizations, projects, or branches', + ), +}); + +export const fetchInputSchema = z.object({ + id: z + .string() + .min(1) + .describe( + 'The ID returned by the search tool to fetch detailed information about the entity', + ), +});