Migrate dev server API from workflows to app-builder queries

sdkennedy2 · sdkennedy2 · commit 7e20d8781b8a · 2026-03-17T14:50:22.000-04:00
diff --git a/packages/plugins/apps/src/backend/vite/dev-server.test.ts b/packages/plugins/apps/src/backend/vite/dev-server.test.ts
@@ -14,7 +14,6 @@ jest.mock('vite', () => ({
     build: (...args) => mockViteBuild(...args),
 }));
 
-const WORKFLOW_ID = '380e7df1-729c-420c-b15e-a3b8e6347d49';
 const DD_SITE = 'datadoghq.com';
 
 const mockFunctions = [
@@ -134,13 +133,13 @@ describe('Dev Server Middleware', () => {
 
             // Mock the Datadog API via nock.
             const apiScope = nock(`https://${DD_SITE}`)
-                .post(`/api/v2/workflows/${WORKFLOW_ID}/single_action_runs`)
-                .reply(200, { data: { id: 'exec-123' } })
-                .get(`/api/v2/workflows/${WORKFLOW_ID}/single_action_runs/exec-123`)
+                .post('/api/v2/app-builder/queries/preview-async')
+                .reply(200, { data: { id: 'receipt-123' } })
+                .get('/api/v2/app-builder/queries/execution-long-polling/receipt-123')
                 .reply(200, {
                     data: {
                         attributes: {
-                            state: 'SUCCEEDED',
+                            done: true,
                             outputs: { result: 'hello' },
                         },
                     },
@@ -156,8 +155,8 @@ describe('Dev Server Middleware', () => {
             middleware(req, res, next);
             expect(next).not.toHaveBeenCalled();
 
-            // Wait for async handler + polling to complete.
-            await new Promise((resolve) => setTimeout(resolve, 2000));
+            // Wait for async handler to complete.
+            await new Promise((resolve) => setTimeout(resolve, 100));
 
             expect(res.statusCode).toBe(200);
             const body = JSON.parse(res.getBody());
@@ -263,7 +262,7 @@ describe('Dev Server Middleware', () => {
             mockViteBuild.mockResolvedValue(mockBuildResult('// code'));
 
             nock(`https://${DD_SITE}`)
-                .post(`/api/v2/workflows/${WORKFLOW_ID}/single_action_runs`)
+                .post('/api/v2/app-builder/queries/preview-async')
                 .reply(403, 'Forbidden');
 
             const req = createMockRequest('/__dd/executeAction', {
@@ -290,11 +289,11 @@ describe('Dev Server Middleware', () => {
                     'DD-APPLICATION-KEY': 'test-app-key',
                 },
             })
-                .post(`/api/v2/workflows/${WORKFLOW_ID}/single_action_runs`)
-                .reply(200, { data: { id: 'exec-1' } })
-                .get(`/api/v2/workflows/${WORKFLOW_ID}/single_action_runs/exec-1`)
+                .post('/api/v2/app-builder/queries/preview-async')
+                .reply(200, { data: { id: 'receipt-1' } })
+                .get('/api/v2/app-builder/queries/execution-long-polling/receipt-1')
                 .reply(200, {
-                    data: { attributes: { state: 'SUCCEEDED', outputs: { value: 42 } } },
+                    data: { attributes: { done: true, outputs: { value: 42 } } },
                 });
 
             const req = createMockRequest('/__dd/executeAction', {
@@ -304,13 +303,68 @@ describe('Dev Server Middleware', () => {
             const res = createMockResponse();
 
             middleware(req, res, jest.fn());
-            await new Promise((resolve) => setTimeout(resolve, 2000));
+            await new Promise((resolve) => setTimeout(resolve, 100));
 
             expect(res.statusCode).toBe(200);
             const body = JSON.parse(res.getBody());
             expect(body.success).toBe(true);
             expect(body.result).toEqual({ value: 42 });
             expect(apiScope.isDone()).toBe(true);
         });
+
+        test('Should handle errors array from long-polling endpoint', async () => {
+            mockViteBuild.mockResolvedValue(mockBuildResult('// code'));
+
+            nock(`https://${DD_SITE}`)
+                .post('/api/v2/app-builder/queries/preview-async')
+                .reply(200, { data: { id: 'receipt-err' } })
+                .get('/api/v2/app-builder/queries/execution-long-polling/receipt-err')
+                .reply(200, {
+                    errors: [{ title: 'ExecutionFailed', detail: 'Script threw an error' }],
+                });
+
+            const req = createMockRequest('/__dd/executeAction', {
+                functionName: 'greet',
+                args: [],
+            });
+            const res = createMockResponse();
+
+            middleware(req, res, jest.fn());
+            await new Promise((resolve) => setTimeout(resolve, 100));
+
+            expect(res.statusCode).toBe(500);
+            const body = JSON.parse(res.getBody());
+            expect(body.success).toBe(false);
+            expect(body.error).toContain('Script threw an error');
+        });
+
+        test('Should retry when long-poll returns done: false', async () => {
+            mockViteBuild.mockResolvedValue(mockBuildResult('// code'));
+
+            const apiScope = nock(`https://${DD_SITE}`)
+                .post('/api/v2/app-builder/queries/preview-async')
+                .reply(200, { data: { id: 'receipt-retry' } })
+                .get('/api/v2/app-builder/queries/execution-long-polling/receipt-retry')
+                .reply(200, { data: { attributes: { done: false } } })
+                .get('/api/v2/app-builder/queries/execution-long-polling/receipt-retry')
+                .reply(200, {
+                    data: { attributes: { done: true, outputs: { ok: true } } },
+                });
+
+            const req = createMockRequest('/__dd/executeAction', {
+                functionName: 'greet',
+                args: [],
+            });
+            const res = createMockResponse();
+
+            middleware(req, res, jest.fn());
+            await new Promise((resolve) => setTimeout(resolve, 100));
+
+            expect(res.statusCode).toBe(200);
+            const body = JSON.parse(res.getBody());
+            expect(body.success).toBe(true);
+            expect(body.result).toEqual({ ok: true });
+            expect(apiScope.isDone()).toBe(true);
+        });
     });
 });
diff --git a/packages/plugins/apps/src/backend/vite/dev-server.ts b/packages/plugins/apps/src/backend/vite/dev-server.ts
@@ -5,6 +5,7 @@
 /* eslint-disable no-await-in-loop */
 
 import type { Logger } from '@dd/core/types';
+import { randomUUID } from 'crypto';
 import type { IncomingMessage, ServerResponse } from 'http';
 import type { RollupOutput } from 'rollup';
 
@@ -16,9 +17,6 @@ import { generateDevVirtualEntryContent } from '../virtual-entry';
 // generate empty chunks when used as an input entry.
 const DEV_VIRTUAL_PREFIX = 'virtual:dd-backend-dev:';
 
-// Hardcoded workflow ID for the development execution environment.
-const WORKFLOW_ID = '380e7df1-729c-420c-b15e-a3b8e6347d49';
-
 interface ExecuteActionRequest {
     functionName: string;
     args?: unknown[];
@@ -136,14 +134,15 @@ async function bundleBackendFunction(
 }
 
 /**
- * Execute a script via Datadog's single_action_runs API.
+ * Execute a script via Datadog's app-builder queries API.
  */
 async function executeScriptViaDatadog(
     scriptBody: string,
+    functionName: string,
     auth: AuthConfig,
     log: Logger,
 ): Promise<unknown> {
-    const endpoint = `https://${auth.site}/api/v2/workflows/${WORKFLOW_ID}/single_action_runs`;
+    const endpoint = `https://${auth.site}/api/v2/app-builder/queries/preview-async`;
 
     log.debug(`Calling Datadog API: ${endpoint}`);
 
@@ -156,10 +155,21 @@ async function executeScriptViaDatadog(
         },
         body: JSON.stringify({
             data: {
-                type: 'single_action_runs',
+                type: 'queries',
                 attributes: {
-                    actionId: 'com.datadoghq.datatransformation.jsFunctionWithActions',
-                    inputs: { script: scriptBody, context: {} },
+                    query: {
+                        id: randomUUID(),
+                        name: functionName,
+                        type: 'action',
+                        properties: {
+                            spec: {
+                                fqn: 'com.datadoghq.datatransformation.jsFunctionWithActions',
+                                inputs: { script: scriptBody },
+                            },
+                            onlyTriggerManually: true,
+                        },
+                    },
+                    template_params: {},
                 },
             },
         }),
@@ -171,31 +181,32 @@ async function executeScriptViaDatadog(
     }
 
     const initialResult = (await response.json()) as { data?: { id?: string } };
-    const executionId = initialResult.data?.id;
+    const receiptId = initialResult.data?.id;
 
-    if (!executionId) {
-        throw new Error('No execution ID returned from Datadog API');
+    if (!receiptId) {
+        throw new Error('No receipt ID returned from Datadog API');
     }
 
-    log.debug(`Action started with ID: ${executionId}`);
+    log.debug(`Query execution started with receipt: ${receiptId}`);
 
-    return pollActionExecution(executionId, auth, log);
+    return pollQueryExecution(receiptId, auth, log);
 }
 
 /**
- * Poll Datadog API until the action execution completes or times out.
+ * Long-poll Datadog API until the query execution completes or times out.
+ * The server holds the connection open until the result is ready or its own timeout expires.
  */
-async function pollActionExecution(
-    executionId: string,
+async function pollQueryExecution(
+    receiptId: string,
     auth: AuthConfig,
     log: Logger,
 ): Promise<unknown> {
-    const endpoint = `https://${auth.site}/api/v2/workflows/${WORKFLOW_ID}/single_action_runs/${executionId}`;
-    const maxAttempts = 30;
-    const pollInterval = 1000;
+    const endpoint = `https://${auth.site}/api/v2/app-builder/queries/execution-long-polling/${receiptId}`;
+    // Each long-poll request waits server-side (~30s). Max retries provides a safety net.
+    const maxRetries = 10;
 
-    for (let attempt = 0; attempt < maxAttempts; attempt++) {
-        log.debug(`Polling attempt ${attempt + 1}/${maxAttempts}...`);
+    for (let attempt = 0; attempt < maxRetries; attempt++) {
+        log.debug(`Long-poll attempt ${attempt + 1}/${maxRetries}...`);
 
         const response = await fetch(endpoint, {
             method: 'GET',
@@ -212,25 +223,27 @@ async function pollActionExecution(
         }
 
         const result = (await response.json()) as {
-            data?: { attributes?: { state?: string; outputs?: unknown; error?: unknown } };
+            data?: { attributes?: { done?: boolean; outputs?: unknown } };
+            errors?: Array<{ detail?: string; title?: string }>;
         };
-        const state = result.data?.attributes?.state;
-
-        log.debug(`Execution state: ${state}`);
 
-        if (state === 'SUCCEEDED') {
-            return result.data!.attributes!.outputs;
+        // Check for error responses.
+        if (result.errors?.length) {
+            const details = result.errors.map((e) => e.detail || e.title).join('; ');
+            throw new Error(`Query execution failed: ${details}`);
         }
 
-        if (state === 'FAILED' || state === 'EXECUTION_FAILED') {
-            const errorDetails = result.data?.attributes?.error || result.data?.attributes;
-            throw new Error(`Action execution failed: ${JSON.stringify(errorDetails)}`);
+        const attrs = result.data?.attributes;
+        log.debug(`Long-poll response, done: ${attrs?.done}`);
+
+        if (attrs?.done) {
+            return attrs.outputs;
         }
 
-        await new Promise((resolve) => setTimeout(resolve, pollInterval));
+        // done === false means server-side long-poll timed out; retry immediately.
     }
 
-    throw new Error('Action execution timed out');
+    throw new Error('Query execution timed out');
 }
 
 /**
@@ -306,7 +319,7 @@ async function handleExecuteAction(
 
         const scriptBody = await bundleBackendFunction(func, args, projectRoot, log);
 
-        const result = await executeScriptViaDatadog(scriptBody, auth, log);
+        const result = await executeScriptViaDatadog(scriptBody, functionName, auth, log);
 
         res.statusCode = 200;
         res.setHeader('Content-Type', 'application/json');
