11"""SQL tools - Execute SQL queries and get table information."""
22
3- from typing import Any , Dict , List , Optional
3+ from typing import Any , Dict , List , Optional , Union
44
55from databricks_tools_core .sql import (
66 execute_sql as _execute_sql ,
1414from ..server import mcp
1515
1616
17+ def _format_results_markdown (rows : List [Dict [str , Any ]]) -> str :
18+ """Format SQL results as a markdown table.
19+
20+ Markdown tables state column names once in the header instead of repeating
21+ them on every row (as JSON does), reducing token usage by ~50%.
22+
23+ Args:
24+ rows: List of row dicts from the SQL executor.
25+
26+ Returns:
27+ Markdown table string, or "(no results)" if empty.
28+ """
29+ if not rows :
30+ return "(no results)"
31+
32+ columns = list (rows [0 ].keys ())
33+
34+ # Build header
35+ header = "| " + " | " .join (columns ) + " |"
36+ separator = "| " + " | " .join ("---" for _ in columns ) + " |"
37+
38+ # Build rows — convert None to empty string, stringify everything
39+ data_lines = []
40+ for row in rows :
41+ cells = []
42+ for col in columns :
43+ val = row .get (col )
44+ cell = "" if val is None else str (val )
45+ # Escape pipe characters inside cell values
46+ cell = cell .replace ("|" , "\\ |" )
47+ cells .append (cell )
48+ data_lines .append ("| " + " | " .join (cells ) + " |" )
49+
50+ parts = [header , separator ] + data_lines
51+ # Append row count for awareness
52+ parts .append (f"\n ({ len (rows )} row{ 's' if len (rows ) != 1 else '' } )" )
53+ return "\n " .join (parts )
54+
55+
1756@mcp .tool
1857def execute_sql (
1958 sql_query : str ,
@@ -22,7 +61,8 @@ def execute_sql(
2261 schema : str = None ,
2362 timeout : int = 180 ,
2463 query_tags : str = None ,
25- ) -> List [Dict [str , Any ]]:
64+ output_format : str = "markdown" ,
65+ ) -> Union [str , List [Dict [str , Any ]]]:
2666 """
2767 Execute a SQL query on a Databricks SQL Warehouse.
2868
@@ -40,18 +80,25 @@ def execute_sql(
4080 timeout: Timeout in seconds (default: 180)
4181 query_tags: Optional query tags for cost attribution (e.g., "team:eng,cost_center:701").
4282 Appears in system.query.history and Query History UI.
83+ output_format: Result format — "markdown" (default) or "json".
84+ Markdown tables are ~50% smaller than JSON because column names appear
85+ only once in the header instead of on every row. Use "json" when you
86+ need machine-parseable output.
4387
4488 Returns:
45- List of dictionaries, each representing a row with column names as keys .
89+ Markdown table string (default) or list of row dictionaries (if output_format="json") .
4690 """
47- return _execute_sql (
91+ rows = _execute_sql (
4892 sql_query = sql_query ,
4993 warehouse_id = warehouse_id ,
5094 catalog = catalog ,
5195 schema = schema ,
5296 timeout = timeout ,
5397 query_tags = query_tags ,
5498 )
99+ if output_format == "json" :
100+ return rows
101+ return _format_results_markdown (rows )
55102
56103
57104@mcp .tool
@@ -63,6 +110,7 @@ def execute_sql_multi(
63110 timeout : int = 180 ,
64111 max_workers : int = 4 ,
65112 query_tags : str = None ,
113+ output_format : str = "markdown" ,
66114) -> Dict [str , Any ]:
67115 """
68116 Execute multiple SQL statements with dependency-aware parallelism.
@@ -82,11 +130,14 @@ def execute_sql_multi(
82130 timeout: Timeout per query in seconds (default: 180)
83131 max_workers: Maximum parallel queries per group (default: 4)
84132 query_tags: Optional query tags for cost attribution (e.g., "team:eng,cost_center:701").
133+ output_format: Result format — "markdown" (default) or "json".
134+ Markdown tables are ~50% smaller than JSON because column names appear
135+ only once in the header instead of on every row.
85136
86137 Returns:
87138 Dictionary with results per query and execution summary.
88139 """
89- return _execute_sql_multi (
140+ result = _execute_sql_multi (
90141 sql_content = sql_content ,
91142 warehouse_id = warehouse_id ,
92143 catalog = catalog ,
@@ -95,6 +146,13 @@ def execute_sql_multi(
95146 max_workers = max_workers ,
96147 query_tags = query_tags ,
97148 )
149+ # Format sample_results in each query result if markdown requested
150+ if output_format != "json" and "results" in result :
151+ for query_result in result ["results" ].values ():
152+ sample = query_result .get ("sample_results" )
153+ if sample and isinstance (sample , list ) and len (sample ) > 0 :
154+ query_result ["sample_results" ] = _format_results_markdown (sample )
155+ return result
98156
99157
100158@mcp .tool
0 commit comments