pggm-mcp-snowflake-server 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

pggm_mcp_snowflake_server/__init__.py

@@ -1,3 +1,111 @@
-from .server import main
+import argparse
+import asyncio
+import os
 
-__all__ = ["main"]
+import dotenv
+import snowflake.connector
+
+from . import server
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Add arguments
+    parser.add_argument(
+        "--allow_write", required=False, default=False, action="store_true", help="Allow write operations on the database"
+    )
+    parser.add_argument("--log_dir", required=False, default=None, help="Directory to log to")
+    parser.add_argument("--log_level", required=False, default="INFO", help="Logging level")
+    parser.add_argument(
+        "--prefetch",
+        action="store_true",
+        dest="prefetch",
+        default=True,
+        help="Prefetch table descriptions (when enabled, list_tables and describe_table are disabled)",
+    )
+    parser.add_argument(
+        "--no-prefetch",
+        action="store_false",
+        dest="prefetch",
+        help="Don't prefetch table descriptions",
+    )
+    parser.add_argument(
+        "--exclude_tools",
+        required=False,
+        default=[],
+        nargs="+",
+        help="List of tools to exclude",
+    )
+
+    # First, get all the arguments we don't know about
+    args, unknown = parser.parse_known_args()
+
+    # Create a dictionary to store our key-value pairs
+    connection_args = {}
+
+    # Iterate through unknown args in pairs
+    for i in range(0, len(unknown), 2):
+        if i + 1 >= len(unknown):
+            break
+
+        key = unknown[i]
+        value = unknown[i + 1]
+
+        # Make sure it's a keyword argument (starts with --)
+        if key.startswith("--"):
+            key = key[2:]  # Remove the '--'
+            connection_args[key] = value
+
+    # Now we can add the known args to kwargs
+    server_args = {
+        "allow_write": args.allow_write,
+        "log_dir": args.log_dir,
+        "log_level": args.log_level,
+        "prefetch": args.prefetch,
+        "exclude_tools": args.exclude_tools,
+    }
+
+    return server_args, connection_args
+
+
+def main():
+    """Main entry point for the package."""
+
+    dotenv.load_dotenv()
+
+    default_connection_args = snowflake.connector.connection.DEFAULT_CONFIGURATION
+
+    connection_args_from_env = {
+        k: os.getenv("SNOWFLAKE_" + k.upper())
+        for k in default_connection_args
+        if os.getenv("SNOWFLAKE_" + k.upper()) is not None
+    }
+
+    server_args, connection_args = parse_args()
+
+    connection_args = {**connection_args_from_env, **connection_args}
+
+    assert (
+        "database" in connection_args
+    ), 'You must provide the account identifier as "--database" argument or "SNOWFLAKE_DATABASE" environment variable.'
+    assert (
+        "schema" in connection_args
+    ), 'You must provide the username as "--schema" argument or "SNOWFLAKE_SCHEMA" environment variable.'
+
+    asyncio.run(
+        server.main(
+            connection_args=connection_args,
+            allow_write=server_args["allow_write"],
+            log_dir=server_args["log_dir"],
+            log_level=server_args["log_level"],
+            exclude_tools=server_args["exclude_tools"],
+        )
+    )
+
+
+# Optionally expose other important items at package level
+__all__ = ["main", "server", "write_detector"]
+
+if __name__ == "__main__":
+    main()

pggm_mcp_snowflake_server/server.py

@@ -23,7 +23,7 @@ logging.basicConfig(
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
     handlers=[logging.StreamHandler()],
 )
-logger = logging.getLogger("pggm_mcp_snowflake_server")
+logger = logging.getLogger("mcp_snowflake_server")
 
 
 def data_to_yaml(data: Any) -> str:
@@ -351,7 +351,7 @@ async def main(
     if log_dir:
         os.makedirs(log_dir, exist_ok=True)
         logger.handlers.append(
-            logging.FileHandler(os.path.join(log_dir, "pggm_mcp_snowflake_server.log"))
+            logging.FileHandler(os.path.join(log_dir, "mcp_snowflake_server.log"))
         )
     if log_level:
         logger.setLevel(log_level)
@@ -614,10 +614,10 @@ async def main(
             write_stream,
             InitializationOptions(
                 server_name="snowflake",
-                server_version=importlib.metadata.version("pggm_mcp_snowflake_server"),
+                server_version=importlib.metadata.version("mcp_snowflake_server"),
                 capabilities=server.get_capabilities(
                     notification_options=NotificationOptions(),
                     experimental_capabilities={},
                 ),
             ),
-        )
+        )

pggm_mcp_snowflake_server/write_detector.py

@@ -1,36 +1,98 @@
-class SQLWriteDetector:
-    """Utility class to detect write operations in SQL queries."""
+import sqlparse
+from sqlparse.sql import Token, TokenList
+from sqlparse.tokens import Keyword, DML, DDL
+from typing import Dict, List, Set, Tuple
+
 
+class SQLWriteDetector:
     def __init__(self):
-        self.write_keywords = [
-            "INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER", "TRUNCATE",
-            "GRANT", "REVOKE", "MERGE", "UPSERT", "REPLACE"
-        ]
-    
-    def analyze_query(self, query: str) -> dict:
+        # Define sets of keywords that indicate write operations
+        self.dml_write_keywords = {"INSERT", "UPDATE", "DELETE", "MERGE", "UPSERT", "REPLACE"}
+
+        self.ddl_keywords = {"CREATE", "ALTER", "DROP", "TRUNCATE", "RENAME"}
+
+        self.dcl_keywords = {"GRANT", "REVOKE"}
+
+        # Combine all write keywords
+        self.write_keywords = self.dml_write_keywords | self.ddl_keywords | self.dcl_keywords
+
+    def analyze_query(self, sql_query: str) -> Dict:
         """
-        Analyze an SQL query to determine if it contains write operations.
-        
+        Analyze a SQL query to determine if it contains write operations.
+
         Args:
-            query: SQL query string to analyze
-            
+            sql_query: The SQL query string to analyze
+
         Returns:
-            Dictionary with analysis results
+            Dictionary containing analysis results
         """
-        result = {
-            "contains_write": False,
-            "write_operations": [],
-            "query_type": "READ"
+        # Parse the SQL query
+        parsed = sqlparse.parse(sql_query)
+        if not parsed:
+            return {"contains_write": False, "write_operations": set(), "has_cte_write": False}
+
+        # Initialize result tracking
+        found_operations = set()
+        has_cte_write = False
+
+        # Analyze each statement in the query
+        for statement in parsed:
+            # Check for write operations in CTEs (WITH clauses)
+            if self._has_cte(statement):
+                cte_write = self._analyze_cte(statement)
+                if cte_write:
+                    has_cte_write = True
+                    found_operations.add("CTE_WRITE")
+
+            # Analyze the main query
+            operations = self._find_write_operations(statement)
+            found_operations.update(operations)
+
+        return {
+            "contains_write": bool(found_operations) or has_cte_write,
+            "write_operations": found_operations,
+            "has_cte_write": has_cte_write,
         }
-        
-        # Convert to uppercase for case-insensitive comparison
-        upper_query = query.upper()
-        
-        # Check for write keywords
-        for keyword in self.write_keywords:
-            if keyword in upper_query.split():
-                result["contains_write"] = True
-                result["write_operations"].append(keyword)
-                result["query_type"] = "WRITE"
-                
-        return result
+
+    def _has_cte(self, statement: TokenList) -> bool:
+        """Check if the statement has a WITH clause."""
+        return any(token.is_keyword and token.normalized == "WITH" for token in statement.tokens)
+
+    def _analyze_cte(self, statement: TokenList) -> bool:
+        """
+        Analyze CTEs (WITH clauses) for write operations.
+        Returns True if any CTE contains a write operation.
+        """
+        in_cte = False
+        for token in statement.tokens:
+            if token.is_keyword and token.normalized == "WITH":
+                in_cte = True
+            elif in_cte:
+                if any(write_kw in token.normalized for write_kw in self.write_keywords):
+                    return True
+        return False
+
+    def _find_write_operations(self, statement: TokenList) -> Set[str]:
+        """
+        Find all write operations in a statement.
+        Returns a set of found write operation keywords.
+        """
+        operations = set()
+
+        for token in statement.tokens:
+            # Skip comments and whitespace
+            if token.is_whitespace or token.ttype in (sqlparse.tokens.Comment,):
+                continue
+
+            # Check if token is a keyword
+            if token.ttype in (Keyword, DML, DDL):
+                normalized = token.normalized.upper()
+                if normalized in self.write_keywords:
+                    operations.add(normalized)
+
+            # Recursively check child tokens
+            if isinstance(token, TokenList):
+                child_ops = self._find_write_operations(token)
+                operations.update(child_ops)
+
+        return operations

pggm_mcp_snowflake_server-0.1.1.dist-info/METADATA

@@ -0,0 +1,67 @@
+Metadata-Version: 2.4
+Name: pggm-mcp-snowflake-server
+Version: 0.1.1
+Summary: Custom Model Context Protocol server for Snowflake
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: mcp
+Requires-Dist: pydantic
+Requires-Dist: snowflake-snowpark-python
+Requires-Dist: pyyaml
+
+# PGGM MCP Snowflake Server
+
+A customized Model Context Protocol (MCP) server for Snowflake integration, allowing AI assistants to interact with Snowflake databases securely and efficiently.
+
+## Features
+
+- Connect to Snowflake databases and execute queries
+- Support for various SQL operations and schema exploration
+- Data insights collection and memoization
+- SQL write operation detection for enhanced security
+- Customizable database, schema, and table filtering
+- Support for authentication through environment variables or command-line arguments
+
+## Installation
+
+### Using pip
+
+```bash
+pip install pggm-mcp-snowflake-server
+```
+
+
+### Tools Available to AI Assistants
+
+The server provides the following tools for AI assistants:
+
+- `list_databases` - List all available databases in Snowflake
+- `list_schemas` - List all schemas in a database
+- `list_tables` - List all tables in a specific database and schema
+- `describe_table` - Get the schema information for a specific table
+- `read_query` - Execute a SELECT query
+- `append_insight` - Add a data insight to the memo
+- `write_query` - Execute an INSERT, UPDATE, or DELETE query (if --allow_write is enabled)
+- `create_table` - Create a new table in the Snowflake database (if --allow_write is enabled)
+
+## Security
+
+By default, the server runs in read-only mode. To enable write operations, you must explicitly pass the `--allow_write` flag.
+
+The server uses SQL parsing to detect and prevent write operations in `read_query` calls, ensuring only approved write operations can be executed.
+
+## Development
+
+### Setup Development Environment
+
+```bash
+git clone https://github.com/yourusername/pggm-mcp-snowflake-server.git
+cd pggm-mcp-snowflake-server
+uv venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+uv pip install -e ".[dev]"
+```
+

pggm_mcp_snowflake_server-0.1.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+pggm_mcp_snowflake_server/__init__.py,sha256=fBi0ejuNSFNA_q0h1Jo6zMlsS0R1_xOW8sne9BygAHQ,3304
+pggm_mcp_snowflake_server/server.py,sha256=hfINhJuS2s44obtKWvxfLek6Y5P6GmBLUgkhnr3DuD0,21959
+pggm_mcp_snowflake_server/write_detector.py,sha256=qLFxghERiZT7-DRT8sGPAqrACTgO61QE29QXtH4CN2w,3635
+pggm_mcp_snowflake_server-0.1.1.dist-info/METADATA,sha256=Okb1U8raBISqfDIvOydkYhL2KJ8cf_Z4Fwc97lHRjHM,2328
+pggm_mcp_snowflake_server-0.1.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+pggm_mcp_snowflake_server-0.1.1.dist-info/entry_points.txt,sha256=kbUmsaZT0hYi6naFrGUO-mIKlwqHi_HKO6HNQKdTJE4,84
+pggm_mcp_snowflake_server-0.1.1.dist-info/top_level.txt,sha256=ouamdLwMWx5aSlAHI_mhoPvm9PEBtovD3qbDvR7x284,26
+pggm_mcp_snowflake_server-0.1.1.dist-info/RECORD,,

pggm_mcp_snowflake_server-0.1.0.dist-info/METADATA

@@ -1,34 +0,0 @@
-Metadata-Version: 2.4
-Name: pggm-mcp-snowflake-server
-Version: 0.1.0
-Summary: Custom Model Context Protocol server for Snowflake
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-Requires-Dist: mcp
-Requires-Dist: pydantic
-Requires-Dist: snowflake-snowpark-python
-Requires-Dist: pyyaml
-
-# PGGM MCP Snowflake Server
-
-A customized Model Context Protocol (MCP) server for Snowflake integration, allowing AI assistants to interact with Snowflake databases.
-
-## Features
-
-- Connect to Snowflake databases and execute queries
-- Support for various SQL operations and schema exploration
-- Data insights collection
-- Customized filters and configurations
-
-## Installation
-
-```bash
-pip install pggm-mcp-snowflake-server
-```
-
-## Usage
-
-This package is designed to be used with MCP-compatible AI assistants for database interactions.

pggm_mcp_snowflake_server-0.1.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-pggm_mcp_snowflake_server/__init__.py,sha256=bq16COjhclXSrjQP18TV21vWwJV24hIBrf4I7OfwZkE,48
-pggm_mcp_snowflake_server/server.py,sha256=4_Jlr_o5i4v_39-LuT9CFcqNwGMlNCXSZNgiiSp2vwg,21976
-pggm_mcp_snowflake_server/write_detector.py,sha256=Zli_U5tnIlCzpVoSYT7jPDBlDUSv4nyS1PnIrXpwZYc,1204
-pggm_mcp_snowflake_server-0.1.0.dist-info/METADATA,sha256=B6Pgk0AKQdZgmgWK19TZ-Z6jAH1g4YiKkGskqp-skXg,1015
-pggm_mcp_snowflake_server-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-pggm_mcp_snowflake_server-0.1.0.dist-info/entry_points.txt,sha256=kbUmsaZT0hYi6naFrGUO-mIKlwqHi_HKO6HNQKdTJE4,84
-pggm_mcp_snowflake_server-0.1.0.dist-info/top_level.txt,sha256=ouamdLwMWx5aSlAHI_mhoPvm9PEBtovD3qbDvR7x284,26
-pggm_mcp_snowflake_server-0.1.0.dist-info/RECORD,,