run-kql-template

Instructions

Implementation Reference

• src/azure_assistant_mcp/server.py:324-385 (handler)

  Main execution logic for the 'run-kql-template' tool within the call_tool dispatcher. Handles argument parsing, template loading and parameterization, scope resolution, query execution, and result formatting.

  if name == "run-kql-template":
      tenant_name = arguments.get("tenant_name")
      subs = arguments.get("subscription_ids")
      use_all = bool(arguments.get("use_all_subscriptions", True))
      top = int(arguments.get("top", 100))
      template_name = arguments.get("template_name")
      params: Dict[str, Any] = arguments.get("params") or {}
      if not template_name:
          return [types.TextContent(type="text", text="Error: template_name is required")]
  
      kql = load_kql_template(template_name)
      if not kql:
          return [types.TextContent(type="text", text=f"Error: template '{template_name}' not found in kql templates.")]
  
      # Simple {{key}} replacement
      try:
          for k, v in params.items():
              kql = kql.replace(f"{{{{{k}}}}}", str(v))
      except Exception:
          pass
  
      # Auto-guess tenant if not provided (best effort: look for tenant-like strings is skipped here)
      cred, default_subs = AZURE_CONFIG.get_credentials(tenant_name)
      mg = AZURE_CONFIG.get_management_group_id(tenant_name)
      if subs:
          pass
      elif use_all and mg:
          subs = []  # force MG usage below
      elif use_all:
          discovered = _enumerate_subscriptions_for_credential(cred)
          subs = discovered or default_subs
      else:
          subs = default_subs
      if not subs and not (use_all and mg):
          return [types.TextContent(type="text", text=(
              "Error: No subscriptions available; provide subscription_ids, set default_subscription_id in config, "
              "or ensure subscription discovery is possible (azure-mgmt-subscription installed and SP has access)."
          ))]
  
      if use_all and mg and not subs:
          result = execute_kql(cred, None, kql, top=top, management_groups=[mg])
      else:
          result = execute_kql(cred, subs, kql, top=top)
      if result["status"] != "success":
          return [types.TextContent(type="text", text=f"ARG query failed: {result.get('error','unknown error')}")]
  
      rows = result["results"]
      scope_line = (
          f"Scope: managementGroup={mg}" if (use_all and mg and not subs) else f"Subscriptions used: {len(subs)}"
      )
      body = [
          f"KQL Template: {template_name}",
          f"Rows: {result['result_count']}",
          f"Tenant: {tenant_name or AZURE_CONFIG.get_default_tenant().get('name')}",
          scope_line,
          "",
          _format_rows(rows),
          "",
          "KQL:",
          kql[:1000],
      ]
      return [types.TextContent(type="text", text="\n".join(body))]

• src/azure_assistant_mcp/server.py:137-152 (registration)

  Registration of the 'run-kql-template' tool in the list_tools() handler, including name, description, and full input schema.

  types.Tool(
      name="run-kql-template",
      description="Execute a KQL template from kql/ by name, with optional {{param}} replacements.",
      inputSchema={
          "type": "object",
          "properties": {
              "template_name": {"type": "string", "description": "Template filename without extension (looks for .md or .kql in kql folder)"},
              "params": {"type": "object", "description": "Optional key/value replacements for {{key}} placeholders in the template"},
              "tenant_name": {"type": "string", "description": "Optional configured tenant name"},
              "subscription_ids": {"type": "array", "items": {"type": "string"}, "description": "Optional explicit subscription IDs"},
              "use_all_subscriptions": {"type": "boolean", "description": "If no subscriptions are provided, attempt to auto-discover all accessible subscriptions (default: true)", "default": True},
              "top": {"type": "integer", "description": "Max rows to return (default: 100)", "default": 100}
          },
          "required": ["template_name"],
      },
  ),

• src/azure_assistant_mcp/kql_loader.py:20-44 (helper)

  Helper function specifically used by 'run-kql-template' to load KQL templates from the 'kql/' directory, supporting .kql files or fenced blocks in .md files.

  def load_kql_template(name: str) -> Optional[str]:
      """Load a KQL template by name.
  
      Supports either plain `.kql` files or Markdown with a fenced ```kql block.
      Returns None if no template is found or parse fails.
      """
      base = _base_dir()
      # Try .kql first
      kql_path = base / f"{name}.kql"
      if kql_path.is_file():
          try:
              return kql_path.read_text(encoding="utf-8").strip()
          except Exception:
              pass
      # Try .md with fenced code block
      md_path = base / f"{name}.md"
      if md_path.is_file():
          try:
              text = md_path.read_text(encoding="utf-8")
              m = _FENCE_RE.search(text)
              if m:
                  return m.group(1).strip()
          except Exception:
              pass
      return None

• src/azure_assistant_mcp/arg.py:7-49 (helper)

  Shared helper function called by 'run-kql-template' (and other tools) to execute the KQL query against Azure Resource Graph service.

  def execute_kql(
      credential,
      subscriptions: Optional[List[str]],
      kql_query: str,
      top: int = 100,
      management_groups: Optional[List[str]] = None,
  ) -> Dict[str, Any]:
      """Execute a KQL query against Azure Resource Graph.
  
      Returns a dict with keys: status, results, result_count, warnings, query
      """
      client = ResourceGraphClient(credential=credential)
  
      # Build typed request for reliability
      options = QueryRequestOptions(result_format="objectArray", top=top)
      # Either subscriptions or management_groups must be provided.
      request = QueryRequest(
          subscriptions=subscriptions,
          management_groups=management_groups,
          query=kql_query,
          options=options,
      )
  
      try:
          response = client.resources(request)
          # response.data is a list[dict] when result_format=objectArray
          rows = list(response.data or [])
          return {
              "status": "success",
              "results": rows,
              "result_count": len(rows),
              "warnings": [],
              "query": kql_query,
          }
      except Exception as e:
          return {
              "status": "error",
              "error": str(e),
              "results": [],
              "result_count": 0,
              "warnings": [],
              "query": kql_query,
          }