Source code for aidputils.agents.tools.sql
import logging
import time
from threading import Lock
from dataclasses import dataclass
from typing import Optional
from aidputils.agents.auth.util import auth_utils
from aidputils.agents.tools.base_tool import BaseTool
from aidputils.agents.tools.sqltool.connection_manager import _ConnectionManager
from aidputils.agents.tools.sqltool.query_executor import _QueryExecutor
from aidputils.agents.tools.utils import SystemUtils
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.toolkit import chat_context
import json
import re
logger = logging.getLogger(__name__)
@dataclass
class _CacheEntry:
manager: _ConnectionManager
expiry: Optional[float] # None when TTL disabled
[docs]
@BaseTool.register
class SQLTool(BaseTool):
metricsUtil = MetricUtility("SQL_TOOL")
success_counter = metricsUtil.create_counter("SQL_TOOL_SUCCESS_COUNTER")
failure_counter = metricsUtil.create_counter("SQL_TOOL_FAILURE_COUNTER")
connection_pool_success_counter = metricsUtil.create_counter("SQL_TOOL_CONNECTION_POOL_SUCCESS_COUNTER")
connection_pool_failure_counter = metricsUtil.create_counter("SQL_TOOL_CONNECTION_POOL_FAILURE_COUNTER")
connection_refresh_success_counter = metricsUtil.create_counter("SQL_TOOL_CONNECTION_REFRESH_SUCCESS_COUNTER")
connection_refresh_failure_counter = metricsUtil.create_counter("SQL_TOOL_CONNECTION_REFRESH_FAILURE_COUNTER")
_cache = {}
_lock = Lock()
@classmethod
def _get_cache_key(cls, dbname: str, datalakeId: str, catalogKey: str) -> str:
return f"{datalakeId}:{catalogKey}:{dbname}"
@classmethod
def __connect(cls, dbname: str, datalakeId: str, catalogKey: str, **context_vars):
key = cls._get_cache_key(dbname, datalakeId, catalogKey)
now = time.monotonic()
# Resolve TTL (seconds): context var > env var > default 3600
ttl_seconds = cls._resolve_ttl_seconds(context_vars)
logger.info(f"TTL resolved for key {key}: ttl_seconds={ttl_seconds}")
# Fast path: read under lock and return if not expired (or TTL disabled)
with cls._lock:
entry = cls._cache.get(key)
if entry is not None and (ttl_seconds <= 0 or (entry.expiry is not None and now < entry.expiry)):
if ttl_seconds <= 0:
logger.info(f"TTL disabled; returning cached manager for key {key}")
else:
logger.info(f"TTL fast-path hit for key {key}: now={now}, expiry={entry.expiry}")
return entry.manager
# Build new manager outside the lock (slow operations)
prev_exp = entry.expiry if entry is not None else None
logger.info(
f"TTL refresh/init required for key {key} (exists={{}} , expiry={{}} , now={{}} , ttl_seconds={{}})".format(
entry is not None, prev_exp, now, ttl_seconds))
try:
new_manager = cls._build_manager(dbname, datalakeId, catalogKey, **context_vars)
except Exception as e:
logger.error(f"Failed to initialize/refresh connection for key {key}: {str(e)}")
# Increment appropriate metric
if entry is None:
cls.metricsUtil.increment_counter(cls.connection_pool_failure_counter)
else:
cls.metricsUtil.increment_counter(cls.connection_refresh_failure_counter)
raise
# Attempt to install the new manager under the lock
with cls._lock:
expires_at = (now + ttl_seconds) if ttl_seconds > 0 else None
installed, old_manager = cls._try_install(key, new_manager, expires_at, now)
if installed:
# Close old manager outside the lock
if old_manager is not None:
try:
old_manager.close()
except Exception:
logger.exception(f"Error closing old manager for key: {key}")
if old_manager is None:
logger.info(f"Initialized engine for key (new): {key}")
cls.metricsUtil.increment_counter(cls.connection_pool_success_counter)
else:
logger.info(f"Refreshed engine for key due to TTL expiry: {key}")
cls.metricsUtil.increment_counter(cls.connection_refresh_success_counter)
logger.info(f"Installed manager for key {key}; expires_at={expires_at}")
return new_manager
else:
# Another thread won the race; discard the new manager to avoid leaks
logger.info(f"Another thread already refreshed key {key}; discarding newly built manager")
try:
new_manager.close()
except Exception:
logger.exception(f"Error closing unused new manager for key: {key}")
# Return current cached manager
with cls._lock:
return cls._cache[key].manager
@classmethod
def __close_manager(cls, dbname: str, datalakeId: str, catalogKey: str):
key = cls._get_cache_key(dbname, datalakeId, catalogKey)
with cls._lock:
entry = cls._cache.pop(key, None)
if entry is not None:
logger.info(f"Closing manager for key {key}")
if entry.manager is not None:
try:
entry.manager.close()
logger.info(f"Closed pool for key: {key}")
except Exception:
logger.exception(f"Error closing pool for key: {key}")
else:
logger.warning(f"conn_manager for key: {key} is None; skipping close")
# Helpers to simplify __connect
@staticmethod
def _resolve_ttl_seconds(context_vars) -> int:
ttl_seconds = context_vars.get("ttl_seconds")
if ttl_seconds is not None:
try:
val = int(ttl_seconds)
except Exception:
val = 3600
logger.info(f"TTL source=context var, value={val}")
return val
env_ttl = SystemUtils.get_env_var("SQL_CONN_TTL_SECONDS")
if env_ttl:
try:
val = int(env_ttl)
except Exception:
val = 3600
logger.info(f"TTL source=env(SQL_CONN_TTL_SECONDS), value={val}")
return val
logger.info("TTL source=default, value=3600")
return 3600
@staticmethod
def _build_manager(dbname: str, datalakeId: str, catalogKey: str, **context_vars):
m = _ConnectionManager(datalakeId, catalogKey, dbname=dbname, **context_vars)
conn_data = m.fetch_connection_data()
m.create_pool(conn_data)
return m
@classmethod
def _try_install(cls, key: str, new_manager, expires_at, now_monotonic):
# Caller must hold cls._lock
cur_entry = cls._cache.get(key)
cur_exp = cur_entry.expiry if cur_entry else None
if cur_entry is None or (cur_exp is not None and now_monotonic >= cur_exp):
logger.info(
f"Installing new manager for key {key}; prev={'none' if cur_entry is None else 'exists'}; prev_exp={cur_exp}; new_exp={expires_at}")
cls._cache[key] = _CacheEntry(new_manager, expires_at)
return True, (cur_entry.manager if cur_entry else None)
logger.info(f"Skip install for key {key}; current not expired (now={now_monotonic}, cur_exp={cur_exp})")
return False, None
@staticmethod
def _build_mcp_structured_result(response: dict) -> dict:
"""
Build MCP-compliant structuredContent for SQLTool.
Returns a dict suitable for "structuredContent" with:
- rows: always present list of row dicts ([] if missing)
- Optional metadata fields only when present:
rows_fetched, rows_affected, query, error, executiontime_ms
"""
metadata = response.get("metadata") or {}
structured_content = {}
# Always include rows
rows = response.get("result")
if rows is None:
rows = []
# New preferred key
structured_content["rows"] = rows
# Optional metadata-derived fields (camelCase keys)
if "rows_fetched" in metadata:
structured_content["rows_fetched"] = metadata["rows_fetched"]
if "rows_affected" in metadata and metadata["rows_affected"] is not None:
structured_content["rows_affected"] = metadata["rows_affected"]
# query can come from metadata or top-level response fallback
if "query" in metadata:
structured_content["query"] = metadata["query"]
elif "query" in response:
structured_content["query"] = response["query"]
# error can come from metadata or top-level response fallback
if "error" in metadata:
structured_content["error"] = metadata["error"]
elif "error" in response:
structured_content["error"] = response["error"]
# execution time (seconds) -> executionTimeMs
if "execution_time" in metadata and metadata["execution_time"] is not None:
try:
structured_content["executiontime_ms"] = float(metadata["execution_time"]) * 1000.0
except Exception:
# Ignore non-numeric values
pass
return structured_content
@classmethod
def _format_output_as_mcp(cls, response: dict) -> dict:
"""
Return an MCP-style result object with both human-readable content and structuredContent.
"""
structured_content = cls._build_mcp_structured_result(response)
# Legacy JSON block for backward compatibility (older MCP format)
metadata = response.get("metadata") or {}
legacy = {
"result": structured_content.get("rows", structured_content.get("result", []))
}
if "rows_fetched" in metadata:
legacy["rows_fetched"] = metadata["rows_fetched"]
if "rows_affected" in metadata and metadata["rows_affected"] is not None:
legacy["rows_affected"] = metadata["rows_affected"]
if "query" in metadata:
legacy["query"] = metadata["query"]
elif "query" in response:
legacy["query"] = response["query"]
if "error" in metadata:
legacy["error"] = metadata["error"]
elif "error" in response:
legacy["error"] = response["error"]
if "execution_time" in metadata and metadata["execution_time"] is not None:
try:
legacy["executiontime_ms"] = float(metadata["execution_time"]) * 1000.0
except Exception:
pass
return {
"result": {
"content": [
{
"type": "text",
"text": json.dumps(structured_content, ensure_ascii=False, separators=(",", ":"))
}
],
"structuredContent": structured_content
},
# Backward compatible top-level content JSON block
"content": [
{
"type": "json",
"json": legacy
}
]
}
@classmethod
def _invoke_tool(cls, conf, runtime_params, **context_vars) -> dict:
is_success = False
tool_name = chat_context.tool_context_var.get()
cls.logger().info(f"context variable for tool name {tool_name}")
try:
if "datalake_id" in context_vars:
datalakeId = context_vars.get("datalake_id")
else:
datalakeId = SystemUtils.get_env_var("DATALAKE_ID")
query_template = conf.get("query")
cls.logger(tool_name="agent-app-sql").info(f"Invoking SQL tool for {query_template}")
dbname = conf.get("schemaKey") or conf.get("schema")
catalogKey = conf.get("catalogKey") or conf.get("catalog")
logger.info(
f"Retrieved arguments: dbname={dbname}, datalakeId={datalakeId}, catalogKey={catalogKey}, query_template={query_template}, runtime_parms: {runtime_params}")
if not dbname or not catalogKey or not query_template:
return {"error": "Missing required parameters"}
if not datalakeId:
return {
"error": "datalakeId is required but not provided in config or environment variable DATALAKE_ID"}
# Enhanced formatting with multi-level binding
query = None
query, bind_params = cls.format_query(query_template, runtime_params, context_vars)
conn_manager = cls.__connect(dbname, datalakeId, catalogKey, **context_vars)
query_executor = _QueryExecutor(conn_manager.pool)
result = query_executor.execute(query, bind_params)
result["query"] = query # This is the bound query
cls.logger(tool_name="agent-app-sql").info(f"Response from SQL tool for {result}")
is_success = True
return result
except Exception as e:
logger.error(f"Query execution failed: {str(e)}", exc_info=True)
return {"error": str(e), "query": query}
finally:
if is_success:
cls.metricsUtil.increment_counter(cls.success_counter)
else:
cls.metricsUtil.increment_counter(cls.failure_counter)
[docs]
@staticmethod
def format_query(query_template: str, runtime_params: dict, context_vars: dict = None) -> tuple[str, dict]:
"""
Query formatting using the regex-based approach with sqlglot enhancement available.
Automatically classifies placeholders as identifiers or values based on regex pattern analysis.
Returns (formatted_query, bind_params)
"""
logger.info(f"Starting query formatting with template: {query_template}")
logger.debug(f"Runtime params: {runtime_params}")
session_vars = chat_context.session_context_var.get()
merged_params = dict(runtime_params) if runtime_params else {}
if isinstance(session_vars, dict):
# Log only keys whose entry either has no 'shouldLog' property or has shouldLog=True
safe_dict = {k: v for k, v in session_vars.items() if not (isinstance(v, dict) and v.get("shouldLog") is False)}
logger.info(f"loggable session_vars(filtered)={safe_dict}")
# Expect new structure: {"sessionvariables.foo": {"value": "...", "name": "...", "isRequired": "..."}}
# Normalize to use the same keys but convert any dots to underscores for bind-safety.
normalized = {sk.replace(".", "_"): sv.get("value") for sk, sv in session_vars.items()}
# Let explicit runtime params override session variables on key collisions
merged_params = {**normalized, **merged_params}
# Rewrite any {{sessionvariables.*}} placeholders by replacing dots with underscores for bind-safety
query_template = re.sub(
r"\{\{\s*(sessionvariables\.[^}]+?)\s*\}\}",
lambda m: "{{" + m.group(1).replace(".", "_") + "}}",
query_template,
)
# Use the regex-based fallback approach as the primary method
return SQLTool._format_query_regex(query_template, merged_params, context_vars)
@staticmethod
def _is_quoted_literal(query_template: str, start: int, end: int) -> bool:
# Find nearest non-space char before
before_char = ''
for i in range(start - 1, -1, -1):
if query_template[i].strip():
before_char = query_template[i]
break
# Find nearest non-space char after
after_char = ''
for i in range(end, len(query_template)):
if query_template[i].strip():
after_char = query_template[i]
break
return before_char == "'" and after_char == "'"
@staticmethod
def _format_query_regex(query_template: str, runtime_params: dict, context_vars: dict = None) -> tuple[str, dict]:
"""
Fallback query formatting using the exact original regex-based approach.
Supports both runtime parameters {{param}} and system parameters [[param]].
"""
logger.debug(f"Runtime params: {runtime_params}")
logger.debug(f"Context vars: {context_vars}")
# Extract runtime and system parameters
runtime_params_set = set(re.findall(r'\{\{(\w+)\}\}', query_template))
system_params_set = set(re.findall(r'\[\[(\w+)\]\]', query_template))
logger.info(f"Extracted runtime parameters: {sorted(runtime_params_set)}")
logger.info(f"Extracted system parameters: {sorted(system_params_set)}")
upper_template = query_template.upper()
# Validate parameters
param_values = {}
for param in runtime_params_set:
if param not in runtime_params:
raise ValueError(f"Missing runtime param: {param}")
param_values[param] = runtime_params[param]
auth_context = auth_utils.get_auth_context()
for param in system_params_set:
if not auth_context or not hasattr(auth_context, param) or not auth_context.__getattribute__(param):
raise ValueError(f"Missing system param: {param}")
param_values[param] = auth_context.__getattribute__(param)
# Collect all replacements
replacements = []
bind_params = {}
# Process runtime parameters {{param}}
for param in runtime_params_set:
value = param_values[param]
logger.debug(f"Processing runtime parameter '{param}' with value: {value}")
pattern = r'\{\{' + re.escape(param) + r'\}\}'
for i, match in enumerate(re.finditer(pattern, query_template)):
pos = match.start()
before = upper_template[:pos]
after = upper_template[pos + len(match.group()):]
if SQLTool._is_quoted_literal(query_template, pos, pos + len(match.group())):
context = 'literal'
logger.debug(
f" Occurrence {i + 1} at pos {pos}: parameter inside single quotes, classified as literal")
else:
context = 'value' # default
if before and before[-1] == '.':
context = 'identifier'
logger.debug(
f" Occurrence {i + 1} at pos {pos}: qualified reference (before ends with '.'), classified as identifier")
else:
# Find all keywords in before
keyword_pattern = r'\b(SELECT|FROM|JOIN|GROUP\s+BY|ORDER\s+BY|INSERT|INSERT\s+INTO|UPDATE|DELETE|WITH|MERGE|CREATE|ALTER|DROP|ON|USING|WHERE|SET|HAVING|VALUES)\b'
keywords_found = re.findall(keyword_pattern, before)
if keywords_found:
last_keyword = keywords_found[-1]
if last_keyword in ['SELECT', 'FROM', 'JOIN', 'GROUP BY', 'ORDER BY', 'INSERT',
'INSERT INTO', 'UPDATE', 'DELETE', 'WITH', 'MERGE', 'CREATE', 'ALTER',
'DROP', 'ON', 'USING']:
context = 'identifier'
logger.debug(
f" Occurrence {i + 1} at pos {pos}: keyword '{last_keyword}' found, classified as identifier")
elif last_keyword in ['WHERE', 'SET', 'HAVING']:
if re.match(r'\s*(\.|=|<>|<=|>=|<|>|!=|LIKE|IN|NOT\s+IN|IS|IS\s+NOT)\s*', after):
context = 'identifier'
logger.debug(
f" Occurrence {i + 1} at pos {pos}: '{last_keyword}' with operator after, classified as identifier")
elif after == '' and before.rstrip().endswith('.'):
context = 'identifier'
logger.debug(
f" Occurrence {i + 1} at pos {pos}: '{last_keyword}' at end with qualified reference, classified as identifier")
else:
logger.debug(
f" Occurrence {i + 1} at pos {pos}: '{last_keyword}' without operator context, classified as value")
else:
logger.debug(
f" Occurrence {i + 1} at pos {pos}: keyword '{last_keyword}' not in identifier list, classified as value")
else:
logger.debug(f" Occurrence {i + 1} at pos {pos}: no keywords found, classified as value")
if context == 'identifier':
# Validate
if not re.match(r'^[a-zA-Z0-9_]+$', str(value)):
logger.error(f"Invalid identifier value for '{param}': {value}")
raise ValueError(f"Invalid identifier value for '{param}': {value}")
replacement = str(value)
logger.debug(f" -> Replacing with identifier value: '{replacement}'")
elif context == 'literal':
replacement = str(value).replace("'", "''")
logger.debug(f" -> Replacing with escaped literal value: '{replacement}'")
else:
bind_params[param] = value
replacement = f':{param}'
logger.debug(f" -> Replacing with bind parameter: '{replacement}'")
replacements.append((match.start(), match.end(), replacement))
# Process system parameters [[param]] - always treated as values
for param in system_params_set:
value = param_values[param]
logger.debug(f"Processing system parameter '{param}' with value: {value}")
pattern = r'\[\[' + re.escape(param) + r'\]\]'
for i, match in enumerate(re.finditer(pattern, query_template)):
# System parameters are always bound values, never identifiers
bind_params[param] = value
replacement = f':{param}'
logger.debug(f" -> Replacing with bind parameter: '{replacement}'")
replacements.append((match.start(), match.end(), replacement))
# Sort replacements by start position ascending
replacements.sort(key=lambda x: x[0])
# Apply replacements
formatted_query = ''
last_end = 0
for start, end, replacement in replacements:
formatted_query += query_template[last_end:start] + replacement
last_end = end
formatted_query += query_template[last_end:]
logger.info(f"Query formatting completed. Final query: {formatted_query[:200]}{'...' if len(formatted_query) > 200 else ''}")
logger.info(f"Bind parameters: {bind_params}")
logger.debug(f"Total replacements made: {len(replacements)}")
return formatted_query, bind_params