mcp-dbutils 0.17.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mcp_dbutils/base.py CHANGED
@@ -18,16 +18,22 @@ from .stats import ResourceStats
18
18
 
19
19
  class ConnectionHandlerError(Exception):
20
20
  """Base exception for connection errors"""
21
+
21
22
  pass
22
23
 
24
+
23
25
  class ConfigurationError(ConnectionHandlerError):
24
26
  """Configuration related errors"""
27
+
25
28
  pass
26
29
 
30
+
27
31
  class ConnectionError(ConnectionHandlerError):
28
32
  """Connection related errors"""
33
+
29
34
  pass
30
35
 
36
+
31
37
  # 常量定义
32
38
  DATABASE_CONNECTION_NAME = "Database connection name"
33
39
  EMPTY_QUERY_ERROR = "SQL query cannot be empty"
@@ -44,14 +50,15 @@ pkg_meta = metadata("mcp-dbutils")
44
50
  LOG_NAME = "dbutils"
45
51
 
46
52
  # MCP日志级别常量
47
- LOG_LEVEL_DEBUG = "debug" # 0
48
- LOG_LEVEL_INFO = "info" # 1
49
- LOG_LEVEL_NOTICE = "notice" # 2
53
+ LOG_LEVEL_DEBUG = "debug" # 0
54
+ LOG_LEVEL_INFO = "info" # 1
55
+ LOG_LEVEL_NOTICE = "notice" # 2
50
56
  LOG_LEVEL_WARNING = "warning" # 3
51
- LOG_LEVEL_ERROR = "error" # 4
52
- LOG_LEVEL_CRITICAL = "critical" # 5
53
- LOG_LEVEL_ALERT = "alert" # 6
54
- LOG_LEVEL_EMERGENCY = "emergency" # 7
57
+ LOG_LEVEL_ERROR = "error" # 4
58
+ LOG_LEVEL_CRITICAL = "critical" # 5
59
+ LOG_LEVEL_ALERT = "alert" # 6
60
+ LOG_LEVEL_EMERGENCY = "emergency" # 7
61
+
55
62
 
56
63
  class ConnectionHandler(ABC):
57
64
  """Abstract base class defining common interface for connection handlers"""
@@ -74,19 +81,18 @@ class ConnectionHandler(ABC):
74
81
 
75
82
  def send_log(self, level: str, message: str):
76
83
  """通过MCP发送日志消息和写入stderr
77
-
84
+
78
85
  Args:
79
86
  level: 日志级别 (debug/info/notice/warning/error/critical/alert/emergency)
80
87
  message: 日志内容
81
88
  """
82
89
  # 本地stderr日志
83
90
  self.log(level, message)
84
-
91
+
85
92
  # MCP日志通知
86
- if self._session and hasattr(self._session, 'request_context'):
93
+ if self._session and hasattr(self._session, "request_context"):
87
94
  self._session.request_context.session.send_log_message(
88
- level=level,
89
- data=message
95
+ level=level, data=message
90
96
  )
91
97
 
92
98
  @property
@@ -119,12 +125,18 @@ class ConnectionHandler(ABC):
119
125
  duration = (datetime.now() - start_time).total_seconds()
120
126
  self.stats.record_query_duration(sql, duration)
121
127
  self.stats.update_memory_usage(result)
122
- self.send_log(LOG_LEVEL_INFO, f"Query executed in {duration*1000:.2f}ms. Resource stats: {json.dumps(self.stats.to_dict())}")
128
+ self.send_log(
129
+ LOG_LEVEL_INFO,
130
+ f"Query executed in {duration * 1000:.2f}ms. Resource stats: {json.dumps(self.stats.to_dict())}",
131
+ )
123
132
  return result
124
133
  except Exception as e:
125
134
  duration = (datetime.now() - start_time).total_seconds()
126
135
  self.stats.record_error(e.__class__.__name__)
127
- self.send_log(LOG_LEVEL_ERROR, f"Query error after {duration*1000:.2f}ms - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}")
136
+ self.send_log(
137
+ LOG_LEVEL_ERROR,
138
+ f"Query error after {duration * 1000:.2f}ms - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}",
139
+ )
128
140
  raise
129
141
 
130
142
  @abstractmethod
@@ -199,12 +211,23 @@ class ConnectionHandler(ABC):
199
211
  """
200
212
  pass
201
213
 
214
+ @abstractmethod
215
+ async def test_connection(self) -> bool:
216
+ """Test database connection
217
+
218
+ Returns:
219
+ bool: True if connection is successful, False otherwise
220
+ """
221
+ pass
222
+
202
223
  @abstractmethod
203
224
  async def cleanup(self):
204
225
  """Cleanup resources"""
205
226
  pass
206
227
 
207
- async def execute_tool_query(self, tool_name: str, table_name: str = "", sql: str = "") -> str:
228
+ async def execute_tool_query(
229
+ self, tool_name: str, table_name: str = "", sql: str = ""
230
+ ) -> str:
208
231
  """Execute a tool query and return formatted result
209
232
 
210
233
  Args:
@@ -217,7 +240,7 @@ class ConnectionHandler(ABC):
217
240
  """
218
241
  try:
219
242
  self.stats.record_query()
220
-
243
+
221
244
  if tool_name == "dbutils-describe-table":
222
245
  result = await self.get_table_description(table_name)
223
246
  elif tool_name == "dbutils-get-ddl":
@@ -234,16 +257,22 @@ class ConnectionHandler(ABC):
234
257
  result = await self.explain_query(sql)
235
258
  else:
236
259
  raise ValueError(f"Unknown tool: {tool_name}")
237
-
260
+
238
261
  self.stats.update_memory_usage(result)
239
- self.send_log(LOG_LEVEL_INFO, f"Resource stats: {json.dumps(self.stats.to_dict())}")
262
+ self.send_log(
263
+ LOG_LEVEL_INFO, f"Resource stats: {json.dumps(self.stats.to_dict())}"
264
+ )
240
265
  return f"[{self.db_type}]\n{result}"
241
-
266
+
242
267
  except Exception as e:
243
268
  self.stats.record_error(e.__class__.__name__)
244
- self.send_log(LOG_LEVEL_ERROR, f"Tool error - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}")
269
+ self.send_log(
270
+ LOG_LEVEL_ERROR,
271
+ f"Tool error - {str(e)}\nResource stats: {json.dumps(self.stats.to_dict())}",
272
+ )
245
273
  raise
246
274
 
275
+
247
276
  class ConnectionServer:
248
277
  """Unified connection server class"""
249
278
 
@@ -259,36 +288,31 @@ class ConnectionServer:
259
288
  # 获取包信息用于服务器配置
260
289
  pkg_meta = metadata("mcp-dbutils")
261
290
  self.logger = create_logger(f"{LOG_NAME}.server", debug)
262
- self.server = Server(
263
- name=LOG_NAME,
264
- version=pkg_meta["Version"]
265
- )
291
+ self.server = Server(name=LOG_NAME, version=pkg_meta["Version"])
266
292
  self._session = None
267
293
  self._setup_handlers()
268
294
  self._setup_prompts()
269
295
 
270
296
  def send_log(self, level: str, message: str):
271
297
  """通过MCP发送日志消息和写入stderr
272
-
298
+
273
299
  Args:
274
300
  level: 日志级别 (debug/info/notice/warning/error/critical/alert/emergency)
275
301
  message: 日志内容
276
302
  """
277
303
  # 本地stderr日志
278
304
  self.logger(level, message)
279
-
305
+
280
306
  # MCP日志通知
281
- if hasattr(self.server, 'session') and self.server.session:
307
+ if hasattr(self.server, "session") and self.server.session:
282
308
  try:
283
- self.server.session.send_log_message(
284
- level=level,
285
- data=message
286
- )
309
+ self.server.session.send_log_message(level=level, data=message)
287
310
  except Exception as e:
288
311
  self.logger("error", f"Failed to send MCP log message: {str(e)}")
289
312
 
290
313
  def _setup_prompts(self):
291
314
  """Setup prompts handlers"""
315
+
292
316
  @self.server.list_prompts()
293
317
  async def handle_list_prompts() -> list[types.Prompt]:
294
318
  """Handle prompts/list request"""
@@ -301,64 +325,79 @@ class ConnectionServer:
301
325
 
302
326
  def _get_config_or_raise(self, connection: str) -> dict:
303
327
  """读取配置文件并验证连接配置
304
-
328
+
305
329
  Args:
306
330
  connection: 连接名称
307
-
331
+
308
332
  Returns:
309
333
  dict: 连接配置
310
-
334
+
311
335
  Raises:
312
336
  ConfigurationError: 如果配置文件格式不正确或连接不存在
313
337
  """
314
- with open(self.config_path, 'r') as f:
338
+ with open(self.config_path, "r") as f:
315
339
  config = yaml.safe_load(f)
316
- if not config or 'connections' not in config:
317
- raise ConfigurationError("Configuration file must contain 'connections' section")
318
- if connection not in config['connections']:
319
- available_connections = list(config['connections'].keys())
320
- raise ConfigurationError(f"Connection not found: {connection}. Available connections: {available_connections}")
321
-
322
- db_config = config['connections'][connection]
323
-
324
- if 'type' not in db_config:
325
- raise ConfigurationError("Database configuration must include 'type' field")
326
-
340
+ if not config or "connections" not in config:
341
+ raise ConfigurationError(
342
+ "Configuration file must contain 'connections' section"
343
+ )
344
+ if connection not in config["connections"]:
345
+ available_connections = list(config["connections"].keys())
346
+ raise ConfigurationError(
347
+ f"Connection not found: {connection}. Available connections: {available_connections}"
348
+ )
349
+
350
+ db_config = config["connections"][connection]
351
+
352
+ if "type" not in db_config:
353
+ raise ConfigurationError(
354
+ "Database configuration must include 'type' field"
355
+ )
356
+
327
357
  return db_config
328
-
329
- def _create_handler_for_type(self, db_type: str, connection: str) -> ConnectionHandler:
358
+
359
+ def _create_handler_for_type(
360
+ self, db_type: str, connection: str
361
+ ) -> ConnectionHandler:
330
362
  """基于数据库类型创建相应的处理器
331
-
363
+
332
364
  Args:
333
365
  db_type: 数据库类型
334
366
  connection: 连接名称
335
-
367
+
336
368
  Returns:
337
369
  ConnectionHandler: 数据库连接处理器
338
-
370
+
339
371
  Raises:
340
372
  ConfigurationError: 如果数据库类型不支持或导入失败
341
373
  """
342
374
  self.send_log(LOG_LEVEL_DEBUG, f"Creating handler for database type: {db_type}")
343
-
375
+
344
376
  try:
345
- if db_type == 'sqlite':
377
+ if db_type == "sqlite":
346
378
  from .sqlite.handler import SQLiteHandler
379
+
347
380
  return SQLiteHandler(self.config_path, connection, self.debug)
348
- elif db_type == 'postgres':
381
+ elif db_type == "postgres":
349
382
  from .postgres.handler import PostgreSQLHandler
383
+
350
384
  return PostgreSQLHandler(self.config_path, connection, self.debug)
351
- elif db_type == 'mysql':
385
+ elif db_type == "mysql":
352
386
  from .mysql.handler import MySQLHandler
387
+
353
388
  return MySQLHandler(self.config_path, connection, self.debug)
354
389
  else:
355
390
  raise ConfigurationError(f"Unsupported database type: {db_type}")
356
391
  except ImportError as e:
357
392
  # 捕获导入错误并转换为ConfigurationError,以保持与现有测试兼容
358
- raise ConfigurationError(f"Failed to import handler for {db_type}: {str(e)}")
393
+ raise ConfigurationError(
394
+ f"Failed to import handler for {db_type}: {str(e)}"
395
+ )
359
396
 
360
397
  @asynccontextmanager
361
- async def get_handler(self, connection: str) -> AsyncContextManager[ConnectionHandler]:
398
+ async def get_handler(
399
+ self, connection: str
400
+ ) -> AsyncContextManager[ConnectionHandler]:
362
401
  """Get connection handler
363
402
 
364
403
  Get appropriate connection handler based on connection name
@@ -371,36 +410,53 @@ class ConnectionServer:
371
410
  """
372
411
  # Read configuration file and validate connection
373
412
  db_config = self._get_config_or_raise(connection)
374
-
413
+
375
414
  # Create appropriate handler based on database type
376
415
  handler = None
377
416
  try:
378
- db_type = db_config['type']
417
+ db_type = db_config["type"]
379
418
  handler = self._create_handler_for_type(db_type, connection)
380
419
 
381
420
  # Set session for MCP logging
382
- if hasattr(self.server, 'session'):
421
+ if hasattr(self.server, "session"):
383
422
  handler._session = self.server.session
384
423
 
385
424
  handler.stats.record_connection_start()
386
- self.send_log(LOG_LEVEL_DEBUG, f"Handler created successfully for {connection}")
387
-
425
+ self.send_log(
426
+ LOG_LEVEL_DEBUG, f"Handler created successfully for {connection}"
427
+ )
428
+
388
429
  yield handler
389
430
  finally:
390
431
  if handler:
391
432
  self.send_log(LOG_LEVEL_DEBUG, f"Cleaning up handler for {connection}")
392
433
  handler.stats.record_connection_end()
393
-
394
- if hasattr(handler, 'cleanup') and callable(handler.cleanup):
434
+
435
+ if hasattr(handler, "cleanup") and callable(handler.cleanup):
395
436
  await handler.cleanup()
396
437
 
397
438
  def _get_available_tools(self) -> list[types.Tool]:
398
439
  """返回所有可用的数据库工具列表
399
-
440
+
400
441
  Returns:
401
442
  list[types.Tool]: 工具列表
402
443
  """
403
444
  return [
445
+ types.Tool(
446
+ name="dbutils-list-connections",
447
+ description="List all available database connections defined in the configuration",
448
+ inputSchema={
449
+ "type": "object",
450
+ "properties": {
451
+ "check_status": {
452
+ "type": "boolean",
453
+ "description": "Whether to check connection status (may be slow with many connections)",
454
+ "default": False,
455
+ }
456
+ },
457
+ "required": [],
458
+ },
459
+ ),
404
460
  types.Tool(
405
461
  name="dbutils-run-query",
406
462
  description="Execute read-only SQL query on database connection",
@@ -409,15 +465,15 @@ class ConnectionServer:
409
465
  "properties": {
410
466
  "connection": {
411
467
  "type": "string",
412
- "description": DATABASE_CONNECTION_NAME
468
+ "description": DATABASE_CONNECTION_NAME,
413
469
  },
414
470
  "sql": {
415
471
  "type": "string",
416
- "description": "SQL query (SELECT only)"
417
- }
472
+ "description": "SQL query (SELECT only)",
473
+ },
418
474
  },
419
- "required": ["connection", "sql"]
420
- }
475
+ "required": ["connection", "sql"],
476
+ },
421
477
  ),
422
478
  types.Tool(
423
479
  name="dbutils-list-tables",
@@ -427,11 +483,11 @@ class ConnectionServer:
427
483
  "properties": {
428
484
  "connection": {
429
485
  "type": "string",
430
- "description": DATABASE_CONNECTION_NAME
486
+ "description": DATABASE_CONNECTION_NAME,
431
487
  }
432
488
  },
433
- "required": ["connection"]
434
- }
489
+ "required": ["connection"],
490
+ },
435
491
  ),
436
492
  types.Tool(
437
493
  name="dbutils-describe-table",
@@ -441,15 +497,15 @@ class ConnectionServer:
441
497
  "properties": {
442
498
  "connection": {
443
499
  "type": "string",
444
- "description": DATABASE_CONNECTION_NAME
500
+ "description": DATABASE_CONNECTION_NAME,
445
501
  },
446
502
  "table": {
447
503
  "type": "string",
448
- "description": "Table name to describe"
449
- }
504
+ "description": "Table name to describe",
505
+ },
450
506
  },
451
- "required": ["connection", "table"]
452
- }
507
+ "required": ["connection", "table"],
508
+ },
453
509
  ),
454
510
  types.Tool(
455
511
  name="dbutils-get-ddl",
@@ -459,15 +515,15 @@ class ConnectionServer:
459
515
  "properties": {
460
516
  "connection": {
461
517
  "type": "string",
462
- "description": DATABASE_CONNECTION_NAME
518
+ "description": DATABASE_CONNECTION_NAME,
463
519
  },
464
520
  "table": {
465
521
  "type": "string",
466
- "description": "Table name to get DDL for"
467
- }
522
+ "description": "Table name to get DDL for",
523
+ },
468
524
  },
469
- "required": ["connection", "table"]
470
- }
525
+ "required": ["connection", "table"],
526
+ },
471
527
  ),
472
528
  types.Tool(
473
529
  name="dbutils-list-indexes",
@@ -477,15 +533,15 @@ class ConnectionServer:
477
533
  "properties": {
478
534
  "connection": {
479
535
  "type": "string",
480
- "description": DATABASE_CONNECTION_NAME
536
+ "description": DATABASE_CONNECTION_NAME,
481
537
  },
482
538
  "table": {
483
539
  "type": "string",
484
- "description": "Table name to list indexes for"
485
- }
540
+ "description": "Table name to list indexes for",
541
+ },
486
542
  },
487
- "required": ["connection", "table"]
488
- }
543
+ "required": ["connection", "table"],
544
+ },
489
545
  ),
490
546
  types.Tool(
491
547
  name="dbutils-get-stats",
@@ -495,15 +551,15 @@ class ConnectionServer:
495
551
  "properties": {
496
552
  "connection": {
497
553
  "type": "string",
498
- "description": DATABASE_CONNECTION_NAME
554
+ "description": DATABASE_CONNECTION_NAME,
499
555
  },
500
556
  "table": {
501
557
  "type": "string",
502
- "description": "Table name to get statistics for"
503
- }
558
+ "description": "Table name to get statistics for",
559
+ },
504
560
  },
505
- "required": ["connection", "table"]
506
- }
561
+ "required": ["connection", "table"],
562
+ },
507
563
  ),
508
564
  types.Tool(
509
565
  name="dbutils-list-constraints",
@@ -513,15 +569,15 @@ class ConnectionServer:
513
569
  "properties": {
514
570
  "connection": {
515
571
  "type": "string",
516
- "description": DATABASE_CONNECTION_NAME
572
+ "description": DATABASE_CONNECTION_NAME,
517
573
  },
518
574
  "table": {
519
575
  "type": "string",
520
- "description": "Table name to list constraints for"
521
- }
576
+ "description": "Table name to list constraints for",
577
+ },
522
578
  },
523
- "required": ["connection", "table"]
524
- }
579
+ "required": ["connection", "table"],
580
+ },
525
581
  ),
526
582
  types.Tool(
527
583
  name="dbutils-explain-query",
@@ -531,15 +587,15 @@ class ConnectionServer:
531
587
  "properties": {
532
588
  "connection": {
533
589
  "type": "string",
534
- "description": DATABASE_CONNECTION_NAME
590
+ "description": DATABASE_CONNECTION_NAME,
535
591
  },
536
592
  "sql": {
537
593
  "type": "string",
538
- "description": "SQL query to explain"
539
- }
594
+ "description": "SQL query to explain",
595
+ },
540
596
  },
541
- "required": ["connection", "sql"]
542
- }
597
+ "required": ["connection", "sql"],
598
+ },
543
599
  ),
544
600
  types.Tool(
545
601
  name="dbutils-get-performance",
@@ -549,11 +605,11 @@ class ConnectionServer:
549
605
  "properties": {
550
606
  "connection": {
551
607
  "type": "string",
552
- "description": DATABASE_CONNECTION_NAME
608
+ "description": DATABASE_CONNECTION_NAME,
553
609
  }
554
610
  },
555
- "required": ["connection"]
556
- }
611
+ "required": ["connection"],
612
+ },
557
613
  ),
558
614
  types.Tool(
559
615
  name="dbutils-analyze-query",
@@ -563,24 +619,108 @@ class ConnectionServer:
563
619
  "properties": {
564
620
  "connection": {
565
621
  "type": "string",
566
- "description": DATABASE_CONNECTION_NAME
622
+ "description": DATABASE_CONNECTION_NAME,
567
623
  },
568
624
  "sql": {
569
625
  "type": "string",
570
- "description": "SQL query to analyze"
571
- }
626
+ "description": "SQL query to analyze",
627
+ },
572
628
  },
573
- "required": ["connection", "sql"]
574
- }
575
- )
629
+ "required": ["connection", "sql"],
630
+ },
631
+ ),
576
632
  ]
577
-
633
+
634
+ async def _handle_list_connections(
635
+ self, check_status: bool = False
636
+ ) -> list[types.TextContent]:
637
+ """处理列出数据库连接工具调用
638
+
639
+ Args:
640
+ check_status: 是否检查连接状态
641
+
642
+ Returns:
643
+ list[types.TextContent]: 数据库连接列表
644
+ """
645
+ connections = []
646
+
647
+ try:
648
+ # 读取配置文件
649
+ with open(self.config_path, "r") as f:
650
+ config = yaml.safe_load(f)
651
+ if not config or "connections" not in config:
652
+ return [
653
+ types.TextContent(
654
+ type="text",
655
+ text="No database connections found in configuration.",
656
+ )
657
+ ]
658
+
659
+ # 获取配置中的所有连接
660
+ for conn_name, conn_config in config["connections"].items():
661
+ db_type = conn_config.get("type", "unknown")
662
+ connection_info = []
663
+
664
+ # 添加基本信息
665
+ connection_info.append(f"Connection: {conn_name}")
666
+ connection_info.append(f"Type: {db_type}")
667
+
668
+ # 根据数据库类型添加特定信息(排除敏感信息)
669
+ if db_type == "sqlite":
670
+ if "path" in conn_config:
671
+ connection_info.append(f"Path: {conn_config['path']}")
672
+ elif "database" in conn_config:
673
+ connection_info.append(
674
+ f"Database: {conn_config['database']}"
675
+ )
676
+ elif db_type in ["mysql", "postgres", "postgresql"]:
677
+ if "host" in conn_config:
678
+ connection_info.append(f"Host: {conn_config['host']}")
679
+ if "port" in conn_config:
680
+ connection_info.append(f"Port: {conn_config['port']}")
681
+ if "database" in conn_config:
682
+ connection_info.append(
683
+ f"Database: {conn_config['database']}"
684
+ )
685
+ if "user" in conn_config:
686
+ connection_info.append(f"User: {conn_config['user']}")
687
+ # 不显示密码
688
+
689
+ # 检查连接状态(如果需要)
690
+ if check_status:
691
+ try:
692
+ async with self.get_handler(conn_name) as handler:
693
+ # 尝试执行一个简单查询来验证连接
694
+ await handler.test_connection()
695
+ connection_info.append("Status: Available")
696
+ except Exception as e:
697
+ connection_info.append(f"Status: Unavailable ({str(e)})")
698
+
699
+ connections.append("\n".join(connection_info))
700
+ except Exception as e:
701
+ self.send_log(LOG_LEVEL_ERROR, f"Error listing connections: {str(e)}")
702
+ return [
703
+ types.TextContent(
704
+ type="text", text=f"Error listing connections: {str(e)}"
705
+ )
706
+ ]
707
+
708
+ if not connections:
709
+ return [
710
+ types.TextContent(
711
+ type="text", text="No database connections found in configuration."
712
+ )
713
+ ]
714
+
715
+ result = "Available database connections:\n\n" + "\n\n".join(connections)
716
+ return [types.TextContent(type="text", text=result)]
717
+
578
718
  async def _handle_list_tables(self, connection: str) -> list[types.TextContent]:
579
719
  """处理列表表格工具调用
580
-
720
+
581
721
  Args:
582
722
  connection: 数据库连接名称
583
-
723
+
584
724
  Returns:
585
725
  list[types.TextContent]: 表格列表
586
726
  """
@@ -588,28 +728,44 @@ class ConnectionServer:
588
728
  tables = await handler.get_tables()
589
729
  if not tables:
590
730
  # 空表列表的情况也返回数据库类型
591
- return [types.TextContent(type="text", text=f"[{handler.db_type}] No tables found")]
592
-
593
- formatted_tables = "\n".join([
594
- f"Table: {table.name}\n" +
595
- f"URI: {table.uri}\n" +
596
- (f"Description: {table.description}\n" if table.description else "") +
597
- "---"
598
- for table in tables
599
- ])
731
+ return [
732
+ types.TextContent(
733
+ type="text", text=f"[{handler.db_type}] No tables found"
734
+ )
735
+ ]
736
+
737
+ formatted_tables = "\n".join(
738
+ [
739
+ f"Table: {table.name}\n"
740
+ + f"URI: {table.uri}\n"
741
+ + (
742
+ f"Description: {table.description}\n"
743
+ if table.description
744
+ else ""
745
+ )
746
+ + "---"
747
+ for table in tables
748
+ ]
749
+ )
600
750
  # 添加数据库类型前缀
601
- return [types.TextContent(type="text", text=f"[{handler.db_type}]\n{formatted_tables}")]
602
-
603
- async def _handle_run_query(self, connection: str, sql: str) -> list[types.TextContent]:
751
+ return [
752
+ types.TextContent(
753
+ type="text", text=f"[{handler.db_type}]\n{formatted_tables}"
754
+ )
755
+ ]
756
+
757
+ async def _handle_run_query(
758
+ self, connection: str, sql: str
759
+ ) -> list[types.TextContent]:
604
760
  """处理运行查询工具调用
605
-
761
+
606
762
  Args:
607
763
  connection: 数据库连接名称
608
764
  sql: SQL查询语句
609
-
765
+
610
766
  Returns:
611
767
  list[types.TextContent]: 查询结果
612
-
768
+
613
769
  Raises:
614
770
  ConfigurationError: 如果SQL为空或非SELECT语句
615
771
  """
@@ -623,81 +779,91 @@ class ConnectionServer:
623
779
  async with self.get_handler(connection) as handler:
624
780
  result = await handler.execute_query(sql)
625
781
  return [types.TextContent(type="text", text=result)]
626
-
627
- async def _handle_table_tools(self, name: str, connection: str, table: str) -> list[types.TextContent]:
782
+
783
+ async def _handle_table_tools(
784
+ self, name: str, connection: str, table: str
785
+ ) -> list[types.TextContent]:
628
786
  """处理表相关工具调用
629
-
787
+
630
788
  Args:
631
789
  name: 工具名称
632
790
  connection: 数据库连接名称
633
791
  table: 表名
634
-
792
+
635
793
  Returns:
636
794
  list[types.TextContent]: 工具执行结果
637
-
795
+
638
796
  Raises:
639
797
  ConfigurationError: 如果表名为空
640
798
  """
641
799
  if not table:
642
800
  raise ConfigurationError(EMPTY_TABLE_NAME_ERROR)
643
-
801
+
644
802
  async with self.get_handler(connection) as handler:
645
803
  result = await handler.execute_tool_query(name, table_name=table)
646
804
  return [types.TextContent(type="text", text=result)]
647
-
648
- async def _handle_explain_query(self, connection: str, sql: str) -> list[types.TextContent]:
805
+
806
+ async def _handle_explain_query(
807
+ self, connection: str, sql: str
808
+ ) -> list[types.TextContent]:
649
809
  """处理解释查询工具调用
650
-
810
+
651
811
  Args:
652
812
  connection: 数据库连接名称
653
813
  sql: SQL查询语句
654
-
814
+
655
815
  Returns:
656
816
  list[types.TextContent]: 查询解释
657
-
817
+
658
818
  Raises:
659
819
  ConfigurationError: 如果SQL为空
660
820
  """
661
821
  if not sql:
662
822
  raise ConfigurationError(EMPTY_QUERY_ERROR)
663
-
823
+
664
824
  async with self.get_handler(connection) as handler:
665
825
  result = await handler.execute_tool_query("dbutils-explain-query", sql=sql)
666
826
  return [types.TextContent(type="text", text=result)]
667
-
827
+
668
828
  async def _handle_performance(self, connection: str) -> list[types.TextContent]:
669
829
  """处理性能统计工具调用
670
-
830
+
671
831
  Args:
672
832
  connection: 数据库连接名称
673
-
833
+
674
834
  Returns:
675
835
  list[types.TextContent]: 性能统计
676
836
  """
677
837
  async with self.get_handler(connection) as handler:
678
838
  performance_stats = handler.stats.get_performance_stats()
679
- return [types.TextContent(type="text", text=f"[{handler.db_type}]\n{performance_stats}")]
680
-
681
- async def _handle_analyze_query(self, connection: str, sql: str) -> list[types.TextContent]:
839
+ return [
840
+ types.TextContent(
841
+ type="text", text=f"[{handler.db_type}]\n{performance_stats}"
842
+ )
843
+ ]
844
+
845
+ async def _handle_analyze_query(
846
+ self, connection: str, sql: str
847
+ ) -> list[types.TextContent]:
682
848
  """处理查询分析工具调用
683
-
849
+
684
850
  Args:
685
851
  connection: 数据库连接名称
686
852
  sql: SQL查询语句
687
-
853
+
688
854
  Returns:
689
855
  list[types.TextContent]: 查询分析结果
690
-
856
+
691
857
  Raises:
692
858
  ConfigurationError: 如果SQL为空
693
859
  """
694
860
  if not sql:
695
861
  raise ConfigurationError(EMPTY_QUERY_ERROR)
696
-
862
+
697
863
  async with self.get_handler(connection) as handler:
698
864
  # First get the execution plan
699
865
  explain_result = await handler.explain_query(sql)
700
-
866
+
701
867
  # Then execute the actual query to measure performance
702
868
  start_time = datetime.now()
703
869
  if sql.lower().startswith("select"):
@@ -705,35 +871,40 @@ class ConnectionServer:
705
871
  await handler.execute_query(sql)
706
872
  except Exception as e:
707
873
  # If query fails, we still provide the execution plan
708
- self.send_log(LOG_LEVEL_ERROR, f"Query execution failed during analysis: {str(e)}")
874
+ self.send_log(
875
+ LOG_LEVEL_ERROR,
876
+ f"Query execution failed during analysis: {str(e)}",
877
+ )
709
878
  duration = (datetime.now() - start_time).total_seconds()
710
-
879
+
711
880
  # Combine analysis results
712
881
  analysis = [
713
882
  f"[{handler.db_type}] Query Analysis",
714
883
  f"SQL: {sql}",
715
884
  "",
716
- f"Execution Time: {duration*1000:.2f}ms",
885
+ f"Execution Time: {duration * 1000:.2f}ms",
717
886
  "",
718
887
  "Execution Plan:",
719
- explain_result
888
+ explain_result,
720
889
  ]
721
-
890
+
722
891
  # Add optimization suggestions
723
892
  suggestions = self._get_optimization_suggestions(explain_result, duration)
724
893
  if suggestions:
725
894
  analysis.append("\nOptimization Suggestions:")
726
895
  analysis.extend(suggestions)
727
-
896
+
728
897
  return [types.TextContent(type="text", text="\n".join(analysis))]
729
-
730
- def _get_optimization_suggestions(self, explain_result: str, duration: float) -> list[str]:
898
+
899
+ def _get_optimization_suggestions(
900
+ self, explain_result: str, duration: float
901
+ ) -> list[str]:
731
902
  """根据执行计划和耗时获取优化建议
732
-
903
+
733
904
  Args:
734
905
  explain_result: 执行计划
735
906
  duration: 查询耗时(秒)
736
-
907
+
737
908
  Returns:
738
909
  list[str]: 优化建议列表
739
910
  """
@@ -745,32 +916,37 @@ class ConnectionServer:
745
916
  if duration > 0.5: # 500ms
746
917
  suggestions.append("- Query is slow, consider optimizing or adding caching")
747
918
  if "temporary" in explain_result.lower():
748
- suggestions.append("- Query creates temporary tables, consider restructuring")
749
-
919
+ suggestions.append(
920
+ "- Query creates temporary tables, consider restructuring"
921
+ )
922
+
750
923
  return suggestions
751
924
 
752
925
  def _setup_handlers(self):
753
926
  """Setup MCP handlers"""
927
+
754
928
  @self.server.list_resources()
755
- async def handle_list_resources(arguments: dict | None = None) -> list[types.Resource]:
756
- if not arguments or 'connection' not in arguments:
929
+ async def handle_list_resources(
930
+ arguments: dict | None = None,
931
+ ) -> list[types.Resource]:
932
+ if not arguments or "connection" not in arguments:
757
933
  # Return empty list when no connection specified
758
934
  return []
759
935
 
760
- connection = arguments['connection']
936
+ connection = arguments["connection"]
761
937
  async with self.get_handler(connection) as handler:
762
938
  return await handler.get_tables()
763
939
 
764
940
  @self.server.read_resource()
765
941
  async def handle_read_resource(uri: str, arguments: dict | None = None) -> str:
766
- if not arguments or 'connection' not in arguments:
942
+ if not arguments or "connection" not in arguments:
767
943
  raise ConfigurationError(CONNECTION_NAME_REQUIRED_ERROR)
768
944
 
769
- parts = uri.split('/')
945
+ parts = uri.split("/")
770
946
  if len(parts) < 3:
771
947
  raise ConfigurationError(INVALID_URI_FORMAT_ERROR)
772
948
 
773
- connection = arguments['connection']
949
+ connection = arguments["connection"]
774
950
  table_name = parts[-2] # URI format: xxx/table_name/schema
775
951
 
776
952
  async with self.get_handler(connection) as handler:
@@ -781,10 +957,17 @@ class ConnectionServer:
781
957
  return self._get_available_tools()
782
958
 
783
959
  @self.server.call_tool()
784
- async def handle_call_tool(name: str, arguments: dict) -> list[types.TextContent]:
960
+ async def handle_call_tool(
961
+ name: str, arguments: dict
962
+ ) -> list[types.TextContent]:
963
+ # Special case for list-connections which doesn't require a connection
964
+ if name == "dbutils-list-connections":
965
+ check_status = arguments.get("check_status", False)
966
+ return await self._handle_list_connections(check_status)
967
+
785
968
  if "connection" not in arguments:
786
969
  raise ConfigurationError(CONNECTION_NAME_REQUIRED_ERROR)
787
-
970
+
788
971
  connection = arguments["connection"]
789
972
 
790
973
  if name == "dbutils-list-tables":
@@ -792,8 +975,13 @@ class ConnectionServer:
792
975
  elif name == "dbutils-run-query":
793
976
  sql = arguments.get("sql", "").strip()
794
977
  return await self._handle_run_query(connection, sql)
795
- elif name in ["dbutils-describe-table", "dbutils-get-ddl", "dbutils-list-indexes",
796
- "dbutils-get-stats", "dbutils-list-constraints"]:
978
+ elif name in [
979
+ "dbutils-describe-table",
980
+ "dbutils-get-ddl",
981
+ "dbutils-list-indexes",
982
+ "dbutils-get-stats",
983
+ "dbutils-list-constraints",
984
+ ]:
797
985
  table = arguments.get("table", "").strip()
798
986
  return await self._handle_table_tools(name, connection, table)
799
987
  elif name == "dbutils-explain-query":
@@ -811,7 +999,5 @@ class ConnectionServer:
811
999
  """Run server"""
812
1000
  async with mcp.server.stdio.stdio_server() as streams:
813
1001
  await self.server.run(
814
- streams[0],
815
- streams[1],
816
- self.server.create_initialization_options()
1002
+ streams[0], streams[1], self.server.create_initialization_options()
817
1003
  )