# Directory Structure
```
├── __init__.py
├── .env.template
├── .gitignore
├── Dockerfile
├── LICENSE
├── pyproject.toml
├── README.md
├── requirements.txt
├── server.py
└── smithery.yaml
```
# Files
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
```
1 | # Sensitive files
2 | .env
3 | *.env
4 | !.env.template
5 |
6 | # Python
7 | __pycache__/
8 | *.py[cod]
9 | *$py.class
10 | venv/
11 | *.egg-info/
12 |
13 | # IDE
14 | .idea/
15 | .vscode/
16 |
17 | # Logs
18 | *.log
```
--------------------------------------------------------------------------------
/__init__.py:
--------------------------------------------------------------------------------
```python
1 | from .server import main
2 |
3 | __all__ = ["main"]
```
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
```
1 | snowflake-connector-python
2 | python-dotenv
3 | mcp
4 |
```
--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
```dockerfile
1 | # Generated by https://smithery.ai. See: https://smithery.ai/docs/config#dockerfile
2 | FROM python:3.10-alpine
3 |
4 | # Set working directory
5 | WORKDIR /app
6 |
7 | # Install system dependencies: add gcc, g++, musl-dev, linux-headers
8 | RUN apk add --no-cache gcc g++ musl-dev linux-headers
9 |
10 | # Copy project files
11 | COPY . /app
12 |
13 | # Install Python dependencies
14 | RUN pip install --no-cache-dir --upgrade pip \
15 | && pip install --no-cache-dir -r requirements.txt
16 |
17 | # Set default environment variables (can be overridden)
18 | ENV SNOWFLAKE_USER=your_username \
19 | SNOWFLAKE_PASSWORD=your_password \
20 | SNOWFLAKE_ACCOUNT=your_account \
21 | SNOWFLAKE_DATABASE=your_database \
22 | SNOWFLAKE_WAREHOUSE=your_warehouse
23 |
24 | # Command to run the MCP server
25 | CMD ["python", "server.py"]
26 |
```
--------------------------------------------------------------------------------
/smithery.yaml:
--------------------------------------------------------------------------------
```yaml
1 | # Smithery configuration file: https://smithery.ai/docs/config#smitheryyaml
2 |
3 | startCommand:
4 | type: stdio
5 | configSchema:
6 | # JSON Schema defining the configuration options for the MCP.
7 | type: object
8 | required:
9 | - snowflakeUser
10 | - snowflakePassword
11 | - snowflakeAccount
12 | - snowflakeDatabase
13 | - snowflakeWarehouse
14 | properties:
15 | snowflakeUser:
16 | type: string
17 | description: Snowflake username
18 | snowflakePassword:
19 | type: string
20 | description: Snowflake password
21 | snowflakeAccount:
22 | type: string
23 | description: Snowflake account identifier
24 | snowflakeDatabase:
25 | type: string
26 | description: Snowflake database name
27 | snowflakeWarehouse:
28 | type: string
29 | description: Snowflake warehouse name
30 | commandFunction:
31 | # A JS function that produces the CLI command based on the given config to start the MCP on stdio.
32 | |-
33 | (config) => ({
34 | command: 'python',
35 | args: ['server.py'],
36 | env: {
37 | SNOWFLAKE_USER: config.snowflakeUser,
38 | SNOWFLAKE_PASSWORD: config.snowflakePassword,
39 | SNOWFLAKE_ACCOUNT: config.snowflakeAccount,
40 | SNOWFLAKE_DATABASE: config.snowflakeDatabase,
41 | SNOWFLAKE_WAREHOUSE: config.snowflakeWarehouse
42 | }
43 | })
44 | exampleConfig:
45 | snowflakeUser: example_user
46 | snowflakePassword: example_password
47 | snowflakeAccount: EXAMPLE12345.US-WEST-2
48 | snowflakeDatabase: EXAMPLE_DB
49 | snowflakeWarehouse: EXAMPLE_WH
50 |
```
--------------------------------------------------------------------------------
/server.py:
--------------------------------------------------------------------------------
```python
1 | #!/usr/bin/env python
2 | import os
3 | import asyncio
4 | import logging
5 | import json
6 | import time
7 | import snowflake.connector
8 | from dotenv import load_dotenv
9 | import mcp.server.stdio
10 | from mcp.server import Server
11 | from mcp.types import Tool, ServerResult, TextContent
12 | from contextlib import closing
13 | from typing import Optional, Any
14 |
15 | # 配置日志 / Configure logging
16 | logging.basicConfig(
17 | level=logging.DEBUG,
18 | format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
19 | )
20 | logger = logging.getLogger('snowflake_server')
21 |
22 | load_dotenv()
23 |
24 | class SnowflakeConnection:
25 | """
26 | Snowflake数据库连接管理类
27 | Snowflake database connection management class
28 | """
29 | def __init__(self):
30 | # 初始化配置信息 / Initialize configuration
31 | self.config = {
32 | "user": os.getenv("SNOWFLAKE_USER"),
33 | "password": os.getenv("SNOWFLAKE_PASSWORD"),
34 | "account": os.getenv("SNOWFLAKE_ACCOUNT"),
35 | "database": os.getenv("SNOWFLAKE_DATABASE"),
36 | "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),
37 | }
38 | self.conn: Optional[snowflake.connector.SnowflakeConnection] = None
39 | logger.info(f"Initialized with config (excluding password): {json.dumps({k:v for k,v in self.config.items() if k != 'password'})}")
40 |
41 | def ensure_connection(self) -> snowflake.connector.SnowflakeConnection:
42 | """
43 | 确保数据库连接可用,如果连接不存在或已断开则重新建立连接
44 | Ensure database connection is available, create new connection if it doesn't exist or is disconnected
45 | """
46 | try:
47 | # 检查连接是否需要重新建立 / Check if connection needs to be re-established
48 | if self.conn is None:
49 | logger.info("Creating new Snowflake connection...")
50 | self.conn = snowflake.connector.connect(
51 | **self.config,
52 | client_session_keep_alive=True,
53 | network_timeout=15,
54 | login_timeout=15
55 | )
56 | self.conn.cursor().execute("ALTER SESSION SET TIMEZONE = 'UTC'")
57 | logger.info("New connection established and configured")
58 |
59 | # 测试连接是否有效 / Test if connection is valid
60 | try:
61 | self.conn.cursor().execute("SELECT 1")
62 | except:
63 | logger.info("Connection lost, reconnecting...")
64 | self.conn = None
65 | return self.ensure_connection()
66 |
67 | return self.conn
68 | except Exception as e:
69 | logger.error(f"Connection error: {str(e)}")
70 | raise
71 |
72 | def execute_query(self, query: str) -> list[dict[str, Any]]:
73 | """
74 | 执行SQL查询并返回结果
75 | Execute SQL query and return results
76 |
77 | Args:
78 | query (str): SQL查询语句 / SQL query statement
79 |
80 | Returns:
81 | list[dict[str, Any]]: 查询结果列表 / List of query results
82 | """
83 | start_time = time.time()
84 | logger.info(f"Executing query: {query[:200]}...") # 只记录前200个字符 / Log only first 200 characters
85 |
86 | try:
87 | conn = self.ensure_connection()
88 | with conn.cursor() as cursor:
89 | # 对于写操作使用事务 / Use transaction for write operations
90 | if any(query.strip().upper().startswith(word) for word in ['INSERT', 'UPDATE', 'DELETE', 'CREATE', 'DROP', 'ALTER']):
91 | cursor.execute("BEGIN")
92 | try:
93 | cursor.execute(query)
94 | conn.commit()
95 | logger.info(f"Write query executed in {time.time() - start_time:.2f}s")
96 | return [{"affected_rows": cursor.rowcount}]
97 | except Exception as e:
98 | conn.rollback()
99 | raise
100 | else:
101 | # 读操作 / Read operations
102 | cursor.execute(query)
103 | if cursor.description:
104 | columns = [col[0] for col in cursor.description]
105 | rows = cursor.fetchall()
106 | results = [dict(zip(columns, row)) for row in rows]
107 | logger.info(f"Read query returned {len(results)} rows in {time.time() - start_time:.2f}s")
108 | return results
109 | return []
110 |
111 | except snowflake.connector.errors.ProgrammingError as e:
112 | logger.error(f"SQL Error: {str(e)}")
113 | logger.error(f"Error Code: {getattr(e, 'errno', 'unknown')}")
114 | raise
115 | except Exception as e:
116 | logger.error(f"Query error: {str(e)}")
117 | logger.error(f"Error type: {type(e).__name__}")
118 | raise
119 |
120 | def close(self):
121 | """
122 | 关闭数据库连接
123 | Close database connection
124 | """
125 | if self.conn:
126 | try:
127 | self.conn.close()
128 | logger.info("Connection closed")
129 | except Exception as e:
130 | logger.error(f"Error closing connection: {str(e)}")
131 | finally:
132 | self.conn = None
133 |
134 | class SnowflakeServer(Server):
135 | """
136 | Snowflake MCP服务器类,处理与客户端的交互
137 | Snowflake MCP server class, handles client interactions
138 | """
139 | def __init__(self):
140 | super().__init__(name="snowflake-server")
141 | self.db = SnowflakeConnection()
142 | logger.info("SnowflakeServer initialized")
143 |
144 | @self.list_tools()
145 | async def handle_tools():
146 | """
147 | 返回可用工具列表
148 | Return list of available tools
149 | """
150 | return [
151 | Tool(
152 | name="execute_query",
153 | description="Execute a SQL query on Snowflake",
154 | inputSchema={
155 | "type": "object",
156 | "properties": {
157 | "query": {
158 | "type": "string",
159 | "description": "SQL query to execute"
160 | }
161 | },
162 | "required": ["query"]
163 | }
164 | )
165 | ]
166 |
167 | @self.call_tool()
168 | async def handle_call_tool(name: str, arguments: dict):
169 | """
170 | 处理工具调用请求
171 | Handle tool call requests
172 |
173 | Args:
174 | name (str): 工具名称 / Tool name
175 | arguments (dict): 工具参数 / Tool arguments
176 |
177 | Returns:
178 | list[TextContent]: 执行结果 / Execution results
179 | """
180 | if name == "execute_query":
181 | start_time = time.time()
182 | try:
183 | result = self.db.execute_query(arguments["query"])
184 | execution_time = time.time() - start_time
185 |
186 | return [TextContent(
187 | type="text",
188 | text=f"Results (execution time: {execution_time:.2f}s):\n{result}"
189 | )]
190 | except Exception as e:
191 | error_message = f"Error executing query: {str(e)}"
192 | logger.error(error_message)
193 | return [TextContent(
194 | type="text",
195 | text=error_message
196 | )]
197 |
198 | def __del__(self):
199 | """
200 | 清理资源,关闭数据库连接
201 | Clean up resources, close database connection
202 | """
203 | if hasattr(self, 'db'):
204 | self.db.close()
205 |
206 | async def main():
207 | """
208 | 主函数,启动服务器并处理请求
209 | Main function, starts server and handles requests
210 | """
211 | try:
212 | server = SnowflakeServer()
213 | initialization_options = server.create_initialization_options()
214 | logger.info("Starting server")
215 |
216 | async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
217 | await server.run(
218 | read_stream,
219 | write_stream,
220 | initialization_options
221 | )
222 | except Exception as e:
223 | logger.critical(f"Server failed: {str(e)}", exc_info=True)
224 | raise
225 | finally:
226 | logger.info("Server shutting down")
227 |
228 | if __name__ == "__main__":
229 | asyncio.run(main())
```