omnigent-ai--omnigent/omnigent/runner/uc_function.py

"""Unity Catalog function execution for the runner.

Executes UC SQL functions declared via ``catalog_path:`` in agent
YAML tool definitions. Functions are called through the Databricks
SQL Statement Execution API (``WorkspaceClient
.statement_execution.execute_statement()``), which supports
parameterized queries and avoids SQL injection.

The ``WorkspaceClient`` is constructed once per profile and cached
for the lifetime of the process — UC function calls are frequent
during agent turns and workspace auth resolution is expensive (SDK
reads ``~/.databrickscfg``, fetches OAuth tokens, etc.).

Design decisions:

- **Statement Execution over direct invoke**: UC functions are SQL
  objects; the canonical invocation path is ``SELECT func(args)``.
  The SDK's ``FunctionsAPI`` has no ``execute`` method — it only
  supports CRUD on function metadata. Statement Execution is the
  only SDK path that runs a function and returns results.
- **Warehouse ID resolution**: Statement Execution needs a SQL
  warehouse. The tool YAML can declare ``warehouse_id:`` per tool;
  when absent, the ``DATABRICKS_WAREHOUSE_ID`` environment variable
  is read at runtime as a convenience fallback.
- **Parameter schema**: The LLM must know what arguments a UC
  function accepts. The YAML's ``parameters:`` block is the source
  of truth (populated by the author or fetched from UC metadata at
  agent-build time — the latter is a future enhancement).
"""

from __future__ import annotations

import json
import logging
import re
from functools import lru_cache
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from databricks.sdk import WorkspaceClient

_logger = logging.getLogger(__name__)

# Validates catalog_path: bare identifier (e.g. "ai_query") or
# dotted three-level name (e.g. "my_catalog.my_schema.func").
# Rejects backticks, semicolons, parens, and other SQL metacharacters.
_CATALOG_PATH_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_.]*$")

# Validates parameter names from the LLM's tool-call arguments.
# Rejects SQL metacharacters (parens, colons, semicolons, etc.).
_PARAM_NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")


@lru_cache(maxsize=32)
def _get_workspace_client(
    profile: str | None,
) -> WorkspaceClient:
    """
    Construct a cached :class:`WorkspaceClient` for the given
    Databricks profile.

    Caches by ``profile`` so repeated tool calls within a session
    reuse the same authenticated client. The ``maxsize=32`` cap
    prevents unbounded growth in multi-profile deployments.

    :param profile: Databricks config profile name from
        ``~/.databrickscfg``, e.g. ``"oss"``. ``None`` uses the
        SDK's default resolution (``DEFAULT`` section, env vars).
    :returns: An authenticated ``WorkspaceClient`` instance.
    :raises ImportError: If ``databricks-sdk`` is not installed.
    """
    from databricks.sdk import WorkspaceClient as _WSC

    if profile is not None:
        return _WSC(profile=profile)
    return _WSC()


def _build_select_statement(
    catalog_path: str,
    args: dict[str, object],
) -> tuple[str, list[dict[str, str]]]:
    """
    Build a parameterized ``SELECT`` statement for a UC function
    call.

    Uses named parameter markers (``:param_name``) instead of
    string interpolation to prevent SQL injection. The Databricks
    Statement Execution API resolves the markers server-side.

    :param catalog_path: UC function name — either a three-level
        qualified name (``"my_catalog.my_schema.classify"``) or a
        bare built-in function name (``"ai_query"``). Three-level
        names are backtick-quoted to handle special characters;
        bare names are emitted unquoted so built-in SQL functions
        like ``ai_query`` resolve correctly.
    :param args: Argument dict from the LLM, e.g.
        ``{"text": "I love Databricks", "lang": "en"}``.
    :returns: A ``(sql, parameters)`` tuple where ``sql`` is the
        parameterized query string and ``parameters`` is a list
        of ``StatementParameterListItem``-compatible dicts with
        ``name`` and ``value`` keys.
    """
    # Validate catalog_path to prevent SQL injection via crafted
    # tool specs. Only alphanumeric, underscore, and dot are allowed.
    if not _CATALOG_PATH_RE.match(catalog_path):
        raise ValueError(
            f"Invalid catalog_path {catalog_path!r}: must contain only "
            f"alphanumeric characters, underscores, and dots."
        )

    # Validate parameter names from the LLM's tool-call arguments.
    # These are interpolated as ``:name`` markers in the SQL string;
    # SQL metacharacters in keys would alter the query structure.
    for name in args:
        if not _PARAM_NAME_RE.match(name):
            raise ValueError(
                f"Invalid parameter name {name!r} for UC function "
                f"{catalog_path!r}: must be alphanumeric/underscore."
            )

    # Bare names (no dots) are built-in SQL functions like ai_query
    # — emit unquoted. Three-level catalog paths are backtick-quoted
    # so dots inside catalog/schema/function names are not misread
    # as identifier separators.
    func_ref = catalog_path if "." not in catalog_path else f"`{catalog_path}`"

    if not args:
        sql = f"SELECT {func_ref}()"
        return sql, []

    param_names = list(args.keys())
    placeholders = ", ".join(f":{name}" for name in param_names)
    sql = f"SELECT {func_ref}({placeholders})"

    parameters = [
        {"name": name, "value": json.dumps(value) if not isinstance(value, str) else value}
        for name, value in args.items()
    ]
    return sql, parameters


async def execute_uc_function(
    catalog_path: str,
    args: dict[str, object],
    *,
    profile: str | None = None,
    warehouse_id: str | None = None,
) -> str:
    """
    Execute a Unity Catalog function and return the result as a
    string.

    Constructs a parameterized ``SELECT catalog.schema.func(:args)``
    query and executes it via the Databricks SQL Statement Execution
    API. The result is extracted from the response's ``data_array``
    and returned as a JSON string.

    :param catalog_path: Three-level UC function name, e.g.
        ``"my_catalog.my_schema.classify_sentiment"``.
    :param args: Argument dict from the LLM, e.g.
        ``{"text": "I love it"}``.
    :param profile: Databricks config profile name, e.g.
        ``"oss"``. ``None`` uses the SDK's default resolution.
    :param warehouse_id: SQL warehouse ID to execute against, e.g.
        ``"abc123def456"``. When ``None``, falls back to the
        ``DATABRICKS_WAREHOUSE_ID`` environment variable.
    :returns: The function's return value as a JSON string. Scalar
        results are returned directly; multi-row results are
        returned as a JSON array.
    :raises ValueError: If no warehouse ID is available from
        either the parameter or the environment variable.
    :raises RuntimeError: If the statement execution fails or
        returns an unexpected status.
    """
    import asyncio
    import os

    if not warehouse_id:
        warehouse_id = os.environ.get("DATABRICKS_WAREHOUSE_ID")
    if not warehouse_id:
        raise ValueError(
            f"UC function {catalog_path!r} requires a warehouse_id. "
            f"Set 'warehouse_id' in the tool's YAML definition or "
            f"the DATABRICKS_WAREHOUSE_ID environment variable."
        )

    client = _get_workspace_client(profile)
    sql, parameters = _build_select_statement(catalog_path, args)

    _logger.info(
        "Executing UC function %s on warehouse %s",
        catalog_path,
        warehouse_id,
    )

    # Statement execution is a blocking HTTP call — run in a thread
    # to avoid blocking the event loop.
    from databricks.sdk.service.sql import StatementParameterListItem

    sdk_params = [StatementParameterListItem(name=p["name"], value=p["value"]) for p in parameters]

    response = await asyncio.to_thread(
        client.statement_execution.execute_statement,
        statement=sql,
        warehouse_id=warehouse_id,
        parameters=sdk_params,
    )

    # Check execution status.
    status = response.status
    if status is None:
        raise RuntimeError(
            f"UC function {catalog_path!r}: statement execution returned no status."
        )
    state = status.state
    if state is None:
        raise RuntimeError(f"UC function {catalog_path!r}: statement status has no state.")

    from databricks.sdk.service.sql import StatementState

    if state == StatementState.FAILED:
        error = status.error
        msg = error.message if error else "unknown error"
        raise RuntimeError(f"UC function {catalog_path!r} execution failed: {msg}")
    if state not in (StatementState.SUCCEEDED,):
        raise RuntimeError(
            f"UC function {catalog_path!r}: unexpected state {state.value!r}. Expected SUCCEEDED."
        )

    # Extract result data.
    result = response.result
    if result is None or result.data_array is None:
        _logger.debug("UC function %s returned no result data", catalog_path)
        return json.dumps(None)

    data = result.data_array
    _logger.debug("UC function %s result: %s", catalog_path, data)
    # Single-row, single-column result (common for scalar functions).
    if len(data) == 1 and len(data[0]) == 1:
        return data[0][0] if data[0][0] is not None else json.dumps(None)

    # Multi-row or multi-column: return as JSON array of arrays.
    return json.dumps(data)