Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:CrewAIInc CrewAI Databricks Query Tool

From Leeroopedia
Knowledge Sources
Domains Database Integration, Data Analytics, Tool Integration
Last Updated 2026-02-11 00:00 GMT

Overview

DatabricksQueryTool is a CrewAI tool that executes SQL queries against Databricks workspace tables using the Databricks SDK, with extensive result processing and error handling for production-grade data analytics workflows.

Description

The tool extends BaseTool and uses the Databricks SDK's statement execution API to run SQL queries. It supports two authentication methods: Databricks CLI profile (DATABRICKS_CONFIG_PROFILE) or direct credentials (DATABRICKS_HOST and DATABRICKS_TOKEN). The tool manages warehouse connections, polls for query completion with a 5-minute timeout, and processes results with sophisticated data reconstruction logic.

A key feature of this tool is its handling of malformed result structures from the Databricks API. The implementation includes:

  • Pattern recognition for detecting incorrectly structured rows (character-by-character data splits, flattened arrays)
  • Data reassembly using regex-based ID detection and title reconstruction
  • Heuristic-based detection of whether data needs special handling (analyzing sample sizes for single-character/digit patterns)
  • Normalized result formatting as readable tables with column alignment

The DatabricksQueryToolSchema validates input parameters, enforces non-empty queries, and automatically appends LIMIT clauses when a row_limit is set and the query lacks one.

Usage

Use this tool when agents need to query Databricks data warehouses, run analytics queries, or retrieve data for analysis. It is essential for enterprise data analytics workflows where CrewAI agents interact with Databricks-hosted data.

Code Reference

Source Location

  • Repository: CrewAI
  • File: lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py
  • Lines: 1-851

Signature

class DatabricksQueryTool(BaseTool):
    name: str = "Databricks SQL Query"
    description: str = (
        "Execute SQL queries against Databricks workspace tables and return the results."
        " Provide a 'query' parameter with the SQL query to execute."
    )
    args_schema: type[BaseModel] = DatabricksQueryToolSchema
    default_catalog: str | None = None
    default_schema: str | None = None
    default_warehouse_id: str | None = None
    package_dependencies: list[str] = ["databricks-sdk"]

    def __init__(
        self,
        default_catalog: str | None = None,
        default_schema: str | None = None,
        default_warehouse_id: str | None = None,
        **kwargs: Any,
    ) -> None: ...

Import

from crewai_tools.tools.databricks_query_tool.databricks_query_tool import DatabricksQueryTool

I/O Contract

Inputs

Name Type Required Description
query str Yes SQL query to execute against Databricks workspace tables
catalog str No Databricks catalog name (defaults to configured catalog)
db_schema str No Databricks schema name (defaults to configured schema)
warehouse_id str No Databricks SQL warehouse ID (defaults to configured warehouse)
row_limit int No Maximum number of rows to return (default: 1000)

Outputs

Name Type Description
return str Formatted query results as an aligned table with column headers and row count, or an error/status message

Environment Variables

Name Required Description
DATABRICKS_CONFIG_PROFILE One of profile or host+token Databricks CLI profile name
DATABRICKS_HOST One of profile or host+token Databricks workspace URL
DATABRICKS_TOKEN One of profile or host+token Databricks access token

Usage Examples

Basic Usage

import os
os.environ["DATABRICKS_HOST"] = "https://my-workspace.databricks.com"
os.environ["DATABRICKS_TOKEN"] = "dapi..."

from crewai_tools.tools.databricks_query_tool.databricks_query_tool import DatabricksQueryTool

tool = DatabricksQueryTool(
    default_catalog="main",
    default_schema="default",
    default_warehouse_id="abc123def456"
)

result = tool._run(query="SELECT * FROM customers WHERE region = 'US'", row_limit=100)
print(result)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment