This is page 2 of 2. Use http://codebase.md/neondatabase-labs/mcp-server-neon?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .bun-version
├── .dockerignore
├── .env.example
├── .github
│ └── workflows
│ ├── koyeb-preview.yml
│ ├── koyeb-prod.yml
│ └── pr.yml
├── .gitignore
├── .npmrc
├── .nvmrc
├── .prettierignore
├── .prettierrc
├── bun.lock
├── CHANGELOG.md
├── Dockerfile
├── eslint.config.js
├── landing
│ ├── .gitignore
│ ├── app
│ │ ├── globals.css
│ │ ├── layout.tsx
│ │ └── page.tsx
│ ├── components
│ │ ├── CodeSnippet.tsx
│ │ ├── CopyableUrl.tsx
│ │ ├── DescriptionItem.tsx
│ │ ├── ExternalIcon.tsx
│ │ ├── ExternalLink.tsx
│ │ ├── Header.tsx
│ │ ├── Introduction.tsx
│ │ ├── ThemeProvider.tsx
│ │ └── ui
│ │ ├── accordion.tsx
│ │ ├── alert.tsx
│ │ └── button.tsx
│ ├── components.json
│ ├── eslint.config.mjs
│ ├── icons
│ │ ├── github.svg
│ │ └── neon.svg
│ ├── lib
│ │ ├── description.ts
│ │ └── utils.ts
│ ├── next.config.ts
│ ├── package-lock.json
│ ├── package.json
│ ├── postcss.config.mjs
│ ├── public
│ │ └── favicon.ico
│ ├── README.md
│ └── tsconfig.json
├── LICENSE
├── mcp-client
│ ├── .env.example
│ ├── package-lock.json
│ ├── package.json
│ ├── README.md
│ ├── src
│ │ ├── bin.ts
│ │ ├── cli-client.ts
│ │ ├── index.ts
│ │ ├── logger.ts
│ │ └── neon-cli-client.ts
│ └── tsconfig.json
├── package-lock.json
├── package.json
├── public
│ └── logo.png
├── PUBLISH.md
├── README.md
├── remote.Dockerfile
├── scripts
│ └── before-publish.ts
├── src
│ ├── analytics
│ │ └── analytics.ts
│ ├── constants.ts
│ ├── describeUtils.ts
│ ├── index.ts
│ ├── initConfig.ts
│ ├── oauth
│ │ ├── client.ts
│ │ ├── cookies.ts
│ │ ├── kv-store.ts
│ │ ├── model.ts
│ │ ├── server.ts
│ │ └── utils.ts
│ ├── resources.ts
│ ├── sentry
│ │ ├── instrument.ts
│ │ └── utils.ts
│ ├── server
│ │ ├── api.ts
│ │ ├── errors.ts
│ │ └── index.ts
│ ├── tools
│ │ ├── definitions.ts
│ │ ├── handlers
│ │ │ └── neon-auth.ts
│ │ ├── index.ts
│ │ ├── state.ts
│ │ ├── tools.ts
│ │ ├── toolsSchema.ts
│ │ ├── types.ts
│ │ └── utils.ts
│ ├── tools-evaluations
│ │ ├── evalUtils.ts
│ │ └── prepare-database-migration.eval.ts
│ ├── transports
│ │ ├── sse-express.ts
│ │ ├── stdio.ts
│ │ └── stream.ts
│ ├── types
│ │ ├── auth.ts
│ │ ├── context.ts
│ │ └── express.d.ts
│ ├── utils
│ │ ├── logger.ts
│ │ └── polyfills.ts
│ └── views
│ ├── approval-dialog.pug
│ └── styles.css
├── tsconfig.json
└── tsconfig.test.json
```
# Files
--------------------------------------------------------------------------------
/mcp-client/src/index.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { Anthropic } from '@anthropic-ai/sdk';
2 |
3 | import {
4 | StdioClientTransport,
5 | StdioServerParameters,
6 | } from '@modelcontextprotocol/sdk/client/stdio.js';
7 | import {
8 | ListToolsResultSchema,
9 | CallToolResultSchema,
10 | } from '@modelcontextprotocol/sdk/types.js';
11 | import { Client } from '@modelcontextprotocol/sdk/client/index.js';
12 | import chalk from 'chalk';
13 | import { Tool } from '@anthropic-ai/sdk/resources/index.mjs';
14 | import { Stream } from '@anthropic-ai/sdk/streaming.mjs';
15 | import { consoleStyles, Logger, LoggerOptions } from './logger.js';
16 |
17 | interface Message {
18 | role: 'user' | 'assistant';
19 | content: string;
20 | }
21 |
22 | type MCPClientOptions = StdioServerParameters & {
23 | loggerOptions?: LoggerOptions;
24 | };
25 |
26 | export class MCPClient {
27 | private anthropicClient: Anthropic;
28 | private messages: Message[] = [];
29 | private mcpClient: Client;
30 | private transport: StdioClientTransport;
31 | private tools: Tool[] = [];
32 | private logger: Logger;
33 |
34 | constructor({ loggerOptions, ...serverConfig }: MCPClientOptions) {
35 | this.anthropicClient = new Anthropic({
36 | apiKey: process.env.ANTHROPIC_API_KEY,
37 | });
38 |
39 | this.mcpClient = new Client(
40 | { name: 'cli-client', version: '1.0.0' },
41 | { capabilities: {} },
42 | );
43 |
44 | this.transport = new StdioClientTransport(serverConfig);
45 | this.logger = new Logger(loggerOptions ?? { mode: 'verbose' });
46 | }
47 |
48 | async start() {
49 | try {
50 | await this.mcpClient.connect(this.transport);
51 | await this.initMCPTools();
52 | } catch (error) {
53 | this.logger.log('Failed to initialize MCP Client: ' + error + '\n', {
54 | type: 'error',
55 | });
56 | process.exit(1);
57 | }
58 | }
59 |
60 | async stop() {
61 | await this.mcpClient.close();
62 | }
63 |
64 | private async initMCPTools() {
65 | const toolsResults = await this.mcpClient.request(
66 | { method: 'tools/list' },
67 | ListToolsResultSchema,
68 | );
69 | this.tools = toolsResults.tools.map(({ inputSchema, ...tool }) => ({
70 | ...tool,
71 | input_schema: inputSchema,
72 | }));
73 | }
74 |
75 | private formatToolCall(toolName: string, args: any): string {
76 | return (
77 | '\n' +
78 | consoleStyles.tool.bracket('[') +
79 | consoleStyles.tool.name(toolName) +
80 | consoleStyles.tool.bracket('] ') +
81 | consoleStyles.tool.args(JSON.stringify(args, null, 2)) +
82 | '\n'
83 | );
84 | }
85 |
86 | private formatJSON(json: string): string {
87 | return json
88 | .replace(/"([^"]+)":/g, chalk.blue('"$1":'))
89 | .replace(/: "([^"]+)"/g, ': ' + chalk.green('"$1"'));
90 | }
91 |
92 | private async processStream(
93 | stream: Stream<Anthropic.Messages.RawMessageStreamEvent>,
94 | ): Promise<void> {
95 | let currentMessage = '';
96 | let currentToolName = '';
97 | let currentToolInputString = '';
98 |
99 | this.logger.log(consoleStyles.assistant);
100 | for await (const chunk of stream) {
101 | switch (chunk.type) {
102 | case 'message_start':
103 | case 'content_block_stop':
104 | continue;
105 |
106 | case 'content_block_start':
107 | if (chunk.content_block?.type === 'tool_use') {
108 | currentToolName = chunk.content_block.name;
109 | }
110 | break;
111 |
112 | case 'content_block_delta':
113 | if (chunk.delta.type === 'text_delta') {
114 | this.logger.log(chunk.delta.text);
115 | currentMessage += chunk.delta.text;
116 | } else if (chunk.delta.type === 'input_json_delta') {
117 | if (currentToolName && chunk.delta.partial_json) {
118 | currentToolInputString += chunk.delta.partial_json;
119 | }
120 | }
121 | break;
122 |
123 | case 'message_delta':
124 | if (currentMessage) {
125 | this.messages.push({
126 | role: 'assistant',
127 | content: currentMessage,
128 | });
129 | }
130 |
131 | if (chunk.delta.stop_reason === 'tool_use') {
132 | const toolArgs = currentToolInputString
133 | ? JSON.parse(currentToolInputString)
134 | : {};
135 |
136 | this.logger.log(
137 | this.formatToolCall(currentToolName, toolArgs) + '\n',
138 | );
139 | const toolResult = await this.mcpClient.request(
140 | {
141 | method: 'tools/call',
142 | params: {
143 | name: currentToolName,
144 | arguments: toolArgs,
145 | },
146 | },
147 | CallToolResultSchema,
148 | );
149 |
150 | const formattedResult = this.formatJSON(
151 | JSON.stringify(toolResult.content.flatMap((c) => c.text)),
152 | );
153 |
154 | this.messages.push({
155 | role: 'user',
156 | content: formattedResult,
157 | });
158 |
159 | const nextStream = await this.anthropicClient.messages.create({
160 | messages: this.messages,
161 | model: 'claude-3-5-sonnet-20241022',
162 | max_tokens: 8192,
163 | tools: this.tools,
164 | stream: true,
165 | });
166 | await this.processStream(nextStream);
167 | }
168 | break;
169 |
170 | case 'message_stop':
171 | break;
172 |
173 | default:
174 | this.logger.log(`Unknown event type: ${JSON.stringify(chunk)}\n`, {
175 | type: 'warning',
176 | });
177 | }
178 | }
179 | }
180 |
181 | async processQuery(query: string) {
182 | try {
183 | this.messages.push({ role: 'user', content: query });
184 |
185 | const stream = await this.anthropicClient.messages.create({
186 | messages: this.messages,
187 | model: 'claude-3-5-sonnet-20241022',
188 | max_tokens: 8192,
189 | tools: this.tools,
190 | stream: true,
191 | });
192 | await this.processStream(stream);
193 |
194 | return this.messages;
195 | } catch (error) {
196 | this.logger.log('\nError during query processing: ' + error + '\n', {
197 | type: 'error',
198 | });
199 | if (error instanceof Error) {
200 | this.logger.log(
201 | consoleStyles.assistant +
202 | 'I apologize, but I encountered an error: ' +
203 | error.message +
204 | '\n',
205 | );
206 | }
207 | }
208 | }
209 | }
210 |
```
--------------------------------------------------------------------------------
/src/oauth/utils.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { NextFunction, Request, Response } from 'express';
2 | import cors from 'cors';
3 | import crypto from 'crypto';
4 | import { model } from './model.js';
5 | import { ApiKeyRecord, apiKeys } from './kv-store.js';
6 | import { createNeonClient } from '../server/api.js';
7 | import { identify } from '../analytics/analytics.js';
8 |
9 | export const ensureCorsHeaders = () =>
10 | cors({
11 | origin: true,
12 | methods: '*',
13 | allowedHeaders: 'Authorization, Origin, Content-Type, Accept, *',
14 | });
15 |
16 | const fetchAccountDetails = async (
17 | accessToken: string,
18 | ): Promise<ApiKeyRecord | null> => {
19 | const apiKeyRecord = await apiKeys.get(accessToken);
20 | if (apiKeyRecord) {
21 | return apiKeyRecord;
22 | }
23 |
24 | try {
25 | const neonClient = createNeonClient(accessToken);
26 | const { data: auth } = await neonClient.getAuthDetails();
27 | if (auth.auth_method === 'api_key_org') {
28 | const { data: org } = await neonClient.getOrganization(auth.account_id);
29 | const record = {
30 | apiKey: accessToken,
31 | authMethod: auth.auth_method,
32 | account: {
33 | id: auth.account_id,
34 | name: org.name,
35 | isOrg: true,
36 | },
37 | };
38 | identify(record.account, { context: { authMethod: record.authMethod } });
39 | await apiKeys.set(accessToken, record);
40 | return record;
41 | }
42 | const { data: user } = await neonClient.getCurrentUserInfo();
43 | const record = {
44 | apiKey: accessToken,
45 | authMethod: auth.auth_method,
46 | account: {
47 | id: user.id,
48 | name: user.name,
49 | email: user.email,
50 | isOrg: false,
51 | },
52 | };
53 | identify(record.account, { context: { authMethod: record.authMethod } });
54 | await apiKeys.set(accessToken, record);
55 | return record;
56 | } catch {
57 | return null;
58 | }
59 | };
60 |
61 | export const requiresAuth =
62 | () => async (request: Request, response: Response, next: NextFunction) => {
63 | const authorization = request.headers.authorization;
64 | if (!authorization) {
65 | response.status(401).json({ error: 'Unauthorized' });
66 | return;
67 | }
68 |
69 | const accessToken = extractBearerToken(authorization);
70 | const token = await model.getAccessToken(accessToken);
71 | if (token) {
72 | if (!token.expires_at || token.expires_at < Date.now()) {
73 | response.status(401).json({ error: 'Access token expired' });
74 | return;
75 | }
76 |
77 | request.auth = {
78 | token: token.accessToken,
79 | clientId: token.client.id,
80 | scopes: Array.isArray(token.scope)
81 | ? token.scope
82 | : (token.scope?.split(' ') ?? []),
83 | extra: {
84 | account: {
85 | id: token.user.id,
86 | name: token.user.name,
87 | email: token.user.email,
88 | isOrg: false,
89 | },
90 | client: {
91 | id: token.client.id,
92 | name: token.client.client_name,
93 | },
94 | },
95 | };
96 |
97 | next();
98 | return;
99 | }
100 |
101 | // If the token is not found, try to resolve the auth headers with Neon for other means of authentication.
102 | const apiKeyRecord = await fetchAccountDetails(accessToken);
103 | if (!apiKeyRecord) {
104 | response.status(401).json({ error: 'Invalid access token' });
105 | return;
106 | }
107 | request.auth = {
108 | token: accessToken,
109 | clientId: 'api-key',
110 | scopes: ['*'],
111 | extra: {
112 | account: apiKeyRecord.account,
113 | },
114 | };
115 | next();
116 | return;
117 | };
118 |
119 | export type DownstreamAuthRequest = {
120 | responseType: string;
121 | clientId: string;
122 | redirectUri: string;
123 | scope: string[];
124 | state: string;
125 | codeChallenge?: string;
126 | codeChallengeMethod?: string;
127 | };
128 |
129 | export const parseAuthRequest = (request: Request): DownstreamAuthRequest => {
130 | const responseType = (request.query.response_type || '') as string;
131 | const clientId = (request.query.client_id || '') as string;
132 | const redirectUri = (request.query.redirect_uri || '') as string;
133 | const scope = (request.query.scope || '') as string;
134 | const state = (request.query.state || '') as string;
135 | const codeChallenge = (request.query.code_challenge as string) || undefined;
136 | const codeChallengeMethod = (request.query.code_challenge_method ||
137 | 'plain') as string;
138 |
139 | return {
140 | responseType,
141 | clientId,
142 | redirectUri,
143 | scope: scope.split(' ').filter(Boolean),
144 | state,
145 | codeChallenge,
146 | codeChallengeMethod,
147 | };
148 | };
149 |
150 | export const decodeAuthParams = (state: string): DownstreamAuthRequest => {
151 | const decoded = atob(state);
152 | return JSON.parse(decoded);
153 | };
154 |
155 | export const generateRandomString = (length: number): string => {
156 | const charset =
157 | 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
158 | const array = new Uint8Array(length);
159 | crypto.getRandomValues(array);
160 | return Array.from(array, (byte) => charset[byte % charset.length]).join('');
161 | };
162 |
163 | export const extractBearerToken = (authorizationHeader: string): string => {
164 | if (!authorizationHeader) return '';
165 | return authorizationHeader.replace(/^Bearer\s+/i, '');
166 | };
167 |
168 | export const extractClientCredentials = (request: Request) => {
169 | const authorization = request.headers.authorization;
170 | if (authorization?.startsWith('Basic ')) {
171 | const credentials = atob(authorization.replace(/^Basic\s+/i, ''));
172 | const [clientId, clientSecret] = credentials.split(':');
173 | return { clientId, clientSecret };
174 | }
175 |
176 | return {
177 | clientId: request.body.client_id,
178 | clientSecret: request.body.client_secret,
179 | };
180 | };
181 |
182 | export const toSeconds = (ms: number): number => {
183 | return Math.floor(ms / 1000);
184 | };
185 |
186 | export const toMilliseconds = (seconds: number): number => {
187 | return seconds * 1000;
188 | };
189 |
190 | export const verifyPKCE = (
191 | codeChallenge: string,
192 | codeChallengeMethod: string,
193 | codeVerifier: string,
194 | ): boolean => {
195 | if (!codeChallenge || !codeChallengeMethod || !codeVerifier) {
196 | return false;
197 | }
198 |
199 | if (codeChallengeMethod === 'S256') {
200 | const hash = crypto
201 | .createHash('sha256')
202 | .update(codeVerifier)
203 | .digest('base64url');
204 | return codeChallenge === hash;
205 | }
206 |
207 | if (codeChallengeMethod === 'plain') {
208 | return codeChallenge === codeVerifier;
209 | }
210 |
211 | return false;
212 | };
213 |
```
--------------------------------------------------------------------------------
/src/tools/utils.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { NEON_DEFAULT_DATABASE_NAME } from '../constants.js';
2 | import { Api, Organization, Branch } from '@neondatabase/api-client';
3 | import { ToolHandlerExtraParams } from './types.js';
4 | import { NotFoundError } from '../server/errors.js';
5 |
6 | export const splitSqlStatements = (sql: string) => {
7 | return sql.split(';').filter(Boolean);
8 | };
9 |
10 | export const DESCRIBE_DATABASE_STATEMENTS = [
11 | `
12 | CREATE OR REPLACE FUNCTION public.show_db_tree()
13 | RETURNS TABLE (tree_structure text) AS
14 | $$
15 | BEGIN
16 | -- First show all databases
17 | RETURN QUERY
18 | SELECT ':file_folder: ' || datname || ' (DATABASE)'
19 | FROM pg_database
20 | WHERE datistemplate = false;
21 |
22 | -- Then show current database structure
23 | RETURN QUERY
24 | WITH RECURSIVE
25 | -- Get schemas
26 | schemas AS (
27 | SELECT
28 | n.nspname AS object_name,
29 | 1 AS level,
30 | n.nspname AS path,
31 | 'SCHEMA' AS object_type
32 | FROM pg_namespace n
33 | WHERE n.nspname NOT LIKE 'pg_%'
34 | AND n.nspname != 'information_schema'
35 | ),
36 |
37 | -- Get all objects (tables, views, functions, etc.)
38 | objects AS (
39 | SELECT
40 | c.relname AS object_name,
41 | 2 AS level,
42 | s.path || ' → ' || c.relname AS path,
43 | CASE c.relkind
44 | WHEN 'r' THEN 'TABLE'
45 | WHEN 'v' THEN 'VIEW'
46 | WHEN 'm' THEN 'MATERIALIZED VIEW'
47 | WHEN 'i' THEN 'INDEX'
48 | WHEN 'S' THEN 'SEQUENCE'
49 | WHEN 'f' THEN 'FOREIGN TABLE'
50 | END AS object_type
51 | FROM pg_class c
52 | JOIN pg_namespace n ON n.oid = c.relnamespace
53 | JOIN schemas s ON n.nspname = s.object_name
54 | WHERE c.relkind IN ('r','v','m','i','S','f')
55 |
56 | UNION ALL
57 |
58 | SELECT
59 | p.proname AS object_name,
60 | 2 AS level,
61 | s.path || ' → ' || p.proname AS path,
62 | 'FUNCTION' AS object_type
63 | FROM pg_proc p
64 | JOIN pg_namespace n ON n.oid = p.pronamespace
65 | JOIN schemas s ON n.nspname = s.object_name
66 | ),
67 |
68 | -- Combine schemas and objects
69 | combined AS (
70 | SELECT * FROM schemas
71 | UNION ALL
72 | SELECT * FROM objects
73 | )
74 |
75 | -- Final output with tree-like formatting
76 | SELECT
77 | REPEAT(' ', level) ||
78 | CASE
79 | WHEN level = 1 THEN '└── :open_file_folder: '
80 | ELSE ' └── ' ||
81 | CASE object_type
82 | WHEN 'TABLE' THEN ':bar_chart: '
83 | WHEN 'VIEW' THEN ':eye: '
84 | WHEN 'MATERIALIZED VIEW' THEN ':newspaper: '
85 | WHEN 'FUNCTION' THEN ':zap: '
86 | WHEN 'INDEX' THEN ':mag: '
87 | WHEN 'SEQUENCE' THEN ':1234: '
88 | WHEN 'FOREIGN TABLE' THEN ':globe_with_meridians: '
89 | ELSE ''
90 | END
91 | END || object_name || ' (' || object_type || ')'
92 | FROM combined
93 | ORDER BY path;
94 | END;
95 | $$ LANGUAGE plpgsql;
96 | `,
97 | `
98 | -- To use the function:
99 | SELECT * FROM show_db_tree();
100 | `,
101 | ];
102 |
103 | /**
104 | * Returns the default database for a project branch
105 | * If a database name is provided, it fetches and returns that database
106 | * Otherwise, it looks for a database named 'neondb' and returns that
107 | * If 'neondb' doesn't exist, it returns the first available database
108 | * Throws an error if no databases are found
109 | */
110 | export async function getDefaultDatabase(
111 | {
112 | projectId,
113 | branchId,
114 | databaseName,
115 | }: {
116 | projectId: string;
117 | branchId: string;
118 | databaseName?: string;
119 | },
120 | neonClient: Api<unknown>,
121 | ) {
122 | const { data } = await neonClient.listProjectBranchDatabases(
123 | projectId,
124 | branchId,
125 | );
126 | const databases = data.databases;
127 | if (databases.length === 0) {
128 | throw new NotFoundError('No databases found in your project branch');
129 | }
130 |
131 | if (databaseName) {
132 | const requestedDatabase = databases.find((db) => db.name === databaseName);
133 | if (requestedDatabase) {
134 | return requestedDatabase;
135 | }
136 | }
137 |
138 | const defaultDatabase = databases.find(
139 | (db) => db.name === NEON_DEFAULT_DATABASE_NAME,
140 | );
141 | return defaultDatabase || databases[0];
142 | }
143 |
144 | /**
145 | * Resolves the organization ID for API calls that require org_id parameter.
146 | *
147 | * For new users (those without billing_account), this function fetches user's organizations and auto-selects only organization managed by console. If there are multiple organizations managed by console, it throws an error asking user to specify org_id.
148 | *
149 | * For existing users (with billing_account), returns undefined to use default behavior.
150 | *
151 | * @param params - The parameters object that may contain org_id
152 | * @param neonClient - The Neon API client
153 | * @returns The organization to use, or undefined for default behavior
154 | */
155 | export async function getOrgByOrgIdOrDefault(
156 | params: { org_id?: string },
157 | neonClient: Api<unknown>,
158 | extra: ToolHandlerExtraParams,
159 | ): Promise<Organization | undefined> {
160 | // 1. If org_id is provided use it
161 | // 2. If using Org API key, use the account id
162 | if (params.org_id || extra.account.isOrg) {
163 | const orgId = params.org_id || extra.account.id;
164 | const { data } = await neonClient.getOrganization(orgId);
165 | return data;
166 | }
167 |
168 | const { data: user } = await neonClient.getCurrentUserInfo();
169 | if (user.billing_account) {
170 | return undefined;
171 | }
172 |
173 | const { data: response } = await neonClient.getCurrentUserOrganizations();
174 | const organizations = response.organizations || [];
175 |
176 | // 1. Filter organizations by managed_by==console, if there is only one organization, return that
177 | const consoleOrganizations = organizations.filter(
178 | (org) => org.managed_by === 'console',
179 | );
180 |
181 | if (consoleOrganizations.length === 1) {
182 | return consoleOrganizations[0];
183 | }
184 |
185 | // 2. If there are no organizations managed by console, and if there is only one organization (unfiltered), then return that organization
186 | if (consoleOrganizations.length === 0 && organizations.length === 1) {
187 | return organizations[0];
188 | }
189 |
190 | // 3. If there are no organizations at all, then throw error saying there are no organizations
191 | if (organizations.length === 0) {
192 | throw new NotFoundError('No organizations found for this user');
193 | }
194 |
195 | // 4. If there are multiple organizations, then throw error mentioning list of all these orgs (unfiltered)
196 | const orgList = organizations
197 | .map(
198 | (org) => `- ${org.name} (ID: ${org.id}) [managed by: ${org.managed_by}]`,
199 | )
200 | .join('\n');
201 | throw new NotFoundError(
202 | `Multiple organizations found. Please specify the org_id parameter with one of the following organization IDs:\n${orgList}`,
203 | );
204 | }
205 |
206 | export function filterOrganizations(
207 | organizations: Organization[],
208 | search?: string,
209 | ) {
210 | if (!search) {
211 | return organizations;
212 | }
213 | const searchLower = search.toLowerCase();
214 | return organizations.filter(
215 | (org) =>
216 | org.name.toLowerCase().includes(searchLower) ||
217 | org.id.toLowerCase().includes(searchLower),
218 | );
219 | }
220 |
221 | /**
222 | * Checks if a string looks like a branch ID based on the neonctl format
223 | * Branch IDs have format like "br-small-term-683261" (br- prefix + haiku pattern)
224 | */
225 | export function looksLikeBranchId(branch: string): boolean {
226 | const HAIKU_REGEX = /^[a-z0-9]+-[a-z0-9]+-[a-z0-9]+$/;
227 | return branch.startsWith('br-') && HAIKU_REGEX.test(branch.substring(3));
228 | }
229 |
230 | /**
231 | * Resolves a branch name or ID to the actual branch ID
232 | * If the input looks like a branch ID, returns it as-is
233 | * Otherwise, searches for a branch with matching name and returns its ID
234 | */
235 | export async function resolveBranchId(
236 | branchNameOrId: string,
237 | projectId: string,
238 | neonClient: Api<unknown>,
239 | ): Promise<{ branchId: string; branches: Branch[] }> {
240 | // Get all branches (we'll need this data anyway)
241 | const branchResponse = await neonClient.listProjectBranches({
242 | projectId,
243 | });
244 | const branches = branchResponse.data.branches;
245 |
246 | if (looksLikeBranchId(branchNameOrId)) {
247 | // Verify the branch ID actually exists
248 | const branch = branches.find((b) => b.id === branchNameOrId);
249 | if (!branch) {
250 | throw new NotFoundError(
251 | `Branch ID "${branchNameOrId}" not found in project ${projectId}`,
252 | );
253 | }
254 | return { branchId: branchNameOrId, branches };
255 | }
256 |
257 | // Search by name
258 | const branch = branches.find((b) => b.name === branchNameOrId);
259 | if (!branch) {
260 | const availableBranches = branches.map((b) => b.name).join(', ');
261 | throw new NotFoundError(
262 | `Branch name "${branchNameOrId}" not found in project ${projectId}.\nAvailable branches: ${availableBranches}`,
263 | );
264 | }
265 |
266 | return { branchId: branch.id, branches };
267 | }
268 |
```
--------------------------------------------------------------------------------
/src/tools-evaluations/prepare-database-migration.eval.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { Eval, EvalCase, Reporter, reportFailures } from 'braintrust';
2 | import { LLMClassifierFromTemplate } from 'autoevals';
3 |
4 | import { createApiClient } from '@neondatabase/api-client';
5 | import { deleteNonDefaultBranches, evaluateTask } from './evalUtils';
6 |
7 | const EVAL_INFO = {
8 | projectId: 'black-recipe-75251165',
9 | roleName: 'neondb_owner',
10 | databaseName: 'neondb',
11 | mainBranchId: 'br-cold-bird-a5icgh5h',
12 | };
13 |
14 | const getMainBranchDatabaseSchema = async () => {
15 | const neonClient = createApiClient({
16 | apiKey: process.env.NEON_API_KEY!,
17 | });
18 |
19 | const dbSchema = await neonClient.getProjectBranchSchema({
20 | projectId: EVAL_INFO.projectId,
21 | branchId: EVAL_INFO.mainBranchId,
22 | db_name: EVAL_INFO.databaseName,
23 | });
24 |
25 | return dbSchema.data.sql;
26 | };
27 |
28 | const factualityAnthropic = LLMClassifierFromTemplate({
29 | name: 'Factuality Anthropic',
30 | promptTemplate: `
31 | You are comparing a submitted answer to an expert answer on a given question. Here is the data:
32 | [BEGIN DATA]
33 | ************
34 | [Question]: {{{input}}}
35 | ************
36 | [Expert]: {{{expected}}}
37 | ************
38 | [Submission]: {{{output}}}
39 | ************
40 | [END DATA]
41 |
42 | Compare the factual content of the submitted answer with the expert answer.
43 | Implementation details like specific IDs, or exact formatting should be considered non-factual differences.
44 |
45 | Ignore the following differences:
46 | - Specific migration IDs or references
47 | - Formatting or structural variations
48 | - Order of presenting the information
49 | - Restatements of the same request/question
50 | - Additional confirmatory language that doesn't add new information
51 |
52 | The submitted answer may either be:
53 | (A) A subset missing key factual information from the expert answer
54 | (B) A superset that FIRST agrees with the expert answer's core facts AND THEN adds additional factual information
55 | (C) Factually equivalent to the expert answer
56 | (D) In factual disagreement with or takes a completely different action than the expert answer
57 | (E) Different only in non-factual implementation details
58 |
59 | Select the most appropriate option, prioritizing the core factual content over implementation specifics.
60 | `,
61 | choiceScores: {
62 | A: 0.4,
63 | B: 0.8,
64 | C: 1,
65 | D: 0,
66 | E: 1,
67 | },
68 | temperature: 0,
69 | useCoT: true,
70 | model: 'claude-3-5-sonnet-20241022',
71 | });
72 |
73 | const mainBranchIntegrityCheck = async (args: {
74 | input: string;
75 | output: string;
76 | expected: string;
77 | metadata?: {
78 | databaseSchemaBeforeRun: string;
79 | databaseSchemaAfterRun: string;
80 | };
81 | }) => {
82 | const databaseSchemaBeforeRun = args.metadata?.databaseSchemaBeforeRun;
83 | const databaseSchemaAfterRun = args.metadata?.databaseSchemaAfterRun;
84 | const databaseSchemaAfterRunResponseIsComplete =
85 | databaseSchemaAfterRun?.includes('PostgreSQL database dump complete') ??
86 | false;
87 |
88 | // sometimes the pg_dump fails to deliver the full responses, which leads to false negatives
89 | // so we must eject
90 | if (!databaseSchemaAfterRunResponseIsComplete) {
91 | return null;
92 | }
93 |
94 | const isSame = databaseSchemaBeforeRun === databaseSchemaAfterRun;
95 |
96 | return {
97 | name: 'Main Branch Integrity Check',
98 | score: isSame ? 1 : 0,
99 | };
100 | };
101 |
102 | Eval('prepare_database_migration', {
103 | data: (): EvalCase<
104 | string,
105 | string,
106 | | {
107 | databaseSchemaBeforeRun: string;
108 | databaseSchemaAfterRun: string;
109 | }
110 | | undefined
111 | >[] => {
112 | return [
113 | // Add column
114 | {
115 | input: `in my ${EVAL_INFO.projectId} project, add a new column Description to the posts table`,
116 | expected: `
117 | I've verified that the Description column has been successfully added to the posts table in a temporary branch. Would you like to commit the migration to the main branch?
118 |
119 | Migration Details:
120 | - Migration ID: <migration_id>
121 | - Temporary Branch Name: <temporary_branch_name>
122 | - Temporary Branch ID: <temporary_branch_id>
123 | - Migration Result: <migration_result>
124 | `,
125 | },
126 |
127 | // Add column with different type
128 | {
129 | input: `in my ${EVAL_INFO.projectId} project, add view_count column to posts table`,
130 | expected: `
131 | I've verified that the view_count column has been successfully added to the posts table in a temporary branch. Would you like to commit the migration to the main branch?
132 |
133 | Migration Details:
134 | - Migration ID: <migration_id>
135 | - Temporary Branch Name: <temporary_branch_name>
136 | - Temporary Branch ID: <temporary_branch_id>
137 | - Migration Result: <migration_result>
138 | `,
139 | },
140 |
141 | // Rename column
142 | {
143 | input: `in my ${EVAL_INFO.projectId} project, rename the content column to body in posts table`,
144 | expected: `
145 | I've verified that the content column has been successfully renamed to body in the posts table in a temporary branch. Would you like to commit the migration to the main branch?
146 |
147 | Migration Details:
148 | - Migration ID: <migration_id>
149 | - Temporary Branch Name: <temporary_branch_name>
150 | - Temporary Branch ID: <temporary_branch_id>
151 | - Migration Result: <migration_result>
152 | `,
153 | },
154 |
155 | // Add index
156 | {
157 | input: `in my ${EVAL_INFO.projectId} project, create an index on title column in posts table`,
158 | expected: `
159 | I've verified that the index has been successfully created on the title column in the posts table in a temporary branch. Would you like to commit the migration to the main branch?
160 |
161 | Migration Details:
162 | - Migration ID: <migration_id>
163 | - Temporary Branch Name: <temporary_branch_name>
164 | - Temporary Branch ID: <temporary_branch_id>
165 | - Migration Result: <migration_result>
166 | `,
167 | },
168 |
169 | // Drop column
170 | {
171 | input: `in my ${EVAL_INFO.projectId} project, drop the content column from posts table`,
172 | expected: `
173 | I've verified that the content column has been successfully dropped from the posts table in a temporary branch. Would you like to commit the migration to the main branch?
174 |
175 | Migration Details:
176 | - Migration ID: <migration_id>
177 | - Temporary Branch Name: <temporary_branch_name>
178 | - Temporary Branch ID: <temporary_branch_id>
179 | - Migration Result: <migration_result>
180 | `,
181 | },
182 |
183 | // Alter column type
184 | {
185 | input: `in my ${EVAL_INFO.projectId} project, change the title column type to text in posts table`,
186 | expected: `
187 | I've verified that the data type of the title column has been successfully changed in the posts table in a temporary branch. Would you like to commit the migration to the main branch?
188 |
189 | Migration Details:
190 | - Migration ID: <migration_id>
191 | - Temporary Branch Name: <temporary_branch_name>
192 | - Temporary Branch ID: <temporary_branch_id>
193 | - Migration Result: <migration_result>
194 | `,
195 | },
196 |
197 | // Add boolean column
198 | {
199 | input: `in my ${EVAL_INFO.projectId} project, add is_published column to posts table`,
200 | expected: `
201 | I've verified that the is_published column has been successfully added to the posts table in a temporary branch. Would you like to commit the migration to the main branch?
202 |
203 | Migration Details:
204 | - Migration ID: <migration_id>
205 | - Temporary Branch Name: <temporary_branch_name>
206 | - Temporary Branch ID: <temporary_branch_id>
207 | - Migration Result: <migration_result>
208 | `,
209 | },
210 |
211 | // Add numeric column
212 | {
213 | input: `in my ${EVAL_INFO.projectId} project, add likes_count column to posts table`,
214 | expected: `
215 | I've verified that the likes_count column has been successfully added to the posts table in a temporary branch. Would you like to commit the migration to the main branch?
216 |
217 | Migration Details:
218 | - Migration ID: <migration_id>
219 | - Temporary Branch Name: <temporary_branch_name>
220 | - Temporary Branch ID: <temporary_branch_id>
221 | - Migration Result: <migration_result>
222 | `,
223 | },
224 |
225 | // Create index
226 | {
227 | input: `in my ${EVAL_INFO.projectId} project, create index on title column in posts table`,
228 | expected: `
229 | I've verified that the index has been successfully created on the title column in the posts table in a temporary branch. Would you like to commit the migration to the main branch?
230 |
231 | Migration Details:
232 | - Migration ID: <migration_id>
233 | - Temporary Branch Name: <temporary_branch_name>
234 | - Temporary Branch ID: <temporary_branch_id>
235 | - Migration Result: <migration_result>
236 | `,
237 | },
238 | ];
239 | },
240 | task: async (input, hooks) => {
241 | const databaseSchemaBeforeRun = await getMainBranchDatabaseSchema();
242 | hooks.metadata.databaseSchemaBeforeRun = databaseSchemaBeforeRun;
243 |
244 | const llmCallMessages = await evaluateTask(input);
245 |
246 | const databaseSchemaAfterRun = await getMainBranchDatabaseSchema();
247 | hooks.metadata.databaseSchemaAfterRun = databaseSchemaAfterRun;
248 | hooks.metadata.llmCallMessages = llmCallMessages;
249 |
250 | deleteNonDefaultBranches(EVAL_INFO.projectId);
251 |
252 | const finalMessage = llmCallMessages[llmCallMessages.length - 1];
253 | return finalMessage.content;
254 | },
255 | trialCount: 20,
256 | maxConcurrency: 2,
257 | scores: [factualityAnthropic, mainBranchIntegrityCheck],
258 | });
259 |
260 | Reporter('Prepare Database Migration Reporter', {
261 | reportEval: async (evaluator, result, { verbose, jsonl }) => {
262 | const { results, summary } = result;
263 | const failingResults = results.filter(
264 | (r: { error: unknown }) => r.error !== undefined,
265 | );
266 |
267 | if (failingResults.length > 0) {
268 | reportFailures(evaluator, failingResults, { verbose, jsonl });
269 | }
270 |
271 | console.log(jsonl ? JSON.stringify(summary) : summary);
272 | return failingResults.length === 0;
273 | },
274 |
275 | // cleanup branches after the run
276 | reportRun: async (evalReports) => {
277 | await deleteNonDefaultBranches(EVAL_INFO.projectId);
278 |
279 | return evalReports.every((r) => r);
280 | },
281 | });
282 |
```
--------------------------------------------------------------------------------
/src/tools/toolsSchema.ts:
--------------------------------------------------------------------------------
```typescript
1 | import {
2 | ListProjectsParams,
3 | ListSharedProjectsParams,
4 | } from '@neondatabase/api-client';
5 | import { z } from 'zod';
6 | import { NEON_DEFAULT_DATABASE_NAME } from '../constants.js';
7 |
8 | type ZodObjectParams<T> = z.ZodObject<{ [key in keyof T]: z.ZodType<T[key]> }>;
9 |
10 | const DATABASE_NAME_DESCRIPTION = `The name of the database. If not provided, the default ${NEON_DEFAULT_DATABASE_NAME} or first available database is used.`;
11 |
12 | export const listProjectsInputSchema = z.object({
13 | cursor: z
14 | .string()
15 | .optional()
16 | .describe(
17 | 'Specify the cursor value from the previous response to retrieve the next batch of projects.',
18 | ),
19 | limit: z
20 | .number()
21 | .default(10)
22 | .describe(
23 | 'Specify a value from 1 to 400 to limit number of projects in the response.',
24 | ),
25 | search: z
26 | .string()
27 | .optional()
28 | .describe(
29 | 'Search by project name or id. You can specify partial name or id values to filter results.',
30 | ),
31 | org_id: z.string().optional().describe('Search for projects by org_id.'),
32 | }) satisfies ZodObjectParams<ListProjectsParams>;
33 |
34 | export const createProjectInputSchema = z.object({
35 | name: z
36 | .string()
37 | .optional()
38 | .describe('An optional name of the project to create.'),
39 | org_id: z
40 | .string()
41 | .optional()
42 | .describe('Create project in a specific organization.'),
43 | });
44 |
45 | export const deleteProjectInputSchema = z.object({
46 | projectId: z.string().describe('The ID of the project to delete'),
47 | });
48 |
49 | export const describeProjectInputSchema = z.object({
50 | projectId: z.string().describe('The ID of the project to describe'),
51 | });
52 |
53 | export const runSqlInputSchema = z.object({
54 | sql: z.string().describe('The SQL query to execute'),
55 | projectId: z
56 | .string()
57 | .describe('The ID of the project to execute the query against'),
58 | branchId: z
59 | .string()
60 | .optional()
61 | .describe(
62 | 'An optional ID of the branch to execute the query against. If not provided the default branch is used.',
63 | ),
64 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
65 | });
66 |
67 | export const runSqlTransactionInputSchema = z.object({
68 | sqlStatements: z.array(z.string()).describe('The SQL statements to execute'),
69 | projectId: z
70 | .string()
71 | .describe('The ID of the project to execute the query against'),
72 | branchId: z
73 | .string()
74 | .optional()
75 | .describe(
76 | 'An optional ID of the branch to execute the query against. If not provided the default branch is used.',
77 | ),
78 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
79 | });
80 |
81 | export const explainSqlStatementInputSchema = z.object({
82 | sql: z.string().describe('The SQL statement to analyze'),
83 | projectId: z
84 | .string()
85 | .describe('The ID of the project to execute the query against'),
86 | branchId: z
87 | .string()
88 | .optional()
89 | .describe(
90 | 'An optional ID of the branch to execute the query against. If not provided the default branch is used.',
91 | ),
92 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
93 | analyze: z
94 | .boolean()
95 | .default(true)
96 | .describe('Whether to include ANALYZE in the EXPLAIN command'),
97 | });
98 | export const describeTableSchemaInputSchema = z.object({
99 | tableName: z.string().describe('The name of the table'),
100 | projectId: z
101 | .string()
102 | .describe('The ID of the project to execute the query against'),
103 | branchId: z
104 | .string()
105 | .optional()
106 | .describe(
107 | 'An optional ID of the branch to execute the query against. If not provided the default branch is used.',
108 | ),
109 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
110 | });
111 |
112 | export const getDatabaseTablesInputSchema = z.object({
113 | projectId: z.string().describe('The ID of the project'),
114 | branchId: z
115 | .string()
116 | .optional()
117 | .describe(
118 | 'An optional ID of the branch. If not provided the default branch is used.',
119 | ),
120 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
121 | });
122 |
123 | export const createBranchInputSchema = z.object({
124 | projectId: z
125 | .string()
126 | .describe('The ID of the project to create the branch in'),
127 | branchName: z.string().optional().describe('An optional name for the branch'),
128 | });
129 |
130 | export const prepareDatabaseMigrationInputSchema = z.object({
131 | migrationSql: z
132 | .string()
133 | .describe('The SQL to execute to create the migration'),
134 | projectId: z
135 | .string()
136 | .describe('The ID of the project to execute the query against'),
137 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
138 | });
139 |
140 | export const completeDatabaseMigrationInputSchema = z.object({
141 | migrationId: z.string(),
142 | });
143 |
144 | export const describeBranchInputSchema = z.object({
145 | projectId: z.string().describe('The ID of the project'),
146 | branchId: z.string().describe('An ID of the branch to describe'),
147 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
148 | });
149 |
150 | export const deleteBranchInputSchema = z.object({
151 | projectId: z.string().describe('The ID of the project containing the branch'),
152 | branchId: z.string().describe('The ID of the branch to delete'),
153 | });
154 |
155 | export const getConnectionStringInputSchema = z.object({
156 | projectId: z
157 | .string()
158 | .describe(
159 | 'The ID of the project. If not provided, the only available project will be used.',
160 | ),
161 | branchId: z
162 | .string()
163 | .optional()
164 | .describe(
165 | 'The ID or name of the branch. If not provided, the default branch will be used.',
166 | ),
167 | computeId: z
168 | .string()
169 | .optional()
170 | .describe(
171 | 'The ID of the compute/endpoint. If not provided, the read-write compute associated with the branch will be used.',
172 | ),
173 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
174 | roleName: z
175 | .string()
176 | .optional()
177 | .describe(
178 | 'The name of the role to connect with. If not provided, the database owner name will be used.',
179 | ),
180 | });
181 |
182 | export const provisionNeonAuthInputSchema = z.object({
183 | projectId: z
184 | .string()
185 | .describe('The ID of the project to provision Neon Auth for'),
186 | database: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
187 | });
188 |
189 | export const prepareQueryTuningInputSchema = z.object({
190 | sql: z.string().describe('The SQL statement to analyze and tune'),
191 | databaseName: z
192 | .string()
193 | .describe('The name of the database to execute the query against'),
194 | projectId: z
195 | .string()
196 | .describe('The ID of the project to execute the query against'),
197 | roleName: z
198 | .string()
199 | .optional()
200 | .describe(
201 | 'The name of the role to connect with. If not provided, the default role (usually "neondb_owner") will be used.',
202 | ),
203 | });
204 |
205 | export const completeQueryTuningInputSchema = z.object({
206 | suggestedSqlStatements: z
207 | .array(z.string())
208 | .describe(
209 | 'The SQL DDL statements to execute to improve performance. These statements are the result of the prior steps, for example creating additional indexes.',
210 | ),
211 | applyChanges: z
212 | .boolean()
213 | .default(false)
214 | .describe('Whether to apply the suggested changes to the main branch'),
215 | tuningId: z
216 | .string()
217 | .describe(
218 | 'The ID of the tuning to complete. This is NOT the branch ID. Remember this ID from the prior step using tool prepare_query_tuning.',
219 | ),
220 | databaseName: z
221 | .string()
222 | .describe('The name of the database to execute the query against'),
223 | projectId: z
224 | .string()
225 | .describe('The ID of the project to execute the query against'),
226 | roleName: z
227 | .string()
228 | .optional()
229 | .describe(
230 | 'The name of the role to connect with. If you have used a specific role in prepare_query_tuning you MUST pass the same role again to this tool. If not provided, the default role (usually "neondb_owner") will be used.',
231 | ),
232 | shouldDeleteTemporaryBranch: z
233 | .boolean()
234 | .default(true)
235 | .describe('Whether to delete the temporary branch after tuning'),
236 | temporaryBranchId: z
237 | .string()
238 | .describe(
239 | 'The ID of the temporary branch that needs to be deleted after tuning.',
240 | ),
241 | branchId: z
242 | .string()
243 | .optional()
244 | .describe(
245 | 'The ID or name of the branch that receives the changes. If not provided, the default (main) branch will be used.',
246 | ),
247 | });
248 |
249 | export const listSlowQueriesInputSchema = z.object({
250 | projectId: z
251 | .string()
252 | .describe('The ID of the project to list slow queries from'),
253 | branchId: z
254 | .string()
255 | .optional()
256 | .describe(
257 | 'An optional ID of the branch. If not provided the default branch is used.',
258 | ),
259 | databaseName: z.string().optional().describe(DATABASE_NAME_DESCRIPTION),
260 | computeId: z
261 | .string()
262 | .optional()
263 | .describe(
264 | 'The ID of the compute/endpoint. If not provided, the read-write compute associated with the branch will be used.',
265 | ),
266 | limit: z
267 | .number()
268 | .optional()
269 | .default(10)
270 | .describe('Maximum number of slow queries to return'),
271 | minExecutionTime: z
272 | .number()
273 | .optional()
274 | .default(1000)
275 | .describe(
276 | 'Minimum execution time in milliseconds to consider a query as slow',
277 | ),
278 | });
279 |
280 | export const listBranchComputesInputSchema = z.object({
281 | projectId: z
282 | .string()
283 | .optional()
284 | .describe(
285 | 'The ID of the project. If not provided, the only available project will be used.',
286 | ),
287 | branchId: z
288 | .string()
289 | .optional()
290 | .describe(
291 | 'The ID of the branch. If provided, endpoints for this specific branch will be listed.',
292 | ),
293 | });
294 |
295 | export const listOrganizationsInputSchema = z.object({
296 | search: z
297 | .string()
298 | .optional()
299 | .describe(
300 | 'Search organizations by name or ID. You can specify partial name or ID values to filter results.',
301 | ),
302 | });
303 |
304 | export const listSharedProjectsInputSchema = z.object({
305 | cursor: z
306 | .string()
307 | .optional()
308 | .describe(
309 | 'Specify the cursor value from the previous response to retrieve the next batch of shared projects.',
310 | ),
311 | limit: z
312 | .number()
313 | .default(10)
314 | .describe(
315 | 'Specify a value from 1 to 400 to limit number of shared projects in the response.',
316 | ),
317 | search: z
318 | .string()
319 | .optional()
320 | .describe(
321 | 'Search by project name or id. You can specify partial name or id values to filter results.',
322 | ),
323 | }) satisfies ZodObjectParams<ListSharedProjectsParams>;
324 |
325 | export const resetFromParentInputSchema = z.object({
326 | projectId: z.string().describe('The ID of the project containing the branch'),
327 | branchIdOrName: z
328 | .string()
329 | .describe('The name or ID of the branch to reset from its parent'),
330 | preserveUnderName: z
331 | .string()
332 | .optional()
333 | .describe(
334 | 'Optional name to preserve the current state under a new branch before resetting',
335 | ),
336 | });
337 |
338 | export const compareDatabaseSchemaInputSchema = z.object({
339 | projectId: z.string().describe('The ID of the project'),
340 | branchId: z.string().describe('The ID of the branch'),
341 | databaseName: z.string().describe(DATABASE_NAME_DESCRIPTION),
342 | });
343 |
```
--------------------------------------------------------------------------------
/src/oauth/server.ts:
--------------------------------------------------------------------------------
```typescript
1 | import {
2 | Request as ExpressRequest,
3 | Response as ExpressResponse,
4 | } from 'express';
5 | import { AuthorizationCode, Client } from 'oauth2-server';
6 | import { model } from './model.js';
7 | import { logger } from '../utils/logger.js';
8 | import express from 'express';
9 | import {
10 | decodeAuthParams,
11 | extractClientCredentials,
12 | generateRandomString,
13 | parseAuthRequest,
14 | toMilliseconds,
15 | toSeconds,
16 | verifyPKCE,
17 | } from './utils.js';
18 | import { exchangeCode, exchangeRefreshToken, upstreamAuth } from './client.js';
19 | import { createNeonClient } from '../server/api.js';
20 | import bodyParser from 'body-parser';
21 | import { SERVER_HOST, COOKIE_SECRET } from '../constants.js';
22 | import {
23 | isClientAlreadyApproved,
24 | updateApprovedClientsCookie,
25 | } from './cookies.js';
26 | import { identify } from '../analytics/analytics.js';
27 |
28 | const SUPPORTED_GRANT_TYPES = ['authorization_code', 'refresh_token'];
29 | const SUPPORTED_RESPONSE_TYPES = ['code'];
30 | const SUPPORTED_AUTH_METHODS = [
31 | 'client_secret_post',
32 | 'client_secret_basic',
33 | 'none',
34 | ];
35 | const SUPPORTED_CODE_CHALLENGE_METHODS = ['S256'];
36 | export const metadata = (req: ExpressRequest, res: ExpressResponse) => {
37 | res.json({
38 | issuer: SERVER_HOST,
39 | authorization_endpoint: `${SERVER_HOST}/authorize`,
40 | token_endpoint: `${SERVER_HOST}/token`,
41 | registration_endpoint: `${SERVER_HOST}/register`,
42 | response_types_supported: SUPPORTED_RESPONSE_TYPES,
43 | response_modes_supported: ['query'],
44 | grant_types_supported: SUPPORTED_GRANT_TYPES,
45 | token_endpoint_auth_methods_supported: SUPPORTED_AUTH_METHODS,
46 | registration_endpoint_auth_methods_supported: SUPPORTED_AUTH_METHODS,
47 | code_challenge_methods_supported: SUPPORTED_CODE_CHALLENGE_METHODS,
48 | });
49 | };
50 |
51 | export const registerClient = async (
52 | req: ExpressRequest,
53 | res: ExpressResponse,
54 | ) => {
55 | const payload = req.body;
56 | logger.info('request to register client: ', {
57 | name: payload.client_name,
58 | client_uri: payload.client_uri,
59 | });
60 |
61 | if (payload.client_name === undefined) {
62 | res
63 | .status(400)
64 | .json({ code: 'invalid_request', error: 'client_name is required' });
65 | return;
66 | }
67 |
68 | if (payload.redirect_uris === undefined) {
69 | res
70 | .status(400)
71 | .json({ code: 'invalid_request', error: 'redirect_uris is required' });
72 | return;
73 | }
74 |
75 | if (
76 | payload.grant_types === undefined ||
77 | !payload.grant_types.every((grant: string) =>
78 | SUPPORTED_GRANT_TYPES.includes(grant),
79 | )
80 | ) {
81 | res.status(400).json({
82 | code: 'invalid_request',
83 | error:
84 | 'grant_types is required and must only include supported grant types',
85 | });
86 | return;
87 | }
88 |
89 | if (
90 | payload.response_types === undefined ||
91 | !payload.response_types.every((responseType: string) =>
92 | SUPPORTED_RESPONSE_TYPES.includes(responseType),
93 | )
94 | ) {
95 | res.status(400).json({
96 | code: 'invalid_request',
97 | error:
98 | 'response_types is required and must only include supported response types',
99 | });
100 | return;
101 | }
102 |
103 | try {
104 | const clientId = generateRandomString(8);
105 | const clientSecret = generateRandomString(32);
106 | const client: Client = {
107 | ...payload,
108 | id: clientId,
109 | secret: clientSecret,
110 | tokenEndpointAuthMethod:
111 | (req.body.token_endpoint_auth_method as string) ?? 'client_secret_post',
112 | registrationDate: Math.floor(Date.now() / 1000),
113 | };
114 |
115 | await model.saveClient(client);
116 | logger.info('new client registered', {
117 | clientId,
118 | client_name: payload.client_name,
119 | redirect_uris: payload.redirect_uris,
120 | client_uri: payload.client_uri,
121 | });
122 |
123 | res.json({
124 | client_id: clientId,
125 | client_secret: clientSecret,
126 | client_name: payload.client_name,
127 | redirect_uris: payload.redirect_uris,
128 | token_endpoint_auth_method: client.tokenEndpointAuthMethod,
129 | });
130 | } catch (error: unknown) {
131 | const message = error instanceof Error ? error.message : 'Unknown error';
132 | logger.error('failed to register client:', {
133 | message,
134 | error,
135 | client: payload.client_name,
136 | client_uri: payload.client_uri,
137 | });
138 | res.status(500).json({ code: 'server_error', error, message });
139 | }
140 | };
141 |
142 | const authRouter = express.Router();
143 | authRouter.get('/.well-known/oauth-authorization-server', metadata);
144 | authRouter.post('/register', bodyParser.json(), registerClient);
145 |
146 | /*
147 | Initiate the authorization code grant flow by validating the request parameters and then redirecting to the upstream authorization server.
148 |
149 | Step 1:
150 | MCP client should invoke this endpoint with the following parameters:
151 | <code>
152 | /authorize?client_id=clientId&redirect_uri=mcp://callback&response_type=code&scope=scope&code_challenge=codeChallenge&code_challenge_method=S256
153 | </code>
154 |
155 | This endpoint will validate the `client_id` and other request parameters and then capture the parameters on `state` param and redirect to the upstream authorization server.
156 | */
157 | authRouter.get(
158 | '/authorize',
159 | bodyParser.urlencoded({ extended: true }),
160 | async (req: ExpressRequest, res: ExpressResponse) => {
161 | const requestParams = parseAuthRequest(req);
162 |
163 | const clientId = requestParams.clientId;
164 | const client = await model.getClient(clientId, '');
165 | if (!client) {
166 | res
167 | .status(400)
168 | .json({ code: 'invalid_request', error: 'invalid client id' });
169 | return;
170 | }
171 |
172 | if (
173 | requestParams.responseType == undefined ||
174 | !client.response_types.includes(requestParams.responseType)
175 | ) {
176 | res
177 | .status(400)
178 | .json({ code: 'invalid_request', error: 'invalid response type' });
179 | return;
180 | }
181 |
182 | if (
183 | requestParams.redirectUri == undefined ||
184 | !client.redirect_uris.includes(requestParams.redirectUri)
185 | ) {
186 | res
187 | .status(400)
188 | .json({ code: 'invalid_request', error: 'invalid redirect uri' });
189 | return;
190 | }
191 |
192 | if (await isClientAlreadyApproved(req, client.id, COOKIE_SECRET)) {
193 | const authUrl = await upstreamAuth(btoa(JSON.stringify(requestParams)));
194 | res.redirect(authUrl.href);
195 | return;
196 | }
197 |
198 | res.render('approval-dialog', {
199 | client,
200 | state: btoa(JSON.stringify(requestParams)),
201 | });
202 | },
203 | );
204 |
205 | authRouter.post(
206 | '/authorize',
207 | bodyParser.urlencoded({ extended: true }),
208 | async (req: ExpressRequest, res: ExpressResponse) => {
209 | const state = req.body.state as string;
210 | if (!state) {
211 | res.status(400).json({ code: 'invalid_request', error: 'invalid state' });
212 | return;
213 | }
214 |
215 | const requestParams = JSON.parse(atob(state));
216 | await updateApprovedClientsCookie(
217 | req,
218 | res,
219 | requestParams.clientId,
220 | COOKIE_SECRET,
221 | );
222 | const authUrl = await upstreamAuth(state);
223 | res.redirect(authUrl.href);
224 | },
225 | );
226 |
227 | /*
228 | Handles the callback from the upstream authorization server and completes the authorization code grant flow with downstream MCP client.
229 |
230 | Step 2:
231 | Upstream authorization server will redirect to `/callback` with the authorization code.
232 | <code>
233 | /callback?code=authorizationCode&state=state
234 | </code>
235 |
236 | - Exchange the upstream authorization code for an access token.
237 | - Generate new authorization code and grant id.
238 | - Save the authorization code and access token in the database.
239 | - Redirect to the MCP client with the new authorization code.
240 | */
241 | authRouter.get(
242 | '/callback',
243 | bodyParser.urlencoded({ extended: true }),
244 | async (req: ExpressRequest, res: ExpressResponse) => {
245 | const tokens = await exchangeCode(req);
246 | const state = req.query.state as string;
247 | const requestParams = decodeAuthParams(state);
248 |
249 | const clientId = requestParams.clientId;
250 | const client = await model.getClient(clientId, '');
251 | if (!client) {
252 | res
253 | .status(400)
254 | .json({ code: 'invalid_request', error: 'invalid client id' });
255 | return;
256 | }
257 |
258 | // Standard authorization code grant
259 | const grantId = generateRandomString(16);
260 | const nonce = generateRandomString(32);
261 | const authCode = `${grantId}:${nonce}`;
262 |
263 | // Get the user's info from Neon
264 | const neonClient = createNeonClient(tokens.access_token);
265 | const { data: user } = await neonClient.getCurrentUserInfo();
266 | const expiresAt = Date.now() + toMilliseconds(tokens.expiresIn() ?? 0);
267 | // Save the authorization code with associated data
268 | const code: AuthorizationCode = {
269 | authorizationCode: authCode,
270 | expiresAt: new Date(Date.now() + 10 * 60 * 1000), // 10 minutes
271 | createdAt: Date.now(),
272 | redirectUri: requestParams.redirectUri,
273 | scope: requestParams.scope.join(' '),
274 | client: client,
275 | user: {
276 | id: user.id,
277 | email: user.email,
278 | name: `${user.name} ${user.last_name}`.trim(),
279 | },
280 | token: {
281 | access_token: tokens.access_token,
282 | access_token_expires_at: expiresAt,
283 | refresh_token: tokens.refresh_token,
284 | id_token: tokens.id_token,
285 | },
286 | code_challenge: requestParams.codeChallenge,
287 | code_challenge_method: requestParams.codeChallengeMethod,
288 | };
289 |
290 | await model.saveAuthorizationCode(code);
291 |
292 | // Redirect back to client with auth code
293 | const redirectUrl = new URL(requestParams.redirectUri);
294 | redirectUrl.searchParams.set('code', authCode);
295 | if (requestParams.state) {
296 | redirectUrl.searchParams.set('state', requestParams.state);
297 | }
298 |
299 | res.redirect(redirectUrl.href);
300 | },
301 | );
302 |
303 | /*
304 | Handles the token exchange for `code` and `refresh_token` grant types with downstream MCP client.
305 |
306 | Step 3:
307 | MCP client should invoke this endpoint after receiving the authorization code to exchange for an access token.
308 | <code>
309 | /token?client_id=clientId&grant_type=code&code=authorizationCode
310 | </code>
311 |
312 | - Verify the authorization code, grant type and client
313 | - Save the access token and refresh token in the database for further API requests verification
314 | - Return with access token and refresh token
315 | */
316 | authRouter.post(
317 | '/token',
318 | bodyParser.urlencoded({ extended: true }),
319 | async (req: ExpressRequest, res: ExpressResponse) => {
320 | const contentType = req.headers['content-type'] as string;
321 | if (contentType !== 'application/x-www-form-urlencoded') {
322 | res
323 | .status(415)
324 | .json({ code: 'invalid_request', error: 'invalid content type' });
325 | return;
326 | }
327 | const { clientId, clientSecret } = extractClientCredentials(req);
328 | if (!clientId) {
329 | res
330 | .status(400)
331 | .json({ code: 'invalid_request', error: 'client_id is required' });
332 | return;
333 | }
334 |
335 | const error = {
336 | error: 'invalid_client',
337 | error_description: 'client not found or invalid client credentials',
338 | };
339 | const client = await model.getClient(clientId, '');
340 | if (!client) {
341 | res.status(400).json({ code: 'invalid_request', ...error });
342 | return;
343 | }
344 |
345 | const isPublicClient = client.tokenEndpointAuthMethod === 'none';
346 | if (!isPublicClient) {
347 | if (clientSecret !== client.secret) {
348 | res.status(400).json({ code: 'invalid_request', ...error });
349 | return;
350 | }
351 | }
352 |
353 | const formData = req.body;
354 | if (formData.grant_type === 'authorization_code') {
355 | const authorizationCode = await model.getAuthorizationCode(formData.code);
356 | if (!authorizationCode) {
357 | res.status(400).json({
358 | code: 'invalid_request',
359 | error: 'invalid authorization code',
360 | });
361 | return;
362 | }
363 |
364 | if (authorizationCode.client.id !== client.id) {
365 | res.status(400).json({
366 | code: 'invalid_request',
367 | error: 'invalid authorization code',
368 | });
369 | return;
370 | }
371 |
372 | if (authorizationCode.expiresAt < new Date()) {
373 | res.status(400).json({
374 | code: 'invalid_request',
375 | error: 'authorization code expired',
376 | });
377 | return;
378 | }
379 |
380 | const isPkceEnabled = authorizationCode.code_challenge !== undefined;
381 | if (
382 | isPkceEnabled &&
383 | !verifyPKCE(
384 | authorizationCode.code_challenge,
385 | authorizationCode.code_challenge_method,
386 | formData.code_verifier,
387 | )
388 | ) {
389 | res.status(400).json({
390 | code: 'invalid_grant',
391 | error: 'invalid PKCE code verifier',
392 | });
393 | return;
394 | }
395 | if (!isPkceEnabled && !formData.redirect_uri) {
396 | res.status(400).json({
397 | code: 'invalid_request',
398 | error: 'redirect_uri is required when not using PKCE',
399 | });
400 | return;
401 | }
402 | if (
403 | formData.redirect_uri &&
404 | !client.redirect_uris.includes(formData.redirect_uri)
405 | ) {
406 | res.status(400).json({
407 | code: 'invalid_request',
408 | error: 'invalid redirect uri',
409 | });
410 | return;
411 | }
412 |
413 | // TODO: Generate fresh tokens and add mapping to database.
414 | const token = await model.saveToken({
415 | accessToken: authorizationCode.token.access_token,
416 | refreshToken: authorizationCode.token.refresh_token,
417 | expires_at: authorizationCode.token.access_token_expires_at,
418 | client: client,
419 | user: authorizationCode.user,
420 | });
421 |
422 | await model.saveRefreshToken({
423 | refreshToken: token.refreshToken ?? '',
424 | accessToken: token.accessToken,
425 | });
426 |
427 | identify(
428 | {
429 | id: authorizationCode.user.id,
430 | name: authorizationCode.user.name,
431 | email: authorizationCode.user.email,
432 | },
433 | {
434 | context: {
435 | client: {
436 | id: client.id,
437 | name: client.client_name,
438 | },
439 | },
440 | },
441 | );
442 |
443 | // Revoke the authorization code, it can only be used once
444 | await model.revokeAuthorizationCode(authorizationCode);
445 | res.json({
446 | access_token: token.accessToken,
447 | expires_in: toSeconds(token.expires_at - Date.now()),
448 | token_type: 'bearer', // TODO: Verify why non-bearer tokens are not working
449 | refresh_token: token.refreshToken,
450 | scope: authorizationCode.scope,
451 | });
452 | return;
453 | } else if (formData.grant_type === 'refresh_token') {
454 | const providedRefreshToken = await model.getRefreshToken(
455 | formData.refresh_token,
456 | );
457 | if (!providedRefreshToken) {
458 | res
459 | .status(400)
460 | .json({ code: 'invalid_request', error: 'invalid refresh token' });
461 | return;
462 | }
463 |
464 | const oldToken = await model.getAccessToken(
465 | providedRefreshToken.accessToken,
466 | );
467 | if (!oldToken) {
468 | // Refresh token is missing its counter access token, delete it
469 | await model.deleteRefreshToken(providedRefreshToken);
470 | res
471 | .status(400)
472 | .json({ code: 'invalid_request', error: 'invalid refresh token' });
473 | return;
474 | }
475 |
476 | if (oldToken.client.id !== client.id) {
477 | res
478 | .status(400)
479 | .json({ code: 'invalid_request', error: 'invalid refresh token' });
480 | return;
481 | }
482 |
483 | const upstreamToken = await exchangeRefreshToken(
484 | providedRefreshToken.refreshToken,
485 | );
486 | const now = Date.now();
487 | const expiresAt = now + toMilliseconds(upstreamToken.expiresIn() ?? 0);
488 | const token = await model.saveToken({
489 | accessToken: upstreamToken.access_token,
490 | refreshToken: upstreamToken.refresh_token ?? '',
491 | expires_at: expiresAt,
492 | client: client,
493 | user: oldToken.user,
494 | });
495 | await model.saveRefreshToken({
496 | refreshToken: token.refreshToken ?? '',
497 | accessToken: token.accessToken,
498 | });
499 |
500 | // Delete the old tokens
501 | await model.deleteToken(oldToken);
502 | await model.deleteRefreshToken(providedRefreshToken);
503 |
504 | res.json({
505 | access_token: token.accessToken,
506 | expires_in: toSeconds(expiresAt - now),
507 | token_type: 'bearer',
508 | refresh_token: token.refreshToken,
509 | scope: oldToken.scope,
510 | });
511 | return;
512 | }
513 | res
514 | .status(400)
515 | .json({ code: 'invalid_request', error: 'invalid grant type' });
516 | },
517 | );
518 |
519 | export { authRouter };
520 |
```
--------------------------------------------------------------------------------
/src/tools/definitions.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { NEON_DEFAULT_DATABASE_NAME } from '../constants.js';
2 | import {
3 | completeDatabaseMigrationInputSchema,
4 | completeQueryTuningInputSchema,
5 | createBranchInputSchema,
6 | createProjectInputSchema,
7 | deleteBranchInputSchema,
8 | deleteProjectInputSchema,
9 | describeBranchInputSchema,
10 | describeProjectInputSchema,
11 | describeTableSchemaInputSchema,
12 | explainSqlStatementInputSchema,
13 | getConnectionStringInputSchema,
14 | getDatabaseTablesInputSchema,
15 | listBranchComputesInputSchema,
16 | listProjectsInputSchema,
17 | prepareDatabaseMigrationInputSchema,
18 | prepareQueryTuningInputSchema,
19 | provisionNeonAuthInputSchema,
20 | runSqlInputSchema,
21 | runSqlTransactionInputSchema,
22 | listSlowQueriesInputSchema,
23 | listOrganizationsInputSchema,
24 | listSharedProjectsInputSchema,
25 | resetFromParentInputSchema,
26 | compareDatabaseSchemaInputSchema,
27 | } from './toolsSchema.js';
28 |
29 | export const NEON_TOOLS = [
30 | {
31 | name: 'list_projects' as const,
32 | description: `Lists the first 10 Neon projects in your account. If you can't find the project, increase the limit by passing a higher value to the \`limit\` parameter. Optionally filter by project name or ID using the \`search\` parameter.`,
33 | inputSchema: listProjectsInputSchema,
34 | },
35 | {
36 | name: 'list_organizations' as const,
37 | description: `Lists all organizations that the current user has access to. Optionally filter by organization name or ID using the \`search\` parameter.`,
38 | inputSchema: listOrganizationsInputSchema,
39 | },
40 | {
41 | name: 'list_shared_projects' as const,
42 | description: `Lists projects that have been shared with the current user. These are projects that the user has been granted access to collaborate on. Optionally filter by project name or ID using the \`search\` parameter.`,
43 | inputSchema: listSharedProjectsInputSchema,
44 | },
45 | {
46 | name: 'create_project' as const,
47 | description:
48 | 'Create a new Neon project. If someone is trying to create a database, use this tool.',
49 | inputSchema: createProjectInputSchema,
50 | },
51 | {
52 | name: 'delete_project' as const,
53 | description: 'Delete a Neon project',
54 | inputSchema: deleteProjectInputSchema,
55 | },
56 | {
57 | name: 'describe_project' as const,
58 | description: 'Describes a Neon project',
59 | inputSchema: describeProjectInputSchema,
60 | },
61 | {
62 | name: 'run_sql' as const,
63 | description: `
64 | <use_case>
65 | Use this tool to execute a single SQL statement against a Neon database.
66 | </use_case>
67 |
68 | <important_notes>
69 | If you have a temporary branch from a prior step, you MUST:
70 | 1. Pass the branch ID to this tool unless explicitly told otherwise
71 | 2. Tell the user that you are using the temporary branch with ID [branch_id]
72 | </important_notes>`,
73 | inputSchema: runSqlInputSchema,
74 | },
75 | {
76 | name: 'run_sql_transaction' as const,
77 | description: `
78 | <use_case>
79 | Use this tool to execute a SQL transaction against a Neon database, should be used for multiple SQL statements.
80 | </use_case>
81 |
82 | <important_notes>
83 | If you have a temporary branch from a prior step, you MUST:
84 | 1. Pass the branch ID to this tool unless explicitly told otherwise
85 | 2. Tell the user that you are using the temporary branch with ID [branch_id]
86 | </important_notes>`,
87 | inputSchema: runSqlTransactionInputSchema,
88 | },
89 | {
90 | name: 'describe_table_schema' as const,
91 | description: 'Describe the schema of a table in a Neon database',
92 | inputSchema: describeTableSchemaInputSchema,
93 | },
94 | {
95 | name: 'get_database_tables' as const,
96 | description: 'Get all tables in a Neon database',
97 | inputSchema: getDatabaseTablesInputSchema,
98 | },
99 | {
100 | name: 'create_branch' as const,
101 | description: 'Create a branch in a Neon project',
102 | inputSchema: createBranchInputSchema,
103 | },
104 | {
105 | name: 'prepare_database_migration' as const,
106 | description: `
107 | <use_case>
108 | This tool performs database schema migrations by automatically generating and executing DDL statements.
109 |
110 | Supported operations:
111 | CREATE operations:
112 | - Add new columns (e.g., "Add email column to users table")
113 | - Create new tables (e.g., "Create posts table with title and content columns")
114 | - Add constraints (e.g., "Add unique constraint on \`users.email\`")
115 |
116 | ALTER operations:
117 | - Modify column types (e.g., "Change posts.views to bigint")
118 | - Rename columns (e.g., "Rename user_name to username in users table")
119 | - Add/modify indexes (e.g., "Add index on \`posts.title\`")
120 | - Add/modify foreign keys (e.g., "Add foreign key from \`posts.user_id\` to \`users.id\`")
121 |
122 | DROP operations:
123 | - Remove columns (e.g., "Drop temporary_field from users table")
124 | - Drop tables (e.g., "Drop the old_logs table")
125 | - Remove constraints (e.g., "Remove unique constraint from posts.slug")
126 |
127 | The tool will:
128 | 1. Parse your natural language request
129 | 2. Generate appropriate SQL
130 | 3. Execute in a temporary branch for safety
131 | 4. Verify the changes before applying to main branch
132 |
133 | Project ID and database name will be automatically extracted from your request.
134 | If the database name is not provided, the default ${NEON_DEFAULT_DATABASE_NAME} or first available database is used.
135 | </use_case>
136 |
137 | <workflow>
138 | 1. Creates a temporary branch
139 | 2. Applies the migration SQL in that branch
140 | 3. Returns migration details for verification
141 | </workflow>
142 |
143 | <important_notes>
144 | After executing this tool, you MUST:
145 | 1. Test the migration in the temporary branch using the \`run_sql\` tool
146 | 2. Ask for confirmation before proceeding
147 | 3. Use \`complete_database_migration\` tool to apply changes to main branch
148 | </important_notes>
149 |
150 | <example>
151 | For a migration like:
152 | \`\`\`sql
153 | ALTER TABLE users ADD COLUMN last_login TIMESTAMP;
154 | \`\`\`
155 |
156 | You should test it with:
157 | \`\`\`sql
158 | SELECT column_name, data_type
159 | FROM information_schema.columns
160 | WHERE table_name = 'users' AND column_name = 'last_login';
161 | \`\`\`
162 |
163 | You can use \`run_sql\` to test the migration in the temporary branch that this tool creates.
164 | </example>
165 |
166 |
167 | <next_steps>
168 | After executing this tool, you MUST follow these steps:
169 | 1. Use \`run_sql\` to verify changes on temporary branch
170 | 2. Follow these instructions to respond to the client:
171 |
172 | <response_instructions>
173 | <instructions>
174 | Provide a brief confirmation of the requested change and ask for migration commit approval.
175 |
176 | You MUST include ALL of the following fields in your response:
177 | - Migration ID (this is required for commit and must be shown first)
178 | - Temporary Branch Name (always include exact branch name)
179 | - Temporary Branch ID (always include exact ID)
180 | - Migration Result (include brief success/failure status)
181 |
182 | Even if some fields are missing from the tool's response, use placeholders like "not provided" rather than omitting fields.
183 | </instructions>
184 |
185 | <do_not_include>
186 | IMPORTANT: Your response MUST NOT contain ANY technical implementation details such as:
187 | - Data types (e.g., DO NOT mention if a column is boolean, varchar, timestamp, etc.)
188 | - Column specifications or properties
189 | - SQL syntax or statements
190 | - Constraint definitions or rules
191 | - Default values
192 | - Index types
193 | - Foreign key specifications
194 |
195 | Keep the response focused ONLY on confirming the high-level change and requesting approval.
196 |
197 | <example>
198 | INCORRECT: "I've added a boolean \`is_published\` column to the \`posts\` table..."
199 | CORRECT: "I've added the \`is_published\` column to the \`posts\` table..."
200 | </example>
201 | </do_not_include>
202 |
203 | <example>
204 | I've verified that [requested change] has been successfully applied to a temporary branch. Would you like to commit the migration \`[migration_id]\` to the main branch?
205 |
206 | Migration Details:
207 | - Migration ID (required for commit)
208 | - Temporary Branch Name
209 | - Temporary Branch ID
210 | - Migration Result
211 | </example>
212 | </response_instructions>
213 |
214 | 3. If approved, use \`complete_database_migration\` tool with the \`migration_id\`
215 | </next_steps>
216 |
217 | <error_handling>
218 | On error, the tool will:
219 | 1. Automatically attempt ONE retry of the exact same operation
220 | 2. If the retry fails:
221 | - Terminate execution
222 | - Return error details
223 | - DO NOT attempt any other tools or alternatives
224 |
225 | Error response will include:
226 | - Original error details
227 | - Confirmation that retry was attempted
228 | - Final error state
229 |
230 | Important: After a failed retry, you must terminate the current flow completely. Do not attempt to use alternative tools or workarounds.
231 | </error_handling>`,
232 | inputSchema: prepareDatabaseMigrationInputSchema,
233 | },
234 | {
235 | name: 'complete_database_migration' as const,
236 | description:
237 | 'Complete a database migration when the user confirms the migration is ready to be applied to the main branch. This tool also lets the client know that the temporary branch created by the `prepare_database_migration` tool has been deleted.',
238 | inputSchema: completeDatabaseMigrationInputSchema,
239 | },
240 | {
241 | name: 'describe_branch' as const,
242 | description:
243 | 'Get a tree view of all objects in a branch, including databases, schemas, tables, views, and functions',
244 | inputSchema: describeBranchInputSchema,
245 | },
246 | {
247 | name: 'delete_branch' as const,
248 | description: 'Delete a branch from a Neon project',
249 | inputSchema: deleteBranchInputSchema,
250 | },
251 | {
252 | name: 'reset_from_parent' as const,
253 | description: `Resets a branch to match its parent's current state, effectively discarding all changes made on the branch. To avoid data loss, provide a name to preserve the changes in a new branch using \`preserveUnderName\` parameter. This tool is commonly used to create fresh development branches from updated parent branch, undo experimental changes, or restore a branch to a known good state. Warning: This operation will discard all changes if \`preserveUnderName\` is not provided.`,
254 | inputSchema: resetFromParentInputSchema,
255 | },
256 | {
257 | name: 'get_connection_string' as const,
258 | description:
259 | 'Get a PostgreSQL connection string for a Neon database with all parameters being optional',
260 | inputSchema: getConnectionStringInputSchema,
261 | },
262 | {
263 | name: 'provision_neon_auth' as const,
264 | inputSchema: provisionNeonAuthInputSchema,
265 | description: `
266 | This tool provisions authentication for a Neon project. It allows developers to easily setup authentication infrastructure by creating a integration with Stack Auth (\`@stackframe/stack\`). Stop using NextAuth.js! use \`@stackframe/stack\`.
267 |
268 | Parameters:
269 | - \`<project_id>\`: The Project ID of the Neon project to provision authentication for.
270 | - \`[database]\`: The database name to setup Neon Auth for. If not provided, the default \`${NEON_DEFAULT_DATABASE_NAME}\` or first available database is used.
271 |
272 | <workflow>
273 | The tool will:
274 | 1. Establish a connection between your Neon Auth project and Stack Auth
275 | 2. Creates a dedicated authentication schema in your database (\`neon_auth\`)
276 | 3. Sets up the user table under the \`neon_auth\` schema. This table is synced with Stack Auth. It does not store user credentials or secrets.
277 | 4. Generates Client Key and Secret Key to connect your application with authentication provider.
278 |
279 | Use the Stack Auth SDK (\`@stackframe/stack\`) on the frontend to connect your application with authentication provider. DO NOT use NextAuth.js! DO NOT use better-auth! Here's some documentation on Stack Auth:
280 | </workflow>
281 |
282 | <use_case>
283 | Stack Auth Guidelines
284 |
285 | <instructions>
286 | If you're building an app with Next.js, to set up Neon Auth and Stack Auth, follow these steps:
287 | 1. Provision a Neon Auth project with this tool
288 | 2. Place the returned credentials in project's \`.env.local\` or \`.env\` file
289 | - \`NEXT_PUBLIC_STACK_PROJECT_ID\`
290 | - \`NEXT_PUBLIC_STACK_PUBLISHABLE_CLIENT_KEY\`
291 | - \`STACK_SECRET_SERVER_KEY\`
292 | 3. To setup Stack Auth, run following command:
293 | \`\`\`shell
294 | npx @stackframe/init-stack . --no-browser
295 | \`\`\`
296 | This command will automaticallysetup the project with -
297 | - It will add \`@stackframe/stack\` dependency to \`package.json\`
298 | - It will create a \`stack.ts\` file in your project to setup \`StackServerApp\`.
299 | - It will wrap the root layout with \`StackProvider\` and \`StackTheme\`
300 | - It will create root Suspense boundary \`app/loading.tsx\` to handle loading state while Stack is fetching user data.
301 | - It will also create \`app/handler/[...stack]/page.tsx\` file to handle auth routes like sign in, sign up, forgot password, etc.
302 | 4. Do not try to manually create any of these files or directories. Do not try to create SignIn, SignUp, or UserButton components manually, instead use the ones provided by \`@stackframe/stack\`.
303 | </instructions>
304 |
305 | <instructions>
306 | Components Guidelines
307 | - Use pre-built components from \`@stackframe/stack\` like \`<UserButton />\`, \`<SignIn />\`, and \`<SignUp />\` to quickly set up auth UI.
308 | - You can also compose smaller pieces like \`<OAuthButtonGroup />\`, \`<MagicLinkSignIn />\`, and \`<CredentialSignIn />\` for custom flows.
309 | <example>
310 | \`\`\`tsx
311 | import { SignIn } from '@stackframe/stack';
312 | export default function Page() {
313 | return <SignIn />;
314 | }
315 | \`\`\`
316 | </example>
317 | </instructions>
318 |
319 | <instructions>
320 | User Management Guidelines
321 | - In Client Components, use the \`useUser()\` hook to retrieve the current user (it returns \`null\` when not signed in).
322 | - Update user details using \`user.update({...})\` and sign out via \`user.signOut()\`.
323 | - For pages that require a user, call \`useUser({ or: "redirect" })\` so unauthorized visitors are automatically redirected.
324 | </instructions>
325 |
326 | <instructions>
327 | Client Component Guidelines
328 | - Client Components rely on hooks like \`useUser()\` and \`useStackApp()\`.
329 |
330 | <example>
331 | \`\`\`tsx
332 | "use client";
333 | import { useUser } from "@stackframe/stack";
334 | export function MyComponent() {
335 | const user = useUser();
336 | return <div>{user ? \`Hello, \${user.displayName}\` : "Not logged in"}</div>;
337 | }
338 | \`\`\`
339 | </example>
340 | </instructions>
341 |
342 | <instructions>
343 | Server Component Guidelines
344 | - For Server Components, use \`stackServerApp.getUser()\` from your \`stack.ts\` file.
345 |
346 | <example>
347 | \`\`\`tsx
348 | import { stackServerApp } from "@/stack";
349 | export default async function ServerComponent() {
350 | const user = await stackServerApp.getUser();
351 | return <div>{user ? \`Hello, \${user.displayName}\` : "Not logged in"}</div>;
352 | }
353 | \`\`\`
354 | </example>
355 | </instructions>
356 |
357 | <instructions>
358 | Page Protection Guidelines
359 | - Protect pages by:
360 | - Using \`useUser({ or: "redirect" })\` in Client Components.
361 | - Using \`await stackServerApp.getUser({ or: "redirect" })\` in Server Components.
362 | - Implementing middleware that checks for a user and redirects to \`/handler/sign-in\` if not found.
363 |
364 | <example>
365 | Example middleware:
366 | \`\`\`tsx
367 | export async function middleware(request: NextRequest) {
368 | const user = await stackServerApp.getUser();
369 | if (!user) {
370 | return NextResponse.redirect(new URL('/handler/sign-in', request.url));
371 | }
372 | return NextResponse.next();
373 | }
374 | export const config = { matcher: '/protected/:path*' };
375 | \`\`\`
376 | </example>
377 | </instructions>
378 |
379 | <workflow>
380 | Example: custom-profile-page
381 | <instructions>
382 | Create a custom profile page that:
383 | - Displays the user's avatar, display name, and email.
384 | - Provides options to sign out.
385 | - Uses Stack Auth components and hooks.
386 | </instructions>
387 | <example>
388 | File: \`app/profile/page.tsx\`
389 | \`\`\`tsx
390 | 'use client';
391 | import { useUser, useStackApp, UserButton } from '@stackframe/stack';
392 | export default function ProfilePage() {
393 | const user = useUser({ or: "redirect" });
394 | const app = useStackApp();
395 | return (
396 | <div>
397 | <UserButton />
398 | <h1>Welcome, {user.displayName || "User"}</h1>
399 | <p>Email: {user.primaryEmail}</p>
400 | <button onClick={() => user.signOut()}>Sign Out</button>
401 | </div>
402 | );
403 | }
404 | \`\`\`
405 | </example>
406 | </workflow>
407 | </use_case>
408 | `,
409 | },
410 | {
411 | name: 'explain_sql_statement' as const,
412 | description:
413 | 'Describe the PostgreSQL query execution plan for a query of SQL statement by running EXPLAIN (ANAYLZE...) in the database',
414 | inputSchema: explainSqlStatementInputSchema,
415 | },
416 | {
417 | name: 'prepare_query_tuning' as const,
418 | description: `
419 | <use_case>
420 | This tool helps developers improve PostgreSQL query performance for slow queries or DML statements by analyzing execution plans and suggesting optimizations.
421 |
422 | The tool will:
423 | 1. Create a temporary branch for testing optimizations and remember the branch ID
424 | 2. Extract and analyze the current query execution plan
425 | 3. Extract all fully qualified table names (\`schema.table\`) referenced in the plan
426 | 4. Gather detailed schema information for each referenced table using \`describe_table_schema\`
427 | 5. Suggest and implement improvements like:
428 | - Adding or modifying indexes based on table schemas and query patterns
429 | - Query structure modifications
430 | - Identifying potential performance bottlenecks
431 | 6. Apply the changes to the temporary branch using \`run_sql\`
432 | 7. Compare performance before and after changes (but ONLY on the temporary branch passing branch ID to all tools)
433 | 8. Continue with next steps using \`complete_query_tuning\` tool (on \`main\` branch)
434 |
435 | Project ID and database name will be automatically extracted from your request.
436 | The temporary branch ID will be added when invoking other tools.
437 | Default database is \`${NEON_DEFAULT_DATABASE_NAME}\` if not specified.
438 |
439 | <important_notes>
440 | This tool is part of the query tuning workflow. Any suggested changes (like creating indexes) must first be applied to the temporary branch using the \`run_sql\` tool.
441 | And then to the main branch using the \`complete_query_tuning\` tool, NOT the \`prepare_database_migration\` tool.
442 | To apply using the \`complete_query_tuning\` tool, you must pass the \`tuning_id\`, NOT the temporary branch ID to it.
443 | </important_notes>
444 | </use_case>
445 |
446 | <workflow>
447 | 1. Creates a temporary branch
448 | 2. Analyzes current query performance and extracts table information
449 | 3. Implements and tests improvements (using tool \`run_sql\` for schema modifications and \`explain_sql_statement\` for performance analysis, but ONLY on the temporary branch created in step 1 passing the same branch ID to all tools)
450 | 4. Returns tuning details for verification
451 | </workflow>
452 |
453 | <important_notes>
454 | After executing this tool, you MUST:
455 | 1. Review the suggested changes
456 | 2. Verify the performance improvements on temporary branch - by applying the changes with \`run_sql\` and running \`explain_sql_statement\` again)
457 | 3. Decide whether to keep or discard the changes
458 | 4. Use \`complete_query_tuning\` tool to apply or discard changes to the main branch
459 |
460 | DO NOT use \`prepare_database_migration\` tool for applying query tuning changes.
461 | Always use \`complete_query_tuning\` to ensure changes are properly tracked and applied.
462 |
463 | Note:
464 | - Some operations like creating indexes can take significant time on large tables
465 | - Table statistics updates (ANALYZE) are NOT automatically performed as they can be long-running
466 | - Table statistics maintenance should be handled by PostgreSQL auto-analyze or scheduled maintenance jobs
467 | - If statistics are suspected to be stale, suggest running ANALYZE as a separate maintenance task
468 | </important_notes>
469 |
470 | <example>
471 | For a query like:
472 | \`\`\`sql
473 | SELECT o.*, c.name
474 | FROM orders o
475 | JOIN customers c ON c.id = o.customer_id
476 | WHERE o.status = 'pending'
477 | AND o.created_at > '2024-01-01';
478 | \`\`\`
479 |
480 | The tool will:
481 | 1. Extract referenced tables: \`public.orders\`, \`public.customers\`
482 | 2. Gather schema information for both tables
483 | 3. Analyze the execution plan
484 | 4. Suggest improvements like:
485 | - Creating a composite index on orders(status, created_at)
486 | - Optimizing the join conditions
487 | 5. If confirmed, apply the suggested changes to the temporary branch using \`run_sql\`
488 | 6. Compare execution plans and performance before and after changes (but ONLY on the temporary branch passing branch ID to all tools)
489 | </example>
490 |
491 | <next_steps>
492 | After executing this tool, you MUST follow these steps:
493 | 1. Review the execution plans and suggested changes
494 | 2. Follow these instructions to respond to the client:
495 |
496 | <response_instructions>
497 | <instructions>
498 | Provide a brief summary of the performance analysis and ask for approval to apply changes on the temporary branch.
499 |
500 | You MUST include ALL of the following fields in your response:
501 | - Tuning ID (this is required for completion)
502 | - Temporary Branch Name
503 | - Temporary Branch ID
504 | - Original Query Cost
505 | - Improved Query Cost
506 | - Referenced Tables (list all tables found in the plan)
507 | - Suggested Changes
508 |
509 | Even if some fields are missing from the tool's response, use placeholders like "not provided" rather than omitting fields.
510 | </instructions>
511 |
512 | <do_not_include>
513 | IMPORTANT: Your response MUST NOT contain ANY technical implementation details such as:
514 | - Exact index definitions
515 | - Internal PostgreSQL settings
516 | - Complex query rewrites
517 | - Table partitioning details
518 |
519 | Keep the response focused on high-level changes and performance metrics.
520 | </do_not_include>
521 |
522 | <example>
523 | I've analyzed your query and found potential improvements that could reduce execution time by [X]%.
524 | Would you like to apply these changes to improve performance?
525 |
526 | Analysis Details:
527 | - Tuning ID: [id]
528 | - Temporary Branch: [name]
529 | - Branch ID: [id]
530 | - Original Cost: [cost]
531 | - Improved Cost: [cost]
532 | - Referenced Tables:
533 | * public.orders
534 | * public.customers
535 | - Suggested Changes:
536 | * Add index for frequently filtered columns
537 | * Optimize join conditions
538 |
539 | To apply these changes, I will use the \`complete_query_tuning\` tool after your approval and pass the \`tuning_id\`, NOT the temporary branch ID to it.
540 | </example>
541 | </response_instructions>
542 |
543 | 3. If approved, use ONLY the \`complete_query_tuning\` tool with the \`tuning_id\`
544 | </next_steps>
545 |
546 | <error_handling>
547 | On error, the tool will:
548 | 1. Automatically attempt ONE retry of the exact same operation
549 | 2. If the retry fails:
550 | - Terminate execution
551 | - Return error details
552 | - Clean up temporary branch
553 | - DO NOT attempt any other tools or alternatives
554 |
555 | Error response will include:
556 | - Original error details
557 | - Confirmation that retry was attempted
558 | - Final error state
559 |
560 | Important: After a failed retry, you must terminate the current flow completely.
561 | </error_handling>
562 | `,
563 | inputSchema: prepareQueryTuningInputSchema,
564 | },
565 | {
566 | name: 'complete_query_tuning' as const,
567 | description: `Complete a query tuning session by either applying the changes to the main branch or discarding them.
568 | <important_notes>
569 | BEFORE RUNNING THIS TOOL: test out the changes in the temporary branch first by running
570 | - \`run_sql\` with the suggested DDL statements.
571 | - \`explain_sql_statement\` with the original query and the temporary branch.
572 | This tool is the ONLY way to finally apply changes after the \`prepare_query_tuning\` tool to the main branch.
573 | You MUST NOT use \`prepare_database_migration\` or other tools to apply query tuning changes.
574 | You MUST pass the \`tuning_id\` obtained from the \`prepare_query_tuning\` tool, NOT the temporary branch ID as \`tuning_id\` to this tool.
575 | You MUST pass the temporary branch ID used in the \`prepare_query_tuning\` tool as TEMPORARY branchId to this tool.
576 | The tool OPTIONALLY receives a second branch ID or name which can be used instead of the main branch to apply the changes.
577 | This tool MUST be called after tool \`prepare_query_tuning\` even when the user rejects the changes, to ensure proper cleanup of temporary branches.
578 | </important_notes>
579 |
580 | This tool:
581 | 1. Applies suggested changes (like creating indexes) to the main branch (or specified branch) if approved
582 | 2. Handles cleanup of temporary branch
583 | 3. Must be called even when changes are rejected to ensure proper cleanup
584 |
585 | Workflow:
586 | 1. After \`prepare_query_tuning\` suggests changes
587 | 2. User reviews and approves/rejects changes
588 | 3. This tool is called to either:
589 | - Apply approved changes to main branch and cleanup
590 | - OR just cleanup if changes are rejected
591 | `,
592 | inputSchema: completeQueryTuningInputSchema,
593 | },
594 | {
595 | name: 'list_slow_queries' as const,
596 | description: `
597 | <use_case>
598 | Use this tool to list slow queries from your Neon database.
599 | </use_case>
600 |
601 | <important_notes>
602 | This tool queries the pg_stat_statements extension to find queries that are taking longer than expected.
603 | The tool will return queries sorted by execution time, with the slowest queries first.
604 | </important_notes>`,
605 | inputSchema: listSlowQueriesInputSchema,
606 | },
607 | {
608 | name: 'list_branch_computes' as const,
609 | description: 'Lists compute endpoints for a project or specific branch',
610 | inputSchema: listBranchComputesInputSchema,
611 | },
612 | {
613 | name: 'compare_database_schema' as const,
614 | description: `
615 | <use_case>
616 | Use this tool to compare the schema of a database between two branches.
617 | The output of the tool is a JSON object with one field: \`diff\`.
618 |
619 | <example>
620 | \`\`\`json
621 | {
622 | "diff": "--- a/neondb\n+++ b/neondb\n@@ -27,7 +27,10 @@\n \n CREATE TABLE public.users (\n id integer NOT NULL,\n- username character varying(50) NOT NULL\n+ username character varying(50) NOT NULL,\n+ is_deleted boolean DEFAULT false NOT NULL,\n+ created_at timestamp with time zone DEFAULT now() NOT NULL,\n+ updated_at timestamp with time zone\n );\n \n \n@@ -79,6 +82,13 @@\n \n \n --\n+-- Name: users_created_at_idx; Type: INDEX; Schema: public; Owner: neondb_owner\n+--\n+\n+CREATE INDEX users_created_at_idx ON public.users USING btree (created_at DESC) WHERE (is_deleted = false);\n+\n+\n+--\n -- Name: DEFAULT PRIVILEGES FOR SEQUENCES; Type: DEFAULT ACL; Schema: public; Owner: cloud_admin\n --\n \n"
623 | }
624 | \`\`\`
625 | </example>
626 |
627 | At this field you will find a difference between two schemas.
628 | The diff represents the changes required to make the parent branch schema match the child branch schema.
629 | The diff field contains a unified diff (git-style patch) as a string.
630 |
631 | You MUST be able to generate a zero-downtime migration from the diff and apply it to the parent branch.
632 | (This branch is a child and has a parent. You can get parent id just querying the branch details.)
633 | </use_case>
634 |
635 | <important_notes>
636 | To generate schema diff, you MUST SPECIFY the \`database_name\`.
637 | If \`database_name\` is not specified, you MUST fall back to the default database name: \`${NEON_DEFAULT_DATABASE_NAME}\`.
638 |
639 | You MUST TAKE INTO ACCOUNT the PostgreSQL version. The PostgreSQL version is the same for both branches.
640 | You MUST ASK user consent before running each generated SQL query.
641 | You SHOULD USE \`run_sql\` tool to run each generated SQL query.
642 | You SHOULD suggest creating a backup or point-in-time restore before running the migration.
643 | Generated queries change the schema of the parent branch and MIGHT BE dangerous to execute.
644 | Generated SQL migrations SHOULD be idempotent where possible (i.e., safe to run multiple times without failure) and include \`IF NOT EXISTS\` / \`IF EXISTS\` where applicable.
645 | You SHOULD recommend including comments in generated SQL linking back to diff hunks (e.g., \`-- from diff @@ -27,7 +27,10 @@\`) to make audits easier.
646 | Generated SQL should be reviewed for dependencies (e.g., foreign key order) before execution.
647 | </important_notes>
648 |
649 | <next_steps>
650 | After executing this tool, you MUST follow these steps:
651 | 1. Review the schema diff and suggest generating a zero-downtime migration.
652 | 2. Follow these instructions to respond to the client:
653 |
654 | <response_instructions>
655 | <instructions>
656 | Provide brief information about the changes:
657 | * Tables
658 | * Views
659 | * Indexes
660 | * Ownership
661 | * Constraints
662 | * Triggers
663 | * Policies
664 | * Extensions
665 | * Schemas
666 | * Sequences
667 | * Tablespaces
668 | * Users
669 | * Roles
670 | * Privileges
671 | </instructions>
672 | </response_instructions>
673 |
674 | 3. If a migration fails, you SHOULD guide the user on how to revert the schema changes, for example by using backups, point-in-time restore, or generating reverse SQL statements (if safe).
675 | </next_steps>
676 |
677 | This tool:
678 | 1. Generates a diff between the child branch and its parent.
679 | 2. Generates a SQL migration from the diff.
680 | 3. Suggest generating zero-downtime migration.
681 |
682 | <workflow>
683 | 1. User asks you to generate a diff between two branches.
684 | 2. You suggest generating a SQL migration from the diff.
685 | 3. Ensure the generated migration is zero-downtime; otherwise, warn the user.
686 | 4. You ensure that your suggested migration is also matching the PostgreSQL version.
687 | 5. You use \`run_sql\` tool to run each generated SQL query and ask the user consent before running it.
688 | Before requesting user consent, present a summary of all generated SQL statements along with their potential impact (e.g., table rewrites, lock risks, validation steps) so the user can make an informed decision.
689 | 6. Propose to rerun the schema diff tool one more time to ensure that the migration is applied correctly.
690 | 7. If the diff is empty, confirm that the parent schema now matches the child schema.
691 | 8. If the diff is not empty after migration, warn the user and assist in resolving the remaining differences.
692 | </workflow>
693 |
694 | <hints>
695 | <hint>
696 | Adding the column with a \`DEFAULT\` static value will not have any locks.
697 | But if the function is called that is not deterministic, it will have locks.
698 |
699 | <example>
700 | \`\`\`sql
701 | -- No table rewrite, minimal lock time
702 | ALTER TABLE users ADD COLUMN status text DEFAULT 'active';
703 | \`\`\`
704 | </example>
705 |
706 | There is an example of a case where the function is not deterministic and will have locks:
707 |
708 | <example>
709 | \`\`\`sql
710 | -- Table rewrite, potentially longer lock time
711 | ALTER TABLE users ADD COLUMN created_at timestamptz DEFAULT now();
712 | \`\`\`
713 |
714 | The fix for this is next:
715 |
716 | \`\`\`sql
717 | -- Adding a nullable column first
718 | ALTER TABLE users ADD COLUMN created_at timestamptz;
719 |
720 | -- Setting the default value because the rows are updated
721 | UPDATE users SET created_at = now();
722 | \`\`\`
723 | </example>
724 | </hint>
725 |
726 | <hint>
727 | Adding constraints in two phases (including foreign keys)
728 |
729 | <example>
730 | \`\`\`sql
731 | -- Step 1: Add constraint without validating existing data
732 | -- Fast - only blocks briefly to update catalog
733 | ALTER TABLE users ADD CONSTRAINT users_age_positive
734 | CHECK (age > 0) NOT VALID;
735 |
736 | -- Step 2: Validate existing data (can take time but doesn't block writes)
737 | -- Uses SHARE UPDATE EXCLUSIVE lock - allows reads/writes
738 | ALTER TABLE users VALIDATE CONSTRAINT users_age_positive;
739 | \`\`\`
740 | </example>
741 |
742 | <example>
743 | \`\`\`sql
744 | -- Step 1: Add foreign key without validation
745 | -- Fast - only updates catalog, doesn't validate existing data
746 | ALTER TABLE orders ADD CONSTRAINT orders_user_id_fk
747 | FOREIGN KEY (user_id) REFERENCES users(id) NOT VALID;
748 |
749 | -- Step 2: Validate existing relationships
750 | -- Can take time but allows concurrent operations
751 | ALTER TABLE orders VALIDATE CONSTRAINT orders_user_id_fk;
752 | \`\`\`
753 | </example>
754 | </hint>
755 |
756 | <hint>
757 | Setting columns to NOT NULL
758 |
759 | <example>
760 | \`\`\`sql
761 | -- Step 1: Add a check constraint (fast with NOT VALID)
762 | ALTER TABLE users ADD CONSTRAINT users_email_not_null
763 | CHECK (email IS NOT NULL) NOT VALID;
764 |
765 | -- Step 2: Validate the constraint (allows concurrent operations)
766 | ALTER TABLE users VALIDATE CONSTRAINT users_email_not_null;
767 |
768 | -- Step 3: Set NOT NULL (fast since constraint guarantees no nulls)
769 | ALTER TABLE users ALTER COLUMN email SET NOT NULL;
770 |
771 | -- Step 4: Drop the redundant check constraint
772 | ALTER TABLE users DROP CONSTRAINT users_email_not_null;
773 | \`\`\`
774 | </example>
775 |
776 | <example>
777 | For PostgreSQL v18+
778 | (to get PostgreSQL version, you can use \`describe_project\` tool or \`run_sql\` tool and execute \`SELECT version();\` query)
779 |
780 | \`\`\`sql
781 | -- PostgreSQL 18+ - Simplified approach
782 | ALTER TABLE users ALTER COLUMN email SET NOT NULL NOT VALID;
783 | ALTER TABLE users VALIDATE CONSTRAINT users_email_not_null;
784 | \`\`\`
785 | </example>
786 | </hint>
787 |
788 | <hint>
789 | In some cases, you need to combine two approaches to achieve a zero-downtime migration.
790 |
791 | <example>
792 | \`\`\`sql
793 | -- Step 1: Adding a nullable column first
794 | ALTER TABLE users ADD COLUMN created_at timestamptz;
795 |
796 | -- Step 2: Updating the all rows with the default value
797 | UPDATE users SET created_at = now() WHERE created_at IS NULL;
798 |
799 | -- Step 3: Creating a not null constraint
800 | ALTER TABLE users ADD CONSTRAINT users_created_at_not_null
801 | CHECK (created_at IS NOT NULL) NOT VALID;
802 |
803 | -- Step 4: Validating the constraint
804 | ALTER TABLE users VALIDATE CONSTRAINT users_created_at_not_null;
805 |
806 | -- Step 5: Setting the column to NOT NULL
807 | ALTER TABLE users ALTER COLUMN created_at SET NOT NULL;
808 |
809 | -- Step 6: Dropping the redundant NOT NULL constraint
810 | ALTER TABLE users DROP CONSTRAINT users_created_at_not_null;
811 |
812 | -- Step 7: Adding the default value
813 | ALTER TABLE users ALTER COLUMN created_at SET DEFAULT now();
814 | \`\`\`
815 | </example>
816 |
817 | For PostgreSQL v18+
818 | <example>
819 | \`\`\`sql
820 | -- Step 1: Adding a nullable column first
821 | ALTER TABLE users ADD COLUMN created_at timestamptz;
822 |
823 | -- Step 2: Updating the all rows with the default value
824 | UPDATE users SET created_at = now() WHERE created_at IS NULL;
825 |
826 | -- Step 3: Creating a not null constraint
827 | ALTER TABLE users ALTER COLUMN created_at SET NOT NULL NOT VALID;
828 |
829 | -- Step 4: Validating the constraint
830 | ALTER TABLE users VALIDATE CONSTRAINT users_created_at_not_null;
831 |
832 | -- Step 5: Adding the default value
833 | ALTER TABLE users ALTER COLUMN created_at SET DEFAULT now();
834 | \`\`\`
835 | </example>
836 | </hint>
837 |
838 | <hint>
839 | Create index CONCURRENTLY
840 |
841 | <example>
842 | \`\`\`sql
843 | CREATE INDEX CONCURRENTLY idx_users_email ON users (email);
844 | \`\`\`
845 | </example>
846 | </hint>
847 |
848 | <hint>
849 | Drop index CONCURRENTLY
850 |
851 | <example>
852 | \`\`\`sql
853 | DROP INDEX CONCURRENTLY idx_users_email;
854 | \`\`\`
855 | </example>
856 | </hint>
857 |
858 | <hint>
859 | Create materialized view WITH NO DATA
860 |
861 | <example>
862 | \`\`\`sql
863 | CREATE MATERIALIZED VIEW mv_users AS SELECT name FROM users WITH NO DATA;
864 | \`\`\`
865 | </example>
866 | </hint>
867 |
868 | <hint>
869 | Refresh materialized view CONCURRENTLY
870 |
871 | <example>
872 | \`\`\`sql
873 | REFRESH MATERIALIZED VIEW CONCURRENTLY mv_users;
874 | \`\`\`
875 | </example>
876 | </hint>
877 | </hints>
878 | `,
879 | inputSchema: compareDatabaseSchemaInputSchema,
880 | },
881 | ];
882 |
```
--------------------------------------------------------------------------------
/src/tools/tools.ts:
--------------------------------------------------------------------------------
```typescript
1 | import {
2 | Api,
3 | Branch,
4 | EndpointType,
5 | ListProjectsParams,
6 | ListSharedProjectsParams,
7 | GetProjectBranchSchemaComparisonParams,
8 | Organization,
9 | ProjectCreateRequest,
10 | } from '@neondatabase/api-client';
11 | import { neon } from '@neondatabase/serverless';
12 | import crypto from 'crypto';
13 | import { InvalidArgumentError, NotFoundError } from '../server/errors.js';
14 |
15 | import { describeTable, formatTableDescription } from '../describeUtils.js';
16 | import { handleProvisionNeonAuth } from './handlers/neon-auth.js';
17 | import { getMigrationFromMemory, persistMigrationToMemory } from './state.js';
18 |
19 | import {
20 | DESCRIBE_DATABASE_STATEMENTS,
21 | getDefaultDatabase,
22 | splitSqlStatements,
23 | getOrgByOrgIdOrDefault,
24 | filterOrganizations,
25 | resolveBranchId,
26 | } from './utils.js';
27 | import { startSpan } from '@sentry/node';
28 | import { ToolHandlerExtraParams, ToolHandlers } from './types.js';
29 |
30 | async function handleListProjects(
31 | params: ListProjectsParams,
32 | neonClient: Api<unknown>,
33 | extra: ToolHandlerExtraParams,
34 | ) {
35 | const organization = await getOrgByOrgIdOrDefault(params, neonClient, extra);
36 |
37 | const response = await neonClient.listProjects({
38 | ...params,
39 | org_id: organization?.id,
40 | });
41 | if (response.status !== 200) {
42 | throw new Error(`Failed to list projects: ${response.statusText}`);
43 | }
44 |
45 | let projects = response.data.projects;
46 |
47 | // If search is provided and no org_id specified, and no projects found in personal account,
48 | // search across all user organizations
49 | if (params.search && !params.org_id && projects.length === 0) {
50 | const organizations = await handleListOrganizations(
51 | neonClient,
52 | extra.account,
53 | );
54 |
55 | // Search projects across all organizations
56 | const allProjects = [];
57 | for (const org of organizations) {
58 | // Skip the default organization
59 | if (organization?.id === org.id) {
60 | continue;
61 | }
62 |
63 | const orgResponse = await neonClient.listProjects({
64 | ...params,
65 | org_id: org.id,
66 | });
67 | if (orgResponse.status === 200) {
68 | allProjects.push(...orgResponse.data.projects);
69 | }
70 | }
71 |
72 | // If we found projects in other organizations, return them
73 | if (allProjects.length > 0) {
74 | projects = allProjects;
75 | }
76 | }
77 |
78 | return projects;
79 | }
80 |
81 | async function handleCreateProject(
82 | params: ProjectCreateRequest,
83 | neonClient: Api<unknown>,
84 | ) {
85 | const response = await neonClient.createProject(params);
86 | if (response.status !== 201) {
87 | throw new Error(`Failed to create project: ${JSON.stringify(response)}`);
88 | }
89 | return response.data;
90 | }
91 |
92 | async function handleDeleteProject(
93 | projectId: string,
94 | neonClient: Api<unknown>,
95 | ) {
96 | const response = await neonClient.deleteProject(projectId);
97 | if (response.status !== 200) {
98 | throw new Error(`Failed to delete project: ${response.statusText}`);
99 | }
100 | return response.data;
101 | }
102 |
103 | async function handleDescribeProject(
104 | projectId: string,
105 | neonClient: Api<unknown>,
106 | ) {
107 | const projectBranches = await neonClient.listProjectBranches({
108 | projectId: projectId,
109 | });
110 | const projectDetails = await neonClient.getProject(projectId);
111 | if (projectBranches.status !== 200) {
112 | throw new Error(
113 | `Failed to get project branches: ${projectBranches.statusText}`,
114 | );
115 | }
116 | if (projectDetails.status !== 200) {
117 | throw new Error(`Failed to get project: ${projectDetails.statusText}`);
118 | }
119 | return {
120 | branches: projectBranches.data,
121 | project: projectDetails.data,
122 | };
123 | }
124 |
125 | async function handleRunSql(
126 | {
127 | sql,
128 | databaseName,
129 | projectId,
130 | branchId,
131 | }: {
132 | sql: string;
133 | databaseName?: string;
134 | projectId: string;
135 | branchId?: string;
136 | },
137 | neonClient: Api<unknown>,
138 | extra: ToolHandlerExtraParams,
139 | ) {
140 | return await startSpan({ name: 'run_sql' }, async () => {
141 | const connectionString = await handleGetConnectionString(
142 | {
143 | projectId,
144 | branchId,
145 | databaseName,
146 | },
147 | neonClient,
148 | extra,
149 | );
150 | const runQuery = neon(connectionString.uri);
151 | const response = await runQuery.query(sql);
152 |
153 | return response;
154 | });
155 | }
156 |
157 | async function handleRunSqlTransaction(
158 | {
159 | sqlStatements,
160 | databaseName,
161 | projectId,
162 | branchId,
163 | }: {
164 | sqlStatements: string[];
165 | databaseName?: string;
166 | projectId: string;
167 | branchId?: string;
168 | },
169 | neonClient: Api<unknown>,
170 | extra: ToolHandlerExtraParams,
171 | ) {
172 | const connectionString = await handleGetConnectionString(
173 | {
174 | projectId,
175 | branchId,
176 | databaseName,
177 | },
178 | neonClient,
179 | extra,
180 | );
181 | const runQuery = neon(connectionString.uri);
182 | const response = await runQuery.transaction(
183 | sqlStatements.map((sql) => runQuery.query(sql)),
184 | );
185 |
186 | return response;
187 | }
188 |
189 | async function handleGetDatabaseTables(
190 | {
191 | projectId,
192 | databaseName,
193 | branchId,
194 | }: {
195 | projectId: string;
196 | databaseName?: string;
197 | branchId?: string;
198 | },
199 | neonClient: Api<unknown>,
200 | extra: ToolHandlerExtraParams,
201 | ) {
202 | const connectionString = await handleGetConnectionString(
203 | {
204 | projectId,
205 | branchId,
206 | databaseName,
207 | },
208 | neonClient,
209 | extra,
210 | );
211 | const runQuery = neon(connectionString.uri);
212 | const query = `
213 | SELECT
214 | table_schema,
215 | table_name,
216 | table_type
217 | FROM information_schema.tables
218 | WHERE table_schema NOT IN ('pg_catalog', 'information_schema')
219 | ORDER BY table_schema, table_name;
220 | `;
221 |
222 | const tables = await runQuery.query(query);
223 | return tables;
224 | }
225 |
226 | async function handleDescribeTableSchema(
227 | {
228 | projectId,
229 | databaseName,
230 | branchId,
231 | tableName,
232 | }: {
233 | projectId: string;
234 | databaseName?: string;
235 | branchId?: string;
236 | tableName: string;
237 | },
238 | neonClient: Api<unknown>,
239 | extra: ToolHandlerExtraParams,
240 | ) {
241 | const connectionString = await handleGetConnectionString(
242 | {
243 | projectId,
244 | branchId,
245 | databaseName,
246 | },
247 | neonClient,
248 | extra,
249 | );
250 |
251 | // Extract table name without schema if schema-qualified
252 | const tableNameParts = tableName.split('.');
253 | const simpleTableName = tableNameParts[tableNameParts.length - 1];
254 |
255 | const description = await describeTable(
256 | connectionString.uri,
257 | simpleTableName,
258 | );
259 | return {
260 | raw: description,
261 | formatted: formatTableDescription(description),
262 | };
263 | }
264 |
265 | async function handleCreateBranch(
266 | {
267 | projectId,
268 | branchName,
269 | }: {
270 | projectId: string;
271 | branchName?: string;
272 | },
273 | neonClient: Api<unknown>,
274 | ) {
275 | const response = await neonClient.createProjectBranch(projectId, {
276 | branch: {
277 | name: branchName,
278 | },
279 | endpoints: [
280 | {
281 | type: EndpointType.ReadWrite,
282 | autoscaling_limit_min_cu: 0.25,
283 | autoscaling_limit_max_cu: 0.25,
284 | provisioner: 'k8s-neonvm',
285 | },
286 | ],
287 | });
288 |
289 | if (response.status !== 201) {
290 | throw new Error(`Failed to create branch: ${response.statusText}`);
291 | }
292 |
293 | return response.data;
294 | }
295 |
296 | async function handleDeleteBranch(
297 | {
298 | projectId,
299 | branchId,
300 | }: {
301 | projectId: string;
302 | branchId: string;
303 | },
304 | neonClient: Api<unknown>,
305 | ) {
306 | const response = await neonClient.deleteProjectBranch(projectId, branchId);
307 | return response.data;
308 | }
309 |
310 | async function handleResetFromParent(
311 | {
312 | projectId,
313 | branchIdOrName,
314 | preserveUnderName,
315 | }: {
316 | projectId: string;
317 | branchIdOrName: string;
318 | preserveUnderName?: string;
319 | },
320 | neonClient: Api<unknown>,
321 | ) {
322 | // Resolve branch name or ID to actual branch ID and get all branches in one call
323 | const { branchId: resolvedBranchId, branches } = await resolveBranchId(
324 | branchIdOrName,
325 | projectId,
326 | neonClient,
327 | );
328 |
329 | const branch = branches.find((b) => b.id === resolvedBranchId);
330 | if (!branch) {
331 | throw new NotFoundError(
332 | `Branch "${branchIdOrName}" not found in project ${projectId}`,
333 | );
334 | }
335 |
336 | // Find the parent branch and validate it exists
337 | const parentBranch = branch.parent_id
338 | ? branches.find((b) => b.id === branch.parent_id)
339 | : undefined;
340 |
341 | if (!parentBranch) {
342 | throw new InvalidArgumentError(
343 | `Branch "${branchIdOrName}" does not have a parent branch and cannot be reset`,
344 | );
345 | }
346 |
347 | // Check if the branch has children
348 | const hasChildren = branches.some((b) => b.parent_id === resolvedBranchId);
349 |
350 | // Auto-generate preserve name if branch has children and none was provided
351 | let finalPreserveName = preserveUnderName;
352 | if (hasChildren && !preserveUnderName) {
353 | const timestamp = new Date()
354 | .toISOString()
355 | .replace(/[:.]/g, '-')
356 | .slice(0, -5);
357 | finalPreserveName = `${branch.name}_old_${timestamp}`;
358 | }
359 |
360 | // Call the restoreProjectBranch API
361 | const response = await neonClient.restoreProjectBranch(
362 | projectId,
363 | resolvedBranchId,
364 | {
365 | source_branch_id: parentBranch.id,
366 | preserve_under_name: finalPreserveName,
367 | },
368 | );
369 |
370 | return {
371 | ...response.data,
372 | preservedBranchName: finalPreserveName,
373 | parentBranch,
374 | };
375 | }
376 |
377 | async function handleGetConnectionString(
378 | {
379 | projectId,
380 | branchId,
381 | computeId,
382 | databaseName,
383 | roleName,
384 | }: {
385 | projectId?: string;
386 | branchId?: string;
387 | computeId?: string;
388 | databaseName?: string;
389 | roleName?: string;
390 | },
391 | neonClient: Api<unknown>,
392 | extra: ToolHandlerExtraParams,
393 | ) {
394 | return await startSpan(
395 | {
396 | name: 'get_connection_string',
397 | },
398 | async () => {
399 | // If projectId is not provided, get the first project but only if there is only one project
400 | if (!projectId) {
401 | const projects = await handleListProjects({}, neonClient, extra);
402 | if (projects.length === 1) {
403 | projectId = projects[0].id;
404 | } else {
405 | throw new NotFoundError(
406 | 'Please provide a project ID or ensure you have only one project in your account.',
407 | );
408 | }
409 | }
410 |
411 | if (!branchId) {
412 | const branches = await neonClient.listProjectBranches({
413 | projectId,
414 | });
415 | const defaultBranch = branches.data.branches.find(
416 | (branch) => branch.default,
417 | );
418 | if (defaultBranch) {
419 | branchId = defaultBranch.id;
420 | } else {
421 | throw new NotFoundError(
422 | 'No default branch found in this project. Please provide a branch ID.',
423 | );
424 | }
425 | }
426 |
427 | // If databaseName is not provided, use default `neondb` or first database
428 | let dbObject;
429 | if (!databaseName) {
430 | dbObject = await getDefaultDatabase(
431 | {
432 | projectId,
433 | branchId,
434 | databaseName,
435 | },
436 | neonClient,
437 | );
438 | databaseName = dbObject.name;
439 |
440 | if (!roleName) {
441 | roleName = dbObject.owner_name;
442 | }
443 | } else if (!roleName) {
444 | const { data } = await neonClient.getProjectBranchDatabase(
445 | projectId,
446 | branchId,
447 | databaseName,
448 | );
449 | roleName = data.database.owner_name;
450 | }
451 |
452 | // Get connection URI with the provided parameters
453 | const connectionString = await neonClient.getConnectionUri({
454 | projectId,
455 | role_name: roleName,
456 | database_name: databaseName,
457 | branch_id: branchId,
458 | endpoint_id: computeId,
459 | });
460 |
461 | return {
462 | uri: connectionString.data.uri,
463 | projectId,
464 | branchId,
465 | databaseName,
466 | roleName,
467 | computeId,
468 | };
469 | },
470 | );
471 | }
472 |
473 | async function handleSchemaMigration(
474 | {
475 | migrationSql,
476 | databaseName,
477 | projectId,
478 | }: {
479 | databaseName?: string;
480 | projectId: string;
481 | migrationSql: string;
482 | },
483 | neonClient: Api<unknown>,
484 | extra: ToolHandlerExtraParams,
485 | ) {
486 | return await startSpan({ name: 'prepare_schema_migration' }, async (span) => {
487 | const newBranch = await handleCreateBranch({ projectId }, neonClient);
488 |
489 | if (!databaseName) {
490 | const dbObject = await getDefaultDatabase(
491 | {
492 | projectId,
493 | branchId: newBranch.branch.id,
494 | databaseName,
495 | },
496 | neonClient,
497 | );
498 | databaseName = dbObject.name;
499 | }
500 |
501 | const result = await handleRunSqlTransaction(
502 | {
503 | sqlStatements: splitSqlStatements(migrationSql),
504 | databaseName,
505 | projectId,
506 | branchId: newBranch.branch.id,
507 | },
508 | neonClient,
509 | extra,
510 | );
511 |
512 | const migrationId = crypto.randomUUID();
513 | span.setAttributes({
514 | projectId,
515 | migrationId,
516 | });
517 | persistMigrationToMemory(migrationId, {
518 | migrationSql,
519 | databaseName,
520 | appliedBranch: newBranch.branch,
521 | });
522 |
523 | return {
524 | branch: newBranch.branch,
525 | migrationId,
526 | migrationResult: result,
527 | };
528 | });
529 | }
530 |
531 | async function handleCommitMigration(
532 | { migrationId }: { migrationId: string },
533 | neonClient: Api<unknown>,
534 | extra: ToolHandlerExtraParams,
535 | ) {
536 | return await startSpan({ name: 'commit_schema_migration' }, async (span) => {
537 | span.setAttributes({
538 | migrationId,
539 | });
540 | const migration = getMigrationFromMemory(migrationId);
541 | if (!migration) {
542 | throw new Error(`Migration not found: ${migrationId}`);
543 | }
544 |
545 | span.setAttributes({
546 | projectId: migration.appliedBranch.project_id,
547 | });
548 | const result = await handleRunSqlTransaction(
549 | {
550 | sqlStatements: splitSqlStatements(migration.migrationSql),
551 | databaseName: migration.databaseName,
552 | projectId: migration.appliedBranch.project_id,
553 | branchId: migration.appliedBranch.parent_id,
554 | },
555 | neonClient,
556 | extra,
557 | );
558 |
559 | await handleDeleteBranch(
560 | {
561 | projectId: migration.appliedBranch.project_id,
562 | branchId: migration.appliedBranch.id,
563 | },
564 | neonClient,
565 | );
566 |
567 | return {
568 | deletedBranch: migration.appliedBranch,
569 | migrationResult: result,
570 | };
571 | });
572 | }
573 |
574 | async function handleDescribeBranch(
575 | {
576 | projectId,
577 | databaseName,
578 | branchId,
579 | }: {
580 | projectId: string;
581 | databaseName?: string;
582 | branchId?: string;
583 | },
584 | neonClient: Api<unknown>,
585 | extra: ToolHandlerExtraParams,
586 | ) {
587 | const connectionString = await handleGetConnectionString(
588 | {
589 | projectId,
590 | branchId,
591 | databaseName,
592 | },
593 | neonClient,
594 | extra,
595 | );
596 | const runQuery = neon(connectionString.uri);
597 | const response = await runQuery.transaction(
598 | DESCRIBE_DATABASE_STATEMENTS.map((sql) => runQuery.query(sql)),
599 | );
600 |
601 | return response;
602 | }
603 |
604 | async function handleExplainSqlStatement(
605 | {
606 | params,
607 | }: {
608 | params: {
609 | sql: string;
610 | databaseName?: string;
611 | projectId: string;
612 | branchId?: string;
613 | analyze: boolean;
614 | };
615 | },
616 | neonClient: Api<unknown>,
617 | extra: ToolHandlerExtraParams,
618 | ) {
619 | const explainPrefix = params.analyze
620 | ? 'EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FILECACHE, FORMAT JSON)'
621 | : 'EXPLAIN (VERBOSE, FORMAT JSON)';
622 |
623 | const explainSql = `${explainPrefix} ${params.sql}`;
624 |
625 | const result = await handleRunSql(
626 | {
627 | sql: explainSql,
628 | databaseName: params.databaseName,
629 | projectId: params.projectId,
630 | branchId: params.branchId,
631 | },
632 | neonClient,
633 | extra,
634 | );
635 |
636 | return {
637 | content: [
638 | {
639 | type: 'text' as const,
640 | text: JSON.stringify(result, null, 2),
641 | },
642 | ],
643 | };
644 | }
645 |
646 | async function createTemporaryBranch(
647 | projectId: string,
648 | neonClient: Api<unknown>,
649 | ): Promise<{ branch: Branch }> {
650 | const result = await handleCreateBranch({ projectId }, neonClient);
651 | if (!result?.branch) {
652 | throw new Error('Failed to create temporary branch');
653 | }
654 | return result;
655 | }
656 |
657 | type QueryTuningParams = {
658 | sql: string;
659 | databaseName: string;
660 | projectId: string;
661 | };
662 |
663 | type CompleteTuningParams = {
664 | suggestedSqlStatements?: string[];
665 | applyChanges?: boolean;
666 | tuningId: string;
667 | databaseName: string;
668 | projectId: string;
669 | temporaryBranch: Branch;
670 | shouldDeleteTemporaryBranch?: boolean;
671 | branch?: Branch;
672 | };
673 |
674 | type QueryTuningResult = {
675 | tuningId: string;
676 | databaseName: string;
677 | projectId: string;
678 | temporaryBranch: Branch;
679 | originalPlan: any;
680 | tableSchemas: any[];
681 | sql: string;
682 | baselineMetrics: QueryMetrics;
683 | };
684 |
685 | type CompleteTuningResult = {
686 | appliedChanges?: string[];
687 | results?: any;
688 | deletedBranches?: string[];
689 | message: string;
690 | };
691 |
692 | async function handleQueryTuning(
693 | params: QueryTuningParams,
694 | neonClient: Api<unknown>,
695 | extra: ToolHandlerExtraParams,
696 | ): Promise<QueryTuningResult> {
697 | let tempBranch: Branch | undefined;
698 | const tuningId = crypto.randomUUID();
699 |
700 | try {
701 | // Create temporary branch
702 | const newBranch = await createTemporaryBranch(params.projectId, neonClient);
703 | if (!newBranch.branch) {
704 | throw new Error('Failed to create temporary branch: branch is undefined');
705 | }
706 | tempBranch = newBranch.branch;
707 |
708 | // Ensure all operations use the temporary branch
709 | const branchParams = {
710 | ...params,
711 | branchId: tempBranch.id,
712 | };
713 |
714 | // First, get the execution plan with table information
715 | const executionPlan = await handleExplainSqlStatement(
716 | {
717 | params: {
718 | sql: branchParams.sql,
719 | databaseName: branchParams.databaseName,
720 | projectId: branchParams.projectId,
721 | branchId: tempBranch.id,
722 | analyze: true,
723 | },
724 | },
725 | neonClient,
726 | extra,
727 | );
728 |
729 | // Extract table names from the plan
730 | const tableNames = extractTableNamesFromPlan(executionPlan);
731 |
732 | if (tableNames.length === 0) {
733 | throw new NotFoundError(
734 | 'No tables found in execution plan. Cannot proceed with optimization.',
735 | );
736 | }
737 |
738 | // Get schema information for all referenced tables in parallel
739 | const tableSchemas = await Promise.all(
740 | tableNames.map(async (tableName) => {
741 | try {
742 | const schema = await handleDescribeTableSchema(
743 | {
744 | tableName,
745 | databaseName: branchParams.databaseName,
746 | projectId: branchParams.projectId,
747 | branchId: newBranch.branch.id,
748 | },
749 | neonClient,
750 | extra,
751 | );
752 | return {
753 | tableName,
754 | schema: schema.raw,
755 | formatted: schema.formatted,
756 | };
757 | } catch (error) {
758 | throw new Error(
759 | `Failed to get schema for table ${tableName}: ${(error as Error).message}`,
760 | );
761 | }
762 | }),
763 | );
764 |
765 | // Get the baseline execution metrics
766 | const baselineMetrics = extractExecutionMetrics(executionPlan);
767 |
768 | // Return the information for analysis
769 | const result: QueryTuningResult = {
770 | tuningId,
771 | databaseName: params.databaseName,
772 | projectId: params.projectId,
773 | temporaryBranch: tempBranch,
774 | originalPlan: executionPlan,
775 | tableSchemas,
776 | sql: params.sql,
777 | baselineMetrics,
778 | };
779 |
780 | return result;
781 | } catch (error) {
782 | // Always attempt to clean up the temporary branch if it was created
783 | if (tempBranch) {
784 | try {
785 | await handleDeleteBranch(
786 | {
787 | projectId: params.projectId,
788 | branchId: tempBranch.id,
789 | },
790 | neonClient,
791 | );
792 | } catch {
793 | // No need to handle cleanup error
794 | }
795 | }
796 |
797 | throw error;
798 | }
799 | }
800 |
801 | // Helper function to extract execution metrics from EXPLAIN output
802 | function extractExecutionMetrics(plan: any): QueryMetrics {
803 | try {
804 | const planJson =
805 | typeof plan.content?.[0]?.text === 'string'
806 | ? JSON.parse(plan.content[0].text)
807 | : plan;
808 |
809 | const metrics: QueryMetrics = {
810 | executionTime: 0,
811 | planningTime: 0,
812 | totalCost: 0,
813 | actualRows: 0,
814 | bufferUsage: {
815 | shared: { hit: 0, read: 0, written: 0, dirtied: 0 },
816 | local: { hit: 0, read: 0, written: 0, dirtied: 0 },
817 | },
818 | };
819 |
820 | // Extract planning and execution time if available
821 | if (planJson?.[0]?.['Planning Time']) {
822 | metrics.planningTime = planJson[0]['Planning Time'];
823 | }
824 | if (planJson?.[0]?.['Execution Time']) {
825 | metrics.executionTime = planJson[0]['Execution Time'];
826 | }
827 |
828 | // Recursively process plan nodes to accumulate costs and buffer usage
829 | function processNode(node: any) {
830 | if (!node || typeof node !== 'object') return;
831 |
832 | // Accumulate costs
833 | if (node['Total Cost']) {
834 | metrics.totalCost = Math.max(metrics.totalCost, node['Total Cost']);
835 | }
836 | if (node['Actual Rows']) {
837 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
838 | metrics.actualRows += node['Actual Rows'];
839 | }
840 |
841 | if (node['Shared Hit Blocks'])
842 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
843 | metrics.bufferUsage.shared.hit += node['Shared Hit Blocks'];
844 | if (node['Shared Read Blocks'])
845 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
846 | metrics.bufferUsage.shared.read += node['Shared Read Blocks'];
847 | if (node['Shared Written Blocks'])
848 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
849 | metrics.bufferUsage.shared.written += node['Shared Written Blocks'];
850 | if (node['Shared Dirtied Blocks'])
851 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
852 | metrics.bufferUsage.shared.dirtied += node['Shared Dirtied Blocks'];
853 |
854 | if (node['Local Hit Blocks'])
855 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
856 | metrics.bufferUsage.local.hit += node['Local Hit Blocks'];
857 | if (node['Local Read Blocks'])
858 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
859 | metrics.bufferUsage.local.read += node['Local Read Blocks'];
860 | if (node['Local Written Blocks'])
861 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
862 | metrics.bufferUsage.local.written += node['Local Written Blocks'];
863 | if (node['Local Dirtied Blocks'])
864 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
865 | metrics.bufferUsage.local.dirtied += node['Local Dirtied Blocks'];
866 |
867 | // Process child nodes
868 | if (Array.isArray(node.Plans)) {
869 | node.Plans.forEach(processNode);
870 | }
871 | }
872 |
873 | if (planJson?.[0]?.Plan) {
874 | processNode(planJson[0].Plan);
875 | }
876 |
877 | return metrics;
878 | } catch {
879 | return {
880 | executionTime: 0,
881 | planningTime: 0,
882 | totalCost: 0,
883 | actualRows: 0,
884 | bufferUsage: {
885 | shared: { hit: 0, read: 0, written: 0, dirtied: 0 },
886 | local: { hit: 0, read: 0, written: 0, dirtied: 0 },
887 | },
888 | };
889 | }
890 | }
891 |
892 | // Types for query metrics
893 | type BufferMetrics = {
894 | hit: number;
895 | read: number;
896 | written: number;
897 | dirtied: number;
898 | };
899 |
900 | type QueryMetrics = {
901 | executionTime: number;
902 | planningTime: number;
903 | totalCost: number;
904 | actualRows: number;
905 | bufferUsage: {
906 | shared: BufferMetrics;
907 | local: BufferMetrics;
908 | };
909 | };
910 |
911 | // Function to extract table names from an execution plan
912 | function extractTableNamesFromPlan(planResult: any): string[] {
913 | const tableNames = new Set<string>();
914 |
915 | function recursivelyExtractFromNode(node: any) {
916 | if (!node || typeof node !== 'object') return;
917 |
918 | // Check if current node has relation information
919 | if (node['Relation Name'] && node.Schema) {
920 | const tableName = `${node.Schema}.${node['Relation Name']}`;
921 | tableNames.add(tableName);
922 | }
923 |
924 | // Recursively process all object properties and array elements
925 | if (Array.isArray(node)) {
926 | node.forEach((item) => {
927 | recursivelyExtractFromNode(item);
928 | });
929 | } else {
930 | Object.values(node).forEach((value) => {
931 | recursivelyExtractFromNode(value);
932 | });
933 | }
934 | }
935 |
936 | try {
937 | // Start with the raw plan result
938 | recursivelyExtractFromNode(planResult);
939 |
940 | // If we have content[0].text, also parse and process that
941 | if (planResult?.content?.[0]?.text) {
942 | try {
943 | const parsedContent = JSON.parse(planResult.content[0].text);
944 | recursivelyExtractFromNode(parsedContent);
945 | } catch {
946 | // No need to handle parse error
947 | }
948 | }
949 | } catch {
950 | // No need to handle extraction error
951 | }
952 |
953 | const result = Array.from(tableNames);
954 | return result;
955 | }
956 |
957 | async function handleCompleteTuning(
958 | params: CompleteTuningParams,
959 | neonClient: Api<unknown>,
960 | extra: ToolHandlerExtraParams,
961 | ): Promise<CompleteTuningResult> {
962 | let results;
963 | const operationLog: string[] = [];
964 |
965 | try {
966 | // Validate branch information
967 | if (!params.temporaryBranch) {
968 | throw new Error(
969 | 'Branch information is required for completing query tuning',
970 | );
971 | }
972 |
973 | // Only proceed with changes if we have both suggestedChanges and branch
974 | if (
975 | params.applyChanges &&
976 | params.suggestedSqlStatements &&
977 | params.suggestedSqlStatements.length > 0
978 | ) {
979 | operationLog.push('Applying optimizations to main branch...');
980 |
981 | results = await handleRunSqlTransaction(
982 | {
983 | sqlStatements: params.suggestedSqlStatements,
984 | databaseName: params.databaseName,
985 | projectId: params.projectId,
986 | branchId: params.branch?.id,
987 | },
988 | neonClient,
989 | extra,
990 | );
991 |
992 | operationLog.push('Successfully applied optimizations to main branch.');
993 | } else {
994 | operationLog.push(
995 | 'No changes were applied (either none suggested or changes were discarded).',
996 | );
997 | }
998 |
999 | // Only delete branch if shouldDeleteTemporaryBranch is true
1000 | if (params.shouldDeleteTemporaryBranch && params.temporaryBranch) {
1001 | operationLog.push('Cleaning up temporary branch...');
1002 |
1003 | await handleDeleteBranch(
1004 | {
1005 | projectId: params.projectId,
1006 | branchId: params.temporaryBranch.id,
1007 | },
1008 | neonClient,
1009 | );
1010 |
1011 | operationLog.push('Successfully cleaned up temporary branch.');
1012 | }
1013 |
1014 | const result: CompleteTuningResult = {
1015 | appliedChanges:
1016 | params.applyChanges && params.suggestedSqlStatements
1017 | ? params.suggestedSqlStatements
1018 | : undefined,
1019 | results,
1020 | deletedBranches:
1021 | params.shouldDeleteTemporaryBranch && params.temporaryBranch
1022 | ? [params.temporaryBranch.id]
1023 | : undefined,
1024 | message: operationLog.join('\n'),
1025 | };
1026 |
1027 | return result;
1028 | } catch (error) {
1029 | throw new Error(
1030 | `Failed to complete query tuning: ${(error as Error).message}`,
1031 | );
1032 | }
1033 | }
1034 |
1035 | async function handleListSlowQueries(
1036 | {
1037 | projectId,
1038 | branchId,
1039 | databaseName,
1040 | computeId,
1041 | limit = 10,
1042 | }: {
1043 | projectId: string;
1044 | branchId?: string;
1045 | databaseName?: string;
1046 | computeId?: string;
1047 | limit?: number;
1048 | },
1049 | neonClient: Api<unknown>,
1050 | extra: ToolHandlerExtraParams,
1051 | ) {
1052 | // Get connection string
1053 | const connectionString = await handleGetConnectionString(
1054 | {
1055 | projectId,
1056 | branchId,
1057 | computeId,
1058 | databaseName,
1059 | },
1060 | neonClient,
1061 | extra,
1062 | );
1063 |
1064 | // Connect to the database
1065 | const sql = neon(connectionString.uri);
1066 |
1067 | // First, check if pg_stat_statements extension is installed
1068 | const checkExtensionQuery = `
1069 | SELECT EXISTS (
1070 | SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements'
1071 | ) as extension_exists;
1072 | `;
1073 |
1074 | const extensionCheck = await sql.query(checkExtensionQuery);
1075 | const extensionExists = extensionCheck[0]?.extension_exists;
1076 |
1077 | if (!extensionExists) {
1078 | throw new NotFoundError(
1079 | `pg_stat_statements extension is not installed on the database. Please install it using the following command: CREATE EXTENSION pg_stat_statements;`,
1080 | );
1081 | }
1082 |
1083 | // Query to get slow queries
1084 | const slowQueriesQuery = `
1085 | SELECT
1086 | query,
1087 | calls,
1088 | total_exec_time,
1089 | mean_exec_time,
1090 | rows,
1091 | shared_blks_hit,
1092 | shared_blks_read,
1093 | shared_blks_written,
1094 | shared_blks_dirtied,
1095 | temp_blks_read,
1096 | temp_blks_written,
1097 | wal_records,
1098 | wal_fpi,
1099 | wal_bytes
1100 | FROM pg_stat_statements
1101 | WHERE query NOT LIKE '%pg_stat_statements%'
1102 | AND query NOT LIKE '%EXPLAIN%'
1103 | ORDER BY mean_exec_time DESC
1104 | LIMIT $1;
1105 | `;
1106 |
1107 | const slowQueries = await sql.query(slowQueriesQuery, [limit]);
1108 |
1109 | // Format the results
1110 | const formattedQueries = slowQueries.map((query: any) => {
1111 | return {
1112 | query: query.query,
1113 | calls: query.calls,
1114 | total_exec_time_ms: query.total_exec_time,
1115 | mean_exec_time_ms: query.mean_exec_time,
1116 | rows: query.rows,
1117 | shared_blocks: {
1118 | hit: query.shared_blks_hit,
1119 | read: query.shared_blks_read,
1120 | written: query.shared_blks_written,
1121 | dirtied: query.shared_blks_dirtied,
1122 | },
1123 | temp_blocks: {
1124 | read: query.temp_blks_read,
1125 | written: query.temp_blks_written,
1126 | },
1127 | io_time: {
1128 | read_ms: query.blk_read_time,
1129 | write_ms: query.blk_write_time,
1130 | },
1131 | wal: {
1132 | records: query.wal_records,
1133 | full_page_images: query.wal_fpi,
1134 | bytes: query.wal_bytes,
1135 | },
1136 | };
1137 | });
1138 |
1139 | return {
1140 | slow_queries: formattedQueries,
1141 | total_queries_found: formattedQueries.length,
1142 | };
1143 | }
1144 |
1145 | async function handleListBranchComputes(
1146 | {
1147 | projectId,
1148 | branchId,
1149 | }: {
1150 | projectId?: string;
1151 | branchId?: string;
1152 | },
1153 | neonClient: Api<unknown>,
1154 | extra: ToolHandlerExtraParams,
1155 | ) {
1156 | // If projectId is not provided, get the first project but only if there is only one project
1157 | if (!projectId) {
1158 | const projects = await handleListProjects({}, neonClient, extra);
1159 | if (projects.length === 1) {
1160 | projectId = projects[0].id;
1161 | } else {
1162 | throw new InvalidArgumentError(
1163 | 'Please provide a project ID or ensure you have only one project in your account.',
1164 | );
1165 | }
1166 | }
1167 |
1168 | let endpoints;
1169 | if (branchId) {
1170 | const response = await neonClient.listProjectBranchEndpoints(
1171 | projectId,
1172 | branchId,
1173 | );
1174 | endpoints = response.data.endpoints;
1175 | } else {
1176 | const response = await neonClient.listProjectEndpoints(projectId);
1177 | endpoints = response.data.endpoints;
1178 | }
1179 |
1180 | return endpoints.map((endpoint) => ({
1181 | compute_id: endpoint.id,
1182 | compute_type: endpoint.type,
1183 | compute_size:
1184 | endpoint.autoscaling_limit_min_cu !== endpoint.autoscaling_limit_max_cu
1185 | ? `${endpoint.autoscaling_limit_min_cu}-${endpoint.autoscaling_limit_max_cu}`
1186 | : endpoint.autoscaling_limit_min_cu,
1187 | last_active: endpoint.last_active,
1188 | ...endpoint,
1189 | }));
1190 | }
1191 |
1192 | async function handleListOrganizations(
1193 | neonClient: Api<unknown>,
1194 | account: ToolHandlerExtraParams['account'],
1195 | search?: string,
1196 | ): Promise<Organization[]> {
1197 | if (account.isOrg) {
1198 | const orgId = account.id;
1199 | const { data } = await neonClient.getOrganization(orgId);
1200 | return filterOrganizations([data], search);
1201 | }
1202 |
1203 | const { data: response } = await neonClient.getCurrentUserOrganizations();
1204 | const organizations = response.organizations || [];
1205 | return filterOrganizations(organizations, search);
1206 | }
1207 |
1208 | async function handleListSharedProjects(
1209 | params: ListSharedProjectsParams,
1210 | neonClient: Api<unknown>,
1211 | ) {
1212 | const response = await neonClient.listSharedProjects(params);
1213 | return response.data.projects;
1214 | }
1215 |
1216 | async function handleCompareDatabaseSchema(
1217 | params: GetProjectBranchSchemaComparisonParams,
1218 | neonClient: Api<unknown>,
1219 | ) {
1220 | const response = await neonClient.getProjectBranchSchemaComparison(params);
1221 | return response.data;
1222 | }
1223 |
1224 | export const NEON_HANDLERS = {
1225 | list_projects: async ({ params }, neonClient, extra) => {
1226 | const organization = await getOrgByOrgIdOrDefault(
1227 | params,
1228 | neonClient,
1229 | extra,
1230 | );
1231 | const projects = await handleListProjects(
1232 | { ...params, org_id: organization?.id },
1233 | neonClient,
1234 | extra,
1235 | );
1236 | return {
1237 | content: [
1238 | {
1239 | type: 'text',
1240 | text: JSON.stringify(
1241 | {
1242 | organization: organization
1243 | ? {
1244 | name: organization.name,
1245 | id: organization.id,
1246 | }
1247 | : undefined,
1248 | projects,
1249 | },
1250 | null,
1251 | 2,
1252 | ),
1253 | },
1254 | ],
1255 | };
1256 | },
1257 |
1258 | create_project: async ({ params }, neonClient, extra) => {
1259 | try {
1260 | const organization = await getOrgByOrgIdOrDefault(
1261 | params,
1262 | neonClient,
1263 | extra,
1264 | );
1265 | const result = await handleCreateProject(
1266 | { project: { name: params.name, org_id: organization?.id } },
1267 | neonClient,
1268 | );
1269 |
1270 | // Get the connection string for the newly created project
1271 | const connectionString = await handleGetConnectionString(
1272 | {
1273 | projectId: result.project.id,
1274 | branchId: result.branch.id,
1275 | databaseName: result.databases[0].name,
1276 | },
1277 | neonClient,
1278 | extra,
1279 | );
1280 |
1281 | return {
1282 | content: [
1283 | {
1284 | type: 'text',
1285 | text: [
1286 | `Your Neon project is created ${organization ? `in organization "${organization.name}"` : ''} and is ready.`,
1287 | `The project_id is "${result.project.id}"`,
1288 | `The branch name is "${result.branch.name}" (ID: ${result.branch.id})`,
1289 | `There is one database available on this branch, called "${result.databases[0].name}",`,
1290 | 'but you can create more databases using SQL commands.',
1291 | '',
1292 | 'Connection string details:',
1293 | `URI: ${connectionString.uri}`,
1294 | `Project ID: ${connectionString.projectId}`,
1295 | `Branch ID: ${connectionString.branchId}`,
1296 | `Database: ${connectionString.databaseName}`,
1297 | `Role: ${connectionString.roleName}`,
1298 | '',
1299 | 'You can use this connection string with any PostgreSQL client to connect to your Neon database.',
1300 | 'For example, with psql:',
1301 | `psql "${connectionString.uri}"`,
1302 | ].join('\n'),
1303 | },
1304 | ],
1305 | };
1306 | } catch (error) {
1307 | const message = error instanceof Error ? error.message : 'Unknown error';
1308 | return {
1309 | isError: true,
1310 | content: [
1311 | {
1312 | type: 'text',
1313 | text: [
1314 | 'An error occurred while creating the project.',
1315 | 'Error details:',
1316 | message,
1317 | 'If you have reached the Neon project limit, please upgrade your account in this link: https://console.neon.tech/app/billing',
1318 | ].join('\n'),
1319 | },
1320 | ],
1321 | };
1322 | }
1323 | },
1324 |
1325 | delete_project: async ({ params }, neonClient) => {
1326 | await handleDeleteProject(params.projectId, neonClient);
1327 | return {
1328 | content: [
1329 | {
1330 | type: 'text',
1331 | text: [
1332 | 'Project deleted successfully.',
1333 | `Project ID: ${params.projectId}`,
1334 | ].join('\n'),
1335 | },
1336 | ],
1337 | };
1338 | },
1339 |
1340 | describe_project: async ({ params }, neonClient) => {
1341 | const result = await handleDescribeProject(params.projectId, neonClient);
1342 | return {
1343 | content: [
1344 | {
1345 | type: 'text',
1346 | text: `This project is called ${result.project.project.name}.`,
1347 | },
1348 | {
1349 | type: 'text',
1350 | text: `It contains the following branches (use the describe branch tool to learn more about each branch): ${JSON.stringify(
1351 | result.branches,
1352 | null,
1353 | 2,
1354 | )}`,
1355 | },
1356 | ],
1357 | };
1358 | },
1359 |
1360 | run_sql: async ({ params }, neonClient, extra) => {
1361 | const result = await handleRunSql(
1362 | {
1363 | sql: params.sql,
1364 | databaseName: params.databaseName,
1365 | projectId: params.projectId,
1366 | branchId: params.branchId,
1367 | },
1368 | neonClient,
1369 | extra,
1370 | );
1371 | return {
1372 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1373 | };
1374 | },
1375 |
1376 | run_sql_transaction: async ({ params }, neonClient, extra) => {
1377 | const result = await handleRunSqlTransaction(
1378 | {
1379 | sqlStatements: params.sqlStatements,
1380 | databaseName: params.databaseName,
1381 | projectId: params.projectId,
1382 | branchId: params.branchId,
1383 | },
1384 | neonClient,
1385 | extra,
1386 | );
1387 | return {
1388 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1389 | };
1390 | },
1391 |
1392 | describe_table_schema: async ({ params }, neonClient, extra) => {
1393 | const result = await handleDescribeTableSchema(
1394 | {
1395 | tableName: params.tableName,
1396 | databaseName: params.databaseName,
1397 | projectId: params.projectId,
1398 | branchId: params.branchId,
1399 | },
1400 | neonClient,
1401 | extra,
1402 | );
1403 | return {
1404 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1405 | };
1406 | },
1407 |
1408 | get_database_tables: async ({ params }, neonClient, extra) => {
1409 | const result = await handleGetDatabaseTables(
1410 | {
1411 | projectId: params.projectId,
1412 | branchId: params.branchId,
1413 | databaseName: params.databaseName,
1414 | },
1415 | neonClient,
1416 | extra,
1417 | );
1418 | return {
1419 | content: [
1420 | {
1421 | type: 'text',
1422 | text: JSON.stringify(result, null, 2),
1423 | },
1424 | ],
1425 | };
1426 | },
1427 |
1428 | create_branch: async ({ params }, neonClient) => {
1429 | const result = await handleCreateBranch(
1430 | {
1431 | projectId: params.projectId,
1432 | branchName: params.branchName,
1433 | },
1434 | neonClient,
1435 | );
1436 | return {
1437 | content: [
1438 | {
1439 | type: 'text',
1440 | text: [
1441 | 'Branch created successfully.',
1442 | `Project ID: ${result.branch.project_id}`,
1443 | `Branch ID: ${result.branch.id}`,
1444 | `Branch name: ${result.branch.name}`,
1445 | `Parent branch: ${result.branch.parent_id}`,
1446 | ].join('\n'),
1447 | },
1448 | ],
1449 | };
1450 | },
1451 |
1452 | prepare_database_migration: async ({ params }, neonClient, extra) => {
1453 | const result = await handleSchemaMigration(
1454 | {
1455 | migrationSql: params.migrationSql,
1456 | databaseName: params.databaseName,
1457 | projectId: params.projectId,
1458 | },
1459 | neonClient,
1460 | extra,
1461 | );
1462 | return {
1463 | content: [
1464 | {
1465 | type: 'text',
1466 | text: `
1467 | <status>Migration created successfully in temporary branch</status>
1468 | <details>
1469 | <migration_id>${result.migrationId}</migration_id>
1470 | <temporary_branch>
1471 | <name>${result.branch.name}</name>
1472 | <id>${result.branch.id}</id>
1473 | </temporary_branch>
1474 | </details>
1475 | <execution_result>${JSON.stringify(result.migrationResult, null, 2)}</execution_result>
1476 |
1477 | <next_actions>
1478 | You MUST follow these steps:
1479 | 1. Test this migration using \`run_sql\` tool on branch \`${result.branch.name}\`
1480 | 2. Verify the changes meet your requirements
1481 | 3. If satisfied, use \`complete_database_migration\` with migration_id: ${result.migrationId}
1482 | </next_actions>
1483 | `,
1484 | },
1485 | ],
1486 | };
1487 | },
1488 |
1489 | complete_database_migration: async ({ params }, neonClient, extra) => {
1490 | const result = await handleCommitMigration(
1491 | {
1492 | migrationId: params.migrationId,
1493 | },
1494 | neonClient,
1495 | extra,
1496 | );
1497 | return {
1498 | content: [
1499 | {
1500 | type: 'text',
1501 | text: `Result: ${JSON.stringify(
1502 | {
1503 | deletedBranch: result.deletedBranch,
1504 | migrationResult: result.migrationResult,
1505 | },
1506 | null,
1507 | 2,
1508 | )}`,
1509 | },
1510 | ],
1511 | };
1512 | },
1513 |
1514 | describe_branch: async ({ params }, neonClient, extra) => {
1515 | const result = await handleDescribeBranch(
1516 | {
1517 | projectId: params.projectId,
1518 | branchId: params.branchId,
1519 | databaseName: params.databaseName,
1520 | },
1521 | neonClient,
1522 | extra,
1523 | );
1524 | return {
1525 | content: [
1526 | {
1527 | type: 'text',
1528 | text: ['Database Structure:', JSON.stringify(result, null, 2)].join(
1529 | '\n',
1530 | ),
1531 | },
1532 | ],
1533 | };
1534 | },
1535 |
1536 | delete_branch: async ({ params }, neonClient) => {
1537 | await handleDeleteBranch(
1538 | {
1539 | projectId: params.projectId,
1540 | branchId: params.branchId,
1541 | },
1542 | neonClient,
1543 | );
1544 | return {
1545 | content: [
1546 | {
1547 | type: 'text',
1548 | text: [
1549 | 'Branch deleted successfully.',
1550 | `Project ID: ${params.projectId}`,
1551 | `Branch ID: ${params.branchId}`,
1552 | ].join('\n'),
1553 | },
1554 | ],
1555 | };
1556 | },
1557 |
1558 | reset_from_parent: async ({ params }, neonClient) => {
1559 | const result = await handleResetFromParent(
1560 | {
1561 | projectId: params.projectId,
1562 | branchIdOrName: params.branchIdOrName,
1563 | preserveUnderName: params.preserveUnderName,
1564 | },
1565 | neonClient,
1566 | );
1567 |
1568 | const parentInfo = `${result.parentBranch.name} (${result.parentBranch.id})`;
1569 |
1570 | const messages = [
1571 | 'Branch reset from parent successfully.',
1572 | `Project: ${params.projectId}`,
1573 | `Branch: ${params.branchIdOrName}`,
1574 | `Reset to parent branch: ${parentInfo}`,
1575 | ];
1576 |
1577 | if (result.preservedBranchName) {
1578 | messages.push(
1579 | params.preserveUnderName
1580 | ? `Previous state preserved as: ${params.preserveUnderName}`
1581 | : `Previous state auto-preserved as: ${result.preservedBranchName} (branch had children)`,
1582 | );
1583 | } else {
1584 | messages.push('Previous state was not preserved');
1585 | }
1586 |
1587 | return {
1588 | content: [
1589 | {
1590 | type: 'text',
1591 | text: messages.join('\n'),
1592 | },
1593 | ],
1594 | };
1595 | },
1596 |
1597 | get_connection_string: async ({ params }, neonClient, extra) => {
1598 | const result = await handleGetConnectionString(
1599 | {
1600 | projectId: params.projectId,
1601 | branchId: params.branchId,
1602 | computeId: params.computeId,
1603 | databaseName: params.databaseName,
1604 | roleName: params.roleName,
1605 | },
1606 | neonClient,
1607 | extra,
1608 | );
1609 | return {
1610 | content: [
1611 | {
1612 | type: 'text',
1613 | text: [
1614 | 'Connection string details:',
1615 | `URI: ${result.uri}`,
1616 | `Project ID: ${result.projectId}`,
1617 | `Database: ${result.databaseName}`,
1618 | `Role: ${result.roleName}`,
1619 | result.branchId
1620 | ? `Branch ID: ${result.branchId}`
1621 | : 'Using default branch',
1622 | result.computeId
1623 | ? `Compute ID: ${result.computeId}`
1624 | : 'Using default compute',
1625 | '',
1626 | 'You can use this connection string with any PostgreSQL client to connect to your Neon database.',
1627 | ].join('\n'),
1628 | },
1629 | ],
1630 | };
1631 | },
1632 |
1633 | provision_neon_auth: async ({ params }, neonClient) => {
1634 | const result = await handleProvisionNeonAuth(
1635 | {
1636 | projectId: params.projectId,
1637 | database: params.database,
1638 | },
1639 | neonClient,
1640 | );
1641 | return result;
1642 | },
1643 |
1644 | explain_sql_statement: async ({ params }, neonClient, extra) => {
1645 | const result = await handleExplainSqlStatement(
1646 | { params },
1647 | neonClient,
1648 | extra,
1649 | );
1650 | return result;
1651 | },
1652 |
1653 | prepare_query_tuning: async ({ params }, neonClient, extra) => {
1654 | const result = await handleQueryTuning(
1655 | {
1656 | sql: params.sql,
1657 | databaseName: params.databaseName,
1658 | projectId: params.projectId,
1659 | },
1660 | neonClient,
1661 | extra,
1662 | );
1663 | return {
1664 | content: [
1665 | {
1666 | type: 'text',
1667 | text: JSON.stringify(
1668 | {
1669 | tuningId: result.tuningId,
1670 | databaseName: result.databaseName,
1671 | projectId: result.projectId,
1672 | temporaryBranch: result.temporaryBranch,
1673 | executionPlan: result.originalPlan,
1674 | tableSchemas: result.tableSchemas,
1675 | sql: result.sql,
1676 | },
1677 | null,
1678 | 2,
1679 | ),
1680 | },
1681 | ],
1682 | };
1683 | },
1684 |
1685 | complete_query_tuning: async ({ params }, neonClient, extra) => {
1686 | const result = await handleCompleteTuning(
1687 | {
1688 | suggestedSqlStatements: params.suggestedSqlStatements,
1689 | applyChanges: params.applyChanges,
1690 | tuningId: params.tuningId,
1691 | databaseName: params.databaseName,
1692 | projectId: params.projectId,
1693 | temporaryBranch: {
1694 | id: params.temporaryBranchId,
1695 | project_id: params.projectId,
1696 | } as Branch,
1697 | shouldDeleteTemporaryBranch: params.shouldDeleteTemporaryBranch,
1698 | branch: params.branchId
1699 | ? ({ id: params.branchId, project_id: params.projectId } as Branch)
1700 | : undefined,
1701 | },
1702 | neonClient,
1703 | extra,
1704 | );
1705 |
1706 | return {
1707 | content: [
1708 | {
1709 | type: 'text',
1710 | text: JSON.stringify(result, null, 2),
1711 | },
1712 | ],
1713 | };
1714 | },
1715 |
1716 | list_slow_queries: async ({ params }, neonClient, extra) => {
1717 | const result = await handleListSlowQueries(
1718 | {
1719 | projectId: params.projectId,
1720 | branchId: params.branchId,
1721 | databaseName: params.databaseName,
1722 | computeId: params.computeId,
1723 | limit: params.limit,
1724 | },
1725 | neonClient,
1726 | extra,
1727 | );
1728 | return {
1729 | content: [
1730 | {
1731 | type: 'text',
1732 | text: JSON.stringify(result, null, 2),
1733 | },
1734 | ],
1735 | };
1736 | },
1737 |
1738 | list_branch_computes: async ({ params }, neonClient, extra) => {
1739 | const result = await handleListBranchComputes(
1740 | {
1741 | projectId: params.projectId,
1742 | branchId: params.branchId,
1743 | },
1744 | neonClient,
1745 | extra,
1746 | );
1747 | return {
1748 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1749 | };
1750 | },
1751 |
1752 | list_organizations: async ({ params }, neonClient, extra) => {
1753 | const organizations = await handleListOrganizations(
1754 | neonClient,
1755 | extra.account,
1756 | params.search,
1757 | );
1758 | return {
1759 | content: [
1760 | {
1761 | type: 'text',
1762 | text: JSON.stringify(organizations, null, 2),
1763 | },
1764 | ],
1765 | };
1766 | },
1767 |
1768 | list_shared_projects: async ({ params }, neonClient) => {
1769 | const sharedProjects = await handleListSharedProjects(params, neonClient);
1770 | return {
1771 | content: [
1772 | {
1773 | type: 'text',
1774 | text: JSON.stringify(
1775 | {
1776 | shared_projects: sharedProjects,
1777 | count: sharedProjects.length,
1778 | },
1779 | null,
1780 | 2,
1781 | ),
1782 | },
1783 | ],
1784 | };
1785 | },
1786 |
1787 | compare_database_schema: async ({ params }, neonClient) => {
1788 | const result = await handleCompareDatabaseSchema(
1789 | {
1790 | projectId: params.projectId,
1791 | branchId: params.branchId,
1792 | db_name: params.databaseName,
1793 | },
1794 | neonClient,
1795 | );
1796 | return {
1797 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1798 | };
1799 | },
1800 | } satisfies ToolHandlers;
1801 |
```