Skip to content

Commit 5848bd5

Browse files
waleedlatif1waleedlatif1
authored
fix(databricks): preserve mutation intent and complete task outputs (#9176)
* fix(databricks): preserve mutation intent and complete task outputs * fix(databricks): preserve blank input and validation semantics
1 parent ca43d9d commit 5848bd5

9 files changed

Lines changed: 608 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: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { filterUndefined } from '@sim/utils/object'
2+
import { UserFailure } from '@/lib/core/errors/user-failure'
23
import type {
34
DatabricksCancelAllRunsParams,
45
DatabricksSuccessResponse,
@@ -24,7 +25,7 @@ export const cancelAllRunsTool: ToolConfig<
2425
type: 'number',
2526
required: false,
2627
visibility: 'user-or-llm',
27-
description: 'Job identifier',
28+
description: 'Job identifier; required unless allQueuedRuns is true',
2829
},
2930
allQueuedRuns: {
3031
type: 'boolean',
@@ -41,6 +42,11 @@ export const cancelAllRunsTool: ToolConfig<
4142
'Content-Type': 'application/json',
4243
}),
4344
body: (params) => {
45+
if (params.jobId == null && params.allQueuedRuns !== true) {
46+
throw new UserFailure(
47+
'Provide jobId or set allQueuedRuns to true to cancel queued workspace runs'
48+
)
49+
}
4450
return filterUndefined({
4551
job_id: params.jobId,
4652
all_queued_runs: params.allQueuedRuns,

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { jsonResponse } from '@sim/testing/helpers/http'
33
import { partialToolRegistry } from '@sim/testing/mocks/tool-registry.mock'
44
import { getErrorMessage } from '@sim/utils/errors'
55
import { describe, expect, it, vi } from 'vitest'
6+
import { adoptToolFailure, classifyFailure } from '@/lib/core/errors/failure-log'
67

78
vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
89

@@ -11,6 +12,7 @@ import { executeSqlTool } from '@/tools/databricks/execute_sql'
1112
import { getRunTool } from '@/tools/databricks/get_run'
1213
import { executeTool } from '@/tools/index'
1314
import { tools } from '@/tools/registry'
15+
import { prepareToolRequest } from '@/tools/request-transport'
1416

1517
/** Registers only this service's configs in the global registry mock; the full one is ~6,000 modules. */
1618
Object.assign(tools, partialToolRegistry(databricksTools))
@@ -48,6 +50,53 @@ const HOST_ALLOWLIST_ERROR =
4850

4951
const DATABRICKS_TOOL_IDS = Object.keys(tools).filter((id) => id.startsWith('databricks_'))
5052

53+
describe('Databricks mutation intent', () => {
54+
const credentials = { host: 'dbc-test.cloud.databricks.com', apiKey: 'test-token' }
55+
56+
it('distinguishes clearing a Genie description from leaving it unchanged', () => {
57+
const update = (description?: string) =>
58+
JSON.parse(
59+
prepareToolRequest(databricksTools.databricksGenieUpdateSpaceTool, {
60+
...credentials,
61+
spaceId: 'space1',
62+
description,
63+
}).body ?? '{}'
64+
)
65+
66+
expect(update('')).toEqual({ description: '' })
67+
expect(update()).not.toHaveProperty('description')
68+
expect(update(' \n\t ')).not.toHaveProperty('description')
69+
})
70+
71+
it.each([undefined, false])(
72+
'rejects cancellation without a job or explicit queued scope (%s)',
73+
async (allQueuedRuns) => {
74+
const result = await executeTool('databricks_cancel_all_runs', {
75+
...credentials,
76+
allQueuedRuns,
77+
})
78+
expect(result.success).toBe(false)
79+
expect(result.error).toContain('Provide jobId or set allQueuedRuns to true')
80+
expect(classifyFailure(adoptToolFailure(new Error(result.error), result))).toBe('user')
81+
}
82+
)
83+
84+
it.each([
85+
[{ jobId: 42 }, { job_id: 42 }],
86+
[{ allQueuedRuns: true }, { all_queued_runs: true }],
87+
[
88+
{ jobId: 42, allQueuedRuns: true },
89+
{ job_id: 42, all_queued_runs: true },
90+
],
91+
])('retains an explicit cancellation scope %j', (scope, body) => {
92+
const request = prepareToolRequest(databricksTools.databricksCancelAllRunsTool, {
93+
...credentials,
94+
...scope,
95+
})
96+
expect(JSON.parse(request.body ?? '{}')).toEqual(body)
97+
})
98+
})
99+
51100
describe('databricks workspace host allowlist', () => {
52101
it.each([
53102
'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 === '' ? '' : optionalDatabricksString(params.description),
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)