andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-11 00:58:30 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #253 from crewAIInc/fix/clear-more-warnings-databricks

Browse Source 
cleanup more warnings from tools
This commit is contained in:
Lorenze Jay
2025-04-01 10:47:25 -07:00
committed by GitHub
parent d842e9df97 6b4453e1b1
commit 5b06a0c189
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py
View File

@@ -16,7 +16,7 @@ class DatabricksQueryToolSchema(BaseModel):
catalog: Optional[str] = Field( catalog: Optional[str] = Field(
None, description="Databricks catalog name (optional, defaults to configured catalog)" None, description="Databricks catalog name (optional, defaults to configured catalog)"
) )
schema: Optional[str] = Field( db_schema: Optional[str] = Field(
None, description="Databricks schema name (optional, defaults to configured schema)" None, description="Databricks schema name (optional, defaults to configured schema)"
) )
warehouse_id: Optional[str] = Field( warehouse_id: Optional[str] = Field(
@@ -168,7 +168,7 @@ class DatabricksQueryTool(BaseTool):
Args: Args:
query (str): SQL query to execute query (str): SQL query to execute
catalog (Optional[str]): Databricks catalog name catalog (Optional[str]): Databricks catalog name
schema (Optional[str]): Databricks schema name db_schema (Optional[str]): Databricks schema name
warehouse_id (Optional[str]): SQL warehouse ID warehouse_id (Optional[str]): SQL warehouse ID
row_limit (Optional[int]): Maximum number of rows to return row_limit (Optional[int]): Maximum number of rows to return
@@ -179,7 +179,7 @@ class DatabricksQueryTool(BaseTool):
# Get parameters with fallbacks to default values # Get parameters with fallbacks to default values
query = kwargs.get("query") query = kwargs.get("query")
catalog = kwargs.get("catalog") or self.default_catalog catalog = kwargs.get("catalog") or self.default_catalog
schema = kwargs.get("schema") or self.default_schema db_schema = kwargs.get("db_schema") or self.default_schema
warehouse_id = kwargs.get("warehouse_id") or self.default_warehouse_id warehouse_id = kwargs.get("warehouse_id") or self.default_warehouse_id
row_limit = kwargs.get("row_limit", 1000) row_limit = kwargs.get("row_limit", 1000)
@@ -187,7 +187,7 @@ class DatabricksQueryTool(BaseTool):
validated_input = DatabricksQueryToolSchema( validated_input = DatabricksQueryToolSchema(
query=query, query=query,
catalog=catalog, catalog=catalog,
schema=schema, db_schema=db_schema,
warehouse_id=warehouse_id, warehouse_id=warehouse_id,
row_limit=row_limit row_limit=row_limit
) )
@@ -195,15 +195,15 @@ class DatabricksQueryTool(BaseTool):
# Extract validated parameters # Extract validated parameters
query = validated_input.query query = validated_input.query
catalog = validated_input.catalog catalog = validated_input.catalog
schema = validated_input.schema db_schema = validated_input.db_schema
warehouse_id = validated_input.warehouse_id warehouse_id = validated_input.warehouse_id
# Setup SQL context with catalog/schema if provided # Setup SQL context with catalog/schema if provided
context = {} context = {}
if catalog: if catalog:
context["catalog"] = catalog context["catalog"] = catalog
if schema: if db_schema:
context["schema"] = schema context["schema"] = db_schema
# Execute query # Execute query
statement = self.workspace_client.statement_execution statement = self.workspace_client.statement_execution