run-arg-kql

Instructions

Implementation Reference

• src/azure_assistant_mcp/server.py:122-136 (registration)

  Registers the 'run-arg-kql' tool in the MCP server's list_tools() with its input schema definition.

  types.Tool(
      name="run-arg-kql",
      description="Run a provided KQL query against Azure Resource Graph.",
      inputSchema={
          "type": "object",
          "properties": {
              "kql_query": {"type": "string", "description": "KQL to execute (must reference valid ARG tables)"},
              "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": ["kql_query"],
      },
  ),

• src/azure_assistant_mcp/server.py:268-322 (handler)

  The handler logic for the 'run-arg-kql' tool within the MCP call_tool dispatcher: parses args, resolves credentials and scope, executes KQL via execute_kql, formats results.

  if name == "run-arg-kql":
      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))
      kql = arguments.get("kql_query")
      if not kql:
          return [types.TextContent(type="text", text="Error: kql_query is required")]
  
      # Auto-guess tenant if not provided
      if not tenant_name:
          tenant_name = _guess_tenant_name_from_text(kql)
  
      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
      # Only error if neither subscriptions nor MG scope is available
      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)."
          ))]
  
      # Try management group scope if no subs and mg configured
      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 = [
          "Azure Resource Graph Query Results:",
          f"Rows: {result['result_count']}",
          f"Tenant: {tenant_name or AZURE_CONFIG.get_default_tenant().get('name')}",
          scope_line,
          "",
          _format_rows(rows),
          "",
          "KQL:",
          result["query"][:1000],
      ]
      return [types.TextContent(type="text", text="\n".join(body))]

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

  Core utility function execute_kql that performs the actual Azure Resource Graph query execution using the ResourceGraphClient.

  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,
          }