This is page 3 of 3. 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
│ ├── claude-code-review.yml
│ ├── claude.yml
│ ├── koyeb-preview.yml
│ ├── koyeb-prod.yml
│ └── pr.yml
├── .gitignore
├── .npmrc
├── .nvmrc
├── .prettierignore
├── .prettierrc
├── AGENTS.md
├── bun.lock
├── CHANGELOG.md
├── CLAUDE.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
│ ├── prompts.ts
│ ├── resources.ts
│ ├── sentry
│ │ ├── instrument.ts
│ │ └── utils.ts
│ ├── server
│ │ ├── api.ts
│ │ ├── errors.ts
│ │ └── index.ts
│ ├── tools
│ │ ├── definitions.ts
│ │ ├── handlers
│ │ │ ├── connection-string.ts
│ │ │ ├── decribe-project.ts
│ │ │ ├── describe-branch.ts
│ │ │ ├── fetch.ts
│ │ │ ├── list-orgs.ts
│ │ │ ├── list-projects.ts
│ │ │ ├── neon-auth.ts
│ │ │ ├── search.ts
│ │ │ ├── urls.ts
│ │ │ └── utils.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
│ │ └── helpers.ts
│ ├── utils
│ │ ├── client-application.ts
│ │ ├── logger.ts
│ │ └── polyfills.ts
│ └── views
│ ├── approval-dialog.pug
│ └── styles.css
├── tsconfig.json
└── tsconfig.test.json
```
# Files
--------------------------------------------------------------------------------
/src/tools/tools.ts:
--------------------------------------------------------------------------------
```typescript
1 | import {
2 | Api,
3 | Branch,
4 | EndpointType,
5 | ListSharedProjectsParams,
6 | GetProjectBranchSchemaComparisonParams,
7 | ProjectCreateRequest,
8 | } from '@neondatabase/api-client';
9 | import { neon } from '@neondatabase/serverless';
10 | import crypto from 'crypto';
11 | import { InvalidArgumentError, NotFoundError } from '../server/errors.js';
12 |
13 | import { describeTable, formatTableDescription } from '../describeUtils.js';
14 | import { handleProvisionNeonAuth } from './handlers/neon-auth.js';
15 | import { handleSearch } from './handlers/search.js';
16 | import { handleFetch } from './handlers/fetch.js';
17 | import { getMigrationFromMemory, persistMigrationToMemory } from './state.js';
18 | import { fetchRawGithubContent, NEON_RESOURCES } from '../resources.js';
19 |
20 | import {
21 | getDefaultDatabase,
22 | splitSqlStatements,
23 | getOrgByOrgIdOrDefault,
24 | resolveBranchId,
25 | } from './utils.js';
26 | import { startSpan } from '@sentry/node';
27 | import { ToolHandlerExtraParams, ToolHandlers } from './types.js';
28 | import { handleListOrganizations } from './handlers/list-orgs.js';
29 | import { handleListProjects } from './handlers/list-projects.js';
30 | import { handleDescribeProject } from './handlers/decribe-project.js';
31 | import { handleGetConnectionString } from './handlers/connection-string.js';
32 | import { handleDescribeBranch } from './handlers/describe-branch.js';
33 |
34 | async function handleCreateProject(
35 | params: ProjectCreateRequest,
36 | neonClient: Api<unknown>,
37 | ) {
38 | const response = await neonClient.createProject(params);
39 | if (response.status !== 201) {
40 | throw new Error(`Failed to create project: ${JSON.stringify(response)}`);
41 | }
42 | return response.data;
43 | }
44 |
45 | async function handleDeleteProject(
46 | projectId: string,
47 | neonClient: Api<unknown>,
48 | ) {
49 | const response = await neonClient.deleteProject(projectId);
50 | if (response.status !== 200) {
51 | throw new Error(`Failed to delete project: ${response.statusText}`);
52 | }
53 | return response.data;
54 | }
55 |
56 | async function handleRunSql(
57 | {
58 | sql,
59 | databaseName,
60 | projectId,
61 | branchId,
62 | }: {
63 | sql: string;
64 | databaseName?: string;
65 | projectId: string;
66 | branchId?: string;
67 | },
68 | neonClient: Api<unknown>,
69 | extra: ToolHandlerExtraParams,
70 | ) {
71 | return await startSpan({ name: 'run_sql' }, async () => {
72 | const connectionString = await handleGetConnectionString(
73 | {
74 | projectId,
75 | branchId,
76 | databaseName,
77 | },
78 | neonClient,
79 | extra,
80 | );
81 | const runQuery = neon(connectionString.uri);
82 |
83 | // If in read-only mode, use transaction with readOnly option
84 | if (extra.readOnly) {
85 | const response = await runQuery.transaction([runQuery.query(sql)], {
86 | readOnly: true,
87 | });
88 | // Return the first result (the actual query result)
89 | return response[0];
90 | }
91 |
92 | const response = await runQuery.query(sql);
93 |
94 | return response;
95 | });
96 | }
97 |
98 | async function handleRunSqlTransaction(
99 | {
100 | sqlStatements,
101 | databaseName,
102 | projectId,
103 | branchId,
104 | }: {
105 | sqlStatements: string[];
106 | databaseName?: string;
107 | projectId: string;
108 | branchId?: string;
109 | },
110 | neonClient: Api<unknown>,
111 | extra: ToolHandlerExtraParams,
112 | ) {
113 | const connectionString = await handleGetConnectionString(
114 | {
115 | projectId,
116 | branchId,
117 | databaseName,
118 | },
119 | neonClient,
120 | extra,
121 | );
122 | const runQuery = neon(connectionString.uri);
123 |
124 | // Use transaction with readOnly option when in read-only mode
125 | const response = await runQuery.transaction(
126 | sqlStatements.map((sql) => runQuery.query(sql)),
127 | extra.readOnly ? { readOnly: true } : undefined,
128 | );
129 |
130 | return response;
131 | }
132 |
133 | async function handleGetDatabaseTables(
134 | {
135 | projectId,
136 | databaseName,
137 | branchId,
138 | }: {
139 | projectId: string;
140 | databaseName?: string;
141 | branchId?: string;
142 | },
143 | neonClient: Api<unknown>,
144 | extra: ToolHandlerExtraParams,
145 | ) {
146 | const connectionString = await handleGetConnectionString(
147 | {
148 | projectId,
149 | branchId,
150 | databaseName,
151 | },
152 | neonClient,
153 | extra,
154 | );
155 | const runQuery = neon(connectionString.uri);
156 | const query = `
157 | SELECT
158 | table_schema,
159 | table_name,
160 | table_type
161 | FROM information_schema.tables
162 | WHERE table_schema NOT IN ('pg_catalog', 'information_schema')
163 | ORDER BY table_schema, table_name;
164 | `;
165 |
166 | const tables = await runQuery.query(query);
167 | return tables;
168 | }
169 |
170 | async function handleDescribeTableSchema(
171 | {
172 | projectId,
173 | databaseName,
174 | branchId,
175 | tableName,
176 | }: {
177 | projectId: string;
178 | databaseName?: string;
179 | branchId?: string;
180 | tableName: string;
181 | },
182 | neonClient: Api<unknown>,
183 | extra: ToolHandlerExtraParams,
184 | ) {
185 | const connectionString = await handleGetConnectionString(
186 | {
187 | projectId,
188 | branchId,
189 | databaseName,
190 | },
191 | neonClient,
192 | extra,
193 | );
194 |
195 | // Extract table name without schema if schema-qualified
196 | const tableNameParts = tableName.split('.');
197 | const simpleTableName = tableNameParts[tableNameParts.length - 1];
198 |
199 | const description = await describeTable(
200 | connectionString.uri,
201 | simpleTableName,
202 | );
203 | return {
204 | raw: description,
205 | formatted: formatTableDescription(description),
206 | };
207 | }
208 |
209 | async function handleCreateBranch(
210 | {
211 | projectId,
212 | branchName,
213 | }: {
214 | projectId: string;
215 | branchName?: string;
216 | },
217 | neonClient: Api<unknown>,
218 | ) {
219 | const response = await neonClient.createProjectBranch(projectId, {
220 | branch: {
221 | name: branchName,
222 | },
223 | endpoints: [
224 | {
225 | type: EndpointType.ReadWrite,
226 | autoscaling_limit_min_cu: 0.25,
227 | autoscaling_limit_max_cu: 0.25,
228 | provisioner: 'k8s-neonvm',
229 | },
230 | ],
231 | });
232 |
233 | if (response.status !== 201) {
234 | throw new Error(`Failed to create branch: ${response.statusText}`);
235 | }
236 |
237 | return response.data;
238 | }
239 |
240 | async function handleDeleteBranch(
241 | {
242 | projectId,
243 | branchId,
244 | }: {
245 | projectId: string;
246 | branchId: string;
247 | },
248 | neonClient: Api<unknown>,
249 | ) {
250 | const response = await neonClient.deleteProjectBranch(projectId, branchId);
251 | return response.data;
252 | }
253 |
254 | async function handleResetFromParent(
255 | {
256 | projectId,
257 | branchIdOrName,
258 | preserveUnderName,
259 | }: {
260 | projectId: string;
261 | branchIdOrName: string;
262 | preserveUnderName?: string;
263 | },
264 | neonClient: Api<unknown>,
265 | ) {
266 | // Resolve branch name or ID to actual branch ID and get all branches in one call
267 | const { branchId: resolvedBranchId, branches } = await resolveBranchId(
268 | branchIdOrName,
269 | projectId,
270 | neonClient,
271 | );
272 |
273 | const branch = branches.find((b) => b.id === resolvedBranchId);
274 | if (!branch) {
275 | throw new NotFoundError(
276 | `Branch "${branchIdOrName}" not found in project ${projectId}`,
277 | );
278 | }
279 |
280 | // Find the parent branch and validate it exists
281 | const parentBranch = branch.parent_id
282 | ? branches.find((b) => b.id === branch.parent_id)
283 | : undefined;
284 |
285 | if (!parentBranch) {
286 | throw new InvalidArgumentError(
287 | `Branch "${branchIdOrName}" does not have a parent branch and cannot be reset`,
288 | );
289 | }
290 |
291 | // Check if the branch has children
292 | const hasChildren = branches.some((b) => b.parent_id === resolvedBranchId);
293 |
294 | // Auto-generate preserve name if branch has children and none was provided
295 | let finalPreserveName = preserveUnderName;
296 | if (hasChildren && !preserveUnderName) {
297 | const timestamp = new Date()
298 | .toISOString()
299 | .replace(/[:.]/g, '-')
300 | .slice(0, -5);
301 | finalPreserveName = `${branch.name}_old_${timestamp}`;
302 | }
303 |
304 | // Call the restoreProjectBranch API
305 | const response = await neonClient.restoreProjectBranch(
306 | projectId,
307 | resolvedBranchId,
308 | {
309 | source_branch_id: parentBranch.id,
310 | preserve_under_name: finalPreserveName,
311 | },
312 | );
313 |
314 | return {
315 | ...response.data,
316 | preservedBranchName: finalPreserveName,
317 | parentBranch,
318 | };
319 | }
320 |
321 | async function handleSchemaMigration(
322 | {
323 | migrationSql,
324 | databaseName,
325 | projectId,
326 | }: {
327 | databaseName?: string;
328 | projectId: string;
329 | migrationSql: string;
330 | },
331 | neonClient: Api<unknown>,
332 | extra: ToolHandlerExtraParams,
333 | ) {
334 | return await startSpan({ name: 'prepare_schema_migration' }, async (span) => {
335 | const newBranch = await handleCreateBranch({ projectId }, neonClient);
336 |
337 | if (!databaseName) {
338 | const dbObject = await getDefaultDatabase(
339 | {
340 | projectId,
341 | branchId: newBranch.branch.id,
342 | databaseName,
343 | },
344 | neonClient,
345 | );
346 | databaseName = dbObject.name;
347 | }
348 |
349 | const result = await handleRunSqlTransaction(
350 | {
351 | sqlStatements: splitSqlStatements(migrationSql),
352 | databaseName,
353 | projectId,
354 | branchId: newBranch.branch.id,
355 | },
356 | neonClient,
357 | extra,
358 | );
359 |
360 | const migrationId = crypto.randomUUID();
361 | span.setAttributes({
362 | projectId,
363 | migrationId,
364 | });
365 | persistMigrationToMemory(migrationId, {
366 | migrationSql,
367 | databaseName,
368 | appliedBranch: newBranch.branch,
369 | });
370 |
371 | return {
372 | branch: newBranch.branch,
373 | migrationId,
374 | migrationResult: result,
375 | };
376 | });
377 | }
378 |
379 | async function handleCommitMigration(
380 | { migrationId }: { migrationId: string },
381 | neonClient: Api<unknown>,
382 | extra: ToolHandlerExtraParams,
383 | ) {
384 | return await startSpan({ name: 'commit_schema_migration' }, async (span) => {
385 | span.setAttributes({
386 | migrationId,
387 | });
388 | const migration = getMigrationFromMemory(migrationId);
389 | if (!migration) {
390 | throw new Error(`Migration not found: ${migrationId}`);
391 | }
392 |
393 | span.setAttributes({
394 | projectId: migration.appliedBranch.project_id,
395 | });
396 | const result = await handleRunSqlTransaction(
397 | {
398 | sqlStatements: splitSqlStatements(migration.migrationSql),
399 | databaseName: migration.databaseName,
400 | projectId: migration.appliedBranch.project_id,
401 | branchId: migration.appliedBranch.parent_id,
402 | },
403 | neonClient,
404 | extra,
405 | );
406 |
407 | await handleDeleteBranch(
408 | {
409 | projectId: migration.appliedBranch.project_id,
410 | branchId: migration.appliedBranch.id,
411 | },
412 | neonClient,
413 | );
414 |
415 | return {
416 | deletedBranch: migration.appliedBranch,
417 | migrationResult: result,
418 | };
419 | });
420 | }
421 |
422 | async function handleExplainSqlStatement(
423 | {
424 | params,
425 | }: {
426 | params: {
427 | sql: string;
428 | databaseName?: string;
429 | projectId: string;
430 | branchId?: string;
431 | analyze: boolean;
432 | };
433 | },
434 | neonClient: Api<unknown>,
435 | extra: ToolHandlerExtraParams,
436 | ) {
437 | const explainPrefix = params.analyze
438 | ? 'EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FILECACHE, FORMAT JSON)'
439 | : 'EXPLAIN (VERBOSE, FORMAT JSON)';
440 |
441 | const explainSql = `${explainPrefix} ${params.sql}`;
442 |
443 | const result = await handleRunSql(
444 | {
445 | sql: explainSql,
446 | databaseName: params.databaseName,
447 | projectId: params.projectId,
448 | branchId: params.branchId,
449 | },
450 | neonClient,
451 | extra,
452 | );
453 |
454 | return {
455 | content: [
456 | {
457 | type: 'text' as const,
458 | text: JSON.stringify(result, null, 2),
459 | },
460 | ],
461 | };
462 | }
463 |
464 | async function createTemporaryBranch(
465 | projectId: string,
466 | neonClient: Api<unknown>,
467 | ): Promise<{ branch: Branch }> {
468 | const result = await handleCreateBranch({ projectId }, neonClient);
469 | if (!result?.branch) {
470 | throw new Error('Failed to create temporary branch');
471 | }
472 | return result;
473 | }
474 |
475 | type QueryTuningParams = {
476 | sql: string;
477 | databaseName: string;
478 | projectId: string;
479 | };
480 |
481 | type CompleteTuningParams = {
482 | suggestedSqlStatements?: string[];
483 | applyChanges?: boolean;
484 | tuningId: string;
485 | databaseName: string;
486 | projectId: string;
487 | temporaryBranch: Branch;
488 | shouldDeleteTemporaryBranch?: boolean;
489 | branch?: Branch;
490 | };
491 |
492 | type QueryTuningResult = {
493 | tuningId: string;
494 | databaseName: string;
495 | projectId: string;
496 | temporaryBranch: Branch;
497 | originalPlan: any;
498 | tableSchemas: any[];
499 | sql: string;
500 | baselineMetrics: QueryMetrics;
501 | };
502 |
503 | type CompleteTuningResult = {
504 | appliedChanges?: string[];
505 | results?: any;
506 | deletedBranches?: string[];
507 | message: string;
508 | };
509 |
510 | async function handleQueryTuning(
511 | params: QueryTuningParams,
512 | neonClient: Api<unknown>,
513 | extra: ToolHandlerExtraParams,
514 | ): Promise<QueryTuningResult> {
515 | let tempBranch: Branch | undefined;
516 | const tuningId = crypto.randomUUID();
517 |
518 | try {
519 | // Create temporary branch
520 | const newBranch = await createTemporaryBranch(params.projectId, neonClient);
521 | if (!newBranch.branch) {
522 | throw new Error('Failed to create temporary branch: branch is undefined');
523 | }
524 | tempBranch = newBranch.branch;
525 |
526 | // Ensure all operations use the temporary branch
527 | const branchParams = {
528 | ...params,
529 | branchId: tempBranch.id,
530 | };
531 |
532 | // First, get the execution plan with table information
533 | const executionPlan = await handleExplainSqlStatement(
534 | {
535 | params: {
536 | sql: branchParams.sql,
537 | databaseName: branchParams.databaseName,
538 | projectId: branchParams.projectId,
539 | branchId: tempBranch.id,
540 | analyze: true,
541 | },
542 | },
543 | neonClient,
544 | extra,
545 | );
546 |
547 | // Extract table names from the plan
548 | const tableNames = extractTableNamesFromPlan(executionPlan);
549 |
550 | if (tableNames.length === 0) {
551 | throw new NotFoundError(
552 | 'No tables found in execution plan. Cannot proceed with optimization.',
553 | );
554 | }
555 |
556 | // Get schema information for all referenced tables in parallel
557 | const tableSchemas = await Promise.all(
558 | tableNames.map(async (tableName) => {
559 | try {
560 | const schema = await handleDescribeTableSchema(
561 | {
562 | tableName,
563 | databaseName: branchParams.databaseName,
564 | projectId: branchParams.projectId,
565 | branchId: newBranch.branch.id,
566 | },
567 | neonClient,
568 | extra,
569 | );
570 | return {
571 | tableName,
572 | schema: schema.raw,
573 | formatted: schema.formatted,
574 | };
575 | } catch (error) {
576 | throw new Error(
577 | `Failed to get schema for table ${tableName}: ${(error as Error).message}`,
578 | );
579 | }
580 | }),
581 | );
582 |
583 | // Get the baseline execution metrics
584 | const baselineMetrics = extractExecutionMetrics(executionPlan);
585 |
586 | // Return the information for analysis
587 | const result: QueryTuningResult = {
588 | tuningId,
589 | databaseName: params.databaseName,
590 | projectId: params.projectId,
591 | temporaryBranch: tempBranch,
592 | originalPlan: executionPlan,
593 | tableSchemas,
594 | sql: params.sql,
595 | baselineMetrics,
596 | };
597 |
598 | return result;
599 | } catch (error) {
600 | // Always attempt to clean up the temporary branch if it was created
601 | if (tempBranch) {
602 | try {
603 | await handleDeleteBranch(
604 | {
605 | projectId: params.projectId,
606 | branchId: tempBranch.id,
607 | },
608 | neonClient,
609 | );
610 | } catch {
611 | // No need to handle cleanup error
612 | }
613 | }
614 |
615 | throw error;
616 | }
617 | }
618 |
619 | // Helper function to extract execution metrics from EXPLAIN output
620 | function extractExecutionMetrics(plan: any): QueryMetrics {
621 | try {
622 | const planJson =
623 | typeof plan.content?.[0]?.text === 'string'
624 | ? JSON.parse(plan.content[0].text)
625 | : plan;
626 |
627 | const metrics: QueryMetrics = {
628 | executionTime: 0,
629 | planningTime: 0,
630 | totalCost: 0,
631 | actualRows: 0,
632 | bufferUsage: {
633 | shared: { hit: 0, read: 0, written: 0, dirtied: 0 },
634 | local: { hit: 0, read: 0, written: 0, dirtied: 0 },
635 | },
636 | };
637 |
638 | // Extract planning and execution time if available
639 | if (planJson?.[0]?.['Planning Time']) {
640 | metrics.planningTime = planJson[0]['Planning Time'];
641 | }
642 | if (planJson?.[0]?.['Execution Time']) {
643 | metrics.executionTime = planJson[0]['Execution Time'];
644 | }
645 |
646 | // Recursively process plan nodes to accumulate costs and buffer usage
647 | function processNode(node: any) {
648 | if (!node || typeof node !== 'object') return;
649 |
650 | // Accumulate costs
651 | if (node['Total Cost']) {
652 | metrics.totalCost = Math.max(metrics.totalCost, node['Total Cost']);
653 | }
654 | if (node['Actual Rows']) {
655 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
656 | metrics.actualRows += node['Actual Rows'];
657 | }
658 |
659 | if (node['Shared Hit Blocks'])
660 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
661 | metrics.bufferUsage.shared.hit += node['Shared Hit Blocks'];
662 | if (node['Shared Read Blocks'])
663 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
664 | metrics.bufferUsage.shared.read += node['Shared Read Blocks'];
665 | if (node['Shared Written Blocks'])
666 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
667 | metrics.bufferUsage.shared.written += node['Shared Written Blocks'];
668 | if (node['Shared Dirtied Blocks'])
669 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
670 | metrics.bufferUsage.shared.dirtied += node['Shared Dirtied Blocks'];
671 |
672 | if (node['Local Hit Blocks'])
673 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
674 | metrics.bufferUsage.local.hit += node['Local Hit Blocks'];
675 | if (node['Local Read Blocks'])
676 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
677 | metrics.bufferUsage.local.read += node['Local Read Blocks'];
678 | if (node['Local Written Blocks'])
679 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
680 | metrics.bufferUsage.local.written += node['Local Written Blocks'];
681 | if (node['Local Dirtied Blocks'])
682 | // eslint-disable-next-line @typescript-eslint/restrict-plus-operands
683 | metrics.bufferUsage.local.dirtied += node['Local Dirtied Blocks'];
684 |
685 | // Process child nodes
686 | if (Array.isArray(node.Plans)) {
687 | node.Plans.forEach(processNode);
688 | }
689 | }
690 |
691 | if (planJson?.[0]?.Plan) {
692 | processNode(planJson[0].Plan);
693 | }
694 |
695 | return metrics;
696 | } catch {
697 | return {
698 | executionTime: 0,
699 | planningTime: 0,
700 | totalCost: 0,
701 | actualRows: 0,
702 | bufferUsage: {
703 | shared: { hit: 0, read: 0, written: 0, dirtied: 0 },
704 | local: { hit: 0, read: 0, written: 0, dirtied: 0 },
705 | },
706 | };
707 | }
708 | }
709 |
710 | // Types for query metrics
711 | type BufferMetrics = {
712 | hit: number;
713 | read: number;
714 | written: number;
715 | dirtied: number;
716 | };
717 |
718 | type QueryMetrics = {
719 | executionTime: number;
720 | planningTime: number;
721 | totalCost: number;
722 | actualRows: number;
723 | bufferUsage: {
724 | shared: BufferMetrics;
725 | local: BufferMetrics;
726 | };
727 | };
728 |
729 | // Function to extract table names from an execution plan
730 | function extractTableNamesFromPlan(planResult: any): string[] {
731 | const tableNames = new Set<string>();
732 |
733 | function recursivelyExtractFromNode(node: any) {
734 | if (!node || typeof node !== 'object') return;
735 |
736 | // Check if current node has relation information
737 | if (node['Relation Name'] && node.Schema) {
738 | const tableName = `${node.Schema}.${node['Relation Name']}`;
739 | tableNames.add(tableName);
740 | }
741 |
742 | // Recursively process all object properties and array elements
743 | if (Array.isArray(node)) {
744 | node.forEach((item) => {
745 | recursivelyExtractFromNode(item);
746 | });
747 | } else {
748 | Object.values(node).forEach((value) => {
749 | recursivelyExtractFromNode(value);
750 | });
751 | }
752 | }
753 |
754 | try {
755 | // Start with the raw plan result
756 | recursivelyExtractFromNode(planResult);
757 |
758 | // If we have content[0].text, also parse and process that
759 | if (planResult?.content?.[0]?.text) {
760 | try {
761 | const parsedContent = JSON.parse(planResult.content[0].text);
762 | recursivelyExtractFromNode(parsedContent);
763 | } catch {
764 | // No need to handle parse error
765 | }
766 | }
767 | } catch {
768 | // No need to handle extraction error
769 | }
770 |
771 | const result = Array.from(tableNames);
772 | return result;
773 | }
774 |
775 | async function handleCompleteTuning(
776 | params: CompleteTuningParams,
777 | neonClient: Api<unknown>,
778 | extra: ToolHandlerExtraParams,
779 | ): Promise<CompleteTuningResult> {
780 | let results;
781 | const operationLog: string[] = [];
782 |
783 | try {
784 | // Validate branch information
785 | if (!params.temporaryBranch) {
786 | throw new Error(
787 | 'Branch information is required for completing query tuning',
788 | );
789 | }
790 |
791 | // Only proceed with changes if we have both suggestedChanges and branch
792 | if (
793 | params.applyChanges &&
794 | params.suggestedSqlStatements &&
795 | params.suggestedSqlStatements.length > 0
796 | ) {
797 | operationLog.push('Applying optimizations to main branch...');
798 |
799 | results = await handleRunSqlTransaction(
800 | {
801 | sqlStatements: params.suggestedSqlStatements,
802 | databaseName: params.databaseName,
803 | projectId: params.projectId,
804 | branchId: params.branch?.id,
805 | },
806 | neonClient,
807 | extra,
808 | );
809 |
810 | operationLog.push('Successfully applied optimizations to main branch.');
811 | } else {
812 | operationLog.push(
813 | 'No changes were applied (either none suggested or changes were discarded).',
814 | );
815 | }
816 |
817 | // Only delete branch if shouldDeleteTemporaryBranch is true
818 | if (params.shouldDeleteTemporaryBranch && params.temporaryBranch) {
819 | operationLog.push('Cleaning up temporary branch...');
820 |
821 | await handleDeleteBranch(
822 | {
823 | projectId: params.projectId,
824 | branchId: params.temporaryBranch.id,
825 | },
826 | neonClient,
827 | );
828 |
829 | operationLog.push('Successfully cleaned up temporary branch.');
830 | }
831 |
832 | const result: CompleteTuningResult = {
833 | appliedChanges:
834 | params.applyChanges && params.suggestedSqlStatements
835 | ? params.suggestedSqlStatements
836 | : undefined,
837 | results,
838 | deletedBranches:
839 | params.shouldDeleteTemporaryBranch && params.temporaryBranch
840 | ? [params.temporaryBranch.id]
841 | : undefined,
842 | message: operationLog.join('\n'),
843 | };
844 |
845 | return result;
846 | } catch (error) {
847 | throw new Error(
848 | `Failed to complete query tuning: ${(error as Error).message}`,
849 | );
850 | }
851 | }
852 |
853 | async function handleListSlowQueries(
854 | {
855 | projectId,
856 | branchId,
857 | databaseName,
858 | computeId,
859 | limit = 10,
860 | }: {
861 | projectId: string;
862 | branchId?: string;
863 | databaseName?: string;
864 | computeId?: string;
865 | limit?: number;
866 | },
867 | neonClient: Api<unknown>,
868 | extra: ToolHandlerExtraParams,
869 | ) {
870 | // Get connection string
871 | const connectionString = await handleGetConnectionString(
872 | {
873 | projectId,
874 | branchId,
875 | computeId,
876 | databaseName,
877 | },
878 | neonClient,
879 | extra,
880 | );
881 |
882 | // Connect to the database
883 | const sql = neon(connectionString.uri);
884 |
885 | // First, check if pg_stat_statements extension is installed
886 | const checkExtensionQuery = `
887 | SELECT EXISTS (
888 | SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements'
889 | ) as extension_exists;
890 | `;
891 |
892 | const extensionCheck = await sql.query(checkExtensionQuery);
893 | const extensionExists = extensionCheck[0]?.extension_exists;
894 |
895 | if (!extensionExists) {
896 | throw new NotFoundError(
897 | `pg_stat_statements extension is not installed on the database. Please install it using the following command: CREATE EXTENSION pg_stat_statements;`,
898 | );
899 | }
900 |
901 | // Query to get slow queries
902 | const slowQueriesQuery = `
903 | SELECT
904 | query,
905 | calls,
906 | total_exec_time,
907 | mean_exec_time,
908 | rows,
909 | shared_blks_hit,
910 | shared_blks_read,
911 | shared_blks_written,
912 | shared_blks_dirtied,
913 | temp_blks_read,
914 | temp_blks_written,
915 | wal_records,
916 | wal_fpi,
917 | wal_bytes
918 | FROM pg_stat_statements
919 | WHERE query NOT LIKE '%pg_stat_statements%'
920 | AND query NOT LIKE '%EXPLAIN%'
921 | ORDER BY mean_exec_time DESC
922 | LIMIT $1;
923 | `;
924 |
925 | const slowQueries = await sql.query(slowQueriesQuery, [limit]);
926 |
927 | // Format the results
928 | const formattedQueries = slowQueries.map((query: any) => {
929 | return {
930 | query: query.query,
931 | calls: query.calls,
932 | total_exec_time_ms: query.total_exec_time,
933 | mean_exec_time_ms: query.mean_exec_time,
934 | rows: query.rows,
935 | shared_blocks: {
936 | hit: query.shared_blks_hit,
937 | read: query.shared_blks_read,
938 | written: query.shared_blks_written,
939 | dirtied: query.shared_blks_dirtied,
940 | },
941 | temp_blocks: {
942 | read: query.temp_blks_read,
943 | written: query.temp_blks_written,
944 | },
945 | io_time: {
946 | read_ms: query.blk_read_time,
947 | write_ms: query.blk_write_time,
948 | },
949 | wal: {
950 | records: query.wal_records,
951 | full_page_images: query.wal_fpi,
952 | bytes: query.wal_bytes,
953 | },
954 | };
955 | });
956 |
957 | return {
958 | slow_queries: formattedQueries,
959 | total_queries_found: formattedQueries.length,
960 | };
961 | }
962 |
963 | async function handleListBranchComputes(
964 | {
965 | projectId,
966 | branchId,
967 | }: {
968 | projectId?: string;
969 | branchId?: string;
970 | },
971 | neonClient: Api<unknown>,
972 | extra: ToolHandlerExtraParams,
973 | ) {
974 | // If projectId is not provided, get the first project but only if there is only one project
975 | if (!projectId) {
976 | const projects = await handleListProjects({}, neonClient, extra);
977 | if (projects.length === 1) {
978 | projectId = projects[0].id;
979 | } else {
980 | throw new InvalidArgumentError(
981 | 'Please provide a project ID or ensure you have only one project in your account.',
982 | );
983 | }
984 | }
985 |
986 | let endpoints;
987 | if (branchId) {
988 | const response = await neonClient.listProjectBranchEndpoints(
989 | projectId,
990 | branchId,
991 | );
992 | endpoints = response.data.endpoints;
993 | } else {
994 | const response = await neonClient.listProjectEndpoints(projectId);
995 | endpoints = response.data.endpoints;
996 | }
997 |
998 | return endpoints.map((endpoint) => ({
999 | compute_id: endpoint.id,
1000 | compute_type: endpoint.type,
1001 | compute_size:
1002 | endpoint.autoscaling_limit_min_cu !== endpoint.autoscaling_limit_max_cu
1003 | ? `${endpoint.autoscaling_limit_min_cu}-${endpoint.autoscaling_limit_max_cu}`
1004 | : endpoint.autoscaling_limit_min_cu,
1005 | last_active: endpoint.last_active,
1006 | ...endpoint,
1007 | }));
1008 | }
1009 |
1010 | async function handleListSharedProjects(
1011 | params: ListSharedProjectsParams,
1012 | neonClient: Api<unknown>,
1013 | ) {
1014 | const response = await neonClient.listSharedProjects(params);
1015 | return response.data.projects;
1016 | }
1017 |
1018 | async function handleLoadResource({ subject }: { subject: string }) {
1019 | const resource = NEON_RESOURCES.find((r) => r.name === subject);
1020 | if (!resource) {
1021 | throw new InvalidArgumentError(`Resource not found: ${subject}`);
1022 | }
1023 |
1024 | try {
1025 | const url = new URL(resource.uri);
1026 | const path = url.pathname;
1027 | const content = await fetchRawGithubContent(path);
1028 | return content;
1029 | } catch (error) {
1030 | const errorMessage =
1031 | error instanceof Error ? error.message : 'Unknown error';
1032 | throw new Error(
1033 | `Failed to load resource "${resource.name}": ${errorMessage}`,
1034 | );
1035 | }
1036 | }
1037 |
1038 | async function handleCompareDatabaseSchema(
1039 | params: GetProjectBranchSchemaComparisonParams,
1040 | neonClient: Api<unknown>,
1041 | ) {
1042 | const response = await neonClient.getProjectBranchSchemaComparison(params);
1043 | return response.data;
1044 | }
1045 |
1046 | export const NEON_HANDLERS = {
1047 | list_projects: async ({ params }, neonClient, extra) => {
1048 | const organization = await getOrgByOrgIdOrDefault(
1049 | params,
1050 | neonClient,
1051 | extra,
1052 | );
1053 | const projects = await handleListProjects(
1054 | { ...params, org_id: organization?.id },
1055 | neonClient,
1056 | extra,
1057 | );
1058 | return {
1059 | content: [
1060 | {
1061 | type: 'text',
1062 | text: JSON.stringify(
1063 | {
1064 | organization: organization
1065 | ? {
1066 | name: organization.name,
1067 | id: organization.id,
1068 | }
1069 | : undefined,
1070 | projects,
1071 | },
1072 | null,
1073 | 2,
1074 | ),
1075 | },
1076 | ],
1077 | };
1078 | },
1079 |
1080 | create_project: async ({ params }, neonClient, extra) => {
1081 | try {
1082 | const organization = await getOrgByOrgIdOrDefault(
1083 | params,
1084 | neonClient,
1085 | extra,
1086 | );
1087 | const result = await handleCreateProject(
1088 | { project: { name: params.name, org_id: organization?.id } },
1089 | neonClient,
1090 | );
1091 |
1092 | // Get the connection string for the newly created project
1093 | const connectionString = await handleGetConnectionString(
1094 | {
1095 | projectId: result.project.id,
1096 | branchId: result.branch.id,
1097 | databaseName: result.databases[0].name,
1098 | },
1099 | neonClient,
1100 | extra,
1101 | );
1102 |
1103 | return {
1104 | content: [
1105 | {
1106 | type: 'text',
1107 | text: [
1108 | `Your Neon project is created ${organization ? `in organization "${organization.name}"` : ''} and is ready.`,
1109 | `The project_id is "${result.project.id}"`,
1110 | `The branch name is "${result.branch.name}" (ID: ${result.branch.id})`,
1111 | `There is one database available on this branch, called "${result.databases[0].name}",`,
1112 | 'but you can create more databases using SQL commands.',
1113 | '',
1114 | 'Connection string details:',
1115 | `URI: ${connectionString.uri}`,
1116 | `Project ID: ${connectionString.projectId}`,
1117 | `Branch ID: ${connectionString.branchId}`,
1118 | `Database: ${connectionString.databaseName}`,
1119 | `Role: ${connectionString.roleName}`,
1120 | '',
1121 | 'You can use this connection string with any PostgreSQL client to connect to your Neon database.',
1122 | 'For example, with psql:',
1123 | `psql "${connectionString.uri}"`,
1124 | ].join('\n'),
1125 | },
1126 | ],
1127 | };
1128 | } catch (error) {
1129 | const message = error instanceof Error ? error.message : 'Unknown error';
1130 | return {
1131 | isError: true,
1132 | content: [
1133 | {
1134 | type: 'text',
1135 | text: [
1136 | 'An error occurred while creating the project.',
1137 | 'Error details:',
1138 | message,
1139 | 'If you have reached the Neon project limit, please upgrade your account in this link: https://console.neon.tech/app/billing',
1140 | ].join('\n'),
1141 | },
1142 | ],
1143 | };
1144 | }
1145 | },
1146 |
1147 | delete_project: async ({ params }, neonClient) => {
1148 | await handleDeleteProject(params.projectId, neonClient);
1149 | return {
1150 | content: [
1151 | {
1152 | type: 'text',
1153 | text: [
1154 | 'Project deleted successfully.',
1155 | `Project ID: ${params.projectId}`,
1156 | ].join('\n'),
1157 | },
1158 | ],
1159 | };
1160 | },
1161 |
1162 | describe_project: async ({ params }, neonClient) => {
1163 | const result = await handleDescribeProject(params.projectId, neonClient);
1164 | return {
1165 | content: [
1166 | {
1167 | type: 'text',
1168 | text: `This project is called ${result.project.name}.`,
1169 | },
1170 | {
1171 | type: 'text',
1172 | text: `It contains the following branches (use the describe branch tool to learn more about each branch): ${JSON.stringify(
1173 | result.branches,
1174 | null,
1175 | 2,
1176 | )}`,
1177 | },
1178 | ],
1179 | };
1180 | },
1181 |
1182 | run_sql: async ({ params }, neonClient, extra) => {
1183 | const result = await handleRunSql(
1184 | {
1185 | sql: params.sql,
1186 | databaseName: params.databaseName,
1187 | projectId: params.projectId,
1188 | branchId: params.branchId,
1189 | },
1190 | neonClient,
1191 | extra,
1192 | );
1193 | return {
1194 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1195 | };
1196 | },
1197 |
1198 | run_sql_transaction: async ({ params }, neonClient, extra) => {
1199 | const result = await handleRunSqlTransaction(
1200 | {
1201 | sqlStatements: params.sqlStatements,
1202 | databaseName: params.databaseName,
1203 | projectId: params.projectId,
1204 | branchId: params.branchId,
1205 | },
1206 | neonClient,
1207 | extra,
1208 | );
1209 | return {
1210 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1211 | };
1212 | },
1213 |
1214 | describe_table_schema: async ({ params }, neonClient, extra) => {
1215 | const result = await handleDescribeTableSchema(
1216 | {
1217 | tableName: params.tableName,
1218 | databaseName: params.databaseName,
1219 | projectId: params.projectId,
1220 | branchId: params.branchId,
1221 | },
1222 | neonClient,
1223 | extra,
1224 | );
1225 | return {
1226 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1227 | };
1228 | },
1229 |
1230 | get_database_tables: async ({ params }, neonClient, extra) => {
1231 | const result = await handleGetDatabaseTables(
1232 | {
1233 | projectId: params.projectId,
1234 | branchId: params.branchId,
1235 | databaseName: params.databaseName,
1236 | },
1237 | neonClient,
1238 | extra,
1239 | );
1240 | return {
1241 | content: [
1242 | {
1243 | type: 'text',
1244 | text: JSON.stringify(result, null, 2),
1245 | },
1246 | ],
1247 | };
1248 | },
1249 |
1250 | create_branch: async ({ params }, neonClient) => {
1251 | const result = await handleCreateBranch(
1252 | {
1253 | projectId: params.projectId,
1254 | branchName: params.branchName,
1255 | },
1256 | neonClient,
1257 | );
1258 | return {
1259 | content: [
1260 | {
1261 | type: 'text',
1262 | text: [
1263 | 'Branch created successfully.',
1264 | `Project ID: ${result.branch.project_id}`,
1265 | `Branch ID: ${result.branch.id}`,
1266 | `Branch name: ${result.branch.name}`,
1267 | `Parent branch: ${result.branch.parent_id}`,
1268 | ].join('\n'),
1269 | },
1270 | ],
1271 | };
1272 | },
1273 |
1274 | prepare_database_migration: async ({ params }, neonClient, extra) => {
1275 | const result = await handleSchemaMigration(
1276 | {
1277 | migrationSql: params.migrationSql,
1278 | databaseName: params.databaseName,
1279 | projectId: params.projectId,
1280 | },
1281 | neonClient,
1282 | extra,
1283 | );
1284 | return {
1285 | content: [
1286 | {
1287 | type: 'text',
1288 | text: `
1289 | <status>Migration created successfully in temporary branch</status>
1290 | <details>
1291 | <migration_id>${result.migrationId}</migration_id>
1292 | <temporary_branch>
1293 | <name>${result.branch.name}</name>
1294 | <id>${result.branch.id}</id>
1295 | </temporary_branch>
1296 | </details>
1297 | <execution_result>${JSON.stringify(result.migrationResult, null, 2)}</execution_result>
1298 |
1299 | <next_actions>
1300 | You MUST follow these steps:
1301 | 1. Test this migration using \`run_sql\` tool on branch \`${result.branch.name}\`
1302 | 2. Verify the changes meet your requirements
1303 | 3. If satisfied, use \`complete_database_migration\` with migration_id: ${result.migrationId}
1304 | </next_actions>
1305 | `,
1306 | },
1307 | ],
1308 | };
1309 | },
1310 |
1311 | complete_database_migration: async ({ params }, neonClient, extra) => {
1312 | const result = await handleCommitMigration(
1313 | {
1314 | migrationId: params.migrationId,
1315 | },
1316 | neonClient,
1317 | extra,
1318 | );
1319 | return {
1320 | content: [
1321 | {
1322 | type: 'text',
1323 | text: `Result: ${JSON.stringify(
1324 | {
1325 | deletedBranch: result.deletedBranch,
1326 | migrationResult: result.migrationResult,
1327 | },
1328 | null,
1329 | 2,
1330 | )}`,
1331 | },
1332 | ],
1333 | };
1334 | },
1335 |
1336 | describe_branch: async ({ params }, neonClient, extra) => {
1337 | return await handleDescribeBranch(
1338 | {
1339 | projectId: params.projectId,
1340 | branchId: params.branchId,
1341 | databaseName: params.databaseName,
1342 | },
1343 | neonClient,
1344 | extra,
1345 | );
1346 | },
1347 |
1348 | delete_branch: async ({ params }, neonClient) => {
1349 | await handleDeleteBranch(
1350 | {
1351 | projectId: params.projectId,
1352 | branchId: params.branchId,
1353 | },
1354 | neonClient,
1355 | );
1356 | return {
1357 | content: [
1358 | {
1359 | type: 'text',
1360 | text: [
1361 | 'Branch deleted successfully.',
1362 | `Project ID: ${params.projectId}`,
1363 | `Branch ID: ${params.branchId}`,
1364 | ].join('\n'),
1365 | },
1366 | ],
1367 | };
1368 | },
1369 |
1370 | reset_from_parent: async ({ params }, neonClient) => {
1371 | const result = await handleResetFromParent(
1372 | {
1373 | projectId: params.projectId,
1374 | branchIdOrName: params.branchIdOrName,
1375 | preserveUnderName: params.preserveUnderName,
1376 | },
1377 | neonClient,
1378 | );
1379 |
1380 | const parentInfo = `${result.parentBranch.name} (${result.parentBranch.id})`;
1381 |
1382 | const messages = [
1383 | 'Branch reset from parent successfully.',
1384 | `Project: ${params.projectId}`,
1385 | `Branch: ${params.branchIdOrName}`,
1386 | `Reset to parent branch: ${parentInfo}`,
1387 | ];
1388 |
1389 | if (result.preservedBranchName) {
1390 | messages.push(
1391 | params.preserveUnderName
1392 | ? `Previous state preserved as: ${params.preserveUnderName}`
1393 | : `Previous state auto-preserved as: ${result.preservedBranchName} (branch had children)`,
1394 | );
1395 | } else {
1396 | messages.push('Previous state was not preserved');
1397 | }
1398 |
1399 | return {
1400 | content: [
1401 | {
1402 | type: 'text',
1403 | text: messages.join('\n'),
1404 | },
1405 | ],
1406 | };
1407 | },
1408 |
1409 | get_connection_string: async ({ params }, neonClient, extra) => {
1410 | const result = await handleGetConnectionString(
1411 | {
1412 | projectId: params.projectId,
1413 | branchId: params.branchId,
1414 | computeId: params.computeId,
1415 | databaseName: params.databaseName,
1416 | roleName: params.roleName,
1417 | },
1418 | neonClient,
1419 | extra,
1420 | );
1421 | return {
1422 | content: [
1423 | {
1424 | type: 'text',
1425 | text: [
1426 | 'Connection string details:',
1427 | `URI: ${result.uri}`,
1428 | `Project ID: ${result.projectId}`,
1429 | `Database: ${result.databaseName}`,
1430 | `Role: ${result.roleName}`,
1431 | result.branchId
1432 | ? `Branch ID: ${result.branchId}`
1433 | : 'Using default branch',
1434 | result.computeId
1435 | ? `Compute ID: ${result.computeId}`
1436 | : 'Using default compute',
1437 | '',
1438 | 'You can use this connection string with any PostgreSQL client to connect to your Neon database.',
1439 | ].join('\n'),
1440 | },
1441 | ],
1442 | };
1443 | },
1444 |
1445 | provision_neon_auth: async ({ params }, neonClient, extra) => {
1446 | const result = await handleProvisionNeonAuth(
1447 | {
1448 | projectId: params.projectId,
1449 | branchId: params.branchId,
1450 | databaseName: params.databaseName,
1451 | },
1452 | neonClient,
1453 | extra,
1454 | );
1455 | return result;
1456 | },
1457 |
1458 | explain_sql_statement: async ({ params }, neonClient, extra) => {
1459 | const result = await handleExplainSqlStatement(
1460 | { params },
1461 | neonClient,
1462 | extra,
1463 | );
1464 | return result;
1465 | },
1466 |
1467 | prepare_query_tuning: async ({ params }, neonClient, extra) => {
1468 | const result = await handleQueryTuning(
1469 | {
1470 | sql: params.sql,
1471 | databaseName: params.databaseName,
1472 | projectId: params.projectId,
1473 | },
1474 | neonClient,
1475 | extra,
1476 | );
1477 | return {
1478 | content: [
1479 | {
1480 | type: 'text',
1481 | text: JSON.stringify(
1482 | {
1483 | tuningId: result.tuningId,
1484 | databaseName: result.databaseName,
1485 | projectId: result.projectId,
1486 | temporaryBranch: result.temporaryBranch,
1487 | executionPlan: result.originalPlan,
1488 | tableSchemas: result.tableSchemas,
1489 | sql: result.sql,
1490 | },
1491 | null,
1492 | 2,
1493 | ),
1494 | },
1495 | ],
1496 | };
1497 | },
1498 |
1499 | complete_query_tuning: async ({ params }, neonClient, extra) => {
1500 | const result = await handleCompleteTuning(
1501 | {
1502 | suggestedSqlStatements: params.suggestedSqlStatements,
1503 | applyChanges: params.applyChanges,
1504 | tuningId: params.tuningId,
1505 | databaseName: params.databaseName,
1506 | projectId: params.projectId,
1507 | temporaryBranch: {
1508 | id: params.temporaryBranchId,
1509 | project_id: params.projectId,
1510 | } as Branch,
1511 | shouldDeleteTemporaryBranch: params.shouldDeleteTemporaryBranch,
1512 | branch: params.branchId
1513 | ? ({ id: params.branchId, project_id: params.projectId } as Branch)
1514 | : undefined,
1515 | },
1516 | neonClient,
1517 | extra,
1518 | );
1519 |
1520 | return {
1521 | content: [
1522 | {
1523 | type: 'text',
1524 | text: JSON.stringify(result, null, 2),
1525 | },
1526 | ],
1527 | };
1528 | },
1529 |
1530 | list_slow_queries: async ({ params }, neonClient, extra) => {
1531 | const result = await handleListSlowQueries(
1532 | {
1533 | projectId: params.projectId,
1534 | branchId: params.branchId,
1535 | databaseName: params.databaseName,
1536 | computeId: params.computeId,
1537 | limit: params.limit,
1538 | },
1539 | neonClient,
1540 | extra,
1541 | );
1542 | return {
1543 | content: [
1544 | {
1545 | type: 'text',
1546 | text: JSON.stringify(result, null, 2),
1547 | },
1548 | ],
1549 | };
1550 | },
1551 |
1552 | list_branch_computes: async ({ params }, neonClient, extra) => {
1553 | const result = await handleListBranchComputes(
1554 | {
1555 | projectId: params.projectId,
1556 | branchId: params.branchId,
1557 | },
1558 | neonClient,
1559 | extra,
1560 | );
1561 | return {
1562 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1563 | };
1564 | },
1565 |
1566 | list_organizations: async ({ params }, neonClient, extra) => {
1567 | const organizations = await handleListOrganizations(
1568 | neonClient,
1569 | extra.account,
1570 | params.search,
1571 | );
1572 | return {
1573 | content: [
1574 | {
1575 | type: 'text',
1576 | text: JSON.stringify(organizations, null, 2),
1577 | },
1578 | ],
1579 | };
1580 | },
1581 |
1582 | list_shared_projects: async ({ params }, neonClient) => {
1583 | const sharedProjects = await handleListSharedProjects(params, neonClient);
1584 | return {
1585 | content: [
1586 | {
1587 | type: 'text',
1588 | text: JSON.stringify(
1589 | {
1590 | shared_projects: sharedProjects,
1591 | count: sharedProjects.length,
1592 | },
1593 | null,
1594 | 2,
1595 | ),
1596 | },
1597 | ],
1598 | };
1599 | },
1600 |
1601 | compare_database_schema: async ({ params }, neonClient) => {
1602 | const result = await handleCompareDatabaseSchema(
1603 | {
1604 | projectId: params.projectId,
1605 | branchId: params.branchId,
1606 | db_name: params.databaseName,
1607 | },
1608 | neonClient,
1609 | );
1610 | return {
1611 | content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
1612 | };
1613 | },
1614 |
1615 | search: async ({ params }, neonClient, extra) => {
1616 | return await handleSearch(params, neonClient, extra);
1617 | },
1618 |
1619 | fetch: async ({ params }, neonClient, extra) => {
1620 | return await handleFetch(params, neonClient, extra);
1621 | },
1622 |
1623 | load_resource: async ({ params }) => {
1624 | const content = await handleLoadResource({ subject: params.subject });
1625 | return {
1626 | content: [
1627 | {
1628 | type: 'text',
1629 | text: content,
1630 | },
1631 | ],
1632 | };
1633 | },
1634 | } satisfies ToolHandlers;
1635 |
```