From bf58535ce1155b8f161484dbbca8b9979907473d Mon Sep 17 00:00:00 2001
From: SurgeODDR <olivier@surgemgmt.com>
Date: Tue, 1 Apr 2025 17:29:33 +0200
Subject: [PATCH] Remove .env.example file, update server script permissions,
 and modify SQL execution parameters for improved functionality

---
 .env.example                        | 11 -----------
 scripts/start_mcp_server.sh         |  2 +-
 src/api/sql.py                      |  8 +++++---
 src/server/databricks_mcp_server.py |  7 ++++++-
 start_mcp_server.sh                 |  0
 5 files changed, 12 insertions(+), 16 deletions(-)
 delete mode 100644 .env.example
 mode change 100644 => 100755 scripts/start_mcp_server.sh
 mode change 100644 => 100755 start_mcp_server.sh

diff --git a/.env.example b/.env.example
deleted file mode 100644
index 40a2da9..0000000
--- a/.env.example
+++ /dev/null
@@ -1,11 +0,0 @@
-# Databricks API configuration
-DATABRICKS_HOST=https://adb-xxxxxxxxxxxx.xx.azuredatabricks.net
-DATABRICKS_TOKEN=your_databricks_token_here
-
-# Server configuration
-SERVER_HOST=0.0.0.0
-SERVER_PORT=8000
-DEBUG=False
-
-# Logging
-LOG_LEVEL=INFO 
\ No newline at end of file
diff --git a/scripts/start_mcp_server.sh b/scripts/start_mcp_server.sh
old mode 100644
new mode 100755
index 4e92645..b49d43e
--- a/scripts/start_mcp_server.sh
+++ b/scripts/start_mcp_server.sh
@@ -27,6 +27,6 @@ if [ -n "$DATABRICKS_HOST" ]; then
     echo "Databricks Host: $DATABRICKS_HOST"
 fi
 
-uv run src.server.databricks_mcp_server
+python -m src.server.databricks_mcp_server
 
 echo "Server stopped at $(date)" 
\ No newline at end of file
diff --git a/src/api/sql.py b/src/api/sql.py
index 43cdb1a..91fde94 100644
--- a/src/api/sql.py
+++ b/src/api/sql.py
@@ -43,9 +43,11 @@ async def execute_statement(
     request_data = {
         "statement": statement,
         "warehouse_id": warehouse_id,
-        "wait_timeout": "0s",  # Wait indefinitely
+        "wait_timeout": "10s",
+        "format": "JSON_ARRAY",
+        "disposition": "INLINE",
         "row_limit": row_limit,
-        "byte_limit": byte_limit,
+        "byte_limit": 16777216,
     }
     
     if catalog:
@@ -57,7 +59,7 @@ async def execute_statement(
     if parameters:
         request_data["parameters"] = parameters
         
-    return make_api_request("POST", "/api/2.0/sql/statements/execute", data=request_data)
+    return make_api_request("POST", "/api/2.0/sql/statements", data=request_data)
 
 
 async def execute_and_wait(
diff --git a/src/server/databricks_mcp_server.py b/src/server/databricks_mcp_server.py
index 4534c97..649df2f 100644
--- a/src/server/databricks_mcp_server.py
+++ b/src/server/databricks_mcp_server.py
@@ -202,7 +202,12 @@ async def execute_sql(params: Dict[str, Any]) -> List[TextContent]:
                 catalog = params.get("catalog")
                 schema = params.get("schema")
                 
-                result = await sql.execute_sql(statement, warehouse_id, catalog, schema)
+                result = await sql.execute_statement(
+                    statement=statement,
+                    warehouse_id=warehouse_id,
+                    catalog=catalog,
+                    schema=schema
+                )
                 return [{"text": json.dumps(result)}]
             except Exception as e:
                 logger.error(f"Error executing SQL: {str(e)}")
diff --git a/start_mcp_server.sh b/start_mcp_server.sh
old mode 100644
new mode 100755