Skip to content

Commit 4fc10e4

Browse files
waleedlatif1waleedlatif1
committed
fix(databricks): preserve mutation intent and complete task outputs
1 parent 3dc0283 commit 4fc10e4

9 files changed

Lines changed: 604 additions & 9 deletions

File tree

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

Lines changed: 330 additions & 2 deletions
Large diffs are not rendered by default.

‎apps/sim/tools/databricks/cancel_all_runs.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ export const cancelAllRunsTool: ToolConfig<
2424
type: 'number',
2525
required: false,
2626
visibility: 'user-or-llm',
27-
description: 'Job identifier',
27+
description: 'Job identifier; required unless allQueuedRuns is true',
2828
},
2929
allQueuedRuns: {
3030
type: 'boolean',
@@ -41,6 +41,11 @@ export const cancelAllRunsTool: ToolConfig<
4141
'Content-Type': 'application/json',
4242
}),
4343
body: (params) => {
44+
if (params.jobId == null && params.allQueuedRuns !== true) {
45+
throw new Error(
46+
'Provide jobId or set allQueuedRuns to true to cancel queued workspace runs'
47+
)
48+
}
4449
return filterUndefined({
4550
job_id: params.jobId,
4651
all_queued_runs: params.allQueuedRuns,

‎apps/sim/tools/databricks/databricks.test.ts‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { executeSqlTool } from '@/tools/databricks/execute_sql'
1111
import { getRunTool } from '@/tools/databricks/get_run'
1212
import { executeTool } from '@/tools/index'
1313
import { tools } from '@/tools/registry'
14+
import { prepareToolRequest } from '@/tools/request-transport'
1415

1516
/** Registers only this service's configs in the global registry mock; the full one is ~6,000 modules. */
1617
Object.assign(tools, partialToolRegistry(databricksTools))
@@ -48,6 +49,51 @@ const HOST_ALLOWLIST_ERROR =
4849

4950
const DATABRICKS_TOOL_IDS = Object.keys(tools).filter((id) => id.startsWith('databricks_'))
5051

52+
describe('Databricks mutation intent', () => {
53+
const credentials = { host: 'dbc-test.cloud.databricks.com', apiKey: 'test-token' }
54+
55+
it('distinguishes clearing a Genie description from leaving it unchanged', () => {
56+
const update = (description?: string) =>
57+
JSON.parse(
58+
prepareToolRequest(databricksTools.databricksGenieUpdateSpaceTool, {
59+
...credentials,
60+
spaceId: 'space1',
61+
description,
62+
}).body ?? '{}'
63+
)
64+
65+
expect(update('')).toEqual({ description: '' })
66+
expect(update()).not.toHaveProperty('description')
67+
})
68+
69+
it.each([undefined, false])(
70+
'rejects cancellation without a job or explicit queued scope (%s)',
71+
(allQueuedRuns) => {
72+
expect(() =>
73+
prepareToolRequest(databricksTools.databricksCancelAllRunsTool, {
74+
...credentials,
75+
allQueuedRuns,
76+
})
77+
).toThrow('Provide jobId or set allQueuedRuns to true')
78+
}
79+
)
80+
81+
it.each([
82+
[{ jobId: 42 }, { job_id: 42 }],
83+
[{ allQueuedRuns: true }, { all_queued_runs: true }],
84+
[
85+
{ jobId: 42, allQueuedRuns: true },
86+
{ job_id: 42, all_queued_runs: true },
87+
],
88+
])('retains an explicit cancellation scope %j', (scope, body) => {
89+
const request = prepareToolRequest(databricksTools.databricksCancelAllRunsTool, {
90+
...credentials,
91+
...scope,
92+
})
93+
expect(JSON.parse(request.body ?? '{}')).toEqual(body)
94+
})
95+
})
96+
5197
describe('databricks workspace host allowlist', () => {
5298
it.each([
5399
'attacker.example.com',

‎apps/sim/tools/databricks/genie_update_space.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const genieUpdateSpaceTool: ToolConfig<
4747
type: 'string',
4848
required: false,
4949
visibility: 'user-or-llm',
50-
description: 'Description of the Genie space',
50+
description: 'Description of the Genie space; an empty string clears the description',
5151
},
5252
warehouseId: {
5353
type: 'string',
@@ -87,7 +87,7 @@ export const genieUpdateSpaceTool: ToolConfig<
8787
return filterUndefined({
8888
serialized_space: optionalDatabricksString(params.serializedSpace),
8989
title: optionalDatabricksString(params.title),
90-
description: optionalDatabricksString(params.description),
90+
description: params.description?.trim(),
9191
warehouse_id: optionalDatabricksString(params.warehouseId),
9292
etag: optionalDatabricksString(params.etag),
9393
parent_path: optionalDatabricksString(params.parentPath),

‎apps/sim/tools/databricks/types.ts‎

Lines changed: 187 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import type { OutputProperty, ToolFileData, ToolResponse } from '@/tools/types'
22

33
/** Base parameters shared by all Databricks tools */
4-
export interface DatabricksBaseParams {
4+
interface DatabricksBaseParams {
55
apiKey: string
66
host: string
77
}
@@ -1195,6 +1195,192 @@ export const JOB_TASK_OUTPUT = {
11951195
base_parameters: { type: 'object', description: 'Notebook parameter map' },
11961196
},
11971197
},
1198+
spark_python_task: {
1199+
type: 'object',
1200+
optional: true,
1201+
properties: {
1202+
python_file: { type: 'string', description: 'Python file path or cloud URI' },
1203+
source: { type: 'string', optional: true, description: 'WORKSPACE or GIT' },
1204+
parameters: { type: 'array', optional: true, items: { type: 'string' } },
1205+
},
1206+
},
1207+
spark_jar_task: {
1208+
type: 'object',
1209+
optional: true,
1210+
properties: {
1211+
main_class_name: { type: 'string', optional: true, description: 'JAR entry point class' },
1212+
parameters: { type: 'array', optional: true, items: { type: 'string' } },
1213+
jar_uri: { type: 'string', optional: true, description: 'Legacy JAR URI' },
1214+
},
1215+
},
1216+
spark_submit_task: {
1217+
type: 'object',
1218+
optional: true,
1219+
description: 'Legacy spark-submit task',
1220+
properties: {
1221+
parameters: { type: 'array', optional: true, items: { type: 'string' } },
1222+
},
1223+
},
1224+
python_wheel_task: {
1225+
type: 'object',
1226+
optional: true,
1227+
properties: {
1228+
package_name: { type: 'string', description: 'Python package name' },
1229+
entry_point: { type: 'string', description: 'Package entry point' },
1230+
parameters: { type: 'array', optional: true, items: { type: 'string' } },
1231+
named_parameters: {
1232+
type: 'object',
1233+
optional: true,
1234+
description: 'Named wheel parameters',
1235+
},
1236+
},
1237+
},
1238+
sql_task: {
1239+
type: 'object',
1240+
optional: true,
1241+
properties: {
1242+
warehouse_id: { type: 'string', description: 'SQL warehouse identifier' },
1243+
parameters: { type: 'object', optional: true, description: 'SQL parameters' },
1244+
query: {
1245+
type: 'object',
1246+
optional: true,
1247+
properties: { query_id: { type: 'string', description: 'Saved query identifier' } },
1248+
},
1249+
file: {
1250+
type: 'object',
1251+
optional: true,
1252+
properties: {
1253+
path: { type: 'string', description: 'SQL file path' },
1254+
source: { type: 'string', optional: true, description: 'WORKSPACE or GIT' },
1255+
},
1256+
},
1257+
alert: {
1258+
type: 'object',
1259+
optional: true,
1260+
properties: {
1261+
alert_id: { type: 'string', description: 'SQL alert identifier' },
1262+
subscriptions: { type: 'array', optional: true, items: { type: 'object' } },
1263+
pause_subscriptions: { type: 'boolean', optional: true },
1264+
},
1265+
},
1266+
dashboard: {
1267+
type: 'object',
1268+
optional: true,
1269+
properties: {
1270+
dashboard_id: { type: 'string', description: 'Legacy SQL dashboard identifier' },
1271+
subscriptions: { type: 'array', optional: true, items: { type: 'object' } },
1272+
custom_subject: { type: 'string', optional: true },
1273+
pause_subscriptions: { type: 'boolean', optional: true },
1274+
},
1275+
},
1276+
},
1277+
},
1278+
run_job_task: {
1279+
type: 'object',
1280+
optional: true,
1281+
properties: {
1282+
job_id: { type: 'number', description: 'Nested job identifier' },
1283+
job_parameters: { type: 'object', optional: true, description: 'Nested job parameters' },
1284+
pipeline_params: {
1285+
type: 'object',
1286+
optional: true,
1287+
description: 'Pipeline refresh settings',
1288+
},
1289+
},
1290+
},
1291+
pipeline_task: {
1292+
type: 'object',
1293+
optional: true,
1294+
properties: {
1295+
pipeline_id: { type: 'string', description: 'Pipeline identifier' },
1296+
parameters: { type: 'object', optional: true },
1297+
full_refresh: { type: 'boolean', optional: true },
1298+
refresh_selection: { type: 'array', optional: true, items: { type: 'string' } },
1299+
full_refresh_selection: { type: 'array', optional: true, items: { type: 'string' } },
1300+
reset_checkpoint_selection: { type: 'array', optional: true, items: { type: 'string' } },
1301+
refresh_flow_selection: { type: 'array', optional: true, items: { type: 'string' } },
1302+
},
1303+
},
1304+
dbt_task: {
1305+
type: 'object',
1306+
optional: true,
1307+
properties: {
1308+
commands: { type: 'array', items: { type: 'string' }, description: 'dbt commands' },
1309+
project_directory: { type: 'string', optional: true },
1310+
profiles_directory: { type: 'string', optional: true },
1311+
warehouse_id: { type: 'string', optional: true },
1312+
catalog: { type: 'string', optional: true },
1313+
schema: { type: 'string', optional: true },
1314+
source: { type: 'string', optional: true, description: 'WORKSPACE or GIT' },
1315+
},
1316+
},
1317+
condition_task: {
1318+
type: 'object',
1319+
optional: true,
1320+
properties: {
1321+
op: { type: 'string', description: 'Comparison operator' },
1322+
left: { type: 'string', description: 'Left operand' },
1323+
right: { type: 'string', description: 'Right operand' },
1324+
outcome: {
1325+
type: 'string',
1326+
optional: true,
1327+
description: 'Evaluation result: true or false',
1328+
},
1329+
},
1330+
},
1331+
for_each_task: {
1332+
type: 'object',
1333+
optional: true,
1334+
properties: {
1335+
inputs: {
1336+
type: 'string',
1337+
description: 'JSON array or parameter reference to iterate over',
1338+
},
1339+
concurrency: { type: 'number', optional: true },
1340+
task: { type: 'object', description: 'Nested task definition' },
1341+
},
1342+
},
1343+
clean_rooms_notebook_task: {
1344+
type: 'object',
1345+
optional: true,
1346+
properties: {
1347+
clean_room_name: { type: 'string' },
1348+
notebook_name: { type: 'string' },
1349+
etag: { type: 'string', optional: true },
1350+
notebook_base_parameters: { type: 'object', optional: true },
1351+
},
1352+
},
1353+
dashboard_task: {
1354+
type: 'object',
1355+
optional: true,
1356+
properties: {
1357+
dashboard_id: { type: 'string', optional: true },
1358+
warehouse_id: { type: 'string', optional: true },
1359+
subscription: { type: 'object', optional: true },
1360+
filters: { type: 'object', optional: true },
1361+
},
1362+
},
1363+
alert_task: {
1364+
type: 'object',
1365+
optional: true,
1366+
properties: {
1367+
alert_id: { type: 'string', optional: true },
1368+
warehouse_id: { type: 'string', optional: true },
1369+
workspace_path: { type: 'string', optional: true },
1370+
subscribers: { type: 'array', optional: true, items: { type: 'object' } },
1371+
},
1372+
},
1373+
power_bi_task: {
1374+
type: 'object',
1375+
optional: true,
1376+
properties: {
1377+
warehouse_id: { type: 'string', optional: true },
1378+
connection_resource_name: { type: 'string', optional: true },
1379+
refresh_after_update: { type: 'boolean', optional: true },
1380+
power_bi_model: { type: 'object', optional: true },
1381+
tables: { type: 'array', optional: true, items: { type: 'object' } },
1382+
},
1383+
},
11981384
state: {
11991385
type: 'object',
12001386
optional: true,

‎apps/sim/tools/generated/tool-metadata.ts‎

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

‎apps/sim/tools/generated/tool-outputs.ts‎

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

‎apps/sim/tools/index.test.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5916,6 +5916,23 @@ describe('Centralized Error Handling', () => {
59165916
return buildBlockExecutionError({ block: apiBlock, error: thrown as Error })
59175917
}
59185918

5919+
it.each(['<html>upstream gateway</html>', ''])(
5920+
'attributes a non-JSON 2xx body to the provider when a transform reads JSON (%j)',
5921+
async (body) => {
5922+
const originalTransform = tools.http_request.transformResponse
5923+
tools.http_request.transformResponse = async (response) => ({
5924+
success: true,
5925+
output: await response.json(),
5926+
})
5927+
try {
5928+
const blockError = await failApiBlock(new Response(body, { status: 200 }))
5929+
expect(classifyFailure(blockError)).toBe('third_party_server')
5930+
} finally {
5931+
tools.http_request.transformResponse = originalTransform
5932+
}
5933+
}
5934+
)
5935+
59195936
it.each([
59205937
[404, 'third_party_client'],
59215938
[503, 'third_party_server'],
@@ -5935,6 +5952,11 @@ describe('Centralized Error Handling', () => {
59355952
'third_party_client',
59365953
],
59375954
['a bug in the transform itself', new TypeError('data.channel is undefined'), 'internal'],
5955+
[
5956+
'a parser bug in the transform itself',
5957+
new SyntaxError('invalid transform configuration'),
5958+
'internal',
5959+
],
59385960
] as const)('attributes %s', async (_name, transformError, kind) => {
59395961
const originalTransform = tools.http_request.transformResponse
59405962
tools.http_request.transformResponse = async () => {

‎apps/sim/tools/index.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3086,7 +3086,15 @@ async function executeToolRequest(
30863086
headers: response.headers,
30873087
url: fullUrl,
30883088
body: response.body,
3089-
json: () => response.json(),
3089+
json: async () => {
3090+
try {
3091+
return await response.json()
3092+
} catch (error) {
3093+
throw error instanceof SyntaxError
3094+
? markFailureKind(error, 'third_party_server')
3095+
: error
3096+
}
3097+
},
30903098
text: () => response.text(),
30913099
arrayBuffer: () => response.arrayBuffer(),
30923100
blob: () => response.blob(),

0 commit comments

Comments
 (0)