Skip to content

Commit 23e4522

Browse files
tianzhouclaude
tianzhou
and
claude
authored
feat: add table filter to search_objects tool (#172) (#177)
Add optional table parameter to search_objects tool for filtering columns and indexes by specific table name. This addresses performance issues when searching in large databases with thousands of columns. Changes: - Add table parameter (exact match, requires schema) - Optimize searchColumns() to only query specified table - Optimize searchIndexes() to only query specified table - Add validation: table requires schema parameter - Add validation: table only applies to column/index object types - Clarify that schema parameter is exact match (not pattern) - Add comprehensive test coverage Closes #172 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude <noreply@anthropic.com>
1 parent 938b255 commit 23e4522

File tree

2 files changed

+262
-7
lines changed

2 files changed

+262
-7
lines changed

‎src/tools/__tests__/search-objects.test.ts‎

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,220 @@ describe('search_database_objects tool', () => {
483483
});
484484
});
485485

486+
describe('table filter', () => {
487+
describe('for columns', () => {
488+
beforeEach(() => {
489+
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);
490+
});
491+
492+
it('should filter columns by table when table parameter is provided', async () => {
493+
const usersColumns: TableColumn[] = [
494+
{ column_name: 'id', data_type: 'INTEGER', is_nullable: 'NO', column_default: null },
495+
{ column_name: 'name', data_type: 'TEXT', is_nullable: 'YES', column_default: null },
496+
];
497+
498+
const ordersColumns: TableColumn[] = [
499+
{ column_name: 'id', data_type: 'INTEGER', is_nullable: 'NO', column_default: null },
500+
{ column_name: 'user_id', data_type: 'INTEGER', is_nullable: 'NO', column_default: null },
501+
];
502+
503+
vi.mocked(mockConnector.getTableSchema).mockImplementation(async (table) => {
504+
if (table === 'users') return usersColumns;
505+
if (table === 'orders') return ordersColumns;
506+
return [];
507+
});
508+
509+
const handler = createSearchDatabaseObjectsToolHandler();
510+
const result = await handler(
511+
{
512+
object_type: 'column',
513+
pattern: '%',
514+
schema: 'public',
515+
table: 'users',
516+
detail_level: 'names',
517+
},
518+
null
519+
);
520+
521+
const parsed = parseToolResponse(result);
522+
expect(parsed.success).toBe(true);
523+
expect(parsed.data.count).toBe(2);
524+
expect(parsed.data.results).toEqual([
525+
{ name: 'id', table: 'users', schema: 'public' },
526+
{ name: 'name', table: 'users', schema: 'public' },
527+
]);
528+
// Verify only users table was queried
529+
expect(mockConnector.getTableSchema).toHaveBeenCalledWith('users', 'public');
530+
expect(mockConnector.getTableSchema).not.toHaveBeenCalledWith('orders', 'public');
531+
});
532+
533+
it('should require schema when table is specified', async () => {
534+
const handler = createSearchDatabaseObjectsToolHandler();
535+
const result = await handler(
536+
{
537+
object_type: 'column',
538+
pattern: '%',
539+
table: 'users',
540+
detail_level: 'names',
541+
},
542+
null
543+
);
544+
545+
expect(result.isError).toBe(true);
546+
const parsed = parseToolResponse(result);
547+
expect(parsed.code).toBe('SCHEMA_REQUIRED');
548+
expect(parsed.error).toContain("'table' parameter requires 'schema'");
549+
});
550+
551+
it('should work with column pattern when table filter is applied', async () => {
552+
const usersColumns: TableColumn[] = [
553+
{ column_name: 'id', data_type: 'INTEGER', is_nullable: 'NO', column_default: null },
554+
{ column_name: 'name', data_type: 'TEXT', is_nullable: 'YES', column_default: null },
555+
{ column_name: 'email', data_type: 'TEXT', is_nullable: 'YES', column_default: null },
556+
];
557+
558+
vi.mocked(mockConnector.getTableSchema).mockResolvedValue(usersColumns);
559+
560+
const handler = createSearchDatabaseObjectsToolHandler();
561+
const result = await handler(
562+
{
563+
object_type: 'column',
564+
pattern: '%e%',
565+
schema: 'public',
566+
table: 'users',
567+
detail_level: 'names',
568+
},
569+
null
570+
);
571+
572+
const parsed = parseToolResponse(result);
573+
expect(parsed.data.count).toBe(2);
574+
expect(parsed.data.results.map((r: any) => r.name)).toEqual(['name', 'email']);
575+
});
576+
});
577+
578+
describe('for indexes', () => {
579+
beforeEach(() => {
580+
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);
581+
});
582+
583+
it('should filter indexes by table when table parameter is provided', async () => {
584+
const usersIndexes: TableIndex[] = [
585+
{
586+
index_name: 'users_pkey',
587+
column_names: ['id'],
588+
is_unique: true,
589+
is_primary: true,
590+
},
591+
{
592+
index_name: 'users_email_idx',
593+
column_names: ['email'],
594+
is_unique: true,
595+
is_primary: false,
596+
},
597+
];
598+
599+
const ordersIndexes: TableIndex[] = [
600+
{
601+
index_name: 'orders_pkey',
602+
column_names: ['id'],
603+
is_unique: true,
604+
is_primary: true,
605+
},
606+
];
607+
608+
vi.mocked(mockConnector.getTableIndexes).mockImplementation(async (table) => {
609+
if (table === 'users') return usersIndexes;
610+
if (table === 'orders') return ordersIndexes;
611+
return [];
612+
});
613+
614+
const handler = createSearchDatabaseObjectsToolHandler();
615+
const result = await handler(
616+
{
617+
object_type: 'index',
618+
pattern: '%',
619+
schema: 'public',
620+
table: 'users',
621+
detail_level: 'names',
622+
},
623+
null
624+
);
625+
626+
const parsed = parseToolResponse(result);
627+
expect(parsed.success).toBe(true);
628+
expect(parsed.data.count).toBe(2);
629+
expect(parsed.data.results.map((r: any) => r.name)).toEqual(['users_pkey', 'users_email_idx']);
630+
// Verify only users table was queried
631+
expect(mockConnector.getTableIndexes).toHaveBeenCalledWith('users', 'public');
632+
expect(mockConnector.getTableIndexes).not.toHaveBeenCalledWith('orders', 'public');
633+
});
634+
});
635+
636+
describe('validation', () => {
637+
it('should reject table parameter for schema object type', async () => {
638+
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);
639+
640+
const handler = createSearchDatabaseObjectsToolHandler();
641+
const result = await handler(
642+
{
643+
object_type: 'schema',
644+
pattern: '%',
645+
schema: 'public',
646+
table: 'users',
647+
detail_level: 'names',
648+
},
649+
null
650+
);
651+
652+
expect(result.isError).toBe(true);
653+
const parsed = parseToolResponse(result);
654+
expect(parsed.code).toBe('INVALID_TABLE_FILTER');
655+
expect(parsed.error).toContain("only applies to object_type 'column' or 'index'");
656+
});
657+
658+
it('should reject table parameter for table object type', async () => {
659+
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);
660+
661+
const handler = createSearchDatabaseObjectsToolHandler();
662+
const result = await handler(
663+
{
664+
object_type: 'table',
665+
pattern: '%',
666+
schema: 'public',
667+
table: 'users',
668+
detail_level: 'names',
669+
},
670+
null
671+
);
672+
673+
expect(result.isError).toBe(true);
674+
const parsed = parseToolResponse(result);
675+
expect(parsed.code).toBe('INVALID_TABLE_FILTER');
676+
});
677+
678+
it('should reject table parameter for procedure object type', async () => {
679+
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);
680+
681+
const handler = createSearchDatabaseObjectsToolHandler();
682+
const result = await handler(
683+
{
684+
object_type: 'procedure',
685+
pattern: '%',
686+
schema: 'public',
687+
table: 'users',
688+
detail_level: 'names',
689+
},
690+
null
691+
);
692+
693+
expect(result.isError).toBe(true);
694+
const parsed = parseToolResponse(result);
695+
expect(parsed.code).toBe('INVALID_TABLE_FILTER');
696+
});
697+
});
698+
});
699+
486700
describe('error handling', () => {
487701
it('should validate schema exists', async () => {
488702
vi.mocked(mockConnector.getSchemas).mockResolvedValue(['public']);

‎src/tools/search-objects.ts‎

Lines changed: 48 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@ export const searchDatabaseObjectsSchema = {
3030
schema: z
3131
.string()
3232
.optional()
33-
.describe("Filter results to a specific schema/database"),
33+
.describe("Filter results to a specific schema/database (exact match)"),
34+
table: z
35+
.string()
36+
.optional()
37+
.describe("Filter to a specific table (exact match). Requires schema parameter. Only applies to columns and indexes."),
3438
detail_level: z
3539
.enum(["names", "summary", "full"])
3640
.default("names")
@@ -227,6 +231,7 @@ async function searchColumns(
227231
connector: Connector,
228232
pattern: string,
229233
schemaFilter: string | undefined,
234+
tableFilter: string | undefined,
230235
detailLevel: DetailLevel,
231236
limit: number
232237
): Promise<any[]> {
@@ -246,9 +251,17 @@ async function searchColumns(
246251
if (results.length >= limit) break;
247252

248253
try {
249-
const tables = await connector.getTables(schemaName);
254+
// Get tables to search
255+
let tablesToSearch: string[];
256+
if (tableFilter) {
257+
// If table filter is specified, only search that table
258+
tablesToSearch = [tableFilter];
259+
} else {
260+
// Otherwise search all tables in the schema
261+
tablesToSearch = await connector.getTables(schemaName);
262+
}
250263

251-
for (const tableName of tables) {
264+
for (const tableName of tablesToSearch) {
252265
if (results.length >= limit) break;
253266

254267
try {
@@ -365,6 +378,7 @@ async function searchIndexes(
365378
connector: Connector,
366379
pattern: string,
367380
schemaFilter: string | undefined,
381+
tableFilter: string | undefined,
368382
detailLevel: DetailLevel,
369383
limit: number
370384
): Promise<any[]> {
@@ -384,9 +398,17 @@ async function searchIndexes(
384398
if (results.length >= limit) break;
385399

386400
try {
387-
const tables = await connector.getTables(schemaName);
401+
// Get tables to search
402+
let tablesToSearch: string[];
403+
if (tableFilter) {
404+
// If table filter is specified, only search that table
405+
tablesToSearch = [tableFilter];
406+
} else {
407+
// Otherwise search all tables in the schema
408+
tablesToSearch = await connector.getTables(schemaName);
409+
}
388410

389-
for (const tableName of tables) {
411+
for (const tableName of tablesToSearch) {
390412
if (results.length >= limit) break;
391413

392414
try {
@@ -437,12 +459,14 @@ export function createSearchDatabaseObjectsToolHandler(sourceId?: string) {
437459
object_type,
438460
pattern = "%",
439461
schema,
462+
table,
440463
detail_level = "names",
441464
limit = 100,
442465
} = args as {
443466
object_type: DatabaseObjectType;
444467
pattern?: string;
445468
schema?: string;
469+
table?: string;
446470
detail_level: DetailLevel;
447471
limit: number;
448472
};
@@ -452,6 +476,22 @@ export function createSearchDatabaseObjectsToolHandler(sourceId?: string) {
452476

453477
// Tool is already registered, so it's enabled (no need to check)
454478

479+
// Validate table parameter
480+
if (table) {
481+
if (!schema) {
482+
return createToolErrorResponse(
483+
"The 'table' parameter requires 'schema' to be specified",
484+
"SCHEMA_REQUIRED"
485+
);
486+
}
487+
if (!["column", "index"].includes(object_type)) {
488+
return createToolErrorResponse(
489+
`The 'table' parameter only applies to object_type 'column' or 'index', not '${object_type}'`,
490+
"INVALID_TABLE_FILTER"
491+
);
492+
}
493+
}
494+
455495
// Validate schema if provided
456496
if (schema) {
457497
const schemas = await connector.getSchemas();
@@ -474,13 +514,13 @@ export function createSearchDatabaseObjectsToolHandler(sourceId?: string) {
474514
results = await searchTables(connector, pattern, schema, detail_level, limit);
475515
break;
476516
case "column":
477-
results = await searchColumns(connector, pattern, schema, detail_level, limit);
517+
results = await searchColumns(connector, pattern, schema, table, detail_level, limit);
478518
break;
479519
case "procedure":
480520
results = await searchProcedures(connector, pattern, schema, detail_level, limit);
481521
break;
482522
case "index":
483-
results = await searchIndexes(connector, pattern, schema, detail_level, limit);
523+
results = await searchIndexes(connector, pattern, schema, table, detail_level, limit);
484524
break;
485525
default:
486526
return createToolErrorResponse(
@@ -493,6 +533,7 @@ export function createSearchDatabaseObjectsToolHandler(sourceId?: string) {
493533
object_type,
494534
pattern,
495535
schema,
536+
table,
496537
detail_level,
497538
count: results.length,
498539
results,

0 commit comments

Comments
 (0)