There are several changes, including:

- added a new environment variables handling 'KBC_SNOWFLAKE_SCHEMA'.
- simplified the logic for the query_table that explicitly requires to include the DB name to make queries running correctly.
- commented the query_table_data which will be fixed tomorrow.
- added extra query_table exception handling too.
- updated the documentation.

.python-version

@@ -1 +1 @@
-3.11
+3.11.11

README.md

@@ -7,6 +7,13 @@
 
 A Model Context Protocol (MCP) server for interacting with Keboola Connection. This server provides tools for listing and accessing data from Keboola Storage API.
 
+## Requirements
+
+- Keboola Storage API token
+- Snowflake Read Only Workspace
+
+> Note: The Snowflake package doesn't work with the latest version of Python. If you're using Python 3.12 and above, you'll need to downgrade to Python 3.11.
+
 ## Installation
 
 ### Installing via Smithery
@@ -18,6 +25,7 @@ npx -y @smithery/cli install keboola-mcp-server --client claude
 ```
 
 ### Manual Installation
+
 First, clone the repository and create a virtual environment:
 
 ```bash
@@ -70,6 +78,7 @@ To use this server with Claude Desktop, follow these steps:
         "KBC_SNOWFLAKE_PASSWORD": "your-snowflake-password",
         "KBC_SNOWFLAKE_WAREHOUSE": "your-snowflake-warehouse",
         "KBC_SNOWFLAKE_DATABASE": "your-snowflake-database",
+        "KBC_SNOWFLAKE_SCHEMA": "your-snowflake-schema",
         "KBC_SNOWFLAKE_ROLE": "your-snowflake-role"
       }
     }
@@ -80,14 +89,17 @@ To use this server with Claude Desktop, follow these steps:
 Replace:
 - `/path/to/keboola-mcp-server` with your actual path to the cloned repository
 - `your-keboola-storage-token` with your Keboola Storage API token
-- `YOUR_REGION` with your Keboola region (e.g., `north-europe.azure`, `connection`, etc.)
+- `YOUR_REGION` with your Keboola region (e.g., `north-europe.azure`, etc.). You can remove it if you region is just `connection` explicitly
 - `your-snowflake-account` with your Snowflake account identifier
 - `your-snowflake-user` with your Snowflake username
 - `your-snowflake-password` with your Snowflake password
 - `your-snowflake-warehouse` with your Snowflake warehouse name
 - `your-snowflake-database` with your Snowflake database name
+- `your-snowflake-schema` with your Snowflake schema name
 - `your-snowflake-role` with your Snowflake role name
 
+> Note: If you are using a specific version of Python (e.g. 3.11 due to some package compatibility issues), you'll need to update the `command` into using that specific version, e.g. `/path/to/keboola-mcp-server/.venv/bin/python3.11`
+
 Note: The Snowflake credentials can be obtained by creating a Read Only Snowflake Workspace in your Keboola project (the same project where you got your Storage Token). The workspace will provide all the necessary Snowflake connection parameters.
 
 3. After updating the configuration:

src/keboola_mcp_server/config.py

@@ -34,6 +34,7 @@ def __init__(
         snowflake_warehouse: Optional[str] = None,
         snowflake_database: Optional[str] = None,
         snowflake_role: Optional[str] = None,
+        snowflake_schema: Optional[str] = None,
         log_level: str = "INFO",
     ):
         self.storage_token = storage_token
@@ -44,6 +45,7 @@ def __init__(
         self.snowflake_warehouse = snowflake_warehouse
         self.snowflake_database = snowflake_database
         self.snowflake_role = snowflake_role
+        self.snowflake_schema = snowflake_schema
         self.log_level = log_level
 
     @classmethod
@@ -57,6 +59,7 @@ def from_env(cls) -> "Config":
             "KBC_SNOWFLAKE_WAREHOUSE",
             "KBC_SNOWFLAKE_DATABASE",
             "KBC_SNOWFLAKE_ROLE",
+            "KBC_SNOWFLAKE_SCHEMA",
         ]:
             logger.debug(f"Reading {env_var}: {'set' if os.getenv(env_var) else 'not set'}")
 
@@ -73,6 +76,7 @@ def from_env(cls) -> "Config":
             snowflake_warehouse=os.getenv("KBC_SNOWFLAKE_WAREHOUSE"),
             snowflake_database=os.getenv("KBC_SNOWFLAKE_DATABASE"),
             snowflake_role=os.getenv("KBC_SNOWFLAKE_ROLE"),
+            snowflake_schema=os.getenv("KBC_SNOWFLAKE_SCHEMA"),
             log_level=os.getenv("KBC_LOG_LEVEL", "INFO"),
         )
 

src/keboola_mcp_server/server.py

@@ -48,23 +48,6 @@ class TableDetail(TypedDict):
     db_identifier: str
 
 
-async def snowflake_connection(config: Config):
-    """Create a Snowflake connection."""
-    if not config.has_snowflake_config():
-        raise ValueError("Snowflake credentials not fully configured")
-
-    conn = snowflake.connector.connect(
-        account=config.snowflake_account,
-        user=config.snowflake_user,
-        password=config.snowflake_password,
-        warehouse=config.snowflake_warehouse,
-        database=config.snowflake_database,
-        schema=config.snowflake_schema,
-        role=config.snowflake_role,
-    )
-    return conn
-
-
 def create_server(config: Optional[Config] = None) -> FastMCP:
     """Create and configure the MCP server.
 
@@ -161,7 +144,8 @@ async def get_table_detail(table_id: str) -> TableDetail:
             "db_identifier": await get_table_db_path(table),
         }
 
-    @mcp.tool()
+    # TODO: fix the implementation of query_table_data
+    # @mcp.tool()
     async def query_table_data(
         table_id: str,
         columns: Optional[List[str]] = None,
@@ -190,45 +174,57 @@ async def query_table_data(
 
         result: str = await query_table(query)
         return result
-
+        
     @mcp.tool()
     async def query_table(sql_query: str) -> str:
-        """Execute a Snowflake SQL query to get data from the Storage."""
-        # Get current database
-        db = await get_current_db()
+        """
+            Execute a Snowflake SQL query to get data from the Storage.
+        
+            Note: SQL queries must include the full path including database name, e.g.:
+            'SELECT * FROM SAPI_10025."in.c-fraudDetection"."test_identify"'
+        """
 
         # Execute query
         if not config.has_snowflake_config():
             raise ValueError("Snowflake credentials not fully configured")
 
-        conn = snowflake.connector.connect(
-            account=config.snowflake_account,
-            user=config.snowflake_user,
-            password=config.snowflake_password,
-            warehouse=config.snowflake_warehouse,
-            database=config.snowflake_database,
-            schema=config.snowflake_schema,
-            role=config.snowflake_role,
-        )
+        conn = None
+        cursor = None
 
         try:
+            conn = snowflake.connector.connect(
+                account=config.snowflake_account,
+                user=config.snowflake_user,
+                password=config.snowflake_password,
+                warehouse=config.snowflake_warehouse,
+                database=config.snowflake_database,
+                schema=config.snowflake_schema,
+                role=config.snowflake_role,
+            )
+
             cursor = conn.cursor()
-            try:
-                cursor.execute(f"USE DATABASE {db}")
-                cursor.execute(sql_query)
-                result = cursor.fetchall()
-                columns = [col[0] for col in cursor.description]
-
-                # Convert to CSV
-                output = StringIO()
-                writer = csv.writer(output)
-                writer.writerow(columns)
-                writer.writerows(result)
-                return output.getvalue()
-            finally:
-                cursor.close()
+            cursor.execute(sql_query)
+            result = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+
+            # Convert to CSV
+            output = StringIO()
+            writer = csv.writer(output)
+            writer.writerow(columns)
+            writer.writerows(result)
+            return output.getvalue()
+
+        except snowflake.connector.errors.ProgrammingError as e:
+            raise ValueError(f"Snowflake query error: {str(e)}")
+
+        except Exception as e:
+            raise ValueError(f"Unexpected error during query execution: {str(e)}")
+
         finally:
-            conn.close()
+            if cursor:
+                cursor.close()
+            if conn:
+                conn.close()
 
     # Tools
     @mcp.tool()