data_id

Author: isaacwasserman

src/mcp_snowflake_server/__init__.py

@@ -25,7 +25,7 @@ def main():
         required=False,
         default=[],
         nargs="+",
-        help="List of tools to exclude from the read_query tool description",
+        help="List of tools to exclude",
     )
 
     dotenv.load_dotenv()

src/mcp_snowflake_server/server.py

@@ -57,9 +57,9 @@ def execute_query(self, query: str) -> list[dict[str, Any]]:
         try:
             result = self.session.sql(query).to_pandas()
             result_rows = result.to_dict(orient="records")
-            query_id = str(uuid.uuid4())
+            data_id = str(uuid.uuid4())
 
-            return result_rows, query_id
+            return result_rows, data_id
 
         except Exception as e:
             logger.error(f'Database error executing "{query}": {e}')
@@ -113,7 +113,7 @@ async def handle_list_tables(arguments, db, *_):
         FROM {db.connection_config['database']}.information_schema.tables 
         WHERE table_schema = '{db.connection_config['schema'].upper()}'
     """
-    results, query_id = db.execute_query(query)
+    results, data_id = db.execute_query(query)
     return [types.TextContent(type="text", text=data_to_yaml(results), artifacts=[{"type": "dataframe", "data": results}])]
 
 
@@ -131,10 +131,10 @@ async def handle_describe_table(arguments, db, *_):
         FROM {database_name}.information_schema.columns 
         WHERE table_schema = '{schema_name}' AND table_name = '{table_name}'
     """
-    results, query_id = db.execute_query(query)
+    results, data_id = db.execute_query(query)
     return [
         types.TextContent(
-            type="text", text=data_to_yaml(results), artifacts=[{"type": "dataframe", "data": results, "query_id": query_id}]
+            type="text", text=data_to_yaml(results), artifacts=[{"type": "dataframe", "data": results, "data_id": data_id}]
         )
     ]
 
@@ -144,16 +144,16 @@ async def handle_read_query(arguments, db, write_detector, *_):
     if write_detector.analyze_query(arguments["query"])["contains_write"]:
         raise ValueError("Calls to read_query should not contain write operations")
 
-    results, query_id = db.execute_query(arguments["query"])
+    results, data_id = db.execute_query(arguments["query"])
     truncate = len(results) > MAX_RESULTS
     results_text = data_to_yaml(results[:MAX_RESULTS])
     if truncate:
         results_text += f"\nResults of query have been truncated. There are {len(results) - MAX_RESULTS} more rows."
-    results_text += f"\nquery_id = {query_id}"
+    results_text += f"\ndata_id = {data_id}"
 
     return [
         types.TextContent(
-            type="text", text=results_text, artifacts=[{"type": "dataframe", "data": results, "query_id": query_id}]
+            type="text", text=results_text, artifacts=[{"type": "dataframe", "data": results, "data_id": data_id}]
         )
     ]
 
@@ -173,7 +173,7 @@ async def handle_write_query(arguments, db, _, allow_write, __):
     if arguments["query"].strip().upper().startswith("SELECT"):
         raise ValueError("SELECT queries are not allowed for write_query")
 
-    results, query_id = db.execute_query(arguments["query"])
+    results, data_id = db.execute_query(arguments["query"])
     return [types.TextContent(type="text", text=str(results))]
 
 
@@ -183,21 +183,21 @@ async def handle_create_table(arguments, db, _, allow_write, __):
     if not arguments["query"].strip().upper().startswith("CREATE TABLE"):
         raise ValueError("Only CREATE TABLE statements are allowed")
 
-    results, query_id = db.execute_query(arguments["query"])
-    return [types.TextContent(type="text", text=f"Table created successfully. query_id = {query_id}")]
+    results, data_id = db.execute_query(arguments["query"])
+    return [types.TextContent(type="text", text=f"Table created successfully. data_id = {data_id}")]
 
 
 async def prefetch_tables(db: SnowflakeDB, credentials: dict) -> str:
     """Prefetch table and column information"""
     try:
         logger.info("Prefetching table descriptions")
-        table_results, query_id = db.execute_query(
+        table_results, data_id = db.execute_query(
             f"""SELECT table_name, comment 
                 FROM {credentials['database']}.information_schema.tables 
                 WHERE table_schema = '{credentials['schema'].upper()}'"""
         )
 
-        column_results, query_id = db.execute_query(
+        column_results, data_id = db.execute_query(
             f"""SELECT table_name, column_name, data_type, comment 
                 FROM {credentials['database']}.information_schema.columns 
                 WHERE table_schema = '{credentials['schema'].upper()}'"""