feat(databricks): expand API coverage and preserve paginated results

Author: waleedlatif1

apps/docs/content/docs/integrations/databricks.mdx

@@ -0,0 +1,58 @@
+import { filterUndefined } from '@sim/utils/object'
+import type {
+  DatabricksCancelAllRunsParams,
+  DatabricksSuccessResponse,
+} from '@/tools/databricks/types'
+import {
+  DATABRICKS_AUTH_PARAMS,
+  databricksUrl,
+  readDatabricksError,
+} from '@/tools/databricks/utils'
+import type { ToolConfig } from '@/tools/types'
+
+export const cancelAllRunsTool: ToolConfig<
+  DatabricksCancelAllRunsParams,
+  DatabricksSuccessResponse
+> = {
+  id: 'databricks_cancel_all_runs',
+  name: 'Databricks Cancel All Runs',
+  description: 'Request cancellation of active runs of a job or queued runs in the workspace.',
+  version: '1.0.0',
+  params: {
+    ...DATABRICKS_AUTH_PARAMS,
+    jobId: {
+      type: 'number',
+      required: false,
+      visibility: 'user-or-llm',
+      description: 'Job identifier',
+    },
+    allQueuedRuns: {
+      type: 'boolean',
+      required: false,
+      visibility: 'user-or-llm',
+      description: 'Cancel queued runs; without a job ID this applies to the workspace',
+    },
+  },
+  request: {
+    url: (params) => databricksUrl(params.host, `/api/2.2/jobs/runs/cancel-all`),
+    method: 'POST',
+    headers: (params) => ({
+      Authorization: `Bearer ${params.apiKey}`,
+      'Content-Type': 'application/json',
+    }),
+    body: (params) => {
+      return filterUndefined({
+        job_id: params.jobId,
+        all_queued_runs: params.allQueuedRuns,
+      })
+    },
+  },
+  transformResponse: async (response) => {
+    if (!response.ok)
+      throw new Error(await readDatabricksError(response, 'Failed to cancel all runs'))
+    return { success: true, output: { success: true } }
+  },
+  outputs: {
+    success: { type: 'boolean', description: 'Whether Databricks accepted the request' },
+  },
+}

apps/sim/blocks/blocks/databricks.ts

@@ -34,7 +34,7 @@ export const cancelRunTool: ToolConfig<DatabricksCancelRunParams, DatabricksCanc
   },
 
   request: {
-    url: (params) => databricksUrl(params.host, '/api/2.1/jobs/runs/cancel'),
+    url: (params) => databricksUrl(params.host, '/api/2.2/jobs/runs/cancel'),
     method: 'POST',
     headers: (params) => ({
       'Content-Type': 'application/json',

apps/sim/tools/databricks/cancel_all_runs.ts

@@ -0,0 +1,61 @@
+import { filterUndefined } from '@sim/utils/object'
+import type {
+  DatabricksGetStatementParams,
+  DatabricksSuccessResponse,
+} from '@/tools/databricks/types'
+import {
+  DATABRICKS_AUTH_PARAMS,
+  databricksUrl,
+  optionalDatabricksString,
+  readDatabricksError,
+} from '@/tools/databricks/utils'
+import type { ToolConfig } from '@/tools/types'
+import { safeUrlPathSegment } from '@/tools/url-path'
+
+export const cancelStatementTool: ToolConfig<
+  DatabricksGetStatementParams,
+  DatabricksSuccessResponse
+> = {
+  id: 'databricks_cancel_statement',
+  name: 'Databricks Cancel Statement',
+  description:
+    'Request cancellation of an executing SQL statement; poll Get Statement to confirm its final status.',
+  version: '1.0.0',
+  params: {
+    ...DATABRICKS_AUTH_PARAMS,
+    statementId: {
+      type: 'string',
+      required: true,
+      visibility: 'user-or-llm',
+      description: 'SQL statement identifier',
+    },
+    warehouseId: {
+      type: 'string',
+      required: false,
+      visibility: 'user-or-llm',
+      description: 'SQL warehouse identifier, when required by the statement endpoint',
+    },
+  },
+  request: {
+    url: (params) =>
+      databricksUrl(
+        params.host,
+        `/api/2.0/sql/statements/${safeUrlPathSegment(params.statementId, 'statementId')}/cancel`
+      ),
+    method: 'POST',
+    headers: (params) => ({
+      Authorization: `Bearer ${params.apiKey}`,
+      'Content-Type': 'application/json',
+    }),
+    body: (params) =>
+      filterUndefined({ warehouse_id: optionalDatabricksString(params.warehouseId) }),
+  },
+  transformResponse: async (response) => {
+    if (!response.ok)
+      throw new Error(await readDatabricksError(response, 'Failed to cancel SQL statement'))
+    return { success: true, output: { success: true } }
+  },
+  outputs: {
+    success: { type: 'boolean', description: 'Whether the cancellation request was accepted' },
+  },
+}