From 7a9f71331252641bd029d6e22ceb7b3bb45f18ae Mon Sep 17 00:00:00 2001
From: Jason Summer <jason.summer@snowflake.com>
Date: Mon, 27 Oct 2025 10:05:32 -0500
Subject: [PATCH] fix: Add request timeout and limit validation

---
 mcp_server_snowflake/cortex_services/tools.py | 47 +++++++++++++++----
 .../semantic_manager/tools.py                 |  7 ++-
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/mcp_server_snowflake/cortex_services/tools.py b/mcp_server_snowflake/cortex_services/tools.py
index 6ace9d6..bdbaa8b 100644
--- a/mcp_server_snowflake/cortex_services/tools.py
+++ b/mcp_server_snowflake/cortex_services/tools.py
@@ -78,8 +78,16 @@ async def query_cortex_agent(
         "tool_choice": {"type": "auto"},
         "stream": False,  # Ignored by Agent API
     }
+    try:
+        response = requests.post(
+            host, headers=headers, json=payload, stream=True, timeout=120
+        )
+    except requests.exceptions.Timeout:
+        raise SnowflakeException(
+            tool="Cortex Agent",
+            message="Request timed out",
+        )
 
-    response = requests.post(host, headers=headers, json=payload, stream=True)
     try:
         response.raise_for_status()
         return response
@@ -150,6 +158,14 @@ async def query_cortex_search(
     if filter_query is None:
         filter_query = {}
 
+    if (
+        limit is not None and not 0 <= int(limit) <= 1000
+    ):  # Cortex Search limits to 1000 server-side
+        raise SnowflakeException(
+            tool="Cortex Search",
+            message="Limit must be between 0 and 1,000",
+        )
+
     payload = {
         "query": query,
         "filter": filter_query,
@@ -158,13 +174,18 @@ async def query_cortex_search(
 
     if isinstance(columns, list) and len(columns) > 0:
         payload["columns"] = columns
+    try:
+        response = requests.post(host, headers=headers, json=payload, timeout=60)
+    except requests.exceptions.Timeout:
+        raise SnowflakeException(
+            tool="Cortex Search",
+            message="Request timed out",
+        )
 
-    response = requests.post(host, headers=headers, json=payload)
-
-    if response.status_code == 200:
+    try:
+        response.raise_for_status()
         return response
-
-    else:
+    except Exception:
         raise SnowflakeException(
             tool="Cortex Search",
             status_code=response.status_code,
@@ -239,12 +260,18 @@ async def query_cortex_analyst(
         "stream": False,
     }
 
-    response = requests.post(host, headers=headers, json=payload)
+    try:
+        response = requests.post(host, headers=headers, json=payload, timeout=120)
+    except requests.exceptions.Timeout:
+        raise SnowflakeException(
+            tool="Cortex Analyst",
+            message="Request timed out",
+        )
 
-    if response.status_code == 200:
+    try:
+        response.raise_for_status()
         return response
-
-    else:
+    except Exception:
         raise SnowflakeException(
             tool="Cortex Analyst",
             status_code=response.status_code,
diff --git a/mcp_server_snowflake/semantic_manager/tools.py b/mcp_server_snowflake/semantic_manager/tools.py
index 1c93ee2..98ff864 100644
--- a/mcp_server_snowflake/semantic_manager/tools.py
+++ b/mcp_server_snowflake/semantic_manager/tools.py
@@ -221,7 +221,12 @@ def write_semantic_view_query(
     if order_by:
         statement += f" ORDER BY {order_by}"
 
-    if limit:
+    if limit is not None:
+        if not 0 <= int(limit) <= 1000:
+            raise SnowflakeException(
+                tool="write_semantic_view_query",
+                message="Limit must be between 0 and 1,000",
+            )
         statement += f" LIMIT {int(limit)}"
 
     try: